[ET-VK] Avoid the near-square local work group on Adreno by msluszniak · Pull Request #22328 · pytorch/executorch

msluszniak · 2026-08-30T11:08:54Z

The tiled GEMM family of shaders produces wrong results on Adreno for some local work group shapes. These shaders have no shared memory, no barriers, and fully independent invocations, so their output cannot legitimately depend on the group shape, but it does.

At M = 1500, K = N = 384 (a Whisper-tiny encoder projection) with the default pick_xy_square_lwg shape of 8x8x1, roughly 10 to 20 percent of executions come back with a handful of wholly corrupted output blocks. The corruption granularity is exactly one work group: tile_m = 4 and div_up_4(N) means each invocation writes a 4x4 output patch, so an 8x8 group covers 32x32. One failing run had five row runs of exactly 32 (work group y 5, 7, 9, 11, 13, every other one) crossed with columns 128 to 255 (work group x 4 to 7). Everything outside those 20 work groups was correct to 5e-07.

Sweeping only the group shape, holding the same 64 invocations and the same dispatch, 60 executions each on a Snapdragon SM8850 (Adreno 840):

local work group	wrong executions
8x8x1 (current default)	7 / 60
4x16x1	10 / 60
16x4x1	0 / 60
32x2x1	0 / 60
64x1x1	0 / 60
2x32x1	0 / 60
1x64x1	0 / 60

Only the two near-square shapes fail. 2x32x1 passing while 4x16x1 fails rules out a simple "wide is good" reading.

This is not specific to linear. pick_xy_square_lwg is shared by thirteen dispatches across Linear, Matmul, Conv2dGemm, Conv1dPW, Conv2dPW, SDPA, Q8taLinear, Q8taConv2dPW, QuantizedLinear and QuantizedConvolution, and conv2d_gemm reproduces it independently (see below). Fixing the shared helper covers all of them.

On Adreno it now picks a 2:1 x:y LwgShape, which the existing D'Hondt allocation turns into 16x4x1 at the default 64 threads and which scales with recommended_lwg_nthreads(). Other vendors keep the square shape.

Cost

Whisper-tiny encoder, same device, median of 60 timed iterations, sweeping only the linear dispatch:

shape	time	vs 8x8x1
8x8x1	150.24 ms
16x4x1	152.55 ms	+1.5%
32x2x1	158.36 ms	+5.4%
64x1x1	200.93 ms	+33.7%
2x32x1	212.64 ms	+41.5%

16x4x1 is the cheapest correct shape by a wide margin. Applied to all thirteen dispatches the whole-encoder cost is +2.9% to +5.3% depending on the variant, which buys a model that is actually correct.

Verification

Distinct outputs over 60 executions with --dump_every_execution, Adreno 840:

model	before	after
Whisper-tiny encoder, conv1d frontend	60 / 60	1 / 60
Whisper-tiny encoder, conv2d im2col frontend (#22330)	9 / 60	1 / 60
Whisper-tiny encoder, conv2d frontend and rank-3 attention	9 / 60	1 / 60
Whisper conv frontend alone, im2col + GEMM	13 / 60	1 / 60

Cosine against XNNPACK on the same input is 0.99999726 after the change; before it fell as low as 0.72.

The last two rows matter for scoping: an earlier revision of this PR patched only Linear.cpp, and the whisper encoder looked clean at 40 executions. It was not. Once the conv frontend moved onto the im2col + GEMM path, conv2d_gemm reproduced the same defect on its own, which is what showed the fix belonged in the shared helper rather than in one caller.

I only have Adreno hardware for this, so the change is gated on device_is_adreno(). If you would rather have one shape everywhere I am happy to drop the gate, but I cannot measure the perf effect on Mali or desktop myself.

cc @SS-JIA @manuelcandales @digantdesai @cbilgin @psiddh @nil-is-all

pytorch-bot · 2026-08-30T11:08:58Z

🔗 Helpful Links

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

📄 Preview Python docs built from this PR

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

❌ 1 New Failure

As of commit 949f806 with merge base f0e8589 ():

NEW FAILURE - The following job has failed:

pull / test-models-linux (mobilebert, portable, mt-l-x86iavx512-8-64) / linux-job (gh)
RuntimeError: Command bash /__w/_temp/exec_script failed with exit code 1

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

Adreno drivers miscompute the tiled GEMM family of shaders for 8x8x1 and 4x16x1 local work groups: entire work groups intermittently write garbage while every other block is bit-exact. These shaders have no shared memory and no barriers and their invocations are independent, so the result cannot legitimately depend on the group shape. Fix it in pick_xy_square_lwg, which all thirteen affected dispatches share: linear, matmul, conv2d_gemm, the pointwise convs, SDPA and the quantized variants. On Adreno use a 2:1 x:y shape, which yields 16x4x1 at the default 64 threads and is the cheapest correct shape measured; other vendors keep the square shape. Fixes pytorch#22327

Adreno drivers miscompute the tiled GEMM family for 8x8x1 and 4x16x1 local work groups: entire work groups intermittently write garbage while every other block is bit-exact. pick_hw_square_wg_size is shared by Linear, Matmul, Conv2dGemm and the pointwise convs, so all of them are affected. Use 16x4x1 on Adreno, which is correct across every run measured and is the cheapest correct shape. Upstream: pytorch/executorch#22327, pytorch/executorch#22328

SS-JIA

AI-assisted review. Inline findings below.

SS-JIA · 2026-09-09T19:13:09Z

  (void)resize_args;
  LocalWorkGroup lwg(
-      kSquareLwg, graph->context()->adapter_ptr()->recommended_lwg_nthreads());
+      graph->device_is_adreno() ? kWideXLwg : kSquareLwg,


[High] The Adreno shape policy is applied before fit_to_global(). For narrow-X grids, fitting can turn the chosen 16x4 group back into 8x8 or 4x16—the exact broken shapes. Reapply the policy after fitting or make fitting preserve it, and add picker tests for those grids.

AI reviewer note: this comment was generated by an AI reviewer.

Fixed in b90bad2: the shape is reapplied after fitting, so a fitted 8x8/4x16 becomes 16x4 or 2x32, with a narrow-X picker test.

SS-JIA

This appears to be another compiler bug specific to Adreno 840. I'd like hold off on implementing this workaround for now. I'll pass it along to Qualcomm to determine if it might be something that has already been fixed in most recent drivers.

msluszniak · 2026-09-09T21:06:30Z

Yeah, sounds like potential problems with drivers.

msluszniak · 2026-09-14T16:44:26Z

@pytorchbot label "module: vulkan"

@SS-JIA

…22330) Fixes #22329. `conv1d.glsl` computes one output element per invocation with no tiling, no shared memory and no register blocking, and every tap re-fetches both a weight texel and an input texel, so nothing is reused between invocations. On top of that `conv1d_gwg` sets `z = div_up_4(batch)`, so the texel's four lanes span the batch dimension and three of every four are dead at batch 1. Measured inside a single Whisper-tiny encoder execution on an Adreno 840, using the shader query pool: | dispatch | useful throughput | | --- | --- | | `conv1d_float`, 80 -> 384, k=3, L=3000 | 31.3 GFLOP/s | | `conv1d_float`, 384 -> 384, k=3, s=2, L=1500 | 27.6 GFLOP/s | | `linear_vec_bias`, M=1500, N=1536, K=384 | 921 GFLOP/s | | `linear_vec_bias`, M=1500, N=384, K=1536 | 900 GFLOP/s | Same GPU, same run. Those two conv1d dispatches were 43% of the encoder's runtime. `conv2d_im2col.glsl` + `conv2d_gemm.glsl` already solve this for 2-D convolutions. This pass rewrites 1-D convolutions into a 2-D convolution over a singleton height dim so the existing conv2d machinery (im2col selection, memory layout tagging, weight prepacking) handles them with no runtime changes. The weight is reshaped in place through `set_param_tensor` rather than with a `view_copy` node, because `add_conv2d_node` reads it as a constant `TensorRef` and a runtime view would hand it a `TENSOR` instead. ### Eligibility Only convolutions that are certain to reach the im2col path are rewritten, so nothing silently moves onto a shader that has not been compared against conv1d: - `groups == 1` and unit dilation, matching `should_use_conv2d_im2col` - `out_channels >= kIm2colMinCOut` (128), so the im2col path is selected on every vendor rather than only on Mali - `kernel_size > 1`, leaving pointwise convs on `conv1d_pw` - batch 1, since conv2d throws on batched input and a rewrite would turn a working conv1d into a hard error - a constant weight, so it can be reshaped in place ### Results Whisper-tiny encoder, Samsung Galaxy S26 Ultra (Adreno 840), fp32, interleaved A/B with the order reversed each round, warm-up iteration discarded, 60 timed iterations per arm. Both Vulkan arms include #22328, without which the encoder is not deterministic on this GPU at all. | arm | median | vs XNNPACK | | --- | --- | --- | | Vulkan before | 153.2 ms | 0.84x | | XNNPACK | 128.0 ms | 1.00x | | Vulkan after | **88.8 ms** | **1.44x** | Cosine against the CPU reference is 0.99999684 after the change. Isolated, just the encoder's two convs plus gelu, GPU time from the query pool: | | before | after | | --- | --- | --- | | conv dispatches | conv1d 17.7 + 47.7 ms | im2col 0.05 + gemm 0.64 + im2col 0.19 + gemm 1.47 ms | | total | 66.5 ms | 3.28 ms | 20x on the frontend, and the output is unchanged at cosine 1.00000000, max abs diff 4.9e-04 against the CPU reference in both. ### Routing and correctness matrix Eight conv1d configurations exported and run on device, with the dispatched shader read back from the query pool: | case | shader after the pass | rewritten | cosine | max abs diff | | --- | --- | --- | --- | --- | | 80 -> 384, k=3 | `conv2d_im2col` + `conv2d_gemm` | yes | 1.00000000 | 2.0e-06 | | 384 -> 384, k=3, s=2 | `conv2d_im2col` + `conv2d_gemm` | yes | 1.00000000 | 4.1e-06 | | 128 -> 256, k=5, p=0 | `conv2d_im2col` + `conv2d_gemm` | yes | 1.00000012 | 3.1e-06 | | 64 -> 64, k=3 | `conv1d_float` | no | 1.00000000 | 8.3e-07 | | 256 -> 256, k=1 | `conv1d_pw_bias_texture` | no | 1.00000000 | 9.5e-07 | | 256 -> 256, k=3, depthwise | `conv1d_dw_bias_texture` | no | 1.00000012 | 4.8e-07 | | 128 -> 256, k=3, dilation 2 | `conv1d_float` | no | 1.00000000 | 1.2e-06 | | 256 -> 256, k=3, groups 4 | `conv1d_float` | no | 1.00000000 | 9.5e-07 | Every ineligible case keeps its original shader and every case matches the CPU reference. `test_vulkan_backend_conv1d_as_conv2d` covers the rewritten path in CI; the two existing conv1d tests are grouped and depthwise, so they only exercise the untouched path. --- ### Depends on #22328 Moving the conv frontend onto the im2col + GEMM path exposes a pre-existing Adreno defect in `conv2d_gemm`, which dispatches with the same near-square local work group as `linear`. Distinct outputs over 60 executions on the Adreno 840, with `Linear.cpp` already fixed: | | distinct / 60 | | --- | --- | | Whisper encoder, conv1d frontend | 1 | | Whisper encoder, this PR | 9 | | Whisper conv frontend alone, this PR | 13 | This is not caused by the rewrite. `conv2d_gemm` is already reached today by any conv2d with `out_channels >= 128`, and it is nondeterministic there for the same reason. #22328 now fixes the shared `pick_xy_square_lwg` helper rather than one caller, and with it every model above is 1 / 60. The 88.8 ms figure above was measured with only `Linear.cpp` fixed. With #22328 applied in full the encoder is 91.4 ms, still 1.34x XNNPACK, and correct on every run. ### Why rewrite rather than write a tiled conv1d shader The 384 -> 384 dispatch moves 21.2 GB of texture traffic for 8.7 MB of compulsory traffic, so the shader is texture bound with essentially no reuse, and the batch-in-lanes waste is a 4x term next to that. A native tiled conv1d would fix the reuse the same way the GEMM does, and would additionally save the im2col materialisation, which is 0.25 ms of the 2.64 ms frontend above, roughly 10%. This pass gets the rest of it by reusing kernels that are already written, tuned and tested, and leaves a dedicated conv1d kernel available later as an optimisation rather than a prerequisite. cc @SS-JIA @manuelcandales @digantdesai @cbilgin

msluszniak · 2026-09-16T09:31:55Z

@nil-is-all I think it will be valuable to pass this one to Qualcomm team as Stephen suggested.

psiddh · 2026-10-07T23:53:11Z

@nil-is-all I think it will be valuable to pass this one to Qualcomm team as Stephen suggested.

cc @qti-horodnic @shewu-quic @winskuo-quic
Tagging few Qualcomm folks

qti-horodnic

@msluszniak
Thanks for the thorough investigation, I'm fine with this workaround. On narrow-X grids fit_to_global() can still refit to 8x8 or 4x16 (+1 to the comment above). Please reapply the shape after fitting, and a small picker test would be great to include too.

Two minor, non-blocking suggestions: could the code comment mention that the fix was only measured on Adreno 840? And since pick_xy_square_lwg now returns a non-square shape on Adreno, it may be worth renaming it or noting that in its comment.

Output that depends on work group shape for fully independent invocations looks like a driver or shader compiler issue, so I'll cut a ticket to the relevant team. I wouldn't block the PR on this, if a driver fix lands we can remove this workaround later.

msluszniak · 2026-10-19T22:07:42Z

Ok, I will address your and Stephen suggestions tomorrow :))

fit_to_global() caps X on narrow grids and moves the spare threads to Y, which turned the 2:1 group back into 8x8x1 (X < 14) or 4x16x1 (X < 6). Reapply the policy after fitting: if the fitted group is one of those, use whichever of the unfitted 2:1 group or 2x32x1 dispatches fewer invocations. Split the Adreno branch into pick_adreno_xy_lwg so it is testable without a device, add a picker test for narrow-X grids, note that the fix was only measured on Adreno 840, and document the non-square shape on pick_xy_square_lwg.

msluszniak · 2026-10-20T12:50:59Z

@qti-horodnic Done in b90bad2: reapplied after fitting with a picker test, noted Adreno 840 in the comment, and documented the non-square shape on pick_xy_square_lwg (kept the name, it has 16 call sites).

qti-horodnic

LGTM, thanks for the change

2x32x1 is correct but slower than the unfitted 2:1 group even when the grid leaves X lanes idle: on an S26 (Adreno 840) multi-qa-MiniLM-L6 runs 34.6 ms with it and 23.0 ms with 16x4x1, with bit-identical output.

nil-is-all · 2026-11-09T16:58:36Z

@msluszniak updated the branch and running CI tests

msluszniak requested a review from SS-JIA as a code owner August 30, 2026 11:08

meta-cla Bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Aug 30, 2026

msluszniak mentioned this pull request Aug 30, 2026

[ET-VK] Lower eligible conv1d as conv2d over a singleton height dim #22330

Merged

msluszniak force-pushed the ms/vulkan-linear-lwg-adreno branch from 4ad33b9 to b31b709 Compare August 30, 2026 13:12

msluszniak changed the title ~~[ET-VK] Avoid the near-square local work group for the tiled linear on Adreno~~ [ET-VK] Avoid the near-square local work group on Adreno Aug 30, 2026

msluszniak mentioned this pull request Aug 30, 2026

[ET-VK][Adreno] fp32 nn.Linear returns wrong results on ~23% of executions at large M #22327

Open

SS-JIA reviewed Sep 9, 2026

View reviewed changes

SS-JIA requested changes Sep 9, 2026

View reviewed changes

pytorch-bot Bot added the module: vulkan Issues related to the Vulkan delegate and code under backends/vulkan/ label Sep 14, 2026

nil-is-all added the release notes: vulkan Changes to the Vulkan backend delegate label Sep 14, 2026

executorch-triage Bot added the community: contribution PRs coming from community (excluding hardware partners) label Sep 22, 2026

psiddh added the module: qnn Issues related to Qualcomm's QNN delegate and code under backends/qualcomm/ label Oct 7, 2026

qti-horodnic suggested changes Oct 19, 2026

View reviewed changes

qti-horodnic approved these changes Oct 20, 2026

View reviewed changes

Vulkan LWG: keep 16x4x1 on narrow Adreno grids

2d136bc

2x32x1 is correct but slower than the unfitted 2:1 group even when the grid leaves X lanes idle: on an S26 (Adreno 840) multi-qa-MiniLM-L6 runs 34.6 ms with it and 23.0 ms with 16x4x1, with bit-identical output.

msluszniak mentioned this pull request Nov 4, 2026

[ET-VK] Keep 16x4x1 on narrow Adreno grids software-mansion-labs/executorch#31

Merged

Merge branch 'main' into ms/vulkan-linear-lwg-adreno

949f806

Conversation

msluszniak commented Aug 30, 2026 • edited by pytorch-bot Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Cost

Verification

Uh oh!

pytorch-bot Bot commented Aug 30, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

🔗 Helpful Links

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

❌ 1 New Failure

Uh oh!

SS-JIA left a comment

Choose a reason for hiding this comment

Uh oh!

SS-JIA Sep 9, 2026

Choose a reason for hiding this comment

Uh oh!

qti-horodnic Oct 19, 2026

Choose a reason for hiding this comment

Uh oh!

msluszniak Oct 20, 2026

Choose a reason for hiding this comment

Uh oh!

SS-JIA left a comment

Choose a reason for hiding this comment

Uh oh!

msluszniak commented Sep 9, 2026

Uh oh!

msluszniak commented Sep 14, 2026

Uh oh!

msluszniak commented Sep 16, 2026

Uh oh!

psiddh commented Oct 7, 2026

Uh oh!

qti-horodnic left a comment

Choose a reason for hiding this comment

Uh oh!

msluszniak commented Oct 19, 2026

Uh oh!

msluszniak commented Oct 20, 2026

Uh oh!

qti-horodnic left a comment

Choose a reason for hiding this comment

Uh oh!

nil-is-all commented Nov 9, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

This branch has not been deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

msluszniak commented Aug 30, 2026 •

edited by pytorch-bot Bot

Loading

pytorch-bot Bot commented Aug 30, 2026 •

edited

Loading

nil-is-all commented Nov 9, 2026 •

edited

Loading