[CUDA] Keep CudaAsyncBuffer's staging buffer alive across a graph replay

onnxruntime/core/providers/cuda/cuda_execution_provider.cc

@@ -555,6 +555,11 @@ Status CUDAExecutionProvider::ReplayGraph(int graph_annotation_id, bool sync) {
   return GetPerThreadContext().ReplayGraph(graph_annotation_id, sync);
 }
 
+void CUDAExecutionProvider::RetainBufferForGraphCapture(std::shared_ptr<void> buffer) const {
+  std::lock_guard<std::mutex> lock(captured_host_buffers_mutex_);
+  captured_host_buffers_.push_back(std::move(buffer));
+}
+
 namespace cuda {
 
 template <>

onnxruntime/core/providers/cuda/cuda_execution_provider.h

@@ -125,7 +125,16 @@ class CUDAExecutionProvider : public IExecutionProvider {
   OrtDevice GetOrtDeviceByMemType(OrtMemType mem_type) const override;
   std::vector<AllocatorPtr> CreatePreferredAllocators() override;
 
+  // Takes ownership of a host buffer that a kernel staged a device copy from while the
+  // stream was capturing.  The copy is a node of the captured graph and re-reads the
+  // buffer on every replay, so it must not go back to the allocator and be overwritten
+  // by a later run.  Held until the provider is destroyed.
+  void RetainBufferForGraphCapture(std::shared_ptr<void> buffer) const;
+
  private:
+  mutable std::mutex captured_host_buffers_mutex_;
+  mutable std::vector<std::shared_ptr<void>> captured_host_buffers_;
+
   CUDAExecutionProviderInfo info_;
   cudaDeviceProp device_prop_;
   bool external_stream_ = false;

onnxruntime/core/providers/cuda/cuda_kernel.h

@@ -109,6 +109,17 @@
     return IAllocator::MakeUniquePtr<T>(Info().GetAllocator(OrtMemType::OrtMemTypeCPU), count_or_bytes);
   }
 
+  // Hands a host buffer to the provider to hold forever.  For buffers a captured
+  // graph copies from: the graph re-reads them on every replay, so they must not be
+  // recycled by the allocator.
+  template <typename T>
+  inline void RetainBufferForGraphCapture(IAllocatorUniquePtr<T> buffer) const {
+    auto deleter = buffer.get_deleter();
+    provider_->RetainBufferForGraphCapture(
+        std::shared_ptr<void>(buffer.release(),
+                              [deleter](void* p) { deleter(static_cast<T*>(p)); }));
+  }
+
   const cudaDeviceProp& GetDeviceProp() const { return provider_->GetDeviceProp(); }
   int GetCudnnConvAlgo() const { return provider_->GetCudnnConvAlgo(); }
   bool GetCudnnConvUseMaxWorkspace() const { return provider_->GetCudnnConvUseMaxWorkspace(); }
@@ -221,7 +232,20 @@
         cudaStream_t cuda_stream = stream ? static_cast<cudaStream_t>(stream->GetHandle()) : nullptr;
         CUDA_RETURN_IF_ERROR(cudaMemcpyAsync(gpu_copy_.get(), cpu_pinned_copy_.get(), count_ * sizeof(T), cudaMemcpyHostToDevice,
                                              cuda_stream));
-        op_kernel_->AddDeferredReleaseCPUPtr(cpu_pinned_copy_.release(), stream);
+        // While the stream is capturing, the copy above becomes a node of the graph
+        // and reads this host buffer again on every replay.  Releasing it here would
+        // let a later run of a different shape take the block and overwrite it, and
+        // the replay would then copy that run's bytes instead -- for the pointer
+        // arrays this class mostly carries, addresses that belong to nothing.
+        cudaStreamCaptureStatus capture_status = cudaStreamCaptureStatusNone;
+        if (cuda_stream != nullptr) {
+          CUDA_RETURN_IF_ERROR(cudaStreamIsCapturing(cuda_stream, &capture_status));
+        }
+        if (capture_status != cudaStreamCaptureStatusNone) {
+          op_kernel_->RetainBufferForGraphCapture(std::move(cpu_pinned_copy_));
+        } else {
+          op_kernel_->AddDeferredReleaseCPUPtr(cpu_pinned_copy_.release(), stream);
+        }
       }
       return Status::OK();
     }

onnxruntime/core/providers/cuda/plugin/cuda_kernel_adapter.h

@@ -539,6 +539,13 @@ struct CudaKernelAdapterRuntimeConfig {
   bool do_copy_in_default_stream = true;
   cudaDeviceProp device_prop{};
   onnxruntime::AttentionKernelOptions attention_kernel_options;
+  std::mutex captured_host_buffers_mutex;
+  std::vector<std::shared_ptr<void>> captured_host_buffers;
+
+  void RetainBufferForGraphCapture(std::shared_ptr<void> buffer) {
+    std::lock_guard<std::mutex> lock(captured_host_buffers_mutex);
+    captured_host_buffers.push_back(std::move(buffer));
+  }
 };
 template <typename T>
 struct SizeOf {
@@ -1257,6 +1264,14 @@ class CudaKernel : public OpKernel {
     return IAllocatorUniquePtr<T>(static_cast<T*>(p), [](T* ptr) { if (ptr) cudaFreeHost(ptr); });
   }
 
+  template <typename T>
+  inline void RetainBufferForGraphCapture(IAllocatorUniquePtr<T> buffer) const {
+    auto deleter = buffer.get_deleter();
+    runtime_config_->RetainBufferForGraphCapture(
+        std::shared_ptr<void>(buffer.release(),
+                              [deleter](void* p) { deleter(static_cast<T*>(p)); }));
+  }
+
   template <typename T>
   class CudaAsyncBuffer {
    public:
@@ -1286,7 +1301,15 @@ class CudaKernel : public OpKernel {
           ORT_THROW("CUDA async buffer copy size overflow for ", count_, " elements");
         }
         if (cudaMemcpyAsync(gpu_.get(), cpu_.get(), bytes, cudaMemcpyHostToDevice, static_cast<cudaStream_t>(s)) != cudaSuccess) return Status(onnxruntime::common::ONNXRUNTIME, onnxruntime::common::FAIL, "Memcpy fail");
-        op_kernel_->AddDeferredReleaseCPUPtr(cpu_.release(), s);
+        cudaStreamCaptureStatus capture_status = cudaStreamCaptureStatusNone;
+        if (s != nullptr) {
+          CUDA_RETURN_IF_ERROR(cudaStreamIsCapturing(static_cast<cudaStream_t>(s), &capture_status));
+        }
+        if (capture_status != cudaStreamCaptureStatusNone) {
+          op_kernel_->RetainBufferForGraphCapture(std::move(cpu_));
+        } else {
+          op_kernel_->AddDeferredReleaseCPUPtr(cpu_.release(), s);
+        }
       }
       return Status::OK();
     }

onnxruntime/test/framework/dynamic_plugin_ep_test.cc

@@ -151,6 +151,18 @@ TEST(DynamicPluginEpInfraTest, CudaKernelAdapterTryBytesForCountNormalCase) {
   ASSERT_TRUE(onnxruntime::cuda::detail::TryBytesForCount(10, 4, bytes));
   EXPECT_EQ(bytes, size_t{40});
 }
+
+TEST(DynamicPluginEpInfraTest, CudaKernelAdapterRetainsGraphCaptureHostBuffer) {
+  std::weak_ptr<int> retained_buffer;
+  {
+    onnxruntime::cuda::detail::CudaKernelAdapterRuntimeConfig config;
+    auto buffer = std::make_shared<int>(42);
+    retained_buffer = buffer;
+    config.RetainBufferForGraphCapture(std::move(buffer));
+    EXPECT_FALSE(retained_buffer.expired());
+  }
+  EXPECT_TRUE(retained_buffer.expired());
+}
 #endif
 
 }  // namespace onnxruntime::test