Fix unsqueeze permutation adaptation in RemovePermutesAroundElementwiseOps (#21551)

Differential Revision: D114508242

Pull Request resolved: #21551

Author: mcremon-meta

backends/transforms/remove_permutes_around_elementwise_ops.py

@@ -149,8 +149,15 @@ def _adapt_permute_across_view(
             dim = cast(int, node.args[1])
             rank = len(permute)
             index = dim if dim >= 0 else dim + rank + 1
-            new_perm = [x + 1 if x >= index else x for x in permute]
-            new_perm.insert(index, index)
+            # `index` is a POSITION in the permuted output; the un-permuted
+            # position the new dim lands at is the permutation VALUE there
+            # (or the end, when appending). permute_subgraph rewrites the dim
+            # arg to exactly this value, so the two must agree -- using `index`
+            # here instead silently desynchronises them whenever
+            # permute[index] != index.
+            inserted_value = permute[index] if index < rank else rank
+            new_perm = [x + 1 if x >= inserted_value else x for x in permute]
+            new_perm.insert(index, inserted_value)
             return new_perm
 
         # Handle explicit squeeze_copy(dim)

backends/transforms/test/test_permute_optimization_passes.py

@@ -1034,6 +1034,39 @@ 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