Fold non-matching end permutes instead of discarding the region (#22293)

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:

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
@@ -548,7 +554,17 @@ def visit(  # noqa: C901
                                 subgraph.edges_out.add((users_source, user))
                                 # Mark the view for inclusion so it gets preserved
                                 continue
-                    return False
+                    # 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),
+                        )
+                    )
             elif user.op == "output":
                 return False
             elif self._is_permutation_sink_view(user):
@@ -794,6 +810,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 +823,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

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,102 @@ def test_no_permutes_is_noop(self) -> None:
             count_node(result.graph_module, exir_ops.edge.aten.permute_copy.default), 0
         )
 
+    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.