Handle rank-changing views in FuseCascadedTransposeOrPermuteOps (#19539)

Summary:

Extend FuseCascadedTransposeOrPermuteOps to fuse permute→view_copy→permute patterns where a squeeze/unsqueeze view sits between two permutes. The pass computes the combined effect of permute+view+permute: if the permutation components cancel out (identity), the entire chain is replaced with a single view_copy.

This handles patterns like permute_3D([0,2,1]) → view(unsqueeze) → permute_4D([0,2,3,1]) which composes to a simple view_copy (the permutations cancel, leaving only the reshape).

Differential Revision: D104775245

Author: mcremon-meta

backends/transforms/fuse_cascaded_transpose_or_permute_ops.py

@@ -20,28 +20,40 @@
 class FuseCascadedTransposeOrPermuteOps(RemoveOrReplacePassInterface):
     """
     Fuse a chain of transpose and permute ops into a single permute or a no-op.
-    Handles branches and chains permutes.
+    Handles branches and chains of permutes, including permute-view-permute
+    patterns where a squeeze/unsqueeze view sits between two permutes.
     """
 
     transpose_or_permute_target = {
         exir_ops.edge.aten.transpose_copy.int,
         exir_ops.edge.aten.permute_copy.default,
     }
 
+    _VIEW_OPS = {
+        exir_ops.edge.aten.view_copy.default,
+        exir_ops.edge.aten.view.default,
+    }
+
     @property
     def targets(self) -> list[EdgeOpOverload]:
         return list(self.transpose_or_permute_target)
 
     def maybe_remove_or_replace(self, node: Node) -> bool:
-        # Fuse with the parent node if it's also a permute or a transpose. Since the
-        # pass interface traverses all ops in order the pass will properly fuse a chain
-        # of permutes.
         parent_node = get_arg(node, "input", Node)
-        if parent_node.target not in self.transpose_or_permute_target:
-            return False
-        input_of_parent = get_arg(parent_node, "input", Node)
 
-        # Compute combined effect of permutes.
+        # Case 1: Direct permute/transpose → permute/transpose
+        if parent_node.target in self.transpose_or_permute_target:
+            return self._fuse_direct(node, parent_node)
+
+        # Case 2: permute → view_copy(squeeze/unsqueeze) → permute
+        if parent_node.target in self._VIEW_OPS:
+            return self._fuse_across_view(node, parent_node)
+
+        return False
+
+    def _fuse_direct(self, node: Node, parent_node: Node) -> bool:
+        """Fuse two adjacent permute/transpose ops."""
+        input_of_parent = get_arg(parent_node, "input", Node)
         dims = list(range(node.meta["val"].ndim))
 
         if parent_node.target == exir_ops.edge.aten.transpose_copy.int:
@@ -54,7 +66,6 @@ def maybe_remove_or_replace(self, node: Node) -> bool:
         else:
             dims = get_permuted_dims(node, dims)
 
-        # If combined effect is identity replace the node with input.
         if dims == sorted(dims):
             node.replace_all_uses_with(input_of_parent)
         else:
@@ -67,3 +78,97 @@ def maybe_remove_or_replace(self, node: Node) -> bool:
             node.replace_all_uses_with(new_permute)
 
         return True
+
+    def _fuse_across_view(self, node: Node, view_node: Node) -> bool:
+        """Fuse permute -> view(squeeze/unsqueeze) -> permute into a view_copy."""
+        # view_node must have exactly one user (this permute node)
+        if len(list(view_node.users.keys())) != 1:
+            return False
+        # view_node's parent must be a permute/transpose
+        view_input = get_arg(view_node, "input", Node)
+        if view_input.target not in self.transpose_or_permute_target:
+            return False
+        # The view must be a squeeze or unsqueeze (rank differs by 1)
+        if view_node.meta.get("val") is None or view_input.meta.get("val") is None:
+            return False
+        view_in_shape = list(view_input.meta["val"].shape)
+        view_out_shape = list(view_node.meta["val"].shape)
+        if abs(len(view_in_shape) - len(view_out_shape)) != 1:
+            return False
+
+        # Get the input before the first permute
+        input_of_first_permute = get_arg(view_input, "input", Node)
+        if input_of_first_permute.meta.get("val") is None:
+            return False
+
+        # Compute the combined effect on the original input dimensions
+        # Start with identity dims for the original input
+        original_ndim = input_of_first_permute.meta["val"].ndim
+        dims = list(range(original_ndim))
+
+        # Apply first permute
+        if view_input.target == exir_ops.edge.aten.transpose_copy.int:
+            dims = get_transposed_dims(view_input, dims)
+        else:
+            dims = get_permuted_dims(view_input, dims)
+
+        # Apply the view (squeeze/unsqueeze)
+        if len(view_out_shape) == len(view_in_shape) + 1:
+            # unsqueeze: insert a new dim
+            index = self._find_extra_one(view_out_shape, view_in_shape)
+            if index == -1:
+                return False
+            dims = [x + 1 if x >= index else x for x in dims]
+            dims.insert(index, -1)  # -1 marks the inserted dim
+        elif len(view_in_shape) == len(view_out_shape) + 1:
+            # squeeze: remove a dim
+            index = self._find_extra_one(view_in_shape, view_out_shape)
+            if index == -1:
+                return False
+            if dims[index] != -1:
+                # Safe: permutation preserves dimension sizes, so a size-1
+                # intermediate dim necessarily originated from a size-1 input dim.
+                pass
+            del dims[index]
+
+        # Apply second permute (node)
+        if node.target == exir_ops.edge.aten.transpose_copy.int:
+            node_dims = list(range(len(dims)))
+            node_dims = get_transposed_dims(node, node_dims)
+            dims = [dims[d] for d in node_dims]
+        else:
+            perm = list(node.args[1])
+            dims = [dims[d] for d in perm]
+
+        # Check if the combined effect (ignoring -1 inserted dims) is identity
+        real_dims = [d for d in dims if d != -1]
+        output_shape = list(node.meta["val"].shape)
+
+        if real_dims == sorted(real_dims):
+            # Combined permutations are identity — replace with view_copy
+            # (the only remaining effect is the squeeze/unsqueeze reshape)
+            if output_shape == list(input_of_first_permute.meta["val"].shape):
+                # Total no-op: replace with input
+                node.replace_all_uses_with(input_of_first_permute)
+            else:
+                with node.graph.inserting_before(node):
+                    new_view = node.graph.call_function(
+                        exir_ops.edge.aten.view_copy.default,
+                        args=(input_of_first_permute, output_shape),
+                    )
+                    new_view.meta = node.meta
+                node.replace_all_uses_with(new_view)
+            return True
+
+        return False
+
+    @staticmethod
+    def _find_extra_one(longer, shorter):
+        if len(longer) != len(shorter) + 1:
+            return -1
+        for i in range(len(shorter)):
+            if longer[i] != shorter[i]:
+                if longer[i] == 1 and shorter[i:] == longer[i + 1:]:
+                    return i
+                return -1
+        return len(shorter) if longer[-1] == 1 else -1

backends/transforms/test/test_permute_optimization_passes.py

@@ -19,12 +19,12 @@
 from executorch.backends.transforms.postpone_permute_below_squeeze_view import (
     PostponePermuteOpBelowSqueezeOrUnsqueezeLikeView,
 )
-from executorch.backends.transforms.replace_nop_transpose_or_permute_with_view import (
-    ReplaceNopTransposeOrPermuteWithViewPass,
-)
 from executorch.backends.transforms.remove_permutes_around_elementwise_ops import (
     RemovePermutesAroundElementwiseOps,
 )
+from executorch.backends.transforms.replace_nop_transpose_or_permute_with_view import (
+    ReplaceNopTransposeOrPermuteWithViewPass,
+)
 from executorch.exir.dialects._ops import ops as exir_ops
 from executorch.exir.pass_base import PassResult
 from torch.utils import _pytree as pytree
@@ -122,15 +122,12 @@ def test_cascaded_permutes_multiple_users(self) -> None:
         permute1 = builder.call_operator(
             op=exir_ops.edge.aten.permute_copy.default, args=(x, [0, 2, 3, 1])
         )
-        # permute2 reverses permute1 => identity
         permute2 = builder.call_operator(
             op=exir_ops.edge.aten.permute_copy.default, args=(permute1, [0, 3, 1, 2])
         )
-        # permute3: different permutation
         permute3 = builder.call_operator(
             op=exir_ops.edge.aten.permute_copy.default, args=(permute1, [0, 2, 1, 3])
         )
-        # permute4 -> permute5: chained
         permute4 = builder.call_operator(
             op=exir_ops.edge.aten.permute_copy.default, args=(permute1, [3, 2, 0, 1])
         )
@@ -151,6 +148,38 @@ def test_cascaded_permutes_multiple_users(self) -> None:
             "FuseCascadedTransposeOrPermuteOps",
         )
 
+    def test_permute_view_permute_fuse(self) -> None:
+        """permute_3D([0,2,1]) → view(unsqueeze) → permute_4D([0,2,3,1]) should
+        be replaced with a single view_copy (permutations cancel out)."""
+        builder = GraphBuilder()
+        x_data = torch.randn(1, 40, 18)
+        x = builder.placeholder("x", x_data)
+        p1 = builder.call_operator(
+            op=exir_ops.edge.aten.permute_copy.default, args=(x, [0, 2, 1])
+        )
+        v = builder.call_operator(
+            op=exir_ops.edge.aten.view_copy.default, args=(p1, [1, 18, 1, 40])
+        )
+        p2 = builder.call_operator(
+            op=exir_ops.edge.aten.permute_copy.default, args=(v, [0, 2, 3, 1])
+        )
+        builder.output([p2])
+        original = builder.get_graph_module()
+        gm_before = copy.deepcopy(original)
+
+        p = FuseCascadedTransposeOrPermuteOps()
+        result = cast(PassResult, p(original))
+        self.assertTrue(result.modified)
+        gm = result.graph_module
+
+        self.assertEqual(count_node(gm, exir_ops.edge.aten.permute_copy.default), 0)
+        self.assertGreaterEqual(
+            count_node(gm, exir_ops.edge.aten.view_copy.default), 1
+        )
+        validate_numerics(
+            gm_before, gm, [x_data], "FuseCascadedAcrossView",
+        )
+
 
 # ──────────────────────────────────────────────────────────────────────
 # Tests for FuseCascadedViewOps
@@ -250,7 +279,6 @@ def test_permute3_view4_chains(self) -> None:
 
         self.assertEqual(count_node(gm, exir_ops.edge.aten.view_copy.default), 2)
         self.assertEqual(count_node(gm, exir_ops.edge.aten.permute_copy.default), 2)
-        # Verify order: views before permutes
         targets = get_compute_nodes(gm)
         view_indices = [
             i
@@ -350,7 +378,6 @@ def test_negative_not_squeeze_like(self) -> None:
             count_node(result.graph_module, exir_ops.edge.aten.permute_copy.default),
             2,
         )
-        # Order unchanged: view, permute, view, permute
         targets = get_compute_nodes(result.graph_module)
         self.assertEqual(targets[0], exir_ops.edge.aten.view_copy.default)
         self.assertEqual(targets[1], exir_ops.edge.aten.permute_copy.default)
@@ -383,6 +410,67 @@ def test_replace_nop_transpose_with_view_float(self) -> None:
             gm_before, gm_after, [x], "ReplaceNopTransposeOrPermuteWithViewPass"
         )
 
+    def test_replace_nop_transpose_with_view_int(self) -> None:
+        x = torch.randint(low=0, high=100, size=(2, 1, 5), dtype=torch.int64)
+        gm = single_op_builder(
+            placeholders=(x,),
+            op=exir_ops.edge.aten.transpose_copy.int,
+            args=(x, 1, 0),
+        )
+        gm_before = copy.deepcopy(gm)
+
+        p = ReplaceNopTransposeOrPermuteWithViewPass()
+        result = cast(PassResult, p(gm))
+        self.assertTrue(result.modified)
+        gm_after = result.graph_module
+        self.assertEqual(count_node(gm_after, exir_ops.edge.aten.transpose_copy.int), 0)
+        self.assertEqual(count_node(gm_after, exir_ops.edge.aten.view_copy.default), 1)
+        validate_numerics(
+            gm_before, gm_after, [x], "ReplaceNopTransposeOrPermuteWithViewPass"
+        )
+
+    def test_replace_nop_permute_5d(self) -> None:
+        x = torch.randn(3, 1, 3, 1, 4)
+        gm = single_op_builder(
+            placeholders=(x,),
+            op=exir_ops.edge.aten.permute_copy.default,
+            args=(x, [0, 2, 4, 1, 3]),
+        )
+        gm_before = copy.deepcopy(gm)
+
+        p = ReplaceNopTransposeOrPermuteWithViewPass()
+        result = cast(PassResult, p(gm))
+        self.assertTrue(result.modified)
+        gm_after = result.graph_module
+        self.assertEqual(
+            count_node(gm_after, exir_ops.edge.aten.permute_copy.default), 0
+        )
+        self.assertEqual(count_node(gm_after, exir_ops.edge.aten.view_copy.default), 1)
+        validate_numerics(
+            gm_before, gm_after, [x], "ReplaceNopTransposeOrPermuteWithViewPass"
+        )
+
+    def test_replace_nop_permute_3d(self) -> None:
+        x = torch.randn(1, 3, 4)
+        gm = single_op_builder(
+            placeholders=(x,),
+            op=exir_ops.edge.aten.permute_copy.default,
+            args=(x, [1, 2, 0]),
+        )
+        gm_before = copy.deepcopy(gm)
+
+        p = ReplaceNopTransposeOrPermuteWithViewPass()
+        result = cast(PassResult, p(gm))
+        self.assertTrue(result.modified)
+        gm_after = result.graph_module
+        self.assertEqual(
+            count_node(gm_after, exir_ops.edge.aten.permute_copy.default), 0
+        )
+        self.assertEqual(count_node(gm_after, exir_ops.edge.aten.view_copy.default), 1)
+        validate_numerics(
+            gm_before, gm_after, [x], "ReplaceNopTransposeOrPermuteWithViewPass"
+        )
+
 
 # ──────────────────────────────────────────────────────────────────────
 # Tests for RemovePermutesAroundElementwiseOps cross-view handling
@@ -458,64 +546,3 @@ def test_4d_permute_squeeze_clamp_3d_permute(self) -> None:
             gm_before, result.graph_module, [x_data],
             "4D_permute_squeeze_clamp_3D_permute",
         )
-
-    def test_replace_nop_transpose_with_view_int(self) -> None:
-        x = torch.randint(low=0, high=100, size=(2, 1, 5), dtype=torch.int64)
-        gm = single_op_builder(
-            placeholders=(x,),
-            op=exir_ops.edge.aten.transpose_copy.int,
-            args=(x, 1, 0),
-        )
-        gm_before = copy.deepcopy(gm)
-
-        p = ReplaceNopTransposeOrPermuteWithViewPass()
-        result = cast(PassResult, p(gm))
-        self.assertTrue(result.modified)
-        gm_after = result.graph_module
-        self.assertEqual(count_node(gm_after, exir_ops.edge.aten.transpose_copy.int), 0)
-        self.assertEqual(count_node(gm_after, exir_ops.edge.aten.view_copy.default), 1)
-        validate_numerics(
-            gm_before, gm_after, [x], "ReplaceNopTransposeOrPermuteWithViewPass"
-        )
-
-    def test_replace_nop_permute_5d(self) -> None:
-        x = torch.randn(3, 1, 3, 1, 4)
-        gm = single_op_builder(
-            placeholders=(x,),
-            op=exir_ops.edge.aten.permute_copy.default,
-            args=(x, [0, 2, 4, 1, 3]),
-        )
-        gm_before = copy.deepcopy(gm)
-
-        p = ReplaceNopTransposeOrPermuteWithViewPass()
-        result = cast(PassResult, p(gm))
-        self.assertTrue(result.modified)
-        gm_after = result.graph_module
-        self.assertEqual(
-            count_node(gm_after, exir_ops.edge.aten.permute_copy.default), 0
-        )
-        self.assertEqual(count_node(gm_after, exir_ops.edge.aten.view_copy.default), 1)
-        validate_numerics(
-            gm_before, gm_after, [x], "ReplaceNopTransposeOrPermuteWithViewPass"
-        )
-
-    def test_replace_nop_permute_3d(self) -> None:
-        x = torch.randn(1, 3, 4)
-        gm = single_op_builder(
-            placeholders=(x,),
-            op=exir_ops.edge.aten.permute_copy.default,
-            args=(x, [1, 2, 0]),
-        )
-        gm_before = copy.deepcopy(gm)
-
-        p = ReplaceNopTransposeOrPermuteWithViewPass()
-        result = cast(PassResult, p(gm))
-        self.assertTrue(result.modified)
-        gm_after = result.graph_module
-        self.assertEqual(
-            count_node(gm_after, exir_ops.edge.aten.permute_copy.default), 0
-        )
-        self.assertEqual(count_node(gm_after, exir_ops.edge.aten.view_copy.default), 1)
-        validate_numerics(
-            gm_before, gm_after, [x], "ReplaceNopTransposeOrPermuteWithViewPass"
-        )