gitadmin/foxhunt
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix(cuda): reduction kernel bugs — missing shmem fold + count reduction

Browse Source 
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) <noreply@anthropic.com>
This commit is contained in:
jgrusewski
2026-03-18 14:01:50 +01:00
parent 3da65f2804
commit 2d703a2f9b
1 changed files with 32 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files

41
crates/ml-core/src/cuda_autograd/reductions.rs
View File

@@ -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;