Keep the CUDA memory pool warm between delegates

backends/cuda/runtime/cuda_allocator.cpp

@@ -12,6 +12,11 @@
 #include <executorch/extension/cuda/runtime_api.h>
 #include <executorch/runtime/platform/log.h>
 
+#if !defined(EXECUTORCH_USE_HIP)
+#include <mutex>
+#include <unordered_map>
+#endif
+
 namespace executorch::backends::cuda {
 
 using executorch::runtime::Error;
@@ -21,6 +26,109 @@ using executorch::runtime::etensor::DeviceType;
 
 namespace {
 
+#if !defined(EXECUTORCH_USE_HIP)
+// The pool hands physical memory back to the driver whenever a synchronization
+// observes a pending free, so with the default threshold of zero it 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
+// threshold caps what the pool may keep rather than granting headroom, so any
+// fixed value stops helping once a model's footprint passes it. The maximum
+// keeps the win independent of model size, and release_cached_memory gives the
+// memory back when the last delegate goes away.
+constexpr uint64_t kMemPoolReleaseThreshold = UINT64_MAX;
+
+struct MemPoolState {
+  std::mutex mutex;
+  // Previous threshold per configured device, so a pool this code did not
+  // create can be put back the way its owner left it.
+  std::unordered_map<int, uint64_t> configured;
+};
+
+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, in which case
+// the pool is left alone.
+int resolve_device(DeviceIndex index) {
+  if (index >= 0) {
+    return static_cast<int>(index);
+  }
+  int current = 0;
+  if (cudaGetDevice(&current) != cudaSuccess) {
+    return -1;
+  }
+  return current;
+}
+
+// Raises the release threshold on first use of a device. Failing to configure
+// the pool only costs speed, so this reports success either way and leaves the
+// pool at its defaults. The device is recorded as configured only after the
+// calls succeed, so a transient failure is retried on the next allocation.
+void ensure_mem_pool_configured(int device) {
+  auto& state = mem_pool_state();
+  {
+    const std::lock_guard<std::mutex> lock(state.mutex);
+    if (state.configured.count(device) != 0) {
+      return;
+    }
+  }
+
+  cudaMemPool_t pool = nullptr;
+  // cudaMallocAsync allocates from the device's current pool, which is the
+  // default pool unless something replaced it. That pool is shared with every
+  // other user of the async allocator in this process.
+  cudaError_t err = cudaDeviceGetMemPool(&pool, device);
+  if (err != cudaSuccess) {
+    ET_LOG(
+        Error,
+        "cudaDeviceGetMemPool failed for device %d: %s. Using pool defaults.",
+        device,
+        cudaGetErrorString(err));
+    (void)cudaGetLastError();
+    return;
+  }
+
+  uint64_t previous = 0;
+  err =
+      cudaMemPoolGetAttribute(pool, cudaMemPoolAttrReleaseThreshold, &previous);
+  if (err != cudaSuccess) {
+    ET_LOG(
+        Error,
+        "Reading the pool release threshold failed for device %d: %s. Using "
+        "pool defaults.",
+        device,
+        cudaGetErrorString(err));
+    (void)cudaGetLastError();
+    return;
+  }
+
+  // Only ever raise it. Another user of this pool may have asked for more, and
+  // lowering theirs would shrink a cache this code does not own.
+  if (previous < kMemPoolReleaseThreshold) {
+    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. Using "
+          "pool defaults.",
+          device,
+          cudaGetErrorString(err));
+      (void)cudaGetLastError();
+      return;
+    }
+  }
+
+  const std::lock_guard<std::mutex> lock(state.mutex);
+  state.configured.emplace(device, previous);
+}
+
+#endif // !EXECUTORCH_USE_HIP
+
 Error copy_impl(
     void* dst,
     const void* src,
@@ -302,6 +410,13 @@ Result<void*> CudaAllocator::allocate_async(
     size_t nbytes,
     DeviceIndex index,
     cudaStream_t stream) {
+#if !defined(EXECUTORCH_USE_HIP)
+  const int device = resolve_device(index);
+  if (device >= 0) {
+    ensure_mem_pool_configured(device);
+  }
+#endif
+
   void* ptr = nullptr;
   cudaError_t err = cudaMallocAsync(&ptr, nbytes, stream);
   if (err != cudaSuccess) {
@@ -313,6 +428,7 @@ Result<void*> CudaAllocator::allocate_async(
         static_cast<int>(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,91 @@ void CudaAllocator::deallocate_async(
         cudaGetErrorString(err),
         ptr,
         static_cast<int>(index));
+    return;
+  }
+}
+
+void CudaAllocator::release_cached_memory(DeviceIndex index) {
+#if defined(EXECUTORCH_USE_HIP)
+  (void)index;
+#else
+  std::unordered_map<int, uint64_t> targets;
+  {
+    auto& state = mem_pool_state();
+    const std::lock_guard<std::mutex> lock(state.mutex);
+    if (index >= 0) {
+      const auto it = state.configured.find(static_cast<int>(index));
+      if (it == state.configured.end()) {
+        return;
+      }
+      targets.emplace(it->first, it->second);
+      state.configured.erase(it);
+    } else {
+      // A caller asking for everything gets every pool this code raised, not
+      // whichever device the calling thread happens to be current on. The
+      // delegate that ran is often not on that device.
+      targets.swap(state.configured);
+    }
+  }
+
+  for (const auto& [device, previous] : targets) {
+    cudaMemPool_t pool = nullptr;
+    if (cudaDeviceGetMemPool(&pool, device) != cudaSuccess) {
+      (void)cudaGetLastError();
+      continue;
+    }
+
+    // The pool can only release a free the driver has already observed, and
+    // nothing on the teardown path waits for the frees this delegate queued, so
+    // without this the trim below finds nothing to give back.
+    int prev_device = 0;
+    if (cudaGetDevice(&prev_device) != cudaSuccess) {
+      (void)cudaGetLastError();
+      continue;
+    }
+    if (cudaSetDevice(device) != cudaSuccess) {
+      (void)cudaGetLastError();
+      continue;
+    }
+    cudaError_t err = cudaDeviceSynchronize();
+    if (err != cudaSuccess) {
+      ET_LOG(
+          Error,
+          "cudaDeviceSynchronize failed for device %d: %s. The pool keeps its "
+          "memory.",
+          device,
+          cudaGetErrorString(err));
+      (void)cudaGetLastError();
+      (void)cudaSetDevice(prev_device);
+      continue;
+    }
+
+    err = cudaMemPoolTrimTo(pool, 0);
+    if (err != cudaSuccess) {
+      ET_LOG(
+          Error,
+          "cudaMemPoolTrimTo failed for device %d: %s.",
+          device,
+          cudaGetErrorString(err));
+      (void)cudaGetLastError();
+    }
+
+    // The pool is shared, so hand it back the way its owner left it.
+    uint64_t restore = previous;
+    err = cudaMemPoolSetAttribute(
+        pool, cudaMemPoolAttrReleaseThreshold, &restore);
+    if (err != cudaSuccess) {
+      ET_LOG(
+          Error,
+          "Restoring the pool release threshold failed for device %d: %s.",
+          device,
+          cudaGetErrorString(err));
+      (void)cudaGetLastError();
+    }
+
+    (void)cudaSetDevice(prev_device);
   }
+#endif
 }
 
 Error CudaAllocator::memcpy_async(

backends/cuda/runtime/cuda_allocator.h

@@ -68,6 +68,20 @@ class CudaAllocator final : public executorch::runtime::DeviceAllocator {
       executorch::runtime::etensor::DeviceIndex index,
       cudaStream_t stream);
 
+  /**
+   * Return memory the 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. Call this once the
+   * work on the device is finished and the caller has synchronized, so a
+   * long-lived process does not keep memory it no longer needs. Safe to call at
+   * any time: allocations that are still live are unaffected.
+   *
+   * @param index Device to release on, or a negative value for the current one.
+   */
+  static void release_cached_memory(
+      executorch::runtime::etensor::DeviceIndex index);
+
   /**
    * 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
   }
 
@@ -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<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.