Seven-step reduction: 30x from one kernel
From the naive version all the way to warp shuffle, with a clear reason for each speedup
- Master parallel reduction, the most important GPU primitive
- Work through the cost of divergence, bank conflict, idle threads, and loop overhead, one by one
- Learn to skip shared memory entirely with warp shuffle
- Build the habit of naming the exact bottleneck each optimization removes
Sum, max, the softmax denominator, L2 norm, the row-sum in attention: these are all reduction. It is the classic GPU teaching case because it lines up almost every optimization trick on one path. The seven versions below come from Mark Harris's classic talk. We will take them apart one by one and name the bottleneck each step actually removes.
V1: interleaved addressing. It runs, but every step diverges
1__global__ void reduce_v1(const float* in, float* out, int n) {2 extern __shared__ float sdata[];3 int tid = threadIdx.x;4 int i = blockIdx.x * blockDim.x + tid;5 6 sdata[tid] = (i < n) ? in[i] : 0.0f;7 __syncthreads();8 9 for (int s = 1; s < blockDim.x; s *= 2) {10 if (tid % (2 * s) == 0) { // the problem is right here11 sdata[tid] += sdata[tid + s];12 }13 __syncthreads();14 }15 if (tid == 0) out[blockIdx.x] = sdata[0];16}tid % (2*s) == 0 leaves only some threads in the warp doing work, and those threads are interleaved: round one is 0, 2, 4, 6, so of 32 threads in the warp half work and half idle, 100% divergence. Worse, % is a modulo, which is relatively expensive on a GPU.
V2: change the indexing. Remove divergence
1for (int s = 1; s < blockDim.x; s *= 2) {2 int index = 2 * s * tid; // switch from "who works" to "which slot"3 if (index < blockDim.x) {4 sdata[index] += sdata[index + s];5 }6 __syncthreads();7}The amount of work is unchanged, but the active threads become a consecutive run of tid starting at 0. The first few warps are fully busy and the later ones are fully idle, so there is no intra-warp divergence. The cost is a new problem: index = 2*s*tid makes the shared-memory stride 2, 4, 8, and bank conflicts show up.
V3: sequential addressing. Remove bank conflicts
1for (int s = blockDim.x / 2; s > 0; s >>= 1) { // walk the stride from large to small2 if (tid < s) {3 sdata[tid] += sdata[tid + s]; // active threads touch consecutive addresses4 }5 __syncthreads();6}V4: add once during the load. Half the threads stop sitting idle
Look at round one of V3: only half of the blockDim.x threads are working. That means half of the threads you launched spend the entire kernel doing one thing: moving data into shared memory. If that is the case, have each thread read two elements and add them during the load, and the number of blocks is cut in half.
1int tid = threadIdx.x;2int i = blockIdx.x * (blockDim.x * 2) + tid; // each block covers twice as much data3 4sdata[tid] = (i < n ? in[i] : 0.0f) + (i + blockDim.x < n ? in[i + blockDim.x] : 0.0f);5__syncthreads();6 7for (int s = blockDim.x / 2; s > 0; s >>= 1) {8 if (tid < s) sdata[tid] += sdata[tid + s];9 __syncthreads();10}V5: unroll the last warp. Drop the extra barriers
When s <= 32, only one warp is still working. At that point __syncthreads() is pure waste: it is a block-level barrier, but only one warp needs to synchronize. Take the last 6 rounds out and unroll them.
1for (int s = blockDim.x / 2; s > 32; s >>= 1) { // loop until 322 if (tid < s) sdata[tid] += sdata[tid + s];3 __syncthreads();4}5 6if (tid < 32) {7 // note: do not rely on implicit intra-warp sync; after Volta you must be explicit8 float v = sdata[tid] + sdata[tid + 32];9 for (int offset = 16; offset > 0; offset >>= 1) {10 v += __shfl_down_sync(0xffffffff, v, offset); // exchange data directly in registers11 }12 if (tid == 0) out[blockIdx.x] = v;13}V6: fully unroll with a template
1template <unsigned int blockSize>2__global__ void reduce_v6(const float* in, float* out, int n) {3 extern __shared__ float sdata[];4 unsigned tid = threadIdx.x;5 unsigned i = blockIdx.x * (blockSize * 2) + tid;6 unsigned gridSize = blockSize * 2 * gridDim.x;7 8 float sum = 0.0f;9 while (i < n) { // grid-stride: one block consumes as much data as it needs10 sum += in[i];11 if (i + blockSize < n) sum += in[i + blockSize];12 i += gridSize;13 }14 sdata[tid] = sum;15 __syncthreads();16 17 // blockSize is a compile-time constant, so these ifs fold away and the loop fully unrolls18 if (blockSize >= 512) { if (tid < 256) sdata[tid] += sdata[tid + 256]; __syncthreads(); }19 if (blockSize >= 256) { if (tid < 128) sdata[tid] += sdata[tid + 128]; __syncthreads(); }20 if (blockSize >= 128) { if (tid < 64) sdata[tid] += sdata[tid + 64]; __syncthreads(); }21 22 if (tid < 32) {23 float v = sdata[tid] + sdata[tid + 32];24 for (int off = 16; off > 0; off >>= 1) v += __shfl_down_sync(0xffffffff, v, off);25 if (tid == 0) out[blockIdx.x] = v;26 }27}28 29// instantiate the template explicitly at the call site30reduce_v6<256><<<blocks, 256, 256 * sizeof(float)>>>(d_in, d_out, n);V7: all warp shuffle. Shared memory is only for cross-warp
1__inline__ __device__ float warpReduceSum(float v) {2 for (int off = warpSize / 2; off > 0; off >>= 1)3 v += __shfl_down_sync(0xffffffff, v, off); // stays in registers the whole way4 return v;5}6 7__global__ void reduce_v7(const float* in, float* out, int n) {8 float sum = 0.0f;9 for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < n;10 i += blockDim.x * gridDim.x) sum += in[i];11 12 sum = warpReduceSum(sum); // stage 1: intra-warp reduction, zero shared memory13 14 __shared__ float warpSums[32]; // at most 1024/32 = 32 warps15 int lane = threadIdx.x % warpSize;16 int wid = threadIdx.x / warpSize;17 if (lane == 0) warpSums[wid] = sum; // each warp writes a single value18 __syncthreads();19 20 // stage 2: warp 0 reduces the per-warp results21 sum = (threadIdx.x < blockDim.x / warpSize) ? warpSums[lane] : 0.0f;22 if (wid == 0) sum = warpReduceSum(sum);23 24 if (threadIdx.x == 0) atomicAdd(out, sum); // could also write out[blockIdx.x]25}| Version | Key change | Bottleneck removed | vs V1 |
|---|---|---|---|
| V1 | tid % (2*s) | none (baseline) | 1× |
| V2 | switch to index = 2*s*tid | warp divergence | ~2× |
| V3 | stride from large to small | bank conflict | ~4× |
| V4 | add once during the load | half the threads idle | ~8× |
| V5 | unroll the last warp | extra block barriers | ~12× |
| V6 | template + full unroll + grid-stride | loop and address-math overhead | ~20× |
| V7 | two-level warp-shuffle reduction | shared-memory round trips | ~30×, near the bandwidth ceiling |
1#include <cub/cub.cuh>2 3void* d_temp = nullptr;4size_t temp_bytes = 0;5 6// first call only asks how much temp space is needed7cub::DeviceReduce::Sum(d_temp, temp_bytes, d_in, d_out, n);8CUDA_CHECK(cudaMalloc(&d_temp, temp_bytes));9 10// the second call actually runs11cub::DeviceReduce::Sum(d_temp, temp_bytes, d_in, d_out, n);Self-check
Answer in your head first, then open the solution. If you cannot, reread that section.