[pssrawat]

Summary:
RemovePermutesAroundElementwiseOps removes the permutes around a region of elementwise ops, so the region runs in its input's original layout. A constant feeding the region was written for the permuted layout, so it has to be converted with the region's end permutation (node_end_perm). Constants with the same rank as the region get a permute_copy, which is correct. This diff fixes constants with a lower rank, which broadcast into the region.

The bug is in the shared pass, so every backend that uses it is affected: the Cortex-M explicit-layout pipeline, Arm (RemovePermutesAroundElementwiseTosaOps subclasses it without overriding this code) and Cadence.

Before

The constant's shape was padded with leading 1s, reordered by the permutation, trimmed back to the constant's rank, and applied with a single view_copy. A reshape keeps the elements in memory order and only changes the shape, so it is only equivalent to a transpose when no real data has to move. Trimming also dropped leading dims without checking that they were 1.

Region [1, 6, 4], node_end_perm = [0, 2, 1]:

Constant	padded	target_shape → trimmed	moved (not checked)	Op emitted	Result
[6, 1] (per-channel)	[1, 6, 1]	[1, 1, 6] → [1, 6]	[1]	view_copy → [1, 6]	correct
[4]	[1, 1, 4]	[1, 4, 1] → [1]	[2]	view_copy → [1]	crash: shape '[1]' is invalid for input of size 4
[6, 4]	[1, 6, 4]	[1, 4, 6] → [4, 6]	[2, 1]	view_copy → [4, 6]	silently wrong output: the data is reshaped, not transposed

After

The constant is padded to the region's rank. The pass then checks whether the dims that hold data (size ≠ 1) keep their relative order under the permutation:

Step	Code	Meaning
1	padded = [1] * (permute_rank - const_rank) + original_shape	The constant's shape at the region's rank, as broadcasting sees it
2	moved = [d for d in node_end_perm if padded[d] != 1]	The non-unit dims, in output order
3	moved == sorted(moved)	Only unit dims move, so a reshape is exact: a single view_copy to the permuted shape (permuted), with leading 1s popped until it is back at the constant's rank or a non-unit dim is reached
4	otherwise	Data has to move: a view_copy to padded, which only adds leading 1s, followed by a permute_copy by node_end_perm

Only leading 1s are dropped, so per-channel constants keep the same reshape as before and a non-unit dim is never cut off; broadcasting restores the dropped 1s.

Constant	padded	permuted	moved	Sorted?	Ops emitted	Result
[6, 1] (per-channel)	[1, 6, 1]	[1, 1, 6] → popped to [1, 6]	[1]	yes	view_copy → [1, 6]	correct, same graph as before
[4]	[1, 1, 4]	[1, 4, 1] → popped to [4, 1]	[2]	yes	view_copy → [4, 1]	correct
[6, 4]	[1, 6, 4]	not used	[2, 1]	no	view_copy → [1, 6, 4], permute_copy → [1, 4, 6]	correct

Same-rank constants, single-element inputs and non-constant inputs are handled as before.

Differential Revision: D128417497

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

✅ No Failures

As of commit 3320f09 with merge base 7586c16 ():
💚 Looks good so far! There are no failures yet. 💚

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

[meta-codesync]

@pssrawat has exported this pull request. If you are a Meta employee, you can view the originating Diff in D128417497.