Remove or move permute after mean

[DrJessop]

Summary:
If we have a permute -> unary chain -> mean, based on the reduction dims of the mean, we can either fully remove the preceding permute or move the permute after the mean.

Case 1: Dims after permute are still in same order with respect to each other, we can fully get rid of the permute and just update the reduction dims of the mean.

Case 2: Not case 1. In this case, it's better to move the permute after the mean, since the permute will operate on less data.

Differential Revision: D102268214

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❗ 1 Active SEVs

There are 1 currently active SEVs. If your PR is affected, please view them below:

• Rolling out OSDC (ARC) runners on pull & trunk workflows in PyTorch main

❌ 2 Cancelled Jobs, 2 Unrelated Failures

As of commit fe0d927 with merge base c3f3d12 ():

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

• periodic (gh)
• pull / test-coreml-bc-macos (macos-m2-stable) / macos-job (gh)

BROKEN TRUNK - The following jobs failed but were present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

• pull / unittest / windows / windows-job (gh) (trunk failure)
  ##[error]The operation was canceled.
• pull / unittest-editable / windows / windows-job (gh) (trunk failure)
  ##[error]The operation was canceled.

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

[meta-codesync]

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

[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.

[mergennachin]

@claude Review this PR

[claude]

Claude finished @mergennachin's task in 3m 44s —— View job

---

Review of PR #19103: Remove or move permute after mean

• Read changed files and understand the optimization passes
• Review RemovePermuteBeforeMeanPass correctness
• Review test coverage
• Post review feedback

---

Summary

This PR adds RemovePermuteBeforeMeanPass, a well-structured optimization for the Cadence backend that handles permute → [unary chain] → mean patterns. The two-case approach (full removal vs. sinking) is sound, and the tests include numerical validation which is good.

Overall the code is correct and well-tested. I have a few observations and one potential issue below.

Issues

1. Potential issue with negative reduction_dims in _can_fully_remove (minor)

remove_ops.py:486 — The remapping new_reduction_dims = [perm[d % ndim] for d in reduction_dims] correctly handles negative dims via Python's modulo. However, _can_fully_remove at line 448 computes canonical_reduction = {d % ndim for d in new_reduction_dims} on the already-remapped dims, which are guaranteed non-negative, so the % ndim is redundant but harmless.

The actual concern is the original reduction_dims — if they contain negative values, the remapping is correct (perm[-1 % 4] = perm[3]), but the _can_fully_remove check operates on new_reduction_dims which are always non-negative after the perm lookup. This is fine.

No action needed — just noting for clarity.

2. Missing meta on newly created permute node — not a bug (remove_ops.py:499-503)

The new_permute node created via graph.create_node has no meta["val"] set. This is handled correctly by the pass infrastructure: HierarchicalInplacePassInterface.call() invokes ExportPass.call() after mutations, which re-interprets the entire graph and populates metadata on all nodes. No issue here.

Suggestions

3. The _find_permute_through_unary_chain user check is redundant with maybe_remove_or_replace (remove_ops.py:430 vs remove_ops.py:480)

The method _find_permute_through_unary_chain checks len(current.users) != 1 for unary ops in the chain (line 430), ensuring each intermediate node has exactly one consumer. Then maybe_remove_or_replace separately checks len(permute_node.users) != 1 (line 480). This is correct — but note that when the unary chain is empty (permute directly feeds mean), the permute's user count is checked only at line 480. When there IS a unary chain, the permute's user count is NOT checked inside _find_permute_through_unary_chain — only the unary nodes' user counts are. The check at line 480 then catches the multi-user permute case. This separation works but could be slightly confusing to future readers.

4. Consider a test with quantize/dequantize in the unary chain

All current tests use relu as the unary op between permute and mean. Since _UNARY_TARGETS includes quantize/dequantize ops (which are arguably the primary real-world use case for this pass), a test with those ops in the chain would strengthen coverage. Not a blocker.

5. No test with negative reduction dims

The code handles negative reduction dims via d % ndim, but there's no test exercising this path. A test like mean(relu, [-1], False) with a preceding permute would cover this. Not a blocker.

6. Pass ordering in CommonRemovePasses looks correct (remove_ops.py:770-771)

RemovePermuteBeforeMeanPass is placed before RemovePermutesAroundElementwiseOps and FuseTransposeOrPermuteOpPairsPass. This is the right order: if the mean pass sinks a permute, subsequent passes can potentially optimize it further (e.g., fuse with another permute or cancel it out with a surrounding permute pair).

Verdict

The logic is correct, the math for dim remapping and post-mean permutation computation checks out, and all test cases include numerical validation. The code is clean and follows existing patterns in the file. This looks good to land.

---