diff --git a/backends/cuda/runtime/cuda_allocator.cpp b/backends/cuda/runtime/cuda_allocator.cpp index 4c7d6aec288..68f5d0c6650 100644 --- a/backends/cuda/runtime/cuda_allocator.cpp +++ b/backends/cuda/runtime/cuda_allocator.cpp @@ -12,6 +12,13 @@ #include #include +#if !defined(EXECUTORCH_USE_HIP) +#include +#include +#include +#include +#endif + namespace executorch::backends::cuda { using executorch::runtime::Error; @@ -21,6 +28,101 @@ using executorch::runtime::etensor::DeviceType; namespace { +#if !defined(EXECUTORCH_USE_HIP) +// The stream ordered allocator hands physical memory back to the driver +// whenever a synchronization observes a pending free, so with the default +// release threshold of zero a pool is emptied repeatedly during one inference +// and every allocation has to map memory again, which measured three orders of +// magnitude slower on an embedded board. +// +// The delegate allocates from a pool it creates rather than the device default +// pool, because the default one is shared with every other user of the async +// allocator in this process. Raising a threshold there caps what that user's +// cache may keep, and trimming it on teardown throws their cached blocks away. +// Owning the pool means the threshold and the trim only ever affect this +// backend, with no attempt to remember and restore somebody else's setting. +constexpr uint64_t kMemPoolReleaseThreshold = UINT64_MAX; + +struct MemPoolState { + std::mutex mutex; + std::unordered_map pools; +}; + +MemPoolState& mem_pool_state() { + static MemPoolState state; + return state; +} + +// Resolves the "current device" sentinel that callers are allowed to pass. +// Returns a negative value when the device cannot be determined. +int resolve_device(DeviceIndex index) { + if (index >= 0) { + return static_cast(index); + } + int current = 0; + const cudaError_t err = cudaGetDevice(¤t); + if (err != cudaSuccess) { + ET_LOG( + Error, + "cudaGetDevice failed: %s. Using pool defaults.", + cudaGetErrorString(err)); + (void)cudaGetLastError(); + return -1; + } + return current; +} + +// The pool this backend allocates from on a device, creating it on first use. +// Returns nullptr when the pool cannot be created, in which case the caller +// falls back to the device default pool and only loses speed. +cudaMemPool_t mem_pool_for(int device) { + auto& state = mem_pool_state(); + const std::lock_guard lock(state.mutex); + const auto it = state.pools.find(device); + if (it != state.pools.end()) { + return it->second; + } + + cudaMemPoolProps props{}; + props.allocType = cudaMemAllocationTypePinned; + props.handleTypes = cudaMemHandleTypeNone; + props.location.type = cudaMemLocationTypeDevice; + props.location.id = device; + + cudaMemPool_t pool = nullptr; + cudaError_t err = cudaMemPoolCreate(&pool, &props); + if (err != cudaSuccess) { + ET_LOG( + Error, + "cudaMemPoolCreate failed for device %d: %s. Using the default pool.", + device, + cudaGetErrorString(err)); + (void)cudaGetLastError(); + // Recorded so a permanent failure is not retried and re-logged on every + // allocation. + state.pools.emplace(device, nullptr); + return nullptr; + } + + uint64_t threshold = kMemPoolReleaseThreshold; + err = cudaMemPoolSetAttribute( + pool, cudaMemPoolAttrReleaseThreshold, &threshold); + if (err != cudaSuccess) { + ET_LOG( + Error, + "Setting the pool release threshold failed for device %d: %s. Keeping " + "the pool at its defaults.", + device, + cudaGetErrorString(err)); + (void)cudaGetLastError(); + } + + state.pools.emplace(device, pool); + return pool; +} + +#endif // !EXECUTORCH_USE_HIP + Error copy_impl( void* dst, const void* src, @@ -303,7 +405,20 @@ Result CudaAllocator::allocate_async( DeviceIndex index, cudaStream_t stream) { void* ptr = nullptr; - cudaError_t err = cudaMallocAsync(&ptr, nbytes, stream); + cudaError_t err = cudaErrorInvalidValue; +#if defined(EXECUTORCH_USE_HIP) + err = cudaMallocAsync(&ptr, nbytes, stream); +#else + // Allocating from this backend's own pool keeps its retained memory out of + // the device default pool, which other users of the async allocator share. + const int device = resolve_device(index); + cudaMemPool_t pool = device >= 0 ? mem_pool_for(device) : nullptr; + if (pool != nullptr) { + err = cudaMallocFromPoolAsync(&ptr, nbytes, pool, stream); + } else { + err = cudaMallocAsync(&ptr, nbytes, stream); + } +#endif if (err != cudaSuccess) { ET_LOG( Error, @@ -313,6 +428,7 @@ Result CudaAllocator::allocate_async( static_cast(index)); return Error::MemoryAllocationFailed; } + return ptr; } @@ -323,6 +439,7 @@ void CudaAllocator::deallocate_async( if (ptr == nullptr) { return; } + cudaError_t err = cudaFreeAsync(ptr, stream); if (err != cudaSuccess) { ET_LOG( @@ -331,7 +448,69 @@ void CudaAllocator::deallocate_async( cudaGetErrorString(err), ptr, static_cast(index)); + return; + } +} + +cudaMemPool_t CudaAllocator::pool_for_device(DeviceIndex index) { +#if defined(EXECUTORCH_USE_HIP) + (void)index; + return nullptr; +#else + const int device = resolve_device(index); + if (device < 0) { + return nullptr; + } + auto& state = mem_pool_state(); + const std::lock_guard lock(state.mutex); + const auto it = state.pools.find(device); + return it == state.pools.end() ? nullptr : it->second; +#endif +} + +void CudaAllocator::release_cached_memory(DeviceIndex index) { +#if defined(EXECUTORCH_USE_HIP) + (void)index; +#else + std::vector> targets; + { + auto& state = mem_pool_state(); + const std::lock_guard lock(state.mutex); + if (index >= 0) { + const auto it = state.pools.find(static_cast(index)); + if (it == state.pools.end()) { + return; + } + targets.emplace_back(it->first, it->second); + } else { + // A caller asking for everything gets every pool this backend created, + // not whichever device the calling thread happens to be current on, since + // the delegate that ran is often not on that device. + targets.assign(state.pools.begin(), state.pools.end()); + } + } + + // The pools stay in the map. Trimming empties one without invalidating it, so + // a later load reuses it rather than paying to create it again. + for (const auto& [device, pool] : targets) { + if (pool == nullptr) { + continue; + } + // Only frees the driver has already observed can be released, so a caller + // that has not synchronized gets less back. That is why this does not + // synchronize itself: a device wide barrier here would wait on every stream + // on the device, including work this backend never queued. + const cudaError_t err = cudaMemPoolTrimTo(pool, 0); + if (err != cudaSuccess) { + ET_LOG( + Error, + "cudaMemPoolTrimTo failed for device %d: %s.", + device, + cudaGetErrorString(err)); + (void)cudaGetLastError(); + } } +#endif } Error CudaAllocator::memcpy_async( diff --git a/backends/cuda/runtime/cuda_allocator.h b/backends/cuda/runtime/cuda_allocator.h index b0a76a51f6d..4c802a7e00b 100644 --- a/backends/cuda/runtime/cuda_allocator.h +++ b/backends/cuda/runtime/cuda_allocator.h @@ -68,6 +68,43 @@ class CudaAllocator final : public executorch::runtime::DeviceAllocator { executorch::runtime::etensor::DeviceIndex index, cudaStream_t stream); + /** + * Return memory this backend's device pool is holding for reuse back to the + * driver. + * + * The pool keeps freed memory so that repeated allocations do not have to map + * it again, which is what makes delegate execution cheap, so a long-lived + * process should call this once its work on the device is finished. Only + * frees the driver has already observed can be released, so a caller that has + * not synchronized simply gets less back. Allocations that are still live are + * unaffected either way. + * + * The pool belongs to this backend rather than being the device default pool, + * so this never affects memory another user of the async allocator is + * holding. + * + * Does nothing on ROCm, where the pool calls have no equivalent. + * + * @param index Device to release on, or a negative value to release every + * device this backend has allocated on. A delegate is often torn down + * from a thread that is not current on the device it ran on, so releasing + * only the current device would leave that memory held. + */ + static void release_cached_memory( + executorch::runtime::etensor::DeviceIndex index); + + /** + * The memory pool this backend allocates from on a device, or nullptr if it + * has not allocated there or the pool could not be created. + * + * Exposed so a test can observe what the pool is holding, which is not + * visible through the device default pool. + * + * @param index Device to query, or a negative value for the current one. + */ + static cudaMemPool_t pool_for_device( + executorch::runtime::etensor::DeviceIndex index); + /** * Copy memory asynchronously on the given CUDA stream. * Supports H2D, D2H, and D2D based on src/dst device types. diff --git a/backends/cuda/runtime/cuda_backend.cpp b/backends/cuda/runtime/cuda_backend.cpp index 29cba8b5ada..39499638f78 100644 --- a/backends/cuda/runtime/cuda_backend.cpp +++ b/backends/cuda/runtime/cuda_backend.cpp @@ -470,6 +470,8 @@ class ET_EXPERIMENTAL CudaBackend final mutable_state_note_handle(handle); + live_handles_.fetch_add(1, std::memory_order_acq_rel); + return (DelegateHandle*)handle; // Return the handle post-processing } @@ -892,6 +894,14 @@ class ET_EXPERIMENTAL CudaBackend final } delete handle; + + // The allocator lets the device pool keep freed memory so that repeated + // delegate execution does not pay to map it again. Nothing is running on + // this backend once the last handle is gone, so hand that memory back + // rather than hold it for the life of the process. + if (live_handles_.fetch_sub(1, std::memory_order_acq_rel) == 1) { + CudaAllocator::release_cached_memory(-1); + } } private: @@ -906,6 +916,10 @@ class ET_EXPERIMENTAL CudaBackend final mutable std::mutex cuda_stream_mutex_; std::shared_ptr shared_cuda_stream_ = nullptr; + // Delegates alive right now. The device memory pool is shared, so it can only + // be released once none of them are left. + mutable std::atomic live_handles_{0}; + // Whether to enable cross-method caching for legacy dense-blob artifacts. // Toggled by the kWeightSharingAcrossMethods runtime backend option. Default // OFF; versioned FQN artifacts do not consult this option. diff --git a/backends/cuda/runtime/test/test_cuda_allocator.cpp b/backends/cuda/runtime/test/test_cuda_allocator.cpp index a872099c4f7..02ba9c909a0 100644 --- a/backends/cuda/runtime/test/test_cuda_allocator.cpp +++ b/backends/cuda/runtime/test/test_cuda_allocator.cpp @@ -181,3 +181,134 @@ TEST_F(CudaAllocatorTest, CopyDeviceToHostOnMissingDeviceFails) { a.deallocate(dptr, 0); } + +// The pool attributes these exercise have no HIP equivalent in the +// compatibility header, and the allocator's pool code is compiled out on ROCm +// for the same reason, so there is nothing to test there. +#if !defined(EXECUTORCH_USE_HIP) + +namespace { +uint64_t reserved_bytes(cudaMemPool_t pool) { + uint64_t reserved = 0; + EXPECT_EQ( + cudaMemPoolGetAttribute( + pool, cudaMemPoolAttrReservedMemCurrent, &reserved), + cudaSuccess); + return reserved; +} +} // namespace + +// The delegate allocates from a pool it owns, so its retained memory must not +// land in the device default pool that other users of the async allocator +// share. +TEST_F(CudaAllocatorTest, AllocatesFromItsOwnPool) { + cudaStream_t stream; + ASSERT_EQ(cudaStreamCreate(&stream), cudaSuccess); + + constexpr size_t kBytes = 8u << 20; + auto res = CudaAllocator::allocate_async(kBytes, 0, stream); + ASSERT_TRUE(res.ok()); + + cudaMemPool_t owned = CudaAllocator::pool_for_device(0); + ASSERT_NE(owned, nullptr) << "the allocator should have created its own pool"; + cudaMemPool_t default_pool = nullptr; + ASSERT_EQ(cudaDeviceGetMemPool(&default_pool, 0), cudaSuccess); + EXPECT_NE(owned, default_pool) << "the pool must not be the device default"; + + // The live block is reserved in the owned pool, which is what identifies it + // as the pool actually serving this allocation. + EXPECT_GE(reserved_bytes(owned), kBytes); + + CudaAllocator::deallocate_async(res.get(), 0, stream); + ASSERT_EQ(cudaStreamSynchronize(stream), cudaSuccess); + ASSERT_EQ(cudaStreamDestroy(stream), cudaSuccess); +} + +// Freed memory is kept so repeated allocation stays cheap, which means a plain +// free no longer shrinks the pool. Without an explicit release a long lived +// process would hold that memory after every program was gone. +TEST_F(CudaAllocatorTest, ReleaseCachedMemoryReturnsPoolMemory) { + cudaStream_t stream; + ASSERT_EQ(cudaStreamCreate(&stream), cudaSuccess); + + constexpr size_t kBytes = 8u << 20; + auto res = CudaAllocator::allocate_async(kBytes, 0, stream); + ASSERT_TRUE(res.ok()); + CudaAllocator::deallocate_async(res.get(), 0, stream); + ASSERT_EQ(cudaStreamSynchronize(stream), cudaSuccess); + + cudaMemPool_t owned = CudaAllocator::pool_for_device(0); + ASSERT_NE(owned, nullptr); + // Freed and synchronized, and still held, which is the point of the change. + ASSERT_GT(reserved_bytes(owned), 0u) + << "the pool should hold the freed block for reuse"; + + CudaAllocator::release_cached_memory(0); + + EXPECT_EQ(reserved_bytes(owned), 0u) + << "released memory should go back to the driver"; + + ASSERT_EQ(cudaStreamDestroy(stream), cudaSuccess); +} + +// Releasing must not disturb allocations that are still in use. +TEST_F(CudaAllocatorTest, ReleaseCachedMemoryKeepsLiveAllocations) { + cudaStream_t stream; + ASSERT_EQ(cudaStreamCreate(&stream), cudaSuccess); + + // Large enough that the two blocks land in separate driver reservations. At a + // few megabytes they share one, so nothing can be released while either is + // live. + constexpr size_t kBytes = 64u << 20; + auto live = CudaAllocator::allocate_async(kBytes, 0, stream); + ASSERT_TRUE(live.ok()); + auto temp = CudaAllocator::allocate_async(kBytes, 0, stream); + ASSERT_TRUE(temp.ok()); + CudaAllocator::deallocate_async(temp.get(), 0, stream); + ASSERT_EQ(cudaStreamSynchronize(stream), cudaSuccess); + + cudaMemPool_t owned = CudaAllocator::pool_for_device(0); + ASSERT_NE(owned, nullptr); + const uint64_t before = reserved_bytes(owned); + + CudaAllocator::release_cached_memory(0); + + // The freed block goes back and the live one stays reserved, so the pool + // gives up only what is not in use. + const uint64_t after = reserved_bytes(owned); + EXPECT_LT(after, before) << "the freed block should have been released"; + EXPECT_GE(after, kBytes) << "the live block must still be reserved"; + + EXPECT_EQ(cudaMemsetAsync(live.get(), 0, kBytes, stream), cudaSuccess); + EXPECT_EQ(cudaStreamSynchronize(stream), cudaSuccess); + + CudaAllocator::deallocate_async(live.get(), 0, stream); + ASSERT_EQ(cudaStreamSynchronize(stream), cudaSuccess); + ASSERT_EQ(cudaStreamDestroy(stream), cudaSuccess); +} + +// A negative index means whichever device is current, which the allocator +// resolves rather than passing on to the driver. +TEST_F(CudaAllocatorTest, ReleaseCachedMemoryAcceptsCurrentDeviceSentinel) { + cudaStream_t stream; + ASSERT_EQ(cudaStreamCreate(&stream), cudaSuccess); + + constexpr size_t kBytes = 8u << 20; + auto res = CudaAllocator::allocate_async(kBytes, 0, stream); + ASSERT_TRUE(res.ok()); + CudaAllocator::deallocate_async(res.get(), 0, stream); + ASSERT_EQ(cudaStreamSynchronize(stream), cudaSuccess); + + cudaMemPool_t owned = CudaAllocator::pool_for_device(-1); + ASSERT_NE(owned, nullptr) << "the sentinel should resolve to this device"; + ASSERT_GT(reserved_bytes(owned), 0u) + << "the pool should hold the freed block for reuse"; + + CudaAllocator::release_cached_memory(-1); + + EXPECT_EQ(reserved_bytes(owned), 0u); + + ASSERT_EQ(cudaStreamDestroy(stream), cudaSuccess); +} + +#endif // !EXECUTORCH_USE_HIP diff --git a/extension/cuda/runtime_api.h b/extension/cuda/runtime_api.h index bae5c6a79bf..af404423548 100644 --- a/extension/cuda/runtime_api.h +++ b/extension/cuda/runtime_api.h @@ -18,6 +18,7 @@ using cudaError_t = hipError_t; using cudaGraph_t = hipGraph_t; using cudaGraphExec_t = hipGraphExec_t; using cudaMemcpyKind = hipMemcpyKind; +using cudaMemPool_t = hipMemPool_t; using cudaMemoryType = hipMemoryType; using cudaStreamCaptureMode = hipStreamCaptureMode; using cudaStream_t = hipStream_t;