feat: Add preliminary support for freezing tensors in Dynamo

Author: gs-olive

py/torch_tensorrt/dynamo/backend/backends.py

@@ -3,6 +3,7 @@
 import torch
 from functools import partial
 import torch._dynamo as td
+from torch._guards import TracingContext
 
 from torch_tensorrt.dynamo import CompilationSettings
 from torch_tensorrt.dynamo.lowering._decompositions import (
@@ -15,10 +16,12 @@
     partition,
     get_submod_inputs,
 )
+from torch_tensorrt.dynamo.lowering._freeze_aot_graph import freeze_autograd_gm
 from torch_tensorrt.dynamo.utils import parse_dynamo_kwargs
 from torch_tensorrt.dynamo.conversion import convert_module
 
-from torch._functorch.aot_autograd import aot_module_simplified, make_boxed_compiler
+from torch._functorch.aot_autograd import make_boxed_compiler
+from .aot_module import aot_module
 
 
 logger = logging.getLogger(__name__)
@@ -30,6 +33,8 @@ def torch_tensorrt_backend(
 ):
     DEFAULT_BACKEND = aot_torch_tensorrt_aten_backend
 
+    TracingContext.get().fake_mode.allow_non_fake_inputs = True
+
     return DEFAULT_BACKEND(gm, sample_inputs, **kwargs)
 
 
@@ -48,7 +53,7 @@ def aot_torch_tensorrt_aten_backend(
     gm = pre_aot_substitutions(gm)
 
     # Invoke AOTAutograd to translate operators to aten
-    return aot_module_simplified(
+    return aot_module(
         gm,
         sample_inputs,
         fw_compiler=make_boxed_compiler(custom_backend),
@@ -73,9 +78,16 @@ def _pretraced_backend(
     try:
         logger.debug("Post-AOT Autograd graph:\n" + str(gm.graph))
 
+        frozen_gm, unfrozen_indices = freeze_autograd_gm(gm, sample_inputs)
+        nonfrozen_inputs = [sample_inputs[idx] for idx in unfrozen_indices]
+
+        frozen_gm.graph.eliminate_dead_code()
+        frozen_gm.graph.lint()
+        frozen_gm.recompile()
+
         trt_compiled = _compile_module(
-            gm,
-            sample_inputs,
+            frozen_gm,
+            nonfrozen_inputs,
             settings=settings,
         )
         return trt_compiled

py/torch_tensorrt/dynamo/conversion/trt_interpreter.py

@@ -22,6 +22,8 @@
     unified_dtype_converter,
     Frameworks,
 )
+from torch.utils._python_dispatch import _disable_current_modes
+
 
 _LOGGER: logging.Logger = logging.getLogger(__name__)
 
@@ -296,6 +298,21 @@ def call_function(self, target, args, kwargs):
         assert self._cur_node_name is not None
         return converter(self.network, target, args, kwargs, self._cur_node_name)
 
+    def get_attr(self, target, args, kwargs):
+        with _disable_current_modes():
+            from torch_tensorrt.fx.converters import to_numpy
+
+            frozen_attr = self.fetch_attr(target)
+
+            if isinstance(frozen_attr, torch.nn.Parameter):
+                constant_tensor = frozen_attr.data
+            else:
+                constant_tensor = frozen_attr
+
+            network_constant = to_numpy(constant_tensor)
+
+        return network_constant
+
     def call_method(self, target, args, kwargs):
         assert isinstance(target, str)
         converter = CONVERTERS.get(target)
@@ -317,6 +334,17 @@ def output(self, target, args, kwargs):
         else:
             outputs = (args[0],)
 
+        for output_idx in range(len(outputs)):
+            from torch_tensorrt.fx.converters import get_trt_tensor
+
+            output = outputs[output_idx]
+
+            if not isinstance(output, trt.tensorrt.ITensor):
+                new_output = get_trt_tensor(self.network, output, target)
+                outputs = (
+                    outputs[:output_idx] + (new_output,) + outputs[output_idx + 1 :]
+                )
+
         if not all(isinstance(output, trt.tensorrt.ITensor) for output in outputs):
             raise RuntimeError("TensorRT requires all outputs to be Tensor!")
 
@@ -356,3 +384,5 @@ def output(self, target, args, kwargs):
             elif self.output_fp16 and output.dtype == trt.float32:
                 output.dtype = trt.float16
             self._output_names.append(name)
+
+        return list(outputs)

py/torch_tensorrt/dynamo/lowering/__init__.py

@@ -8,3 +8,4 @@
 from ._partition import partition, get_submod_inputs, DEFAULT_SINGLE_NODE_PARTITIONS
 from .substitutions import *
 from ._fusers import *
+from ._freeze_aot_graph import *

py/torch_tensorrt/dynamo/lowering/_partition.py

@@ -125,8 +125,8 @@ def is_node_supported(
 
         if (
             node.target in CONVERTERS.keys()
-            and node_name not in self.torch_executed_ops
-        ):
+            or (node.op == "get_attr" and "frozen" in node_name)
+        ) and node_name not in self.torch_executed_ops:
             # If node is a proper, supported computational node, store the operator
             if not node.is_impure():
                 self.supported_operators.add(node_name)