From 0f51a4af4d2272f0f8e871b3e5abc684da06cbd3 Mon Sep 17 00:00:00 2001 From: Andrew Pullin Date: Mon, 10 Aug 2026 08:37:20 -0700 Subject: [PATCH 1/5] Cache param/buffer/constant names in RemovePermutesAroundElementwiseTosaOps Summary: RemovePermutesAroundElementwiseTosaOps._is_constant called is_param_node() for every node reached during recursive visit() walk. Each is_param_node rebuilds immutable dict over full input_specs via uncached graph_signature properties. This made pass O(nodes * inputs). Precompute union of param/buffer/lifted-constant placeholder names once in __init__ and use O(1) set membership. Semantics unchanged. Differential Revision: D114224782 --- ...remove_permutes_around_elementwise_tosa_ops.py | 15 ++++++++++++--- 1 file changed, 12 insertions(+), 3 deletions(-) diff --git a/backends/arm/_passes/remove_permutes_around_elementwise_tosa_ops.py b/backends/arm/_passes/remove_permutes_around_elementwise_tosa_ops.py index b241038f7a9..02d8cad0817 100644 --- a/backends/arm/_passes/remove_permutes_around_elementwise_tosa_ops.py +++ b/backends/arm/_passes/remove_permutes_around_elementwise_tosa_ops.py @@ -5,7 +5,6 @@ import torch -from executorch.backends.arm._passes.arm_pass_utils import is_param_node from executorch.backends.arm._passes.insert_table_ops import TableOps from executorch.backends.transforms.remove_permutes_around_elementwise_ops import ( RemovePermutesAroundElementwiseOps, @@ -25,10 +24,20 @@ def __init__(self, exported_program: ExportedProgram) -> None: } ) self.exported_program = exported_program + # Precompute parameter/buffer/lifted-constant placeholder names once. + # is_param_node() rebuilds these graph_signature maps on every call, so + # calling it per node made the visit() walk O(nodes * inputs). + gs = exported_program.graph_signature + self._constant_input_names: set[str] = ( + set(gs.inputs_to_parameters) + | set(gs.inputs_to_buffers) + | set(gs.inputs_to_lifted_tensor_constants) + ) def _is_constant(self, node: torch.fx.Node) -> bool: - # Override fragile string match check with exported program check - return super()._is_constant(node) or is_param_node(self.exported_program, node) + # get_attr nodes are handled by super()._is_constant; set membership + # here is equivalent to is_param_node for placeholder inputs. + return super()._is_constant(node) or node.name in self._constant_input_names def permute_subgraph(self, subgraph) -> bool: # TABLE lookup inputs are already tied to the table layout. From 3a540e0517f5393ade9942f1bd6329be53fedef5 Mon Sep 17 00:00:00 2001 From: Andrew Pullin Date: Mon, 10 Aug 2026 08:37:20 -0700 Subject: [PATCH 2/5] Speed up permute propagation cleanup Summary: Speed up permute propagation cleanup. ARM lowering spends significant time in permute propagation. Reduce avoidable repeated work while keeping same optimizations: - Canonicalize view/permute chain collection now uses deque + membership set instead of list pop(0) and remove(), avoiding quadratic bookkeeping. - FuseDuplicateUsersPass deduplicates pending producer revisits while preserving same revisit behavior after fusions. - PropagateViewCopyPermutePass still retraces after each moved transform for metadata safety, but defers horizontal/vertical cleanup until full scan finds no more direct propagation moves. Preserves fixed-point behavior while avoiding expensive cleanup after every single moved transform. Behavior-preserving speedup, all existing pass tests pass. Differential Revision: D114224762 --- .../canonicalize_view_copy_permute_pass.py | 16 +++++++++---- .../arm/_passes/fuse_duplicate_users_pass.py | 12 ++++++++-- .../propagate_view_copy_permute_pass.py | 23 +++++++++++-------- 3 files changed, 35 insertions(+), 16 deletions(-) diff --git a/backends/arm/_passes/canonicalize_view_copy_permute_pass.py b/backends/arm/_passes/canonicalize_view_copy_permute_pass.py index 1930ba0b5f9..00790079373 100644 --- a/backends/arm/_passes/canonicalize_view_copy_permute_pass.py +++ b/backends/arm/_passes/canonicalize_view_copy_permute_pass.py @@ -5,6 +5,7 @@ from __future__ import annotations +from collections import deque from typing import cast, Sequence, Set, Type import torch @@ -91,20 +92,25 @@ def _collect_chains(self, graph_module: GraphModule) -> list[list[Node]]: """Returns a list of linear chains of view/permutes in the graph.""" chains: list[list[Node]] = [] - view_permute_nodes = [ + view_permute_nodes = deque( node for node in graph_module.graph.nodes if node.target in self._TARGETS - ] + ) + remaining = set(view_permute_nodes) while view_permute_nodes: - node = view_permute_nodes.pop(0) + node = view_permute_nodes.popleft() + if node not in remaining: + continue + remaining.remove(node) + chain = [node] current = node while len(current.users) == 1: user = next(iter(current.users)) - if user.target not in self._TARGETS: + if user.target not in self._TARGETS or user not in remaining: break - view_permute_nodes.remove(user) + remaining.remove(user) chain.append(user) current = user diff --git a/backends/arm/_passes/fuse_duplicate_users_pass.py b/backends/arm/_passes/fuse_duplicate_users_pass.py index 9bd21112569..746dade94bd 100644 --- a/backends/arm/_passes/fuse_duplicate_users_pass.py +++ b/backends/arm/_passes/fuse_duplicate_users_pass.py @@ -45,9 +45,17 @@ def call(self, graph_module: GraphModule) -> PassResult: node_order = {node: index for index, node in enumerate(graph.nodes)} producers: Deque[Node] = deque(node for node in graph.nodes) + queued_producers: Set[Node] = set(producers) + + def enqueue_producer(node: Node) -> None: + if node.graph is None or node in queued_producers: + return + producers.append(node) + queued_producers.add(node) while producers: producer = producers.popleft() + queued_producers.discard(producer) if producer.graph is None: # Node was deleted by a previous rewrite while still queued. @@ -84,8 +92,8 @@ def call(self, graph_module: GraphModule) -> PassResult: # Revisit the current producer and the surviving user so that # newly formed duplicate chains can be fused in later # iterations. - producers.append(producer) - producers.append(representative) + enqueue_producer(producer) + enqueue_producer(representative) if modified: graph_module.recompile() diff --git a/backends/arm/_passes/propagate_view_copy_permute_pass.py b/backends/arm/_passes/propagate_view_copy_permute_pass.py index a54c8f28312..d38dd202fe1 100644 --- a/backends/arm/_passes/propagate_view_copy_permute_pass.py +++ b/backends/arm/_passes/propagate_view_copy_permute_pass.py @@ -108,20 +108,25 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult: continue if self._propagate(node): iteration_modified = True + graph_module = self._retrace(graph_module) break if iteration_modified: - graph_module = self._retrace(graph_module) - result = self.fuse_horizontal(graph_module) - graph_module = result.graph_module - iteration_modified |= result.modified - result = self.fuse_vertical(graph_module) - graph_module = result.graph_module - iteration_modified |= result.modified + modified = True + continue + + result = self.fuse_horizontal(graph_module) + graph_module = result.graph_module + iteration_modified |= result.modified + result = self.fuse_vertical(graph_module) + graph_module = result.graph_module + iteration_modified |= result.modified modified |= iteration_modified - if not iteration_modified: - break + if iteration_modified: + graph_module = self._retrace(graph_module) + continue + break if modified: graph_module = self._retrace(graph_module) From d882dee4aa30b51679acf9ac74ec8911430a57bf Mon Sep 17 00:00:00 2001 From: Andrew Pullin Date: Mon, 10 Aug 2026 08:37:20 -0700 Subject: [PATCH 3/5] Skip redundant recompile in identical-input transform fusion Summary: Skip redundant recompile in identical-input transform fusion. Remove redundant GraphModule.recompile() immediately before super().call() in FuseIdenticalInputTransformsPass. Pass already eliminates dead code and lints before handing to ARM framework for retracing. Recompiling Python GraphModule at that point is unnecessary because ExportPass call reconstructs/interprets graph rather than relying on just-compiled code object. Behavior-preserving speedup. Differential Revision: D114790128 --- backends/arm/_passes/fuse_identical_input_transforms_pass.py | 1 - 1 file changed, 1 deletion(-) diff --git a/backends/arm/_passes/fuse_identical_input_transforms_pass.py b/backends/arm/_passes/fuse_identical_input_transforms_pass.py index 2f174b3abba..fd488abf836 100644 --- a/backends/arm/_passes/fuse_identical_input_transforms_pass.py +++ b/backends/arm/_passes/fuse_identical_input_transforms_pass.py @@ -180,7 +180,6 @@ def call(self, graph_module: GraphModule) -> PassResult: if modified: graph_module.graph.eliminate_dead_code() graph_module.graph.lint() - graph_module.recompile() graph_module = super().call(graph_module).graph_module return PassResult(graph_module, modified) From 9329c4d9432ab4314d21570b2966e9cb3ef27691 Mon Sep 17 00:00:00 2001 From: Andrew Pullin Date: Mon, 10 Aug 2026 08:37:20 -0700 Subject: [PATCH 4/5] Drop redundant recompiles before ARM pass retracing Summary: Drop redundant recompile() before ARM pass retracing. Remove redundant GraphModule.recompile() calls immediately before super().call(graph_module) in CanonicalizeViewCopyPermutePass and FuseDuplicateUsersPass. The following ExportPass retrace/interpreter does not consume compiled Python code object, so recompile adds wall time without changing graph. Keep graph.lint() in modified paths for invariant checking. Behavior-preserving lowering speedup. Differential Revision: D114790155 --- backends/arm/_passes/canonicalize_view_copy_permute_pass.py | 2 +- backends/arm/_passes/fuse_duplicate_users_pass.py | 1 - 2 files changed, 1 insertion(+), 2 deletions(-) diff --git a/backends/arm/_passes/canonicalize_view_copy_permute_pass.py b/backends/arm/_passes/canonicalize_view_copy_permute_pass.py index 00790079373..e45b05feaae 100644 --- a/backends/arm/_passes/canonicalize_view_copy_permute_pass.py +++ b/backends/arm/_passes/canonicalize_view_copy_permute_pass.py @@ -83,7 +83,7 @@ def call(self, graph_module: GraphModule) -> PassResult: if modified: graph_module.graph.eliminate_dead_code() - graph_module.recompile() + graph_module.graph.lint() graph_module = super().call(graph_module).graph_module return PassResult(graph_module, modified) diff --git a/backends/arm/_passes/fuse_duplicate_users_pass.py b/backends/arm/_passes/fuse_duplicate_users_pass.py index 746dade94bd..596bd12d9e0 100644 --- a/backends/arm/_passes/fuse_duplicate_users_pass.py +++ b/backends/arm/_passes/fuse_duplicate_users_pass.py @@ -96,7 +96,6 @@ def enqueue_producer(node: Node) -> None: enqueue_producer(representative) if modified: - graph_module.recompile() graph_module.graph.lint() graph_module = super().call(graph_module).graph_module From 8b7aa90d5f186acbc8ea767a9a2fc336255d2ece Mon Sep 17 00:00:00 2001 From: Andrew Pullin Date: Mon, 10 Aug 2026 08:37:20 -0700 Subject: [PATCH 5/5] Skip redundant recompile after permute propagation retrace (#21696) Summary: Remove redundant GraphModule.recompile() at end of PropagateViewCopyPermutePass.call(). When pass modifies graph it already calls _retrace(), which eliminates dead code, lints, and runs ARM ExportPass retracing to return rebuilt graph module. Immediately following recompile() recompiles Python code not needed before pass manager continues, adding wall time. Keep _retrace() and drop extra compile step. Differential Revision: D114790177 --- backends/arm/_passes/propagate_view_copy_permute_pass.py | 1 - 1 file changed, 1 deletion(-) diff --git a/backends/arm/_passes/propagate_view_copy_permute_pass.py b/backends/arm/_passes/propagate_view_copy_permute_pass.py index d38dd202fe1..5e625dc8d0c 100644 --- a/backends/arm/_passes/propagate_view_copy_permute_pass.py +++ b/backends/arm/_passes/propagate_view_copy_permute_pass.py @@ -130,7 +130,6 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult: if modified: graph_module = self._retrace(graph_module) - graph_module.recompile() return PassResult(graph_module, modified)