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

⏳ No Failures, 2 Pending

As of commit 12f2694 with merge base 4628b2b ():
💚 Looks good so far! There are no failures yet. 💚

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.

[nil-is-all]

@borys3kk rebasing with main should resolve the failing CI tests

[borys3kk]

Merged main in 12f2694, thanks.

[JakeStevens]

The shared-weight fix in c22288b7 holds — I checked it across several configs. Two blocking items are inline: test coverage of the lowered path, and the external-constant tag.

[JakeStevens]

the lowered tests pass whether or not the decomposition happens

_check_lowered checks for one delegate, no top-level convolution, and matching outputs. All three already held before this PR: ConvolutionConfig delegates an fp32 depthwise transposed conv and it serializes as XNNConvTranspose2d. The structural checks run the pass on its own through RunPasses, so nothing exercises it in the default XNNPACKPassManager list that preprocess uses. If it falls out of that list, gets reordered, or _plan stops matching inside preprocess, all 17 tests stay green and the speedup quietly disappears. test_deconv_relu only calls _check_lowered, so the behaviour in its comment is never checked.

Repro: stub the pass out during lowering with mock.patch.object(DecomposeDepthwiseConvTransposePass, "call", lambda self, gm: PassResult(gm, False)). DwDeconv(4, 2, 1) still lowers to 1 delegate with 0 top-level convs, and _decompose is never called.

Suggested fix: assert in _check_lowered that the pass actually ran during lowering, e.g.

with mock.patch.object(
    DecomposeDepthwiseConvTransposePass,
    "_decompose",
    autospec=True,
    side_effect=DecomposeDepthwiseConvTransposePass._decompose,
) as decompose:
    tester.to_edge_transform_and_lower()
self.assertEqual(decompose.call_count, num_deconvs)

For test_deconv_relu, add a pass-level check: run RunPasses([DecomposeDepthwiseConvTransposePass, FuseActivationPass]) and assert the single relu reads the depth_to_space_2d output. That pins down the ordering the description relies on. If the pass ever moves after FuseActivationPass, the ReLU gets fused into the conv the pass deletes. A pass list in the default order, as in test_fuse_pad_into_conv_pass.py, would let the same test check that the phase pads fold.

[JakeStevens]

phase weights fall out of the external .ptd

create_constant_placeholder sets only meta["val"] on the new placeholder. So the phase weights lose the original weight's meta["custom"]["delegate_constant_tag"], and line 225 then deletes the tagged original. NodeVisitor uses that tag to route constants to the external file (operators/node_visitor.py:650). Without it, the phase weights silently land in the .pte. PropagateCustomMetaPass only handles q/dq nodes and runs before this pass, so nothing puts the tag back.

Repro: tag DwDeconv(4, 2, 1) with delegate_external_constants_pass_unlifted(gen_tag_fn=lambda _: "model") (same setup as test_propagate_custom_meta_pass.py), then lower and serialize.

• as is: 4 named-data entries in the .pte (the phase weights) and 1 in model.ptd (the bias)
• with the tag copied: 0 in the .pte and 5 in model.ptd

Suggested fix: copy it right after creating each phase placeholder

if "custom" in weight_node.meta:
    phase_weight_node.meta["custom"] = weight_node.meta["custom"]

and add a test modeled on test_propagate_custom_meta_pass.py that asserts the .pte has no named data. FuseBatchNormPass has the same gap, so having create_constant_placeholder take a node to copy meta from would be the general fix, but that can wait for a follow-up.