diff --git a/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.cc b/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.cc index 1eabe1561e..27ed2be881 100644 --- a/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.cc +++ b/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.cc @@ -1243,7 +1243,7 @@ common::Status TensorrtExecutionProvider::Compile(const std::vector& fuse &engines_[context->node_name], &contexts_[context->node_name], &builders_[context->node_name], &networks_[context->node_name], input_info_[context->node_name], output_info_[context->node_name], input_shape_ranges_[context->node_name], &tensorrt_mu_, &fp16_enable_, &int8_enable_, &max_workspace_size_, - trt_node_name_with_precision, engine_cache_enable_, cache_path_, runtime_, + trt_node_name_with_precision, engine_cache_enable_, cache_path_, runtime_, trt_profile_, allocator_, dynamic_range_map, engine_decryption_enable_, engine_decryption_}; *state = p.release(); return 0; @@ -1254,7 +1254,6 @@ common::Status TensorrtExecutionProvider::Compile(const std::vector& fuse if (state) delete static_cast(state); }; - // Create compute function compute_info.compute_func = [this](FunctionState state, const OrtCustomOpApi* api, OrtKernelContext* context) { Ort::CustomOpApi ort{*api}; @@ -1267,13 +1266,14 @@ common::Status TensorrtExecutionProvider::Compile(const std::vector& fuse auto trt_builder = trt_state->builder->get(); auto trt_engine = trt_state->engine->get(); auto trt_context = trt_state->context->get(); + auto trt_profile = trt_state->trt_profile; auto alloc = trt_state->scratch_allocator; int num_inputs = input_indexes.size(); int num_outputs = output_indexes.size(); bool engine_update = false; - std::unordered_map dimension_update; + std::unordered_set input_names; std::unordered_map> tensor_shape_values; - nvinfer1::IOptimizationProfile* trt_profile = nullptr; + cudaStream_t stream = static_cast(this->GetComputeStream()); // Load serialized engine @@ -1348,9 +1348,8 @@ common::Status TensorrtExecutionProvider::Compile(const std::vector& fuse const std::string& input_name = input->getName(); nvinfer1::Dims dims = input->getDimensions(); int nb_dims = dims.nbDims; - // Check and update shape ranges for dynamic shape inputs - dimension_update[input_name] = false; + input_names.insert(input_name); if (shape_ranges.find(input_name) != shape_ranges.end()) { int input_index = 0; const auto& iter = input_indexes.find(input_name); @@ -1411,14 +1410,14 @@ common::Status TensorrtExecutionProvider::Compile(const std::vector& fuse if (tensor_shape_value < shape_range[j].first) { shape_range[j].first = tensor_shape_value; shapes_min[j] = tensor_shape_value; - dimension_update[input_name] = true; + engine_update = true; } // Update shape range upper bound if (tensor_shape_value > shape_range[j].second) { shape_range[j].second = tensor_shape_value; shapes_max[j] = tensor_shape_value; shapes_opt[j] = tensor_shape_value; - dimension_update[input_name] = true; + engine_update = true; } } } else { @@ -1431,7 +1430,7 @@ common::Status TensorrtExecutionProvider::Compile(const std::vector& fuse shapes_opt[j] = tensor_shape_value; shapes_max[j] = tensor_shape_value; } - dimension_update[input_name] = true; + engine_update = true; } if (trt_profile == nullptr) { @@ -1440,6 +1439,7 @@ common::Status TensorrtExecutionProvider::Compile(const std::vector& fuse trt_profile->setShapeValues(input_name.c_str(), nvinfer1::OptProfileSelector::kMIN, &shapes_min[0], shape_size); trt_profile->setShapeValues(input_name.c_str(), nvinfer1::OptProfileSelector::kOPT, &shapes_opt[0], shape_size); trt_profile->setShapeValues(input_name.c_str(), nvinfer1::OptProfileSelector::kMAX, &shapes_max[0], shape_size); + } else { // Execution tensor nvinfer1::Dims dims_min(dims), dims_opt(dims), dims_max(dims); for (int j = 0, end = nb_dims; j < end; ++j) { @@ -1453,14 +1453,14 @@ common::Status TensorrtExecutionProvider::Compile(const std::vector& fuse if (tensor_shape < shape_range[j].first) { shape_range[j].first = tensor_shape; dims_min.d[j] = tensor_shape; - dimension_update[input_name] = true; + engine_update = true; } // Update maximum dimension if (tensor_shape > shape_range[j].second) { shape_range[j].second = tensor_shape; dims_max.d[j] = tensor_shape; dims_opt.d[j] = tensor_shape; - dimension_update[input_name] = true; + engine_update = true; } } } @@ -1474,10 +1474,6 @@ common::Status TensorrtExecutionProvider::Compile(const std::vector& fuse } ort.ReleaseTensorTypeAndShapeInfo(tensor_info); } - - if (!engine_update && dimension_update[input_name]) { - engine_update = true; - } } // Regenerate engine @@ -1568,7 +1564,7 @@ common::Status TensorrtExecutionProvider::Compile(const std::vector& fuse // Set dynamic shapes nvinfer1::Dims dimensions = trt_engine->getBindingDimensions(static_cast(binding_index)); int nb_dims = dimensions.nbDims; - if (dimension_update.find(input_name) != dimension_update.end()) { + if (input_names.count(input_name) == 1) { if (trt_engine->isShapeBinding(binding_index)) { trt_context->setInputShapeBinding(binding_index, &tensor_shape_values[input_name][0]); } else { diff --git a/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.h b/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.h index 5974139f90..29b03954b0 100644 --- a/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.h +++ b/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.h @@ -100,6 +100,8 @@ struct TensorrtFuncState { bool engine_cache_enable; std::string engine_cache_path; nvinfer1::IRuntime* runtime = nullptr; + + nvinfer1::IOptimizationProfile* trt_profile = nullptr; AllocatorPtr scratch_allocator; std::unordered_map dynamic_range_map; bool engine_decryption_enable; @@ -148,6 +150,7 @@ class TensorrtExecutionProvider : public IExecutionProvider { bool engine_cache_enable_ = false; std::string cache_path_; nvinfer1::IRuntime* runtime_ = nullptr; + nvinfer1::IOptimizationProfile* trt_profile_ = nullptr; OrtMutex tensorrt_mu_; int device_id_; AllocatorPtr allocator_;