MLSYS ENGINEERING

7.4. Tiling memory optimization

We now introduce a more advanced tiling strategy with more flexibility for tile sizes.

The key idea is not to load entire rows of A into the shared memory at once. Instead, we can divide the rows of A into smaller tiles, and load one tile of A into the shared memory at a time. And do the same for B.

Let's see a visualized example. Again, we use the 4 by 4 matrix multiplication C = AB as an example, shown in the following figure.

A B C t=1 t=2 A B C
Figure 26. Memory-optimized matrix multiplication tiling strategy.

To compute one 2 by 2 block of C, we can divide the two rows of A, a 2 by 4 tile, into two 2 by 2 tiles. And do the same for the two columns of B. The 2 by 2 blocks from A and B are loaded into the shared memory at different time steps.

How could this even work? We need a full row of A and a full column of B to compute one element of C, which is the dot product of the two vectors. However, we can divide the dot product into two parts. For example, if a row of A is [1, 2, 3, 4] and a column of B is [5, 6, 7, 8], the dot product would be 1*5 + 2*6 + 3*7 + 4*8. Now, we can divide it into two parts: (1*5 + 2*6) and (3*7 + 4*8). At t=1, we do c = 1*5 + 2*6. At t=2, we do c += 3*7 + 4*8.

In this way, we have split the vectors from A and B into two segments. Therefore, there is no need to load the entire row of A and entire column of B into the shared memory at once. We can arbitrarily decide the tile size we access each time.

Regarding the arithmetic intensity, it is the same as the previous tiling strategy: the same amount of computation and the same total bytes transferred, just split across two time steps instead of one. But unlike the previous strategy, this approach is not constrained by the matrix dimensions. We can freely choose tile sizes that fill shared memory well, which improves memory utilization, reduces the number of batches, and gives the pipeline more time to overlap computation with loading. The result is better performance despite identical arithmetic intensity, confirming the lesson from Shared memory utilization: equal arithmetic intensity does not guarantee equal performance.