[export] Update dynamo_graph_capture_for_export to return GraphModule.

Author: zhxchen17

test/distributed/tensor/test_dtensor_export.py

@@ -323,7 +323,7 @@ def test_annotate_aot_export_joint_with_descriptors_alone(self):
         [
             (
                 graph_capture_and_aot_export_joint_with_descriptors_v2,
-                "[[4, 10], [4], [10, 4], [10], [4, 10], [4], [10, 4], [10], [s64, 10], [s64, 10]]",
+                "[[4, 10], [4], [10, 4], [10], [s49, 10], [s49, 10]]",
             ),
             (
                 graph_capture_and_aot_export_joint_with_descriptors,

test/export/test_experimental.py

@@ -7,6 +7,7 @@
 
 import torch
 import torch._dynamo
+from torch._dynamo.functional_export import dynamo_graph_capture_for_export
 from torch._dynamo.test_case import run_tests, TestCase
 from torch._functorch.aot_autograd import aot_export_module
 from torch.export import export
@@ -403,8 +404,6 @@ def forward(self, x):
         self.assertEqual(res_export, res_eager)
 
     def test_dynamo_graph_capture(self):
-        from torch._dynamo.functional_export import dynamo_graph_capture_for_export
-
         class Foo(torch.nn.Module):
             def forward(self, dct, lst, bleh):
                 x = dct["a"] * lst[1][0]
@@ -439,6 +438,22 @@ def make_inputs():
         test_inputs = make_inputs()
         self.assertEqual(gm(*test_inputs), foo(*test_inputs))
 
+    def test_dynamo_graph_capture_closure(self):
+        from torch.export import Dim
+
+        N = 3
+
+        class MyModel(torch.nn.Module):
+            def forward(self, x):
+                y = x[:-1, :]  # [s0 - 1, 32]
+                stacked = torch.stack([y] * N, dim=0)  # [N * (s0 - 1), 32]
+                reshaped = stacked.reshape(-1, N, 32)  # [(s0 - 1), N, 32]
+                return reshaped
+
+        inps = (torch.randn(10, 32),)
+        ep = dynamo_graph_capture_for_export(MyModel())(*inps)
+        self.assertEqual(ep(*inps), MyModel()(*inps))
+
 
 if __name__ == "__main__":
     run_tests()

torch/_dynamo/convert_frame.py

@@ -894,6 +894,7 @@ def graph_capture_output(self) -> GraphCaptureOutput:
             output_graph.import_sources,
             output_graph.traced_code,
             self.bytecode,
+            self.tracer_output.closure,
         )
 
 
@@ -925,6 +926,7 @@ class GraphCaptureOutput:
     import_sources: dict[str, str]
     traced_code: list[CodeType]
     bytecode: CodeType
+    closure: Optional[tuple[Any, ...]]
 
     def build_guards(
         self,
@@ -979,7 +981,7 @@ def forward_callable(self) -> Callable[..., Any]:
         return types.FunctionType(
             self.graph_capture_output.bytecode,
             f_globals,
-            closure=(),
+            closure=self.graph_capture_output.closure,
         )
 
 

torch/_dynamo/eval_frame.py

@@ -1211,7 +1211,9 @@ def __call__(self, fn: Callable[..., Any]) -> Callable[..., Any]:
 
 # Make dynamo graph to have same input/output spec as user code
 def argument_names(
-    f_sig: inspect.Signature, args: list[Any], kwargs: dict[str, Any]
+    f_sig: inspect.Signature,
+    args: Union[list[Any], tuple[Any, ...]],
+    kwargs: dict[str, Any],
 ) -> list[str]:
     def signature_to_fullargspec(sig: inspect.Signature) -> inspect.FullArgSpec:
         # Get a list of Parameter objects from the Signature object

torch/_dynamo/functional_export.py

@@ -1,9 +1,8 @@
-import copy
 import inspect
 import logging
 import traceback
-import types
 from collections import namedtuple
+from dataclasses import dataclass
 from typing import Any, Callable, Optional, TYPE_CHECKING, Union
 
 import sympy
@@ -23,12 +22,12 @@
     DimDynamic,
     StatelessSymbolicContext,
 )
-from torch.fx.graph import _PyTreeCodeGen, _PyTreeInfo
+from torch.fx.graph import _ExportCodeGen, _PyTreeCodeGen, _PyTreeInfo
+from torch.utils._pytree import TreeSpec
 
 
 if TYPE_CHECKING:
     from torch._subclasses.fake_tensor import FakeTensorMode
-    from torch.utils._pytree import TreeSpec
 
 
 log = logging.getLogger(__name__)
@@ -449,9 +448,20 @@ def _suggest_or_raise_constraint_violation(
         raise constraint_violation_error
 
 
+@dataclass(frozen=True)
+class PyTreeifyOutput:
+    graph_module: torch.fx.GraphModule
+    in_spec: TreeSpec
+    in_shuffle_graph: torch.fx.GraphModule
+    num_flat_args: int
+    out_spec: TreeSpec
+    out_shuffle_graph: torch.fx.GraphModule
+    root: Optional[torch.nn.Module] = None
+
+
 def pytreeify(
     out: CaptureOutput, mod: Any, args: tuple[Any, ...], kwargs: dict[str, Any]
-) -> Any:
+) -> PyTreeifyOutput:
     """
     Given a dynamo capture output, return a callable graph module that
     contain the following information:
@@ -469,10 +479,13 @@ def pytreeify(
     backend_input = out.backend_input
     backend = out.backend_input.graph_module
 
+    root = None
     if isinstance(mod, torch.nn.Module):
         args = (mod,) + args
+        root = mod
     elif inspect.ismethod(mod):
         args = (mod.__self__,) + args
+        root = mod.__self__
 
     flat_real_args, in_spec = pytree.tree_flatten((args, kwargs))
 
@@ -538,47 +551,73 @@ def backend_dummy(*example_inputs):
     out_shuffle = OutShuffle()
     out_shuffle_graph = torch.fx.symbolic_trace(out_shuffle)
 
-    def pytree_call(*args, **kwargs):
-        import torch.export._unlift
+    assert out_shuffle.out_spec is not None
+    return PyTreeifyOutput(
+        backend_input.graph_module,
+        in_spec,
+        in_shuffle_graph,
+        len(flat_real_args),
+        out_shuffle.out_spec,
+        out_shuffle_graph,
+        root=root,  # type: ignore[arg-type]
+    )
 
-        flat_args, in_spec_runtime = pytree.tree_flatten((args, kwargs))
-        if not torch.export._unlift.eq_spec(in_spec_runtime, in_spec):
-            raise RuntimeError(
-                f"Model input mismatch. Expected input spec: {in_spec}. Actual input spec: {in_spec_runtime}"
-            )
-        flat_outs = backend_input.graph_module(*in_shuffle_graph(*flat_args))
-        assert out_shuffle.out_spec is not None
-        return pytree.tree_unflatten(
-            out_shuffle_graph(*flat_args, *flat_outs), out_shuffle.out_spec
-        )
 
-    if isinstance(mod, torch.nn.Module):
-        compiled_mod = copy.copy(mod)
-        compiled_mod.forward = types.MethodType(pytree_call, compiled_mod)
-        if not hasattr(compiled_mod, "meta"):
-            compiled_mod.meta = {}  # type: ignore[attr-defined]
-        if isinstance(compiled_mod.meta, dict) and "fake_mode" not in compiled_mod.meta:
-            compiled_mod.meta["fake_mode"] = out.backend_input.fake_mode
-        return compiled_mod
-    elif inspect.ismethod(mod):
-        return types.MethodType(pytree_call, mod.__self__)
-    else:
-        return pytree_call
+def normalize_graph_module(gm):
+    for node in gm.graph.nodes:
+        if node.op == "placeholder":
+            node.meta["val"] = node.meta["example_value"]
 
 
 def dynamo_graph_capture_for_export(
     mod: Callable[..., Any],
+    constraints: Optional[list[Constraint]] = None,
 ) -> Callable[..., Any]:
     def inner(*args: Any, **kwargs: Any) -> Any:
         with (
             get_metrics_context(),
             dynamo_timed("fullgraph_capture"),
         ):
-            out = fullgraph_capture(mod, args, kwargs)
+            out = fullgraph_capture(
+                mod,
+                args,
+                kwargs,
+                constraints=constraints,
+                _is_export_deprecated_do_not_use=True,
+            )
 
         # TODO filter out side effects.
-
-        return pytreeify(out, mod, args, kwargs)
+        pyt = pytreeify(out, mod, args, kwargs)
+
+        graph_module = pyt.graph_module
+        tree_leaf_names = [
+            graph_module.graph._graph_namespace.create_name(f"_tree_leaf_{i}", None)
+            for i in range(pyt.num_flat_args)
+        ]
+        graph_module.graph._codegen = _ExportCodeGen(
+            _PyTreeInfo(
+                argument_names(inspect.signature(mod), args, kwargs),
+                pyt.in_spec,
+                pyt.out_spec,
+            ),
+            pyt.in_shuffle_graph,
+            pyt.out_shuffle_graph,
+            tree_leaf_names,
+            pyt.root,
+        )  # type: ignore[attr-defined]
+        normalize_graph_module(graph_module)
+        graph_module._in_spec = pyt.in_spec
+        graph_module._out_spec = pyt.out_spec
+        graph_module._in_shuffle_graph = pyt.in_shuffle_graph
+        graph_module._out_shuffle_graph = pyt.out_shuffle_graph
+        object.__setattr__(graph_module, "_root", pyt.root)  # type: ignore[attr-defined]
+        graph_module.recompile()
+        graph_module.meta["module_call_specs"] = (
+            out.graph_capture_output.output_graph.export_metadata.module_call_spec
+        )
+        assert out.backend_input is not None
+        graph_module.meta["fake_mode"] = out.backend_input.fake_mode  # type: ignore[attr-defined]
+        return graph_module
 
     return inner
 

torch/_dynamo/output_graph.py

@@ -2741,12 +2741,14 @@ class DynamoTracerOutput:
     error_on_graph_break: bool
     is_tracing_resume_prologue: bool
     output_graph: Optional[OutputGraph]
+    closure: Optional[tuple[Any, ...]]
 
     def __init__(
         self, tracer: "InstructionTranslatorBase", error: Optional[Any] = None
     ) -> None:
         self.error_on_graph_break = tracer.error_on_graph_break
         self.is_tracing_resume_prologue = tracer.is_tracing_resume_prologue
+        self.closure = tracer.closure
         if error:
             self.output_graph = None
         else:

torch/_dynamo/symbolic_convert.py

@@ -4119,6 +4119,7 @@ def __init__(
         self.f_builtins: dict[str, Any] = f_builtins
         self.code_options: dict[str, Any] = code_options
         self.f_code: types.CodeType = f_code
+        self.closure = closure
 
         # Execution record for replaying errors
         if closure is not None and config.replay_record_enabled:

torch/export/_trace.py

@@ -97,7 +97,7 @@
     GuardOnDataDependentSymNode,
     ShapeEnv,
 )
-from torch.fx.graph import _PyTreeCodeGen, _PyTreeInfo
+from torch.fx.graph import _PyTreeInfo
 from torch.utils._pytree import TreeSpec
 from torch.utils._sympy.value_ranges import ValueRangeError
 
@@ -1486,12 +1486,10 @@ def _strict_export(
 
     orig_arg_names = gm_torch_level.graph._codegen.pytree_info.orig_args  # type: ignore[attr-defined]
 
-    gm_torch_level.graph._codegen = _PyTreeCodeGen(
-        _PyTreeInfo(
-            orig_arg_names,
-            gm_torch_level._in_spec,
-            out_spec,
-        )
+    gm_torch_level.graph._codegen.pytree_info = _PyTreeInfo(
+        orig_arg_names,
+        gm_torch_level._in_spec,
+        out_spec,
     )
     gm_torch_level.recompile()