Tiled GEMM and data reuse

Calculate reuse and arithmetic intensity for one block tile and identify its capacity limits

Before you begin

Follow a linked prerequisite when you cannot explain it from memory. Background skills without links are stated as abilities rather than hidden terminology.

C = AB indexing
Shared-memory tiles
Roofline

A transformer feed-forward layer receives 2048 token vectors with 4096 features and multiplies them by a 4096 × 11008 weight matrix. PyTorch normally dispatches the operation to a mature matrix library, and Nsight Systems labels the GPU launch with a cuBLAS or cuBLASLt kernel name. A teaching implementation can replace the library call so Nsight Compute can compare a one-output-per-thread kernel with a shared-memory implementation on the same M=2048, N=11008, K=4096 problem.

The first kernel repeatedly requests the same matrix entries because neighboring output elements share input rows and weight columns. Nsight Compute exposes the consequence as high device-memory traffic and low achieved arithmetic throughput. The improved kernel copies a rectangular group of input values into on-chip shared memory once, then many threads reuse those values. The rectangular group is called a tile only after the program and the measured reuse are visible.

The three loops of matrix multiplication hide a reuse opportunity. Every element of A contributes to many columns of C; every element of B contributes to many rows. In CUDA, global memory is the device-wide address space used for large tensor allocations. On the GPU declared for this example, those allocations are physically backed by off-chip high-bandwidth memory (HBM), although hardware caches can satisfy some accesses. A naive thread-per-output implementation can request the same values from that address space repeatedly, turning a compute-rich algorithm into a bandwidth problem. Tiling changes where the loops live so a block deliberately reuses panels in on-chip storage close to the arithmetic.

The matrix equation contains reuse, but the execution schedule determines whether the hardware realizes it. Values must remain in registers, shared memory, or cache long enough to serve several output calculations. A blocked kernel assigns an output tile, loads the corresponding K panels, performs many products from those panels, and retains accumulators until the tile is complete.

Begin with one output element C[i,j]. The output element is the dot product of row i of A and column j of B. A thread can compute that dot product correctly with a loop over K. The performance problem appears when neighboring threads compute C[i,j+1] or C[i+1,j]: they need many of the same A and B values, but a naive schedule does not define where that shared data should remain.

The repeated loads in a one-output-per-thread kernel

Assume A is M×K, B is K×N, and C is M×N. A thread assigned to C[i,j] loads K values from row i of A and K values from column j of B. The adjacent thread for C[i,j+1] loads the same K values from A again. The thread for C[i+1,j] loads the same K values from B again. Hardware caches can recover part of this reuse, but the program has not guaranteed it or arranged accesses for the complete block.

The first useful calculation is bytes per output. Ignoring cache reuse, one FP16 output dot product with length K reads about 4K bytes from A and B and performs 2K floating-point operations. The ratio is about 0.5 floating-point operation per byte before the output store. Modern matrix engines can execute far more arithmetic than this traffic can supply, so a schedule that realizes shared input reuse is required.

One shared-memory tile step

Choose an output region of Mblock×Nblock for one thread block. For a K chunk of width Ktile, the block needs an A panel of Mblock×Ktile and a B panel of Ktile×Nblock. Threads cooperatively request those panels through CUDA global-memory addresses. The requested bytes originate in HBM unless an on-device cache already holds them. The load places the panels in the streaming multiprocessor's on-chip shared memory, which is visible to every thread in the block. A block barrier confirms that the complete panels are ready before any thread consumes them. The block then computes all products from that K chunk before reusing the shared-memory storage for the next panels.

An A value in the panel can contribute to Nblock output columns. A B value can contribute to Mblock output rows. The tile therefore converts one HBM load into many multiply-accumulates. A second barrier is needed before threads overwrite shared storage for the next K chunk, unless asynchronous pipeline primitives provide an equivalent producer-consumer protocol.

Ownership inside the thread block

The complete block tile is too large for one thread. Warps own smaller warp tiles, and each lane owns some accumulator elements within its warp tile. The accumulator values usually remain in registers across the complete K loop. A lane also loads fragments of A and B from shared memory in a pattern that permits reuse across several accumulator updates.

Larger per-warp or per-thread output tiles increase reuse and instruction-level parallelism, but they also increase the number of live accumulators. The compiler allocates those accumulators as registers. Excessive register demand can reduce the number of resident warps or cause spills to local memory. Shared-memory tile size, register count, threads per block, and available output tiles must therefore be evaluated together.

Tile-size limits and edge dimensions

For one K step, approximate work is 2·Mblock·Nblock·Ktile FLOPs. Input data is (Mblock·Ktile + Ktile·Nblock) elements. Increasing Mblock or Nblock raises work faster than panel bytes, but capacity and parallelism set upper limits. A large tile can also leave too few blocks when M or N is small.

Dimensions are rarely exact multiples of the chosen tile. The kernel can predicate loads and stores, pad matrices, or dispatch to a tail implementation. Predication preserves correctness but may leave lanes inactive. Padding performs extra arithmetic and consumes memory. A specialized tail path adds code and dispatch cost. Measure useful FLOPs separately from executed FLOPs so edge waste remains visible.

Worked example: Estimating reuse in a blocked GEMM

A block computes a 128×128 output tile using K chunks of 32 in FP16.

For one K chunk, load 128×32 elements of A and 32×128 of B, about 16 KiB total in FP16.
The staged data supports 128×128×32 multiply-accumulates, conventionally about 1.05 million floating-point operations. The shared tile enables substantial work per high-bandwidth-memory byte.
Divide 1.05 million FLOPs by 16 KiB to obtain about 64 FLOPs per input byte for this tile step, before output traffic. Compare this number with the naive approximation of 0.5 FLOP per byte.
Retain 16,384 output accumulators across threads or warps. Check register distribution and whether the block remains resident with the shared allocation.
For a 256-thread block, the output tile averages 64 accumulator elements per thread. If accumulators are FP32, they alone require roughly 64 registers per thread before address and operand registers are counted.
Sweep smaller and larger tiles against square, narrow, and tail shapes. Record HBM traffic, achieved matrix throughput, occupancy, and wasted padded work.

Conclusion: Blocking realizes the algorithm's reuse, but the best tile is the one that balances locality with resource capacity and enough independent output tiles.

Common errors

Assuming any shared-memory use is tiling. Reuse must be counted; a staged value consumed once only adds a copy.
Choosing the largest possible tile. Resource cliffs and insufficient grid parallelism can erase its intensity advantage.
Comparing only aligned square matrices. Production GEMM families include narrow, batched, grouped, and ragged shapes.
Counting cache hits as guaranteed reuse without examining the launch order and working set. Explicit shared-memory tiling provides a defined reuse scope.

Reference summary

A naive kernel can assign one output element to each thread and repeatedly load a row of A and column of B. Adjacent output threads reuse much of the same data, but without on-chip staging that reuse may depend on caches. A blocked kernel cooperatively loads A and B tiles, synchronizes, performs many multiply-accumulates, and advances along K.

text

for each K tile:
  cooperatively load A[Mblock, Ktile] and B[Ktile, Nblock]
  synchronize the block
  update register accumulators for C[Mblock, Nblock]
  synchronize before shared storage is reused

The block retains output accumulators while it advances through K panels.

tile work ≈ 2BM·BN·BK; tile input ≈ (BM·BK + BK·BN)·bytes

Larger BM and BN increase reuse but consume registers and shared memory.

CTA tile: the output region owned by one thread block.
Warp tile: a subdivision assigned to a warp.
Instruction tile: the shape consumed by scalar FMA or a matrix instruction.
K tile: the reduction slice staged at each pipeline step.
Epilogue: scaling, bias, activation, quantization, or layout applied before the final write.

Knowledge check

Why does increasing a square tile initially improve GEMM but eventually reduce or reverse the gain?

Show the answer guide

Increasing a square GEMM tile initially improves reuse: each loaded A and B value participates in more multiply-accumulate operations, so high-bandwidth-memory bytes per output element fall.
Larger tiles also improve the ratio of interior work to indexing, synchronization, and boundary overhead when the matrix dimensions fit the tile well.
The gain stops when the tile requires too much shared memory, too many registers, or too many threads. Excess resource use can reduce the number of resident blocks and the scheduler's ability to hide memory and instruction latency.
A large tile can also create wasted work on edge tiles, poor warp-level shapes, bank conflicts, or insufficient parallel blocks for small matrices. The best tile is therefore shape- and device-specific rather than the largest tile that compiles.

Practice

Find the tile-size optimum for one GEMM family

Implement or parameterize a tiled FP32 GEMM for M=N=K=2048 with square block tiles of 8, 16, 32, and 64 where hardware limits permit. Add boundary handling and repeat the sweep for M=2053, N=2039, K=2048. Compare outputs against a framework GEMM before benchmarking at least 100 warmed iterations.

Evidence to keep: Submit kernel code, maximum-error correctness tests for both shapes, a benchmark table with latency and achieved GFLOP/s, and profiler data listing shared memory, registers, occupancy, and global-memory traffic for each tile. Include a plot of performance versus tile size and identify the resource or edge effect that reverses the gain.

A simple tiled kernel still alternates load and compute. Software pipelines overlap those phases so data for the next K tile arrives before arithmetic needs it.

References

Sources and further reading

The lesson explanations synthesize primary papers and official references. Follow the originals for proofs, exact APIs, architecture details, and version-sensitive behavior.