Few subroutines have traveled as improbable a road as the general matrix-matrix multiply (GEMM, after the nomenclature pioneered in the decades-old BLAS library). For most of computing history, GEMM was a workhorse of numerical linear algebra – the third level of BLAS, the operation underneath LINPACK, and the yardstick by which the TOP500 ranks the fastest computers on Earth. It was deemed so important that even before deep learning arrived, every CPU and GPU vendor employed a small army (dozens to hundreds of full-time employees) dedicated to optimizing its performance.
With the advent of deep learning, though, matrix multiplication has taken an even more central role in the industry. A fully connected layer is a matrix multiply. The attention blocks that define modern language models are matrix multiplies. Convolutions are lowered to matrix multiplies. Training a large model is, to a first approximation, an enormous sequence of GEMMs run billions of times, and just the electricity bill is a line item that moves markets. If NVIDIA had not added Tensor Cores to their GPUs starting with Volta, dedicated GEMM hardware would have absconded with deep learning clock cycles long ago. As things stand, implementing GEMM tiles in hardware, albeit with heretofore-unheard-of low precisions, has become de rigueur for deep learning hardware designs, including GPUs.
This chapter will take us on an optimization journey for single-precision GEMM (SGEMM), from a naive kernel to one that drives the Tensor Cores. The point is not to beat cuBLAS – rather, to educate about the various strategies needed to best exploit the hardware.
GEMM always has been a workload whose implementations were, not only
architecture-specific, but implementation-specific. Even on CPUs, GEMM
implementations have long been tuned to the exact cache hierarchy of
whichever CPU they were running on. So it is worth noting that with the
introduction of Tensor Cores, NVIDIA abandoned all pretense at
delivering performance portability from one GPU generation to
another. The ergonomic combination of threadIdx/blockIdx
allusions to data parallelism in scalar-looking code and shared memory
with its well-defined access rules, coupled with PTX to empower NVIDIA
to drastically change the machine instruction set without disrupting
application portability, had fueled more than a decade of performance
portability far superior to contemporaneous CPU technologies.
Starting with Tensor Cores, even PTX’s promises of portability were abandoned. The machine instructions at the heart of Tensor Cores – whose complexity would have boggled the minds of RISC proponents of the 1980s – have changed too much too fast for PTX to keep up. Complicated things, the operations performed by the instructions have been changing, not only with ever-narrowing bit widths, but with memory-moving instructions that can be invoked by running kernels, because the near-memory in the SMs simply couldn’t cycle through Tensor Core operands fast enough with the existing hardware and software abstractions.
This chapter only covers up to Ampere capabilities. We’ll add Hopper and Blackwell sections in due course, and each of those topics will deserve another section because both GPU generations added substantive features specifically for GEMM.
All measurements in this chapter were taken on a GeForce RTX 3060
(compute capability 8.6, Ampere) multiplying rectangular
single-precision matrices (M = 4096, N = 2048,
K = 4096), and throughput is reported in GFLOP/s, computed as
2MNK divided by the kernel time. cuBLAS is included throughout
as the professionally tuned reference. One caveat worth stating up
front: on CUDA 12, cuBLAS Sgemm computes in true FP32 by
default – the faster Tensor Core (TF32) path must be requested
explicitly – so the tables below distinguish the two.