From 30a450dcf8f90006b18a826199aefd0e7a6908b2 Mon Sep 17 00:00:00 2001
From: Chi Lo <54722500+chilo-ms@users.noreply.github.com>
Date: Thu, 31 Aug 2023 13:32:15 -0700
Subject: [PATCH] Fix TRT EP's cuda graph feature (#17355)

When users run inference with cuda graph enable with multithreading,
only the main thread creating the inference session will successfully
initialize cuda graph instance, for other threads executing the
inference run directly, they will hit segfault due to not calling
allocation/initialization for cuda graph instance.
This PR fixes this issue.
---
 .../tensorrt/tensorrt_execution_provider.cc    | 18 +++++-------------
 .../tensorrt/tensorrt_execution_provider.h     |  8 ++++----
 2 files changed, 9 insertions(+), 17 deletions(-)

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<CUDAGraph>();
-}
-
 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<std::string, ShapeRangesMap> 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<CUDAGraph> 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<const TensorrtExecutionProvider*, std::weak_ptr<PerThreadContext>>;