diff --git a/backends/arm/_passes/canonicalize_view_copy_permute_pass.py b/backends/arm/_passes/canonicalize_view_copy_permute_pass.py index 1930ba0b5f9..e45b05feaae 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 @@ -82,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) @@ -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..596bd12d9e0 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,11 +92,10 @@ 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() graph_module.graph.lint() graph_module = super().call(graph_module).graph_module 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) diff --git a/backends/arm/_passes/propagate_view_copy_permute_pass.py b/backends/arm/_passes/propagate_view_copy_permute_pass.py index a54c8f28312..5e625dc8d0c 100644 --- a/backends/arm/_passes/propagate_view_copy_permute_pass.py +++ b/backends/arm/_passes/propagate_view_copy_permute_pass.py @@ -108,24 +108,28 @@ 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) - graph_module.recompile() return PassResult(graph_module, modified) 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.