fix: Refactor tensor freezing in Dynamo

Author: gs-olive

py/torch_tensorrt/dynamo/backend/aot_module.py

@@ -1,6 +1,8 @@
 import logging
 from typing import Sequence
 import torch
+from torch._dynamo.utils import detect_fake_mode
+import unittest
 from functools import partial
 import torch._dynamo as td
 from torch._guards import TracingContext
@@ -23,8 +25,7 @@
     repair_long_or_double_inputs,
 )
 
-from torch._functorch.aot_autograd import make_boxed_compiler
-from .aot_module import aot_module
+from torch._functorch.aot_autograd import aot_export_joint_simple
 
 
 logger = logging.getLogger(__name__)
@@ -47,21 +48,25 @@ def aot_torch_tensorrt_aten_backend(
 ):
     settings = parse_dynamo_kwargs(kwargs)
 
-    custom_backend = partial(
-        _pretraced_backend,
-        settings=settings,
-    )
-
     # Perform Pre-AOT Lowering for Module-Level Replacement
     gm = pre_aot_substitutions(gm)
 
-    # Invoke AOTAutograd to translate operators to aten
-    return aot_module(
-        gm,
-        sample_inputs,
-        fw_compiler=make_boxed_compiler(custom_backend),
-        decompositions=get_decompositions(),
-    )
+    fake_mode = detect_fake_mode(sample_inputs)
+
+    # Place backend tracing within FakeTensor context allowing nonfake Tensors
+    with unittest.mock.patch.object(
+        fake_mode, "allow_non_fake_inputs", True
+    ), fake_mode:
+
+        # Invoke AOTAutograd to translate operators to aten
+        graph_module = aot_export_joint_simple(
+            gm,
+            sample_inputs,
+            trace_joint=False,
+            decompositions=get_decompositions(),
+        )
+
+        return _pretraced_backend(graph_module, sample_inputs, settings)
 
 
 def _pretraced_backend(
@@ -81,16 +86,9 @@ 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(
-            frozen_gm,
-            nonfrozen_inputs,
+            gm,
+            sample_inputs,
             settings=settings,
         )
         return trt_compiled

py/torch_tensorrt/dynamo/backend/backends.py

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