mirror of
https://github.com/ggml-org/llama.cpp.git
synced 2026-09-14 18:02:52 +02:00
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:
@@ -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;
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user