diff --git a/exir/pass_base.py b/exir/pass_base.py index 6071aae2be8..cf7f0117c29 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,32 @@ class ExportedProgramPassBase(ABC): def __call__(self, exported_program: ExportedProgram) -> ExportedProgramPassResult: """ Runs the precondition check, the pass itself, and the postcondition check. + + Prefer the node replacement APIs (``replace_all_uses_with``, + ``replace_input_with``) in ``call``: a replace hook keeps the signature + valid as the graph changes. Output specs are realigned with the graph + afterwards regardless, so ``ensures`` sees a self-consistent program. """ self.requires(exported_program) - res = self.call(exported_program) - self.ensures(exported_program) - return res + signature = exported_program.graph_signature + replace_hook = signature.get_replace_hook() + hook_modified = False + + def tracking_hook(old: torch.fx.Node, new: str, user: torch.fx.Node) -> None: + nonlocal hook_modified + if user.op == "output" and old.name != new: + hook_modified = True + replace_hook(old=old, new=new, user=user) + + with exported_program.graph_module._set_replace_hook(tracking_hook): + res = self.call(exported_program) + signature_modified = _sync_output_specs(res.exported_program) + self.ensures(res.exported_program) + return ExportedProgramPassResult( + res.exported_program, + res.modified or hook_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..2f59b26c741 100644 --- a/exir/tests/test_pass_infra.py +++ b/exir/tests/test_pass_infra.py @@ -579,3 +579,85 @@ 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()) + + def test_the_replace_hook_updates_the_signature_during_the_pass(self) -> None: + """A pass using the replacement APIs sees a valid signature as it runs.""" + + signature_during_pass = [] + + class ReplaceViaApiPass(ExportedProgramPassBase): + def call(self, ep: ExportedProgram) -> ExportedProgramPassResult: + graph = ep.graph_module.graph + (old,) = graph.output_node().args[0] + with graph.inserting_before(graph.output_node()): + new = graph.call_function( + exir_ops.edge.aten.mul.Tensor, (old.args[0], old.args[1]) + ) + new.meta = dict(old.meta) + old.replace_all_uses_with(new) + signature_during_pass.append( + ep.graph_signature.output_specs[0].arg.name + ) + return ExportedProgramPassResult(ep, True) + + result = ReplaceViaApiPass()(self._program()) + + graph_output_name = result.exported_program.graph.output_node().args[0][0].name + self.assertEqual(signature_during_pass, [graph_output_name])