diff --git a/backends/cadence/aot/tests/test_remove_ops_passes.py b/backends/cadence/aot/tests/test_remove_ops_passes.py index afd777719b7..216a80970d9 100644 --- a/backends/cadence/aot/tests/test_remove_ops_passes.py +++ b/backends/cadence/aot/tests/test_remove_ops_passes.py @@ -477,11 +477,21 @@ def test_keep_permutes_around_elemwise_ops_add(self) -> None: ) builder.output([permute]) original = builder.get_graph_module() + gm_before = copy.deepcopy(original) p = RemovePermutesAroundElementwiseOps() graph_after_passes = cast(PassResult, p(original)).graph_module - # Ensure no permutes were removed, since the dimensions don't fit the expected pattern + # The end permute is not the inverse of the start one, so it cannot be + # dropped. It absorbs the start permute instead, leaving one behind. self.assertEqual( - count_node(graph_after_passes, exir_ops.edge.aten.permute_copy.default), 2 + count_node(graph_after_passes, exir_ops.edge.aten.permute_copy.default), 1 + ) + + sample_inputs = [torch.randn(1, 8, 4, 4, dtype=torch.float32)] + validate( + gm_before, + graph_after_passes, + sample_inputs, + "RemovePermutesAroundElementwiseOps", ) def test_remove_permutes_around_elemwise_ops_add_mean(self) -> None: diff --git a/backends/transforms/remove_permutes_around_elementwise_ops.py b/backends/transforms/remove_permutes_around_elementwise_ops.py index 77e9f94c3d6..1a3a6ada995 100644 --- a/backends/transforms/remove_permutes_around_elementwise_ops.py +++ b/backends/transforms/remove_permutes_around_elementwise_ops.py @@ -41,6 +41,12 @@ class Subgraph: edges_in: set[tuple[torch.fx.Node, torch.fx.Node]] = field(default_factory=set) # Outgoing edges of the subgraph to permute nodes. edges_out: set[tuple[torch.fx.Node, torch.fx.Node]] = field(default_factory=set) + # Outgoing edges to permutes that do not match end_permute. Those are + # kept and their permutation rewritten to absorb the removed start + # permute, as (producer, permute node, new permutation). + edges_out_to_update: set[ + tuple[torch.fx.Node, torch.fx.Node, tuple[int, ...]] + ] = field(default_factory=set) # Incoming edges from constant nodes that need a compensating permute. constant_edges_in: set[tuple[torch.fx.Node, torch.fx.Node]] = field( default_factory=set @@ -372,7 +378,6 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult: # noqa: C901 # Expected end permutation for the subgraph. end_permute = [start_permute.index(i) for i in range(len(start_permute))] - # Try direct users first (same-rank matching) for user in node.users: if ( not self.is_node_permutable(user) @@ -385,50 +390,6 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult: # noqa: C901 for n in subgraph.nodes: processed_nodes.add(n) - # Also try: permute → view(squeeze/unsqueeze) → chain → ... - # If the permute's sole user is a squeeze/unsqueeze view, - # adapt the permutation across the view and search for a - # matching end permute at the new rank. - users = list(node.users.keys()) - if ( - len(users) == 1 - and self._is_squeeze_unsqueeze_view(users[0]) - and node not in processed_nodes - ): - view_node = users[0] - adapted_start = self._adapt_permute_across_view( - start_permute, view_node - ) - if adapted_start is not None: - adapted_end = [ - adapted_start.index(i) for i in range(len(adapted_start)) - ] - for view_user in view_node.users: - if ( - not self.is_node_permutable(view_user) - and self._interleave_triple(view_user) is None - ): - continue - subgraph = self.Subgraph(adapted_start, adapted_end) - # Include the view in the subgraph - subgraph.nodes.add(view_node) - subgraph.node_end_permute[view_node] = adapted_end - # Use the ORIGINAL start_permute for the view node - # so update_view_copy can remap its shape correctly - subgraph.node_start_permute[view_node] = start_permute - # The start permute feeds into the view - subgraph.edges_in.add((node, view_node)) - if self.visit( - view_user, - subgraph, - processed_nodes, - adapted_end, - adapted_start, - ): - subgraphs_found.append(subgraph) - for n in subgraph.nodes: - processed_nodes.add(n) - modified = False for subgraph in subgraphs_found: if self.permute_subgraph(subgraph): @@ -526,29 +487,17 @@ def visit( # noqa: C901 if user_perm == downstream_end: subgraph.edges_out.add((users_source, user)) else: - # Check if permute → view(squeeze/unsqueeze) forms an - # end boundary at a different rank. - user_users = list(user.users.keys()) - if len(user_users) == 1 and self._is_squeeze_unsqueeze_view( - user_users[0] - ): - view_after: torch.fx.Node = user_users[0] - # Adapt the start permute across the view and derive - # the expected end permute as its inverse. - adapted_start_after = self._adapt_permute_across_view( - downstream_start, view_after + # Non-matching permute: keep it and fold the start permute into it + # rather than discarding the region. + if user_perm is None or len(user_perm) != len(downstream_start): + return False + subgraph.edges_out_to_update.add( + ( + users_source, + user, + tuple(downstream_start[d] for d in user_perm), ) - if adapted_start_after is not None: - adapted = [ - adapted_start_after.index(i) - for i in range(len(adapted_start_after)) - ] - if user_perm == adapted: - # Include both the permute and the view as end edges - subgraph.edges_out.add((users_source, user)) - # Mark the view for inclusion so it gets preserved - continue - return False + ) elif user.op == "output": return False elif self._is_permutation_sink_view(user): @@ -682,11 +631,27 @@ def _is_constant_pad(self, node: torch.fx.Node) -> bool: return True + def _removes_a_permute(self, subgraph: Subgraph) -> bool: + """Whether rewriting this region reduces the number of permutes.""" + if subgraph.edges_out: + return True + rewired = set(subgraph.edges_in) + for permute, _ in subgraph.edges_in: + if all((permute, user) in rewired for user in permute.users): + return True + return False + def permute_subgraph(self, subgraph: Subgraph) -> bool: # noqa: C901 # Ensure that the subgraph's edges have not been modified by an earlier rewrite before applying changes. if not self._subgraph_edges_are_current(subgraph): return False + # Folding an end permute only pays for itself if some permute goes away. + # Otherwise the region is rewritten for nothing, and the composed + # permutation is a worse fusion candidate for the passes downstream. + if subgraph.edges_out_to_update and not self._removes_a_permute(subgraph): + return False + # Nodes belonging to a repeat_interleave triple are rewritten as a unit # below, so they must skip the per-node dim handling and the view rank # check (the triple's interior ranks intentionally differ from the @@ -794,6 +759,11 @@ def permute_subgraph(self, subgraph: Subgraph) -> bool: # noqa: C901 assert out.target in PERMUTE_COPY_TARGETS out.replace_all_uses_with(inp) + # Update outgoing permutes that can't be eliminated. + for _, out, new_permutation in subgraph.edges_out_to_update: + assert out.target in PERMUTE_COPY_TARGETS + set_arg(out, "dims", list(new_permutation)) + return True def _subgraph_edges_are_current(self, subgraph: Subgraph) -> bool: @@ -802,10 +772,20 @@ def _subgraph_edges_are_current(self, subgraph: Subgraph) -> bool: if inp.target not in PERMUTE_COPY_TARGETS or inp not in out.all_input_nodes: return False + # edges_out_to_update can rewrite a permute in place, leaving it wired. + if self.get_permutation(inp) != subgraph.node_start_permute.get( + out, subgraph.start_permute + ): + return False + for inp, out in subgraph.edges_out: if out.target not in PERMUTE_COPY_TARGETS or out not in inp.users: return False + for inp, out, _ in subgraph.edges_out_to_update: + if out.target not in PERMUTE_COPY_TARGETS or out not in inp.users: + return False + for const_node, user_node in subgraph.constant_edges_in: if const_node not in user_node.all_input_nodes: return False diff --git a/backends/transforms/test/test_permute_optimization_passes.py b/backends/transforms/test/test_permute_optimization_passes.py index 9aceab743ef..1f357deb171 100644 --- a/backends/transforms/test/test_permute_optimization_passes.py +++ b/backends/transforms/test/test_permute_optimization_passes.py @@ -1308,12 +1308,18 @@ def test_4d_permute_squeeze_view_slice_mul_3d_permute(self) -> None: gm_before = copy.deepcopy(original) p = RemovePermutesAroundElementwiseOps() - # With the fix, the adapted permutation becomes identity [0,1,2], - # so no matching end permute is found and the graph is unchanged. - # Before the fix, the wrong adapted permutation [1,0,2] would match - # the end permute and create an invalid subgraph, causing a crash. + # The adapted permutation is the identity [0,1,2], so the end permute + # [1,0,2] does not match it. Rather than discarding the region, the end + # permute is kept with its permutation composed against the adapted + # start -- the identity here, so it stays [1,0,2] -- while the leading + # permute is removed. Before the fix, the wrong adapted permutation + # [1,0,2] would match the end permute and create an invalid subgraph. result = cast(PassResult, p(original)) - self.assertFalse(result.modified) + self.assertTrue(result.modified) + self.assertEqual( + count_node(result.graph_module, exir_ops.edge.aten.permute_copy.default), + 1, + ) validate_numerics( gm_before, result.graph_module, @@ -2167,6 +2173,187 @@ def test_no_permutes_is_noop(self) -> None: count_node(result.graph_module, exir_ops.edge.aten.permute_copy.default), 0 ) + def test_region_ending_in_permute_then_view(self) -> None: + """permute -> mul -> permute -> view(squeeze). + + The end permute is expressed at the region's rank while the view that + follows changes rank, so no boundary spanning both can ever match. The + end permute is folded on its own instead, absorbing the start permute as + [1,0,2].""" + builder = GraphBuilder() + x_data = torch.randn(1, 3, 4) + x = builder.placeholder("x", x_data) + p1 = builder.call_operator( + op=exir_ops.edge.aten.permute_copy.default, args=(x, [1, 2, 0]) + ) + mul = builder.call_operator(op=exir_ops.edge.aten.mul.Tensor, args=(p1, p1)) + # Not the inverse of [1,2,0] (that would be [2,0,1]). + p2 = builder.call_operator( + op=exir_ops.edge.aten.permute_copy.default, args=(mul, [0, 2, 1]) + ) + view = builder.call_operator( + op=exir_ops.edge.aten.view_copy.default, args=(p2, [3, 4]) + ) + builder.output([view]) + original = builder.get_graph_module() + gm_before = copy.deepcopy(original) + + p = RemovePermutesAroundElementwiseOps() + result = cast(PassResult, p(original)) + self.assertTrue(result.modified) + + permutes = result.graph_module.graph.find_nodes( + op="call_function", target=exir_ops.edge.aten.permute_copy.default + ) + self.assertEqual(len(permutes), 1) + self.assertEqual(list(permutes[0].args[1]), [1, 0, 2]) + + validate_numerics( + gm_before, + result.graph_module, + [x_data], + "region_ending_in_permute_then_view", + ) + + def test_view_seeded_region_with_shared_view(self) -> None: + """permute -> squeeze view -> two slices, only one of which is visited. + + The view-seeded path requires the permute to have a single user, but not + the view. It adds the view to the region and rewrites its shape arg, + while traversing only one of the view's users, so a sibling consumer is + left indexing the pre-rewrite layout: here slice_b keeps slicing dim 0 + after the view's shape arg becomes [6, 4], yielding [2, 4] instead of + [2, 6].""" + builder = GraphBuilder() + x_data = torch.randn(1, 6, 4) + x = builder.placeholder("x", x_data) + # Sole user of the permute is a squeeze view -> takes the view-seeded path. + p1 = builder.call_operator( + op=exir_ops.edge.aten.permute_copy.default, args=(x, [0, 2, 1]) + ) + view = builder.call_operator( + op=exir_ops.edge.aten.view_copy.default, args=(p1, [4, 6]) + ) + # Two consumers of the same view; each is independently permutable. + slice_a = builder.call_operator( + op=exir_ops.edge.aten.slice_copy.Tensor, args=(view, 0, 0, 2) + ) + slice_b = builder.call_operator( + op=exir_ops.edge.aten.slice_copy.Tensor, args=(view, 0, 2, 4) + ) + # A matching end permute, so slice_a's region is applied. + pa = builder.call_operator( + op=exir_ops.edge.aten.permute_copy.default, args=(slice_a, [1, 0]) + ) + builder.output([pa, slice_b]) + original = builder.get_graph_module() + gm_before = copy.deepcopy(original) + + p = RemovePermutesAroundElementwiseOps() + result = cast(PassResult, p(original)) + validate_numerics( + gm_before, + result.graph_module, + [x_data], + "view_seeded_region_with_shared_view", + ) + + def test_non_matching_end_permute_is_composed(self) -> None: + """permute([1,2,0]) → mul → permute([1,0,2]), where the end permute is + not the inverse of the start one. + + The end permute cannot be dropped, but it can absorb the start permute + instead of the region being discarded: with the start permute gone the + producer is back in the original layout, so composing the two as + R[k] = start[end[k]] = [2,1,0] leaves the consumer's tensor unchanged + and still removes one permute.""" + builder = GraphBuilder() + # Distinct dim sizes so a wrong permutation cannot go unnoticed. + x_data = torch.randn(2, 3, 4) + x = builder.placeholder("x", x_data) + p1 = builder.call_operator( + op=exir_ops.edge.aten.permute_copy.default, args=(x, [1, 2, 0]) + ) + mul = builder.call_operator(op=exir_ops.edge.aten.mul.Tensor, args=(p1, p1)) + # [1,0,2] is not the inverse of [1,2,0] (that would be [2,0,1]). + p2 = builder.call_operator( + op=exir_ops.edge.aten.permute_copy.default, args=(mul, [1, 0, 2]) + ) + builder.output([p2]) + original = builder.get_graph_module() + gm_before = copy.deepcopy(original) + + p = RemovePermutesAroundElementwiseOps() + result = cast(PassResult, p(original)) + self.assertTrue(result.modified) + + permutes = result.graph_module.graph.find_nodes( + op="call_function", target=exir_ops.edge.aten.permute_copy.default + ) + self.assertEqual(len(permutes), 1) + self.assertEqual(list(permutes[0].args[1]), [2, 1, 0]) + + validate_numerics( + gm_before, + result.graph_module, + [x_data], + "non_matching_end_permute_is_composed", + ) + + def test_chained_regions_absorb_into_last_permute(self) -> None: + """permute(A) -> mul -> permute(B) -> add -> permute(C), no pair matching. + + Each region folds its start permute into its end permute, so the chain + collapses into the trailing permute carrying A∘B∘C = [2,0,1]. One region + folds per invocation: rewriting B in place invalidates the candidate the + second region was planned against, so it is skipped until the next run + rediscovers it against the rewritten graph.""" + builder = GraphBuilder() + # Distinct dim sizes so a wrong permutation cannot go unnoticed. + x_data = torch.randn(2, 3, 4) + x = builder.placeholder("x", x_data) + # inverse([1,2,0]) is [2,0,1], so B does not match A's region. + pa = builder.call_operator( + op=exir_ops.edge.aten.permute_copy.default, args=(x, [1, 2, 0]) + ) + mul = builder.call_operator(op=exir_ops.edge.aten.mul.Tensor, args=(pa, pa)) + pb = builder.call_operator( + op=exir_ops.edge.aten.permute_copy.default, args=(mul, [1, 0, 2]) + ) + add = builder.call_operator(op=exir_ops.edge.aten.add.Tensor, args=(pb, pb)) + pc = builder.call_operator( + op=exir_ops.edge.aten.permute_copy.default, args=(add, [0, 2, 1]) + ) + builder.output([pc]) + original = builder.get_graph_module() + gm_before = copy.deepcopy(original) + + # One region folds per invocation, so the three-permute chain needs a + # few rounds; the bound just keeps a non-converging pass from hanging. + graph_module = original + for _ in range(5): + result = cast( + PassResult, RemovePermutesAroundElementwiseOps()(graph_module) + ) + graph_module = result.graph_module + if not result.modified: + break + else: + self.fail("pass did not reach a fixed point") + + permutes = graph_module.graph.find_nodes( + op="call_function", target=exir_ops.edge.aten.permute_copy.default + ) + self.assertEqual(len(permutes), 1) + self.assertEqual(list(permutes[0].args[1]), [2, 0, 1]) + + validate_numerics( + gm_before, + graph_module, + [x_data], + "chained_regions_absorb_into_last_permute", + ) + class LayoutPermuteVisibilityTest(unittest.TestCase): """The data-movement passes must see both permute dialects.