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(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)
+        if view_node.meta.get("val") is None or view_input.meta.get("val") is None:
+            return False
+        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)
+        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 = get_arg(node, "dims")
+            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]
+
+        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, 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