diff --git a/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.cc b/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.cc index e721683ba9..a223ba729b 100644 --- a/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.cc +++ b/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.cc @@ -464,11 +464,11 @@ SubGraphCollection_t TensorrtExecutionProvider::GetSupportedList(SubGraphCollect // Get supported node list recursively SubGraphCollection_t parser_nodes_list; TensorrtLogger& trt_logger = GetTensorrtLogger(); - auto trt_builder = unique_pointer(nvinfer1::createInferBuilder(trt_logger)); + auto trt_builder = tensorrt_ptr::unique_pointer(nvinfer1::createInferBuilder(trt_logger)); const auto explicitBatch = 1U << static_cast(nvinfer1::NetworkDefinitionCreationFlag::kEXPLICIT_BATCH); - auto trt_network = unique_pointer(trt_builder->createNetworkV2(explicitBatch)); + auto trt_network = tensorrt_ptr::unique_pointer(trt_builder->createNetworkV2(explicitBatch)); - auto trt_parser = unique_pointer(nvonnxparser::createParser(*trt_network, trt_logger)); + auto trt_parser = tensorrt_ptr::unique_pointer(nvonnxparser::createParser(*trt_network, trt_logger)); trt_parser->supportsModel(string_buf.data(), string_buf.size(), parser_nodes_list); SubGraphCollection_t next_nodes_list; @@ -682,11 +682,11 @@ common::Status TensorrtExecutionProvider::Compile(const std::vector(nvinfer1::createInferBuilder(trt_logger)); + auto trt_builder = tensorrt_ptr::unique_pointer(nvinfer1::createInferBuilder(trt_logger)); const auto explicitBatch = 1U << static_cast(nvinfer1::NetworkDefinitionCreationFlag::kEXPLICIT_BATCH); - auto trt_network = unique_pointer(trt_builder->createNetworkV2(explicitBatch)); - auto trt_config = unique_pointer(trt_builder->createBuilderConfig()); - auto trt_parser = unique_pointer(nvonnxparser::createParser(*trt_network, trt_logger)); + auto trt_network = tensorrt_ptr::unique_pointer(trt_builder->createNetworkV2(explicitBatch)); + auto trt_config = tensorrt_ptr::unique_pointer(trt_builder->createBuilderConfig()); + auto trt_parser = tensorrt_ptr::unique_pointer(nvonnxparser::createParser(*trt_network, trt_logger)); trt_parser->parse(string_buf.data(), string_buf.size()); trt_config->setMaxWorkspaceSize(max_workspace_size_); @@ -734,14 +734,14 @@ common::Status TensorrtExecutionProvider::Compile(const std::vectorsetFlag(nvinfer1::BuilderFlag::kFP16); } - auto trt_engine = unique_pointer(trt_builder->buildEngineWithConfig(*trt_network, *trt_config)); + auto trt_engine = tensorrt_ptr::unique_pointer(trt_builder->buildEngineWithConfig(*trt_network, *trt_config)); if (trt_engine == nullptr) { return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, "TensorRT EP could not build Engine for fused node: " + fused_node->Name()); } // Build TensorRT context - auto trt_context = unique_pointer(trt_engine->createExecutionContext()); + auto trt_context = tensorrt_ptr::unique_pointer(trt_engine->createExecutionContext()); if (trt_context == nullptr) { return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, "TensorRT EP could not build Execution Context for fused node: " + fused_node->Name()); @@ -815,7 +815,7 @@ common::Status TensorrtExecutionProvider::Compile(const std::vector p = onnxruntime::make_unique(); *p = {context->allocate_func, context->release_func, context->allocator_handle, parsers_[context->node_name].get(), - engines_[context->node_name].get(), contexts_[context->node_name].get(), builders_[context->node_name].get(), + &engines_[context->node_name], &contexts_[context->node_name], builders_[context->node_name].get(), networks_[context->node_name].get(), input_info_[context->node_name], output_info_[context->node_name], input_shape_ranges_[context->node_name], output_shapes_[context->node_name], &tensorrt_mu_, &fp16_enable_, &max_workspace_size_}; @@ -845,7 +845,7 @@ common::Status TensorrtExecutionProvider::Compile(const std::vectorcontext; + auto trt_context = trt_state->context->get(); auto trt_builder = trt_state->builder; nvinfer1::IOptimizationProfile* trt_profile = nullptr; for (int i = 0, end = num_binding_inputs; i < end; ++i) { @@ -913,21 +913,26 @@ common::Status TensorrtExecutionProvider::Compile(const std::vector(trt_builder->createBuilderConfig()); + auto trt_config = tensorrt_ptr::unique_pointer(trt_builder->createBuilderConfig()); trt_config->setMaxWorkspaceSize(*(trt_state->max_workspace_size_ptr)); trt_config->addOptimizationProfile(trt_profile); if (*(trt_state->fp16_enable_ptr) && trt_builder->platformHasFastFp16()) { trt_config->setFlag(nvinfer1::BuilderFlag::kFP16); } - trt_state->engine = trt_builder->buildEngineWithConfig(*trt_state->network, *trt_config); - if (trt_state->engine == nullptr) { + trt_state->context->reset(); + trt_state->engine->reset(); + *(trt_state->engine) = tensorrt_ptr::unique_pointer( + trt_builder->buildEngineWithConfig(*trt_state->network, *trt_config)); + + if (trt_state->engine->get() == nullptr) { return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, "TensorRT EP Failed to Build Engine."); } - trt_state->context = trt_state->engine->createExecutionContext(); - if (trt_state->context == nullptr) { + *(trt_state->context) = tensorrt_ptr::unique_pointer( + trt_state->engine->get()->createExecutionContext()); + if (trt_state->context->get() == nullptr) { return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, "TensorRT EP Failed to Create Context."); } - trt_context = trt_state->context; + trt_context = trt_state->context->get(); } // Set input shapes and assign input buffers diff --git a/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.h b/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.h index c3a50ae0de..33d7eb1532 100644 --- a/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.h +++ b/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.h @@ -38,6 +38,21 @@ class TensorrtLogger : public nvinfer1::ILogger { } }; +namespace tensorrt_ptr { + + struct TensorrtInferDeleter { + template + void operator()(T* obj) const { + if (obj) { + obj->destroy(); + } + } + }; + + template + using unique_pointer = std::unique_ptr; +}; + // Information needed to construct trt execution providers. struct TensorrtExecutionProviderInfo { int device_id{0}; @@ -45,12 +60,13 @@ struct TensorrtExecutionProviderInfo { // Information to construct kernel function state. struct TensorrtFuncState { + AllocateFunc test_allocate_func = nullptr; DestroyFunc test_release_func = nullptr; AllocatorHandle allocator = nullptr; nvonnxparser::IParser* parser = nullptr; - nvinfer1::ICudaEngine* engine = nullptr; - nvinfer1::IExecutionContext* context = nullptr; + tensorrt_ptr::unique_pointer * engine = nullptr; + tensorrt_ptr::unique_pointer * context = nullptr; nvinfer1::IBuilder* builder = nullptr; nvinfer1::INetworkDefinition* network = nullptr; std::vector> input_info; @@ -89,25 +105,14 @@ class TensorrtExecutionProvider : public IExecutionProvider { bool fp16_enable_ = false; bool dump_subgraphs_ = false; - struct InferDeleter { - template - void operator()(T* obj) const { - if (obj) { - obj->destroy(); - } - } - }; - - template - using unique_pointer = std::unique_ptr; OrtMutex tensorrt_mu_; int device_id_; - std::unordered_map> parsers_; - std::unordered_map> engines_; - std::unordered_map> contexts_; - std::unordered_map> builders_; - std::unordered_map> networks_; + 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_; std::unordered_map>> output_info_; std::unordered_map>>> input_shape_ranges_;