Optimizing matmul on CPU

Introduction and Setup

This blog starts with a naive implementation of matmul in C and optimizes it one step at a time.

Benchmarking is done on the following machine with a 10 core CPU (4 Performance cores+6 Efficiency cores).

Algorithmic complexity of Matrix Multiplication: Calculating the FLOPs required

Matrix multiplication is ubiquitous in many areas of computer sciene. As matrices grow in size, the total FLOPs required to calculate the matmul grow cubically¹. Let’s see how.

Consider two matrices $A (i \times k)$ and $B (k \times j)$. The product of $A$ and $B$, $AB$ is a matrix $C$ of shape $(i \times j)$.

For simplicity and without loss of generality, let’s consider the matrices to be square so we have $A (n \times n)$, $B (n \times n)$ and their product $C (n \times n)$. One element $(c_1, c_2)$ of $C$ is defined as:

$$c_{c1,c2} = \sum_{x=1}^{n} a_{c1,x} b_{x,c2}$$

This is a total of 2n - 1 floating point operations (FLOPs) required to calculate one element of C – $n$ multiplication ops + $(n-1)$ addition ops. A total of $n^2$ elements need to be calculated in $C$ and hence the total FLOPs:

$$(2n - 1) n^2 = 2n^3 - n^2$$

Asymptomatically, $n^2$ becomes pretty negligible in comparison to $2n^3$ and hence can be ignored so the total FLOPs required is roughly $2n^3$. Therefore, the computational complexity is $O(n^3)$.

For this blog, I am considering two matrix sizes $N=4096$ and $N=8192$ which is roughly 137 and 1099 total GFLOPs respectively. To benchmark, I calculate the time required to perform this matmul using numpy arrays of those sizes (np.matmul):

Naive implementation

The matmul loop is this:

for (int i = 0; i < N; i++) {
    for (int j = 0; j < N; j++) {
        for (int k = 0; k < N; k++) {
            C[i][j] += A[i][k] * B[k][j];
        }
    }
}

When complied with the -O3 flag² which is the maximum level with safe optimizations, the latencies are as follows:

The speedup column for all further tables is calculated with respect to the row (optimization technique) just above so it simply gives an idea of the amount of improvement we get with a new intervention.

Full compilation command below:

gcc -Wall -O3 sgemm-cpu/matmuls/naive.c -o sgemm-cpu/matmuls/naive

All further optimizations use the same compilation flags.

Minor optimization: Avoiding unnecessary memory accesses

A minor optimization to the naive implementation is avoiding multiple reads and writes of intermediate partial sums from and to the memory/cache, like so:

for (int i = 0; i < N; i++) {
    for (int j = 0; j < N; j++) {
        float running_sum = 0.0;
        for (int k = 0; k < N; k++) {
            running_sum += A[i][k] * B[k][j];
        }
        C[i][j] = running_sum;
    }
}

That helps very little for $N=4096$ but quite a lot for $N=8192$ ³:

This optimization simply causes the compiler to keep the partial sum in the registers only and write back to the memory only once when the full sum is done. Meaning in this case the compiler doesn’t issue separate instructions to store the partial sum back to the memory after every loop iteration and to load it back on the next iteration; and that reduces some latency ⁴.

Loop reordering

The naive implementation follows the most natural mathy way to calculate a matmul $C = AB$ — element $C[0][0]$ is fully calculated first via a scalar product of the first row of $A$ with the first column of $B$, element $C[0][1]$ is fully calculated next via the scalar product of the first row of $A$ with the second column of $B$, and so on.

There are two key points here:

1. Languages like C store matrices in the memory in a row major format like this:

If we now follow the calculation of $C[0][0]$ in code, it’s equivalent to completing the inner-most $k$ loop for $i=0, j=0$.

The first part of the figure below uses 3 shades of red to color the values from $A, B, C$ that are accessed in the code during the calculation of $C[0][0]$. Focus on the location of these values in the memory by following the color.

The figure quickly makes it clear that the accessed values for $A, C$ are close to each other in the memory while they’re further apart for $B$. Now, remember that data is fetched from the memory in the caches in the granularity of cache lines. One cache line on my machine is of size 128 bytes which is equivalent to 32 single precision values.

So on the very first loop iteration $i=0, j=0, k=0$, when $A[0][0], B[0][0], C[0][0]$ are fetched from the memory, the cache looks something like:

This is simply because a cache line loads contiguous values from the memory where the matrices are stored in a row major format. On the second iteration $i=0, j=0, k=1$, we need $A[0][1], B[1][0], C[0][0]$ and while $A[0][1]$ and $C[0][0]$ are found in the cache, we have a miss for $B[1][0]$ that we need to fetch from the memory. And this holds for each iteration of the k-loop for $i=0, j=0$. Easy to figure out why – our cache line is only 128 bytes (32 floats) while each subsequent loop iteration accesses a value from B that’s 4096 values apart from the value accessed in the last iteration. And this high cache miss rate explains the high latency we saw w the naive implementation.

What about the subsequent iterations for different values of j and i? Say $i=0, j=1, k=0$ when we need $B[0][1]$ that was a part of the cache line we loaded on $i=0, j=0, k=0$ — that cache line most probably would get evicted from the cache till now. Hence, the cache miss rate still remains high.

👉 👉 👉 This is actually a central recurring point of discussion: Some data that we load once gets evicted by the time we need it again making things go slower. In a few upcoming sections, we’ll be revolving around this point making optimizations that aim to reuse the data in subsequent iterations by using it as much as needed (and possible) before its eviction.

2. The second insight is not too difficult to understand – changing the loop orders (eg. $jik$ or $jki$ etc. instead of the most natural $ijk$) still gives the correct matmul. Given that, realise that some loop orders have a better overall cache hit rate as compared to the naive $ijk$ order (btw some orders also have a worse rate than $ijk$!). And hence just by changing the orders, we’ll be able to reduce the latency by not having to make as many high latency accesses to the memory. Experimenting w different orders, the lowest latency corresponds to order ikj as follows:

for (int i = 0; i < N; i++) {
    for (int k = 0; k < N; k++) {
        float a_ik = A[i][k];
        for (int j = 0; j < N; j++) {
            C[i][j] += a_ik * B[k][j];
        }
    }
}

Auto-vectorization

Apart from a better cache hit rate, the above implementation also allows the compiler to auto-vectorize different iters of the inner-most j-loop using SIMD instructions. This follows the fact that different iterations of the j-loop are completely independent of each other.

Cache Blocking / Tiling

The long story short about tiling is the same as loop reordering: it’s all about enabling a better cache reuse.

Let’s try to understand tiling by doing something boring. For the best loop order $ikj$, we start by looking at the values that are accessed by the inner $k, j$ loops for different values of $i$.

$(i=0, k=0-4095, j=0-4095)$

• $A[0][0], A[0][1], … A[0][4095]$ –> A’s 1st row
• $C[0][0], C[0][1], … C[0][4095]$ –> C’s 1st row
• $B[0][0], B[0][1], … B[0][4095]$ –> B’s 1st row

…

• $B[4095][0], B[4095][1], … B[4095][4095]$ –> B’s 4095th row

$(i=1, k=0-4095, j=0-4095)$

• $A[1][0], A[1][1], … A[1][4095]$ –> A’s 2nd row
• $C[1][0], C[1][1], … C[1][4095]$ –> C’s 2nd row
• $B[0][0], B[0][1], … B[0][4095]$ –> B’s 1st row

…

• $B[4095][0], B[4095][1], … B[4095][4095]$ –> B’s 4095th row

What do we infer?

👉 At every iteration of the $i$ loop, the code accesses a different row of A and C but all the rows of B are accessed at each iteration.

So far so good. Let’s now take an example with smaller 8x8 matrices and see what caches actually look like with our best loop order. For the example, assume a fully associative cache following an LRU (least recently used) policy with a cache line of 32 bytes = 8 floats, and a cache size of 160 bytes = 5 cache lines.

The (terrible) figure below shows the values accessed in sequence for $i=0$, and the state of the cache at each step. Red denotes a cache miss requiring a cache update and blue denotes a cache hit with no cache update required. Note that by the time we end $k=2$ for, our cache is full.

The figure below shows the program resumed from $k=3$. Also shown is the cache state at the end of $i=0$ after multiple evictions. Note how the first 5 rows of B had to be evicted even before reaching $i=1$.

Now what? Start $i=1$. The same story gets repeated and below are the various cache states; clearly the entire B matrix needs to be fetched from the memory to be able to finish $i=1$! Not even a single row found in cache when needed despite the fact that all of them were loaded in the previous iteration $i=0$.

This was a toy example where 1 cache line = 1 row, and no single row of B was found in the cache. In practice, the number of caches, their sizes, eviction policies, cache line sizes, associativity all matter in determining what rows/values are found but the general gist stated below remains the same.

The gist: At each new value of i, we’re having to load all or some rows (/partial rows depending on the cache line) of B from the high latency memory despite having them loaded in the previous iteration of i. We should definitely do better!

k-tiling

Think about this: how about there was no eviction and hence no need to reload values of B at each new iteration of i? It’d have been great except that caches are limited and we cannot increase their size. How do we make peace with eviction? By using the data as much as we want before it’s evicted. Meaning we load a few rows of B in the cache and use them for all the values of i so that we never require them again? Then load the next few and again use them for all the values of i. The same until we’re done with all the rows of B. This is where the word tile comes from – we’re loading and utilizing one tile fully before evicting it.

What I’ve described is tiling on the k-loop only. In code:

for (int k_tile=0; k_tile<N; k_tile+=TILE_SIZE) {
    for (int i=0; i<N; i++) {
        int kend = (k_tile + TILE_SIZE > N) ? N : k_tile + TILE_SIZE; // check if this is expensive than an if
        for (int k=k_tile; k<kend; k++) {
            float a_ik = A[i][k];
            for (int j=0; j<N; j++) {
                C[i][j] += a_ik * B[k][j];
            }
        }
    }
}

After doing a good amount of search over different values of TILE_SIZE, I found k-tiling didn’t really yield any benefit on my machine both for $N=4096$ and $N=8192$. This is probably because of the already large cache sizes of my machime.

ijk tiling

Let’s now separately understand the purpose of tiling rest of the two loops - i and j.

Tiling on the j-loop : With k-tiling, rather than loading all the rows of $B$ in the cache for each iteration of $i$, we’re covering a particular number of rows at a time that fit in the cache. But for these rows, we’re still loading all the columns of $B$ (all values of $j$)! In the oversimplified example above, 1 row = 1 cache line but that obviously isn’t the case with bigger matrices and hence, it’d further benefit for the cache hit rate to also tile on the $j$ loop.

Tiling on the i-loop : Further with both $k$ and $j$ tiled, we’d still be needing all the rows of C (all values of i, but ofc not full rows due to j-tiling) in the cache multiple times for different values of k_tile and j_tile. And that might again lead to cache misses for certain values of $C$. And hence, it also benefits to tile on the $i$ loop which, combined with $j$ tiling, effectively translates to taking a sub-matrix of C and finishing all the calculations for it (covering all values of k) before proceeding to another sub-matrix.

In code, this looks like:

for (int i_tile = 0; i_tile < N; i_tile += TILE_I) {
        int iend = (i_tile + TILE_I < N) ? i_tile + TILE_I : N;

        for (int j_tile = 0; j_tile < N; j_tile += TILE_J) {
            int jend = (j_tile + TILE_J < N) ? j_tile + TILE_J : N;

            /* for a specific tile of C (i_tile:iend, j_tile:jend) we now cover all values of k in our already
            discussed k-tiled manner */

            for (int k_tile = 0; k_tile < N; k_tile += TILE_K) {
                int kend = (k_tile + TILE_K < N) ? k_tile + TILE_K : N;

                for (int i = i_tile; i < iend; i++) {
                    for (int k = k_tile; k < kend; k++) {
                        float a_ik = A[i][k];
                        for (int j = j_tile; j < jend; j++) {
                            C[i][j] += a_ik * B[k][j];
                        }
                    }
                }
            }
        }
    }

With ijk tiling, the results are as follows:

The best tile sizes for $N=4096$ is 128, 256, 128 for ikj respectively, and for $N=8192$ is 128 for all ikj.

Also note that the inner most j-loop can be vectorized - another optimization the compiler does by itself (check assembly).

Register Blocking

Tiling enables excellent cache hit for all three matrices. But let’s focus on C[i][j] - at each new iteration of the $k$ and $k_tile$ loops, $C[i][j]$ needs to be loaded from and stored back to the cache. A better way can be to work out an even smaller tile (micro-tile) of C such that it fits entirely into the registers eliminating the need to do multiple repetitive reads from and writes to the cache. Basically, just how we make the code cache-optimized reducing main memory reads/write, we now aim to do the same thing for registers that are closest to the ALUs. Consequently, this is known as register blocking and a simple implementation of it goes:

for (int i_tile = 0; i_tile < N; i_tile += TILE_I) {
        int iend = (i_tile + TILE_I < N) ? i_tile + TILE_I : N;

        for (int j_tile = 0; j_tile < N; j_tile += TILE_J) {
            int jend = (j_tile + TILE_J < N) ? j_tile + TILE_J : N;

            for (int k_tile = 0; k_tile < N; k_tile += TILE_K) {
                int kend = (k_tile + TILE_K < N) ? k_tile + TILE_K : N;

                for (int i = i_tile; i < iend; i += IR) {
                    for (int j = j_tile; j < jend; j += JR) {
                        float c00 = C[i + 0][j + 0];
                        float c01 = C[i + 0][j + 1];
                        float c02 = C[i + 0][j + 2];
                        float c03 = C[i + 0][j + 3];

                        float c10 = C[i + 1][j + 0];
                        float c11 = C[i + 1][j + 1];
                        float c12 = C[i + 1][j + 2];
                        float c13 = C[i + 1][j + 3];

                        float c20 = C[i + 2][j + 0];
                        float c21 = C[i + 2][j + 1];
                        float c22 = C[i + 2][j + 2];
                        float c23 = C[i + 2][j + 3];

                        float c30 = C[i + 3][j + 0];
                        float c31 = C[i + 3][j + 1];
                        float c32 = C[i + 3][j + 2];
                        float c33 = C[i + 3][j + 3];
                    
                        for (int k = k_tile; k < kend; k++) {
                            c00 += A[i + 0][k] * B[k][j + 0];
                            c01 += A[i + 0][k] * B[k][j + 1];
                            c02 += A[i + 0][k] * B[k][j + 2];
                            c03 += A[i + 0][k] * B[k][j + 3];

                            c10 += A[i + 1][k] * B[k][j + 0];
                            c11 += A[i + 1][k] * B[k][j + 1];
                            c12 += A[i + 1][k] * B[k][j + 2];
                            c13 += A[i + 1][k] * B[k][j + 3];

                            c20 += A[i + 2][k] * B[k][j + 0];
                            c21 += A[i + 2][k] * B[k][j + 1];
                            c22 += A[i + 2][k] * B[k][j + 2];
                            c23 += A[i + 2][k] * B[k][j + 3];

                            c30 += A[i + 3][k] * B[k][j + 0];
                            c31 += A[i + 3][k] * B[k][j + 1];
                            c32 += A[i + 3][k] * B[k][j + 2];
                            c33 += A[i + 3][k] * B[k][j + 3];
                        }
                        C[i + 0][j + 0] = c00;
                        C[i + 0][j + 1] = c01;
                        C[i + 0][j + 2] = c02;
                        C[i + 0][j + 3] = c03;

                        C[i + 1][j + 0] = c10;
                        C[i + 1][j + 1] = c11;
                        C[i + 1][j + 2] = c12;
                        C[i + 1][j + 3] = c13;

                        C[i + 2][j + 0] = c20;
                        C[i + 2][j + 1] = c21;
                        C[i + 2][j + 2] = c22;
                        C[i + 2][j + 3] = c23;

                        C[i + 3][j + 0] = c30;
                        C[i + 3][j + 1] = c31;
                        C[i + 3][j + 2] = c32;
                        C[i + 3][j + 3] = c33;
                }
            }
        }
    }
}

This gives an avg latency of 4.3s for $N=4096$ - worse than the previous tiled version. Why did this happen? Two major reasons: First, we lost the nicley vectorizable inner-most j loop in this version. Second, the cache hit-rate for B values becomes v poor. Let’s fix both of them:

SIMD

In the current implementation, a lot of scalar calculations inside the k-loop can be performed as vector ops. Specifically, the 4x4 microtile of C can be put into 4 128-bit NEON registers (on apple silicon) each containing one row of C (4floats=128bit). Another register can hold the four contiguous B values B[k][j] through B[k][j+3] and so the 16 scalar FMA instructions now become four 4-wide vector FMA instructions, like so:

for (int i_tile = 0; i_tile < N; i_tile += TILE_I) {
        int iend = (i_tile + TILE_I < N) ? i_tile + TILE_I : N;

        for (int j_tile = 0; j_tile < N; j_tile += TILE_J) {
            int jend = (j_tile + TILE_J < N) ? j_tile + TILE_J : N;

            for (int k_tile = 0; k_tile < N; k_tile += TILE_K) {
                int kend = (k_tile + TILE_K < N) ? k_tile + TILE_K : N;

                for (int i = i_tile; i < iend; i += IR) {
                    for (int j = j_tile; j < jend; j += JR) {

                        float32x4_t c0 = vld1q_f32(&C[i+0][j]);
                        float32x4_t c1 = vld1q_f32(&C[i+1][j]);
                        float32x4_t c2 = vld1q_f32(&C[i+2][j]);
                        float32x4_t c3 = vld1q_f32(&C[i+3][j]);

                        for (int k = k_tile; k < kend; k++) {
                            float32x4_t b = vld1q_f32(&B[k][j]);

                            c0 = vfmaq_n_f32(c0, b, A[i+0][k]);
                            c1 = vfmaq_n_f32(c1, b, A[i+1][k]);
                            c2 = vfmaq_n_f32(c2, b, A[i+2][k]);
                            c3 = vfmaq_n_f32(c3, b, A[i+3][k]);
                        }

                        vst1q_f32(&C[i+0][j], c0);
                        vst1q_f32(&C[i+1][j], c1);
                        vst1q_f32(&C[i+2][j], c2);
                        vst1q_f32(&C[i+3][j], c3);
                    }
                }
            }
        }
    }

(Note this isn’t SIMD vectorizing the k-loop, like how the compiler did for j-loop in ijk-tiled version, but rather the instructions inside one iteration of it.)

Packing

To fix the poor cache hit rate, a simple solution is to arrange values of A and B in the pattern they’re accessed in the inner loops (the microkernel as it’s called) - this is called packing and while it introduces some extra read/write costs, the computation benefits vastly from improved data reuse.

One implementation can look like:

for (int i_tile = 0; i_tile < N; i_tile += TILE_I) {
        int iend = (i_tile + TILE_I < N) ? i_tile + TILE_I : N;

        for (int j_tile = 0; j_tile < N; j_tile += TILE_J) {
            int jend = (j_tile + TILE_J < N) ? j_tile + TILE_J : N;

            for (int k_tile = 0; k_tile < N; k_tile += TILE_K) {
                int kend = (k_tile + TILE_K < N) ? k_tile + TILE_K : N;

                // pack A:
                int a_pos = 0;
                for (int i = i_tile; i < iend; i += IR) {
                    for (int k = k_tile; k < kend; k++) {
                        A_pack[a_pos++] = A[i + 0][k];
                        A_pack[a_pos++] = A[i + 1][k];
                        A_pack[a_pos++] = A[i + 2][k];
                        A_pack[a_pos++] = A[i + 3][k];
                    }
                }

                // pack B
                int b_pos = 0;
                for (int j = j_tile; j < jend; j += JR) {
                    for (int k = k_tile; k < kend; k++) {
                        B_pack[b_pos++] = B[k][j + 0];
                        B_pack[b_pos++] = B[k][j + 1];
                        B_pack[b_pos++] = B[k][j + 2];
                        B_pack[b_pos++] = B[k][j + 3];
                    }
                }

                int a_micro_start = 0;
                for (int i = i_tile; i < iend; i += IR) {
                    int b_micro_start = 0;
                    for (int j = j_tile; j < jend; j += JR) {
                        float32x4_t c0 = vld1q_f32(&C[i+0][j]);
                        float32x4_t c1 = vld1q_f32(&C[i+1][j]);
                        float32x4_t c2 = vld1q_f32(&C[i+2][j]);
                        float32x4_t c3 = vld1q_f32(&C[i+3][j]);

                        float *a_ptr = &A_pack[a_micro_start];
                        float *b_ptr = &B_pack[b_micro_start];

                        for (int k = k_tile; k < kend; k++) {

                            float32x4_t a = vld1q_f32(a_ptr);
                            float32x4_t b = vld1q_f32(b_ptr);

                            c0 = vfmaq_n_f32(c0, b, vgetq_lane_f32(a, 0));
                            c1 = vfmaq_n_f32(c1,b, vgetq_lane_f32(a, 1));
                            c2 = vfmaq_n_f32(c2, b, vgetq_lane_f32(a, 2));
                            c3 = vfmaq_n_f32(c3, b, vgetq_lane_f32(a, 3));

                            a_ptr += IR;
                            b_ptr += JR;
                        }
                        vst1q_f32(&C[i+0][j], c0);
                        vst1q_f32(&C[i+1][j], c1);
                        vst1q_f32(&C[i+2][j], c2);
                        vst1q_f32(&C[i+3][j], c3);

                        b_micro_start += (kend - k_tile) * JR;
                    }
                    a_micro_start += (kend - k_tile) * IR;
                }
            }
        }
    }

Best tile sizes for this implementation:

Multithreading

Until now, the implementations leveraged a single CPU cores. To make use of multiple available cores, the job now is to find ways to divide (parallelize) the computation of C among threads that the OS could then schedule on different cores. This can be done by letting one thread handle the comptation of one specific C tile (realise that the calculation of different C tiles is independent).

#pragma omp parallel for collapse(2) num_threads(8) default(none) shared(A, B, C, A_pack, B_pack)
    for (int i_tile = 0; i_tile < N; i_tile += TILE_I) {
        int iend = (i_tile + TILE_I < N) ? i_tile + TILE_I : N;

        for (int j_tile = 0; j_tile < N; j_tile += TILE_J) {
            int jend = (j_tile + TILE_J < N) ? j_tile + TILE_J : N;

            for (int k_tile = 0; k_tile < N; k_tile += TILE_K) {
                int kend = (k_tile + TILE_K < N) ? k_tile + TILE_K : N;

                // pack A: 
                int a_pos = 0;
                for (int i = i_tile; i < iend; i += IR) {
                    for (int k = k_tile; k < kend; k++) {
                        A_pack[a_pos++] = A[i + 0][k];
                        A_pack[a_pos++] = A[i + 1][k];
                        A_pack[a_pos++] = A[i + 2][k];
                        A_pack[a_pos++] = A[i + 3][k];
                    }
                }

                // pack B
                int b_pos = 0;
                for (int j = j_tile; j < jend; j += JR) {
                    for (int k = k_tile; k < kend; k++) {
                        B_pack[b_pos++] = B[k][j + 0];
                        B_pack[b_pos++] = B[k][j + 1];
                        B_pack[b_pos++] = B[k][j + 2];
                        B_pack[b_pos++] = B[k][j + 3];
                    }
                }

                int a_micro_start = 0;
                for (int i = i_tile; i < iend; i += IR) {
                    int b_micro_start = 0;
                    for (int j = j_tile; j < jend; j += JR) {
                        float32x4_t c0 = vld1q_f32(&C[i+0][j]);
                        float32x4_t c1 = vld1q_f32(&C[i+1][j]);
                        float32x4_t c2 = vld1q_f32(&C[i+2][j]);
                        float32x4_t c3 = vld1q_f32(&C[i+3][j]);

                        float *a_ptr = &A_pack[a_micro_start];
                        float *b_ptr = &B_pack[b_micro_start];

                        for (int k = k_tile; k < kend; k++) {

                            float32x4_t a = vld1q_f32(a_ptr);
                            float32x4_t b = vld1q_f32(b_ptr);

                            c0 = vfmaq_n_f32(c0, b, vgetq_lane_f32(a, 0));
                            c1 = vfmaq_n_f32(c1,b, vgetq_lane_f32(a, 1));
                            c2 = vfmaq_n_f32(c2, b, vgetq_lane_f32(a, 2));
                            c3 = vfmaq_n_f32(c3, b, vgetq_lane_f32(a, 3));

                            a_ptr += IR;
                            b_ptr += JR;
                        }
                        vst1q_f32(&C[i+0][j], c0);
                        vst1q_f32(&C[i+1][j], c1);
                        vst1q_f32(&C[i+2][j], c2);
                        vst1q_f32(&C[i+3][j], c3);

                        b_micro_start += (kend - k_tile) * JR;
                    }
                    a_micro_start += (kend - k_tile) * IR;
                }
            }
        }
    }

This parallelization strategy also doesn’t require any communication among the several threads.

Best tile sizes: