I'm the primary author of many Julia libraries geared toward "architecture-specific optimizations", including LoopVectorization.jl and Octavian.jl.
For BLAS-like operations, one of the most important optimizations LoopVectorization.jl does is "register tiling".
While CPUs may have a huge number of actual registers (used for register renaming to enable out of order execution), they have only a small number of named registers. Below, when I talk about number, I'm specifically referring to how many are named.
For floating point vector registers:
- The Apple M1 (ARM) CPU has 32x 128-bit registers (v0-v31)
- x86_64 CPUs with AVX have 16x 256-bit registers (ymm0-ymm15)
- x86_65 CPUs with AVX512 have 32x 512-bit registers (zmm0-zmm31)
The first obvious optimization we can do is use the widest vector registers available. On most CPUs, the throughput in terms of instructions/cycle doesn't change as a function of register width. So using 256-bit registers instead of 128-bit will almost give you 2x performance "for free".
I'll let W mean register width here. With double precision (64-bit) floating point numbers, W would respectively be 2, 4, and 8 for the above CPUs.
Also, not all CPUs support "fused multiply add" (fma) instructions. These perform both multiplication and addition. A low-hanging optimization is to use these when they're available. Most recent CPUs have them, so I'll assume we have then in the below discussion.
When I write intervals below, the intervals are closed-closed intervals, like in Julia (and unlike Python). So m:m+W-1 would mean m, m+1, ..., m+W-1.
I'll also assume that A is an M x K column-major matrix, and that B is a K x N column-major matrix (so that C = A*B is M x N).
For the CPU to be able to operate on data, it must be loaded into a register**. Even if the memory is available in the L1 cache, the highest cache level, loading is likely to still be relatively slow. Many CPUs can perform 2x fused multiply add instructions per clock cycle, vs 1-2 loads per clock cycle (e.g., 2 loads loads that don't cross a cacheline boundary, or 1 load that does).
So if you want to compute A[m:m+W-1,k] * B[k,n] + C[m:m+W-1,n] and A[m+W:m+2W-1,k] * B[k,n] + C[m+W:m+2W-1,n], these 4W arithmetic operations (2 fused multiply adds) will take an average** of 1 clock cycle.
However, these 6 loads could take 3-6 cycles, even if all the memory is in the L1 cache!
So, what can we do here?
The answer is register tiling. Let's take AVX2 CPUs as an example.
We can initialize an 8x6 block of C:
# 8 elements from column 1
C00 = C[m :m+ W-1,n]
C10 = C[m+W:m+2W-1,n]
# 8 elements from column 2
C01 = C[m :m+ W-1,n+1]
C11 = C[m+W:m+2W-1,n+1]
# 8 elements from column 3
C02 = C[m :m+ W-1,n+2]
C12 = C[m+W:m+2W-1,n+2]
# 8 elements from column 4
C03 = C[m :m+ W-1,n+3]
C13 = C[m+W:m+2W-1,n+3]
# 8 elements from column 5
C04 = C[m :m+ W-1,n+4]
C14 = C[m+W:m+2W-1,n+4]
# 8 elements from column 6
C05 = C[m :m+ W-1,n+5]
C15 = C[m+W:m+2W-1,n+5]
This uses up 12 of our 16 named registers, leaving us with 4 unused registers.
Now we loop over the inner-product dimension, K.
On each iteration, we load 8 elements from a column of A, using 2 more registers:
A0 = A[m :m+ W-1,k]
A1 = A[m+W:m+2W-1,k]
And then we unroll across n:n+5:
B0 = B[k,n] # broadcast load
C00 = A0 * B0 + C00
C10 = A1 * B0 + C10
B0 = B[k,n+1] # broadcast load
C01 = A0 * B0 + C01
C11 = A1 * B0 + C11
B0 = B[k,n+2] # broadcast load
C02 = A0 * B0 + C02
C12 = A1 * B0 + C12
B0 = B[k,n+3] # broadcast load
C03 = A0 * B0 + C03
C13 = A1 * B0 + C13
B0 = B[k,n+4] # broadcast load
C04 = A0 * B0 + C04
C14 = A1 * B0 + C14
B0 = B[k,n+5] # broadcast load
C05 = A0 * B0 + C05
C15 = A1 * B0 + C15
So on each iteration of the k loop we're now performing 2 loads from A and 6 loads from B, for 8 loads total.
With these 8 loads, we get to do 12 fused multiply add instructions.
Now, that is enough to keep the fma units busy (instead of waiting on loods), and actually hit close to peak theoretical FLOPS!
That is, it'd be enough to keep them busy if the data is available in a high cache level.
This is an architecture specific optimization because it relies on details like the vector width of the host CPU and the number of registers it actually has available. LoopVectorization.jl has a cost model that it solves using Lagrange multipliers (with the constraints being the number of registers available), finding a continuous solution in the neighborhood of the exact integer solution. From there it can check the neighboring integers to find the values that minimize cost while satisfying the constraints on available registers.
Ideally, the macro kernel would be as large as possible, because the larger it is, the fewer times we need to pass over our data, reducing the amount of memory bandwidth through the cache needed.
Something LoopVectorization doesn't do yet (but will some day!) is cache tiling. Octavian.jl, OpenBLAS, BLIS, IntelMKL, etc -- all real matrix multiplication implementations -- do this. If you look at the benchmark chart in the Octavian README, you can see how LoopVectorization's performance becomes highly erratic as the CPU starts to run out of cache, before eventually dropping like a rock after 1000x1000, due to the lack of cache tiling.
Cache tiling follows a similar idea to register tiling, except it's about memory in your CPU's cache levels, L1, L2, and L3.
With register tiling, by operating in blocks, we were able to reuse the memory loaded from A and B many times (e.g., we did 2 loads from A, and multiplied them with 6 loads from B). Similarly, we operate in a block-wise pattern, so that we're only using blocks from A and B that (on a 3-cache level X86_64 CPU) fit in the L2 and L3 cache. We will need to reload the entirety of the loaded block of the A matrix into registers for every n:n+5 strip of the B-block. That can be a lot of reloads, but because the memory is already in the L2 cache, this should be relatively quick. If we used a larger block of A, we'd constantly be fetching from the L3 cache or the RAM instead!
Sizes of the caches are again a CPU-specific detail.
** x86_64 CPUs let you load and operate in the same instruction, but the loading must still happen.
*** By average I'm referring to throughput. The actual time to complete will be several clock cycles, generally 4 or 5. But many fused multiply adds can be in progress at the same time, so that most CPUs can complete an average of 2 of them per clock cycle. The Apple M1 can do 4/cycle.
Another, somewhat different class of CPU optimizations is taking advantage of specific instructions available. For example, I wrote an exp2 implementation specialized for AVX512. Compared to the AVX version, aside from using larger vectors, it's able to use the AVX512-specific instructions like vscalefpd, and imperm2pd for a fast 16-element lookup table, while the AVX version uses slow gather instructions and a 256-element table. Shuffle instructions (as well as wide enough vectors to use them as a table) are an AVX512-feature, so "small tables" as an implementation strategy is only an option for AVX512.
It needs 2 fma units to really improve performance over the larger table gather-based version, though. A smaller table means you need a bigger polynomial, which requires a lot of fma.
Things like vscalefpd, vgetmantpd, vgetexppd are a bit like "fma" in that they're a faster way of doing what you have to do with multiple instructions on CPUs without AVX512.
