fix/feat: Add support for 64bit Tensor inputs FX

- Add `truncate_long_and_double` argument in FX settings to allow 64bit
inputs
- Utilize existing Dynamo functionality to repair FX aten graphs with
64bit inputs
- Refactor imports in Dynamo to avoid circular import issue from new
dependency issues
- Add test cases to validate new feature

Author: gs-olive

py/torch_tensorrt/dynamo/backend/__init__.py

@@ -5,7 +5,8 @@
 from functools import partial
 
 from typing import Any, Optional, Sequence
-from torch_tensorrt import EngineCapability, Device
+from torch_tensorrt._Device import Device
+from torch_tensorrt._enums import EngineCapability
 from torch_tensorrt.fx.utils import LowerPrecision
 
 from torch_tensorrt.dynamo.common import CompilationSettings

py/torch_tensorrt/dynamo/backend/utils.py

@@ -4,7 +4,8 @@
 
 from torch_tensorrt.dynamo.common import CompilationSettings
 from typing import Any, Union, Sequence, Dict
-from torch_tensorrt import _Input, Device
+from torch_tensorrt import _Input
+from torch_tensorrt._Device import Device
 
 
 logger = logging.getLogger(__name__)

py/torch_tensorrt/dynamo/fx_ts_compat/lower.py

@@ -31,6 +31,7 @@
     VERSION_COMPATIBLE,
     OPTIMIZATION_LEVEL,
     USE_EXPERIMENTAL_RT,
+    TRUNCATE_LONG_AND_DOUBLE,
 )
 
 logger = logging.getLogger(__name__)
@@ -51,7 +52,7 @@ def compile(
     dla_local_dram_size=1073741824,
     dla_global_dram_size=536870912,
     calibrator=None,
-    truncate_long_and_double=False,
+    truncate_long_and_double=TRUNCATE_LONG_AND_DOUBLE,
     require_full_compilation=False,
     explicit_batch_dimension=False,
     debug=DEBUG,
@@ -86,6 +87,7 @@ def compile(
         max_aux_streams: max number of aux stream to use
         version_compatible: enable version compatible feature
         optimization_level: builder optimization level
+        truncate_long_and_double: Whether to truncate long and double inputs to TRT engines automatically
     Returns:
         A torch.nn.Module lowered by TensorRT.
     """
@@ -144,6 +146,7 @@ def compile(
         max_aux_streams=max_aux_streams,
         version_compatible=version_compatible,
         optimization_level=optimization_level,
+        truncate_long_and_double=truncate_long_and_double,
     )
     lowerer = Lowerer.create(lower_setting=lower_setting)
     return lowerer(module, inputs)
@@ -222,6 +225,7 @@ def default_split_function(
     splitter_setting.use_implicit_batch_dim = not lower_setting.explicit_batch_dimension
     splitter_setting.min_block_size = lower_setting.min_block_size
     splitter_setting.use_experimental_rt = lower_setting.use_experimental_rt
+    splitter_setting.truncate_long_and_double = lower_setting.truncate_long_and_double
     splitter = TRTSplitter(model, inputs, settings=splitter_setting)
     splitter.node_support_preview()
     return splitter.generate_split_results()

py/torch_tensorrt/dynamo/fx_ts_compat/lower_setting.py

@@ -73,6 +73,7 @@ class LowerSetting(LowerSettingBasic):
     max_aux_streams: max number of aux stream to use
     version_compatible: enable version compatible feature
     optimization_level: builder optimization level
+    truncate_long_and_double: Whether to truncate long and double inputs to TRT engines automatically
     """
 
     input_specs: List[InputTensorSpec] = dc.field(default_factory=list)
@@ -102,3 +103,4 @@ class LowerSetting(LowerSettingBasic):
     max_aux_streams: Optional[int] = None
     version_compatible: bool = False
     optimization_level: Optional[int] = None
+    truncate_long_and_double: bool = False

py/torch_tensorrt/dynamo/fx_ts_compat/passes/lower_pass_manager_builder.py

@@ -10,7 +10,10 @@
 from torch.fx.passes.splitter_base import generate_inputs_for_submodules, SplitResult
 from torch_tensorrt.fx.utils import LowerPrecision
 from torch_tensorrt import _Input
-from torch_tensorrt.dynamo.common import InputTensorSpec
+from torch_tensorrt.dynamo.common import (
+    InputTensorSpec,
+    repair_long_or_double_inputs,
+)
 
 from ..lower_setting import LowerSetting
 from torch_tensorrt.fx.observer import Observer
@@ -196,6 +199,14 @@ def lower_func(split_result: SplitResult) -> nn.Module:
                     _LOGGER.info(f"Now lowering submodule {submod_name}")
                     lowering_start_time = datetime.datetime.now()
 
+                    if self.lower_setting.truncate_long_and_double:
+                        submod_inputs = repair_long_or_double_inputs(
+                            parent_graph=split_result.split_module,
+                            submodule=submod,
+                            submodule_inputs=submod_inputs,
+                            submodule_name=submod_name,
+                        )
+
                     self.lower_setting.input_specs = self._trt_input
 
                     lowered_module = self._lower_func(

py/torch_tensorrt/fx/lower.py

@@ -43,6 +43,7 @@ def compile(
     use_experimental_fx_rt=False,
     correctness_atol=1e-1,
     correctness_rtol=1e-1,
+    truncate_long_and_double=False,
 ) -> nn.Module:
     """
     Takes in original module, input and lowering setting, run lowering workflow to turn module
@@ -62,6 +63,7 @@ def compile(
         cuda_graph_batch_size: Cuda graph batch size, default to be -1.
         dynamic_batch: batch dimension (dim=0) is dynamic.
         use_experimental_fx_rt: Uses the next generation TRTModule which supports both Python and TorchScript based execution (including in C++).
+        truncate_long_and_double: Whether to truncate long and double inputs to TRT engines automatically
     Returns:
         A torch.nn.Module lowered by TensorRT.
     """
@@ -85,6 +87,7 @@ def compile(
         use_experimental_rt=use_experimental_fx_rt,
         correctness_atol=correctness_atol,
         correctness_rtol=correctness_rtol,
+        truncate_long_and_double=truncate_long_and_double,
     )
     lowerer = Lowerer.create(lower_setting=lower_setting)
     return lowerer(module, input)
@@ -159,6 +162,7 @@ def default_split_function(
     splitter_setting.use_implicit_batch_dim = not lower_setting.explicit_batch_dimension
     splitter_setting.min_acc_module_size = lower_setting.min_acc_module_size
     splitter_setting.use_experimental_rt = lower_setting.use_experimental_rt
+    splitter_setting.truncate_long_and_double = lower_setting.truncate_long_and_double
     splitter = TRTSplitter(model, inputs, settings=splitter_setting)
     splitter.node_support_preview()
     return splitter.generate_split_results()

py/torch_tensorrt/fx/lower_setting.py

@@ -74,6 +74,7 @@ class LowerSetting(LowerSettingBasic):
     correctness_atol: absolute tolerance for correctness check
     correctness_rtol: relative tolerance for correctness check
     use_experimental_rt: Uses the next generation TRTModule which supports both Python and TorchScript based execution (including in C++).
+    truncate_long_and_double: Whether to truncate long and double inputs to TRT engines automatically
     """
 
     input_specs: List[InputTensorSpec] = dc.field(default_factory=list)
@@ -101,3 +102,4 @@ class LowerSetting(LowerSettingBasic):
     correctness_atol: float = 0.1
     correctness_rtol: float = 0.1
     use_experimental_rt: bool = False
+    truncate_long_and_double: bool = False

py/torch_tensorrt/fx/passes/lower_pass_manager_builder.py

@@ -10,6 +10,9 @@
 from torch.fx.passes.splitter_base import generate_inputs_for_submodules, SplitResult
 from torch_tensorrt.fx.passes.pass_utils import apply_bfloat_float_conversion
 from torch_tensorrt.fx.utils import LowerPrecision
+from torch_tensorrt.dynamo.common import (
+    repair_long_or_double_inputs,
+)
 
 from ..input_tensor_spec import generate_input_specs
 
@@ -193,6 +196,14 @@ def lower_func(split_result: SplitResult) -> nn.Module:
                     _LOGGER.info(f"Now lowering submodule {submod_name}")
                     lowering_start_time = datetime.datetime.now()
 
+                    if self.lower_setting.truncate_long_and_double:
+                        submod_inputs = repair_long_or_double_inputs(
+                            parent_graph=split_result.split_module,
+                            submodule=submod,
+                            submodule_inputs=submod_inputs,
+                            submodule_name=submod_name,
+                        )
+
                     self.lower_setting.input_specs = generate_input_specs(
                         submod_inputs,
                         self.lower_setting,

py/torch_tensorrt/fx/test/tracer/test_aten_long_and_double_inputs.py

@@ -0,0 +1,75 @@
+import unittest
+
+import torch
+
+from torch_tensorrt.fx.lower import compile
+from torch_tensorrt.fx.utils import LowerPrecision
+
+
+class LongInputTest(unittest.TestCase):
+    def test_long_input(self):
+        class Model(torch.nn.Module):
+            def forward(self, x):
+                out = x + 1
+                out = out * 2
+                out = out - 1
+                return out
+
+        mod = Model().cuda().eval()
+
+        inputs = [torch.randint(-40, 40, (3, 4, 7)).cuda().long()]
+
+        aten_mod = compile(
+            mod,
+            inputs,
+            min_acc_module_size=3,
+            explicit_batch_dimension=True,
+            verbose_log=True,
+            lower_precision=LowerPrecision.FP16,
+            truncate_long_and_double=True,
+            dynamic_batch=False,
+            is_aten=True,
+        )
+
+        aten_output = aten_mod(*inputs)[0].detach().cpu()
+        torch_output = mod(*inputs).detach().cpu()
+
+        max_diff = float(torch.max(torch.abs(aten_output - torch_output)))
+
+        self.assertAlmostEqual(
+            max_diff, 0, 4, msg="Torch outputs don't match with TRT outputs"
+        )
+
+
+class DoubleInputTest(unittest.TestCase):
+    def test_double_input(self):
+        class Model(torch.nn.Module):
+            def forward(self, x):
+                out = x + 1
+                out = out * 2
+                return torch.mean(out, dim=-1)
+
+        mod = Model().cuda().eval()
+
+        inputs = [torch.rand((3, 4, 1)).cuda().double()]
+
+        aten_mod = compile(
+            mod,
+            inputs,
+            min_acc_module_size=3,
+            explicit_batch_dimension=True,
+            verbose_log=True,
+            lower_precision=LowerPrecision.FP32,
+            truncate_long_and_double=True,
+            dynamic_batch=False,
+            is_aten=True,
+        )
+
+        aten_output = aten_mod(*inputs)[0].detach().cpu()
+        torch_output = mod(*inputs).detach().cpu()
+
+        max_diff = float(torch.max(torch.abs(aten_output - torch_output)))
+
+        self.assertAlmostEqual(
+            max_diff, 0, 4, msg="Torch outputs don't match with TRT outputs"
+        )

py/torch_tensorrt/fx/tools/trt_splitter.py

@@ -19,6 +19,7 @@
 def create_trt_operator_support(
     use_implicit_batch_dim=True,
     exclude_support_node_name: set = (),
+    truncate_long_and_double: bool = False,
 ) -> ops.OperatorSupportBase:
     """Creates an `OperatorSupportBase` instance used for TRT splitting purpose."""
     # Create an `OperatorSupport` that declares a node supported if it
@@ -32,14 +33,17 @@ def create_trt_operator_support(
             support_dict[get_acc_ops_name(k)] = None
     supported_if_converter_registered = ops.OperatorSupport(support_dict=support_dict)
 
-    return ops.chain(
-        ops.OpSupports.decline_if_node_in_names(exclude_support_node_name),
-        # 1. Node is not supported if it has args with int64 or float64 dtype:
-        ops.OpSupports.decline_if_input_dtype(torch.int64),
-        ops.OpSupports.decline_if_input_dtype(torch.float64),
-        # 2. Node is supported if it has TRT converter:
-        supported_if_converter_registered,
-    )
+    op_support_checks = [
+        ops.OpSupports.decline_if_node_in_names(exclude_support_node_name)
+    ]
+
+    if not truncate_long_and_double:
+        op_support_checks.append(ops.OpSupports.decline_if_input_dtype(torch.int64))
+        op_support_checks.append(ops.OpSupports.decline_if_input_dtype(torch.float64))
+
+    op_support_checks.append(supported_if_converter_registered)
+
+    return ops.chain(*op_support_checks)
 
 
 class TRTSplitterSetting(splitter_base._SplitterSettingBase):
@@ -52,6 +56,7 @@ def __init__(self):
         self.use_implicit_batch_dim: bool = True
         self.exclude_support_node_name: set = set()
         self.use_experimental_rt: bool = False
+        self.truncate_long_and_double: bool = False
 
         if self.use_experimental_rt and self.use_implicit_batch_dim:
             raise ValueError(
@@ -71,7 +76,9 @@ def __init__(
             settings = TRTSplitterSetting()
         if not operator_support:
             operator_support = create_trt_operator_support(
-                settings.use_implicit_batch_dim, settings.exclude_support_node_name
+                settings.use_implicit_batch_dim,
+                settings.exclude_support_node_name,
+                settings.truncate_long_and_double,
             )
         super().__init__(
             module,