From 73da3f3705003bf4e4d872047a4fabb8c1856243 Mon Sep 17 00:00:00 2001
From: Kevin Chen <45886021+kevinch-nv@users.noreply.github.com>
Date: Mon, 15 Aug 2022 08:22:50 -0700
Subject: [PATCH] Add TRT uint8 support (#12570)

* uint8 support

Signed-off-by: Kevin Chen <kevinch@nvidia.com>

* Handle outputs as well

Signed-off-by: Kevin Chen <kevinch@nvidia.com>

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

diff --git a/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.cc b/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.cc
index 03a8f6ac94..c8d426f1ec 100644
--- a/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.cc
+++ b/onnxruntime/core/providers/tensorrt/tensorrt_execution_provider.cc
@@ -1737,6 +1737,16 @@ common::Status TensorrtExecutionProvider::Compile(const std::vector<FusedNodeAnd
             }
             break;
           }
+          case ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT8: {
+            auto input_tensor_ptr = ort.GetTensorData<uint8_t>(input_tensor);
+            if (input_tensor_ptr == nullptr) {
+              scratch_buffers.push_back(IAllocator::MakeUniquePtr<void>(alloc, sizeof(uint8_t)));
+              buffers[binding_index] = scratch_buffers.back().get();
+            } else {
+              buffers[binding_index] = const_cast<uint8_t*>(input_tensor_ptr);
+            }
+            break;
+          }
           case ONNX_TENSOR_ELEMENT_DATA_TYPE_INT32: {
             auto input_tensor_ptr = ort.GetTensorData<int32_t>(input_tensor);
             if (input_tensor_ptr == nullptr) {
@@ -1870,6 +1880,16 @@ common::Status TensorrtExecutionProvider::Compile(const std::vector<FusedNodeAnd
             }
             break;
           }
+          case ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT8: {
+            auto output_tensor_ptr = ort.GetTensorMutableData<uint8_t>(output_tensor[i]);
+            if (output_tensor_ptr == nullptr) {
+              scratch_buffers.push_back(IAllocator::MakeUniquePtr<void>(alloc, sizeof(uint8_t)));
+              buffers[binding_index] = scratch_buffers.back().get();
+            } else {
+              buffers[binding_index] = output_tensor_ptr;
+            }
+            break;
+          }
           case ONNX_TENSOR_ELEMENT_DATA_TYPE_INT32: {
             auto output_tensor_ptr = ort.GetTensorMutableData<int32_t>(output_tensor[i]);
             if (output_tensor_ptr == nullptr) {