Repository navigation
[ET-VK][pixel_shuffle] Rewrite q8ta_pixel_shuffle to write whole output words - #24457
Merged
Merged
Conversation
…ut words
## Ulterior Motive
Reduce the inference latency of int8 segmentation CNNs that upsample with pixel shuffle on mobile GPUs. Per-op profiling of such a 512x512 model found `q8ta_pixel_shuffle` to be its most expensive operator, running far below memory bandwidth.
## Rationale
**What**: Rewrites `q8ta_pixel_shuffle` so that each invocation turns the 4 input words holding 16 channels at one input pixel into 4 full output words with a 4x4 byte transpose. The global work grid and local workgroup size change to match. The test sets add the `PACKED_INT8_4C` layout and every pair of channels-packed layouts, and check equal-qparams cases bit-exactly.
**Why**: The old shader ran one invocation per output word, with output width as the fastest-varying index and the channel block as the slowest. Channels-packed int8x4 layouts store channels innermost, so neighbouring invocations wrote words 64 B apart for a `PACKED_INT8_4C` output, or 16 B apart for `PACKED_INT8_4W4C`. Each 64 B output line was therefore completed over 16 (4C) or 4 (4W4C) separate passes over the whole tensor. Once the output no longer fits in cache, every pass writes partial lines back to memory. In addition, the 4 lanes of an output word come from 4 different input words, so each invocation issued 4 loads to produce 4 bytes.
On Adreno 830, the old shader took the following for [1,256,256,256] -> [1,64,512,512]:
| Layouts (in -> out) | Time |
| --- | --- |
| 4W4C -> 4C | 10.65 ms |
| 4W4C -> 4W4C | 2.93 ms |
| 4C1W -> 4C1W (width innermost, which matches the old thread order) | 0.80 ms |
Its 4W4C -> 4C time grew faster than the data size: 0.26 / 2.0 / 10.7 ms at 64^2 / 128^2 / 256^2 input planes. The new shader takes 0.55 ms for that 4W4C -> 4C shape (19x faster) and is 2.8-4.1x faster on smaller shapes.
## Details
### Shader
With upscale factor 2, `out[n, c, 2h + dy, 2w + dx] = in[n, 4c + 2dy + dx, h, w]`. Each invocation handles output channel block `k` at one input pixel `(h, w)`. Input word `i` holds input channels `16k + 4i` to `16k + 4i + 3`, which are output channel `4k + i` at the 4 output pixels of the 2x2 block. The 4 input words form a 4x4 byte matrix whose transpose is the 4 output words:
```
input word i = [ out ch 4k+i at (dy,dx) = (0,0) | (0,1) | (1,0) | (1,1) ]
| transpose_int8x4_block
v
output word j = [ out ch 4k | 4k+1 | 4k+2 | 4k+3 ] at pixel (2h + j/2, 2w + j%2)
```
Every loaded and stored byte is used exactly once. The requantize formula is unchanged and is applied per word before the transpose. `passthrough` (input and output qparams match) is now a specialization constant, so the requantize path compiles out in that case. The `upscale_factor` push constant is removed; the dispatch already rejects any factor other than 2.
### Dispatch
The global work grid is (output channel blocks, input width, input height x batch). The local workgroup shape depends on the layouts, so that one workgroup's stores cover contiguous memory. Values below are for 64 output channels and were chosen from sweeps on Adreno 830:
- `PACKED_INT8_4C` output: channel blocks first (16x4x1), so 16 invocations write one 64 B line.
- `PACKED_INT8_4C1W` on both sides: width (1x64x1).
- Other layout pairs: 4x16x1.
### Tests and benchmark sets
- The `PACKED_INT8_4C` layout is added. Correctness covers all 9 pairs of `PACKED_INT8_4W4C`, `PACKED_INT8_4C` and `PACKED_INT8_4C1W`.
- New correctness shapes: C_out = 1 (three padding lanes) and batch 2.
- Equal-qparams cases use integer-valued inputs with scale 1 and zero tolerance, so they check bit-exactness. Requantize cases keep a one-step tolerance: the GPU multiplies by the reciprocal scale and GLSL `round()` breaks ties in an implementation-defined direction, while the CPU reference divides.
- The performance set has 81 cases (9 layout pairs x 9 shapes), including decoder-style [1,256,H,W] inputs with H = W = 16 to 256.
Suggested review order: `q8ta_pixel_shuffle.glsl`, `Q8taPixelShuffle.cpp`, then the test files.
Authored with AI assistance (Claude Code).
Differential Revision: [D129089325](https://our.internmc.facebook.com/intern/diff/D129089325/)
**NOTE FOR REVIEWERS**: This PR has internal Meta-specific changes or comments, please review them on [Phabricator](https://our.internmc.facebook.com/intern/diff/D129089325/)!
ghstack-source-id: 462609407
Pull-Request: #24440
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/24457
Note: Links to docs will display an error until the docs builds have been completed. ⏳ No Failures, 163 PendingAs of commit 1b31697 with merge base 05092c1 ( This comment was automatically generated by Dr. CI and updates every 15 minutes. |
meta-cla
Bot
added
the
CLA Signed
This PR needs a
|
SS-JIA
approved these changes
Nov 2, 2026
This branch was previously 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.
This PR was created by the merge bot to help merge the original PR into the main branch.
ghstack PR number: #24440 by @SS-JIA
^ Please use this as the source of truth for the PR details, comments, and reviews
ghstack PR base: https://github.com/pytorch/executorch/tree/gh/SS-JIA/680/base
ghstack PR head: https://github.com/pytorch/executorch/tree/gh/SS-JIA/680/head
Merge bot PR base: https://github.com/pytorch/executorch/tree/main
Merge bot PR head: https://github.com/pytorch/executorch/tree/gh/SS-JIA/680/orig
Differential Revision: D129089325
@diff-train-skip-merge