Move permute optimization passes to shared transforms location (#19002)

[mcremon-meta]

Summary:

Move 6 permute optimization passes and their shared infrastructure from
executorch/backends/cadence/aot/ to executorch/backends/transforms/ so
they can be shared between the Cadence and Arm backends without a
cross-backend dependency.

New files:

• permute_pass_utils.py: base classes (HierarchicalInplacePassInterface,
  RemoveOrReplacePassInterface, FuseOpPairsAcrossBranchesPass) and
  utilities (get_arg, set_arg, get_transposed_dims, get_permuted_dims,
  get_shape, get_edge_overload_packet)
• fuse_cascaded_transpose_or_permute_ops.py
• fuse_cascaded_view_ops.py
• fuse_transpose_or_permute_op_pairs_pass.py
• remove_permutes_around_elementwise_ops.py
• postpone_permute_below_squeeze_view.py
• replace_nop_transpose_or_permute_with_view.py

The shared versions omit register_cadence_pass decorators and
cadence-specific ops from default op sets. Cadence files will subclass
these and re-add the decorators and ops.

Added OSS tests (test_permute_optimization_passes.py) for the 4 passes
that can be imported without quantized op registration:
FuseCascadedTransposeOrPermuteOps, FuseCascadedViewOps,
PostponePermuteOpBelowSqueezeOrUnsqueezeLikeView, and
ReplaceNopTransposeOrPermuteWithViewPass. These run in GitHub CI via
pytest and are discovered automatically through pytest.ini testpaths.

Reviewed By: ethansfng

Differential Revision: D101459577

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❌ 1 New Failure, 2 Unrelated Failures

As of commit 59c06be with merge base d408a10 ():

NEW FAILURE - The following job has failed:

• pull / test-llama-runner-qnn-linux (fp32, qnn_16a16w, qnn) / linux-job (gh)
  RuntimeError: Command docker exec -t 9d460379ea2d70e2f39fc716895215f7000bdfcf83dda90de5f1f254d93982a5 /exec failed with exit code 92

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]

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

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

@claude review this PR.

[claude]

Claude finished @digantdesai's task in 3m 9s —— View job

---

Review of PR #19002: Move permute optimization passes to shared transforms location

• Read new shared transform files
• Read modified cadence files (subclasses)
• Read modified arm backend file
• Read test file
• Read build files (BUCK, targets.bzl)
• Analyze CI failures
• Post review feedback

---

Overall this is a well-structured refactor that cleanly separates shared permute optimization logic from cadence-specific decorators/ops. The inheritance pattern is sound. I have one blocking issue and several non-blocking items.

Blocking: Python 3.10 compatibility (typing.override)

backends/transforms/permute_pass_utils.py:17 — override is imported from typing, but it was only added in Python 3.12. The CI runs on Python 3.10 and fails with:

ImportError: cannot import name 'override' from 'typing' (/opt/conda/envs/py_3.10/lib/python3.10/typing.py)

This breaks all 5 unittest jobs. The fix is to import from typing_extensions instead:

from typing_extensions import override

Or conditionally:

import sys
if sys.version_info >= (3, 12):
    from typing import override
else:
    from typing_extensions import override

Fix this →

Non-blocking: Lint issues (UFMT + F401)

The linter reports several issues introduced by this PR in the cadence files:

1. Import ordering (UFMT) in fuse_ops.py, remove_ops.py, reorder_ops.py, replace_ops.py — the new shared imports are not in the correct sorted position. Running lintrunner -a should auto-fix these.

2. Unused imports (F401) in the cadence files — after extracting code to shared modules, several imports became unused:

   • fuse_ops.py: collections.deque, typing.Callable, compiler_utils.get_permuted_dims, compiler_utils.get_transposed_dims, ExportPass
   • remove_ops.py: dataclasses.dataclass, dataclasses.field
   • reorder_ops.py: copy, typing.cast

   These should be cleaned up to keep the code tidy.

3. Formatting in fuse_ops.py and remove_ops.py — the set union expressions for bypass_ops and permutable_ops in the cadence subclasses need wrapping in parentheses per the formatter:

   # Current (lint fails):
   bypass_ops: set[EdgeOpOverload] = _SharedFuseTransposeOrPermuteOpPairsPass.bypass_ops | {...}
   # Expected:
   bypass_ops: set[EdgeOpOverload] = (
       _SharedFuseTransposeOrPermuteOpPairsPass.bypass_ops
       | {...}
   )

Minor observations

1. get_shape name collision — permute_pass_utils.py defines a get_shape(graph_module, node) function, while backends/transforms/__init__.py already exports a different get_shape(input_node) from addmm_mm_to_linear.py. These have different signatures and semantics. Since neither is re-exported from the same path at import time, this isn't a runtime conflict, but it could be confusing. Consider renaming one of them in a follow-up.

2. Test coverage — The PR notes that 2 of the 6 passes (FuseTransposeOrPermuteOpPairsPass and RemovePermutesAroundElementwiseOps) are not tested in OSS due to quantized op registration requirements. This is a reasonable pragmatic choice, but it would be good to track adding those tests when possible.

3. arm/process_node.py change — The identity dim_order short-circuit (if dim_order == tuple(range(len(dim_order))): return np_tensor) is a nice optimization bundled in the same PR. It's correct but logically separate from the permute pass refactor.

Summary

The refactoring approach is sound — extracting shared base classes into backends/transforms/ and having cadence subclass them to add decorators and ops. The test file is thorough for the passes it covers. The main issue to fix before merging is the Python 3.10 compatibility problem with typing.override, plus running the linter to fix import ordering and unused imports.

---

[digantdesai]

Thanks, these move makes sense. Running claude review to catch nits.