From b692994a88baebc723a0ec54390bd25a4fb6e686 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mateusz=20S=C5=82uszniak?= Date: Sun, 30 Aug 2026 11:30:07 +0200 Subject: [PATCH 1/2] [ET-VK] Do not return early before the barrier in the reduce shader 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. --- .../vulkan/runtime/graph/ops/glsl/reduce.glsl | 81 +++++++++++-------- 1 file changed, 46 insertions(+), 35 deletions(-) diff --git a/backends/vulkan/runtime/graph/ops/glsl/reduce.glsl b/backends/vulkan/runtime/graph/ops/glsl/reduce.glsl index 209440cec6a..a90b5e57cc0 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/reduce.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/reduce.glsl @@ -86,19 +86,24 @@ int tid_to_smi(const ivec2 tid) { * This case is simpler because each element of a texel belongs to a separate * reduction "group", meaning we don't have to perform reduction along a texel. */ -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); - scan_pos[reduce_dim] = 0; - vec4 accum = INIT_ACCUM(load_texel(tin, scan_pos)); - - scan_pos[reduce_dim] = tid.x; - // Partially accumulate over elements i, i + NWORKERS, i + 2*NWORKERS, ... of - // the reduction row - for (int i = tid.x; i < safe_idx(tin_sizes, reduce_dim); - i += NWORKERS, scan_pos[reduce_dim] += NWORKERS) { - accum = UPDATE_ACCUM(accum, load_texel(tin, scan_pos)); + // 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. + vec4 accum = vec4(0); + if (in_bounds) { + scan_pos[reduce_dim] = 0; + accum = INIT_ACCUM(load_texel(tin, scan_pos)); + + scan_pos[reduce_dim] = tid.x; + // Partially accumulate over elements i, i + NWORKERS, i + 2*NWORKERS, ... of + // the reduction row + for (int i = tid.x; i < safe_idx(tin_sizes, reduce_dim); + i += NWORKERS, scan_pos[reduce_dim] += NWORKERS) { + accum = UPDATE_ACCUM(accum, load_texel(tin, scan_pos)); + } } // Write partial output to shared memory and synchronize work group shared_vecs[smi] = accum; @@ -106,7 +111,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) { // Iterate over the partial outputs to obtain the overall output int group_i = tid.y * NWORKERS; accum = shared_vecs[group_i++]; @@ -141,7 +146,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); @@ -151,23 +156,28 @@ 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; - scan_pos[reduce_dim] = 0; - vec4 accum = INIT_ACCUM(vec4(load_texel(tin, scan_pos).x)); - - // 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) { - accum = UPDATE_ACCUM(accum, load_texel(tin, scan_pos)); - } - // 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 intex = load_texel(tin, scan_pos); - for (int i = 0; i < nspill; i++) { - accum.x = UPDATE_ACCUM(accum.x, intex[i]); + // 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. + vec4 accum = vec4(0); + if (in_bounds) { + scan_pos[reduce_dim] = 0; + accum = INIT_ACCUM(vec4(load_texel(tin, scan_pos).x)); + + // 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) { + accum = UPDATE_ACCUM(accum, load_texel(tin, scan_pos)); + } + // 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 intex = load_texel(tin, scan_pos); + for (int i = 0; i < nspill; i++) { + accum.x = UPDATE_ACCUM(accum.x, intex[i]); + } } } // Write partial output to shared memory and synchronize work group @@ -176,7 +186,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) { // Iterate over the partial maximums to obtain the overall maximum int group_i = tid.y * NWORKERS; accum = shared_vecs[group_i++]; @@ -203,13 +213,14 @@ void main() { gl_LocalInvocationID[reduce_dim], gl_LocalInvocationID[group_dim]); - if (any(greaterThanEqual(scan_pos, tin_limits))) { - return; - } + // Do not return early here. Both reduction routines contain a barrier(), and + // returning would leave it 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); } } From dbabf91473dd3b4023254e19ca13224dae9312dc Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mateusz=20S=C5=82uszniak?= Date: Sun, 30 Aug 2026 12:33:21 +0200 Subject: [PATCH 2/2] [ET-VK] Apply the same barrier fix to softmax, var and reduce2d 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. --- .../runtime/graph/ops/glsl/reduce2d.glsl | 40 +++++++------- .../runtime/graph/ops/glsl/softmax.glsl | 53 +++++++++++-------- .../runtime/graph/ops/glsl/var_texture3d.glsl | 35 ++++++------ 3 files changed, 71 insertions(+), 57 deletions(-) diff --git a/backends/vulkan/runtime/graph/ops/glsl/reduce2d.glsl b/backends/vulkan/runtime/graph/ops/glsl/reduce2d.glsl index bd55025f534..db4e3631728 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/reduce2d.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/reduce2d.glsl @@ -59,23 +59,26 @@ int tid_to_smi(const ivec2 tid) { // with the accumulator. #define POSTPROCESS(accum) ${POSTPROCESS} -void reduce_2d_non_packed_dim(const ivec2 tid, ivec3 scan_pos) { +void reduce_2d_non_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); - scan_pos[reduce_dim1] = 0; - scan_pos[reduce_dim2] = 0; - vec4 accum = INIT_ACCUM(load_texel(tin, scan_pos)); - - // First dimension reduction - scan_pos[reduce_dim1] = tid.x; - for (int i = tid.x; i < safe_idx(tin_sizes, reduce_dim1); - i += NWORKERS, scan_pos[reduce_dim1] += NWORKERS) { - - // Second dimension reduction + vec4 accum = vec4(0); + if (in_bounds) { + scan_pos[reduce_dim1] = 0; scan_pos[reduce_dim2] = 0; - for (int j = 0; j < safe_idx(tin_sizes, reduce_dim2); j++, scan_pos[reduce_dim2]++) { - accum = UPDATE_ACCUM(accum, load_texel(tin, scan_pos)); + accum = INIT_ACCUM(load_texel(tin, scan_pos)); + + // First dimension reduction + scan_pos[reduce_dim1] = tid.x; + for (int i = tid.x; i < safe_idx(tin_sizes, reduce_dim1); + i += NWORKERS, scan_pos[reduce_dim1] += NWORKERS) { + + // Second dimension reduction + scan_pos[reduce_dim2] = 0; + for (int j = 0; j < safe_idx(tin_sizes, reduce_dim2); j++, scan_pos[reduce_dim2]++) { + accum = UPDATE_ACCUM(accum, load_texel(tin, scan_pos)); + } } } @@ -84,7 +87,7 @@ void reduce_2d_non_packed_dim(const ivec2 tid, ivec3 scan_pos) { barrier(); // Main thread aggregates results - if (tid.x == 0) { + if (in_bounds && tid.x == 0) { // Iterate over the partial outputs to obtain the overall output int group_i = tid.y * NWORKERS; accum = shared_vecs[group_i++]; @@ -121,9 +124,10 @@ void main() { gl_LocalInvocationID[reduce_dim1], 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)); - reduce_2d_non_packed_dim(tid, scan_pos); + reduce_2d_non_packed_dim(tid, scan_pos, in_bounds); } \ No newline at end of file diff --git a/backends/vulkan/runtime/graph/ops/glsl/softmax.glsl b/backends/vulkan/runtime/graph/ops/glsl/softmax.glsl index ce9d2477795..9e14d51bccf 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/softmax.glsl +++ b/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); } } diff --git a/backends/vulkan/runtime/graph/ops/glsl/var_texture3d.glsl b/backends/vulkan/runtime/graph/ops/glsl/var_texture3d.glsl index 1ee938c58dc..2d462661f63 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/var_texture3d.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/var_texture3d.glsl @@ -65,7 +65,7 @@ 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); @@ -73,13 +73,15 @@ void reduce_nonpacked_dim(const ivec2 tid, ivec3 scan_pos) { 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); @@ -175,7 +177,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 +213,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); } }