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

[TensorRT EP] Back out the PerThreadContext (#17690)

Browse source 
Current TRT EP's PerthreadContext allows more than one IExecutionContext
instance to be created by one engine instance.
But, it's possible to hit an error that caused by TRT API
context.setBindingDimensions() in our TRT EP code
[here](https://nam06.safelinks.protection.outlook.com/?url=https%3A%2F%2Fgithub.com%2Fmicrosoft%2Fonnxruntime%2Fblob%2Fmain%2Fonnxruntime%2Fcore%2Fproviders%2Ftensorrt%2Ftensorrt_execution_provider.cc%23L2775&data=05%7C01%7CChi.Lo%40microsoft.com%7Cd8b23c3a4c0b4dcce9b408dbbd9309de%7C72f988bf86f141af91ab2d7cd011db47%7C1%7C0%7C638312211465211140%7CUnknown%7CTWFpbGZsb3d8eyJWIjoiMC4wLjAwMDAiLCJQIjoiV2luMzIiLCJBTiI6Ik1haWwiLCJXVCI6Mn0%3D%7C3000%7C%7C%7C&sdata=5EZoAoXgWFSuz%2BIRMH%2FXZaO%2BfKNP%2FZDZYEZg3W%2Ff30w%3D&reserved=0)
under the case of the input shape changes ( meaning engine being
rebuilt) with multithreading.
From the
[doc](https://nam06.safelinks.protection.outlook.com/?url=https%3A%2F%2Fdocs.nvidia.com%2Fdeeplearning%2Ftensorrt%2Fapi%2Fc_api%2Fclassnvinfer1_1_1_i_execution_context.html%23ada050e88320bcc40987b0acadc2ef962&data=05%7C01%7CChi.Lo%40microsoft.com%7Cd8b23c3a4c0b4dcce9b408dbbd9309de%7C72f988bf86f141af91ab2d7cd011db47%7C1%7C0%7C638312211465211140%7CUnknown%7CTWFpbGZsb3d8eyJWIjoiMC4wLjAwMDAiLCJQIjoiV2luMzIiLCJBTiI6Ik1haWwiLCJXVCI6Mn0%3D%7C3000%7C%7C%7C&sdata=%2BmVZU5iLD97B3YBPdHZP7jOQ2dGoleI3R0mSMVgopG4%3D&reserved=0)
and the
[discussion](https://nam06.safelinks.protection.outlook.com/?url=https%3A%2F%2Fgithub.com%2FNVIDIA%2FTensorRT%2Fissues%2F846&data=05%7C01%7CChi.Lo%40microsoft.com%7Cd8b23c3a4c0b4dcce9b408dbbd9309de%7C72f988bf86f141af91ab2d7cd011db47%7C1%7C0%7C638312211465211140%7CUnknown%7CTWFpbGZsb3d8eyJWIjoiMC4wLjAwMDAiLCJQIjoiV2luMzIiLCJBTiI6Ik1haWwiLCJXVCI6Mn0%3D%7C3000%7C%7C%7C&sdata=c8v%2FK2UkQ%2FNbf8w1sHNDGsB2kxw4sSmkyQ2QuCs8Fs8%3D&reserved=0),
it seems we should have different OptimizationProfile for different
IExecutionContext which our current TRT EP doesn’t support regardless of
using PerThreadContext implementation.
Back out the PerThreadContext until we completely solve this issue.
This commit is contained in:
Chi Lo 2023-09-26 09:28:17 -07:00 committed by GitHub
parent aed43f429a
commit 7572e6055c
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
2 changed files with 55 additions and 82 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -1143,46 +1143,35 @@ bool TensorrtExecutionProvider::IsGraphCaptureEnabled() const {
return cuda_graph_enable_;
}
bool TensorrtExecutionProvider::IsGraphCaptured() const {
return GetPerThreadContext().IsGraphCaptured();
}
Status TensorrtExecutionProvider::ReplayGraph() {
return GetPerThreadContext().ReplayGraph();
}
void TensorrtExecutionProvider::PerThreadContext::SetGraphStream(cudaStream_t stream) {
cuda_graph_.SetStream(stream);
}
bool TensorrtExecutionProvider::PerThreadContext::IsGraphCaptureAllowed() const {
bool TensorrtExecutionProvider::IsGraphCaptureAllowed() const {
return regular_run_count_before_graph_capture_ >= min_num_runs_before_cuda_graph_capture_;
}
void TensorrtExecutionProvider::PerThreadContext::CaptureBegin() {
void TensorrtExecutionProvider::CaptureBegin() {
cuda_graph_.Reset();
cuda_graph_.CaptureBegin();
}
void TensorrtExecutionProvider::PerThreadContext::CaptureEnd() {
void TensorrtExecutionProvider::CaptureEnd() {
cuda_graph_.CaptureEnd();
is_graph_captured_ = true;
}
bool TensorrtExecutionProvider::PerThreadContext::IsGraphCaptured() const {
bool TensorrtExecutionProvider::IsGraphCaptured() const {
return is_graph_captured_;
}
Status TensorrtExecutionProvider::PerThreadContext::ReplayGraph() {
Status TensorrtExecutionProvider::ReplayGraph() {
ORT_ENFORCE(IsGraphCaptured());
// Please note that CUDAGraph::Replay() is not thread safe.
// The cuda graph object is maintained by a per thread basis,
// ORT TRT calls ReplayGraph() in compute_func() where synchromization is enforced due to lock_guard(),
// therefore calling CUDAGraph::Replay() here is guaranteed to be thread safe.
return cuda_graph_.Replay();
}
void TensorrtExecutionProvider::PerThreadContext::IncrementRegularRunCountBeforeGraphCapture() {
// The cuda graph object is maintained by a per thread basis,
void TensorrtExecutionProvider::IncrementRegularRunCountBeforeGraphCapture() {
// Please note that this function is not thread safe.
// ORT TRT calls this function in compute_func() where synchronization is enforced due to lock_guard(),
// therefore following increment is guaranteed to be thread safe.
++regular_run_count_before_graph_capture_;
}
@ -1213,18 +1202,6 @@ Status TensorrtExecutionProvider::OnRunEnd(bool sync_stream) {
if (sync_stream && external_stream_) {
CUDA_RETURN_IF_ERROR(cudaStreamSynchronize(stream_));
}
// The reason of !IsGraphCaptureEnabled():
// If cuda graph is enabled, the per thread context will not be released
// because the per thread cuda graph needs to be maintained and replayed for
// the next run.
// The reason of PerThreadContextCache()->find(this) != PerThreadContextCache()->end():
// In extreme cases (e.g., 1-op graph and that op fallbacks to CPU),
// PerThreadContext won't be created and there is nothing to release.
if (!IsGraphCaptureEnabled() &&
PerThreadContextCache()->find(this) != PerThreadContextCache()->end()) {
ReleasePerThreadContext();
}
return Status::OK();
}
@ -2384,6 +2361,7 @@ common::Status TensorrtExecutionProvider::Compile(const std::vector<FusedNodeAnd
// Save TRT engine, other TRT objects and input/output info to map
parsers_.emplace(fused_node.Name(), std::move(trt_parser));
engines_.emplace(fused_node.Name(), std::move(trt_engine));
contexts_.emplace(fused_node.Name(), std::move(trt_context));
builders_.emplace(fused_node.Name(), std::move(trt_builder));
networks_.emplace(fused_node.Name(), std::move(trt_network));
input_info_[fused_node.Name()].push_back(input_indexes);
@ -2392,14 +2370,6 @@ common::Status TensorrtExecutionProvider::Compile(const std::vector<FusedNodeAnd
input_shape_ranges_[fused_node.Name()] = input_implicit_shape_ranges;
profiles_.emplace(fused_node.Name(), std::move(trt_profiles));
// Save TRT context to PerThreadContext map since maintaining execution context in a per thread basis is suggested by TRT doc to avoid synchronization issue
if (trt_context) {
auto context_status = GetPerThreadContext().UpdateTensorRTContext(fused_node.Name(), std::move(trt_context));
if (!context_status) {
return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, "TensorRT EP failed to create context.");
}
}
// Create function state
// TODO: remove default capture
NodeComputeInfo compute_info;
@ -2411,7 +2381,7 @@ common::Status TensorrtExecutionProvider::Compile(const std::vector<FusedNodeAnd
tactics = GetTacticSourceFromString(tactic_sources_);
}
*p = {context->allocate_func, context->release_func, context->allocator_handle, context->node_name,
&parsers_[context->node_name], &engines_[context->node_name], &builders_[context->node_name],
&parsers_[context->node_name], &engines_[context->node_name], &contexts_[context->node_name], &builders_[context->node_name],
&networks_[context->node_name], input_info_[context->node_name], output_info_[context->node_name],
input_shape_ranges_[context->node_name], &tensorrt_mu_, fp16_enable_, int8_enable_, int8_calibration_cache_available_,
dla_enable_, dla_core_, &max_workspace_size_, trt_node_name_with_precision, engine_cache_enable_, cache_path_,
@ -2445,6 +2415,7 @@ common::Status TensorrtExecutionProvider::Compile(const std::vector<FusedNodeAnd
auto& shape_ranges = trt_state->input_shape_ranges;
auto trt_builder = trt_state->builder->get();
auto trt_engine = trt_state->engine->get();
auto trt_context = trt_state->context->get();
auto trt_profiles = trt_state->profiles;
auto max_context_mem_size_ptr = trt_state->max_context_mem_size_ptr;
int num_inputs = static_cast<int>(input_indexes.size());
@ -2502,7 +2473,7 @@ common::Status TensorrtExecutionProvider::Compile(const std::vector<FusedNodeAnd
trt_state->engine->reset();
*(trt_state->engine) = std::unique_ptr<nvinfer1::ICudaEngine>(
trt_state->runtime->deserializeCudaEngine(engine_buf.get(), engine_size, nullptr));
if (*(trt_state->engine) == nullptr) {
if (!(*(trt_state->engine))) {
return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, "TensorRT EP Failed to Build Engine.");
}
LOGS_DEFAULT(VERBOSE) << "[TensorRT EP] DeSerialized " + engine_cache_path;
@ -2527,7 +2498,7 @@ common::Status TensorrtExecutionProvider::Compile(const std::vector<FusedNodeAnd
// https://docs.nvidia.com/deeplearning/tensorrt/developer-guide/index.html#threading
trt_state->engine->reset();
*(trt_state->engine) = std::unique_ptr<nvinfer1::ICudaEngine>(trt_state->runtime->deserializeCudaEngine(engine_buf.get(), engine_size, nullptr));
if (*(trt_state->engine) == nullptr) {
if (!(*(trt_state->engine))) {
return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL,
"TensorRT EP could not deserialize engine from encrypted cache: " + encrypted_engine_cache_path);
}
@ -2556,10 +2527,7 @@ common::Status TensorrtExecutionProvider::Compile(const std::vector<FusedNodeAnd
// Regenerate engine
if (engine_update) {
// Destroy the IExecutionContext objects before destroying an engine object, otherwise it will lead to undefined behavior.
if (GetPerThreadContext().IsTensorRTContextInMap(fused_node_name)) {
GetPerThreadContext().ResetTensorRTContext(fused_node_name);
}
trt_state->context->reset();
trt_state->engine->reset();
auto trt_config = std::unique_ptr<nvinfer1::IBuilderConfig>(trt_builder->createBuilderConfig());
trt_config->setMaxWorkspaceSize(*(trt_state->max_workspace_size_ptr));
@ -2660,7 +2628,7 @@ common::Status TensorrtExecutionProvider::Compile(const std::vector<FusedNodeAnd
LOGS_DEFAULT(INFO) << "TensorRT engine build for " << trt_state->trt_node_name_with_precision << " took: " << std::chrono::duration_cast<std::chrono::milliseconds>(engine_build_stop - engine_build_start).count() << "ms" << std::endl;
}
}
if (*(trt_state->engine) == nullptr) {
if (!(*(trt_state->engine))) {
return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, "TensorRT EP Failed to Build Engine.");
}
trt_engine = trt_state->engine->get();
@ -2706,32 +2674,20 @@ common::Status TensorrtExecutionProvider::Compile(const std::vector<FusedNodeAnd
context_update = true;
}
// Build execution context if either of the following conditions is true:
// (1) The engine is built or updated by this thread.
// (2) The first inference run for this thread where there is no IExecutionContext object yet.
// (3) The engine is updated by another thread. (We compare the profile shapes maintained by the PerThreadContext to the profile shapes maintained by TRT EP)
//
// Note: Creating an execution context from an engine is thread safe per TRT doc
// https://docs.nvidia.com/deeplearning/tensorrt/developer-guide/index.html#threading
if (context_update ||
!GetPerThreadContext().IsTensorRTContextInMap(fused_node_name) ||
GetPerThreadContext().CompareProfileShapes(fused_node_name, shape_ranges)) {
std::unique_ptr<nvinfer1::IExecutionContext> new_context;
if (context_update) {
if (trt_state->context_memory_sharing_enable) {
new_context.reset(trt_state->engine->get()->createExecutionContextWithoutDeviceMemory());
*(trt_state->context) = std::unique_ptr<nvinfer1::IExecutionContext>(
trt_state->engine->get()->createExecutionContextWithoutDeviceMemory());
} else {
new_context.reset(trt_state->engine->get()->createExecutionContext());
*(trt_state->context) = std::unique_ptr<nvinfer1::IExecutionContext>(
trt_state->engine->get()->createExecutionContext());
}
auto context_status = GetPerThreadContext().UpdateTensorRTContext(fused_node_name, std::move(new_context));
if (!context_status) {
if (!(*(trt_state->context))) {
return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, "TensorRT EP failed to create context.");
}
GetPerThreadContext().UpdateProfileShapes(fused_node_name, shape_ranges);
trt_context = trt_state->context->get();
}
// Get the reference to the IExecutionContext object that is maintained on a per thread basis.
nvinfer1::IExecutionContext& trt_context = GetPerThreadContext().GetTensorRTContext(fused_node_name);
// Get input and output binding names
int total_bindings = trt_engine->getNbBindings();
std::vector<void*> buffers(total_bindings);
@ -2767,12 +2723,12 @@ common::Status TensorrtExecutionProvider::Compile(const std::vector<FusedNodeAnd
int nb_dims = dimensions.nbDims;
if (input_names.count(input_name) == 1) {
if (trt_engine->isShapeBinding(binding_index)) {
trt_context.setInputShapeBinding(binding_index, &tensor_shape_values[input_name][0]);
trt_context->setInputShapeBinding(binding_index, &tensor_shape_values[input_name][0]);
} else {
for (int j = 0, end = nb_dims; j < end; ++j) {
dimensions.d[j] = static_cast<int32_t>(tensor_shapes[j]);
}
const bool status = trt_context.setBindingDimensions(binding_index, dimensions);
const bool status = trt_context->setBindingDimensions(binding_index, dimensions);
if (!status) {
ORT_THROW_IF_ERROR(ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL,
"TensorRT EP cannot set the dynamic dimensions of a binding"));
@ -2911,7 +2867,7 @@ common::Status TensorrtExecutionProvider::Compile(const std::vector<FusedNodeAnd
if (index_iter != output_indexes.end()) {
output_index = index_iter->second;
}
nvinfer1::Dims dimensions = trt_context.getBindingDimensions(static_cast<int>(binding_index));
nvinfer1::Dims dimensions = trt_context->getBindingDimensions(static_cast<int>(binding_index));
int nb_dims = dimensions.nbDims;
std::vector<int64_t> output_shapes(nb_dims);
for (int j = 0, end = nb_dims; j < end; ++j) {
@ -3045,20 +3001,20 @@ common::Status TensorrtExecutionProvider::Compile(const std::vector<FusedNodeAnd
if (mem_size > *max_context_mem_size_ptr) {
*max_context_mem_size_ptr = mem_size;
}
trt_context.setDeviceMemory(IAllocator::MakeUniquePtrFromOrtAllocator<void>(alloc, *max_context_mem_size_ptr).get());
trt_context->setDeviceMemory(IAllocator::MakeUniquePtrFromOrtAllocator<void>(alloc, *max_context_mem_size_ptr).get());
}
// Start CUDA graph capture.
// Note: The reason we don't put graph capture in OnRunStart() like CUDA EP does is because
// current ORT TRT doesn't get cuda stream until compute time and graph capture requires cuda stream.
if (cuda_graph_enable_ && GetPerThreadContext().IsGraphCaptureAllowed() && !GetPerThreadContext().IsGraphCaptured()) {
if (cuda_graph_enable_ && IsGraphCaptureAllowed() && !IsGraphCaptured()) {
LOGS_DEFAULT(INFO) << "Capturing the cuda graph for this model";
GetPerThreadContext().SetGraphStream(stream);
GetPerThreadContext().CaptureBegin();
cuda_graph_.SetStream(stream);
CaptureBegin();
}
// Run TRT inference
if (!trt_context.enqueueV2(&buffers[0], stream, nullptr)) {
if (!trt_context->enqueueV2(&buffers[0], stream, nullptr)) {
return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "TensorRT EP execution context enqueue failed.");
}
@ -3089,14 +3045,14 @@ common::Status TensorrtExecutionProvider::Compile(const std::vector<FusedNodeAnd
// Note: One reason we don't put end of graph capture in OnRunEnd() like CUDA EP does is because of cuda stream mentioned in graph capture
// above, another reason is because OnRunEnd() is not synchronized with OnRunStart() and ExecuteGraph() per inference_session.cc.
// It's safe to start/end CUDA graph capture in compute_func() here since cuda graph object is maintained by a per thread basis.
if (cuda_graph_enable_ && !GetPerThreadContext().IsGraphCaptured()) {
if (GetPerThreadContext().IsGraphCaptureAllowed()) {
GetPerThreadContext().CaptureEnd();
if (cuda_graph_enable_ && !IsGraphCaptured()) {
if (IsGraphCaptureAllowed()) {
CaptureEnd();
// CUDA work issued to a capturing stream doesn’t actually run on the GPU,
// so run the captured graph here to actually execute the work.
ORT_RETURN_IF_ERROR(GetPerThreadContext().ReplayGraph());
ORT_RETURN_IF_ERROR(ReplayGraph());
} else {
GetPerThreadContext().IncrementRegularRunCountBeforeGraphCapture();
IncrementRegularRunCountBeforeGraphCapture();
}
}
@ -3116,8 +3072,8 @@ void TensorrtExecutionProvider::RegisterStreamHandlers(IStreamCommandHandleRegis
true /* release_cpu_buffer_on_cuda_stream */,
stream_,
external_stream_ /* use_existing_stream */,
GetPerThreadContext().CudnnHandle(),
GetPerThreadContext().CublasHandle());
external_cudnn_handle_,
external_cublas_handle_);
}
OrtDevice TensorrtExecutionProvider::GetOrtDeviceByMemType(OrtMemType mem_type) const {

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

@ -105,6 +105,7 @@ struct TensorrtFuncState {
std::string fused_node_name;
tensorrt_ptr::unique_pointer<nvonnxparser::IParser>* parser = nullptr;
std::unique_ptr<nvinfer1::ICudaEngine>* engine = nullptr;
std::unique_ptr<nvinfer1::IExecutionContext>* context = nullptr;
std::unique_ptr<nvinfer1::IBuilder>* builder = nullptr;
std::unique_ptr<nvinfer1::INetworkDefinition>* network = nullptr;
std::vector<std::unordered_map<std::string, size_t>> input_info;
@ -246,6 +247,7 @@ class TensorrtExecutionProvider : public IExecutionProvider {
// For those non thread safe operations, TRT EP uses (1) lock_guard or (2) PerThreadContext to make sure synchronization.
std::unordered_map<std::string, tensorrt_ptr::unique_pointer<nvonnxparser::IParser>> parsers_;
std::unordered_map<std::string, std::unique_ptr<nvinfer1::ICudaEngine>> engines_;
std::unordered_map<std::string, std::unique_ptr<nvinfer1::IExecutionContext>> contexts_;
std::unordered_map<std::string, std::unique_ptr<nvinfer1::IBuilder>> builders_;
std::unordered_map<std::string, std::unique_ptr<nvinfer1::INetworkDefinition>> networks_;
std::unordered_map<std::string, std::vector<std::unordered_map<std::string, size_t>>> input_info_;
@ -256,6 +258,21 @@ class TensorrtExecutionProvider : public IExecutionProvider {
std::unordered_map<std::string, ShapeRangesMap> input_shape_ranges_; // The profile shape ranges that the engine is built with
std::unordered_map<std::string, std::vector<nvinfer1::IOptimizationProfile*>> profiles_;
// for external stream, we need to create its cudnn/cublass handle before cuda EP enable cuda graph capture
cudnnHandle_t external_cudnn_handle_ = nullptr;
cublasHandle_t external_cublas_handle_ = nullptr;
CUDAGraph cuda_graph_;
bool is_graph_captured_ = false;
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_ = 1; // required min regular runs before graph capture for the necessary memory allocations.
// [Note] We don't use PerThreadContext for now since it has issue with multithreading
//
// TRT or CUDA objects that must be maintained on a per thread basis will be put under this PerThreadContext data structure.
// For example, TensorRT execution context and CUDA graph are the ones to be put here.
class PerThreadContext final {