[borys3kk]

Summary

Profiling the MI-GAN inpainting model (a conv encoder-decoder, fp16) on XNNPACK showed the graph split into 7-8 delegate partitions around portable upsample_nearest2d ops. This PR lowers upsample_nearest2d x2.

XNNPACK has no nearest resize, and its static_broadcast is defined but returns xnn_status_unsupported_parameter at reshape time. Nearest x2 is instead lowered as concatenate4 of the input with itself along channels ([N, H, W, 4C]), followed by depth_to_space_2d(block_size=2). This adds an XNNDepthToSpace2d node, appended at the end of the XNodeUnion in schema.fbs / runtime_schema.fbs so existing node IDs don't change, plus defineDepthToSpace2dNode in XNNCompiler.cpp. UpsampleNearest2dConfig only partitions static 4D inputs with an exact x2 output. An explicit scale_factors that is not exactly 2.0 stays portable as well: PyTorch indexes with the given scale rather than the size ratio, so scale_factor=2.004 doubles the size but picks different pixels. Static quantization is supported: the quantizer annotates the op with input/output shared qparams, and the intermediate value carries them, since neither node changes values.

The ChannelsLastTaggedReshapePass change that was originally part of this PR (keeping NHWC across scalar and low-rank broadcast operands) moved to its own PR at review request, together with the fix for the dynamic-shape low-rank operand case reported below.

Measured on MI-GAN 256/512 fp16 with the host runtime, together with that pass change, on an ExecuTorch 1.4.1-based branch: XNNPACK partitions went from 7/8 to 1, 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).

The new node needs an updated runtime. A program that uses it fails to load on an older runtime.

Review order: the serialization schema and XNNCompiler.cpp, then op_upsample_nearest2d.py with UpsampleNearest2dConfig and the quantizer annotator, then the tests.

Test plan

Tested on macOS arm64, with pybindings built from this branch (cmake --preset pybind, target portable_lib) so that the run-on-delegate stages use the new runtime node.

backends/xnnpack/test/ops/test_upsample_nearest2d.py covers x2 by scale and by size, fp16, fp32 and qs8 conv-upsample-conv lowered to one partition, and x3 and scale_factor=2.004 staying portable.

python -m pytest backends/xnnpack/test/ops/test_upsample_nearest2d.py backends/xnnpack/test/passes/test_channels_last_tagged_reshape.py backends/xnnpack/test/test_xnnpack_partitioner.py
# 35 passed

lintrunner reports no issues on the changed files.

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

⚠️ 15 Awaiting Approval

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

AWAITING APPROVAL - The following workflows need approval before CI can run:

• Build documentation (gh)
• Build Presets (gh)
• Cadence Build & Test (gh)
• Lint (gh)
• periodic (gh)
• pull (gh)
• Test ARM Backend (gh)
• Test CoreML Backend (gh)
• Test Cortex-M Backend (gh)
• Test NXP Backend (gh)
• Test QNN Backend (gh)
• Test QNN Windows MSVC build (gh)
• Test Vulkan Backend (gh)
• Test WebGPU Backend (gh)
• Test XNNPACK Backend (gh)

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

[borys3kk]

@pytorchbot label "release notes: xnnpack"

[JakeStevens]

sorry, can we split this into two PRs? to keep the changes atomic per pr instead of squishing them together. the upsample changes are additive and lower risk.

that being said, a lower-rank operand with symbolic dims breaks lowering with these changes. Before, a dynamic-shape model with conv_out + y, where y is [H, W] or [C, H, W] with dynamic H/W, now fails in preprocess. On main it lowered fine by falling back to NCHW.

[borys3kk]

sorry, can we split this into two PRs? to keep the changes atomic per pr instead of squishing them together. the upsample changes are additive and lower risk.

that being said, a lower-rank operand with symbolic dims breaks lowering with these changes. Before, a dynamic-shape model with conv_out + y, where y is [H, W] or [C, H, W] with dynamic H/W, now fails in preprocess. On main it lowered fine by falling back to NCHW.

I will look into that tomorrow

[borys3kk]

Done, split in two: this PR now only lowers upsample_nearest2d, the ChannelsLastTaggedReshapePass change moved to #24183.

Thanks for the dynamic-shape report, reproduced it: unsqueeze_to_4d put the symbolic dims into the view_copy shape, which XNNStaticReshape rejects. In #24183 an operand with symbolic dims keeps the NCHW path as on main, with a test for conv_out + y with y of shape [H, W] and dynamic H/W that also runs at a second input size.