From 2d703a2f9b0c8540394351d2c59627a6710cab3a Mon Sep 17 00:00:00 2001 From: jgrusewski Date: Wed, 18 Mar 2026 14:01:50 +0100 Subject: [PATCH] =?UTF-8?q?fix(cuda):=20reduction=20kernel=20bugs=20?= =?UTF-8?q?=E2=80=94=20missing=20shmem=20fold=20+=20count=20reduction?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Bug 1: Tree reduction stopped at s>32, leaving shmem[32..63] unfolded before warp shuffle. Added explicit fold: thread i merges shmem[i+32] before __shfl_down_sync. Affected all 5 kernels (stats, argmax, sum, argmax_rows, col_sum). Caused max=96 instead of 100 for N=100. Bug 2: fused_stats_reduce count only atomicAdd'd thread 0's local_count. Added 5th shared memory slot for count through full tree reduction. All 5 reduction tests pass. Co-Authored-By: Claude Opus 4.6 (1M context) --- .../ml-core/src/cuda_autograd/reductions.rs | 41 +++++++++++++++---- 1 file changed, 32 insertions(+), 9 deletions(-) diff --git a/crates/ml-core/src/cuda_autograd/reductions.rs b/crates/ml-core/src/cuda_autograd/reductions.rs index 120bdcfbd..ca6ac5785 100644 --- a/crates/ml-core/src/cuda_autograd/reductions.rs +++ b/crates/ml-core/src/cuda_autograd/reductions.rs @@ -99,7 +99,7 @@ impl ReductionKernels { // Cap blocks to avoid excessive atomic contention let blocks = blocks.min(1024); let n_i32 = n as i32; - let shared_bytes = threads * 4 * 4; // 4 floats per thread for shared reduction + let shared_bytes = threads * 5 * 4; // 5 floats per thread for shared reduction (min, max, sum, sum_sq, count) let cfg = LaunchConfig { grid_dim: (blocks, 1, 1), @@ -413,6 +413,7 @@ void fused_stats_reduce(const float* __restrict__ data, shmem[B + threadIdx.x] = local_max; shmem[2*B + threadIdx.x] = local_sum; shmem[3*B + threadIdx.x] = local_sum_sq; + shmem[4*B + threadIdx.x] = (float)local_count; __syncthreads(); // Tree reduction in shared memory @@ -422,27 +423,29 @@ void fused_stats_reduce(const float* __restrict__ data, shmem[B + threadIdx.x] = fmaxf(shmem[B + threadIdx.x], shmem[B + threadIdx.x + s]); shmem[2*B + threadIdx.x] += shmem[2*B + threadIdx.x + s]; shmem[3*B + threadIdx.x] += shmem[3*B + threadIdx.x + s]; + shmem[4*B + threadIdx.x] += shmem[4*B + threadIdx.x + s]; } __syncthreads(); } // Final warp reduction (no __syncthreads needed within a warp) if (threadIdx.x < 32) { - float w_min = shmem[threadIdx.x]; - float w_max = shmem[B + threadIdx.x]; - float w_sum = shmem[2*B + threadIdx.x]; - float w_ssq = shmem[3*B + threadIdx.x]; + // Fold the second set of 32 elements (indices 32..63) into 0..31 + float w_min = fminf(shmem[threadIdx.x], shmem[threadIdx.x + 32]); + float w_max = fmaxf(shmem[B + threadIdx.x], shmem[B + threadIdx.x + 32]); + float w_sum = shmem[2*B + threadIdx.x] + shmem[2*B + threadIdx.x + 32]; + float w_ssq = shmem[3*B + threadIdx.x] + shmem[3*B + threadIdx.x + 32]; + float w_cnt = shmem[4*B + threadIdx.x] + shmem[4*B + threadIdx.x + 32]; // Reduce within the warp using shuffle w_min = warp_reduce_min(w_min); w_max = warp_reduce_max(w_max); w_sum = warp_reduce_sum(w_sum); w_ssq = warp_reduce_sum(w_ssq); + w_cnt = warp_reduce_sum(w_cnt); if (threadIdx.x == 0) { // Atomic merge to global result - // result[0] = min (atomicMin via int reinterpretation for positive floats, - // but simpler: use atomicCAS loop for general floats) // For simplicity and correctness with negative values, use CAS loops. { // atomicMin for float via CAS loop @@ -470,7 +473,7 @@ void fused_stats_reduce(const float* __restrict__ data, } atomicAdd(&result[2], w_sum); atomicAdd(&result[3], w_ssq); - atomicAdd(&result[4], (float)local_count); + atomicAdd(&result[4], w_cnt); } } } @@ -514,8 +517,15 @@ void argmax_flat(const float* __restrict__ data, // Final warp if (threadIdx.x < 32) { + // Fold shmem[32..63] into shmem[0..31] float w_val = shmem[threadIdx.x]; float w_idx = shmem[B + threadIdx.x]; + float other_val32 = shmem[threadIdx.x + 32]; + float other_idx32 = shmem[B + threadIdx.x + 32]; + if (other_val32 > w_val) { + w_val = other_val32; + w_idx = other_idx32; + } for (int offset = 16; offset > 0; offset >>= 1) { float other_val = __shfl_down_sync(0xFFFFFFFF, w_val, offset); @@ -590,6 +600,15 @@ void argmax_rows(const float* __restrict__ data, if (threadIdx.x < 32) { float w_val = shmem[threadIdx.x]; float w_idx = shmem[B + threadIdx.x]; + // Fold shmem[32..63] into shmem[0..31] when block is large enough + if (B >= 64) { + float other_val32 = shmem[threadIdx.x + 32]; + float other_idx32 = shmem[B + threadIdx.x + 32]; + if (other_val32 > w_val) { + w_val = other_val32; + w_idx = other_idx32; + } + } for (int offset = 16; offset > 0; offset >>= 1) { float other_val = __shfl_down_sync(0xFFFFFFFF, w_val, offset); @@ -633,7 +652,8 @@ void sum_reduce(const float* __restrict__ data, } if (threadIdx.x < 32) { - float w_sum = shmem[threadIdx.x]; + // Fold shmem[32..63] into shmem[0..31] + float w_sum = shmem[threadIdx.x] + shmem[threadIdx.x + 32]; w_sum = warp_reduce_sum(w_sum); if (threadIdx.x == 0) { atomicAdd(result, w_sum); @@ -672,7 +692,10 @@ void col_sum_reduce(const float* __restrict__ data, } if (threadIdx.x < 32) { + // Fold shmem[32..63] into shmem[0..31] when block is large enough float w_sum = shmem[threadIdx.x]; + if (blockDim.x >= 64) + w_sum += shmem[threadIdx.x + 32]; w_sum = warp_reduce_sum(w_sum); if (threadIdx.x == 0) { result[col] = w_sum;