4
$\begingroup$

I keep coming across phrases like "highly optimized BLAS kernels" with "architecture-specific optimizations", but have never been able to find what exactly these optimizations are, and how one could go about writing their own architecture-specific optimizations.

What is a good way of getting to know different BLAS implementations, their low-level optimizations and methods for architecture specific performance tuning and benchmarking?

I am fairly new to the field, having done only one undergrad course that required us to implement simple sgemm/dgemm subroutines using OpenMP, MPI and CUDA C. Please let me know if I can improve the question in any way.

$\endgroup$
11
  • 1
    $\begingroup$ Robert van de Geijn had an amazing series of blog posts hosted on his academic website, but apparently it is dead now. Though there is the post "How to optimize GEMM" on the github of the research group FLAME (or shpc) which discusses the same thing: github.com/flame/how-to-optimize-gemm/wiki $\endgroup$ May 19, 2021 at 8:59
  • $\begingroup$ You could take a look a the OpenBLAS project for examples. $\endgroup$ May 19, 2021 at 14:35
  • $\begingroup$ github.com/xianyi/OpenBLAS/blob/develop/kernel/x86_64/axpy.S $\endgroup$
    – user14717
    May 21, 2021 at 13:36
  • $\begingroup$ there's one: It's straight assembler. Ironically to my eyes it doesn't look very optimized; I at least expected an fma instruction in there, and moreover thought I might get it does in parallel in ymm registers. $\endgroup$
    – user14717
    May 21, 2021 at 13:37
  • 1
    $\begingroup$ @user14717 While OpenBLAS is a good open-source implementation of BLAS, it doesn't explain what is done in the name of optimization. Michael Lehn of Ulm University has a course on basics of HPC in which he goes through many techniques (low and high level) used to optimize various kernels. Unfortunately, it is in German. van de Geijn's documentation was nice, now his research group has MOOCs and they provide the lecture notes for free if I am not wrong. But those are long, I don't know if the asker is willing to go through a whole course. $\endgroup$ May 21, 2021 at 15:51

2 Answers 2

7
$\begingroup$

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.

$\endgroup$
3
$\begingroup$

Abdullah, thank you for the plug for our materials. We have repackaged these as a Massive Open Online Course (MOOC) on edX titled "LAFF-On Programming for High Performance". It is free for auditors. For info, see http://ulaff.net.

The premier open source implementation of the BLAS is the BLAS-like Library Instantiation Software (BLIS). There is a high level discussion of it in an article in this April's SIAM News: https://sinews.siam.org/Details-Page/blis-blas-and-so-much-more

$\endgroup$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge that you have read and understand our privacy policy and code of conduct.

Not the answer you're looking for? Browse other questions tagged or ask your own question.