Keep the CUDA memory pool warm between delegates

backends/cuda/runtime/cuda_allocator.cpp

@@ -12,6 +12,12 @@
 #include <executorch/extension/cuda/runtime_api.h>
 #include <executorch/runtime/platform/log.h>
 
+#if !defined(EXECUTORCH_USE_HIP)
+#include <mutex>
+#include <unordered_map>
+#include <vector>
+#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<int, cudaMemPool_t> 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<int>(index);
+  }
+  int current = 0;
+  const cudaError_t err = cudaGetDevice(&current);
+  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<std::mutex> 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<void*> 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<int>(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<int>(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<std::mutex> 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<std::pair<int, cudaMemPool_t>> targets;
+  {
+    auto& state = mem_pool_state();
+    const std::lock_guard<std::mutex> lock(state.mutex);
+    if (index >= 0) {
+      const auto it = state.pools.find(static_cast<int>(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,

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.

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
   }
 
@@ -864,6 +866,21 @@ class ET_EXPERIMENTAL CudaBackend final
 
     mutable_state_forget_handle(handle);
 
+    // Waited on before the stream reference goes, so the frees this handle
+    // queued are observed by the driver. Without this the pool trim at the end
+    // of teardown sees them as still pending and gives back nothing, and the
+    // work is being abandoned anyway.
+    if (handle->cuda_stream != nullptr && *handle->cuda_stream != nullptr) {
+      const cudaError_t sync_err = cudaStreamSynchronize(*handle->cuda_stream);
+      if (sync_err != cudaSuccess) {
+        ET_LOG(
+            Error,
+            "cudaStreamSynchronize failed during teardown: %s.",
+            cudaGetErrorString(sync_err));
+        (void)cudaGetLastError();
+      }
+    }
+
     // The CUDA stream is managed by shared_ptr in the handle.
     // It will be automatically destroyed when the last handle using it
     // is destroyed. Just reset our reference.
@@ -892,6 +909,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 +931,10 @@ class ET_EXPERIMENTAL CudaBackend final
   mutable std::mutex cuda_stream_mutex_;
   std::shared_ptr<cudaStream_t> 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<size_t> 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.