Parallel Reduction — infra.bea.ai

Reduction — computing a single value (sum, max, min) from an array — is the simplest kernel that requires inter-thread communication. It’s also the foundation for LayerNorm, Softmax, loss computation, and gradient aggregation. Getting reduction right teaches the core GPU optimization skills: avoiding divergence, eliminating bank conflicts, and using warp primitives.

1. The Problem

Given an array of $N$ elements, compute the sum:

\text{result} = \sum_{i=0}^{N-1} a[i]

On CPU: one loop, $O(N)$ sequential. On GPU: we want $O(N/P)$ time with $P$ parallel threads, but threads need to combine their partial results — this requires a tree reduction.

Performance Target

Reduction is memory-bound — we read $N$ elements and do $N-1$ additions. Arithmetic intensity ≈ 1. The target is: achieve peak HBM bandwidth.

For H100 (3,350 GB/s), reducing 100M floats (400 MB): theoretical minimum = $400 \text{ MB} / 3350 \text{ GB/s} = 0.12$ ms.

2. V0: Naive Parallel Reduction

Reduction Tree

Each thread block reduces a chunk. Within the block, use a tree pattern in shared memory:

__global__ void reduceV0(float* input, float* output, int N) {
    __shared__ float sdata[256];
    int tid = threadIdx.x;
    int idx = blockIdx.x * blockDim.x + threadIdx.x;

    sdata[tid] = (idx < N) ? input[idx] : 0.0f;
    __syncthreads();

    // Tree reduction: stride starts at 1, doubles each step
    for (int stride = 1; stride < blockDim.x; stride *= 2) {
        if (tid % (2 * stride) == 0) {
            sdata[tid] += sdata[tid + stride];
        }
        __syncthreads();
    }

    if (tid == 0) output[blockIdx.x] = sdata[0];
}

Problem: tid % (2 * stride) == 0 causes warp divergence — in the first iteration, only even threads are active (50% utilization). In later iterations, even fewer threads work.

3. V1: Eliminate Warp Divergence

Replace the modulo condition with contiguous thread mapping:

for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) {
    if (tid < stride) {
        sdata[tid] += sdata[tid + stride];
    }
    __syncthreads();
}

Now threads 0 to stride-1 are active — contiguous threads in the same warps. In the first iteration, the first 128 threads (4 warps) are all active with no divergence. Only the last few iterations have partial warps.

Speedup: ~1.5-2× over V0.

4. V2: Eliminate Bank Conflicts

V1 has bank conflicts when stride is a power of 2: threads access sdata[tid] and sdata[tid + stride], which may map to the same bank.

Fix: Reverse the reduction direction — start with large stride and halve:

for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) {
    if (tid < stride) {
        sdata[tid] += sdata[tid + stride];
    }
    __syncthreads();
}

Actually V1 already does this correctly. The bank conflict issue is more subtle — it appears in V0’s increasing-stride pattern where sdata[tid] and sdata[tid + stride] with small stride hit nearby banks. V1’s decreasing-stride pattern naturally avoids the worst conflicts.

5. V3: First-Add During Load

In V0-V2, half the threads are idle from the start — they just load data and immediately become inactive. Fix: Each thread loads and adds two elements during the load phase:

__global__ void reduceV3(float* input, float* output, int N) {
    __shared__ float sdata[256];
    int tid = threadIdx.x;
    int idx = blockIdx.x * (blockDim.x * 2) + threadIdx.x;

    // Each thread adds two elements during load
    float val = 0.0f;
    if (idx < N) val += input[idx];
    if (idx + blockDim.x < N) val += input[idx + blockDim.x];
    sdata[tid] = val;
    __syncthreads();

    for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) {
        if (tid < stride) {
            sdata[tid] += sdata[tid + stride];
        }
        __syncthreads();
    }

    if (tid == 0) output[blockIdx.x] = sdata[0];
}

Speedup: ~2× over V2 — we halved the number of blocks needed.

6. V4: Warp Shuffle for Final Steps

The last 5 iterations (stride ≤ 16) only involve a single warp. We don’t need __syncthreads() or shared memory — use warp shuffle instead:

// After shared memory reduction gets down to 32 elements:
if (tid < 32) {
    float val = sdata[tid];

    // Warp-level reduction using shuffle
    val += __shfl_down_sync(0xFFFFFFFF, val, 16);
    val += __shfl_down_sync(0xFFFFFFFF, val, 8);
    val += __shfl_down_sync(0xFFFFFFFF, val, 4);
    val += __shfl_down_sync(0xFFFFFFFF, val, 2);
    val += __shfl_down_sync(0xFFFFFFFF, val, 1);

    if (tid == 0) output[blockIdx.x] = val;
}

Why it’s faster: __shfl_down_sync is register-to-register — no shared memory access, no bank conflicts, ~1 cycle per operation.

7. V5: Full Warp Shuffle Reduction

Take it further — do the entire block reduction using warp shuffle:

__global__ void reduceV5(float* input, float* output, int N) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    float val = (idx < N) ? input[idx] : 0.0f;

    // Step 1: Warp-level reduction (32 → 1)
    for (int offset = 16; offset > 0; offset >>= 1) {
        val += __shfl_down_sync(0xFFFFFFFF, val, offset);
    }

    // Step 2: Each warp's lane 0 writes to shared memory
    __shared__ float warp_sums[8];  // max 8 warps per block (256 threads)
    int warp_id = threadIdx.x / 32;
    int lane = threadIdx.x % 32;
    if (lane == 0) warp_sums[warp_id] = val;
    __syncthreads();

    // Step 3: First warp reduces the warp sums
    if (warp_id == 0) {
        val = (lane < blockDim.x / 32) ? warp_sums[lane] : 0.0f;
        for (int offset = 16; offset > 0; offset >>= 1) {
            val += __shfl_down_sync(0xFFFFFFFF, val, offset);
        }
        if (lane == 0) output[blockIdx.x] = val;
    }
}

This is the production pattern: warp shuffle → shared memory (one value per warp) → warp shuffle. Minimal shared memory, minimal synchronization.

8. Two-Level Reduction: Block → Global

Each block produces one partial sum. To get the final result, we need a second reduction:

Option A: Launch a second kernel to reduce the block results (simple, always works).

Option B: Use atomicAdd for the final aggregation (one atomic per block — low contention).

// At the end of each block's reduction:
if (threadIdx.x == 0) {
    atomicAdd(output, block_sum);  // one atomic per block
}

For 400 blocks, this is 400 atomics — negligible overhead.

9. Performance Summary

Version	Key optimization	Relative speedup
V0	Naive tree	1.0× (baseline)
V1	Eliminate divergence	~1.5×
V2	Fix bank conflicts	~1.7×
V3	First-add during load	~3.0×
V4	Warp shuffle for tail	~3.3×
V5	Full warp shuffle	~3.5×
+ two-level	Block atomics	~3.5× (same, less overhead)
Target	Peak HBM bandwidth	~4.0×

Key lesson: Most of the speedup comes from V3 (eliminating idle threads) and V5 (replacing shared memory with warp shuffle). The pattern — load efficiently, reduce in registers, minimize shared memory, use atomics for global aggregation — applies to every reduction-like operation in AI.

Self-Check

Explain why reduction is memory-bound and what the bandwidth target is
Implement V0 (naive) and explain the warp divergence problem
Implement the first-add-during-load optimization (V3)
Implement a full warp shuffle reduction (V5)
Explain the two-level reduction pattern (block → atomic)
Calculate theoretical minimum time for reducing 100M floats on H100

References

NVIDIA — Optimizing Parallel Reduction (Harris) — the classic tutorial
CUB Library — Device-wide Reduction