From e6e26183bc70dea3ee2ee63b45dfd8ddb0c514e2 Mon Sep 17 00:00:00 2001 From: Matthias Cremon Date: Tue, 15 Sep 2026 22:07:22 -0700 Subject: [PATCH] Keep output specs in sync in ExportedProgramPassBase Summary: A pass that replaces the graph's output node leaves `output_specs` naming the node that is no longer there. Nothing catches it at the time; it surfaces later as a signature that disagrees with the graph, and every pass author has to remember to fix it up by hand. `ExportedProgramPassBase.__call__` now realigns the specs with the graph between the pass and the postcondition check: ``` self.requires(exported_program) res = self.call(exported_program) signature_modified = _sync_output_specs(res.exported_program) self.ensures(res.exported_program) ``` Rewriting a spec to name the current output node is always safe. Changing an output between a node and a literal is not, so `_sync_output_specs` raises rather than guessing, and a length mismatch between outputs and specs raises too -- a pass that adds or removes an output is expected to maintain its own signature. Two consequences worth calling out: - A pass that only renames an output now reports `modified=True` even if it returned False, because the signature did change. That is what callers driving a pass to fixpoint need to see. - `ensures()` now receives the pass's result rather than the program that went in. The previous behaviour looks like an oversight -- a postcondition check that inspects the pre-pass program cannot check much -- and no pass in the tree relies on it: of the subclasses of `ExportedProgramPassBase`, exactly one overrides `ensures()`, and it calls `exported_program.validate()`, which wants the result. This was previously implemented as a `fused_quant`-local subclass. It is not specific to that package, so it moves here; the subclass goes away in the diff above this one. Differential Revision: D120259092 --- exir/pass_base.py | 58 ++++++++++++++++++++++++++++++++-- exir/tests/test_pass_infra.py | 59 +++++++++++++++++++++++++++++++++++ 2 files changed, 114 insertions(+), 3 deletions(-) diff --git a/exir/pass_base.py b/exir/pass_base.py index 6071aae2be8..ee967eef30c 100644 --- a/exir/pass_base.py +++ b/exir/pass_base.py @@ -10,7 +10,7 @@ import traceback from abc import ABC, abstractmethod from contextlib import nullcontext -from dataclasses import dataclass +from dataclasses import dataclass, replace from typing import ( Any, Callable, @@ -38,6 +38,7 @@ from torch._subclasses.fake_tensor import FakeTensor from torch._subclasses.functional_tensor import FunctionalTensor, FunctionalTensorMode from torch.export import ExportedProgram +from torch.export.graph_signature import ConstantArgument from torch.fx import traceback as fx_traceback from torch.fx.experimental.proxy_tensor import PythonKeyTracer from torch.fx.graph import CodeGen @@ -237,6 +238,50 @@ class ExportedProgramPassResult: modified: bool +def _sync_output_specs(exported_program: ExportedProgram) -> bool: + """Realign ``output_specs`` with the graph's output node. + + A pass that replaces an output node leaves the signature naming the old one, + which later stages then disagree with. Rewriting the spec to match is always + safe; changing an output between a node and a literal is not, so that raises + rather than guessing. + + Returns whether any spec changed. + """ + output_node = exported_program.graph_module.graph.output_node() + outputs = output_node.args[0] + assert isinstance(outputs, (tuple, list)) + + output_specs = exported_program.graph_signature.output_specs + if len(outputs) != len(output_specs): + raise ExportPassBaseError( + f"Graph has {len(outputs)} outputs, but its signature has " + f"{len(output_specs)} output specs" + ) + + modified = False + for output, output_spec in zip(outputs, output_specs): + if isinstance(output_spec.arg, ConstantArgument): + if isinstance(output, torch.fx.Node): + raise ExportPassBaseError( + f"Output {output.name} replaced a literal output; changing output " + "representation is not supported" + ) + if output_spec.arg.value != output: + output_spec.arg = replace(output_spec.arg, value=output) + modified = True + continue + if not isinstance(output, torch.fx.Node): + raise ExportPassBaseError( + f"Output {output_spec.arg.name} became a literal; changing output " + "representation is not supported" + ) + if output_spec.arg.name != output.name: + output_spec.arg = replace(output_spec.arg, name=output.name) + modified = True + return modified + + class ExportedProgramPassBase(ABC): """ Base interface for implementing passes that operate on ExportedProgram. @@ -245,12 +290,19 @@ class ExportedProgramPassBase(ABC): def __call__(self, exported_program: ExportedProgram) -> ExportedProgramPassResult: """ Runs the precondition check, the pass itself, and the postcondition check. + + Output specs are realigned with the graph between the pass and the + postcondition check, so ``ensures`` sees a self-consistent program. """ self.requires(exported_program) res = self.call(exported_program) - self.ensures(exported_program) - return res + signature_modified = _sync_output_specs(res.exported_program) + self.ensures(res.exported_program) + return ExportedProgramPassResult( + res.exported_program, + res.modified or signature_modified, + ) @abstractmethod def call(self, exported_program: ExportedProgram) -> ExportedProgramPassResult: diff --git a/exir/tests/test_pass_infra.py b/exir/tests/test_pass_infra.py index 16ed5af4180..f5367ed668b 100644 --- a/exir/tests/test_pass_infra.py +++ b/exir/tests/test_pass_infra.py @@ -579,3 +579,62 @@ def placeholder( new_input = self._find_input_node(new_graph_module) self.assertNotEqual(self._symbolic_input_shape(new_input), original_snapshot) + + +class ExportedProgramPassBaseOutputSpecTest(unittest.TestCase): + """__call__ realigns output specs with the graph before ensures() runs.""" + + class _Model(torch.nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + return x + x + + def _program(self) -> ExportedProgram: + return to_edge( + export(self._Model(), (torch.randn(2, 2),)) + ).exported_program() + + def test_replacing_the_output_node_updates_the_signature(self) -> None: + class ReplaceOutputPass(ExportedProgramPassBase): + def call(self, ep: ExportedProgram) -> ExportedProgramPassResult: + graph = ep.graph_module.graph + output_node = graph.output_node() + (old,) = output_node.args[0] + with graph.inserting_before(output_node): + new = graph.call_function( + exir_ops.edge.aten.mul.Tensor, (old.args[0], old.args[1]) + ) + new.meta = dict(old.meta) + output_node.args = ((new,),) + return ExportedProgramPassResult(ep, True) + + program = self._program() + result = ReplaceOutputPass()(program) + + (spec,) = result.exported_program.graph_signature.output_specs + graph_output_name = result.exported_program.graph.output_node().args[0][0].name + self.assertEqual(spec.arg.name, graph_output_name) + result.exported_program.validate() + + def test_a_signature_only_change_is_reported_as_modified(self) -> None: + """A pass that renames the output reports modified even if it says False.""" + + class RenameOutputPass(ExportedProgramPassBase): + def call(self, ep: ExportedProgram) -> ExportedProgramPassResult: + ep.graph.output_node().args[0][0].name = "renamed_output" + return ExportedProgramPassResult(ep, False) + + result = RenameOutputPass()(self._program()) + + self.assertTrue(result.modified) + (spec,) = result.exported_program.graph_signature.output_specs + self.assertEqual(spec.arg.name, "renamed_output") + + def test_output_count_mismatch_is_rejected(self) -> None: + class DropOutputPass(ExportedProgramPassBase): + def call(self, ep: ExportedProgram) -> ExportedProgramPassResult: + output_node = ep.graph.output_node() + output_node.args = ((*output_node.args[0], output_node.args[0][0]),) + return ExportedProgramPassResult(ep, True) + + with self.assertRaisesRegex(ExportPassBaseError, "output specs"): + DropOutputPass()(self._program())