Register Blocking (Thread Coarsening)

Definition

Register blocking (a.k.a. thread coarsening) is the GEMM optimization rung that relieves SMEM bandwidth by having each thread compute a small TM×TN micro-tile of outputs instead of one, staging a slice of the SMEM tiles into registers and reusing those across many FMAs. It is the same cache-and-reuse move as tiling, applied one tier down the memory hierarchy (SMEM → registers), and it is the single biggest jump on the GEMM ladder (~13% → ~69% of cuBLAS).

How it works

After SMEM tiling, the inner loop issues 2 SMEM loads per FMA — the shared-memory pipe is the bottleneck (~1 FLOP/load). Fix: fewer, fatter threads, each accumulating a micro-tile in the register file (~256 KB/SM, ~1-cycle latency, effectively unlimited bandwidth to the ALUs).

The TM×TN outer product. Per $K$-step (dotIdx), load a length-TM column of As (regM) and a length-TN row of Bs (regN) into registers, then do all TM·TN FMAs of their outer product into a register accumulator:

for i in TM: regM[i] = As[...]     // TM SMEM loads
for j in TN: regN[j] = Bs[...]     // TN SMEM loads
for m,n:     acc[m][n] += regM[m]*regN[n]   // TM*TN FMAs, 0 SMEM loads

SMEM arithmetic intensity jumps from ~1 to
$\frac{2TM\cdot TN}{TM+TN}\stackrel{TM=TN=8}{\to }\frac{2\cdot 64}{16}=8\frac{\text{FLOP}}{\text{SMEM load}}$
— an 8× cut in SMEM traffic per FLOP (16 loads feed 64 FMAs). Now the ALUs, not SMEM, bound the kernel. 1D blocking (TM only) first takes ~13%→37%; the 2D outer product takes ~37%→69%. This micro-tile is exactly the shape of a Tensor Core MMA, which is why it is the bridge to CUTLASS.

Cost: ~90–110 regs/thread (64 accumulators + regM/regN + addressing) caps occupancy (often 2 blocks/SM, ~32%), but a coarsened thread hides latency with instruction-level parallelism (many independent FMAs), tolerating low occupancy. Push TM,TN too far → register spills to local memory (HBM) — catastrophic. Check ptxas -v; sweep the shape.

Warptiling. Adds a block → warp → thread level between block-tile and thread-tile. The warp is the real scheduling and SMEM/register-access unit, so a contiguous warp-tile improves locality, enables register reuse within the warp, and matches per-warp MMA issue — the last ~15 points to ~94% of cuBLAS (the structure CUTLASS uses).

Why it matters

Register blocking is the least-obvious, highest-payoff rung — more speedup than the entire rest of the ladder combined. It is the general recursion of GPU performance: push data reuse down to the fastest tier that still has capacity. The outer-product formulation is the mental model for every Tensor-Core kernel and for the compute core of transformer training.

Taught in

• 05-optimizing-gemm — §6 1D/2D register blocking and the outer product; §7 warptiling.

See also

• shared-memory-tiling — the tier above; register blocking’s prerequisite
• gemm — the workload
• tensor-core — the micro-tile maps onto MMAs
• gpu-memory-hierarchy — the register tier
• occupancy — the trade-off register blocking makes
• shared-memory-bank-conflicts — conflict-free inner-loop SMEM reads
• arithmetic-intensity — per-tier intensity
• gpu-systems-for-llms