[Vulkan] reduce shader returns early before barrier(), hanging the GPU (VK_ERROR_DEVICE_LOST on Mali-G76)

[msluszniak]

🐛 Describe the bug

backends/vulkan/runtime/graph/ops/glsl/reduce.glsl bounds-checks the global position in main() and returns early:

if (any(greaterThanEqual(scan_pos, tin_limits))) {
  return;
}

if (reduce_dim != packed_dim) {
  reduce_nonpacked_dim(tid, scan_pos);
} else {
  reduce_packed_dim(tid, scan_pos);
}

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 any invocation that takes the early return 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 in reduce_gwg (impl/Reduce.cpp):

constexpr uint32_t nworkers_per_group = 4u;
constexpr uint32_t ngroups = 4u;
utils::uvec3 lwg_extents{1u, 1u, 1u};
lwg_extents[reduce_dim_whcn] = nworkers_per_group;
lwg_extents[group_dim_whcn] = ngroups;

but group_dim is picked as the larger of the two non-reduce dims of the output, which is frequently 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 and the remaining 4 hang.

Reproduction

import torch, torch.nn as nn
class M(nn.Module):
    def forward(self, x): return torch.sum(x, dim=1, keepdim=True)
# lower M with VulkanPartitioner, run on device with input (1, 384)

Measured on a Samsung Galaxy S10+ (Mali-G76), torch.sum(x, dim=1, keepdim=True) over (b, 384):

b	result
1	VK_ERROR_DEVICE_LOST
2	VK_ERROR_DEVICE_LOST
3	passes
4	passes
5	VK_ERROR_DEVICE_LOST
8	passes

The pattern is not a clean cutoff because divergent barrier() is undefined rather than specified to hang, so whether the driver survives depends on which invocations vanish. torch.mean(dim=1), torch.linalg.vector_norm(dim=1) and torch.sqrt(torch.sum(x*x, dim=1)) all reproduce it identically. Reducing dim 0, or reducing a 3D tensor along dim 2, does not, because the group dim then has extent >= 4.

Impact

Any model ending in an L2 normalization over a vector hits this. sentence-transformers/all-MiniLM-L6-v2 finishes with F.normalize(x, p=2, dim=1) on a (1, 384) tensor, so it loses the device on the very first execution. I bisected it by exporting progressive prefixes: all 6 BERT encoder layers run correctly, masked sum and mean pooling run correctly, and adding the final normalize is what loses the device.

Versions

• ExecuTorch main (c27baa8031)
• Device: Samsung Galaxy S10+ (SM-G975F), Mali-G76, Android 12, Vulkan driver v1.r32p1

Fix

Carry the bounds check as a flag instead of returning, keeping barrier() in uniform control flow. Out of bounds invocations skip the loads and 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, since within a work group the bounds check varies only along group_dim, which is exactly tid.y, and each group aggregates only its own slots.

With that change every b above passes and matches the CPU reference, and MiniLM runs with cosine 0.99999720, bit-identical across 10 executions. No regression on the models I have on hand: the Whisper-tiny encoder is unchanged at cosine 0.99999702 and selfie segmentation is bit-exact.

PR follows.

cc @SS-JIA @manuelcandales @digantdesai @cbilgin

---

Update: this affects four shaders, not one

I audited every shader in backends/vulkan/runtime/graph/ops/glsl that calls barrier() (26 of them) and checked, for each return in main(), whether the guard is uniform across the work group. Three more shaders have the same defect as reduce.glsl:

shader	dispatch	why it diverges
reduce.glsl	reduce_gwg	4 groups along a group_dim whose extent can be 1
reduce2d.glsl	reduce_gwg (same function)	same
var_texture3d.glsl	var_texture_gwg	same construction, lwg_extents[group_dim] = 4u
softmax.glsl	pick_softmax_gwg, texture path	same construction

All four now carry the bounds check as a flag rather than returning.

The rest are correct

The other barrier-using shaders were checked and their guards are uniform across the work group:

• coopmat_mm.glsl tests gl_WorkGroupID, which is uniform by construction. Its comment already notes the branch cannot trigger under the current dispatch.
• reduce_per_row_buffer.glsl, native_layer_norm_buffer.glsl, rms_norm_buffer.glsl, fused_ce.glsl test a global id component whose local size is 1 ({1, rows, 1} global with (64, 1, 1) local), so the guard is uniform.
• softmax_buffer.glsl zeroes reduce_dim before the bounds check, and reduce_dim is the only dim with local size above 1.
• sdpa_attn_weights_softmax.glsl, sdpa_compute_attn_weights_coop.glsl, sdpa_compute_out_coop.glsl guard on gl_GlobalInvocationID.x and .z only, and the varying dimension for each is not in the guard.
• linear_q4gsw_coop.glsl dispatches (1, 1, 64) local, so the linear_idx_from_gid() used in its guard is uniform.
• q4gsw_linear_gemv_coop__w_4x8.glsl returns only when the entire work group is out of bounds, and says so in a comment.

One I could not settle

quantize_and_pack_4h4w_with_group_sums.glsl guards on gl_GlobalInvocationID.x while dispatching (4, 1, 16) or (2, 1, 32) local, so .x varies within the work group and the guard looks divergent ahead of two barrier() calls. I have left it alone: it is an int4 quantized path I have no model on hand to exercise, and I would rather not ship an untested change to it. Worth a look from someone who can run that path.

Verification of the three added here

Galaxy S26 Ultra (Adreno 840) and Galaxy S10+ (Mali-G76), before and after:

• all-MiniLM-L6-v2, whose attention exercises softmax: unchanged at cosine 0.99999243, 1 distinct output over 20 executions.
• Whisper-tiny encoder: no change in behaviour. It remains nondeterministic on Adreno, but that is [Vulkan][Adreno] fp32 nn.Linear returns wrong results on ~23% of executions at large M #22327 (the fp32 linear path), not this. Its best replay still reaches cosine 0.999989.