Repository: pytorch/TensorRT

Issue:
✨[Feature] Support INT64 inputs at the graph input level
**Is your feature request related to a problem? Please describe.**
Currently if the graph is given int64 inputs, the compilation fails with unsupported type. 


**Describe the solution you'd like**
Solution we are expecting : Cast INT64 inputs to INT32 and issue a warning to users about this instead of breaking compilation.

**Describe alternatives you've considered**
<!--A clear and concise description of any alternative solutions or features you've considered.-->

**Additional context**
<!--Add any other context or screenshots about the feature request here.-->



Discussion:
This feature may warrant a multi-phase detailed RFC, as it relates closely to Issue #1346, and the current (Draft) implementation can cause issues for users requiring `aten::scatter` support or other Torch-only Long-only operations. The linked PR #1551 implements autocasting in a way which is efficient (no-copying) but can break existing models due to issues with Int64-required operations, since it in-place casts input tensors for the duration of inference. I think additional design work is needed to track data types through different engines and implement an efficient casting system.


RFC: #1553 
Feature split between MVP & some optional phases (detailed in RFC), trying to complete MVP for this sprint.

Resolving patch:
diff --git a/py/torch_tensorrt/_Input.py b/py/torch_tensorrt/_Input.py
--- a/py/torch_tensorrt/_Input.py
+++ b/py/torch_tensorrt/_Input.py
@@ -219,7 +219,9 @@ def _supported_input_size_type(input_size: Any) -> bool:
     @staticmethod
     def _parse_dtype(dtype: Any) -> _enums.dtype:
         if isinstance(dtype, torch.dtype):
-            if dtype == torch.int32:
+            if dtype == torch.long:
+                return _enums.dtype.long
+            elif dtype == torch.int32:
                 return _enums.dtype.int32
             elif dtype == torch.half:
                 return _enums.dtype.half
@@ -229,7 +231,7 @@ def _parse_dtype(dtype: Any) -> _enums.dtype:
                 return _enums.dtype.bool
             else:
                 raise TypeError(
-                    "Provided an unsupported data type as an input data type (support: bool, int32, half, float), got: "
+                    "Provided an unsupported data type as an input data type (support: bool, int32, long, half, float), got: "
                     + str(dtype)
                 )
 
@@ -242,6 +244,9 @@ def _parse_dtype(dtype: Any) -> _enums.dtype:
                 + str(type(dtype))
             )
 
+    def is_trt_dtype(self) -> bool:
+        return self.dtype != _enums.dtype.long
+
     @staticmethod
     def _parse_format(format: Any) -> _enums.TensorFormat:
         if isinstance(format, torch.memory_format):
diff --git a/py/torch_tensorrt/ts/_compile_spec.py b/py/torch_tensorrt/ts/_compile_spec.py
--- a/py/torch_tensorrt/ts/_compile_spec.py
+++ b/py/torch_tensorrt/ts/_compile_spec.py
@@ -214,6 +214,14 @@ def _parse_input_signature(input_signature: Any):
             if isinstance(input_signature, torch.Tensor)
             else input_signature
         )
+
+        if not i.is_trt_dtype():
+            raise TypeError(
+                "Using non-TRT input types with input_signature is not currently "
+                + "supported. Please specify inputs individually to use "
+                + "non-TRT types."
+            )
+
         clone = _internal_input_to_torch_class_input(i._to_internal())
         return clone
     else:


Test patch:
diff --git a/tests/cpp/test_collections.cpp b/tests/cpp/test_collections.cpp
--- a/tests/cpp/test_collections.cpp
+++ b/tests/cpp/test_collections.cpp
@@ -45,6 +45,50 @@ TEST(CppAPITests, TestCollectionStandardTensorInput) {
   ASSERT_TRUE(torch_tensorrt::tests::util::cosineSimEqual(out.toTensor(), trt_out.toTensor()));
 }
 
+TEST(CppAPITests, TestCollectionStandardTensorInputLongDtype) {
+  std::string path = "tests/modules/standard_tensor_input_scripted.jit.pt";
+  torch::Tensor in0 = torch::randn({1, 3, 512, 512}, torch::kCUDA).to(torch::kLong);
+  std::vector<at::Tensor> inputs;
+  inputs.push_back(in0);
+  inputs.push_back(in0);
+
+  torch::jit::Module mod;
+  try {
+    // Deserialize the ScriptModule from a file using torch::jit::load().
+    mod = torch::jit::load(path);
+  } catch (const c10::Error& e) {
+    std::cerr << "error loading the model\n";
+  }
+  mod.eval();
+  mod.to(torch::kCUDA);
+
+  std::vector<torch::jit::IValue> inputs_;
+
+  for (auto in : inputs) {
+    inputs_.push_back(torch::jit::IValue(in.clone()));
+  }
+
+  auto out = mod.forward(inputs_);
+
+  std::vector<torch_tensorrt::Input> input_range;
+
+  // Specify Long input tensor type
+  input_range.push_back({in0.sizes(), torch::kLong});
+  input_range.push_back({in0.sizes(), torch::kLong});
+  torch_tensorrt::ts::CompileSpec compile_settings(input_range);
+  compile_settings.min_block_size = 1;
+
+  // // FP32 execution with long and double truncation
+  compile_settings.enabled_precisions = {torch::kFloat};
+  compile_settings.truncate_long_and_double = true;
+  // // Compile module
+  auto trt_mod = torch_tensorrt::torchscript::compile(mod, compile_settings);
+  auto trt_out = trt_mod.forward(inputs_);
+
+  ASSERT_TRUE(torch_tensorrt::tests::util::cosineSimEqual(
+      out.toTensor().to(torch::kFloat), trt_out.toTensor().to(torch::kFloat)));
+}
+
 TEST(CppAPITests, TestCollectionTupleInput) {
   std::string path = "tests/modules/tuple_input_scripted.jit.pt";
   torch::Tensor in0 = torch::randn({1, 3, 512, 512}, torch::kCUDA).to(torch::kHalf);
diff --git a/tests/py/api/test_collections.py b/tests/py/api/test_collections.py
--- a/tests/py/api/test_collections.py
+++ b/tests/py/api/test_collections.py
@@ -50,6 +50,36 @@ def test_compile(self):
         )
 
 
+class TestStandardTensorInputLong(unittest.TestCase):
+    def test_compile(self):
+
+        self.input = torch.randn((1, 3, 224, 224)).to("cuda")
+        self.model = (
+            torch.jit.load(MODULE_DIR + "/standard_tensor_input_scripted.jit.pt")
+            .eval()
+            .to("cuda")
+        )
+
+        compile_spec = {
+            "inputs": [
+                torchtrt.Input(self.input.shape, dtype=torch.long),
+                torchtrt.Input(self.input.shape, dtype=torch.long),
+            ],
+            "device": torchtrt.Device("gpu:0"),
+            "enabled_precisions": {torch.float},
+            "truncate_long_and_double": True,
+        }
+
+        trt_mod = torchtrt.ts.compile(self.model, **compile_spec)
+        cos_sim = cosine_similarity(
+            self.model(self.input, self.input), trt_mod(self.input, self.input)
+        )
+        self.assertTrue(
+            cos_sim > COSINE_THRESHOLD,
+            msg=f"standard_tensor_input_long_scripted TRT outputs don't match with the original model. Cosine sim score: {cos_sim} Threshold: {COSINE_THRESHOLD}",
+        )
+
+
 class TestTupleInput(unittest.TestCase):
     def test_compile(self):
 
diff --git a/tests/util/run_graph_engine.cpp b/tests/util/run_graph_engine.cpp
--- a/tests/util/run_graph_engine.cpp
+++ b/tests/util/run_graph_engine.cpp
@@ -20,7 +20,7 @@ namespace util {
 std::vector<core::ir::Input> toInputs(std::vector<at::Tensor> ten) {
   std::vector<core::ir::Input> a;
   for (auto i : ten) {
-    a.push_back(core::ir::Input(core::util::toVec(i.sizes()), core::util::ScalarTypeToTRTDataType(i.scalar_type())));
+    a.push_back(core::ir::Input(core::util::toVec(i.sizes()), i.scalar_type()));
   }
   return a;
 }
@@ -30,7 +30,7 @@ std::vector<core::ir::Input> toInputsDynamic(std::vector<at::Tensor> ten, bool d
 
   for (auto i : ten) {
     auto opt = core::util::toVec(i.sizes());
-    auto dtype = core::util::ScalarTypeToTRTDataType(i.scalar_type());
+    auto dtype = i.scalar_type();
 
     if (dynamic_batch) {
       std::vector<int64_t> min_range(opt);