Preserve distinct mutation output names in graph signatures (#24105)

Summary:

Problem
`ExportGraphSignature.buffers_to_mutate` and `user_inputs_to_mutate` index writes by output node name. No-op folding and CSE can merge two writes into one node, so one of the mutation entries is lost from the signature.

Solution
Rebuild the signature from `output_specs` (`_get_updated_graph_signature` in `exir/_program_utils.py`). Copy a computed result only when two writes collide in the same mutation map, and keep positional user outputs and ordinary graph uses shared. Select ATen, edge or dim-order cloning from the program's verifier. Preserve value/debug metadata and FX provenance on the copies without inheriting allocation or custom-kernel placement annotations. Mutation output copies keep independent storage, as required by the existing output pointer contract.

This is the ExecuTorch half of D125489522, split out so that ExecuTorch owners can review it on its own. The Helios half (removing the Helios-side duplicated-output splitter that this change makes unnecessary) stays in D125489522, stacked on top. fbcode and xplat copies are identical.

 {F2001970137}

Differential Revision: D127072149

Author: ThomasJannaud

exir/BUCK

@@ -296,6 +296,7 @@ fbcode_target(_kind = runtime.python_library,
     ],
     deps = [
         "//caffe2:torch",
+        "//executorch/exir/dialects:lib",
     ],
 )
 

exir/_program_utils.py

@@ -7,12 +7,15 @@
 # pyre-unsafe
 
 import torch
+from executorch.exir.dialects._ops import ops as exir_ops
 from torch.export.exported_program import (
     ConstantArgument,
     ExportGraphSignature,
     InputSpec,
     OutputSpec,
 )
+from torch.export.graph_signature import OutputKind
+from torch.fx.traceback import NodeSource, NodeSourceAction
 
 
 def _get_updated_range_constraints(gm, old_range_constraints=None):
@@ -54,9 +57,12 @@ def get_shape_env(gm):
 def _get_updated_graph_signature(
     old_signature: ExportGraphSignature,
     new_gm: torch.fx.GraphModule,
+    *,
+    use_edge_ops: bool = False,
+    use_dim_order: bool = False,
 ) -> ExportGraphSignature:
     """
-    Update the graph signature's user_input/user_outputs.
+    Update the positional signature, disambiguating mutation writes in-place.
     """
     new_input_specs = []
     i = 0
@@ -87,6 +93,51 @@ def _get_updated_graph_signature(
     output_node = new_gm.graph.output_node()
     assert output_node.op == "output"
 
+    output_args = output_node.args[0]
+    assert isinstance(output_args, (tuple, list))
+    assert len(output_args) == len(
+        old_signature.output_specs
+    ), "Number of outputs changed after transformation"
+    seen_mutations: set[tuple[OutputKind, torch.fx.Node]] = set()
+    new_output_args = list(output_args)
+    for i, (node, spec) in enumerate(zip(output_args, old_signature.output_specs)):
+        if spec.kind not in (
+            OutputKind.BUFFER_MUTATION,
+            OutputKind.USER_INPUT_MUTATION,
+        ):
+            continue
+        assert isinstance(node, torch.fx.Node)
+        mutation = (spec.kind, node)
+        if mutation in seen_mutations:
+            # Mutation dictionaries are keyed by node name, while user outputs
+            # are positional and may share a node. Copy only colliding writes.
+            with new_gm.graph.inserting_before(output_node):
+                clone = new_gm.graph.call_function(
+                    (
+                        (
+                            exir_ops.edge.dim_order_ops._clone_dim_order.default
+                            if use_dim_order
+                            else exir_ops.edge.aten.clone.default
+                        )
+                        if use_edge_ops
+                        else torch.ops.aten.clone.default
+                    ),
+                    args=(node,),
+                )
+            clone.meta = {
+                key: node.meta[key]
+                for key in ("val", "stack_trace", "nn_module_stack", "debug_handle")
+                if key in node.meta
+            }
+            clone.meta["from_node"] = [
+                NodeSource(node, "_get_updated_graph_signature", NodeSourceAction.CREATE)
+            ]
+            new_output_args[i] = clone
+        seen_mutations.add(mutation)
+    if new_output_args != list(output_args):
+        output_node.args = (tuple(new_output_args),) + output_node.args[1:]
+        new_gm.recompile()
+
     new_output_specs = []
     for i, node in enumerate(output_node.args[0]):
         assert i < len(

exir/pass_manager.py

@@ -225,6 +225,10 @@ def __call__(  # noqa: C901
                             new_graph_signature = _get_updated_graph_signature(
                                 exported_program.graph_signature,
                                 res.graph_module,
+                                use_edge_ops=exported_program.dialect != "ATEN",
+                                use_dim_order=getattr(
+                                    exported_program.verifier(), "use_dim_order", False
+                                ),
                             )
                             # ExportedProgram.__init__ runs this, but this path mutates
                             # the program in place and never reconstructs it. Duplicate

exir/program/_program.py

@@ -206,7 +206,14 @@ def _update_exported_program_graph_module(
         root=gm,
         graph=gm.graph,
         graph_signature=_get_updated_graph_signature(
-            exported_program.graph_signature, gm
+            exported_program.graph_signature,
+            gm,
+            use_edge_ops=exported_program.dialect != "ATEN",
+            use_dim_order=getattr(
+                (override_verifiers or [exported_program.verifier])[0](),
+                "use_dim_order",
+                False,
+            ),
         ),
         state_dict=exported_program.state_dict,
         range_constraints=_get_updated_range_constraints(
@@ -711,7 +718,10 @@ def _to_edge(ep, config: EdgeCompileConfig) -> "ExirExportedProgram":
         root=new_gm,
         graph=new_gm.graph,
         graph_signature=_get_updated_graph_signature(
-            new_ep.exported_program.graph_signature, new_gm
+            new_ep.exported_program.graph_signature,
+            new_gm,
+            use_edge_ops=True,
+            use_dim_order=not config._skip_dim_order,
         ),
         state_dict=new_ep.exported_program.state_dict,
         range_constraints=new_ep.exported_program.range_constraints,
@@ -851,7 +861,12 @@ def _generate_edge_program(
     edge_program = ExportedProgram(
         root=gm,
         graph=gm.graph,
-        graph_signature=_get_updated_graph_signature(program.graph_signature, gm),
+        graph_signature=_get_updated_graph_signature(
+            program.graph_signature,
+            gm,
+            use_edge_ops=True,
+            use_dim_order=not config._skip_dim_order,
+        ),
         state_dict=program.state_dict,
         range_constraints=program.range_constraints,
         module_call_graph=program.module_call_graph,

exir/tests/targets.bzl

@@ -220,6 +220,7 @@ def define_common_targets(is_fbcode = False):
             ":models",
             "//caffe2:torch",
             "//caffe2/functorch:functorch_src",
+            "//executorch/exir:_program_utils",
             "//executorch/exir:graph_module",
             "//executorch/exir:lib",
             "//executorch/exir:memory",