Fuse RaBitQ block reductions into one templated blockReduceSumN<N>

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,33 +922,32 @@ 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)
 {
   int tid              = threadIdx.x;
   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 <unsigned int BlockSize>
@@ -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);

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<int8_t>(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) {

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 <typename T>
-__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 <int N, typename T>
+__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<N != 1> used elsewhere in the same kernel.
+template <typename T>
+__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