reinplace_pass resyncs output_specs after in-place rewrites

Differential Revision: D111753298

Pull Request resolved: #20893

Author: metascroy

backends/mlx/passes.py

@@ -99,49 +99,8 @@ def call(self, exported_program) -> ExportedProgramPassResult:
         }
         if ops_to_inplace:
             reinplace_pass(exported_program, ops_to_inplace=ops_to_inplace)
-            self._resync_output_specs(exported_program)
         return ExportedProgramPassResult(exported_program, True)
 
-    @staticmethod
-    def _resync_output_specs(exported_program) -> None:
-        """Re-sync graph-signature output names after reinplace.
-
-        ``reinplace_pass`` rewrites an output-producing node (e.g. the final
-        ``exp`` -> ``exp_``) via ``replace_all_uses_with`` + erase, but does not
-        update ``graph_signature.output_specs``. Output order is preserved, so we
-        positionally re-sync each spec's argument name to the current output node
-        arg; otherwise ``ExportedProgram.validate()`` (run by the pass manager)
-        raises a SpecViolationError.
-
-        This positional pairing is only valid because reinplace does a 1:1
-        ``replace_all_uses_with`` + erase and never drops, adds, or reorders
-        outputs. We assert ``len(output_specs) == len(out_args)`` so that a
-        future change violating that invariant fails loudly here instead of
-        silently mis-pairing names (``zip`` would otherwise truncate). Specs
-        whose ``arg`` is not a named tensor (e.g. ``ConstantArgument``) carry no
-        ``name`` and are skipped by the ``getattr`` guard below.
-        """
-        out_node = next(
-            n for n in reversed(exported_program.graph.nodes) if n.op == "output"
-        )
-        out_args = out_node.args[0]
-        if not isinstance(out_args, (tuple, list)):
-            out_args = (out_args,)
-        output_specs = exported_program.graph_signature.output_specs
-        assert len(output_specs) == len(out_args), (
-            "reinplace changed graph output count: "
-            f"{len(output_specs)} output_specs vs {len(out_args)} output args. "
-            "Positional output-spec re-sync assumes a 1:1, order-preserving "
-            "rewrite."
-        )
-        for spec, arg in zip(output_specs, out_args):
-            if (
-                isinstance(arg, torch.fx.Node)
-                and getattr(spec.arg, "name", None) is not None
-                and spec.arg.name != arg.name
-            ):
-                spec.arg.name = arg.name
-
 
 @dataclass
 class RMSNormMatch(PatternMatch):

exir/passes/reinplace.py

@@ -507,4 +507,26 @@ def reinplace_pass(  # noqa: C901
             continue
 
         seen_nodes.update(node.all_input_nodes)
+
+    # Reinplace rewrites may rename output nodes (e.g. aten_relu_default →
+    # aten_relu__default) but replace_all_uses_with does not update the
+    # graph signature. Fix output_specs to match the actual output node args.
+    output_node = next((n for n in ep.graph.nodes if n.op == "output"), None)
+    if output_node is not None:
+        out_args = output_node.args[0]
+        if not isinstance(out_args, (tuple, list)):
+            out_args = (out_args,)
+        output_specs = ep.graph_signature.output_specs
+        assert len(output_specs) == len(out_args), (
+            f"reinplace: output spec count changed: "
+            f"{len(output_specs)} specs vs {len(out_args)} output args"
+        )
+        for spec, arg in zip(output_specs, out_args):
+            if (
+                isinstance(arg, torch.fx.Node)
+                and getattr(spec.arg, "name", None) is not None
+                and spec.arg.name != arg.name
+            ):
+                spec.arg.name = arg.name
+
     return ep

exir/tests/test_reinplace_pass.py

@@ -326,6 +326,25 @@ def forward(
         # kernels in the other tests in this file.
         edge.to_executorch()
 
+    def test_output_specs_updated_after_reinplace(self) -> None:
+        """reinplace_pass must update graph_signature.output_specs to match
+        the renamed output nodes (e.g. relu -> relu_)."""
+
+        class M(torch.nn.Module):
+            def forward(self, x: torch.Tensor) -> torch.Tensor:
+                return torch.relu(x + 1.0)
+
+        ep = export(M(), (torch.randn(4),), strict=True)
+        edge_program = to_edge(ep).exported_program()
+
+        custom_set = {edge_ops.edge.aten.relu.default}
+        ep = reinplace_pass(edge_program, ops_to_inplace=custom_set)
+
+        output_node = next(n for n in ep.graph.nodes if n.op == "output")
+        actual_names = [arg.name for arg in output_node.args[0] if hasattr(arg, "name")]
+        spec_names = [s.arg.name for s in ep.graph_signature.output_specs]
+        self.assertEqual(actual_names, spec_names)
+
     def test_broadcasting_self_not_reinplaced(self) -> None:
         """An op whose mutated arg (self) broadcasts up to a larger output
         must NOT be reinplaced: the in-place form cannot grow self