[ET-VK] Do not return early before the barrier in the reduce shader - #22326
Open
msluszniak wants to merge 2 commits into
Open
[ET-VK] Do not return early before the barrier in the reduce shader#22326msluszniak wants to merge 2 commits into
msluszniak wants to merge 2 commits into
Conversation
main() bounds-checks the global position and returns early:
if (any(greaterThanEqual(scan_pos, tin_limits))) {
return;
}
Both reduce_nonpacked_dim() and reduce_packed_dim() then call barrier().
Vulkan requires barrier() to be reached by every invocation in the work
group under uniform control flow, so an invocation that returns leaves
the remaining ones waiting on a barrier that can never complete. On a
Mali-G76 this hangs the GPU and the submit fails with
VK_ERROR_DEVICE_LOST.
The work group is always sized with ngroups = 4 along group_dim, but
group_dim is picked as the larger of the two non-reduce dims of the
output, which can be smaller than 4. Reducing a 2D tensor along dim 1
produces an output whose two candidate group dims both have extent 1,
so 12 of the 16 invocations return early.
Carry the bounds check as a flag instead of returning, so the barrier
stays in uniform control flow. Out of bounds invocations skip the loads
and the accumulation, still write their (unused) shared memory slot,
reach the barrier, and skip the output write. Their shared memory
contents are never read by an in bounds group, because within a work
group the bounds check varies only along group_dim, which is exactly
tid.y, and a group aggregates only its own slots.
Verified on a Mali-G76: torch.sum(x, dim=1, keepdim=True) over (b, 384)
lost the device for b = 1, 2 and 5 and now passes for every b, matching
the CPU reference. all-MiniLM-L6-v2, whose final F.normalize reduces a
(1, 384) tensor, previously lost the device on load and now runs with
cosine 0.99999720 against the CPU reference, bit-identical across 10
executions.
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/22326
Note: Links to docs will display an error until the docs builds have been completed.
|
meta-cla
Bot
added
the
CLA Signed
This PR needs a
|
msluszniak
added a commit
to software-mansion-labs/executorch
that referenced
this pull request
Aug 30, 2026
main() bounds-checks the global position and returns early, but both reduce_nonpacked_dim() and reduce_packed_dim() call barrier(). Vulkan requires barrier() to be reached by every invocation in the work group under uniform control flow, so an invocation that returns leaves the rest waiting on a barrier that can never complete. On a Mali-G76 this hangs the GPU and the submit fails with VK_ERROR_DEVICE_LOST. The work group is always sized with ngroups = 4 along group_dim, but group_dim is the larger of the two non-reduce dims of the output and is often smaller than 4. Reducing a 2D tensor along dim 1 gives an output whose two candidate group dims both have extent 1, so 12 of the 16 invocations return early. Carry the bounds check as a flag instead. all-MiniLM-L6-v2 ends in F.normalize over a (1, 384) tensor and previously lost the device on its first execution. Backport of upstream pytorch/executorch#22326.
Auditing every shader that calls barrier() turned up three more with the same defect as reduce.glsl: main() bounds-checks the global position and returns early, while the routine it calls contains a barrier(). All three are dispatched through the same style of work group sizing, which sets 4 groups along a group_dim whose extent can be 1, so invocations diverge and the survivors wait on a barrier that never completes. softmax.glsl pick_softmax_gwg, texture path var_texture3d.glsl var_texture_gwg reduce2d.glsl reduce_gwg, the same function reduce.glsl uses Each now carries the bounds check as a flag instead of returning, as reduce.glsl already does. The remaining barrier-using shaders were checked and are correct: their guards are uniform across the work group, either because they test gl_WorkGroupID (coopmat_mm) or because they test a global id component whose local size is 1 (reduce_per_row_buffer, softmax_buffer, rms_norm_buffer, native_layer_norm_buffer, fused_ce, the three sdpa shaders, linear_q4gsw_coop). q4gsw_linear_gemv_coop__w_4x8 returns only when the whole work group is out of bounds and says so in a comment.
msluszniak
added a commit
to software-mansion-labs/executorch
that referenced
this pull request
Aug 30, 2026
reduce2d, var_texture3d and softmax build their work group the same way reduce.glsl does, with 4 groups along a group_dim whose extent can be 1, so they return early from a subset of invocations before a barrier(). Carry the bounds check as a flag instead, as already done for reduce.glsl. Upstream: pytorch/executorch#22325, pytorch/executorch#22326
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
2 participants
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.
Fixes #22325.
main()inreduce.glslbounds-checks the global position and returns early, but bothreduce_nonpacked_dim()andreduce_packed_dim()callbarrier(). Vulkan requiresbarrier()to be reached by every invocation in the work group under uniform control flow, so an invocation that returns leaves the remaining ones waiting on a barrier that can never complete. On a Mali-G76 that hangs the GPU and the submit fails withVK_ERROR_DEVICE_LOST.The work group is always sized with
ngroups = 4alonggroup_dim, butgroup_dimis the larger of the two non-reduce dims of the output and is often smaller than 4. Reducing a 2D tensor along dim 1 gives an output whose two candidate group dims both have extent 1, so 12 of 16 invocations return early.This carries the bounds check as a flag instead. Out of bounds invocations skip the loads and the accumulation, still write their unused shared memory slot, reach the barrier, and skip the output write. Their shared memory contents are never read by an in bounds group: within a work group the bounds check varies only along
group_dim, which is exactlytid.y, and a group aggregates only its own slots.Verification
Samsung Galaxy S10+ (Mali-G76),
torch.sum(x, dim=1, keepdim=True)over(b, 384):sentence-transformers/all-MiniLM-L6-v2ends inF.normalize(x, p=2, dim=1)over a(1, 384)tensor and previously lost the device on its first execution. It now runs with cosine 0.99999720 against the CPU reference, bit-identical across 10 executions.No regression on the other models I have on device: the Whisper-tiny encoder is unchanged at cosine 0.99999702, and selfie segmentation stays bit-exact at 1.19e-07.
Update: extended to every shader with the same defect
I audited all 26 shaders that call
barrier()and checked, for eachreturninmain(), whether the guard is uniform across the work group. Three more had the same defect and are now fixed the same way:reduce2d.glslreduce_gwg, the same functionreduce.glslusesgroup_dimwhose extent can be 1var_texture3d.glslvar_texture_gwglwg_extents[group_dim] = 4usoftmax.glslpick_softmax_gwg, texture pathThe remaining barrier-using shaders are correct and are left alone. Their guards are uniform across the work group, either because they test
gl_WorkGroupID(coopmat_mm), because they test a global id component whose local size is 1 (reduce_per_row_buffer,native_layer_norm_buffer,rms_norm_buffer,fused_ce,linear_q4gsw_coop), because the tested dimension is zeroed first (softmax_buffer), or because the varying dimension is not in the guard (the threesdpashaders).q4gsw_linear_gemv_coop__w_4x8returns only when the whole work group is out of bounds.One shader is deliberately not touched:
quantize_and_pack_4h4w_with_group_sums.glslguards ongl_GlobalInvocationID.xwhile dispatching(4, 1, 16)or(2, 1, 32)local, so its guard looks divergent ahead of two barriers. It is an int4 quantized path I have no model to exercise, and I would rather not ship an untested change to it. Flagged in the issue.Verification of the three added here
On a Galaxy S26 Ultra (Adreno 840):
all-MiniLM-L6-v2, whose attention exercisessoftmax: unchanged at cosine 0.99999243 against the CPU reference, 1 distinct output over 20 executions.linearpath, unrelated to this; its best replay still reaches cosine 0.999989.