diff --git a/onnxruntime/core/providers/cuda/cuda_execution_provider.cc b/onnxruntime/core/providers/cuda/cuda_execution_provider.cc index f624c3f66b73c..9729b495dc1f8 100755 --- a/onnxruntime/core/providers/cuda/cuda_execution_provider.cc +++ b/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 buffer) const { + std::lock_guard lock(captured_host_buffers_mutex_); + captured_host_buffers_.push_back(std::move(buffer)); +} + namespace cuda { template <> diff --git a/onnxruntime/core/providers/cuda/cuda_execution_provider.h b/onnxruntime/core/providers/cuda/cuda_execution_provider.h index 3d617794b1fbf..ef546273f3165 100644 --- a/onnxruntime/core/providers/cuda/cuda_execution_provider.h +++ b/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 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 buffer) const; + private: + mutable std::mutex captured_host_buffers_mutex_; + mutable std::vector> captured_host_buffers_; + CUDAExecutionProviderInfo info_; cudaDeviceProp device_prop_; bool external_stream_ = false; diff --git a/onnxruntime/core/providers/cuda/cuda_kernel.h b/onnxruntime/core/providers/cuda/cuda_kernel.h index 6152b5f888eb2..b9d48e4925334 100644 --- a/onnxruntime/core/providers/cuda/cuda_kernel.h +++ b/onnxruntime/core/providers/cuda/cuda_kernel.h @@ -109,6 +109,17 @@ class CudaKernel : public OpKernel { return IAllocator::MakeUniquePtr(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 + inline void RetainBufferForGraphCapture(IAllocatorUniquePtr buffer) const { + auto deleter = buffer.get_deleter(); + provider_->RetainBufferForGraphCapture( + std::shared_ptr(buffer.release(), + [deleter](void* p) { deleter(static_cast(p)); })); + } + const cudaDeviceProp& GetDeviceProp() const { return provider_->GetDeviceProp(); } int GetCudnnConvAlgo() const { return provider_->GetCudnnConvAlgo(); } bool GetCudnnConvUseMaxWorkspace() const { return provider_->GetCudnnConvUseMaxWorkspace(); } @@ -221,7 +232,20 @@ class CudaKernel : public OpKernel { cudaStream_t cuda_stream = stream ? static_cast(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(); } diff --git a/onnxruntime/core/providers/cuda/plugin/cuda_kernel_adapter.h b/onnxruntime/core/providers/cuda/plugin/cuda_kernel_adapter.h index ed698aaf74eb1..6781eddab35b3 100644 --- a/onnxruntime/core/providers/cuda/plugin/cuda_kernel_adapter.h +++ b/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> captured_host_buffers; + + void RetainBufferForGraphCapture(std::shared_ptr buffer) { + std::lock_guard lock(captured_host_buffers_mutex); + captured_host_buffers.push_back(std::move(buffer)); + } }; template struct SizeOf { @@ -1257,6 +1264,14 @@ class CudaKernel : public OpKernel { return IAllocatorUniquePtr(static_cast(p), [](T* ptr) { if (ptr) cudaFreeHost(ptr); }); } + template + inline void RetainBufferForGraphCapture(IAllocatorUniquePtr buffer) const { + auto deleter = buffer.get_deleter(); + runtime_config_->RetainBufferForGraphCapture( + std::shared_ptr(buffer.release(), + [deleter](void* p) { deleter(static_cast(p)); })); + } + template 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(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(s), &capture_status)); + } + if (capture_status != cudaStreamCaptureStatusNone) { + op_kernel_->RetainBufferForGraphCapture(std::move(cpu_)); + } else { + op_kernel_->AddDeferredReleaseCPUPtr(cpu_.release(), s); + } } return Status::OK(); } diff --git a/onnxruntime/test/framework/dynamic_plugin_ep_test.cc b/onnxruntime/test/framework/dynamic_plugin_ep_test.cc index 2453096fa8094..3f6d1ec1654b2 100644 --- a/onnxruntime/test/framework/dynamic_plugin_ep_test.cc +++ b/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 retained_buffer; + { + onnxruntime::cuda::detail::CudaKernelAdapterRuntimeConfig config; + auto buffer = std::make_shared(42); + retained_buffer = buffer; + config.RetainBufferForGraphCapture(std::move(buffer)); + EXPECT_FALSE(retained_buffer.expired()); + } + EXPECT_TRUE(retained_buffer.expired()); +} #endif } // namespace onnxruntime::test