vulkan: fix data race and OOB access in argsort(large) (#28705)

argsort had a data race in the inner loop, which VVL caught. But I don't think
this was causing failures in practice.

argsort_large has OOB accesses which might explain the failures in CI, but I
couldn't reproduce it locally and I don't think it's a convincing explanation
of the failures.
This commit is contained in:
Jeff Bolz
2026-09-11 08:44:13 +03:00
committed by GitHub
parent df03399b88
commit 481c65f091
2 changed files with 21 additions and 12 deletions
@@ -33,7 +33,11 @@ void argsort(bool needs_bounds_check, const uint row) {
const uint row_offset = row * p.ncols;
// initialize indices
dst_row[col] = ivec2(col, floatBitsToInt(data_a[row_offset + col]));
ivec2 value = ivec2(col, 0);
if (!needs_bounds_check || col < p.ncols) {
value.y = floatBitsToInt(data_a[row_offset + col]);
}
dst_row[col] = value;
barrier();
uint num_outer_loop_iters = NCOLS_PADDED_LOG2;
@@ -42,18 +46,20 @@ void argsort(bool needs_bounds_check, const uint row) {
[[unroll]] for (uint j = k / 2, inner_idx = 0; inner_idx < num_inner_loop_iters; j /= 2, inner_idx++) {
const int ixj = int(col ^ j);
int idx_0 = (col & k) == 0 ? col : ixj;
int idx_1 = (col & k) == 0 ? ixj : col;
if (ixj > col) {
int idx_0 = (col & k) == 0 ? col : ixj;
int idx_1 = (col & k) == 0 ? ixj : col;
ivec2 sh_idx_0 = dst_row[idx_0];
ivec2 sh_idx_1 = dst_row[idx_1];
bool idx_0_oob = needs_bounds_check ? sh_idx_0.x >= p.ncols : false;
bool idx_1_oob = needs_bounds_check ? sh_idx_1.x >= p.ncols : false;
ivec2 sh_idx_0 = dst_row[idx_0];
ivec2 sh_idx_1 = dst_row[idx_1];
bool idx_0_oob = needs_bounds_check ? sh_idx_0.x >= p.ncols : false;
bool idx_1_oob = needs_bounds_check ? sh_idx_1.x >= p.ncols : false;
if ((idx_0_oob ||
(!idx_1_oob && intBitsToFloat(sh_idx_0.y) > intBitsToFloat(sh_idx_1.y))) && (ixj > col)) {
dst_row[idx_0] = sh_idx_1;
dst_row[idx_1] = sh_idx_0;
if (idx_0_oob ||
(!idx_1_oob && intBitsToFloat(sh_idx_0.y) > intBitsToFloat(sh_idx_1.y))) {
dst_row[idx_0] = sh_idx_1;
dst_row[idx_1] = sh_idx_0;
}
}
barrier();
@@ -42,7 +42,10 @@ void argsort(bool needs_bounds_check, const uint row) {
[[unroll]] for (int u = 0; u < WG_UNROLL_FACTOR; ++u) {
uint c = u*BLOCK_SIZE + col;
if (c < p.ncols_padded) {
ivec2 v = ivec2(c, floatBitsToInt(data_a[row_offset + c]));
ivec2 v = ivec2(c, 0);
if (!needs_bounds_check || c < p.ncols) {
v.y = floatBitsToInt(data_a[row_offset + c]);
}
tmp_idx[idx_offset + c] = v;
}
}