chapter five

5 Higher-dimensional kernels

 

This chapter covers

  • Row-major and column-major memory layout, strides, and memory coalescing
  • Writing and optimizing a matrix transpose kernel using shared memory tiling
  • Building aligned and boundary-safe tiled matrix products
  • Reaching Tensor Core matrix instructions through tl.dot
  • Applying grouped L2 ordering, autotuning, correctness tests, and roofline analysis

If you were asked to name the main components of any modern deep learning architecture, you would most likely say the attention mechanism, or a linear layer. Linear layers and attention projections express their main calculation as matrix multiplication. Many convolution implementations reduce their work to matrix multiplication, while softmax, normalization, and activation kernels retain different arithmetic and memory-access patterns. Matrix multiplication therefore provides a practical route from multidimensional indexing to the specialized matrix hardware in a current GPU.

5.1 Row-major and column-major ordering

5.2 Strides: the language of tensor layout

5.2.1 A concrete example

5.2.2 The same memory, two tensors

5.2.3 Stride interpretation and special values

5.2.4 Memory coalescing: why access patterns affect performance

5.3 Matrix addition

5.3.1 Building the first 2D kernel

5.3.2 2D program ID decomposition and tile indexing

5.3.3 Boundary masking for non-tile-aligned shapes

5.3.4 The complete matrix-addition kernel

5.4 Matrix transpose

5.4.1 Naive transpose kernel

5.4.2 Why uncoalesced writes hurt

5.4.3 Staging the transpose through shared memory

5.4.4 Bank conflicts and the one-column padding fix

5.5 Matrix multiplication

5.5.1 Extending matrix multiplication with scaling and accumulation

5.5.2 Naive matrix multiplication kernel

5.5.3 Arithmetic intensity and the roofline view

5.5.4 Reusing values across output elements