Fix unsqueeze permutation adaptation in RemovePermutesAroundElementwiseOps (#21551)

[mcremon-meta]

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

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/21551

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❌ 7 Cancelled Jobs, 47 Pending

As of commit 958a9b7 with merge base e0a91fe ():

CANCELLED JOBS - The following jobs were cancelled. Please retry:

• Test RISC-V Backend / test-riscv (add, true, false) / run / linux-job (gh)
  ##[error]The operation was canceled.
• Test RISC-V Backend / test-riscv (add, true, true) / run / linux-job (gh)
  ##[error]The operation was canceled.
• Test RISC-V Backend / test-riscv (llama2, true, false) / run / linux-job (gh)
  ##[error]The operation was canceled.
• Test RISC-V Backend / test-riscv (llama2, true, true) / run / linux-job (gh)
  ##[error]The operation was canceled.
• Test RISC-V Backend / test-riscv (mv2, true, true) / run / linux-job (gh)
  ##[error]The operation was canceled.
• Test RISC-V Backend / test-riscv (resnet18, true, false) / run / linux-job (gh)
  ##[error]The operation was canceled.
• Test RISC-V Backend / test-riscv (yolo26, true, false) / run / linux-job (gh)
  ##[error]The operation was canceled.

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[meta-codesync]

@mcremon-meta has exported this pull request. If you are a Meta employee, you can view the originating Diff in D114508242.

[github-actions]

This PR needs a release notes: label

If your change should be included in the release notes (i.e. would users of this library care about this change?), please use a label starting with release notes:. This helps us keep track and include your important work in the next release notes.

To add a label, you can comment to pytorchbot, for example
@pytorchbot label "release notes: none"

For more information, see
https://github.com/pytorch/pytorch/wiki/PyTorch-AutoLabel-Bot#why-categorize-for-release-notes-and-how-does-it-work.

[digantdesai]

perhaps a naive question, the permute is done on the input x, and not on the scalar constants, how was the pass able to remove it?

[mcremon-meta]

we used to add compensating permutes to the [1, 1, 1, 1] constants, which is now layout invariant and treated as a nop. So those 3 are not added anymore, bringing the total count to 0

[digantdesai]

Thanks.