From f8134151bc01d986deb9ba04759d6476be8da397 Mon Sep 17 00:00:00 2001 From: Matthias Cremon Date: Mon, 24 Aug 2026 15:41:43 -0700 Subject: [PATCH] Remove pass-local ExportPass replays (#22091) Summary: Remove full-graph `ExportPass` replay from four Tier 1 pass-local transformations. `BindOptionalArgsPass` and `RemoveBranchedQuantDequant` now update existing nodes in place, `ReplaceEmptyTensorsWithFullPass` copies metadata to its equivalent replacement node and recompiles modified nested modules, and `CombinePermuteWithMatmul` recomputes fake tensors only across its affected frontier. The shared cascaded-permute transform exposes an in-place entry point so Helios retains immediate cleanup without reconstructing the graph. These changes preserve graph, unaffected-node, and metadata identity while keeping no-op `modified=False` behavior. Each modifying invocation avoids one full replay; `CombinePermuteWithMatmul` can avoid two when its follow-up fusion also modifies the graph. Based on existing replay measurements, each avoided replay saves approximately 4-8 ms for 10 nodes, 40-75 ms for 100 nodes, and 0.4-0.75 s for 1,000 nodes. Already-normalized no-op invocations are unchanged. Reviewed By: aliafzal, ethansfng Differential Revision: D115472301 --- backends/cadence/aot/passes.py | 7 ++-- backends/cadence/aot/remove_ops.py | 12 +++--- backends/cadence/aot/replace_ops.py | 39 +++++++------------ backends/cadence/aot/simplify_ops.py | 17 ++++---- .../aot/tests/test_remove_ops_passes.py | 19 +++++++++ .../aot/tests/test_replace_ops_passes.py | 22 +++++++++++ .../aot/tests/test_simplify_ops_passes.py | 39 ++++++++++++++++++- backends/transforms/permute_pass_utils.py | 5 ++- .../test/test_permute_optimization_passes.py | 36 +++++++++++++++++ 9 files changed, 153 insertions(+), 43 deletions(-) diff --git a/backends/cadence/aot/passes.py b/backends/cadence/aot/passes.py index b0798a4a4e5..fb4df7d12c1 100644 --- a/backends/cadence/aot/passes.py +++ b/backends/cadence/aot/passes.py @@ -68,6 +68,7 @@ from executorch.exir.passes.scalar_to_tensor_pass import ScalarToTensorPass from executorch.exir.passes.spec_prop_pass import SpecPropPass from torch.export.exported_program import ExportedProgram +from torch.fx.passes.infra.pass_base import PassBase class InitializePipeline(ExportPass): @@ -104,7 +105,7 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult: # The passes that must run for the graph to be legal on the target, regardless # of compile mode. Everything else in the pipeline is an optimization. -REQUIRED_PASSES: frozenset[Type[ExportPass]] = frozenset( +REQUIRED_PASSES: frozenset[Type[PassBase]] = frozenset( { InitializePipeline, FinalizePipeline, @@ -139,7 +140,7 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult: ) -def _get_pipeline() -> list[Type[ExportPass]]: +def _get_pipeline() -> list[Type[PassBase]]: """The full ordered pass pipeline. Order is load-bearing and levels are interleaved, so this list is the single @@ -168,7 +169,7 @@ def _get_pipeline() -> list[Type[ExportPass]]: def get_passes( mode: CompileMode | str, edge_passes_config: Optional[EdgePassesConfig] = None, -) -> list[Type[ExportPass]]: +) -> list[Type[PassBase]]: # Coerce at the choke point: modes arrive from CLIs and JSON as plain # strings, and a str silently matches none of the checks below. mode = CompileMode(mode) diff --git a/backends/cadence/aot/remove_ops.py b/backends/cadence/aot/remove_ops.py index b12314d0ec3..7603db75bf1 100644 --- a/backends/cadence/aot/remove_ops.py +++ b/backends/cadence/aot/remove_ops.py @@ -40,6 +40,7 @@ from torch.export import ExportedProgram from torch.export.graph_signature import InputKind, OutputKind from torch.fx.node import Node +from torch.fx.passes.infra.pass_base import PassBase from torch.utils import _pytree as pytree @@ -727,7 +728,7 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult: return PassResult(graph_module, False) -class RemoveBranchedQuantDequant(ExportPass): +class RemoveBranchedQuantDequant(PassBase): """ This pass looks for adjacent quant and dequant nodes with identical parameters, where the quant node has other users in addition to the @@ -755,10 +756,9 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult: if modified: graph_module.graph.eliminate_dead_code() - result = super().call(graph_module) - return result + graph_module.recompile() - return PassResult(graph_module, False) + return PassResult(graph_module, modified) def remove_branched( self, @@ -848,7 +848,7 @@ def maybe_remove_or_replace(self, node: Node) -> bool: class CommonRemovePasses: - passes: List[Type[ExportPass]] = [ + passes: List[Type[PassBase]] = [ # Canonicalise squeeze/unsqueeze to view_copy first: the nop-view and # permute passes below both reason about view_copy only. ReplaceSqueezeAndUnsqueezeWithViewPassImported, @@ -964,7 +964,7 @@ def call(self, exported_program: ExportedProgram) -> ExportedProgramPassResult: class CadenceRemoveNops: - passes: List[Type[ExportPass]] = CommonRemovePasses.passes + [ + passes: List[Type[PassBase]] = CommonRemovePasses.passes + [ SimplifySliceOpPass, RemoveNopRequantizeOpPass, RemoveZeroSizedConstantPadNd, diff --git a/backends/cadence/aot/replace_ops.py b/backends/cadence/aot/replace_ops.py index e0d837e53ef..93624540c5b 100644 --- a/backends/cadence/aot/replace_ops.py +++ b/backends/cadence/aot/replace_ops.py @@ -37,7 +37,8 @@ ) from executorch.exir.dialects._ops import ops as exir_ops from executorch.exir.dialects.edge._ops import EdgeOpOverload -from executorch.exir.pass_base import ExportPass, PassResult +from executorch.exir.pass_base import PassResult +from torch.fx.passes.infra.pass_base import PassBase # A map to represent ops that: # (a) are functionally equivalent; and @@ -1981,47 +1982,37 @@ def maybe_remove_or_replace(self, node: torch.fx.Node) -> bool: return True -class ReplaceEmptyTensorsWithFullPass(ExportPass): +class ReplaceEmptyTensorsWithFullPass(PassBase): """Replaces nodes that produce empty tensors with full nodes.""" - def call_operator(self, op, args, kwargs, meta): - val = meta.data.get("val", None) - if isinstance(val, torch.Tensor) and val.numel() == 0: - return super().call_operator( - exir_ops.edge.aten.full.default, - args=(val.shape, 0), - kwargs={"dtype": val.dtype}, - meta=meta, - ) - return super().call_operator(op, args, kwargs, meta) - def call(self, graph_module: torch.fx.GraphModule) -> PassResult: changed = False for module in filter( lambda m: isinstance(m, torch.fx.GraphModule), graph_module.modules() ): - module = cast(torch.fx.GraphModule, module) - for node in module.graph.nodes: + graph_module_to_update = cast(torch.fx.GraphModule, module) + module_changed = False + for node in graph_module_to_update.graph.nodes: if node.op != "call_function": continue val = node.meta.get("val", None) if isinstance(val, torch.Tensor) and val.numel() == 0: - with module.graph.inserting_before(node): - new_node = module.graph.call_function( + with graph_module_to_update.graph.inserting_before(node): + new_node = graph_module_to_update.graph.call_function( exir_ops.edge.aten.full.default, args=(val.shape, 0), kwargs={"dtype": val.dtype}, ) - new_node.meta = node.meta + new_node.meta = node.meta.copy() node.replace_all_uses_with(new_node) - changed = True + module_changed = True - if changed: - graph_module.graph.eliminate_dead_code() - graph_module.recompile() - return super().call(graph_module) + if module_changed: + graph_module_to_update.graph.eliminate_dead_code() + graph_module_to_update.recompile() + changed = True - return PassResult(graph_module, False) + return PassResult(graph_module, changed) class ReplaceWhereWithFullArgsWithWhereScalar(RemoveOrReplacePassInterface): diff --git a/backends/cadence/aot/simplify_ops.py b/backends/cadence/aot/simplify_ops.py index 83c974a898b..72a108bddde 100644 --- a/backends/cadence/aot/simplify_ops.py +++ b/backends/cadence/aot/simplify_ops.py @@ -17,8 +17,9 @@ from executorch.backends.cadence.aot.utils import rebind from executorch.exir.dialects._ops import ops as exir_ops from executorch.exir.dialects.edge._ops import EdgeOpOverload -from executorch.exir.pass_base import ExportPass, PassResult +from executorch.exir.pass_base import PassResult from torch.fx import Node +from torch.fx.passes.infra.pass_base import PassBase class SimplifySliceOpPass(RemoveOrReplacePassInterface): @@ -169,7 +170,7 @@ def maybe_remove_or_replace(self, node: Node) -> bool: return True -class BindOptionalArgsPass(ExportPass): +class BindOptionalArgsPass(PassBase): """Bind all optional args and kwargs.""" def call(self, graph_module: torch.fx.GraphModule) -> PassResult: @@ -181,7 +182,9 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult: for module in filter( lambda m: isinstance(m, torch.fx.GraphModule), graph_module.modules() ): - for node in cast(torch.fx.GraphModule, module).graph.nodes: + graph_module_to_update = cast(torch.fx.GraphModule, module) + module_modified = False + for node in graph_module_to_update.graph.nodes: if node.op != "call_function": continue if not isinstance(node.target, EdgeOpOverload): @@ -198,11 +201,11 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult: if new_args != node.args or new_kwargs != node.kwargs: node.args = new_args node.kwargs = new_kwargs - modified = True + module_modified = True - if modified: - graph_module.recompile() - return super().call(graph_module) + if module_modified: + graph_module_to_update.recompile() + modified = True return PassResult(graph_module, modified) diff --git a/backends/cadence/aot/tests/test_remove_ops_passes.py b/backends/cadence/aot/tests/test_remove_ops_passes.py index 3931b697266..1bce74d9ab9 100644 --- a/backends/cadence/aot/tests/test_remove_ops_passes.py +++ b/backends/cadence/aot/tests/test_remove_ops_passes.py @@ -913,9 +913,28 @@ def test_remove_dequant_on_branch(self) -> None: ) builder.output([x1_output, y1_output]) original = builder.get_graph_module() + original_nodes = list(original.graph.nodes) + original_abs_nodes = original.graph.find_nodes( + op="call_function", target=exir_ops.edge.aten.abs.default + ) + x1_output_node = original_abs_nodes[1] + x1_output_node.meta["remove_branched_test"] = "preserved" + x1_output_meta = x1_output_node.meta pass_result = cast(PassResult, RemoveBranchedQuantDequant()(original)) self.assertTrue(pass_result.modified) graph_after_passes = pass_result.graph_module + self.assertIs(graph_after_passes, original) + self.assertIs(x1_output_node.meta, x1_output_meta) + self.assertEqual(x1_output_node.meta["remove_branched_test"], "preserved") + self.assertIs( + next( + node + for node in graph_after_passes.graph.nodes + if node.name == x1_output_node.name + ), + x1_output_node, + ) + self.assertLess(len(list(graph_after_passes.graph.nodes)), len(original_nodes)) self.assertEqual( count_node( graph_after_passes, diff --git a/backends/cadence/aot/tests/test_replace_ops_passes.py b/backends/cadence/aot/tests/test_replace_ops_passes.py index 99bb56b652f..a5a712d6349 100644 --- a/backends/cadence/aot/tests/test_replace_ops_passes.py +++ b/backends/cadence/aot/tests/test_replace_ops_passes.py @@ -3179,10 +3179,32 @@ def test_empty_slice(self) -> None: # Deepcopy before the pass gm_before = copy.deepcopy(gm) + original_graph_module = gm + original_nodes = list(gm.graph.nodes) + empty_slice = gm.graph.find_nodes( + op="call_function", target=exir_ops.edge.aten.slice_copy.Tensor + )[0] + empty_slice.meta["replace_empty_test"] = "preserved" + empty_slice_meta = empty_slice.meta + unchanged_slice = gm.graph.find_nodes( + op="call_function", target=exir_ops.edge.aten.slice_copy.Tensor + )[1] + unchanged_slice_meta = unchanged_slice.meta result = ReplaceEmptyTensorsWithFullPass().call(gm) self.assertTrue(result.modified) updated_gm = result.graph_module + self.assertIs(updated_gm, original_graph_module) + self.assertIs(unchanged_slice.meta, unchanged_slice_meta) + self.assertIn(unchanged_slice, updated_gm.graph.nodes) + self.assertNotIn(empty_slice, updated_gm.graph.nodes) + full = updated_gm.graph.find_nodes( + op="call_function", target=exir_ops.edge.aten.full.default + )[0] + self.assertIsNot(full.meta, empty_slice_meta) + self.assertEqual(full.meta, empty_slice_meta) + self.assertEqual(full.meta["replace_empty_test"], "preserved") + self.assertEqual(len(list(updated_gm.graph.nodes)), len(original_nodes)) # Validate numerical accuracy inputs = [x_input] diff --git a/backends/cadence/aot/tests/test_simplify_ops_passes.py b/backends/cadence/aot/tests/test_simplify_ops_passes.py index 79d89c24713..30fd3a55336 100644 --- a/backends/cadence/aot/tests/test_simplify_ops_passes.py +++ b/backends/cadence/aot/tests/test_simplify_ops_passes.py @@ -147,14 +147,49 @@ def test_simplify_slice_op_args(self) -> None: args=(x, 1), kwargs={"end": 3}, ) + original_graph_module = gm original_slice_copy = list(gm.graph.nodes)[1] + original_metadata = original_slice_copy.meta + original_slice_copy.meta["bind_optional_args_test"] = "preserved" self.assertEqual(original_slice_copy.args[1:], (1,)) self.assertEqual(original_slice_copy.kwargs, {"end": 3}) + result = transform_and_check_numerics( gm, (x,), BindOptionalArgsPass(), "BindOptionalArgsPass" ) + self.assertTrue(result.modified) - gm = result.graph_module - modified_slice_copy = list(gm.graph.nodes)[1] + self.assertIs(result.graph_module, original_graph_module) + modified_slice_copy = list(result.graph_module.graph.nodes)[1] + self.assertIs(modified_slice_copy, original_slice_copy) + self.assertIs(modified_slice_copy.meta, original_metadata) + self.assertEqual( + modified_slice_copy.meta["bind_optional_args_test"], "preserved" + ) self.assertEqual(modified_slice_copy.args[1:], (1, None, 3, 1)) self.assertEqual(modified_slice_copy.kwargs, {}) + + def test_bind_optional_args_no_change(self) -> None: + x = torch.rand(4, 5) + gm = single_op_builder( + placeholders=(x,), + op=exir_ops.edge.aten.slice_copy.Tensor, + args=(x, 1, None, 3, 1), + ) + original_nodes = list(gm.graph.nodes) + slice_copy = original_nodes[1] + original_metadata = slice_copy.meta + slice_copy.meta["bind_optional_args_test"] = "preserved" + + result = BindOptionalArgsPass().call(gm) + + self.assertFalse(result.modified) + self.assertIs(result.graph_module, gm) + self.assertEqual( + len(list(result.graph_module.graph.nodes)), len(original_nodes) + ) + self.assertIs(list(result.graph_module.graph.nodes)[1], slice_copy) + self.assertIs(slice_copy.meta, original_metadata) + self.assertEqual(slice_copy.meta["bind_optional_args_test"], "preserved") + self.assertEqual(slice_copy.args[1:], (1, None, 3, 1)) + self.assertEqual(slice_copy.kwargs, {}) diff --git a/backends/transforms/permute_pass_utils.py b/backends/transforms/permute_pass_utils.py index fca8946165e..b91c90d26f8 100644 --- a/backends/transforms/permute_pass_utils.py +++ b/backends/transforms/permute_pass_utils.py @@ -157,8 +157,11 @@ def _apply_hierarchical_inplace(self, graph_module: torch.fx.GraphModule) -> boo modified |= self._apply_flat_inplace(module) return modified + def apply_inplace(self, graph_module: torch.fx.GraphModule) -> bool: + return self._apply_hierarchical_inplace(graph_module) + def call(self, graph_module: torch.fx.GraphModule) -> PassResult: - modified = self._apply_hierarchical_inplace(graph_module) + modified = self.apply_inplace(graph_module) if modified: graph_module.graph.eliminate_dead_code() graph_module.recompile() diff --git a/backends/transforms/test/test_permute_optimization_passes.py b/backends/transforms/test/test_permute_optimization_passes.py index cf4583357d7..1264f5266c2 100644 --- a/backends/transforms/test/test_permute_optimization_passes.py +++ b/backends/transforms/test/test_permute_optimization_passes.py @@ -125,6 +125,42 @@ def test_permute_transpose_fusion(self) -> None: "FuseCascadedTransposeOrPermuteOps", ) + def test_apply_inplace_preserves_identity_and_output_metadata(self) -> None: + builder = GraphBuilder() + x = builder.placeholder("x", torch.randn(2, 3, 4)) + permute1 = builder.call_operator( + op=exir_ops.edge.aten.permute_copy.default, args=(x, [0, 2, 1]) + ) + permute2 = builder.call_operator( + op=exir_ops.edge.aten.permute_copy.default, args=(permute1, [2, 1, 0]) + ) + builder.output([permute2]) + graph_module = builder.get_graph_module() + placeholder = next( + node for node in graph_module.graph.nodes if node.op == "placeholder" + ) + placeholder.meta["apply_inplace_test"] = "preserved" + placeholder_meta = placeholder.meta + output_permute = graph_module.graph.find_nodes( + op="call_function", target=exir_ops.edge.aten.permute_copy.default + )[1] + output_permute.meta["apply_inplace_test"] = "output" + output_meta = output_permute.meta + + modified = FuseCascadedTransposeOrPermuteOps().apply_inplace(graph_module) + graph_module.graph.eliminate_dead_code() + graph_module.recompile() + + self.assertTrue(modified) + self.assertIn(placeholder, graph_module.graph.nodes) + self.assertIs(placeholder.meta, placeholder_meta) + self.assertEqual(placeholder.meta["apply_inplace_test"], "preserved") + fused_permute = graph_module.graph.find_nodes( + op="call_function", target=exir_ops.edge.aten.permute_copy.default + )[0] + self.assertIs(fused_permute.meta, output_meta) + self.assertEqual(fused_permute.meta["apply_inplace_test"], "output") + def test_cascaded_permutes_multiple_users(self) -> None: builder = GraphBuilder() x = builder.placeholder("x", torch.randn(2, 3, 4, 5))