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).

Reviewed By: DrJessop

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,104 @@ def maybe_remove_or_replace(self, node: Node) -> bool:
             node.replace_all_uses_with(new_permute)
 
         return True
+
+    def _apply_view_to_dims(
+        self, dims: list[int], view_in_shape, view_out_shape
+    ) -> list[int] | None:
+        """Apply a squeeze or unsqueeze view to dimension mapping.
+
+        Returns the updated dims, or None if the view cannot be mapped.
+        """
+        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 None
+            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 None
+            dims = list(dims)
+            del dims[index]
+        return dims
+
+    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(view_node.users) != 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)
+        view_in_shape = view_input.meta["val"].shape
+        view_out_shape = 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)
+
+        # 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)
+        dims = self._apply_view_to_dims(dims, view_in_shape, view_out_shape)
+        if dims is None:
+            return False
+
+        # 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]
+        elif node.target == exir_ops.edge.aten.permute_copy.default:
+            perm = get_arg(node, "dims")
+            dims = [dims[d] for d in perm]
+        else:
+            raise ValueError(f"Unexpected target: {node.target}")
+
+        # Check if the combined effect (ignoring -1 inserted dims) is identity
+        real_dims = [d for d in dims if d != -1]
+
+        if real_dims == sorted(real_dims):
+            # Combined permutations are identity — replace with view_copy
+            # (the only remaining effect is the squeeze/unsqueeze reshape)
+            output_shape = node.meta["val"].shape
+            if output_shape == 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, list(output_shape)),
+                    )
+                    new_view.meta = node.meta
+                node.replace_all_uses_with(new_view)
+            return True
+
+        return False
+
+    @staticmethod
+    def _find_extra_one(longer: list[int], shorter: list[int]) -> int:
+        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/remove_permutes_around_elementwise_ops.py

@@ -12,7 +12,7 @@
 
 import torch
 import torch.fx
-from executorch.backends.transforms.permute_pass_utils import get_arg
+from executorch.backends.transforms.permute_pass_utils import get_arg, set_arg
 from executorch.exir.dialects._ops import ops as exir_ops
 from executorch.exir.pass_base import ExportPass, PassResult
 
@@ -106,11 +106,8 @@ def _check_squeeze_unsqueeze_view(self, node: torch.fx.Node) -> bool:
             return True
         if node.target not in self._VIEW_OPS:
             return False
-        if node.meta.get("val") is None:
-            return False
         inp = node.args[0]
-        if not isinstance(inp, torch.fx.Node) or inp.meta.get("val") is None:
-            return False
+        assert isinstance(inp, torch.fx.Node)
         in_shape = inp.meta["val"].shape
         out_shape = node.meta["val"].shape
         if len(out_shape) == len(in_shape) + 1:
@@ -377,9 +374,9 @@ def _is_constant(self, node: torch.fx.Node) -> bool:
     def _get_node_rank(self, node: torch.fx.Node) -> int | None:
         """Return the tensor rank of a node's output, or None if unknown."""
         val = node.meta.get("val")
-        if val is not None and hasattr(val, "shape"):
-            return len(val.shape)
-        return None
+        if val is None:
+            return None
+        return len(val.shape)
 
     @staticmethod
     def _is_pointwise(target) -> bool:
@@ -432,10 +429,8 @@ def permute_subgraph(self, subgraph: Subgraph) -> None:  # noqa: C901
                     node.update_arg(1, index)
             elif node.target in self._SQUEEZE_OPS:
                 # squeeze dim is in input space (rank)
-                dim = cast(int, node.args[1])
-                rank = len(node_start_perm)
-                index = dim if dim >= 0 else dim + rank
-                node.update_arg(1, node_start_perm[index])
+                dim = get_arg(node, "dim", int)
+                set_arg(node, "dim", node_start_perm[dim])
 
         # Skip incoming permutes.
         for inp, out in subgraph.edges_in:
@@ -486,42 +481,25 @@ def permute_subgraph(self, subgraph: Subgraph) -> None:  # noqa: C901
             out.replace_all_uses_with(inp)
 
     def update_cat(self, node: torch.fx.Node, start_permute: list[int]) -> None:
-        if len(node.args) >= 2:
-            node.update_arg(1, start_permute[cast(int, node.args[1])])
-        elif "dim" in node.kwargs:
-            node.update_kwarg("dim", start_permute[cast(int, node.kwargs["dim"])])
-        else:
-            # Default cat dim is 0.
-            node.update_kwarg("dim", start_permute[0])
+        dim = get_arg(node, "dim", int)
+        set_arg(node, "dim", start_permute[dim])
 
     def update_mean_dim(self, node: torch.fx.Node, start_permute: list[int]) -> None:
-        if len(node.args) >= 2:
-            node.update_arg(
-                1, [start_permute[dim] for dim in cast(list[int], node.args[1])]
-            )
-        else:
-            node.update_kwarg(
-                "dim",
-                [start_permute[dim] for dim in cast(list[int], node.kwargs["dim"])],
-            )
+        dims = get_arg(node, "dim")
+        set_arg(node, "dim", [start_permute[d] for d in cast(list[int], dims)])
 
     def update_slice_copy(self, node: torch.fx.Node, start_permute: list[int]) -> None:
-        if len(node.args) >= 2:
-            node.update_arg(1, start_permute[cast(int, node.args[1])])
-        else:
-            node.update_kwarg("dim", start_permute[cast(int, node.kwargs["dim"])])
+        dim = get_arg(node, "dim", int)
+        set_arg(node, "dim", start_permute[dim])
 
     def update_view_copy(self, node: torch.fx.Node, start_permute: list[int]) -> None:
         """Adjust view_copy shape arg after permute removal.
 
         After removing the start permute, the view's input is in the original
         (un-permuted) layout. Recompute the view's target shape accordingly.
         """
-        if node.meta.get("val") is None:
-            return
         inp = node.args[0]
-        if not isinstance(inp, torch.fx.Node) or inp.meta.get("val") is None:
-            return
+        assert isinstance(inp, torch.fx.Node)
 
         in_shape = inp.meta["val"].shape
         out_shape = node.meta["val"].shape