From 73037978f85d0a5d4b1f601101dbf65e04a43456 Mon Sep 17 00:00:00 2001 From: Chi Lo <54722500+chilo-ms@users.noreply.github.com> Date: Tue, 8 Aug 2023 13:02:34 -0700 Subject: [PATCH] Add PerThreadContext for TRT EP (#16599) Maintaining one execution context on a per thread basis is suggested per TRT [doc](https://docs.nvidia.com/deeplearning/tensorrt/developer-guide/index.html#threading) to avoid synchronization issue. For previous TRT EP, we did see synchronization issues when running multithreading on some models, for example, FasterRCNN. This PR leverages per thread context implementation from CUDA EP. Followings are the modifications: - Move CUDA graph and IExecutionContext objects to per thread context. - Remove lock_gruad that previously placed for the whole compute_func() and put lock_gruad in the blocks where multiple threads may update kernel function state, access one builder, create/serialize/save engine, save profile and serialize/save timing cache. - On CentOS, don't unload TRT EP shared library and leave it around, so that destructor of thread local data is still accessible upon thread exits. Note: Tested this PR with onnxruntime_perf_test and the overhead of PerThreadContext is small. --- .../tensorrt/tensorrt_execution_provider.cc | 731 +++++++++++------- .../tensorrt/tensorrt_execution_provider.h | 131 +++- .../core/session/provider_bridge_ort.cc | 24 +- .../providers/tensorrt/tensorrt_basic_test.cc | 74 +- 4 files changed, 670 insertions(+), 290 deletions(-) diff --git a/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.cc b/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.cc index ca6b25b75b..3f0cfdac8a 100644 --- a/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.cc +++ b/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.cc @@ -376,7 +376,7 @@ bool ApplyProfileShapesFromProviderOptions(std::vector>>& profile_min_shapes, std::unordered_map>>& profile_max_shapes, std::unordered_map>>& profile_opt_shapes, - std::unordered_map>>>& input_explicit_shape_ranges) { + ShapeRangesMap& input_explicit_shape_ranges) { if (trt_profiles.size() == 0) { LOGS_DEFAULT(WARNING) << "[TensorRT EP] Number of optimization profiles should be greater than 0, but it's 0."; return false; @@ -485,7 +485,7 @@ bool ApplyProfileShapesFromProviderOptions(std::vector& trt_profiles, Ort::KernelContext ctx, nvinfer1::ITensor* input, - std::unordered_map>>>& shape_ranges, + ShapeRangesMap& shape_ranges, const std::unordered_map& input_indexes, std::unordered_map>& tensor_shape_values, cudaStream_t stream, @@ -648,6 +648,142 @@ Status ApplyProfileShapesFromInputTensorValue(std::vector 0) { + if (input_shape_ranges_[fused_node] != shape_ranges) { + LOGS_DEFAULT(VERBOSE) << "[TensorRT EP] The shape ranges maintained by the PerThreadContext is different from the shape ranges maintained by TRT EP. \ + This means the engine is updated and will need to update the execution context as well."; + return true; + } + } + return false; +} + +/* + * Updates the shape ranges maintained by the PerThreadContext. + * As long as the execution context maintained by the PerThreadContext is updated, the associated shape ranges should be updated as well. + * + */ +void TensorrtExecutionProvider::PerThreadContext::UpdateProfileShapes(std::string fused_node, ShapeRangesMap& shape_ranges) { + input_shape_ranges_[fused_node] = shape_ranges; +} + +void TensorrtExecutionProvider::PerThreadContext::ResetTensorRTContext(std::string fused_node) { + auto it = trt_context_map_.find(fused_node); + if (it != trt_context_map_.end()) { + trt_context_map_[fused_node].reset(); + } +} + +bool TensorrtExecutionProvider::PerThreadContext::UpdateTensorRTContext(std::string fused_node, std::unique_ptr context) { + if (!context) { + context = std::make_unique(); + } + trt_context_map_[fused_node] = std::move(context); + + if (trt_context_map_[fused_node]) { + return true; + } + return false; +} + +bool TensorrtExecutionProvider::PerThreadContext::IsTensorRTContextInMap(std::string fused_node) { + auto it = trt_context_map_.find(fused_node); + if (it != trt_context_map_.end()) { + return true; + } + return false; +} + +nvinfer1::IExecutionContext& TensorrtExecutionProvider::PerThreadContext::GetTensorRTContext(std::string fused_node) { + auto it = trt_context_map_.find(fused_node); + if (it != trt_context_map_.end()) { + return *(it->second); // dereference shared pointer + } + auto context = std::make_unique(); + trt_context_map_[fused_node] = std::move(context); + return *(trt_context_map_[fused_node]); // dereference shared pointer +} + +void TensorrtExecutionProvider::ReleasePerThreadContext() const { + const auto& per_thread_context_cache = PerThreadContextCache(); + + auto cached_context_it = per_thread_context_cache->find(this); + ORT_ENFORCE(cached_context_it != per_thread_context_cache->end()); + auto cached_context = cached_context_it->second.lock(); + ORT_ENFORCE(cached_context); + + { + std::lock_guard lock(context_state_.mutex); + context_state_.active_contexts.erase(cached_context); + context_state_.retired_context_pool.push_back(cached_context); + } + + per_thread_context_cache->erase(cached_context_it); +} + +TensorrtExecutionProvider::PerThreadContext& TensorrtExecutionProvider::GetPerThreadContext() const { + const auto& per_thread_context_cache = PerThreadContextCache(); + + // try to use cached context + auto cached_context_it = per_thread_context_cache->find(this); + if (cached_context_it != per_thread_context_cache->end()) { + auto cached_context = cached_context_it->second.lock(); + ORT_ENFORCE(cached_context); + return *cached_context; + } + + // get context and update cache + std::shared_ptr context; + { + std::lock_guard lock(context_state_.mutex); + + // get or create a context + if (context_state_.retired_context_pool.empty()) { + context = std::make_shared(info_.device_id, info_.has_user_compute_stream, stream_); + } else { + context = context_state_.retired_context_pool.back(); + context_state_.retired_context_pool.pop_back(); + } + + // insert into active_contexts, should not already be present + const auto active_contexts_insert_result = context_state_.active_contexts.insert(context); + ORT_ENFORCE(active_contexts_insert_result.second); + + // insert into caches_to_update_on_destruction, may already be present + ORT_IGNORE_RETURN_VALUE(context_state_.caches_to_update_on_destruction.insert(per_thread_context_cache)); + } + + per_thread_context_cache->insert(std::make_pair(this, context)); + + return *context; +} + TensorrtExecutionProvider::TensorrtExecutionProvider(const TensorrtExecutionProviderInfo& info) : IExecutionProvider{onnxruntime::kTensorrtExecutionProvider, OrtDevice(OrtDevice::GPU, OrtDevice::MemType::DEFAULT, info.device_id), true}, info_(info), device_id_(info.device_id) { InitProviderOrtApi(); @@ -656,10 +792,6 @@ TensorrtExecutionProvider::TensorrtExecutionProvider(const TensorrtExecutionProv if (info.has_user_compute_stream) { external_stream_ = true; stream_ = static_cast(info.user_compute_stream); - ORT_IGNORE_RETURN_VALUE(CUBLAS_CALL(cublasCreate(&external_cublas_handle_))); - ORT_IGNORE_RETURN_VALUE(CUBLAS_CALL(cublasSetStream(external_cublas_handle_, stream_))); - ORT_IGNORE_RETURN_VALUE(CUDNN_CALL(cudnnCreate(&external_cudnn_handle_))); - ORT_IGNORE_RETURN_VALUE(CUDNN_CALL(cudnnSetStream(external_cudnn_handle_, stream_))); } std::string profile_min_shapes, profile_max_shapes, profile_opt_shapes; @@ -903,7 +1035,7 @@ TensorrtExecutionProvider::TensorrtExecutionProvider(const TensorrtExecutionProv } if (cuda_graph_enable_) { - cuda_graph_ = std::make_unique(); + GetPerThreadContext().InitCUDAGraph(); } /* @@ -984,10 +1116,16 @@ TensorrtExecutionProvider::TensorrtExecutionProvider(const TensorrtExecutionProv } TensorrtExecutionProvider::~TensorrtExecutionProvider() { - if (external_stream_) { - ORT_IGNORE_RETURN_VALUE(CUBLAS_CALL(cublasDestroy(external_cublas_handle_))); - ORT_IGNORE_RETURN_VALUE(CUDNN_CALL(cudnnDestroy(external_cudnn_handle_))); + // clean up thread local context caches + { + std::lock_guard lock(context_state_.mutex); + for (const auto& cache_weak : context_state_.caches_to_update_on_destruction) { + const auto cache = cache_weak.lock(); + if (!cache) continue; + ORT_IGNORE_RETURN_VALUE(cache->erase(this)); + } } + if (!external_stream_ && stream_) { ORT_IGNORE_RETURN_VALUE(CUDA_CALL(cudaStreamDestroy(stream_))); } @@ -1004,35 +1142,50 @@ bool TensorrtExecutionProvider::IsGraphCaptureEnabled() const { return cuda_graph_enable_; } -bool TensorrtExecutionProvider::IsGraphCaptureAllowed() const { +bool TensorrtExecutionProvider::IsGraphCaptured() const { + return GetPerThreadContext().IsGraphCaptured(); +} + +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); +} + +bool TensorrtExecutionProvider::PerThreadContext::IsGraphCaptureAllowed() const { return regular_run_count_before_graph_capture_ >= min_num_runs_before_cuda_graph_capture_; } -void TensorrtExecutionProvider::CaptureBegin() { +void TensorrtExecutionProvider::PerThreadContext::CaptureBegin() { cuda_graph_->Reset(); cuda_graph_->CaptureBegin(); } -void TensorrtExecutionProvider::CaptureEnd() { +void TensorrtExecutionProvider::PerThreadContext::CaptureEnd() { cuda_graph_->CaptureEnd(); is_graph_captured_ = true; } -bool TensorrtExecutionProvider::IsGraphCaptured() const { +bool TensorrtExecutionProvider::PerThreadContext::IsGraphCaptured() const { return is_graph_captured_; } -Status TensorrtExecutionProvider::ReplayGraph() { +Status TensorrtExecutionProvider::PerThreadContext::ReplayGraph() { ORT_ENFORCE(IsGraphCaptured()); // Please note that CUDAGraph::Replay() is not thread safe. - // ORT TRT calls ReplayGraph() in compute_func() where synchromization is enforced due to lock_guard(), + // 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(); } -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(), +void TensorrtExecutionProvider::PerThreadContext::IncrementRegularRunCountBeforeGraphCapture() { + // The cuda graph object is maintained by a per thread basis, // therefore following increment is guaranteed to be thread safe. ++regular_run_count_before_graph_capture_; } @@ -1063,6 +1216,18 @@ 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(); } @@ -1847,8 +2012,8 @@ common::Status TensorrtExecutionProvider::Compile(const std::vector>>> input_explicit_shape_ranges; - std::unordered_map>>> input_implicit_shape_ranges; + ShapeRangesMap input_explicit_shape_ranges; + ShapeRangesMap input_implicit_shape_ranges; if ((!profile_min_shapes_.empty()) && (!profile_max_shapes_.empty()) && (!profile_opt_shapes_.empty())) { has_explicit_profile = true; @@ -2174,6 +2339,8 @@ common::Status TensorrtExecutionProvider::Compile(const std::vectorgetDeviceMemorySize(); if (mem_size > max_ctx_mem_size_) { @@ -2183,7 +2350,7 @@ common::Status TensorrtExecutionProvider::Compile(const std::vector(trt_engine->createExecutionContext()); } - if (trt_context == nullptr) { + if (!trt_context) { return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, "TensorRT EP could not build execution context for fused node: " + fused_node.Name()); } @@ -2211,10 +2378,9 @@ common::Status TensorrtExecutionProvider::Compile(const std::vectorallocate_func, context->release_func, context->allocator_handle, &parsers_[context->node_name], - &engines_[context->node_name], &contexts_[context->node_name], &builders_[context->node_name], + *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], &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_, @@ -2256,19 +2430,19 @@ common::Status TensorrtExecutionProvider::Compile(const std::vector(state); - std::lock_guard lock(*(trt_state->tensorrt_mu_ptr)); const std::unordered_map& input_indexes = (trt_state->input_info)[0]; const std::unordered_map& output_indexes = (trt_state->output_info)[0]; const std::unordered_map& output_types = (trt_state->output_info)[1]; + auto fused_node_name = trt_state->fused_node_name; 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(input_indexes.size()); int num_outputs = static_cast(output_indexes.size()); bool engine_update = false; + bool context_update = false; std::unordered_set input_names; std::unordered_map> tensor_shape_values; @@ -2288,7 +2462,7 @@ common::Status TensorrtExecutionProvider::Compile(const std::vectorengine_cache_path, trt_state->trt_node_name_with_precision); const std::string engine_cache_path = cache_path + "_sm" + compute_capability + ".engine"; const std::string profile_cache_path = cache_path + "_sm" + compute_capability + ".profile"; @@ -2296,248 +2470,261 @@ common::Status TensorrtExecutionProvider::Compile(const std::vectorengine_cache_enable && trt_engine == nullptr) { - std::ifstream engine_file(engine_cache_path, std::ios::binary | std::ios::in); - std::ifstream profile_file(profile_cache_path, std::ios::binary | std::ios::in); - if (engine_file && profile_file) { - // Deserialize profile - shape_ranges = DeserializeProfileV2(profile_file); - LOGS_DEFAULT(VERBOSE) << "[TensorRT EP] DeSerialized " + profile_cache_path; - // Deserialize engine - trt_state->context->reset(); + + // Following block is the critical section where multiple threads may update kernel function state, access one builder, create/serialize/save engine, + // save profile and serialize/save timing cache. Therefore, those operations should be synchronized across different threads when ORT is using multithreading. + // More details here, https://docs.nvidia.com/deeplearning/tensorrt/developer-guide/index.html#threading + { + std::lock_guard lock(*(trt_state->tensorrt_mu_ptr)); + + // Load serialized engine + if (trt_state->engine_cache_enable && trt_engine == nullptr) { + std::ifstream engine_file(engine_cache_path, std::ios::binary | std::ios::in); + std::ifstream profile_file(profile_cache_path, std::ios::binary | std::ios::in); + if (engine_file && profile_file) { + // Deserialize profile + shape_ranges = DeserializeProfileV2(profile_file); + LOGS_DEFAULT(VERBOSE) << "[TensorRT EP] DeSerialized " + profile_cache_path; + + // Prepare buffer + engine_file.seekg(0, std::ios::end); + size_t engine_size = engine_file.tellg(); + engine_file.seekg(0, std::ios::beg); + std::unique_ptr engine_buf{new char[engine_size]}; + engine_file.read((char*)engine_buf.get(), engine_size); + + // Deserialize engine + // Note: Deserializing an engine from a TensorRT runtime is thread safe per TRT doc + // https://docs.nvidia.com/deeplearning/tensorrt/developer-guide/index.html#threading + trt_state->engine->reset(); + *(trt_state->engine) = std::unique_ptr( + trt_state->runtime->deserializeCudaEngine(engine_buf.get(), engine_size, nullptr)); + if (*(trt_state->engine) == nullptr) { + return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, "TensorRT EP Failed to Build Engine."); + } + LOGS_DEFAULT(VERBOSE) << "[TensorRT EP] DeSerialized " + engine_cache_path; + trt_engine = trt_state->engine->get(); + context_update = true; + } else if (trt_state->engine_decryption_enable && !engine_file && profile_file) { + shape_ranges = DeserializeProfileV2(profile_file); + LOGS_DEFAULT(VERBOSE) << "[TensorRT EP] DeSerialized " + profile_cache_path; + // Decrypt engine + size_t engine_size = 0; + if (!trt_state->engine_decryption(engine_cache_path.c_str(), nullptr, &engine_size)) { + return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, + "TensorRT EP could not get engine buffer size"); + } + std::unique_ptr engine_buf{new char[engine_size]}; + if (!trt_state->engine_decryption(engine_cache_path.c_str(), &engine_buf[0], &engine_size)) { + return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, + "TensorRT EP could not call engine decryption function decrypt"); + } + // Deserialize engine + // Note: Deserializing an engine from a TensorRT runtime is thread safe per TRT doc + // https://docs.nvidia.com/deeplearning/tensorrt/developer-guide/index.html#threading + trt_state->engine->reset(); + *(trt_state->engine) = std::unique_ptr(trt_state->runtime->deserializeCudaEngine(engine_buf.get(), engine_size, nullptr)); + if (*(trt_state->engine) == nullptr) { + return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, + "TensorRT EP could not deserialize engine from encrypted cache: " + engine_cache_path); + } + LOGS_DEFAULT(VERBOSE) << "[TensorRT EP] DeSerialized " + engine_cache_path; + trt_engine = trt_state->engine->get(); + context_update = true; + } + } + + // Check and update shape ranges for dynamic shape inputs. + for (int i = 0, end = num_inputs; i < end; ++i) { + auto input = trt_state->network->get()->getInput(i); + const std::string& input_name = input->getName(); + input_names.insert(input_name); + + // If there is any input tensor in shape_ranges, it means this input tensor has dynamic shape and its profile shape values have not yet resolved. + // TRT EP will help determine the min/max/opt profile values based on current input tensor value. + if (shape_ranges.find(input_name) != shape_ranges.end()) { + auto status = ApplyProfileShapesFromInputTensorValue(trt_profiles, ctx, input, shape_ranges, input_indexes, tensor_shape_values, stream, &engine_update); + if (status != Status::OK()) { + return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, "TensorRT EP failed to parse input tensor and generate optimization profiles."); + } + } + } + + // 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->engine->reset(); - engine_file.seekg(0, std::ios::end); - size_t engine_size = engine_file.tellg(); - engine_file.seekg(0, std::ios::beg); - std::unique_ptr engine_buf{new char[engine_size]}; - engine_file.read((char*)engine_buf.get(), engine_size); - *(trt_state->engine) = std::unique_ptr( - trt_state->runtime->deserializeCudaEngine(engine_buf.get(), engine_size, nullptr)); + auto trt_config = std::unique_ptr(trt_builder->createBuilderConfig()); + trt_config->setMaxWorkspaceSize(*(trt_state->max_workspace_size_ptr)); + for (auto trt_profile : trt_profiles) { + trt_config->addOptimizationProfile(trt_profile); + } + + // Set INT8 Per Tensor Dynamic range + if (trt_state->int8_enable && trt_builder->platformHasFastInt8() && trt_state->int8_calibration_cache_available) { + trt_config->setInt8Calibrator(nullptr); + if (!SetDynamicRange(*trt_state->network->get(), trt_state->dynamic_range_map)) { + return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, "TensorRT EP failed to set INT8 dynamic range."); + } + } + + // Set precision + if (trt_state->fp16_enable && trt_state->int8_enable) { + trt_config->setFlags(1U << static_cast(nvinfer1::BuilderFlag::kFP16) | 1U << static_cast(nvinfer1::BuilderFlag::kINT8)); + } else if (trt_state->fp16_enable) { + trt_config->setFlag(nvinfer1::BuilderFlag::kFP16); + } else if (trt_state->int8_enable) { + trt_config->setFlag(nvinfer1::BuilderFlag::kINT8); + } + + // Set DLA (DLA can only run with FP16 or INT8) + if ((trt_state->fp16_enable || trt_state->int8_enable) && trt_state->dla_enable) { + LOGS_DEFAULT(VERBOSE) << "[TensorRT EP] use DLA core " << trt_state->dla_core; + trt_config->setFlag(nvinfer1::BuilderFlag::kGPU_FALLBACK); + trt_config->setDefaultDeviceType(nvinfer1::DeviceType::kDLA); + trt_config->setDLACore(trt_state->dla_core); + } + + // enable sparse weights + if (trt_state->sparsity_enable) { + trt_config->setFlag(nvinfer1::BuilderFlag::kSPARSE_WEIGHTS); + LOGS_DEFAULT(VERBOSE) << "[TensorRT EP] Sparse weights are allowed"; + } + + // enable builder heuristics + if (trt_state->build_heuristics_enable) { + trt_config->setFlag(nvinfer1::BuilderFlag::kENABLE_TACTIC_HEURISTIC); + LOGS_DEFAULT(VERBOSE) << "[TensorRT EP] Builder heuristics are enabled"; + } +#if NV_TENSORRT_MINOR > 5 && NV_TENSORRT_MAJOR >= 8 + // switch optimizaion level + if (trt_state->builder_optimization_level != 3) { + trt_config->setBuilderOptimizationLevel(trt_state->builder_optimization_level); + LOGS_DEFAULT(VERBOSE) << "[TensorRT EP] Builder optimization level is set to " << builder_optimization_level_; + } + + // limit auxiliary streams + if (trt_state->auxiliary_streams >= 0) { + trt_config->setMaxAuxStreams(trt_state->auxiliary_streams); + LOGS_DEFAULT(VERBOSE) << "[TensorRT EP] Auxiliary streams are se to " << trt_state->auxiliary_streams; + } +#else + if (trt_state->builder_optimization_level != 3) { + LOGS_DEFAULT(WARNING) << "[TensorRT EP] Builder optimization level can only be used on TRT 8.6 onwards!"; + } + if (trt_state->auxiliary_streams >= 0) { + LOGS_DEFAULT(WARNING) << "[TensorRT EP] Auxiliary streams can only be set on TRT 8.6 onwards!"; + } +#endif + // limit used tactic sources + if (trt_state->filter_tactic_sources) { + nvinfer1::TacticSources tactics = trt_config->getTacticSources(); + tactics |= trt_state->tactic_sources; + trt_config->setTacticSources(tactics); + LOGS_DEFAULT(VERBOSE) << "[TensorRT EP] Tactic sources are limited using bitmask " << tactics; + } + + // Load timing cache from file. Create a fresh cache if the file doesn't exist + std::unique_ptr timing_cache = nullptr; + if (trt_state->timing_cache_enable) { + std::vector loaded_timing_cache = loadTimingCacheFile(timing_cache_path); + timing_cache.reset(trt_config->createTimingCache(static_cast(loaded_timing_cache.data()), loaded_timing_cache.size())); + if (timing_cache == nullptr) { + return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, + "TensorRT EP could not create timing cache: " + timing_cache_path); + } + trt_config->setTimingCache(*timing_cache, force_timing_cache_match_); + if (detailed_build_log_) { + LOGS_DEFAULT(VERBOSE) << "[TensorRT EP] Deserialized timing cache from " + timing_cache_path; + } + } + + // Build engine + { + auto lock = GetApiLock(); + std::chrono::steady_clock::time_point engine_build_start; + if (detailed_build_log_) { + engine_build_start = std::chrono::steady_clock::now(); + } + *(trt_state->engine) = std::unique_ptr( + trt_builder->buildEngineWithConfig(*trt_state->network->get(), *trt_config)); + if (detailed_build_log_) { + auto engine_build_stop = std::chrono::steady_clock::now(); + LOGS_DEFAULT(INFO) << "TensorRT engine build for " << trt_state->trt_node_name_with_precision << " took: " << std::chrono::duration_cast(engine_build_stop - engine_build_start).count() << "ms" << std::endl; + } + } if (*(trt_state->engine) == nullptr) { return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, "TensorRT EP Failed to Build Engine."); } - LOGS_DEFAULT(VERBOSE) << "[TensorRT EP] DeSerialized " + engine_cache_path; trt_engine = trt_state->engine->get(); - if (trt_state->context_memory_sharing_enable) { - *(trt_state->context) = std::unique_ptr( - trt_state->engine->get()->createExecutionContextWithoutDeviceMemory()); - } else { - *(trt_state->context) = std::unique_ptr( - trt_state->engine->get()->createExecutionContext()); - } - if (*(trt_state->context) == nullptr) { - return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, "TensorRT EP failed to create context."); - } - trt_context = trt_state->context->get(); - } else if (trt_state->engine_decryption_enable && !engine_file && profile_file) { - shape_ranges = DeserializeProfileV2(profile_file); - LOGS_DEFAULT(VERBOSE) << "[TensorRT EP] DeSerialized " + profile_cache_path; - // Decrypt engine - size_t engine_size = 0; - if (!trt_state->engine_decryption(engine_cache_path.c_str(), nullptr, &engine_size)) { - return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, - "TensorRT EP could not get engine buffer size"); - } - std::unique_ptr engine_buf{new char[engine_size]}; - if (!trt_state->engine_decryption(engine_cache_path.c_str(), &engine_buf[0], &engine_size)) { - return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, - "TensorRT EP could not call engine decryption function decrypt"); - } - // Deserialize engine - trt_state->context->reset(); - trt_state->engine->reset(); - *(trt_state->engine) = std::unique_ptr(trt_state->runtime->deserializeCudaEngine(engine_buf.get(), engine_size, nullptr)); - if (*(trt_state->engine) == nullptr) { - return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, - "TensorRT EP could not deserialize engine from encrypted cache: " + engine_cache_path); - } - LOGS_DEFAULT(VERBOSE) << "[TensorRT EP] DeSerialized " + engine_cache_path; - trt_engine = trt_state->engine->get(); - if (trt_state->context_memory_sharing_enable) { - *(trt_state->context) = std::unique_ptr( - trt_state->engine->get()->createExecutionContextWithoutDeviceMemory()); - } else { - *(trt_state->context) = std::unique_ptr( - trt_state->engine->get()->createExecutionContext()); - } - if (*(trt_state->context) == nullptr) { - return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, "TensorRT EP failed to create context."); - } - trt_context = trt_state->context->get(); - } - } + if (trt_state->engine_cache_enable) { + // Serialize engine profile + SerializeProfileV2(profile_cache_path, shape_ranges); + LOGS_DEFAULT(VERBOSE) << "[TensorRT EP] Serialized " + profile_cache_path; - // Check and update shape ranges for dynamic shape inputs. - for (int i = 0, end = num_inputs; i < end; ++i) { - auto input = trt_state->network->get()->getInput(i); - const std::string& input_name = input->getName(); - input_names.insert(input_name); - - // If there is any input tensor in shape_ranges, it means this input tensor has dynamic shape and its profile shape values have not yet resolved. - // TRT EP will help determine the min/max/opt profile values based on current input tensor value. - if (shape_ranges.find(input_name) != shape_ranges.end()) { - auto status = ApplyProfileShapesFromInputTensorValue(trt_profiles, ctx, input, shape_ranges, input_indexes, tensor_shape_values, stream, &engine_update); - if (status != Status::OK()) { - return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, "TensorRT EP failed to parse input tensor and generate optimization profiles."); - } - } - } - - // Regenerate engine - if (engine_update) { - trt_state->context->reset(); - trt_state->engine->reset(); - auto trt_config = std::unique_ptr(trt_builder->createBuilderConfig()); - trt_config->setMaxWorkspaceSize(*(trt_state->max_workspace_size_ptr)); - for (auto trt_profile : trt_profiles) { - trt_config->addOptimizationProfile(trt_profile); - } - - // Set INT8 Per Tensor Dynamic range - if (trt_state->int8_enable && trt_builder->platformHasFastInt8() && trt_state->int8_calibration_cache_available) { - trt_config->setInt8Calibrator(nullptr); - if (!SetDynamicRange(*trt_state->network->get(), trt_state->dynamic_range_map)) { - return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, "TensorRT EP failed to set INT8 dynamic range."); - } - } - - // Set precision - if (trt_state->fp16_enable && trt_state->int8_enable) { - trt_config->setFlags(1U << static_cast(nvinfer1::BuilderFlag::kFP16) | 1U << static_cast(nvinfer1::BuilderFlag::kINT8)); - } else if (trt_state->fp16_enable) { - trt_config->setFlag(nvinfer1::BuilderFlag::kFP16); - } else if (trt_state->int8_enable) { - trt_config->setFlag(nvinfer1::BuilderFlag::kINT8); - } - - // Set DLA (DLA can only run with FP16 or INT8) - if ((trt_state->fp16_enable || trt_state->int8_enable) && trt_state->dla_enable) { - LOGS_DEFAULT(VERBOSE) << "[TensorRT EP] use DLA core " << trt_state->dla_core; - trt_config->setFlag(nvinfer1::BuilderFlag::kGPU_FALLBACK); - trt_config->setDefaultDeviceType(nvinfer1::DeviceType::kDLA); - trt_config->setDLACore(trt_state->dla_core); - } - - // enable sparse weights - if (trt_state->sparsity_enable) { - trt_config->setFlag(nvinfer1::BuilderFlag::kSPARSE_WEIGHTS); - LOGS_DEFAULT(VERBOSE) << "[TensorRT EP] Sparse weights are allowed"; - } - - // enable builder heuristics - if (trt_state->build_heuristics_enable) { - trt_config->setFlag(nvinfer1::BuilderFlag::kENABLE_TACTIC_HEURISTIC); - LOGS_DEFAULT(VERBOSE) << "[TensorRT EP] Builder heuristics are enabled"; - } -#if NV_TENSORRT_MINOR > 5 && NV_TENSORRT_MAJOR >= 8 - // switch optimizaion level - if (trt_state->builder_optimization_level != 3) { - trt_config->setBuilderOptimizationLevel(trt_state->builder_optimization_level); - LOGS_DEFAULT(VERBOSE) << "[TensorRT EP] Builder optimization level is set to " << builder_optimization_level_; - } - - // limit auxiliary streams - if (trt_state->auxiliary_streams >= 0) { - trt_config->setMaxAuxStreams(trt_state->auxiliary_streams); - LOGS_DEFAULT(VERBOSE) << "[TensorRT EP] Auxiliary streams are se to " << trt_state->auxiliary_streams; - } -#else - if (trt_state->builder_optimization_level != 3) { - LOGS_DEFAULT(WARNING) << "[TensorRT EP] Builder optimization level can only be used on TRT 8.6 onwards!"; - } - if (trt_state->auxiliary_streams >= 0) { - LOGS_DEFAULT(WARNING) << "[TensorRT EP] Auxiliary streams can only be set on TRT 8.6 onwards!"; - } -#endif - // limit used tactic sources - if (trt_state->filter_tactic_sources) { - nvinfer1::TacticSources tactics = trt_config->getTacticSources(); - tactics |= trt_state->tactic_sources; - trt_config->setTacticSources(tactics); - LOGS_DEFAULT(VERBOSE) << "[TensorRT EP] Tactic sources are limited using bitmask " << tactics; - } - - // Load timing cache from file. Create a fresh cache if the file doesn't exist - std::unique_ptr timing_cache = nullptr; - if (trt_state->timing_cache_enable) { - std::vector loaded_timing_cache = loadTimingCacheFile(timing_cache_path); - timing_cache.reset(trt_config->createTimingCache(static_cast(loaded_timing_cache.data()), loaded_timing_cache.size())); - if (timing_cache == nullptr) { - return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, - "TensorRT EP could not create timing cache: " + timing_cache_path); - } - trt_config->setTimingCache(*timing_cache, force_timing_cache_match_); - if (detailed_build_log_) { - LOGS_DEFAULT(VERBOSE) << "[TensorRT EP] Deserialized timing cache from " + timing_cache_path; - } - } - - // Build engine - { - auto lock = GetApiLock(); - std::chrono::steady_clock::time_point engine_build_start; - if (detailed_build_log_) { - engine_build_start = std::chrono::steady_clock::now(); - } - *(trt_state->engine) = std::unique_ptr( - trt_builder->buildEngineWithConfig(*trt_state->network->get(), *trt_config)); - if (detailed_build_log_) { - auto engine_build_stop = std::chrono::steady_clock::now(); - LOGS_DEFAULT(INFO) << "TensorRT engine build for " << trt_state->trt_node_name_with_precision << " took: " << std::chrono::duration_cast(engine_build_stop - engine_build_start).count() << "ms" << std::endl; - } - } - if (*(trt_state->engine) == nullptr) { - return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, "TensorRT EP Failed to Build Engine."); - } - trt_engine = trt_state->engine->get(); - if (trt_state->engine_cache_enable) { - // Serialize engine profile - SerializeProfileV2(profile_cache_path, shape_ranges); - LOGS_DEFAULT(VERBOSE) << "[TensorRT EP] Serialized " + profile_cache_path; - - // Serialize engine - std::unique_ptr serializedModel(trt_engine->serialize()); - size_t engine_size = serializedModel->size(); - if (trt_state->engine_decryption_enable) { - // Encrypt engine - if (!trt_state->engine_encryption(engine_cache_path.c_str(), reinterpret_cast(serializedModel->data()), engine_size)) { - return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, - "TensorRT EP could not call engine encryption function encrypt"); + // Serialize engine + std::unique_ptr serializedModel(trt_engine->serialize()); + size_t engine_size = serializedModel->size(); + if (trt_state->engine_decryption_enable) { + // Encrypt engine + if (!trt_state->engine_encryption(engine_cache_path.c_str(), reinterpret_cast(serializedModel->data()), engine_size)) { + return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, + "TensorRT EP could not call engine encryption function encrypt"); + } + } else { + std::ofstream file(engine_cache_path, std::ios::binary | std::ios::out); + file.write(reinterpret_cast(serializedModel->data()), engine_size); } - } else { - std::ofstream file(engine_cache_path, std::ios::binary | std::ios::out); - file.write(reinterpret_cast(serializedModel->data()), engine_size); + LOGS_DEFAULT(VERBOSE) << "[TensorRT EP] Serialized " + engine_cache_path; } - } - // serialize and save timing cache - if (trt_state->timing_cache_enable) { - auto timing_cache = trt_config->getTimingCache(); - std::unique_ptr timingCacheHostData{timing_cache->serialize()}; - if (timingCacheHostData == nullptr) { - return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, - "TensorRT EP could not serialize timing cache: " + timing_cache_path); - } - saveTimingCacheFile(timing_cache_path, timingCacheHostData.get()); - if (detailed_build_log_) { - LOGS_DEFAULT(VERBOSE) << "[TensorRT EP] Serialized timing cache " + timing_cache_path; + // serialize and save timing cache + if (trt_state->timing_cache_enable) { + auto timing_cache = trt_config->getTimingCache(); + std::unique_ptr timingCacheHostData{timing_cache->serialize()}; + if (timingCacheHostData == nullptr) { + return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, + "TensorRT EP could not serialize timing cache: " + timing_cache_path); + } + saveTimingCacheFile(timing_cache_path, timingCacheHostData.get()); + if (detailed_build_log_) { + LOGS_DEFAULT(VERBOSE) << "[TensorRT EP] Serialized timing cache " + timing_cache_path; + } } + context_update = true; } + } - // Build context + // 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 new_context; if (trt_state->context_memory_sharing_enable) { - *(trt_state->context) = std::unique_ptr( - trt_state->engine->get()->createExecutionContextWithoutDeviceMemory()); + new_context.reset(trt_state->engine->get()->createExecutionContextWithoutDeviceMemory()); } else { - *(trt_state->context) = std::unique_ptr( - trt_state->engine->get()->createExecutionContext()); + new_context.reset(trt_state->engine->get()->createExecutionContext()); } - if (*(trt_state->context) == nullptr) { + auto context_status = GetPerThreadContext().UpdateTensorRTContext(fused_node_name, std::move(new_context)); + if (!context_status) { return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, "TensorRT EP failed to create context."); } - trt_context = trt_state->context->get(); + GetPerThreadContext().UpdateProfileShapes(fused_node_name, shape_ranges); } + // 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 buffers(total_bindings); @@ -2573,12 +2760,12 @@ common::Status TensorrtExecutionProvider::Compile(const std::vectorisShapeBinding(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(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")); @@ -2717,7 +2904,7 @@ common::Status TensorrtExecutionProvider::Compile(const std::vectorsecond; } - nvinfer1::Dims dimensions = trt_context->getBindingDimensions(static_cast(binding_index)); + nvinfer1::Dims dimensions = trt_context.getBindingDimensions(static_cast(binding_index)); int nb_dims = dimensions.nbDims; std::vector output_shapes(nb_dims); for (int j = 0, end = nb_dims; j < end; ++j) { @@ -2851,20 +3038,20 @@ common::Status TensorrtExecutionProvider::Compile(const std::vector *max_context_mem_size_ptr) { *max_context_mem_size_ptr = mem_size; } - trt_context->setDeviceMemory(IAllocator::MakeUniquePtrFromOrtAllocator(alloc, *max_context_mem_size_ptr).get()); + trt_context.setDeviceMemory(IAllocator::MakeUniquePtrFromOrtAllocator(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_ && IsGraphCaptureAllowed() && !IsGraphCaptured()) { + if (cuda_graph_enable_ && GetPerThreadContext().IsGraphCaptureAllowed() && !GetPerThreadContext().IsGraphCaptured()) { LOGS_DEFAULT(INFO) << "Capturing the cuda graph for this model"; - cuda_graph_->SetStream(stream); - CaptureBegin(); + GetPerThreadContext().SetGraphStream(stream); + GetPerThreadContext().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."); } @@ -2893,17 +3080,16 @@ common::Status TensorrtExecutionProvider::Compile(const std::vector using unique_pointer = std::unique_ptr; }; // namespace tensorrt_ptr +using ShapeRangesMap = std::unordered_map>>>; + // Information to construct kernel function state. struct TensorrtFuncState { AllocateFunc test_allocate_func = nullptr; DestroyFunc test_release_func = nullptr; AllocatorHandle allocator = nullptr; + std::string fused_node_name; tensorrt_ptr::unique_pointer* parser = nullptr; std::unique_ptr* engine = nullptr; - std::unique_ptr* context = nullptr; std::unique_ptr* builder = nullptr; std::unique_ptr* network = nullptr; std::vector> input_info; @@ -153,6 +155,14 @@ class TensorrtExecutionProvider : public IExecutionProvider { explicit TensorrtExecutionProvider(const TensorrtExecutionProviderInfo& info); virtual ~TensorrtExecutionProvider(); + cublasHandle_t PerThreadDefaultCublasHandle() { + return GetPerThreadContext().CublasHandle(); + } + + cudnnHandle_t PerThreadDefaultCudnnHandle() { + return GetPerThreadContext().CudnnHandle(); + } + virtual std::shared_ptr GetKernelRegistry() const override; std::unique_ptr GetDataTransfer() const override; @@ -227,16 +237,15 @@ class TensorrtExecutionProvider : public IExecutionProvider { // and should be kept for the lifetime of TRT EP object. OrtAllocator* alloc_ = nullptr; - std::unique_ptr cuda_graph_; // 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) - bool is_graph_captured_ = false; - int regular_run_count_before_graph_capture_ = 0; - const int min_num_runs_before_cuda_graph_capture_ = 1; // required min regular runs before graph capture for the necessary memory allocations. - std::unordered_set control_flow_op_set_ = {"If", "Loop", "Scan"}; mutable std::unordered_map> subgraph_context_map_; + + // Following maps that hold TRT objects will be accessible by different threads if ORT is using multithreading. + // In general, TensorRT objects are not thread safe; accesses to an object from different threads must be serialized by the client. + // But there are still some thread safe operations, please see here https://docs.nvidia.com/deeplearning/tensorrt/developer-guide/index.html#threading + // For those non thread safe operations, TRT EP uses (1) lock_guard or (2) PerThreadContext to make sure synchronization. std::unordered_map> parsers_; std::unordered_map> engines_; - std::unordered_map> contexts_; std::unordered_map> builders_; std::unordered_map> networks_; std::unordered_map>> input_info_; @@ -244,12 +253,112 @@ class TensorrtExecutionProvider : public IExecutionProvider { std::unordered_map>> profile_min_shapes_; std::unordered_map>> profile_max_shapes_; std::unordered_map>> profile_opt_shapes_; - std::unordered_map>>>> input_shape_ranges_; + std::unordered_map input_shape_ranges_; // The profile shape ranges that the engine is built with std::unordered_map> 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; + // 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 { + public: + PerThreadContext(OrtDevice::DeviceId device_id, bool has_user_compute_stream, cudaStream_t stream); + ~PerThreadContext(); + + cublasHandle_t CublasHandle() const { + return external_cublas_handle_; + } + + cudnnHandle_t CudnnHandle() const { + return external_cudnn_handle_; + } + + bool IsTensorRTContextInMap(std::string fused_node); + nvinfer1::IExecutionContext& GetTensorRTContext(std::string fused_node); + bool UpdateTensorRTContext(std::string fused_node, std::unique_ptr context); + void ResetTensorRTContext(std::string fused_node); + bool CompareProfileShapes(std::string fused_node, ShapeRangesMap& shape_ranges); + void UpdateProfileShapes(std::string fused_node, ShapeRangesMap& shape_ranges); + + void InitCUDAGraph(); + void SetGraphStream(cudaStream_t stream); + bool IsGraphCaptureAllowed() const; + void CaptureBegin(); + void CaptureEnd(); + bool IsGraphCaptured() const; + Status ReplayGraph(); + void IncrementRegularRunCountBeforeGraphCapture(); + + private: + cudnnHandle_t external_cudnn_handle_ = nullptr; + cublasHandle_t external_cublas_handle_ = nullptr; + + // Maintaining execution context on a per thread basis is suggested by TRT doc. + // Also, for enqueueV2() in execution context, to perform inference concurrently in multiple streams, use one execution context per stream. + // ORT multi-streams feature uses one stream for one thread, therefore maintaining execution context on a per thread basis is necessary for TRT EP, + // otherwise it may result in undefined behavior or synchronization issues. + // + // See more details here: + // https://docs.nvidia.com/deeplearning/tensorrt/developer-guide/index.html#threading + // https://docs.nvidia.com/deeplearning/tensorrt/api/c_api/classnvinfer1_1_1_i_execution_context.html#a63cd95430852038ce864e17c670e0b36 + std::unordered_map> trt_context_map_; + + // The profile shape ranges for the engine that the execution context maintained by the PerThreadContext is built with. + // TRT EP needs this info to determine whether to rebuild the execution context. + 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_; + bool is_graph_captured_ = false; + int regular_run_count_before_graph_capture_ = -1; + // 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. + }; + + using PerThreadContextMap = std::unordered_map>; + // thread local PerThreadContext cache + + struct ContextCacheHolder { + ContextCacheHolder() { + // Keep a weak pointer to the object, if the weak pointer can be locked, then the shared pointer is still around, so we can reset it + RunOnUnload([&, weak_p_ = std::weak_ptr(p)] { + if (auto lock = weak_p_.lock()) { + p.reset(); + } + }); + } + + std::shared_ptr p = std::make_shared(); + }; + + static const std::shared_ptr& PerThreadContextCache() { + thread_local const ContextCacheHolder per_thread_context_cache; + return per_thread_context_cache.p; + } + + struct PerThreadContextState { + // contexts that are currently active + std::set, std::owner_less>> active_contexts; + // contexts available for reuse + std::vector> retired_context_pool; + // weak references to thread local caches from which this TensorrtExecutionProvider instance's entry should be removed + // upon destruction + std::set, std::owner_less>> + caches_to_update_on_destruction; + // synchronizes access to PerThreadContextState members + OrtMutex mutex; + }; + + // The execution provider maintains the PerThreadContexts in this structure. + // Synchronization is required to update the contained structures. + // On the other hand, access to an individual PerThreadContext is assumed to be from a single thread at a time, + // so synchronization is not required for that. + mutable PerThreadContextState context_state_; + + PerThreadContext& GetPerThreadContext() const; + void ReleasePerThreadContext() const; /**Get IndexedSubGraph based on node list of the subgraph*/ std::unique_ptr GetSubGraph(SubGraph_t graph_nodes_index, diff --git a/onnxruntime/core/session/provider_bridge_ort.cc b/onnxruntime/core/session/provider_bridge_ort.cc index fe3fae718d..7357feccbf 100644 --- a/onnxruntime/core/session/provider_bridge_ort.cc +++ b/onnxruntime/core/session/provider_bridge_ort.cc @@ -1263,7 +1263,29 @@ static ProviderLibrary s_library_rocm(LIBRARY_PREFIX ORT_TSTR("onnxruntime_provi ); static ProviderLibrary s_library_dnnl(LIBRARY_PREFIX ORT_TSTR("onnxruntime_providers_dnnl") LIBRARY_EXTENSION); static ProviderLibrary s_library_openvino(LIBRARY_PREFIX ORT_TSTR("onnxruntime_providers_openvino") LIBRARY_EXTENSION); -static ProviderLibrary s_library_tensorrt(LIBRARY_PREFIX ORT_TSTR("onnxruntime_providers_tensorrt") LIBRARY_EXTENSION); +static ProviderLibrary s_library_tensorrt(LIBRARY_PREFIX ORT_TSTR("onnxruntime_providers_tensorrt") LIBRARY_EXTENSION +#ifndef _WIN32 + , + /* + * unload - On CentOS if we unload the TensorRT shared provider we crash. + * + * The reason is TensorRT EP holds a thread local data which won't be destroyed + * until thread exits which happens after TRT EP destruction. Upon thread exits, + * the destructor of thread local data will be called but the address of destructor + * is invalid since the destructor is defined in TRT EP which is already been + * removed from the address space. Therefore, we will hit a segmentation fault. + * So, here we won't unload the TensorRT shared provider and leave the library around. + * This way the OS/CRT/etc doesn't ever clean up until the process exits. + * + * Interestingly, TensorRT shared library won't crash on Ubuntu and Windows when being unloaded. + * The destructor of thread local data can be successfully called upon thread exits. + * One thing worth attention is, on Unix, the function to unload a library is allowed to do nothing. + * Please see here: https://pubs.opengroup.org/onlinepubs/007904975/functions/dlclose.html + * + */ + false +#endif +); static ProviderLibrary s_library_migraphx(LIBRARY_PREFIX ORT_TSTR("onnxruntime_providers_migraphx") LIBRARY_EXTENSION); void UnloadSharedProviders() { diff --git a/onnxruntime/test/providers/tensorrt/tensorrt_basic_test.cc b/onnxruntime/test/providers/tensorrt/tensorrt_basic_test.cc index f84ab3c89b..288cdfca2b 100644 --- a/onnxruntime/test/providers/tensorrt/tensorrt_basic_test.cc +++ b/onnxruntime/test/providers/tensorrt/tensorrt_basic_test.cc @@ -38,7 +38,8 @@ void VerifyOutputs(const std::vector& fetches, const std::vector& fetches, const std::vector dims) { +void CreateBaseModel(std::string model_name, + std::string graph_name, + std::vector dims, + bool add_non_zero_node = false) { onnxruntime::Model model(graph_name, false, DefaultLoggingManager().DefaultLogger()); auto& graph = model.MainGraph(); std::vector inputs; @@ -80,10 +95,27 @@ void CreateBaseModel(std::string model_name, std::string graph_name, std::vector inputs.clear(); inputs.push_back(&output_arg); inputs.push_back(&input_arg_3); - auto& output_arg_2 = graph.GetOrCreateNodeArg("M", &float_tensor); - outputs.clear(); - outputs.push_back(&output_arg_2); - graph.AddNode("node_2", "Add", "node 2.", inputs, outputs); + + if (add_non_zero_node) { + auto& output_arg_2 = graph.GetOrCreateNodeArg("node_2_out_1", &float_tensor); + outputs.clear(); + outputs.push_back(&output_arg_2); + graph.AddNode("node_2", "Add", "node 2.", inputs, outputs); + + inputs.clear(); + inputs.push_back(&output_arg_2); + ONNX_NAMESPACE::TypeProto int_tensor; + int_tensor.mutable_tensor_type()->set_elem_type(ONNX_NAMESPACE::TensorProto_DataType_INT64); + auto& output_arg_3 = graph.GetOrCreateNodeArg("M", &int_tensor); + outputs.clear(); + outputs.push_back(&output_arg_3); + graph.AddNode("node_3", "NonZero", "node 3.", inputs, outputs); + } else { + auto& output_arg_2 = graph.GetOrCreateNodeArg("M", &float_tensor); + outputs.clear(); + outputs.push_back(&output_arg_2); + graph.AddNode("node_2", "Add", "node 2.", inputs, outputs); + } auto status = graph.Resolve(); ASSERT_TRUE(status.IsOK()); @@ -102,6 +134,18 @@ void RunSession(InferenceSession& session_object, VerifyOutputs(fetches, expected_dims, expected_values); } +void RunSession2(InferenceSession& session_object, + RunOptions& run_options, + NameMLValMap& feeds, + std::vector output_names, + std::vector expected_dims, + std::vector expected_values) { + std::vector fetches; + auto status = session_object.Run(run_options, feeds, output_names, &fetches); + ASSERT_TRUE(status.IsOK()); + VerifyOutputs(fetches, expected_dims, expected_values); +} + void RunWithOneSessionSingleThreadInference(std::string model_name, std::string sess_log_id) { SessionOptions so; so.session_logid = sess_log_id; @@ -184,7 +228,7 @@ void RunWithOneSessionSingleThreadInference(std::string model_name, std::string RunSession(session_object, run_options, feeds, output_names, expected_dims_mul_m, expected_values_mul_m); } -void RunWithOneSessionMultiThreadsInference(std::string model_name, std::string sess_log_id) { +void RunWithOneSessionMultiThreadsInference(std::string model_name, std::string sess_log_id, bool has_non_zero_node = false) { SessionOptions so; so.session_logid = sess_log_id; RunOptions run_options; @@ -212,6 +256,8 @@ void RunWithOneSessionMultiThreadsInference(std::string model_name, std::string // prepare expected inputs and outputs std::vector expected_dims_mul_m = {1, 3, 2}; std::vector expected_values_mul_m = {3.0f, 6.0f, 9.0f, 12.0f, 15.0f, 18.0f}; + std::vector expected_dims_nonzero_m = {3, 6}; + std::vector expected_values_nonzero_m = {0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 2, 2, 0, 1, 0, 1, 0, 1}; OrtTensorRTProviderOptionsV2 params{ 0, @@ -266,8 +312,12 @@ void RunWithOneSessionMultiThreadsInference(std::string model_name, std::string std::vector threads; int num_thread = 5; - for (int i = 0; i < num_thread; ++i) - threads.push_back(std::thread(RunSession, std::ref(session_object), std::ref(run_options), std::ref(feeds), std::ref(output_names), std::ref(expected_dims_mul_m), std::ref(expected_values_mul_m))); + for (int i = 0; i < num_thread; ++i) { + if (has_non_zero_node) + threads.push_back(std::thread(RunSession2, std::ref(session_object), std::ref(run_options), std::ref(feeds), std::ref(output_names), std::ref(expected_dims_nonzero_m), std::ref(expected_values_nonzero_m))); + else + threads.push_back(std::thread(RunSession, std::ref(session_object), std::ref(run_options), std::ref(feeds), std::ref(output_names), std::ref(expected_dims_mul_m), std::ref(expected_values_mul_m))); + } for (auto& th : threads) th.join(); @@ -298,6 +348,12 @@ TEST(TensorrtExecutionProviderTest, SessionCreationWithSingleThreadAndInferenceW CreateBaseModel(model_name, graph_name, dims); RunWithOneSessionMultiThreadsInference(model_name, sess_log_id); + + // In addition to the test case that whole model can be run by TRT, we also need to test the case where + // the model is partitioned into TRT EP and CUDA EP subgraphs. + // We did observe synchronization issue for TRT EP without PerContextThread implementation running those models. + CreateBaseModel(model_name, graph_name, dims, true); + RunWithOneSessionMultiThreadsInference(model_name, sess_log_id, true); } // Test loading same model in different way, when hash id is generated via model name/model content/env metadata