Fix IndexError when a squeeze view is reached upstream (#21620)

Differential Revision: D115002472

Pull Request resolved: #21620

Author: mcremon-meta

backends/transforms/remove_permutes_around_elementwise_ops.py

@@ -144,6 +144,19 @@ def _adapt_permute_across_view(
         Adapts from input-rank to output-rank space (downstream direction).
         Returns the adjusted permutation, or None if not possible.
         """
+        inp = node.args[0]
+        assert isinstance(inp, torch.fx.Node)
+        inp_val = inp.meta.get("val")
+        if inp_val is None:
+            return None
+        in_shape = inp_val.shape
+        # ``permute`` must live in the view's input-rank space. It does not when
+        # the view is reached by upstream traversal, where the permutation is
+        # expressed at the view's output rank; bail out so the caller drops the
+        # subgraph instead of indexing out of range or silently mis-adapting.
+        if len(permute) != len(in_shape):
+            return None
+
         # Handle explicit unsqueeze_copy(dim)
         if node.target in self._UNSQUEEZE_OPS:
             dim = cast(int, node.args[1])
@@ -176,9 +189,6 @@ def _adapt_permute_across_view(
             return new_perm
 
         # Handle view_copy (squeeze/unsqueeze-like reshape)
-        inp = node.args[0]
-        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:
@@ -294,6 +304,18 @@ def visit(  # noqa: C901
             return True
         if node in processed_nodes or not self.is_node_permutable(node):
             return False
+        # A permutable op can still change rank via broadcasting (e.g.
+        # [8, 1] + [1, 1, 1] -> [1, 8, 1]), which would leave the node carrying
+        # a permutation of the wrong rank. Downstream rewrites index the
+        # permutation by dim (update_cat / update_mean_dim / update_slice_copy),
+        # so bail out rather than mis-permute or index out of range.
+        # Squeeze/unsqueeze views are exempt: they intentionally carry their
+        # input-rank permutation and are rank-checked in
+        # _adapt_permute_across_view.
+        if not self._is_squeeze_unsqueeze_view(node):
+            node_shape = getattr(node.meta.get("val"), "shape", None)
+            if node_shape is not None and len(node_shape) != len(current_start_permute):
+                return False
         subgraph.nodes.add(node)
         subgraph.node_end_permute[node] = current_end_permute
         subgraph.node_start_permute[node] = current_start_permute

backends/transforms/test/test_permute_optimization_passes.py

@@ -1034,39 +1034,6 @@ def test_permute_unsqueeze_copy_neg_dim_mul_squeeze_copy_permute(self) -> None:
             "permute_unsqueeze_copy_neg_dim_mul_squeeze_copy_permute",
         )
 
-    def test_unsqueeze_at_moved_position(self) -> None:
-        """The permutation moves the unsqueeze position (P[index] != index), so
-        the adapted permutation must be built from P[index], not index."""
-        builder = GraphBuilder()
-        x_data = torch.randn(1, 8, 16)
-        x = builder.placeholder("x", x_data)
-        p1 = builder.call_operator(
-            op=exir_ops.edge.aten.permute_copy.default, args=(x, [0, 2, 1])
-        )
-        u = builder.call_operator(
-            op=exir_ops.edge.aten.unsqueeze_copy.default, args=(p1, 1)
-        )
-        mul = builder.call_operator(op=exir_ops.edge.aten.mul.Tensor, args=(u, u))
-        sq = builder.call_operator(
-            op=exir_ops.edge.aten.squeeze_copy.dim, args=(mul, 1)
-        )
-        p2 = builder.call_operator(
-            op=exir_ops.edge.aten.permute_copy.default, args=(sq, [0, 2, 1])
-        )
-        builder.output([p2])
-        original = builder.get_graph_module()
-        gm_before = copy.deepcopy(original)
-
-        p = RemovePermutesAroundElementwiseOps()
-        result = cast(PassResult, p(original))
-        self.assertTrue(result.modified)
-        self.assertEqual(
-            count_node(result.graph_module, exir_ops.edge.aten.permute_copy.default), 0
-        )
-        validate_numerics(
-            gm_before, result.graph_module, [x_data], "UnsqueezeAtMovedPosition"
-        )
-
     def test_upstream_view_rank_mismatch_no_crash(self) -> None:
         """Regression test for IndexError when a squeeze/unsqueeze view_copy
         is reached via upstream traversal with a permutation whose rank does
@@ -1170,6 +1137,124 @@ def test_permutation_sink_view_splitting_the_non_unit_dim(self) -> None:
             "permutation_sink_view_splitting_the_non_unit_dim",
         )
 
+    def test_upstream_squeeze_view_rank_mismatch_no_crash(self) -> None:
+        """Regression test for IndexError when a squeeze view_copy is reached
+        via upstream traversal with a permutation at the view's output rank.
+
+        Graph:
+            y [8, 4, 1]                    x [4, 8]
+                    |                            |
+            view_copy (squeeze 3D→2D)      permute [1, 0]
+            [8, 4]                         [8, 4]
+                    \\                          /
+                     ---- add (2D) -----------
+                            |
+                      permute [1, 0]
+                            |
+                         output
+
+        `visit` reaches the view_copy from `add` with the 2D permutation
+        [1, 0], but the view's input is 3D, so the squeezed position (2) is
+        out of range for the permutation. Before the fix,
+        _adapt_permute_across_view crashed with IndexError: list index out
+        of range."""
+        builder = GraphBuilder()
+        x_data = torch.randn(4, 8)
+        y_data = torch.randn(8, 4, 1)
+        x = builder.placeholder("x", x_data)
+        y = builder.placeholder("y", y_data)
+        # Squeeze via view_copy: [8, 4, 1] → [8, 4]
+        view_sq = builder.call_operator(
+            op=exir_ops.edge.aten.view_copy.default, args=(y, [8, 4])
+        )
+        # Start permute: [4, 8] → [8, 4]
+        p1 = builder.call_operator(
+            op=exir_ops.edge.aten.permute_copy.default, args=(x, [1, 0])
+        )
+        add = builder.call_operator(
+            op=exir_ops.edge.aten.add.Tensor, args=(p1, view_sq)
+        )
+        # End permute: [8, 4] → [4, 8]
+        p2 = builder.call_operator(
+            op=exir_ops.edge.aten.permute_copy.default, args=(add, [1, 0])
+        )
+        builder.output([p2])
+        original = builder.get_graph_module()
+        gm_before = copy.deepcopy(original)
+
+        # Should not crash, and should skip the subgraph due to rank mismatch
+        p = RemovePermutesAroundElementwiseOps()
+        result = cast(PassResult, p(original))
+        self.assertFalse(result.modified)
+        self.assertEqual(
+            count_node(result.graph_module, exir_ops.edge.aten.permute_copy.default), 2
+        )
+        validate_numerics(
+            gm_before,
+            result.graph_module,
+            [x_data, y_data],
+            "upstream_squeeze_view_rank_mismatch_no_crash",
+        )
+
+    def test_broadcast_rank_increase_no_crash(self) -> None:
+        """Regression test for IndexError when broadcasting raises a node's
+        rank above the rank of the permutation it is visited with.
+
+        Graph:
+            x [1, 8]                full([1, 1, 1])
+                |                        |
+            permute [1, 0]               |
+            [8, 1]                       |
+                \\                       /
+                 ------ add ------------
+                    [1, 8, 1]  (rank 3)
+                        |
+                slice_copy(dim=2)
+                        |
+                view_copy -> [8]   (permutation sink)
+
+        `add` is visited with the rank-2 permutation [1, 0] but broadcasts to
+        rank 3. The numel-1 `full` input and the sink `view_copy` both let
+        traversal terminate without ever meeting a rank-matched permute, so the
+        subgraph closed with a rank-2 permutation on rank-3 nodes. Before the
+        fix, update_slice_copy did `start_permute[2]` and raised
+        IndexError: list index out of range."""
+        builder = GraphBuilder()
+        x_data = torch.randn(1, 8)
+        x = builder.placeholder("x", x_data)
+        p1 = builder.call_operator(
+            op=exir_ops.edge.aten.permute_copy.default, args=(x, [1, 0])
+        )
+        ones = builder.call_operator(
+            op=exir_ops.edge.aten.full.default, args=([1, 1, 1], 1.0)
+        )
+        # Broadcast [8, 1] + [1, 1, 1] -> [1, 8, 1]: rank 3 under a rank 2 permute
+        add = builder.call_operator(op=exir_ops.edge.aten.add.Tensor, args=(p1, ones))
+        sl = builder.call_operator(
+            op=exir_ops.edge.aten.slice_copy.Tensor, args=(add, 2, 0, 1)
+        )
+        # Sink view: input [1, 8, 1] has a single non-unit dim
+        sink = builder.call_operator(
+            op=exir_ops.edge.aten.view_copy.default, args=(sl, [8])
+        )
+        builder.output([sink])
+        original = builder.get_graph_module()
+        gm_before = copy.deepcopy(original)
+
+        # Should not crash, and should skip the subgraph due to rank mismatch
+        p = RemovePermutesAroundElementwiseOps()
+        result = cast(PassResult, p(original))
+        self.assertFalse(result.modified)
+        self.assertEqual(
+            count_node(result.graph_module, exir_ops.edge.aten.permute_copy.default), 1
+        )
+        validate_numerics(
+            gm_before,
+            result.graph_module,
+            [x_data],
+            "broadcast_rank_increase_no_crash",
+        )
+
 
 # ──────────────────────────────────────────────────────────────────────
 # Tests for RemovePermutesAroundElementwiseOps