[borys3kk]

Summary

Split out of #24113 at review request: that PR now only lowers upsample_nearest2d, this one carries the ChannelsLastTaggedReshapePass change.

Profiling the MI-GAN inpainting model (a conv encoder-decoder, fp16) on XNNPACK showed about 25% of the run time spent in NCHW<->NHWC transposes around elementwise ops whose second operand is a Python scalar or a low-rank tensor. ChannelsLastTaggedReshapePass now keeps NHWC across scalar and low-rank broadcast operands of add/sub/mul/div/maximum/minimum:

• A single-element operand, such as a lifted Python scalar (the gain in x * 1.4142), no longer takes part in the layout decision. Before, the NHWC activation was transposed to NCHW for the multiply and back for the next convolution. The same filter is used by the re-converge loop from Fix XNNPACK channels-last reshape for per-channel binary ops under dynamic quant #20376.
• A rank < 4 operand that broadcasts over the trailing dims (a [H, W] noise map added to [N, C, H, W], or a [C, 1, 1] scale) is reshaped to 4D with leading ones, so it can be converted to NHWC itself instead of dragging the activation to NCHW.
• Dequantized operands are left at their own rank, because a view between the dequantize and the op would carry no quant params. In qs8 graphs they take the existing NCHW path.
• Operands with symbolic dims also keep the NCHW path, because XNNStaticReshape cannot express their shape. This covers the dynamic-shape conv_out + y case with y of shape [H, W] or [C, H, W] reported on [XNNPACK] Lower upsample_nearest2d x2 as concatenate4 + depth_to_space_2d #24113, which lowers again as it does on main.

Measured on MI-GAN 256/512 fp16 with the host runtime, together with the upsample_nearest2d lowering from #24113, on an ExecuTorch 1.4.1-based branch: layout conversions per forward went from 128 to 26, the delegated share of run time from 87/82% to 98/95%, and Method::execute dropped by 29%/28%. The output was unchanged (max abs diff 2e-3, fp16 noise).

Review order: channels_last_tagged_reshape_pass.py, then the tests.

Test plan

Tested on macOS arm64 with locally built pybindings.

New tests in backends/xnnpack/test/passes/test_channels_last_tagged_reshape.py cover a scalar operand, [H, W] and [C, 1, 1] operands, a lower-rank left operand, qs8 lower-rank operands that keep their rank, and a [H, W] operand with dynamic H and W that falls back to NCHW and runs at a second input size.

python -m pytest backends/xnnpack/test/passes/test_channels_last_tagged_reshape.py backends/xnnpack/test/passes/test_remove_redundant_copy_pass.py backends/xnnpack/test/ops/test_{add,sub,multiply,div,maximum,minimum}.py
# 62 passed

lintrunner reports no issues on the changed files.

cc @GregoryComer @digantdesai @cbilgin @JakeStevens @nil-is-all

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❌ 1 Awaiting Approval, 1 New Failure

As of commit bdf14d2 with merge base 422cc39 ():

AWAITING APPROVAL - The following workflow needs approval before CI can run:

• periodic (gh)

NEW FAILURE - The following job has failed:

• pull / test-arm-cortex-m-size-test (bare_metal) / linux-job (gh)
  RuntimeError: Command bash /__w/_temp/exec_script failed with exit code 1

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

[borys3kk]

@pytorchbot label "release notes: xnnpack"

[JakeStevens]

thanks for the split!

a few things:

A single-element operand with more than 4 dims breaks lowering eg with conv(x) * buf where buf is [1,1,1,1,1], the scalar is skipped, the only remaining input is NHWC, and the op is tagged NHWC but its output is 5D and lowering then fails with RuntimeError: Converting to nhwc requires 4d tensor. It's unlikely to show up in real models, but the fix is one line:
return val is not None and val.numel() == 1 and val.dim() <= 4

Lower-rank constants now cost a reshape and transpose on every inference. we could add a pass to do this AOT.

We should include a full to_edge_transform_and_lower → to_executorch → serialize → run test for at least one static lower-rank case