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

@@ -6,13 +6,105 @@
 
 # pyre-unsafe
 
+from collections.abc import Callable
+
 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
+
+_MUTATION_OUTPUT_CLONE_OP = "_mutation_output_clone_op"
+
+
+def _get_mutation_output_clone_op(
+    program: torch.export.ExportedProgram,
+    clone_op: Callable[..., torch.Tensor] | None = None,
+) -> Callable[..., torch.Tensor] | None:
+    if clone_op is not None:
+        return clone_op
+    if program.dialect == "ATEN":
+        return torch.ops.aten.clone.default
+    kind = program.graph_module.meta.get(_MUTATION_OUTPUT_CLONE_OP)
+    if kind == "aten":
+        return torch.ops.aten.clone.default
+    if kind == "edge":
+        return exir_ops.edge.aten.clone.default
+    if kind == "edge_dim_order":
+        return exir_ops.edge.dim_order_ops._clone_dim_order.default
+    return None
+
+
+def _set_mutation_output_clone_op(
+    gm: torch.fx.GraphModule, clone_op: Callable[..., torch.Tensor] | None
+) -> None:
+    """Record the clone dialect and layout as metadata preserved by deepcopy."""
+    if clone_op is None:
+        return
+    if clone_op is torch.ops.aten.clone.default:
+        kind = "aten"
+    elif clone_op is exir_ops.edge.aten.clone.default:
+        kind = "edge"
+    elif clone_op is exir_ops.edge.dim_order_ops._clone_dim_order.default:
+        kind = "edge_dim_order"
+    else:
+        raise ValueError(f"Unsupported mutation output clone operator: {clone_op}")
+    gm.meta[_MUTATION_OUTPUT_CLONE_OP] = kind
+
+
+def _disambiguate_mutation_outputs(
+    signature: ExportGraphSignature,
+    gm: torch.fx.GraphModule,
+    *,
+    clone_op: Callable[..., torch.Tensor] | None,
+) -> bool:
+    """Copy colliding mutation outputs in-place, preserving positional user outputs."""
+    output_node = gm.graph.output_node()
+    output_args = output_node.args[0]
+    assert isinstance(output_args, (tuple, list))
+    assert len(output_args) == len(
+        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, 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:
+            if clone_op is None:
+                raise RuntimeError(
+                    "Mutation output repair requires an explicit clone operator "
+                    "for programs without an Edge compile configuration"
+                )
+            # Mutation dictionaries are keyed by node name, while user outputs
+            # are positional and may share a node. Copy only colliding writes.
+            with gm.graph.inserting_before(output_node):
+                clone = gm.graph.call_function(clone_op, 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, "_disambiguate_mutation_outputs", NodeSourceAction.CREATE)
+            ]
+            new_output_args[i] = clone
+        seen_mutations.add(mutation)
+    if new_output_args == list(output_args):
+        return False
+    output_node.args = (tuple(new_output_args),) + output_node.args[1:]
+    gm.recompile()
+    return True
 
 
 def _get_updated_range_constraints(gm, old_range_constraints=None):
@@ -55,9 +147,7 @@ def _get_updated_graph_signature(
     old_signature: ExportGraphSignature,
     new_gm: torch.fx.GraphModule,
 ) -> ExportGraphSignature:
-    """
-    Update the graph signature's user_input/user_outputs.
-    """
+    """Rebuild the positional signature without modifying the graph module."""
     new_input_specs = []
     i = 0
     for node in new_gm.graph.nodes:
@@ -87,6 +177,11 @@ 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"
     new_output_specs = []
     for i, node in enumerate(output_node.args[0]):
         assert i < len(

exir/pass_manager.py

@@ -15,8 +15,11 @@
 import torch.fx.passes.infra.pass_manager as fx
 import torch.utils._pytree as pytree
 from executorch.exir._program_utils import (
+    _disambiguate_mutation_outputs,
+    _get_mutation_output_clone_op,
     _get_updated_graph_signature,
     _get_updated_range_constraints,
+    _set_mutation_output_clone_op,
 )
 from executorch.exir.error import ExportError, ExportErrorType
 from executorch.exir.pass_base import ExportedProgramPassBase, ExportedProgramPassResult
@@ -179,6 +182,8 @@ def __call__(  # noqa: C901
         self,
         exported_program: ExportedProgram,
         override_verifiers: Optional[list[Type[Verifier]]] = None,
+        *,
+        mutation_output_clone_op: Callable[..., torch.Tensor] | None = None,
     ) -> ExportedProgramPassResult:
         """
         Runs passes on an ExportedProgram.
@@ -195,6 +200,9 @@ def __call__(  # noqa: C901
             self.solve_constraints()
 
         exported_program = copy.copy(exported_program)
+        clone_op = _get_mutation_output_clone_op(
+            exported_program, mutation_output_clone_op
+        )
 
         if override_verifiers:
             exported_program._verifiers = override_verifiers
@@ -222,9 +230,14 @@ def __call__(  # noqa: C901
                             # possible that the verifier will fail upon new ExportedProgram construction,
                             # and we should only run verification after each pass if
                             # run_checks_after_each_pass is True.
-                            new_graph_signature = _get_updated_graph_signature(
+                            _disambiguate_mutation_outputs(
                                 exported_program.graph_signature,
                                 res.graph_module,
+                                clone_op=clone_op,
+                            )
+                            _set_mutation_output_clone_op(res.graph_module, clone_op)
+                            new_graph_signature = _get_updated_graph_signature(
+                                exported_program.graph_signature, res.graph_module
                             )
                             # ExportedProgram.__init__ runs this, but this path mutates
                             # the program in place and never reconstructs it. Duplicate

exir/program/BUCK

@@ -35,6 +35,7 @@ fbcode_target(_kind = runtime.python_library,
         "//executorch/exir/backend:partitioner",
         "//executorch/exir/backend/canonical_partitioners:canonical_partitioner_lib",
         "//executorch/exir/capture:config",
+        "//executorch/exir/dialects:lib",
         "//executorch/exir/emit:emit",
         "//executorch/exir/emit:lib",
         "//executorch/exir/operator:util",

exir/program/_program.py

@@ -11,7 +11,7 @@
 import logging
 import os
 from collections import defaultdict
-from typing import Any, Dict, List, Optional, Sequence, Set, TextIO, Type, Union
+from typing import Any, Callable, Dict, List, Optional, Sequence, Set, TextIO, Type, Union
 
 import torch
 import torch._export
@@ -30,6 +30,7 @@
 from executorch.exir.backend.partitioner import Partitioner
 from executorch.exir.capture._config import EdgeCompileConfig, ExecutorchBackendConfig
 from executorch.exir.delegate import executorch_call_delegate, is_lowered_module
+from executorch.exir.dialects._ops import ops as exir_ops
 from executorch.exir.emit import emit_program, EmitterOutput
 from executorch.exir.emit._emitter import _DelegateDebugIdentifierMap
 from executorch.exir.error import ExportError
@@ -132,15 +133,19 @@ def wrapper(*args: Any, **kwargs: Any) -> Any:
 
 
 from executorch.exir._program_utils import (  # noqa: E402
+    _disambiguate_mutation_outputs,
+    _get_mutation_output_clone_op,
     _get_updated_graph_signature,
     _get_updated_range_constraints,
+    _set_mutation_output_clone_op,
 )
 
 
 def _transform(
     self,
     *passes: PassType,
     override_verifiers: None | list[Type[Verifier]] = None,
+    mutation_output_clone_op: Callable[..., torch.Tensor] | None = None,
 ) -> "ExportedProgram":
     """
     Transforms the program according to the provided passes.
@@ -150,6 +155,8 @@ def _transform(
         *passes: A sequence of passes to apply to the program
         override_verifiers: Optional list of verifier classes to use instead of the default verifiers.
             This is needed if the transforms yields illegal graph that the default verifier cannot handle.
+        mutation_output_clone_op: Clone operator for the transformed graph's dialect and layout mode.
+            Defaults to ATen cloning or the contract recorded during Edge conversion.
 
     Returns:
         ExportedProgram: A new ExportedProgram with the transformations applied, or self if no changes were made
@@ -160,14 +167,19 @@ def _transform(
     ), f"Expected all passes to be of PassType, not list or Verifier. Use override_verifiers kwarg instead. Got: {list(passes)}"
 
     return _transform_with_pass_manager(
-        self, ExportedProgramPassManager(list(passes)), override_verifiers
+        self,
+        ExportedProgramPassManager(list(passes)),
+        override_verifiers,
+        mutation_output_clone_op=mutation_output_clone_op,
     )
 
 
 def _transform_with_pass_manager(
     self: ExportedProgram,
     pass_manager: Union[ExportedProgramPassManager, GraphModulePassManager],
     override_verifiers: None | list[Type[Verifier]] = None,
+    *,
+    mutation_output_clone_op: Callable[..., torch.Tensor] | None = None,
 ) -> "ExportedProgram":
     """
     Transforms the program using the provided pass_manager.
@@ -184,7 +196,11 @@ def _transform_with_pass_manager(
         ExportedProgram: A new ExportedProgram with the transformations applied, or self if no changes were made
     """
     if isinstance(pass_manager, ExportedProgramPassManager):
-        res = pass_manager(self, override_verifiers)
+        res = pass_manager(
+            self,
+            override_verifiers,
+            mutation_output_clone_op=mutation_output_clone_op,
+        )
         if not res.modified:
             return self
         return res.exported_program
@@ -193,20 +209,30 @@ def _transform_with_pass_manager(
         if not res.modified:
             return self
         return _update_exported_program_graph_module(
-            self, res.graph_module, override_verifiers
+            self,
+            res.graph_module,
+            override_verifiers,
+            mutation_output_clone_op=mutation_output_clone_op,
         )
 
 
 def _update_exported_program_graph_module(
     exported_program: ExportedProgram,
     gm: torch.fx.GraphModule,
     override_verifiers: None | list[Type[Verifier]] = None,
+    *,
+    mutation_output_clone_op: Callable[..., torch.Tensor] | None = None,
 ) -> "ExportedProgram":
+    clone_op = _get_mutation_output_clone_op(exported_program, mutation_output_clone_op)
+    _disambiguate_mutation_outputs(
+        exported_program.graph_signature, gm, clone_op=clone_op
+    )
     transformed_ep = ExportedProgram(
         root=gm,
         graph=gm.graph,
         graph_signature=_get_updated_graph_signature(
-            exported_program.graph_signature, gm
+            exported_program.graph_signature,
+            gm,
         ),
         state_dict=exported_program.state_dict,
         range_constraints=_get_updated_range_constraints(
@@ -219,6 +245,7 @@ def _update_exported_program_graph_module(
     )
     transformed_ep.graph_module.meta.update(exported_program.graph_module.meta)
     transformed_ep.graph_module.meta.update(gm.meta)
+    _set_mutation_output_clone_op(transformed_ep.graph_module, clone_op)
     return transformed_ep
 
 
@@ -707,11 +734,21 @@ def _to_edge(ep, config: EdgeCompileConfig) -> "ExirExportedProgram":
         assert new_gm_res is not None
         new_gm = new_gm_res.graph_module
 
+    clone_op = (
+        exir_ops.edge.aten.clone.default
+        if config._skip_dim_order
+        else exir_ops.edge.dim_order_ops._clone_dim_order.default
+    )
+    _disambiguate_mutation_outputs(
+        new_ep.exported_program.graph_signature, new_gm, clone_op=clone_op
+    )
+    _set_mutation_output_clone_op(new_gm, clone_op)
     new_ep.exported_program = ExportedProgram(
         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,
         ),
         state_dict=new_ep.exported_program.state_dict,
         range_constraints=new_ep.exported_program.range_constraints,
@@ -848,10 +885,20 @@ def _generate_edge_program(
         assert gm_res is not None
         gm = gm_res.graph_module
 
+    clone_op = (
+        exir_ops.edge.aten.clone.default
+        if config._skip_dim_order
+        else exir_ops.edge.dim_order_ops._clone_dim_order.default
+    )
+    _disambiguate_mutation_outputs(program.graph_signature, gm, clone_op=clone_op)
+    _set_mutation_output_clone_op(gm, clone_op)
     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,
+        ),
         state_dict=program.state_dict,
         range_constraints=program.range_constraints,
         module_call_graph=program.module_call_graph,
@@ -1612,6 +1659,11 @@ def transform(
         """
 
         compile_config = compile_config or self.compile_config
+        clone_op = (
+            exir_ops.edge.aten.clone.default
+            if compile_config._skip_dim_order
+            else exir_ops.edge.dim_order_ops._clone_dim_order.default
+        )
         new_programs: Dict[str, ExportedProgram] = {}
 
         # Cast passes parameter upfront.
@@ -1640,11 +1692,17 @@ def transform(
 
             # Depending on the passes parameter, call the corresponding transform function.
             if passes_seq is not None:
-                new_programs[name] = _transform(program, *passes_seq)
+                new_programs[name] = _transform(
+                    program, *passes_seq, mutation_output_clone_op=clone_op
+                )
             elif passes_dict is not None:
-                new_programs[name] = _transform(program, *passes_dict[name])
+                new_programs[name] = _transform(
+                    program, *passes_dict[name], mutation_output_clone_op=clone_op
+                )
             elif pass_manager is not None:
-                new_programs[name] = _transform_with_pass_manager(program, pass_manager)
+                new_programs[name] = _transform_with_pass_manager(
+                    program, pass_manager, mutation_output_clone_op=clone_op
+                )
 
             # Verify the correctness of model graph after each transformation.
             EXIREdgeDialectVerifier(edge_compile_config=compile_config)(