diff --git a/backends/arm/test/passes/test_propagate_permutes_views_pass.py b/backends/arm/test/passes/test_propagate_permutes_views_pass.py index 90702e4bacd..163f367b2df 100644 --- a/backends/arm/test/passes/test_propagate_permutes_views_pass.py +++ b/backends/arm/test/passes/test_propagate_permutes_views_pass.py @@ -1168,6 +1168,39 @@ def test_up_pass_fuses_equivalent_output_permutations_before_fan_out() -> None: assert targets.index(PERMUTE) < targets.index(RELU) < targets.index(ADD) +def test_up_pass_fusion_enables_further_propagation() -> None: + graph = torch.fx.Graph() + x = graph.placeholder("x") + x.meta["val"] = torch.empty((1, 2, 3, 4)) + relu = graph.call_function(RELU, args=(x,)) + relu.meta["val"] = torch.empty((1, 2, 3, 4)) + neg = graph.call_function(NEG, args=(relu,)) + neg.meta["val"] = torch.empty((1, 2, 3, 4)) + first_permute = graph.call_function(PERMUTE, args=(neg, [0, 2, 3, 1])) + first_permute.meta["val"] = torch.empty((1, 3, 4, 2)) + second_permute = graph.call_function(PERMUTE, args=(relu, [0, 2, 3, 1])) + second_permute.meta["val"] = torch.empty((1, 3, 4, 2)) + add = graph.call_function(ADD, args=(first_permute, second_permute)) + add.meta["val"] = torch.empty((1, 3, 4, 2)) + graph.output(add) + + graph_module = _run_pass_on_graph_module(graph, PropagateViewCopyPermuteUpPass) + call_nodes = [ + node for node in graph_module.graph.nodes if node.op == "call_function" + ] + permutes = [node for node in call_nodes if node.target == PERMUTE] + x = next(node for node in graph_module.graph.nodes if node.op == "placeholder") + relu = next(node for node in call_nodes if node.target == RELU) + neg = next(node for node in call_nodes if node.target == NEG) + add = next(node for node in call_nodes if node.target == ADD) + + assert len(permutes) == 1 + assert permutes[0].args[0] is x + assert relu.args[0] is permutes[0] + assert neg.args[0] is relu + assert set(add.all_input_nodes) == {relu, neg} + + def test_propagate_moves_before_dtype_changing_rescale() -> None: graph = torch.fx.Graph() x = graph.placeholder("x") diff --git a/backends/transforms/canonicalize_view_copy_permute_pass.py b/backends/transforms/canonicalize_view_copy_permute_pass.py index c29a246b04a..dcf6355c90e 100644 --- a/backends/transforms/canonicalize_view_copy_permute_pass.py +++ b/backends/transforms/canonicalize_view_copy_permute_pass.py @@ -5,6 +5,7 @@ from __future__ import annotations +from collections import deque from typing import Any, cast, Iterable, Sequence, Set, Type import torch @@ -104,20 +105,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/transforms/fuse_duplicate_users_pass.py b/backends/transforms/fuse_duplicate_users_pass.py index b3989e76c94..5e91e722a60 100644 --- a/backends/transforms/fuse_duplicate_users_pass.py +++ b/backends/transforms/fuse_duplicate_users_pass.py @@ -43,11 +43,20 @@ 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) + erased_nodes: Set[Node] = set() + + def enqueue_producer(node: Node) -> None: + if node in erased_nodes 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: + if producer in erased_nodes: # Node was deleted by a previous rewrite while still queued. continue @@ -56,7 +65,9 @@ def call(self, graph_module: GraphModule) -> PassResult: if len(user_nodes) < 2: continue - candidate_groups = self._get_candidate_groups(node_order, user_nodes) + candidate_groups = self._get_candidate_groups( + node_order, user_nodes, erased_nodes + ) signature_to_user: Dict[Tuple[Hashable, ...], Node] = {} for group in candidate_groups: @@ -77,13 +88,14 @@ def call(self, graph_module: GraphModule) -> PassResult: user.replace_all_uses_with(representative) graph.erase_node(user) + erased_nodes.add(user) modified = True # 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: if self._recompile_before_retrace: @@ -93,10 +105,10 @@ def call(self, graph_module: GraphModule) -> PassResult: return PassResult(graph_module, modified) - def _get_candidate_groups(self, node_order, user_nodes): + def _get_candidate_groups(self, node_order, user_nodes, erased_nodes): users_by_target: Dict[Tuple[str, Hashable], List[Node]] = {} for user in user_nodes: - if user.graph is None: + if user in erased_nodes: # User might already have been removed by a prior rewrite. continue