MLSYS ENGINEERING

7.2. Optimized tiling

Now, let's see if we can do better with a more optimized tiling strategy. We are still multiplying two 4 by 4 matrices A and B to get C = AB. We will use arithmetic intensity to measure our improvement. Arithmetic intensity is the total FLOPs divided by the total bytes transferred, introduced in Arithmetic intensity. Given that different tiling strategies should have the same amount of computation for the same matrix multiplication, the fewer bytes transferred, the higher the arithmetic intensity is.

As shown in the following figure, we have P1 to P4 representing four different processes. They all run in parallel on different SMs. We tile the 4 by 4 matrix C into 4 non-overlapping 2 by 2 blocks. Each process is responsible for computing one block of C.

To compute one block of C, each process needs to load two rows of A (8 floats) and two columns of B (8 floats) to shared memory. Then, each process computes the 2 by 2 block of C (4 floats) using the data in shared memory, and writes the block back to global memory.

P1 P2 P3 P4 A B C
Figure 23. Optimized matrix multiplication tiling strategy.

For each process, there are exactly 20 (= 8 + 8 + 4) floats transferred, down from 24 in the naive strategy. Each process computes the same 4 output elements, each requiring 4 multiply-adds, for a total of 32 FLOPs. The arithmetic intensity for the naive strategy is 32 FLOPs / (24 floats × 4 bytes) ≈ 0.33 FLOPs/byte, and for the optimized strategy it is 32 FLOPs / (20 floats × 4 bytes) = 0.40 FLOPs/byte.

The improvement is modest on this small example, but the benefit grows with block size. For a large N×N matrix with a B×B block, the arithmetic intensity is roughly B/4 FLOPs/byte. With a 64×64 block, that is about 16 FLOPs/byte, enough to push matmul close to the compute-bound regime on a modern GPU.