Skip to content

[ET-VK] Do not return early before the barrier in the reduce shader #22326

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
msluszniak wants to merge 2 commits into
base: main
Choose a base branch
from
+117 −92
Open

[ET-VK] Do not return early before the barrier in the reduce shader #22326

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
b692994
[ET-VK] Do not return early before the barrier in the reduce shader
msluszniak Aug 30, 2026
dbabf91
[ET-VK] Apply the same barrier fix to softmax, var and reduce2d
msluszniak Aug 30, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
81 changes: 46 additions & 35 deletions backends/vulkan/runtime/graph/ops/glsl/reduce.glsl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -86,27 +86,32 @@ 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;
barrier();

// 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++];
Expand Down Expand Up @@ -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);

Expand All @@ -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
Expand All @@ -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++];
Expand All @@ -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);
}
}
40 changes: 22 additions & 18 deletions backends/vulkan/runtime/graph/ops/glsl/reduce2d.glsl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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));
}
}
}

Expand All @@ -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++];
Expand Down Expand Up @@ -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);
}
53 changes: 30 additions & 23 deletions backends/vulkan/runtime/graph/ops/glsl/softmax.glsl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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();
Expand All @@ -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();
Expand All @@ -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));
Expand Down Expand Up @@ -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;

Expand All @@ -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);
Expand All @@ -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);
Expand All @@ -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) {
Expand All @@ -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);
}
}
35 changes: 19 additions & 16 deletions backends/vulkan/runtime/graph/ops/glsl/var_texture3d.glsl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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];
Expand Down Expand Up @@ -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);

Expand Down Expand Up @@ -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];
Expand Down Expand Up @@ -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);
}
}
Loading