From 4d8fca0a4eedfbf2a07241d410e5dcf8ec9b5031 Mon Sep 17 00:00:00 2001
From: stevenlix <38092805+stevenlix@users.noreply.github.com>
Date: Fri, 10 Jan 2020 14:07:40 -0800
Subject: [PATCH] Fix memory leak in TRT (#2815)

* fix memory leak issue

* revert EP_FAIL on enueueV2
---
 .../tensorrt/tensorrt_execution_provider.cc      | 16 +++++++++++-----
 1 file changed, 11 insertions(+), 5 deletions(-)

diff --git a/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.cc b/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.cc
index 028a34e09d..1a56b8229e 100644
--- a/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.cc
+++ b/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.cc
@@ -676,7 +676,7 @@ common::Status TensorrtExecutionProvider::Compile(const std::vector<onnxruntime:
       // Update shape ranges
       bool dimension_update = false;
       auto trt_builder = trt_state->builder;
-      auto trt_profile = trt_builder->createOptimizationProfile();
+      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);
@@ -714,6 +714,9 @@ common::Status TensorrtExecutionProvider::Compile(const std::vector<onnxruntime:
           }
 
           if (dimension_update) {
+            if (trt_profile == nullptr) {
+              trt_profile = trt_builder->createOptimizationProfile();
+            }
             if (engine.isShapeBinding(i)) {
               std::vector<int32_t> shapes_min(nb_dims), shapes_opt(nb_dims), shapes_max(nb_dims);
               for (int j = 0, end = nb_dims; j < end; ++j) {
@@ -733,7 +736,7 @@ common::Status TensorrtExecutionProvider::Compile(const std::vector<onnxruntime:
         }
 
         // TensorRT6 requires optimization profile to be defined for all inputs if any input dimension is symbolic
-        if (dynamic_shape) {
+        if (dimension_update && dynamic_shape) {
           trt_profile->setDimensions(input->getName(), nvinfer1::OptProfileSelector::kMIN, dims_min);
           trt_profile->setDimensions(input->getName(), nvinfer1::OptProfileSelector::kOPT, dims_opt);
           trt_profile->setDimensions(input->getName(), nvinfer1::OptProfileSelector::kMAX, dims_max);
@@ -746,10 +749,13 @@ common::Status TensorrtExecutionProvider::Compile(const std::vector<onnxruntime:
         auto trt_config = unique_pointer<nvinfer1::IBuilderConfig>(trt_builder->createBuilderConfig());
         trt_config->addOptimizationProfile(trt_profile);
         trt_state->engine = trt_builder->buildEngineWithConfig(*trt_state->network, *trt_config);
-        ORT_ENFORCE(trt_state->engine != nullptr);
-
+        if (trt_state->engine == nullptr) {
+          return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, "TensorRT EP Failed to Build Engine.");
+        }
         trt_state->context = trt_state->engine->createExecutionContext();
-        ORT_ENFORCE(trt_state->context != nullptr);
+        if (trt_state->context == nullptr) {
+          return ORT_MAKE_STATUS(ONNXRUNTIME, EP_FAIL, "TensorRT EP Failed to Create Context.");
+        }
         trt_context = trt_state->context;
       }