diff --git a/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.cc b/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.cc index 36ab2f62b6..e90417a6d1 100644 --- a/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.cc +++ b/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.cc @@ -1039,10 +1039,6 @@ TensorrtExecutionProvider::TensorrtExecutionProvider(const TensorrtExecutionProv int8_calibration_cache_available_ = !int8_calibration_cache_name_.empty(); } - if (cuda_graph_enable_) { - GetPerThreadContext().InitCUDAGraph(); - } - /* * Parse explicit min/max/opt profile shapes from provider options. * @@ -1155,12 +1151,8 @@ Status TensorrtExecutionProvider::ReplayGraph() { return GetPerThreadContext().ReplayGraph(); } -void TensorrtExecutionProvider::PerThreadContext::InitCUDAGraph() { - cuda_graph_ = std::make_unique(); -} - void TensorrtExecutionProvider::PerThreadContext::SetGraphStream(cudaStream_t stream) { - cuda_graph_->SetStream(stream); + cuda_graph_.SetStream(stream); } bool TensorrtExecutionProvider::PerThreadContext::IsGraphCaptureAllowed() const { @@ -1168,12 +1160,12 @@ bool TensorrtExecutionProvider::PerThreadContext::IsGraphCaptureAllowed() const } void TensorrtExecutionProvider::PerThreadContext::CaptureBegin() { - cuda_graph_->Reset(); - cuda_graph_->CaptureBegin(); + cuda_graph_.Reset(); + cuda_graph_.CaptureBegin(); } void TensorrtExecutionProvider::PerThreadContext::CaptureEnd() { - cuda_graph_->CaptureEnd(); + cuda_graph_.CaptureEnd(); is_graph_captured_ = true; } @@ -1186,7 +1178,7 @@ Status TensorrtExecutionProvider::PerThreadContext::ReplayGraph() { // Please note that CUDAGraph::Replay() is not thread safe. // The cuda graph object is maintained by a per thread basis, // therefore calling CUDAGraph::Replay() here is guaranteed to be thread safe. - return cuda_graph_->Replay(); + return cuda_graph_.Replay(); } void TensorrtExecutionProvider::PerThreadContext::IncrementRegularRunCountBeforeGraphCapture() { diff --git a/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.h b/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.h index 619285bd5f..e00e5df581 100644 --- a/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.h +++ b/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.h @@ -306,15 +306,15 @@ class TensorrtExecutionProvider : public IExecutionProvider { std::unordered_map input_shape_ranges_; // Cuda graph with multi threads will be supported in the future, so cuda_graph_ is put under PerThreadContext. - // ORT TRT only supports CUDA graph when whole model is supported by TRT, so simply maintaining a CUDAGraph pointer is enough (no need to maintain one CUDAGraph pointer per TRT subgraph) - std::unique_ptr cuda_graph_; + // ORT TRT only supports CUDA graph when whole model is supported by TRT, so simply maintaining a CUDAGraph instance is enough (no need to maintain one CUDAGraph instance per TRT subgraph) + CUDAGraph cuda_graph_; bool is_graph_captured_ = false; - int regular_run_count_before_graph_capture_ = -1; + int regular_run_count_before_graph_capture_ = 0; // There is chance (currently only happens in CUDA EP) that the second regular run allocates GPU memory for causes like: // (1) memory pattern is enabled. (2) arena allocation for stream. // Since no GPU memory allocation is allowed during graph capturing, we need at least two regular runs // to allocate enough memory in Arena before graph capturing. - const int min_num_runs_before_cuda_graph_capture_ = 0; // required min regular runs before graph capture for the necessary memory allocations. + const int min_num_runs_before_cuda_graph_capture_ = 1; // required min regular runs before graph capture for the necessary memory allocations. }; using PerThreadContextMap = std::unordered_map>;