fix: Add truncate_long_and_double to Dynamo

- Add default, setting, and function arguments for
`truncate_long_and_double` in Dynamo
- Add utilities for repairing long/double inputs to TRT engines,
including support for autocasting back to long/double after the
computation completes
- Add multiple helper functions to enable easy testing and diagnosis of
long/double IO to TRT engines
- Add necessary compiler code to enable usage of the `truncate_long_and_double`
argument as a switch for the feature
- Add Dynamo compile support for `truncate_long_and_double` compilation
argument by intercepting long/double type inputs and casting them to
their 32-bit counterparts prior to usage in TRT-accelerated subgraphs,
then casting back if necessary
- Add robust logic to handle 64-bit inputs and outputs
- Add test cases for long and double scenarios
- Centralize truncation utility for later use in Dynamo export path

Author: gs-olive

py/torch_tensorrt/dynamo/_defaults.py

@@ -10,3 +10,4 @@
 VERSION_COMPATIBLE = False
 OPTIMIZATION_LEVEL = None
 USE_EXPERIMENTAL_RT = False
+TRUNCATE_LONG_AND_DOUBLE = False

py/torch_tensorrt/dynamo/backend/__init__.py

@@ -21,6 +21,7 @@
     VERSION_COMPATIBLE,
     OPTIMIZATION_LEVEL,
     USE_EXPERIMENTAL_RT,
+    TRUNCATE_LONG_AND_DOUBLE,
 )
 
 
@@ -44,7 +45,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,
     min_block_size=MIN_BLOCK_SIZE,
     torch_executed_ops=[],
@@ -62,7 +63,7 @@ def compile(
         "The Dynamo backend is an experimental feature, for which only the "
         + "following arguments are supported: "
         + "{enabled_precisions, debug, workspace_size, min_block_size, "
-        + "torch_executed_ops, pass_through_build_failures}"
+        + "truncate_long_and_double, torch_executed_ops, pass_through_build_failures}"
     )
 
     if "use_experimental_fx_rt" in kwargs:
@@ -103,6 +104,7 @@ def compile(
         version_compatible=version_compatible,
         optimization_level=optimization_level,
         use_experimental_rt=use_experimental_rt,
+        truncate_long_and_double=truncate_long_and_double,
         **kwargs,
     )
 
@@ -130,6 +132,7 @@ def create_backend(
     version_compatible: bool = VERSION_COMPATIBLE,
     optimization_level: Optional[int] = OPTIMIZATION_LEVEL,
     use_experimental_rt: bool = USE_EXPERIMENTAL_RT,
+    truncate_long_and_double: bool = TRUNCATE_LONG_AND_DOUBLE,
     **kwargs,
 ):
     """Create torch.compile backend given specified arguments
@@ -163,6 +166,7 @@ def create_backend(
         version_compatible=version_compatible,
         optimization_level=optimization_level,
         use_experimental_rt=use_experimental_rt,
+        truncate_long_and_double=truncate_long_and_double,
     )
 
     return partial(

py/torch_tensorrt/dynamo/backend/backends.py

@@ -16,6 +16,7 @@
     get_submod_inputs,
 )
 from torch_tensorrt.dynamo.backend.utils import parse_dynamo_kwargs
+from torch_tensorrt.dynamo.common import repair_long_or_double_inputs
 from torch_tensorrt.dynamo.backend.conversion import convert_module
 
 from torch._dynamo.backends.common import fake_tensor_unsupported
@@ -141,6 +142,16 @@ def _compile_module(
             partitioned_module, submodule, sample_inputs
         )
 
+        # Ensure all submodule inputs do not require a gradient
+        for param in submodule_inputs:
+            param.requires_grad = False
+
+        # Handle long/double inputs if requested by the user
+        if settings.truncate_long_and_double:
+            submodule_inputs = repair_long_or_double_inputs(
+                partitioned_module, submodule, submodule_inputs, name
+            )
+
         # Create TRT Module from submodule
         trt_mod = convert_module(
             submodule,

py/torch_tensorrt/dynamo/backend/test/test_backend_compiler.py

@@ -169,5 +169,116 @@ def forward(self, x, y):
         )
 
 
+class Test64BitInput(TestCase):
+    def test_float64_input_full_support(self):
+        class FullySupportedMultiOp(torch.nn.Module):
+            def forward(self, x, y):
+                return torch.ops.aten.mean.dim(
+                    torch.ops.aten.mul.Tensor(torch.ops.aten.add.Tensor(x, y), 2), [0]
+                )
+
+        fx_graph = torch.fx.symbolic_trace(FullySupportedMultiOp())
+        partitioned_graph = partition(deepcopy(fx_graph), min_block_size=3)
+
+        self.assertEquals(
+            len(list(partitioned_graph.named_children())),
+            1,
+            "All operators are supported, there should be one segment",
+        )
+
+        inputs = [
+            torch.randint(-5, 5, (16, 7), dtype=torch.double).cuda(),
+            torch.randint(-5, 5, (16, 7), dtype=torch.double).cuda(),
+        ]
+
+        torch._dynamo.reset()
+
+        # Validate that the results between Torch and Torch-TRT are similar
+        optimized_model = compile(
+            fx_graph,
+            inputs,
+            min_block_size=1,
+            pass_through_build_failures=True,
+            truncate_long_and_double=True,
+            debug=True,
+        )
+        optimized_model_results = optimized_model(*inputs).detach().cpu()
+        torch_model_results = fx_graph(*inputs).detach().cpu()
+
+        max_diff = float(
+            torch.max(torch.abs(optimized_model_results - torch_model_results))
+        )
+        self.assertAlmostEqual(
+            max_diff,
+            0,
+            DECIMALS_OF_AGREEMENT,
+            f"TRT outputs don't match with the original model.",
+        )
+
+    def test_int64_input_partial_support(self):
+        class PartiallySupportedMultiOp(torch.nn.Module):
+            def forward(self, x, y):
+                return torch.ops.aten.div.Tensor_mode(
+                    x, torch.ops.aten.add.Tensor(y, y), rounding_mode="floor"
+                )
+
+        fx_graph = torch.fx.symbolic_trace(PartiallySupportedMultiOp())
+        unexpected_ops = {torch.ops.aten.add.Tensor}
+
+        inputs = [
+            torch.randint(-40, 40, (16, 7, 5), dtype=torch.long).cuda(),
+            torch.randint(1, 40, (16, 7, 5), dtype=torch.long).cuda(),
+        ]
+
+        (unexpected_ops_seen, _, partitioned_graphs,) = lower_graph_testing(
+            fx_graph,
+            inputs,
+            unexpected_ops=unexpected_ops,
+            min_block_size=1,
+            torch_executed_ops={"torch.ops.aten.add.Tensor"},
+            testing_partitioning=True,
+        )
+
+        self.assertEquals(
+            len(unexpected_ops_seen),
+            0,
+            f"The following unexpected ops were encountered: {unexpected_ops_seen}",
+        )
+        self.assertEquals(
+            len(partitioned_graphs),
+            1,
+            "Without control flow breaks, there should only be a single graph",
+        )
+        self.assertEquals(
+            len(list(partitioned_graphs[0].named_children())),
+            1,
+            "Certain operators are set to run in Torch, expected 1 segment",
+        )
+
+        torch._dynamo.reset()
+
+        # Validate that the results between Torch and Torch-TRT are similar
+        optimized_model = compile(
+            fx_graph,
+            inputs,
+            min_block_size=1,
+            pass_through_build_failures=True,
+            truncate_long_and_double=True,
+            debug=True,
+        )
+        optimized_model_results = optimized_model(*inputs).detach().cpu()
+        torch_model_results = fx_graph(*inputs).detach().cpu()
+
+        max_diff = float(
+            torch.max(torch.abs(optimized_model_results - torch_model_results))
+        )
+        self.assertAlmostEqual(
+            max_diff,
+            0,
+            DECIMALS_OF_AGREEMENT,
+            f"TRT outputs don't match with the original model.",
+        )
+
+
 if __name__ == "__main__":
     run_tests()

py/torch_tensorrt/dynamo/backend/utils.py

@@ -3,6 +3,7 @@
 from dataclasses import replace, fields
 
 from torch_tensorrt.dynamo.backend._settings import CompilationSettings
+from torch.fx.node import _get_qualified_name
 from typing import Any, Union, Sequence, Dict
 from torch_tensorrt import _Input, Device
 
@@ -73,7 +74,6 @@ def prepare_device(device: Union[Device, torch.device]) -> torch.device:
 
     return device
 
-
 def parse_dynamo_kwargs(kwargs: Dict) -> CompilationSettings:
     """Parses the kwargs field of a Dynamo backend
 

py/torch_tensorrt/dynamo/common/__init__.py

@@ -1,3 +1,4 @@
 from ._settings import CompilationSettings
 from .input_tensor_spec import InputTensorSpec
 from .fx2trt import TRTInterpreter, TRTInterpreterResult
+from .truncate_long_and_double import repair_long_or_double_inputs

py/torch_tensorrt/dynamo/common/_settings.py

@@ -12,6 +12,7 @@
     VERSION_COMPATIBLE,
     OPTIMIZATION_LEVEL,
     USE_EXPERIMENTAL_RT,
+    TRUNCATE_LONG_AND_DOUBLE,
 )
 
 
@@ -27,3 +28,4 @@ class CompilationSettings:
     version_compatible: bool = VERSION_COMPATIBLE
     optimization_level: Optional[int] = OPTIMIZATION_LEVEL
     use_experimental_rt: bool = USE_EXPERIMENTAL_RT
+    truncate_long_and_double: bool = TRUNCATE_LONG_AND_DOUBLE