Dense neural networks, attention projections, convolution lowering, and many embedding transformations reduce to matrix multiplication or closely related tensor contractions. The mathematical operation is simple: many multiply-accumulate operations produce an output tile. The implementation problem is that the input values must be delivered to arithmetic units at a rate high enough to keep those units busy.
Data Reuse for Matrix Multiplication
The optimization of matrix multiplication is a well studied problem in data engineering. Here, we go through some principles that come from dealing with that problem.
For \(C = A \times B\), each element \(A_{i,k}\) is reused across many columns of \(B\), and each element \(B_{k,j}\) is reused across many rows of \(A\). A straightforward implementation that reloads these values from memory for every output element wastes bandwidth. Efficient implementations therefore compute a tile of \(C\) at a time.
The tile has three important properties:
- It is small enough for on-chip storage such as registers, shared memory, or a hardware tile register file.
- It is large enough to expose many independent multiply-accumulate operations.
- It is arranged so each loaded input value is used multiple times before being evicted.
This is the same locality idea as cache blocking in CPU algorithms, but matrix units make the blocked operation a hardware-level primitive.
Tensor Cores
Modern GPUs expose matrix units through instructions that operate on small fragments, such as warp-level matrix multiply-accumulate instructions. The programmer or compiler arranges fragments of $A$, $B$, and $C$ in registers or shared memory, then issues a matrix instruction that performs many fused operations.
The instruction is not a complete large GEMM. It is a tile primitive. A high-performance GEMM kernel surrounds the primitive with:
- Global memory loads for large matrices.
- Shared-memory staging to coalesce and reuse data.
- Register tiling for per-thread fragments.
- Synchronization between producer and consumer phases within a thread block.
- Epilogue operations such as bias, scaling, activation, or quantization.
Intel AMX follows a different programming model but a similar objective. It provides tile registers and tile operations so CPU cores can execute matrix-style work with reduced instruction overhead. The operating system must save and restore the extended tile state during context switches, which makes AMX partly an ISA feature and partly an OS integration feature.