Fast BLAS-like matrix multiplication in Rust

Question

I have some matrix multiplication code I am writing for a linear algebra library I'm working on. The goal is to emulate BLAS (and achieve BLAS-like performance) without using a BLAS library, as my use-case only needs 2 or 3 BLAS operations and using a BLAS library would be pulling in too many unnecessary dependencies. The code is as follows:

// Performs the operation A.T -> B
// where A dims = M x N
fn transpose(a: &[f64], m: usize, n: usize) -> Vec<f64> {
    let mut b = vec![0.0; n * m];

    for i in 0..n {
        for j in 0..m {
            b[(m * i) + j] = a[(n * j) + i]
        }
    }
    b
}

// Optimized matrix multiplication
// Performs AB -> C
// where A dims = N x M
// and B dims = M x P
// and C dims N x P
fn sgemm(n: usize, m: usize, p: usize, a: &[f64], b: &[f64]) -> Vec<f64> {
    let mut c = vec![0.0; n * p];

    // Matrix multiply by transpose
    // to speed up performance
    let b_transpose = transpose(&b, m, p);

    for i in 0..n {
        for j in 0..p {
            for k in 0..m {
                c[(p * i) + j] += a[(m * i) + k] * b_transpose[(p * k) + j];
            }
        }
    }
    c
}

fn main() {
    let a = vec![1.,2.,3.,4.,5.,6.]; // 3 x 2 matrix
    let b = vec![1.,1.,1.,1.,1.,1.]; // 2 x 3 matrix
    let c = sgemm(3, 2, 3, &a, &b);
    println!("{:?}", c);
}

I am aware that performance can likely be improved, as my level of optimizations is not very high. Please let me know of any suggested improvements to the code.

Answer 1

sgemm

The s in sgemm stands for single-precision. This code works with f64, so it's a dgemm.

Performance

Unfortunately there won't be anything left of this code when you're done upgrading it. Everything will change. So as far as reviewing goes, that was it. Everything after this point, is a description of what you need to do to engineer the replacement.

First decide whether to focus on small matrices, such as in your example, or big ones. They are completely different in nature. For small matrices, the challenge is in writing "bespoke SIMD code" for every particular size that you care about, making every instruction count (there won't be many instructions, so each one is individually more important).

Medium to large matrices are a whole different game, which the rest of this answer is about.

By the way don't expect to compete with eg Intel MKL, on the other hand you can get 80% of the way there or so no problem. Some people treat this as some sort of mythical problem that mere mortals cannot solve, which may be true if your goal is to get all the way to the end, but getting most of the way there is not so bad.

The ingredients

Here's a shopping list for efficient matrix multiplication:

• Tiling, with repacking (copying a tile into its own memory). The tiling is to make good use of cache, the repacking is to avoid TLB misses. Transposing b helps a bit, but the real solution is tiling, and actually transposing gets in the way of other things that you'll end up wanting to do.
  An obvious question that comes up with tiling is what dimensions the tiles should have. The best choice varies depending on cache size and such..
• A sufficient amount of unrolling (with independent accumulators, otherwise it still doesn't help) to keep the multipliers/FMA units busy. Note that this unrolling cannot be purely of the inner loop. Unrolling the j-loop is necessary to be able to vectorize properly at all, and you will need to unroll it more to reach the amount of independent chains of computation that you need to keep the FMA units busy.
• SIMD. A somewhat common mistake is trying to apply SIMD to speed up individual dot-products. That has a disadvantage: it requires a horizontal addition (adding up the components of a single SIMD-vector, henceforth just "vector") in the end. Alternatively, you can apply SIMD to compute several different dot-products: you can load a small sub-row from b into a vector (which is why it shouldn't be transposed), broadcast an entry of a into a vector, multiply them together, then add the whole resulting vector into c (even though I describe it like that, in practice you'll want to load some chunks from c into registers first, add to the registers, then store back the results). That way that are no horizontal additions, there is a broadcast which feels a little horizontal but it's much cheaper.

Here's another question where I worked it out a bit more (but in C, so adjust as required), not showing the tiler/repacker that would sit at a higher level to juggle the tiles and make calls to that function. In that answer I didn't go into hard-core optimization, it's just some basic implementation of the principles that you need.

Even more detail is available in the classic paper Anatomy of High-Performance Matrix Multiplication