[ET-VK] Avoid the near-square local work group on Adreno

[msluszniak]

Fixes #22327.

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.

[pytorch-bot]

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

⚠️ 16 Awaiting Approval

As of commit b31b709 with merge base c27baa8 ():

AWAITING APPROVAL - The following workflows need approval before CI can run:

• Build documentation (gh)
• Build Presets (gh)
• Cadence Build & Test (gh)
• Lint (gh)
• periodic (gh)
• pull (gh)
• Test ARM Backend (gh)
• Test CoreML Backend (gh)
• Test Cortex-M Backend (gh)
• Test NXP Backend (gh)
• Test QNN Backend (gh)
• Test QNN Windows MSVC build (gh)
• Test Vulkan Backend (gh)
• Test Vulkan Backend (specialized runners) (gh)
• Test WebGPU Backend (gh)
• Test XNNPACK Backend (gh)

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

[github-actions]

This PR needs a release notes: label

If your change should be included in the release notes (i.e. would users of this library care about this change?), please use a label starting with release notes:. This helps us keep track and include your important work in the next release notes.

To add a label, you can comment to pytorchbot, for example
@pytorchbot label "release notes: none"

For more information, see
https://github.com/pytorch/pytorch/wiki/PyTorch-AutoLabel-Bot#why-categorize-for-release-notes-and-how-does-it-work.