[ET-VK] Do not return early before the barrier in softmax and var (#22326)

Fixes #22325.

`main()` in `reduce.glsl` 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 remaining ones waiting on a barrier
that can never complete. On a Mali-G76 that 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 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 exactly `tid.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)`:

| b | before | after |
| --- | --- | --- |
| 1 | device lost | passes, max abs diff 0 |
| 2 | device lost | passes, 1.9e-06 |
| 3 | passes | passes, 3.8e-06 |
| 4 | passes | passes, 2.9e-06 |
| 5 | device lost | passes, 3.8e-06 |
| 8 | passes | passes, 2.9e-06 |

`sentence-transformers/all-MiniLM-L6-v2` ends in `F.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 each
`return` in `main()`, whether the guard is uniform across the work
group. Three more had the same defect and are now fixed the same way:

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

The 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 three `sdpa` shaders). `q4gsw_linear_gemv_coop__w_4x8` returns only
when the whole work group is out of bounds.

One shader is deliberately not touched:
`quantize_and_pack_4h4w_with_group_sums.glsl` guards on
`gl_GlobalInvocationID.x` while 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 exercises `softmax`: unchanged at
cosine 0.99999243 against the CPU reference, 1 distinct output over 20
executions.
- Whisper-tiny encoder: unchanged. It stays nondeterministic on Adreno,
but that is #22327 in the fp32 `linear` path, unrelated to this; its
best replay still reaches cosine 0.999989.


---

## Update: rebased

`reduce.glsl` and `reduce2d.glsl` landed with #22348, so this is now
only `softmax.glsl` and `var_texture3d.glsl`. Rebased on main.


cc @SS-JIA @manuelcandales @digantdesai @cbilgin

Author: msluszniak

backends/vulkan/runtime/graph/ops/glsl/softmax.glsl

@@ -51,15 +51,18 @@ int tid_to_smi(const ivec2 tid) {
  * This case is simpler because each element of a texel belongs to a separate
  * reduction dim, meaning we don't have to perform reduction along a texel.
  */
-void softmax_nonpacked_dim(const ivec2 tid, ivec3 scan_pos) {
+void softmax_nonpacked_dim(const ivec2 tid, ivec3 scan_pos, const bool in_bounds) {
   const int smi = tid_to_smi(tid);
   int group_i;
 
-  scan_pos[reduce_dim] = tid.x;
-  vec4 max_elements = texelFetch(tin, scan_pos, 0);
-  for (int i = tid.x; i < safe_idx(in_meta.sizes, reduce_dim);
-       i += NWORKERS, scan_pos[reduce_dim] += NWORKERS) {
-    max_elements = max(max_elements, texelFetch(tin, scan_pos, 0));
+  vec4 max_elements = vec4(-1.0 / 0.0);
+  if (in_bounds) {
+    scan_pos[reduce_dim] = tid.x;
+    max_elements = texelFetch(tin, scan_pos, 0);
+    for (int i = tid.x; i < safe_idx(in_meta.sizes, reduce_dim);
+         i += NWORKERS, scan_pos[reduce_dim] += NWORKERS) {
+      max_elements = max(max_elements, texelFetch(tin, scan_pos, 0));
+    }
   }
   shared_max[smi] = max_elements;
   barrier();
@@ -69,11 +72,13 @@ void softmax_nonpacked_dim(const ivec2 tid, ivec3 scan_pos) {
     max_elements = max(max_elements, shared_max[group_i]);
   }
 
-  scan_pos[reduce_dim] = tid.x;
   vec4 denominators = vec4(0);
-  for (int i = tid.x; i < safe_idx(in_meta.sizes, reduce_dim);
-       i += NWORKERS, scan_pos[reduce_dim] += NWORKERS) {
-    denominators += exp(texelFetch(tin, scan_pos, 0) - max_elements);
+  if (in_bounds) {
+    scan_pos[reduce_dim] = tid.x;
+    for (int i = tid.x; i < safe_idx(in_meta.sizes, reduce_dim);
+         i += NWORKERS, scan_pos[reduce_dim] += NWORKERS) {
+      denominators += exp(texelFetch(tin, scan_pos, 0) - max_elements);
+    }
   }
   shared_sum[smi] = denominators;
   barrier();
@@ -88,7 +93,8 @@ void softmax_nonpacked_dim(const ivec2 tid, ivec3 scan_pos) {
       scan_pos[packed_dim] == (safe_idx(out_meta.limits, packed_dim) - 1);
 
   scan_pos[reduce_dim] = tid.x;
-  for (int i = tid.x; i < safe_idx(in_meta.sizes, reduce_dim);
+  for (int i = in_bounds ? tid.x : safe_idx(in_meta.sizes, reduce_dim);
+       i < safe_idx(in_meta.sizes, reduce_dim);
        i += NWORKERS, scan_pos[reduce_dim] += NWORKERS) {
     const vec4 numerators = op1(texelFetch(tin, scan_pos, 0) - max_elements);
     const vec4 safe_denom = max(denominators, vec4(1e-37));
@@ -120,7 +126,7 @@ void softmax_nonpacked_dim(const ivec2 tid, ivec3 scan_pos) {
  * elements in texels (which occur when the size of the packed dim is not a
  * multiple of 4) so that they do not influence the output of reduction.
  */
-void softmax_packed_dim(const ivec2 tid, ivec3 scan_pos) {
+void softmax_packed_dim(const ivec2 tid, ivec3 scan_pos, const bool in_bounds) {
   const int smi = tid_to_smi(tid);
   int group_i;
 
@@ -129,11 +135,11 @@ void softmax_packed_dim(const ivec2 tid, ivec3 scan_pos) {
 
   scan_pos[reduce_dim] = tid.x;
   vec4 max_elements = vec4(-3.402823e+38);
-  for (int i = tid.x * 4; i < reduce_len;
+  for (int i = in_bounds ? tid.x * 4 : reduce_len; i < reduce_len;
        i += NWORKERS * 4, scan_pos[reduce_dim] += NWORKERS) {
     max_elements = max(max_elements, texelFetch(tin, scan_pos, 0));
   }
-  if (scan_pos[reduce_dim] == safe_idx(out_meta.limits, reduce_dim) - 1 && nspill > 0) {
+  if (in_bounds && scan_pos[reduce_dim] == safe_idx(out_meta.limits, reduce_dim) - 1 && nspill > 0) {
     const vec4 intex = texelFetch(tin, scan_pos, 0);
     for (int i = 0; i < nspill; ++i) {
       max_elements.x = max(intex[i], max_elements.x);
@@ -153,11 +159,11 @@ void softmax_packed_dim(const ivec2 tid, ivec3 scan_pos) {
 
   scan_pos[reduce_dim] = tid.x;
   vec4 denominators = vec4(0);
-  for (int i = tid.x * 4; i < reduce_len;
+  for (int i = in_bounds ? tid.x * 4 : reduce_len; i < reduce_len;
        i += NWORKERS * 4, scan_pos[reduce_dim] += NWORKERS) {
     denominators += exp(texelFetch(tin, scan_pos, 0) - max_element);
   }
-  if (nspill > 0 && scan_pos[reduce_dim] == safe_idx(out_meta.limits, reduce_dim) - 1) {
+  if (in_bounds && nspill > 0 && scan_pos[reduce_dim] == safe_idx(out_meta.limits, reduce_dim) - 1) {
     const vec4 intex = texelFetch(tin, scan_pos, 0);
     for (int i = 0; i < nspill; ++i) {
       denominators.x += exp(intex[i] - max_element);
@@ -177,12 +183,12 @@ void softmax_packed_dim(const ivec2 tid, ivec3 scan_pos) {
   const float safe_denominator = max(denominator, 1e-37);
 
   scan_pos[reduce_dim] = tid.x;
-  for (int i = tid.x * 4; i < reduce_len;
+  for (int i = in_bounds ? tid.x * 4 : reduce_len; i < reduce_len;
        i += NWORKERS * 4, scan_pos[reduce_dim] += NWORKERS) {
     const vec4 numerators = op1(texelFetch(tin, scan_pos, 0) - max_element);
     imageStore(tout, scan_pos, op2(numerators, safe_denominator));
   }
-  if (nspill > 0 && scan_pos[reduce_dim] == safe_idx(out_meta.limits, reduce_dim) - 1) {
+  if (in_bounds && nspill > 0 && scan_pos[reduce_dim] == safe_idx(out_meta.limits, reduce_dim) - 1) {
     const vec4 numerator = op1(texelFetch(tin, scan_pos, 0) - max_element);
     vec4 outtex = op2(numerator, safe_denominator);
     [[unroll]] for (int i = nspill; i < 4; ++i) {
@@ -200,13 +206,14 @@ void main() {
       gl_LocalInvocationID[reduce_dim],
       gl_LocalInvocationID[group_dim]);
 
-  if (any(greaterThanEqual(scan_pos, out_meta.limits))) {
-    return;
-  }
+  // Do not return early here. The routines below contain barrier() calls, and
+  // returning would leave them in non-uniform control flow, which is undefined
+  // and hangs the GPU on some drivers. Carry the bounds check instead.
+  const bool in_bounds = !any(greaterThanEqual(scan_pos, out_meta.limits));
 
   if (reduce_dim != packed_dim) {
-    softmax_nonpacked_dim(tid, scan_pos);
+    softmax_nonpacked_dim(tid, scan_pos, in_bounds);
   } else {
-    softmax_packed_dim(tid, scan_pos);
+    softmax_packed_dim(tid, scan_pos, in_bounds);
   }
 }

backends/vulkan/runtime/graph/ops/glsl/var_texture3d.glsl

@@ -65,21 +65,23 @@ VEC4_T calculate_variance(VEC4_T sum, VEC4_T sum_sq, int count) {
   return variance;
 }
 
-void reduce_nonpacked_dim(const ivec2 tid, ivec3 scan_pos) {
+void reduce_nonpacked_dim(const ivec2 tid, ivec3 scan_pos, const bool in_bounds) {
   // shared memory index of this thread
   const int smi = tid_to_smi(tid);
 
   VEC4_T sum = VEC4_T(0);
   VEC4_T sum_sq = VEC4_T(0);
   int count = 0;
 
-  scan_pos[reduce_dim] = tid.x;
-  for (int i = tid.x; i < safe_idx(tin_sizes, reduce_dim);
-       i += NWORKERS, scan_pos[reduce_dim] += NWORKERS) {
-    VEC4_T val = load_texel(tin, scan_pos);
-    sum += val;
-    sum_sq += val * val;
-    count += 1;
+  if (in_bounds) {
+    scan_pos[reduce_dim] = tid.x;
+    for (int i = tid.x; i < safe_idx(tin_sizes, reduce_dim);
+         i += NWORKERS, scan_pos[reduce_dim] += NWORKERS) {
+      VEC4_T val = load_texel(tin, scan_pos);
+      sum += val;
+      sum_sq += val * val;
+      count += 1;
+    }
   }
   // Write partial output to shared memory and synchronize work group
   shared_sum[smi] = sum;
@@ -89,7 +91,7 @@ void reduce_nonpacked_dim(const ivec2 tid, ivec3 scan_pos) {
 
   // Since the reduction row is reduced to only one element, only the "main"
   // thread in the group needs aggregate the partial outputs
-  if (tid.x == 0) {
+  if (in_bounds && tid.x == 0) {
     int group_i = tid.y * NWORKERS;
     sum = shared_sum[group_i];
     sum_sq = shared_sum_sq[group_i];
@@ -132,7 +134,7 @@ void reduce_nonpacked_dim(const ivec2 tid, ivec3 scan_pos) {
  * elements in texels (which occur when the size of the packed dim is not a
  * multiple of 4) so that they do not influence the output of reduction.
  */
-void reduce_packed_dim(const ivec2 tid, ivec3 scan_pos) {
+void reduce_packed_dim(const ivec2 tid, ivec3 scan_pos, const bool in_bounds) {
   // shared memory index of this thread
   const int smi = tid_to_smi(tid);
 
@@ -142,29 +144,36 @@ void reduce_packed_dim(const ivec2 tid, ivec3 scan_pos) {
   // handled specially if it has padding elements.
   const int reduce_len = safe_idx(tin_sizes, packed_dim) - nspill;
 
+  // Out of bounds threads must not load or accumulate, but they must still
+  // reach the barrier below, so the work is guarded rather than returned from.
+  // They contribute a zero partial, which the aggregating thread of an in
+  // bounds row never reads: `in_bounds` does not depend on tid.x, so a row is
+  // either wholly in bounds or wholly out of it.
   VEC4_T sum = VEC4_T(0);
   VEC4_T sum_sq = VEC4_T(0);
   int count = 0;
 
-  // Partially accumulate over elements i, i + NWORKERS, i + 2*NWORKERS, ... of
-  // the reduction row
-  scan_pos[reduce_dim] = tid.x;
-  for (int i = tid.x * 4; i < reduce_len;
-       i += NWORKERS * 4, scan_pos[reduce_dim] += NWORKERS) {
-    VEC4_T val = load_texel(tin, scan_pos);
-    sum += val;
-    sum_sq += val * val;
-    count += 4;
-  }
-  // For the last texel in the dim, if there are padding elements then each
-  // element of the texel needs to be processed individually such that the
-  // padding elements are ignored
-  if (scan_pos[reduce_dim] == safe_idx(tin_limits, reduce_dim) - 1 && nspill > 0) {
-    const VEC4_T val = load_texel(tin, scan_pos);
-    for (int i = 0; i < nspill; i++) {
-      sum.x += val[i];
-      sum_sq.x += val[i] * val[i];
-      count += 1;
+  if (in_bounds) {
+    // Partially accumulate over elements i, i + NWORKERS, i + 2*NWORKERS, ... of
+    // the reduction row
+    scan_pos[reduce_dim] = tid.x;
+    for (int i = tid.x * 4; i < reduce_len;
+         i += NWORKERS * 4, scan_pos[reduce_dim] += NWORKERS) {
+      VEC4_T val = load_texel(tin, scan_pos);
+      sum += val;
+      sum_sq += val * val;
+      count += 4;
+    }
+    // For the last texel in the dim, if there are padding elements then each
+    // element of the texel needs to be processed individually such that the
+    // padding elements are ignored
+    if (scan_pos[reduce_dim] == safe_idx(tin_limits, reduce_dim) - 1 && nspill > 0) {
+      const VEC4_T val = load_texel(tin, scan_pos);
+      for (int i = 0; i < nspill; i++) {
+        sum.x += val[i];
+        sum_sq.x += val[i] * val[i];
+        count += 1;
+      }
     }
   }
   // Write partial output to shared memory and synchronize work group
@@ -175,7 +184,7 @@ void reduce_packed_dim(const ivec2 tid, ivec3 scan_pos) {
 
   // Since the reduction row is reduced to only one element, only the "main"
   // thread in the group needs aggregate the partial outputs
-  if (tid.x == 0) {
+  if (in_bounds && tid.x == 0) {
     sum = shared_sum[tid.y * NWORKERS];
     sum_sq = shared_sum_sq[tid.y * NWORKERS];
     count = shared_count[tid.y * NWORKERS];
@@ -211,13 +220,14 @@ void main() {
       gl_LocalInvocationID[reduce_dim],
       gl_LocalInvocationID[group_dim]);
 
-  if (any(greaterThanEqual(scan_pos, tin_limits))) {
-    return;
-  }
+  // Do not return early here. The routines below contain barrier() calls, and
+  // returning would leave them in non-uniform control flow, which is undefined
+  // and hangs the GPU on some drivers. Carry the bounds check instead.
+  const bool in_bounds = !any(greaterThanEqual(scan_pos, tin_limits));
 
   if (reduce_dim != packed_dim) {
-    reduce_nonpacked_dim(tid, scan_pos);
+    reduce_nonpacked_dim(tid, scan_pos, in_bounds);
   } else {
-    reduce_packed_dim(tid, scan_pos);
+    reduce_packed_dim(tid, scan_pos, in_bounds);
   }
 }