Replace tosa_dim_order with explicit NCHW↔NHWC permutes (#19015)

[mcremon-meta]

Summary:

Replace implicit tosa_dim_order-based layout handling with explicit
permute_copy ops around TOSA operators that require NHWC layout.

Rewrite passes insert explicit NCHW↔NHWC permutes

RewriteConvPass, RewriteAvgPool2dPass, and RewriteMaxPool2dPass
now insert aten.permute_copy nodes (NCHW→NHWC before the TOSA op,
NHWC→NCHW after) instead of relying on ToTosaMemoryFormatPass for
layout conversion. This makes layout transitions visible in the graph.

Grouped conv decomposition in NHWC

RewriteConvPass decomposes grouped convolutions (non-depthwise) into
per-group TOSA.CONV2D ops operating entirely in NHWC, with a single
input/output permute pair wrapping the whole group. Supports INT8,
INT16 (with and without bias) quantisation paths, including the full
INT16+bias chain: CONV2D → RESCALE(INT48→INT32) → ADD(bias) →
RESCALE(INT32→INT16).

ToTosaMemoryFormatPass scoped down

Now only assigns non-identity dim_order to parameter/buffer
placeholders (for weight serialisation) and graph I/O. Inserts
permute_copy instead of tosa.TRANSPOSE. Skips users that already
carry a matching permute (inserted by the rewrite passes).

TOSA dialect op metas expect NHWC

All TOSA op meta functions (CONV2D, CONV3D, DEPTHWISE_CONV2D,
AVG_POOL2D, MAX_POOL2D, TRANSPOSE_CONV2D) now assume NHWC
input layout and produce NHWC output shapes.

Removed tosa_dim_order shape remapping

tosa_shape() no longer reorders dimensions—just resolves symints.
_get_matching_fake_tensor() returns node.meta["val"] directly.
Serialisation mapping always uses identity dim_order.

Operator serialisation simplified

op_amax, op_amin, op_any, op_cat, op_sum, and op_permute
no longer remap reduction/concat axes through dim_order since
tensors are already in the layout expected by TOSA.

Permute optimisation passes added

Six shared passes from executorch/backends/transforms/ are now run
after TOSA lowering to fuse, cancel, and simplify the permutes
introduced above:

• RemovePermutesAroundElementwiseOps (extended for RESCALE)
• FuseTransposeOrPermuteOpPairsPass (extended for RESCALE)
• ReplaceNopTransposeOrPermuteWithViewPass
• PostponePermuteOpBelowSqueezeOrUnsqueezeLikeView
• FuseCascadedTransposeOrPermuteOps
• FuseCascadedViewOps

Removed passes

DecomposeConvWithInt16ActivationPass and DecomposeGroupedConvPass
are removed—their logic is now handled inline by RewriteConvPass.
RewriteSlicePass is repositioned after the permute optimisations.

Ethos-U55 partitioner simplified

The dual NCHW/NHWC permute constraint check is removed since tensors
are always in the expected layout at partition time.

Differential Revision: D100712787

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❌ 10 New Failures, 1 Cancelled Job, 4 Unrelated Failures

As of commit 1736792 with merge base 89600b3 ():

NEW FAILURES - The following jobs have failed:

• Lint / lintrunner (gh)
  >>> Lint for backends/arm/tosa/utils.py:
• Lint / lintrunner-mypy (gh)
  >>> Lint for backends/arm/_passes/rewrite_conv_pass.py:
• pull / android / run-emulator (gh)
  The process '/usr/local/lib/android/sdk/platform-tools/adb' failed with exit code 224
• pull / unittest-arm-backend-with-no-deps (test_pytest_models_tosa) / linux-job (gh)
  RuntimeError: Command docker exec -t beb4f5791d0a6e79d4326dd30508ebe9c958c52dda6c52280c0ef60bc40e8485 /exec failed with exit code 1
• pull / unittest-arm-backend-with-no-deps (test_pytest_ops_tosa) / linux-job (gh)
  RuntimeError: Command docker exec -t 69eef4fed105fac44a3fc5bf965eb8c661c31c7d0a5fdcb4aeef89c3cd236e71 /exec failed with exit code 1
• trunk / test-arm-backend-ethos-u (test_pytest_models_ethos_u55) / linux-job (gh)
  RuntimeError: Command docker exec -t cd340fd364a56e83bd2f35292f79c49e94f6fcda813a27895220be3c867aec79 /exec failed with exit code 1
• trunk / test-arm-backend-ethos-u (test_pytest_models_ethos_u85) / linux-job (gh)
  RuntimeError: Command docker exec -t e3195270272e66a896a389599a70ce47b5c6245eb5289d2b619f7349468e56a8 /exec failed with exit code 1
• trunk / test-arm-backend-ethos-u (test_pytest_ops_ethos_u55) / linux-job (gh)
  RuntimeError: Command docker exec -t 118d1a0516a3a6cd331b30cb289579fce980732672537d3b04b1de796685633f /exec failed with exit code 1
• trunk / test-arm-backend-ethos-u (test_pytest_ops_ethos_u85) / linux-job (gh)
  RuntimeError: Command docker exec -t a84652221eb88a117527cda49d30211c0c5a18b26fe7507b7412a7f97fc095dc /exec failed with exit code 1
• trunk / test-arm-backend-vkml (test_pytest_ops_vkml) / linux-job (gh)
  RuntimeError: Command docker exec -t e55d9312873d30cb197c4939800eead02c74391d28cd7d5a607f70e4671fde5a /exec failed with exit code 1

CANCELLED JOB - The following job was cancelled. Please retry:

• pull / unittest / macos / macos-job (gh)
  ##[error]The operation was canceled.

FLAKY - The following job failed but was likely due to flakiness present on trunk:

• pull / test-qnn-testsuite-linux / test-backend-linux (qnn, models) / linux-job (gh) (detected as infra flaky with no log or failing log classifier)

BROKEN TRUNK - The following jobs failed but was 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.
• trunk / unittest-release / 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 D100712787.

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

[AdrianLundell]

Hi, great work! I have written some comments below, but overall I think it might be easier for us to land the PR and let you run a final internal check rather than vice versa, just since we have a bit easier access to CI and since we have a broader scope in what the backend needs to handle (e.g. dim-order input). The resulting behavior should be very similar from your perspective. Does that sound good to you?

• I am getting an increased number of transposes on some graphs in our transpose-count suite, mainly due to the following issues:

  • Transposes not fused over elementwise branching in all
  • Transposes not fusing in upwards/downward forks
  • As long as we don't see regression on important models and have good regression tests I'm fine with landing current behaviour for now and follow up with improvements later.

• The override and beartype dependencies will have to be removed from permute_pass_utils

• A number of of tests failing:

  • The current PR does not consider inputs with non-contiguous dim-order and gives numerical diffs for such graphs. This is fixed with the
    NormalizeDelegateIoPass from https://github.com/pytorch/executorch/pull/18948/changes#diff-6f629d0126cf621c8e72045a9bf407b6ae8bd362ee1672d86014c0d081df6869. (Should be placed earlier in the pass pipeline than in that patch to avoid some cases where passes rewrite the dim-order before it has been handled.)
  • Many fails related to upsampling
  • A few other conv and other test-cases

• No reason to remove the various vaidations done for the max/avg-pool2d tosa ops.

• Some aesthetic nits:

  • Rewrite max_pool2d/avg_pool are completely rewritten as call-passes, would be preferrable with a minimal diff only inserting tranposes in the call_operator pass.
  • The way the new passes are introduced into the arm-backend differs from how we do it normally, would be good to follow the general structure.
  • You can remove the ToTosaMemoryFormat and tosa.TRANSPOSE traces and permute op helpers completely since they are not used anymore