From e6fff47d2b7e5c3551d23a06ce0e78f946c696f5 Mon Sep 17 00:00:00 2001
From: Kevin Chen <kevinch@nvidia.com>
Date: Thu, 13 Jun 2024 15:33:51 -0700
Subject: [PATCH 1/3] Add support for INT64 types in TensorRT constant layer
 calibration

Signed-off-by: Kevin Chen <kevinch@nvidia.com>
---
 .../providers/tensorrt/tensorrt_execution_provider.cc | 11 ++++++++++-
 1 file changed, 10 insertions(+), 1 deletion(-)

diff --git a/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.cc b/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.cc
index dff74a404a456..a0cdda99a29be 100644
--- a/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.cc
+++ b/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.cc
@@ -71,6 +71,7 @@ bool SetDynamicRange(nvinfer1::INetworkDefinition& network, std::unordered_map<s
     auto dynamic_range_iter = dynamic_range_map.find(tensor_name);
     if (dynamic_range_iter != dynamic_range_map.end()) {
       if (!network.getInput(i)->setDynamicRange(-dynamic_range_iter->second, dynamic_range_iter->second)) {
+        LOGS_DEFAULT(ERROR) << "Failed to set dynamic range for network input " << tensor_name;
         return false;
       }
     }
@@ -84,10 +85,12 @@ bool SetDynamicRange(nvinfer1::INetworkDefinition& network, std::unordered_map<s
       auto dynamic_range_iter = dynamic_range_map.find(tensor_name);
       if (dynamic_range_iter != dynamic_range_map.end()) {
         if (!trt_layer->getOutput(j)->setDynamicRange(-dynamic_range_iter->second, dynamic_range_iter->second)) {
+          LOGS_DEFAULT(ERROR) << "Failed to set dynamic range for tensor " << tensor_name;
           return false;
         }
       } else if (trt_layer->getType() == nvinfer1::LayerType::kCONSTANT) {
         nvinfer1::IConstantLayer* const_layer = static_cast<nvinfer1::IConstantLayer*>(trt_layer);
+        const std::string const_layer_name = const_layer->getName();
         auto trt_weights = const_layer->getWeights();
         double max_weight = std::numeric_limits<double>::min();
         for (int64_t k = 0, end = trt_weights.count; k < end; ++k) {
@@ -108,13 +111,19 @@ bool SetDynamicRange(nvinfer1::INetworkDefinition& network, std::unordered_map<s
             case nvinfer1::DataType::kINT32:
               weight = static_cast<const int32_t*>(trt_weights.values)[k];
               break;
+#if NV_TENSORRT_MAJOR >= 10
+            case nvinfer1::DataType::kINT64:
+              weight = static_cast<const int64_t*>(trt_weights.values)[k];
+              break;
+#endif // NV_TENSORRT_MAJOR >= 10
             default:
-              LOGS_DEFAULT(ERROR) << "Found unsupported datatype!";
+              LOGS_DEFAULT(ERROR) << "Found unsupported datatype for layer " << const_layer_name;
               return false;
           }
           max_weight = std::max(max_weight, std::abs(weight));
         }
         if (!trt_layer->getOutput(j)->setDynamicRange(static_cast<float>(-max_weight), static_cast<float>(max_weight))) {
+          LOGS_DEFAULT(ERROR) << "Failed to set dynamic range for layer " << const_layer_name;
           return false;
         }
       }

From ed1a45e0a7ea6b2196eaf677e7548b9ff69ffbcc Mon Sep 17 00:00:00 2001
From: Kevin Chen <kevinch@nvidia.com>
Date: Thu, 13 Jun 2024 16:28:46 -0700
Subject: [PATCH 2/3] Fix comment whitespace

Signed-off-by: Kevin Chen <kevinch@nvidia.com>
---
 .../core/providers/tensorrt/tensorrt_execution_provider.cc      | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.cc b/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.cc
index a0cdda99a29be..ac9f96fc0ed69 100644
--- a/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.cc
+++ b/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.cc
@@ -115,7 +115,7 @@ bool SetDynamicRange(nvinfer1::INetworkDefinition& network, std::unordered_map<s
             case nvinfer1::DataType::kINT64:
               weight = static_cast<const int64_t*>(trt_weights.values)[k];
               break;
-#endif // NV_TENSORRT_MAJOR >= 10
+#endif  // NV_TENSORRT_MAJOR >= 10
             default:
               LOGS_DEFAULT(ERROR) << "Found unsupported datatype for layer " << const_layer_name;
               return false;

From 521fffaee17456a7775af913fd557b1a2dbaf629 Mon Sep 17 00:00:00 2001
From: Kevin Chen <kevinch@nvidia.com>
Date: Tue, 18 Jun 2024 16:06:54 -0700
Subject: [PATCH 3/3] Fix warning

Signed-off-by: Kevin Chen <kevinch@nvidia.com>
---
 .../core/providers/tensorrt/tensorrt_execution_provider.cc    | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.cc b/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.cc
index ac9f96fc0ed69..2bf3d27928a72 100644
--- a/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.cc
+++ b/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.cc
@@ -113,8 +113,10 @@ bool SetDynamicRange(nvinfer1::INetworkDefinition& network, std::unordered_map<s
               break;
 #if NV_TENSORRT_MAJOR >= 10
             case nvinfer1::DataType::kINT64:
-              weight = static_cast<const int64_t*>(trt_weights.values)[k];
+            {
+              weight = static_cast<double>(static_cast<const int64_t*>(trt_weights.values)[k]);
               break;
+            }
 #endif  // NV_TENSORRT_MAJOR >= 10
             default:
               LOGS_DEFAULT(ERROR) << "Found unsupported datatype for layer " << const_layer_name;