GEMM Worklog II: Registers & Warp Tiling

GEMM Worklog II: Registers & Warp Tiling

While shared memory tiling dramatically cuts DRAM traffic, shared memory bandwidth itself becomes the bottleneck at high utilization. Reaching near-cuBLAS performance requires register tiling and vectorized memory instructions.

Micro-Tiling in Registers

Registers provide ~20 TB/s aggregate bandwidth. Instead of each thread calculating a single output element \(C_{i,j}\), each thread computes a \(TM imes TN\) block of elements (e.g. \(8 imes 8 = 64\) values per thread).

Vectorized Memory Loads with float4

Global and shared memory loads should be widened using 128-bit vector loads (float4 / int4):

// Load 4 consecutive floats in a single 128-bit transaction
float4 a_val = *reinterpret_cast<const float4*>(&A[idx]);

The Warptiling Hierarchy

Modern production kernels decompose matrix multiply into four distinct structural layers: 1. Grid Tile: Handled by the entire GPU grid across all SMs. 2. Block Tile: Handled by a single thread block using shared memory. 3. Warp Tile: Handled by a 32-thread warp executing synchronized sub-tiles. 4. Thread Tile: Handled by individual thread registers.

By structuring the kernel into a \(128 imes 128\) block tile and \(64 imes 64\) warp tiles with double buffering, manual register allocation reaches 93.7% of cuBLAS speed on NVIDIA Ampere hardware.