metalworkingGitHub ↗
/techniques/register-blocking

Register Blocking

Register blocking keeps each simdgroup's patch of the output resident in registers for the kernel's entire lifetime: accumulate across every K step, write to memory exactly once.

CUDA equivalent: the accumulator fragments of any warp-tiled GEMM. The difference is proportion: with 208 KB of registers against 32 KB of threadgroup memory, Apple kernels push more of the working set into registers than CUDA-typical, and the technique's failure mode (spilling) is correspondingly more catastrophic.

The pattern in its simplest real form, from the GEMM case study: each simdgroup owns a 4×4 grid of simdgroup_float8x8 accumulators, a 32×32 patch of 1024 floats in registers,

  simdgroup_float8x8 acc[SIMD_TILE][SIMD_TILE];
  for (ushort i = 0; i < SIMD_TILE; i++)
    for (ushort j = 0; j < SIMD_TILE; j++)
      acc[i][j] = simdgroup_float8x8(0);

— m5-gemm sync_copy.metal:104-107↗

accumulated into on every K iteration, stored once in the epilogue. Steel's BlockMMA is the same idea as a template: a TM × TN fragment grid whose size falls out of BM/BN and the simdgroup layout.

A fragments 1 column / K step B fragments · 1 row / K step acc[0][0] acc[0][1] acc[3][3] 32×32 output patch 16 × simdgroup_float8x8 = 1024 floats, in registers, whole kernel the reuse arithmetic: each A fragment feeds a whole row (4 MMAs) each B fragment feeds a whole column (4 MMAs) 8 fragment loads → 16 multiply-accumulates per K step the cliff edges: too small → threadgroup-memory bandwidth bound too big → spills: SIMD_TILE 4 → 8 ran 10× slower the two levers: max_total_threads_per_threadgroup (allocator can plan) 16-bit fragments (half the budget spent) written to DRAM exactly once, in the epilogue

The reuse geometry: eight fragment loads feed sixteen multiply-accumulates per K step, and the green grid never leaves the register file.

The sizing tension, and the two measured cliff edges:

  • Too small: each fragment of A and B loaded from threadgroup memory feeds few multiplies, so you're bandwidth bound on threadgroup memory instead of DRAM. Bigger accumulator grids amortize every fragment load across more MMAs.
  • Too big: the allocator spills, and the measured cost was 10× slower, not 10% (SIMD_TILE 4 → 8 in the case study). The cliff is sharp because spilled accumulators turn every MMA's operand into a memory round-trip.

Which is why the two register-pressure levers appear in every serious kernel: max_total_threads_per_threadgroup so the allocator knows the real thread count (the case-study author's "single biggest practical win"), and 16-bit fragments to halve the budget spent. metal-flash-attention stakes out the extreme position: it sometimes chooses tile shapes that spill, on the theory that a predictable spill of the right operand beats a smaller tile. It's the boldest register-pressure bet in the case studies, and evidence the cliff edge is worth mapping precisely.

Edit this page on GitHub ↗