fix: Add constant folding utility to freezing

Author: gs-olive

py/torch_tensorrt/dynamo/backend/backends.py

@@ -8,6 +8,7 @@
 import torch._dynamo as td
 from torch._dynamo.utils import detect_fake_mode
 from torch._functorch.aot_autograd import aot_export_joint_simple
+from torch._inductor.freezing import ConstantFolder, replace_node_with_constant
 from torch_tensorrt.dynamo import CompilationSettings
 from torch_tensorrt.dynamo.compile import compile_module
 from torch_tensorrt.dynamo.lowering._decompositions import get_decompositions
@@ -62,6 +63,8 @@ def aot_torch_tensorrt_aten_backend(
             ),
         )
 
+        constant_fold(graph_module)
+
         return _pretraced_backend(graph_module, sample_inputs, settings)
 
 
@@ -105,3 +108,25 @@ def _pretraced_backend(
                 + "specify pass_through_build_failures=False."
             )
             raise
+
+
+@torch.utils._python_dispatch._disable_current_modes()  # type: ignore
+def constant_fold(gm: torch.fx.GraphModule) -> Any:
+    cf = ConstantFolder(gm, skip_constructors=False)
+    cf.run()
+
+    for node, constant in cf.node_replacements.items():
+        replace_node_with_constant(gm, node, constant)
+
+    erased_params = []
+    for node in gm.graph.nodes:
+        if node.op == "get_attr" and len(node.users) == 0:
+            delattr(gm, node.target)
+            erased_params.append(node)
+
+    for node in erased_params:
+        gm.graph.erase_node(node)
+
+    gm.graph.eliminate_dead_code()
+    gm.graph.lint()
+    gm.recompile()

py/torch_tensorrt/dynamo/lowering/__init__.py

@@ -1,5 +1,4 @@
 from ._decompositions import get_decompositions  # noqa: F401
-from ._freeze_aot_graph import *  # noqa: F401
 from ._fusers import *  # noqa: F401
 from ._pre_aot_lowering import SUBSTITUTION_REGISTRY  # noqa: F401
 from ._pre_aot_lowering import register_substitution  # noqa: F401

py/torch_tensorrt/dynamo/lowering/_pre_aot_lowering.py

@@ -81,10 +81,6 @@ def pre_aot_substitutions(gm: torch.fx.GraphModule) -> torch.fx.GraphModule:
     """
     logger.debug("Pre-module replacement graph:\n" + str(gm.graph))
 
-    # Ensure all parameters are in inference mode
-    for param in gm.parameters():
-        param.requires_grad = False
-
     # Iterate over graph nodes, extracting module calls, to check for interceptions
     for n in gm.graph.nodes:
         exists_in_registry = False

py/torch_tensorrt/dynamo/partitioning/_adjacency_partitioner.py

@@ -39,7 +39,9 @@ def is_node_supported(
     ) -> bool:
         node_name = ConverterRegistry.qualified_name_or_str(node.target)
 
-        if node in CONVERTERS and node_name not in self.torch_executed_ops:
+        if (
+            node in CONVERTERS or (node.op == "get_attr" and "constant" 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():
                 if node_name not in self.supported_operators:

py/torch_tensorrt/dynamo/partitioning/_global_partitioner.py

@@ -122,8 +122,7 @@ def is_node_supported(
         node_name = ConverterRegistry.qualified_name_or_str(node.target)
 
         if (
-            node.target in CONVERTERS.keys()
-            or (node.op == "get_attr" and "constant" in node_name)
+            node in CONVERTERS or (node.op == "get_attr" and "constant" 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():