saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-23 19:32:23 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Fix TRT EP's cuda graph feature (#17355)

Browse source 
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.
This commit is contained in:
Chi Lo 2023-08-31 13:32:15 -07:00 committed by GitHub
parent ae90b716ff
commit 30a450dcf8
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
2 changed files with 9 additions and 17 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

18
onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.cc
View file

@ -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() {

8
onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.h
View file

@ -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>>;