From 481c65f091f74c5e7089dd0a3a1cc6b50cced31e Mon Sep 17 00:00:00 2001 From: Jeff Bolz Date: Fri, 11 Sep 2026 00:44:13 -0500 Subject: [PATCH] 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. --- .../ggml-vulkan/vulkan-shaders/argsort.comp | 28 +++++++++++-------- .../vulkan-shaders/argsort_large.comp | 5 +++- 2 files changed, 21 insertions(+), 12 deletions(-) diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/argsort.comp b/ggml/src/ggml-vulkan/vulkan-shaders/argsort.comp index 0fc2b9b725..4ba63f7aee 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/argsort.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/argsort.comp @@ -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(); diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/argsort_large.comp b/ggml/src/ggml-vulkan/vulkan-shaders/argsort_large.comp index 920bac6bb8..b2df441374 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/argsort_large.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/argsort_large.comp @@ -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; } }