From 28f693a8a13c3c1a3b563e7e8e75b633bae1f27a Mon Sep 17 00:00:00 2001
From: stevenlix <38092805+stevenlix@users.noreply.github.com>
Date: Mon, 11 May 2020 16:11:41 -0700
Subject: [PATCH] Update TensorRT dynamic shape profile when input shape
 changed during runtime (#3904)

* Update dynamic shape range when input shape changed during runtime

* Update tensorrt_execution_provider.cc

* Update tensorrt_execution_provider.cc
---
 .../tensorrt/tensorrt_execution_provider.cc   | 21 ++++++++++---------
 1 file changed, 11 insertions(+), 10 deletions(-)

diff --git a/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.cc b/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.cc
index b89adf20c8..e70105afa4 100644
--- a/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.cc
+++ b/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.cc
@@ -695,9 +695,7 @@ common::Status TensorrtExecutionProvider::Compile(const std::vector<onnxruntime:
     for (unsigned int i = 0, end = trt_network->getNbInputs(); i < end; ++i) {
       auto input = trt_network->getInput(i);
       nvinfer1::Dims dims = input->getDimensions();
-      nvinfer1::Dims dims_min = dims;
-      nvinfer1::Dims dims_opt = dims;
-      nvinfer1::Dims dims_max = dims;
+      nvinfer1::Dims dims_min(dims), dims_opt(dims), dims_max(dims);
 
       int nb_dims = dims.nbDims;
       if (input->isShapeTensor()) {  // Shape tensor
@@ -851,16 +849,14 @@ common::Status TensorrtExecutionProvider::Compile(const std::vector<onnxruntime:
       auto trt_builder = trt_state->builder;
       nvinfer1::IOptimizationProfile* trt_profile = nullptr;
       for (int i = 0, end = num_binding_inputs; i < end; ++i) {
-        // TODO: check if getInput indexing is same with binding index
-        auto input = trt_state->network->getInput(i);
-        nvinfer1::Dims dims = input->getDimensions();
-        nvinfer1::Dims dims_min = dims;
-        nvinfer1::Dims dims_opt = dims;
-        nvinfer1::Dims dims_max = dims;
-
         // Check and update shape ranges for dynamic shape inputs
         auto& shape_ranges = trt_state->input_shape_ranges;
         if (shape_ranges.find(i) != shape_ranges.end()) {
+          // TODO: check if getInput indexing is same with binding index
+          auto input = trt_state->network->getInput(i);
+          nvinfer1::Dims dims = input->getDimensions();
+          nvinfer1::Dims dims_min(dims), dims_opt(dims), dims_max(dims);
+
           const OrtValue* input_tensor = ort.KernelContext_GetInput(context, input_indexes[i]);
           auto tensor_info = ort.GetTensorTypeAndShape(input_tensor);
           const auto& tensor_shape = ort.GetTensorShape(tensor_info);
@@ -870,10 +866,15 @@ common::Status TensorrtExecutionProvider::Compile(const std::vector<onnxruntime:
           for (int j = 0, end = nb_dims; j < end; ++j) {
             auto& shape_range = shape_ranges[i];
             if (shape_range.find(j) != shape_range.end()) {
+              dims_min.d[j] = shape_range[j].first;
+              dims_opt.d[j] = shape_range[j].second;
+              dims_max.d[j] = shape_range[j].second;
+
               // Update minimum dimension
               if (tensor_shape[j] < shape_range[j].first) {
                 shape_range[j].first = tensor_shape[j];
                 dims_min.d[j] = tensor_shape[j];
+                dims_opt.d[j] = tensor_shape[j];
                 dimension_update = true;
               }
               // Update maximum dimension