diff --git a/cpp/src/neighbors/ivf_rabitq/gpu_index/quantizer_gpu.cu b/cpp/src/neighbors/ivf_rabitq/gpu_index/quantizer_gpu.cu index af0876acca..002c96440c 100644 --- a/cpp/src/neighbors/ivf_rabitq/gpu_index/quantizer_gpu.cu +++ b/cpp/src/neighbors/ivf_rabitq/gpu_index/quantizer_gpu.cu @@ -263,11 +263,14 @@ __global__ void pack_and_compute_factors_kernel( xu_sq += xu * xu; } - // Perform parallel reduction within the block - l2_sqr = blockReduceSum(l2_sqr); - ip_resi_xucb = blockReduceSum(ip_resi_xucb); - ip_cent_xucb = blockReduceSum(ip_cent_xucb); - xu_sq = blockReduceSum(xu_sq); + // Fused single-pass reduction of the four factor components: one call over its own + // shared scratch, so no back-to-back scratch reuse and no inter-call __syncthreads. + float reduce_buf[4] = {l2_sqr, ip_resi_xucb, ip_cent_xucb, xu_sq}; + blockReduceSumN<4>(reduce_buf); + l2_sqr = reduce_buf[0]; + ip_resi_xucb = reduce_buf[1]; + ip_cent_xucb = reduce_buf[2]; + xu_sq = reduce_buf[3]; // Thread 0 performs the final calculations and writes the factors if (threadIdx.x == 0) { @@ -401,11 +404,14 @@ __global__ void exrabitq_fused_kernel_batch( xu_sq += xu * xu; } - // Perform parallel reductions for all factor components - l2_sqr = blockReduceSum(l2_sqr); - ip_resi_xucb = blockReduceSum(ip_resi_xucb); - ip_cent_xucb = blockReduceSum(ip_cent_xucb); - xu_sq = blockReduceSum(xu_sq); + // Fused single-pass reduction of the four factor components: one call over its own + // shared scratch, so no back-to-back scratch reuse and no inter-call __syncthreads. + float reduce_buf[4] = {l2_sqr, ip_resi_xucb, ip_cent_xucb, xu_sq}; + blockReduceSumN<4>(reduce_buf); + l2_sqr = reduce_buf[0]; + ip_resi_xucb = reduce_buf[1]; + ip_cent_xucb = reduce_buf[2]; + xu_sq = reduce_buf[3]; // Thread 0 computes and writes the final factors if (tid == 0) { @@ -916,6 +922,9 @@ extern __constant__ float d_kTightStart_opt[9] = { 0.81f, }; +// PRECONDITION: s_xp_norm must be block-visible before the call — callers +// fill it with block-strided writes and must issue a __syncthreads() between +// the last write and this call (CUB-style caller-owned input visibility). __device__ float compute_best_rescale_parallel( float* s_xp_norm, int D, int EX_BITS, float* reuse_space, int BlockSize) { @@ -923,26 +932,22 @@ __device__ float compute_best_rescale_parallel( constexpr float kEps = 1e-5f; constexpr int kNEnum = 10; + // Tournament winners are broadcast through a dedicated shared slot (never + // aliased by the workspace and never reused as tournament scratch), so the + // final combine of each argmax tree doubles as the broadcast write and the + // publish/overwrite races need no guard barriers. + __shared__ float s_best_t; + //========================================================================= - // Step 1: Find maximum value using parallel reduction + // Step 1: Block maximum via warp-shuffle all-reduce. fmaxf is exactly + // associative, so this is bit-identical to the former tree reduction; the + // result lands in every thread, which also keeps the early return uniform. //========================================================================= float local_max = 0.0f; for (int i = tid; i < D; i += BlockSize) { local_max = fmaxf(local_max, fabsf(s_xp_norm[i])); } - - // Block-level reduction for max - float* s_reduce; - s_reduce = reuse_space; - s_reduce[tid] = local_max; - __syncthreads(); - - for (int stride = BlockSize / 2; stride > 0; stride >>= 1) { - if (tid < stride) { s_reduce[tid] = fmaxf(s_reduce[tid], s_reduce[tid + stride]); } - __syncthreads(); - } - - float max_o = s_reduce[0]; // full max, already visible to all threads via the reduction's sync + float max_o = blockAllReduceMax(local_max); if (max_o < kEps) return 1.0f; //========================================================================= @@ -989,7 +994,7 @@ __device__ float compute_best_rescale_parallel( s_coarse_t[tid] = best_coarse_t; __syncthreads(); - for (int stride = BlockSize / 2; stride > 0; stride >>= 1) { + for (int stride = BlockSize / 2; stride > 1; stride >>= 1) { if (tid < stride) { if (s_coarse_ip[tid + stride] > s_coarse_ip[tid]) { s_coarse_ip[tid] = s_coarse_ip[tid + stride]; @@ -998,11 +1003,18 @@ __device__ float compute_best_rescale_parallel( } __syncthreads(); } + // Final combine writes the winner straight into the broadcast slot; the + // barrier below both completes the tree and publishes s_best_t, and the + // fine phase may then overwrite the tournament scratch freely. + if (tid == 0) { + s_best_t = (s_coarse_ip[1] > s_coarse_ip[0]) ? s_coarse_t[1] : s_coarse_t[0]; + } + __syncthreads(); //========================================================================= // Phase 2: Fine search around best coarse point //========================================================================= - float center_t = s_coarse_t[0]; + float center_t = s_best_t; float range = (t_end - t_start) / COARSE_SAMPLES; float fine_start = fmaxf(t_start, center_t - range); float fine_end = fminf(t_end, center_t + range); @@ -1041,7 +1053,7 @@ __device__ float compute_best_rescale_parallel( s_fine_t[tid] = best_fine_t; __syncthreads(); - for (int stride = BlockSize / 2; stride > 0; stride >>= 1) { + for (int stride = BlockSize / 2; stride > 1; stride >>= 1) { if (tid < stride) { if (s_fine_ip[tid + stride] > s_fine_ip[tid]) { s_fine_ip[tid] = s_fine_ip[tid + stride]; @@ -1050,8 +1062,15 @@ __device__ float compute_best_rescale_parallel( } __syncthreads(); } - - return s_fine_t[0]; + // Overwriting s_best_t is safe: the fine-publish barrier above ordered + // every thread's read of the coarse center before it. Returning the static + // slot keeps the result immune to workspace reuse in callers (e.g. + // s_tmp_code aliases the workspace in exrabitq_fused_kernel_batch_ori). + if (tid == 0) { + s_best_t = (s_fine_ip[1] > s_fine_ip[0]) ? s_fine_t[1] : s_fine_t[0]; + } + __syncthreads(); + return s_best_t; } template @@ -1090,6 +1109,9 @@ __global__ void exrabitq_fused_kernel_batch_ori( s_xp_norm[j] = d_XP_norm[row * D + j]; // Note: s_bin_xp, s_xp not loaded yet - we'll reuse as workspace } + // compute_best_rescale_parallel cross-reads every element of s_xp_norm: + // its precondition requires the writes above to be block-visible. + __syncthreads(); // Compute scaling factor (reusing unused buffers) float const_scaling_factor = compute_best_rescale_parallel(s_xp_norm, @@ -1145,11 +1167,14 @@ __global__ void exrabitq_fused_kernel_batch_ori( xu_sq += xu * xu; } - // Perform parallel reductions for all factor components - l2_sqr = blockReduceSum(l2_sqr); - ip_resi_xucb = blockReduceSum(ip_resi_xucb); - ip_cent_xucb = blockReduceSum(ip_cent_xucb); - xu_sq = blockReduceSum(xu_sq); + // Fused single-pass reduction of the four factor components: one call over its own + // shared scratch, so no back-to-back scratch reuse and no inter-call __syncthreads. + float reduce_buf[4] = {l2_sqr, ip_resi_xucb, ip_cent_xucb, xu_sq}; + blockReduceSumN<4>(reduce_buf); + l2_sqr = reduce_buf[0]; + ip_resi_xucb = reduce_buf[1]; + ip_cent_xucb = reduce_buf[2]; + xu_sq = reduce_buf[3]; // Thread 0 computes and writes the final factors if (tid == 0) { @@ -1312,6 +1337,9 @@ __global__ void fully_fused_kernel(float* __restrict__ output_factors, for (int i = tid; i < cols; i += block_size) { row_data[i] = fabsf(row_data[i] * inv_norm); } + // compute_best_rescale_parallel cross-reads every element of row_data: + // its precondition requires the writes above to be block-visible. + __syncthreads(); float rescale_factor = compute_best_rescale_parallel(row_data, cols, ex_bits, reuse_space, block_size); diff --git a/cpp/src/neighbors/ivf_rabitq/gpu_index/searcher_gpu_quantize_query.cu b/cpp/src/neighbors/ivf_rabitq/gpu_index/searcher_gpu_quantize_query.cu index 92906fb0e6..7d4c932b00 100644 --- a/cpp/src/neighbors/ivf_rabitq/gpu_index/searcher_gpu_quantize_query.cu +++ b/cpp/src/neighbors/ivf_rabitq/gpu_index/searcher_gpu_quantize_query.cu @@ -108,9 +108,11 @@ __global__ void exrabitq_quantize_query( out_ptr[j] = static_cast(code_val); } - // only thread 0 in the block needs the results, so simply use blockReduceSum - ip_resi_xucb = blockReduceSum(ip_resi_xucb); - xu_sq = blockReduceSum(xu_sq); + // only thread 0 needs the results; fuse both reductions into one race-free pass + float reduce_buf[2] = {ip_resi_xucb, xu_sq}; + blockReduceSumN<2>(reduce_buf); + ip_resi_xucb = reduce_buf[0]; + xu_sq = reduce_buf[1]; // Thread 0 computes and writes the final factors if (tid == 0) { diff --git a/cpp/src/neighbors/ivf_rabitq/utils/reductions.cuh b/cpp/src/neighbors/ivf_rabitq/utils/reductions.cuh index d2bb6641e8..35f1ecbc40 100644 --- a/cpp/src/neighbors/ivf_rabitq/utils/reductions.cuh +++ b/cpp/src/neighbors/ivf_rabitq/utils/reductions.cuh @@ -22,26 +22,84 @@ __inline__ __device__ T warpReduceSum(T val) return val; } -// Block-level sum reduction via per-warp shfl + shared-memory exchange. -// Returns the full block sum in lane 0 of warp 0; other lanes hold partial values. -template -__inline__ __device__ T blockReduceSum(T val) +// Block-level sum reduction of N independent values in a SINGLE pass over a +// private (per-instantiation) shared scratch. Each value gets its own column of +// shared[N][32], so all N are reduced within one write -> __syncthreads() -> read +// window. Because a group of sums is reduced in one call, the scratch is never +// reused by a back-to-back call, so no inter-call barrier is required: the +// shared-memory reuse race is removed by construction rather than patched with a +// guard. On return vals[k] holds the full block sum of the input vals[k] in lane 0 +// of warp 0 (other lanes/warps hold partial values), matching blockReduceSum. +// +// Assumes blockDim.x is a multiple of the warp size and uses at most 32 warps. +template +__inline__ __device__ void blockReduceSumN(T (&vals)[N]) { - __shared__ T shared[32]; // up to 1024 threads -> 32 warps + __shared__ T shared[N][32]; // one column per value; up to 1024 threads -> 32 warps int lane = threadIdx.x & 31; int wid = threadIdx.x >> 5; - val = warpReduceSum(val); - // Wait until every warp has finished reading shared[] in the previous call - // before this call overwrites it. Doing this after the warp-local reduction - // overlaps the extra synchronization with useful work. +#pragma unroll + for (int k = 0; k < N; ++k) { + vals[k] = warpReduceSum(vals[k]); + } + if (lane == 0) { +#pragma unroll + for (int k = 0; k < N; ++k) { + shared[k][wid] = vals[k]; + } + } __syncthreads(); - if (lane == 0) shared[wid] = val; + + bool active = threadIdx.x < blockDim.x / 32; +#pragma unroll + for (int k = 0; k < N; ++k) { + T v = active ? shared[k][lane] : T(0); + if (wid == 0) v = warpReduceSum(v); + vals[k] = v; + } +} + +// Single-value convenience that routes through the fused family (N = 1). Standalone +// reductions (never called back-to-back over the same scratch) use this and stay +// race-free by construction; N = 1 gives each such site its own shared column, +// distinct from any blockReduceSumN used elsewhere in the same kernel. +template +__inline__ __device__ T blockReduceSum(T val) +{ + T v[1] = {val}; + blockReduceSumN<1>(v); + return v[0]; +} + +// Block-level max all-reduce: butterfly shuffles + one shared-memory exchange, +// a single __syncthreads(), and EVERY thread receives the full block maximum +// (no separate broadcast barrier needed). Because fmaxf is exactly associative +// and commutative, the result is bit-identical to a tree reduction over the +// same inputs. Uses its own private scratch; the same one-call-per-kernel +// contract as blockReduceSumN applies. Assumes blockDim.x is a multiple of the +// warp size and uses at most 32 warps. +__inline__ __device__ float blockAllReduceMax(float val) +{ + __shared__ float shared[32]; + int lane = threadIdx.x & 31; + int wid = threadIdx.x >> 5; + int nwarps = blockDim.x >> 5; + const float kNegInf = __int_as_float(0xff800000); // -inf + +#pragma unroll + for (int offset = raft::WarpSize / 2; offset > 0; offset >>= 1) { + val = fmaxf(val, __shfl_xor_sync(0xffffffff, val, offset)); + } + if (lane == 0) { shared[wid] = val; } __syncthreads(); - T out = (threadIdx.x < blockDim.x / 32) ? shared[lane] : T(0); - if (wid == 0) out = warpReduceSum(out); - return out; + float v = (lane < nwarps) ? shared[lane] : kNegInf; +#pragma unroll + for (int offset = raft::WarpSize / 2; offset > 0; offset >>= 1) { + v = fmaxf(v, __shfl_xor_sync(0xffffffff, v, offset)); + } + return v; } } // namespace cuvs::neighbors::ivf_rabitq::detail