diff --git a/backends/cuda/runtime/cuda_allocator.cpp b/backends/cuda/runtime/cuda_allocator.cpp index 4c7d6aec288..6d046764953 100644 --- a/backends/cuda/runtime/cuda_allocator.cpp +++ b/backends/cuda/runtime/cuda_allocator.cpp @@ -12,6 +12,12 @@ #include #include +#if !defined(EXECUTORCH_USE_HIP) +#include +#include +#include +#endif + namespace executorch::backends::cuda { using executorch::runtime::Error; @@ -21,6 +27,102 @@ 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 the threshold there would make that shared +// pool hold on to memory on their behalf, and trimming it on teardown would +// throw 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,16 +405,50 @@ Result CudaAllocator::allocate_async( DeviceIndex index, cudaStream_t stream) { void* ptr = nullptr; - cudaError_t err = cudaMallocAsync(&ptr, nbytes, stream); + cudaError_t err; + // Named for the log below, so a failure points at the call that actually ran. + const char* allocator_name = "cudaMallocAsync"; + int log_device = static_cast(index); +#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. + // + // The pool has to belong to the device the stream runs on. A pool from + // another device returns a pointer that stream cannot touch, and the failure + // surfaces later as an illegal access rather than here. The caller's index + // and the stream can disagree, so the plain async allocation is used unless + // the index names the device that is current for this stream. + const int device = resolve_device(index); + int stream_device = -1; + if (cudaGetDevice(&stream_device) != cudaSuccess) { + (void)cudaGetLastError(); + stream_device = -1; + } + cudaMemPool_t pool = + (device >= 0 && device == stream_device) ? mem_pool_for(device) : nullptr; + if (device >= 0) { + log_device = device; + } + if (pool != nullptr) { + allocator_name = "cudaMallocFromPoolAsync"; + err = cudaMallocFromPoolAsync(&ptr, nbytes, pool, stream); + } else { + err = cudaMallocAsync(&ptr, nbytes, stream); + } +#endif if (err != cudaSuccess) { ET_LOG( Error, - "cudaMallocAsync failed: %s (requested %zu bytes on device %d)", + "%s failed: %s (requested %zu bytes on device %d)", + allocator_name, cudaGetErrorString(err), nbytes, - static_cast(index)); + log_device); return Error::MemoryAllocationFailed; } + return ptr; } @@ -323,6 +459,7 @@ void CudaAllocator::deallocate_async( if (ptr == nullptr) { return; } + cudaError_t err = cudaFreeAsync(ptr, stream); if (err != cudaSuccess) { ET_LOG( @@ -334,6 +471,84 @@ void CudaAllocator::deallocate_async( } } +#if !defined(EXECUTORCH_USE_HIP) +cudaMemPool_t CudaAllocator::pool_for_device(DeviceIndex index) { + 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 // !EXECUTORCH_USE_HIP + +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) { + // Only frees the driver has already observed can be released, and a free + // still pending gives back nothing rather than less. The backend waits on + // its stream during teardown for that reason, before dropping it. This + // cannot wait on anything itself: it holds no stream, and a caller + // reaching it directly is responsible for having synchronized. + const cudaError_t err = cudaMemPoolTrimTo(pool, 0); + if (err != cudaSuccess) { + ET_LOG( + Error, + "cudaMemPoolTrimTo failed for device %d: %s.", + device, + cudaGetErrorString(err)); + (void)cudaGetLastError(); + } + } + + // Memory a method allocated while its CUDA graph was being captured belongs + // to the device graph pool, which the pool trim cannot reach, so a + // graph-enabled method would otherwise hold its footprint for the life of + // the process. Outside the branch above because graph memory is a device + // resource and exists whether or not this backend has a pool here. + // + // Device scoped, unlike everything else in this function: it releases + // unused graph memory cached by every user of the device, so another + // library in this process pays to map its own graph allocations again. + // Nothing breaks, since only unused blocks go. + const cudaError_t graph_err = cudaDeviceGraphMemTrim(device); + if (graph_err != cudaSuccess) { + ET_LOG( + Error, + "cudaDeviceGraphMemTrim failed for device %d: %s.", + device, + cudaGetErrorString(graph_err)); + (void)cudaGetLastError(); + } + } +#endif +} + Error CudaAllocator::memcpy_async( void* dst, const void* src, diff --git a/backends/cuda/runtime/cuda_allocator.h b/backends/cuda/runtime/cuda_allocator.h index b0a76a51f6d..8bf383fcb51 100644 --- a/backends/cuda/runtime/cuda_allocator.h +++ b/backends/cuda/runtime/cuda_allocator.h @@ -68,6 +68,53 @@ 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 the pool trim never affects memory another user of the async allocator + * is holding. The graph memory trim is the exception: it is scoped to the + * device, so it also releases unused graph memory cached by other users in + * this process. + * + * Does nothing on ROCm. HIP has equivalents for all of these calls; this + * repository's CUDA-to-HIP compatibility header does not alias them yet. + * + * @param index Device to release on, or a negative value to release every + * device this backend has allocated on. Note that means every device, not + * the current one, which is what a negative value means elsewhere in this + * class: 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); + +#if !defined(EXECUTORCH_USE_HIP) + /** + * 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. No production caller. + * + * Not declared on ROCm: the pool code is compiled out there, so there is + * nothing to observe and the pool type needs no HIP alias. + * + * @param index Device to query, or a negative value for the current one. + */ + static cudaMemPool_t pool_for_device( + executorch::runtime::etensor::DeviceIndex index); +#endif // !EXECUTORCH_USE_HIP + /** * 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 4322d2da15a..f826633984e 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,28 @@ 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) { + // Only frees the driver has already observed can be released, so without + // this the trim gives back nothing rather than less. A device wait rather + // than a stream wait because the frees went to whichever stream was + // current when the storage was released, which is not necessarily one + // this handle recorded, and teardown is not bound to the thread that ran + // the method either. + const cudaError_t sync_err = cudaDeviceSynchronize(); + if (sync_err != cudaSuccess) { + ET_LOG( + Error, + "cudaDeviceSynchronize before releasing the pool failed: %s.", + cudaGetErrorString(sync_err)); + (void)cudaGetLastError(); + } + CudaAllocator::release_cached_memory(-1); + } } private: @@ -911,6 +935,10 @@ class ET_EXPERIMENTAL CudaBackend final // OFF; versioned FQN artifacts do not consult this option. std::atomic weight_sharing_across_methods_{false}; + // 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}; + // --------------------------------------------------------------- // Per-weight constant cache. // diff --git a/backends/cuda/runtime/test/test_cuda_allocator.cpp b/backends/cuda/runtime/test/test_cuda_allocator.cpp index a872099c4f7..4ffa81a98e0 100644 --- a/backends/cuda/runtime/test/test_cuda_allocator.cpp +++ b/backends/cuda/runtime/test/test_cuda_allocator.cpp @@ -34,6 +34,15 @@ class CudaAllocatorTest : public testing::Test { } } + // The pool is meant to stay warm, so without this a test that measured + // reserved bytes would see whatever an earlier one left behind, and the order + // would matter. + void TearDown() override { + if (device_count_ > 0) { + CudaAllocator::release_cached_memory(-1); + } + } + // One past the last valid device ordinal, so switching to it always fails. // Only the tests that need such an ordinal call this, so the fit check lives // here rather than in SetUp, where it would also skip the device-0 tests. @@ -181,3 +190,206 @@ 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. +// The retention threshold is the whole point of owning a pool: at the default +// of zero the driver empties it on every synchronize. Nothing else in this +// suite notices a smaller value, so it is asserted directly. +TEST_F(CudaAllocatorTest, PoolRetainsMemoryWithoutLimit) { + cudaStream_t stream; + ASSERT_EQ(cudaStreamCreate(&stream), cudaSuccess); + + auto res = CudaAllocator::allocate_async(8u << 20, 0, stream); + ASSERT_TRUE(res.ok()); + + cudaMemPool_t pool = CudaAllocator::pool_for_device(0); + ASSERT_NE(pool, nullptr); + + uint64_t threshold = 0; + ASSERT_EQ( + cudaMemPoolGetAttribute( + pool, cudaMemPoolAttrReleaseThreshold, &threshold), + cudaSuccess); + EXPECT_EQ(threshold, UINT64_MAX) + << "the pool must hold on to freed memory rather than return it"; + + CudaAllocator::deallocate_async(res.get(), 0, stream); + ASSERT_EQ(cudaStreamSynchronize(stream), cudaSuccess); + ASSERT_EQ(cudaStreamDestroy(stream), cudaSuccess); +} + +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 releases every device this backend has allocated on, not the +// current one. This runner has a single GPU, so the two cannot be told apart +// here; what it pins is that the sentinel is resolved rather than passed to the +// driver. +TEST_F(CudaAllocatorTest, ReleaseCachedMemoryAcceptsTheAllDevicesSentinel) { + 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); +} + +// Memory allocated during a graph capture goes to the device graph pool, which +// the pool trim cannot reach, so releasing has to trim that too. Without the +// graph trim this is the only new test that fails. +TEST_F(CudaAllocatorTest, ReleaseCachedMemoryReturnsGraphMemory) { + cudaStream_t stream; + ASSERT_EQ(cudaStreamCreate(&stream), cudaSuccess); + + constexpr size_t kBytes = 64u << 20; + cudaGraph_t graph = nullptr; + cudaGraphExec_t graph_exec = nullptr; + ASSERT_EQ( + cudaStreamBeginCapture(stream, cudaStreamCaptureModeRelaxed), + cudaSuccess); + auto captured = CudaAllocator::allocate_async(kBytes, 0, stream); + ASSERT_TRUE(captured.ok()); + CudaAllocator::deallocate_async(captured.get(), 0, stream); + ASSERT_EQ(cudaStreamEndCapture(stream, &graph), cudaSuccess); + ASSERT_EQ( + cudaGraphInstantiate(&graph_exec, graph, nullptr, nullptr, 0), + cudaSuccess); + ASSERT_EQ(cudaGraphLaunch(graph_exec, stream), cudaSuccess); + ASSERT_EQ(cudaStreamSynchronize(stream), cudaSuccess); + + size_t reserved = 0; + ASSERT_EQ( + cudaDeviceGetGraphMemAttribute( + 0, cudaGraphMemAttrReservedMemCurrent, &reserved), + cudaSuccess); + ASSERT_GT(reserved, 0u) << "the capture should have reserved graph memory"; + + ASSERT_EQ(cudaGraphExecDestroy(graph_exec), cudaSuccess); + ASSERT_EQ(cudaGraphDestroy(graph), cudaSuccess); + + CudaAllocator::release_cached_memory(0); + + ASSERT_EQ( + cudaDeviceGetGraphMemAttribute( + 0, cudaGraphMemAttrReservedMemCurrent, &reserved), + cudaSuccess); + EXPECT_EQ(reserved, 0u) << "graph memory should go back to the driver"; + + ASSERT_EQ(cudaStreamDestroy(stream), cudaSuccess); +} + +#endif // !EXECUTORCH_USE_HIP