Repository navigation
[XNNPACK] One polyphase pass for dense and depthwise transposed convs - #25041
Draft
JakeStevens wants to merge 20 commits into
Draft
JakeStevens wants to merge 20 commits into
JakeStevens wants to merge 20 commits into
Conversation
…ow-rank broadcast operands Profiling a conv encoder-decoder (MI-GAN inpainting, fp16) showed the graph split into 7-8 XNNPACK partitions around portable upsample_nearest2d ops and ~25% of the time spent in NCHW<->NHWC transposes. upsample_nearest2d.vec with an exact x2 scale on a static 4D input is now partitioned and lowered as concatenate4 of the input along channels followed by depth_to_space_2d(block_size=2). XNNPACK has no nearest resize and its static_broadcast is not implemented at run time. This adds an XNNDepthToSpace2d node to the flatbuffer schemas and XNNCompiler. Static quantization is supported: the quantizer annotates the op with shared input/output qparams and the intermediate value carries them. Other scales, including an explicit scale_factors that is not exactly 2.0, stay portable. In ChannelsLastTaggedReshapePass, single-element operands of add/sub/mul/div/maximum/minimum (a lifted Python scalar such as the gain in x * 1.4142) no longer take part in the layout decision, and rank < 4 operands that broadcast over the trailing dims (a [H, W] map added to [N, C, H, W]) are reshaped to 4D with leading ones so they follow the NHWC operand instead of dragging the activation back to NCHW. Dequantized operands keep their rank, since a view between the dequantize and the op would carry no quant params.
Backend passes keep a reference to the ExportedProgram they were built from (XNNPACKPass.exported_program, for one), while ExportedProgramPassManager runs them on a shallow copy of it and rebuilds the signature of that copy afterwards. delete_constant_placeholder replaced the original program's signature object, so the copy kept the old one and its verifier failed with Parameter <name> is not in the state dict. create_constant_placeholder already inserts into the shared input_specs list in place; do the same for deletion. Tests cover a deletion made inside a pass run through the manager for every constant placeholder kind. Authored with Claude Code.
…vs + depth_to_space 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 MI-GAN the twelve 4x4 / stride-2 filters more than double the forward. DecomposeDepthwiseConvTransposePass rewrites such a node, when it is a clean x s upscale ((H-1)s - 2p + k + output_padding == sH, dilation 1, every phase's taps inside the input), into s*s stride-1 depthwise convs on the low-resolution input (phase (ay, ax) correlates with every s-th tap of the kernel, through a possibly asymmetric zero pad), a channel concatenation in (ay, ax) order and a block-major depth_to_space_2d. That last op is XNNPACK's native channel order (aten.pixel_shuffle is channel-major), so it is a new et_xnnpack custom op created only after partitioning, with a reference implementation for tracing and tests, a node visitor emitting XNNDepthToSpace2d, and an NHWC tag in the channels-last pass. The pass runs before FuseActivationPass, so an activation after the transposed conv is never fused into the node it removes. Every convolution is planned before any is rewritten, so deconvs feeding each other work, and the phase weights keep the original weight's input kind (parameter, buffer or constant). Deleting the folded weight relies on delete_constant_placeholder editing the graph signature in place. Authored with Claude Code.
create_constant_placeholder returns an existing placeholder of the same name, so two depthwise deconvs sharing one weight but using different stride or padding ended up with the first one's phase kernels. Include both in the name and add tests for a shared weight with different and with identical parameters.
create_constant_placeholder only sets meta["val"], so the phase weights lost the delegate_constant_tag of the weight they replace and landed in the .pte instead of the .ptd file.
…tionPass, and external constants XNNPACK delegates an fp32 depthwise transposed conv on its own, so the lowered tests passed whether or not the decomposition happened. _check_lowered now counts the _decompose calls and checks the serialized graph for depth_to_space_2d and no transposed conv. test_deconv_relu runs the default pass list and checks that the relu reads the depth_to_space_2d output and that the phase pads fold. A new test lowers with external constants and expects every constant in model.ptd.
…v-transpose-main # Conflicts: # backends/xnnpack/_passes/channels_last_tagged_reshape_pass.py # backends/xnnpack/operators/op_upsample_nearest2d.py # backends/xnnpack/partition/config/generic_node_configs.py # backends/xnnpack/quantizer/xnnpack_quantizer_utils.py # backends/xnnpack/test/ops/test_upsample_nearest2d.py
…ch#24810) Summary: Pull Request resolved: pytorch#24810 Adds `examples/models/phonon`: Gradium's Apache-2.0 Phonon 31M English text-to-speech model exported to one `.pte`, with a backend-agnostic C++ runner (`phonon_runner`) and a Python runner (`run_phonon.py`). It follows `nemotron3_diarization`: one export script, CMake presets with `make phonon-{cpu,cuda,mlx,vulkan}`, CI jobs in `pull.yml`, `mlx.yml` and `cuda.yml`, and a short README in the same shape. CUDA writes `aoti_cuda_blob.ptd` beside the `.pte`. Validated against the official runtime on XNNPACK, Vulkan, WebGPU, CUDA and MLX. Exports default to FP32 on every backend; `--dtype bf16` is available on CUDA and MLX. WebGPU runs in the browser: the README gives the `--backend webgpu --webgpu-compat-ops` export and links the Text to Speech demo of the ExecuTorch WebGPU Playground. Also fixes an XNNPACK delegate use-after-free that Phonon hits: `XNNCompiler::compileModel` frees the named-data buffer of every packed filter/bias right after `xnn_create_runtime_v4`, even when another node still reads that constant. Phonon's `StreamingConvTranspose1d` feeds the transposed-conv bias to the conv and to the held-back tail subtraction, so a program loaded from a file (C++ runner, pybindings) segfaulted in the first `mimi_decode` call. A packed value that another node also reads now stays alive; an unknown node type keeps every constant alive. Same class of bug as the PReLU fix (pytorch#21480). Differential Revision: D128116963
…#24929) Summary: Pull Request resolved: pytorch#24929 Follow-up to D128116963. The C++ runner gets `--stream`, which writes 16-bit mono PCM to stdout as each 80 ms frame is decoded instead of writing the WAV at the end, so playback can start after the first frame (`phonon_runner ... --stream | aplay -q -f S16_LE -r 24000 -c 1`). The runner already decoded frame by frame through a callback; this only routes the samples to stdout and moves the summary line to stderr. `--qlinear` now accepts only the modes that each backend has a kernel for (XNNPACK 8da4w/8da8w, CUDA 4w/8w, Vulkan and WebGPU 4w/8da4w, MLX 4w/8w) and rejects the rest at export instead of failing at runtime. Quantization covers the flow LM's transformer and flow net; the input projection, the one-output EOS head and the Mimi codec stay in the export dtype. The EOS head is skipped because the Vulkan and WebGPU 4-bit linear does not yet handle an output size that is not a multiple of 4. CUDA 4w uses the packed int4 kernel, which needs `--dtype bf16` and packs on the GPU, so two `torch.arange` calls in the attention take the input's device. MLX 4w uses asymmetric int4: symmetric int4 keeps the words but delays the end-of-speech decision, adding up to two seconds of silence at the end of a sentence, while asymmetric int4 keeps the timing of FP32. The README documents both features with the sizes per backend. This diff was written with the help of Claude Code (Claude Opus 5.5). Reviewed By: mergennachin Differential Revision: D131277773
… a multiple of 4 (pytorch#24930) Summary: Pull Request resolved: pytorch#24930 make_linear_q4gsw_op and make_linear_dq8ca_q4gsw_op pad the weight scales and the bias of a 4-bit linear to a multiple of 4 output channels, but they do not pad the packed weights. The CompositeExplicitAutograd implementation of et_vk.linear_q4gsw (also used by et_vk.linear_dq8ca_q4gsw) dequantized the unpadded weights with the padded scales, so any lowering that re-traced the graph failed for such a linear, for example the 512 -> 1 end-of-speech head of a TTS model: RuntimeError: shape '[1, 16, 1]' is invalid for input of size 64 The implementation now drops the padded scale and bias rows. This affects the export-time reference implementation only; the dynamically quantized 8-bit activation, 4-bit weight kernel already handles any number of output channels on device. This diff was written with the help of Claude Code (Claude Opus 5.5). Reviewed By: SS-JIA Differential Revision: D131277774
…#24971) Summary: Pull Request resolved: pytorch#24971 The quantized linear pattern detector matched every linear with a dequantized constant weight, while replace_quantized_linear_patterns only has custom ops for 8-bit statically quantized activations and for 4-bit group-quantized weights (weight-only or with dynamically quantized activations). The partitioner delegates every node of a match, so an 8-bit weight-only linear (torchao IntxWeightOnlyConfig int8) was delegated with its torchao.dequantize_affine node left in place. The Vulkan runtime has no dequantize_affine kernel (the quantize/dequantize op implementations were removed in pytorch#15750), so the program failed to initialize on device. The registry entries for torchao.quantize_affine and torchao.dequantize_affine still claimed both ops as supported, which also delegated them outside of patterns. The detector now reports a match only when a custom op exists for it and the runtime can run it. The weight-only 4-bit kernel packs the weights in blocks of 4 output channels (Q4gswLinear.cpp checks N % 4 == 0 at prepack), so weight-only 4-bit linears with other output sizes are no longer matched; previously they exported and then failed at init. The match now also records the activation's quantize node in the dynamically quantized pattern, which the partitioner previously picked up only through the stale quantize_affine registry entry, and the two torchao registry entries are removed. Unmatched linears now stay outside of the delegate. With torchao quantization this surfaces at export as "Missing out variants: {'torchao::dequantize_affine'}" instead of a failure on device. WebGPU uses the same patterns, so the same applies there. Not changed: the quantized_decomposed quantize_per_channel / quantize_per_token / dequantize_per_channel / dequantize_per_token and choose_qparams (tensor, per_token_asymmetric) registry entries also have no runtime kernel since pytorch#15750. I left them because the VulkanQuantizer dynamic flow may still depend on them being claimed; they are worth a look by the backend owners. This diff was written with the help of Claude Code (Claude Opus 5.5). Reviewed By: SS-JIA Differential Revision: D131277775
…s on XNNPACK Summary: XNNPACK's deconvolution is the most expensive op in Phonon's Mimi decoder: about 45% of each frame on x86 and 38% on arm64. Every transposed convolution in Mimi has kernel = 2 * stride, so each one is exactly a kernel-2 convolution with stride * out_channels outputs, one per output phase, whose phases are then interleaved into the output steps. The XNNPACK export now applies that rewrite (conv_transpose1d_as_conv1d) to the three decoder transposed convolutions, and the existing depthwise_conv_transpose_as_mul to the upsampler. The output is the same up to float rounding: the exported 31M program writes audio within 1 LSB (16-bit PCM) of the program without the rewrite, with the same length and frames. Running the transposed convolutions as 2-D convolutions with the sequence on the width axis (conv_transpose1d_as_2d) does not help: Mimi stays at 8.2 ms per frame, so the gain comes from the convolution kernel rather than the tiling axis. This diff was written with the help of Claude Code (Claude Opus 5.5). Differential Revision: D131605676
The polyphase transposed conv lowering below uses et_xnnpack::depth_to_space_2d from pytorch#24114 and replaces a module swap of pytorch#25004. Conflicts with main resolved in the XNNPACK pass list and the memory-format-sensitive ops of ChannelsLastTaggedReshapePass, which no longer lists mean.dim.
…to space XNNPACK's deconvolution is much slower than its convolution for the upsamplers of audio decoders and vocoders, transposed convs with kernel = 2 * stride. DecomposeConvTransposePass lowers a dense (groups == 1) 1-D or 2-D transposed conv whose kernel is a multiple of its stride as one stride-1 conv that computes every output phase, an interleave of the phases (depth_to_space_2d in 2-D, a view, permute and view in 1-D) and a slice for the padding. The Phonon XNNPACK export drops conv_transpose1d_as_conv1d, which did the same rewrite for the Mimi decoder before export. This change was authored with Claude Code (Claude Opus 5.5).
DecomposeDepthwiseConvTransposePass and DecomposeConvTransposePass were the same polyphase decomposition, each with its own matching, phase math, constants, interleave and place in the pass list. DecomposeConvTransposePass now does both: a dense transposed conv computes all phases in one conv, and a depthwise one computes each phase in its own depthwise conv, since XNNPACK only uses its depthwise microkernels with one output channel per group. The depthwise case no longer needs a clean x s upscale in 2-D. Each phase reads the input through whatever zero pad its taps need, a slice trims an output size that is not a multiple of s, and in 1-D the phases run as 2-D convs on one unsqueezed input, so that XNNPACK converts the layout once. Symmetric phase pads become the conv's padding. This change was authored with Claude Code (Claude Opus 5.5).
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/25041
Note: Links to docs will display an error until the docs builds have been completed. ❗ 1 Active SEVsThere are 1 currently active SEVs. If your PR is affected, please view them below: ❌ 3 Unclassified FailuresAs of commit 2eda7a0 with merge base 6411f1e ( UNCLASSIFIED FAILURES - DrCI could not classify the following jobs because the workflow did not run on the merge base. The failures may be pre-existing on trunk or introduced by this PR:
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
meta-cla
Bot
added
the
CLA Signed
This PR needs a
|
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Stacked on #25037, whose commits (with #24114 and #25004) are included in this branch; only the top commit is new.
DecomposeDepthwiseConvTransposePass(#24114) andDecomposeConvTransposePass(#25037) were the same polyphase decomposition, each with its own matching, phase math, constants, interleave and entry in the pass list. This merges them intoDecomposeConvTransposePass, which keeps two ways of computing the phases. A dense transposed conv (groups == 1) computes all of them in one conv with s * out_channels outputs, as in #25037. A depthwise one computes each phase in its own depthwise conv, as in #24114: XNNPACK only uses its depthwise microkernels with one output channel per group (select_microkernel_typeinconvolution-nhwc.c), so one conv with s² outputs per channel would fall back to the grouped IGEMM path that #24114 avoids. The rest is shared: the per-axis phase taps and pads (the formula of #24114), the phase constants that keep the external-constant tag, the interleave (depth_to_space_2d, or a view, permute and view in 1-D) and the slice. The merged file keeps the Software Mansion copyright of #24114, whose depthwise code it carries.The depthwise case also covers more. It no longer needs a clean x s upscale in 2-D: each phase reads the input through whatever zero pad its taps need, and a slice trims an output size that is not a multiple of s. 1-D depthwise transposed convs work too, as 2-D phase convs on one unsqueezed input that are concatenated before one squeeze, so that XNNPACK converts the layout once rather than around every phase (27 XNNPACK nodes instead of 87 for 16 phases). Symmetric phase pads become the conv's padding instead of a
constant_pad_nd. Phases without taps or that would need a negative pad still stay on XNNPACK's deconvolution.Depthwise transposed convs alone through the XNNPACK delegate on an Apple M4 Max, median µs per call, XNNPACK's deconvolution → this pass (12-thread timings vary by up to 30% between runs on this machine; the ordering does not):
The last row is the upsampler of Phonon's Mimi decoder, which the example runs as a multiply (
depthwise_conv_transpose_as_mul). With the upsampler on this pass instead,mimi_decodetakes as long as with the swap at 1 and 4 threads (1.900 against 1.904 ms, 1.972 against 1.974 ms per frame) but 3.8% longer at 12 threads (4.40 against 4.24 ms), where each of the 16 phase convs adds threadpool dispatch to a single input step. So the example keeps the swap, and its exported program is byte-identical to the one of #25037.Test plan
The tests of #24114 now run against
DecomposeConvTransposePass. Its case of a deconv that is not a clean upscale is now lowered, and the case left alone is a phase that would need a negative pad. New tests lower 1-D depthwise transposed convs without padding (strides 2 and 16, one and five steps), one with an asymmetric pad and an odd output size, and a 2-D one with an odd output size, and compare them with eager through the delegate.lintrunnerreports no issues on the changed files. Tested on macOS arm64 with pybindings built from this branch.This PR was authored with Claude Code (Claude Opus 5.5).