Repository navigation
[XNNPACK] Lower dense transposed convs as one polyphase conv + depth to space - #25037
Draft
JakeStevens wants to merge 19 commits into
Draft
JakeStevens wants to merge 19 commits into
JakeStevens wants to merge 19 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).
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/25037
Note: Links to docs will display an error until the docs builds have been completed. ❌ 1 New Failure, 3 Pending, 1 Unclassified FailureAs of commit aa1d717 with merge base 6411f1e ( UNCLASSIFIED FAILURE - DrCI could not classify the following job because the workflow did not run on the merge base. The failure 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 #24114 and #25004, whose commits are included in this branch; only the top commit is new. It uses the
et_xnnpack::depth_to_space_2dop of #24114, and it moves the Mimi decoder rewrite of #25004 from the Phonon example into the XNNPACK backend.XNNPACK's deconvolution is much slower than its convolution for the upsamplers of audio decoders and vocoders, which are transposed convs with kernel = 2 * stride (Mimi, EnCodec, DAC, HiFi-GAN). #25004 works around it for Phonon with a module swap before export.
DecomposeConvTransposePassdoes the same during lowering, for every model: a dense (groups == 1) 1-D or 2-D transposed conv whose kernel is a multiple of its stride, k = m * s with s > 1, becomes one stride-1 conv with m taps and padding m - 1 that computes every output phase (s * out_channels outputs, s² * out_channels in 2-D), followed by an interleave of the phases and a slice for the padding:The phases are ordered phase-major, so the 2-D interleave is the block-major
depth_to_space_2dof #24114; the 1-D interleave is a view, a permute and a view. The bias is repeated per phase into the conv. The pass runs right afterDecomposeDepthwiseConvTransposePass, which keeps the depthwise ones, and beforeFuseActivationPass. Grouped transposed convs and those with dilation, output_padding > padding, unequal 2-D strides, a dequantized weight or dynamic shapes stay on XNNPACK's deconvolution node.The Phonon XNNPACK export drops
conv_transpose1d_as_conv1d: the backend now lowers the decoder's three transposed convs itself. It keepsdepthwise_conv_transpose_as_mulfor the depthwise upsampler, which this pass leaves alone and which XNNPACK still runs faster as a multiply (RTF 0.0359 without it at one thread).conv_transpose1d_as_conv1ditself stays for the other backends of #25020.Phonon 31M,
--backend xnnpack --quantize 8bit, a 16.6 s English prompt at temperature 0, Apple M4 Max. RTF (lower is faster) and the medianmimi_decodetime per 80 ms frame:phonon_runnerRTF, default threadsrun_phonon.pyRTF, 4 threadsrun_phonon.pyRTF, 1 threadmimi_decode, 4 threadsmimi_decode, 1 threadAll three give the same 207 frames, and the audio of this PR is within 0.06 LSB (16-bit PCM) of the program without any rewrite. The numbers are medians of 9 interleaved runs of
phonon_runner, or 5 ofrun_phonon.py, whose thread count is set with_unsafe_reset_threadpool.phonon_runnerstarts one thread per performance core, and cpuinfo counts all 16 cores of this machine (12 performance, 4 efficiency) as performance cores; with 16 threadsmimi_decodeis dominated by threadpool overhead (5.9 ms per frame without the rewrite, against 3.0 ms with 4 threads), so the gain there is smaller.Review order:
_passes/decompose_conv_transpose.pyand its registration in_passes/__init__.py, then the tests, then the Phonon example.Test plan
The new tests in
backends/xnnpack/test/passes/test_decompose_conv_transpose.pycompare the rewritten graph and the lowered program against eager, and check that the pass ran during lowering and that the serialized XNNPACK graph has noXNNConvTranspose2dleft. They cover 1-D with kernel = 2 * stride at strides 2, 4, 5 and 6 with and without bias, a single input step, padding and output_padding, kernels of one and three strides, 2-D with symmetric, asymmetric and no padding, fp16, two deconvs with an ELU between them, a weight shared by two deconvs, a bias that another op also reads (the overlap-add of a streaming decoder), the pass order, external constants, and the cases that are left alone. The Phonon tests lower the small random model's decoder through the pass and compare it with the reference implementation.lintrunnerreports no issues on the changed files. Tested on macOS arm64 with pybindings andphonon_runnerbuilt from this branch.This PR was authored with Claude Code (Claude Opus 5.5).