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