[borys3kk]

Summary

Stacked on #24113 and #24111: their commits are included in this branch, and only the commit titled like this PR is new here. #24113 adds the XNNDepthToSpace2d node that this PR emits, and #24111 lets the pass delete the folded weight through delete_constant_placeholder.

XNNPACK runs transposed convolutions through its subconvolution IGEMM path one group at a time. With one channel per group (a depthwise ConvTranspose2d, the natural way to write a StyleGAN upfirdn x2 upsampling filter) that is far slower than the depthwise conv microkernels. On the MI-GAN inpainting model, the twelve 4x4 / stride-2 filters more than double the forward.

DecomposeDepthwiseConvTransposePass rewrites a depthwise transposed conv (groups == in_channels == out_channels, weight lifted as a parameter, buffer or constant) into s*s stride-1 depthwise convs on the low-resolution input, a channel concatenation in (ay, ax) order, and a block-major depth_to_space_2d. Phase (ay, ax) correlates the input with every s-th tap of the kernel, through a zero pad that can be asymmetric. The pass only fires when the node is a clean xs upscale on both axes, (H - 1)s - 2p + k + output_padding == sH with the same s, dilation is 1, and every phase's taps fall inside the input. Anything else, including dequantized weights and grouped but not depthwise convs, stays on XNNPACK's own deconvolution node. Single-element stride / padding / output_padding are broadcast the same way op_conv2d.py does, and dynamic spatial dims work.

depth_to_space_2d is block-major: input channel (by * s + bx) * C + c goes to output pixel (h * s + by, w * s + bx), channel c. That is the order XNNPACK's depth-to-space uses, and it matches the phase concatenation directly. aten.pixel_shuffle is channel-major and does not map onto the node. So this is a new et_xnnpack::depth_to_space_2d custom op in backends/xnnpack/custom_ops_lib.py, registered with torch.library the same way backends/vulkan/custom_ops_lib.py does it. The pass creates it only after partitioning. Its reference implementation is there for tracing and tests, op_depth_to_space.py serializes it as XNNDepthToSpace2d, and ChannelsLastTaggedReshapePass tags it NHWC.

The pass runs right before FuseActivationPass, so an activation after the transposed conv is never fused into the node the pass removes. It stays a separate node after depth_to_space_2d. Every convolution is planned before any is rewritten, so deconvs that feed each other work. The phase weights keep the original weight's input kind.

The phase pads are constant_pad_nd nodes on the conv input. With #24112 they fold into the new convs; without it they run as XNNPACK static pads.

Device numbers for MI-GAN fp16 on a Galaxy S23 (Snapdragon 8 Gen 2), measured on an ExecuTorch 1.4.1-based branch with the equivalent of #24113 and #24112 applied. Setup: executor_runner, 4 threads, mean of 10 forwards after 2 warm-ups, 3 interleaved rounds with 15 s cooldowns, round spread under 1%. "Stock deconv" is the depthwise ConvTranspose2d left on XNNPACK's deconvolution node. "Hand-written" is the same filter written by hand as nearest upsample + checkerboard mask + pad + 4x4 depthwise conv.

size	stock deconv	hand-written	this PR	vs hand-written
256	235 ms	108 ms	100 ms	-7%
512	455 ms	189 ms	171 ms	-10%
faces_256	234 ms	110 ms	100 ms	-9%

Outputs are bit-identical to the hand-written variant. The stock deconv differs by 1e-2 (fp16 accumulation order). The polyphase convs do a quarter of the filter MACs, and the upsample and mask passes go away.

Review order: custom_ops_lib.py and operators/op_depth_to_space.py, then _passes/decompose_depthwise_conv_transpose.py with its registration in _passes/__init__.py and the NHWC tag in channels_last_tagged_reshape_pass.py, then the BUCK targets and the tests.

Test plan

Tested on macOS arm64, with pybindings built from #24113 (cmake --preset pybind, target portable_lib) so that the run-on-delegate stages use a runtime with XNNDepthToSpace2d. This PR has no C++ changes.

New tests in backends/xnnpack/test/passes/test_decompose_depthwise_conv_transpose.py check the node counts after the pass and compare against eager, both after the pass (reference depth_to_space_2d) and on the XNNPACK delegate. They cover:

• 4x4 stride 2, 3x3 with output padding, 2x2 without padding, stride 3, and fp16 without bias
• batch 3, a conv before the deconv, two deconvs in a row, deconv + ReLU, dynamic height and width, and single-element spatial params
• non-clean upscales, grouped but not depthwise convs, and dequantized weights, all left alone
• the phase math in plain torch against conv_transpose2d
• one weight shared by two deconvs with different padding (separate phase weights) and with identical parameters (shared phase weights)

python -m pytest backends/xnnpack/test/passes/test_decompose_depthwise_conv_transpose.py
# 17 passed
python -m pytest backends/xnnpack/test/ops backends/xnnpack/test/passes backends/xnnpack/test/quantizer backends/xnnpack/test/serialization backends/xnnpack/test/test_xnnpack_partitioner.py -n 6
# 386 passed
python -m pytest -c /dev/null backends/test/suite/operators/test_convtranspose2d.py backends/test/suite/operators/test_conv2d.py -m backend_xnnpack -n 6
# 96 passed

backends/xnnpack/test/quantizer/test_xnnpack_quantizer.py run on its own gives 31 passed and 1 failed. test_gru fails the same way on the base branch ("quantized engine NoQEngine is not supported"). lintrunner reports no issues on the changed files.

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

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❌ 3 New Failures

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

NEW FAILURES - The following jobs have 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
• pull / test-arm-cortex-m-size-test (zephyr-preset) / linux-job (gh)
  RuntimeError: Command bash /__w/_temp/exec_script failed with exit code 1
• pull / unittest-editable / linux / linux-job (gh)
  [ FAILED ] TextGenerationTest.OpaqueCancelAndResetReleaseBackingBeforeWaitReturns

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

[borys3kk]

@pytorchbot label "release notes: xnnpack"

[JakeStevens]

there is one bug: two deconvs that share one weight can silently produce wrong output

_decompose names each phase weight f"{weight_node.name}_phase_{ay}_{ax}". When a placeholder with that name already exists, create_constant_placeholder returns it instead of making a new one. But, the phase kernels depend on the stride and padding as well as the weight. So the second deconv gets the first one's kernels whenever the weight is shared but the stride or padding differs.

Repro: one Parameter used by two depthwise deconvs, both clean ×s upscales, with different padding.

• k=5, s=3, (p=1, op=0) and (p=2, op=2): the pass crashes. The aten.cat shape check fails during the retrace because the reused kernels have
  the wrong number of taps.
• k=9, s=3, (p=3, op=0) and (p=4, op=2): no error, but the output is wrong. Only 9 p_w_phase_* placeholders exist for 18 phase convs, and the
  max error against eager is 15.3.
• Control, same model with two separate copies of the weight: max error 1.9e-6.

Suggested fix: include s and padding in the name, e.g. f"{weight_node.name}_s{s}p{padding[0]}x{padding[1]}phase{ay}{ax}".
A regression test along the lines of the k=9 case is worth adding.

[borys3kk]

Good catch, thanks. Fixed in c22288b: the phase placeholders are now named {weight}_s{s}_p{py}x{px}_phase_{ay}_{ax}, so a weight shared by deconvs with different stride or padding gets its own set of phase kernels, while identical deconvs still share one set. Added test_shared_weight_different_padding (your k=9, s=3 case: 18 phase placeholders, matches eager) and test_shared_weight_same_params (9 shared placeholders). The first one fails on the previous commit with the output mismatch you saw.

[borys3kk]

Addressed the review in c22288b (phase placeholders keyed by stride and padding, two shared-weight tests). 17 passed in test_decompose_depthwise_conv_transpose.py, lintrunner clean, test plan in the description updated. The branch is still stacked on #24113 and #24111; I will rebase once they land.