Fix unsqueeze permutation adaptation in RemovePermutesAroundElementwiseOps (#21551)

Summary:

`_adapt_permute_across_view` and `permute_subgraph` disagreed about where an
unsqueeze's new dim lands in the un-permuted layout.

When the region walker crosses an unsqueeze, it has to (a) adapt the running
permutation to the new rank and (b) later rewrite the unsqueeze's `dim` arg
into un-permuted space. `permute_subgraph` does (b) as `permute[index]`, but
`_adapt_permute_across_view` did (a) by inserting `index` -- both as the shift
threshold and as the inserted value. The two only coincide when
`permute[index] == index`.

For example with `P = [0, 2, 1]` and an unsqueeze at index 1, the rewrite emits
`dim=2` while the adaptation produces `[0, 1, 3, 2]`; the permutation
consistent with `dim=2` is `[0, 2, 3, 1]`. Downstream nodes then get remapped
against a wrong permutation and the end-permute comparison no longer matches.
In practice that shows up as a missed optimisation -- the subgraph is rejected
-- but nothing guarantees the wrong permutation cannot coincidentally match an
end permute, so this is a correctness fix, not just a missed-match fix.

The squeeze branch immediately below already did this correctly
(`squeezed_value = permute[index]`, with a comment spelling out that the
position and the value are different things); only unsqueeze was inconsistent.

Note this pass is shared: the Cadence (opt_level=2) and Arm/TOSA subclasses
inherit the fixed behaviour. Their suites are covered below.

Reviewed By: digantdesai

Differential Revision: D114508242

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