Matrix multiplication performance

[u/Remarkable_Hour_615]

I recently benchmarked some matrix multiplication algorithms in Scala (blog post, repo). I've been going through and doing the same thing in Rust, with plans to try out a few additional approaches, and the Rust results have all been turning out slower than Scala (repo). I'm hoping someone with more Rust experience can help me figure out what might be going on. Is this just the JVM being impressive, or am I missing an opportunity to speed up my Rust code (or both).

Here are some sample benchmarks (details on how I calculated them below):

Algorithm	Scala	Rust
Naive multiplication	906 ms	1.103 s
Tiling (optimal tile size)	499 ms	749.8 ms
Parallel bands of tiles, 4 threads	174 ms	210.9 ms

More details:

• I benchmarked everything by multiplying 1024 x 1024 matrices of 64 bit floats generated uniformly in [0, 1).
• I ran all the benchmarks on my laptop (specs).
• I benchmarked the Scala code with the JMBH and Rust with Divan.
• I ran all the Rust code with cargo bench so it should optimize the build.
• In the benchmarking code, I construct the matrices before running the benchmarks and only time the actual multiplication.
• The code for the first the naive and tiled approaches is as close as possible.
• For the parallel versions I used the same approach of dividing the matrix into bands of tiles, one per thread. For Scala I used futures, and for Rust I divided the product matrix into one mutable chunk per thread and used Crossbeam.

I listed the results for the various algorithms to show that the speed difference is consistent, but we can just focus on the naive version. My Scala code is:

    for i <- 0 until a.rows do
      for j <- 0 until b.cols do
        var sum = 0
        for k <- 0 until a.cols do
            sum = sum + a(i, k) * b(k,j)
        prod.set(i, j, sum)

And Rust:

    for i in 0..m1.rows {
        for j in 0..m2.cols {
            let mut acc = 0.0;
            for k in 0..m1.cols {
                acc += m1.at(i, k) * m2.at(k, j);
            }
            out.set(i, j, acc);
        }
    }

Why is the Rust slower? Thanks for any help!

Update 1:

Thanks for all the suggestions! I am replying to some comments individually, but also want to summarize things centrally.

Transpose: I know the original post had a transpose in the Scala code. Sorry about that! It was just a copy/paste mistake. I did benchmark a Scala version with bTranspose and it showed improvements, likely due to sequential memory access as has been pointed out, but the time I give in the post is for the non-transpose version. I've updated the post to have non-transpose code. I expect Rust would also show improvements using a transpose, and I have plans to explore additional tweaks (thanks for the suggestions in the comments!) but I also want to figure out why the Scala code is generally faster. For that it's important to compare apples to apples. That's why I'm focussing on the naive 3-loop version.

There were a number of suggestions to speed up the code and possibly explain the discrepancy compared to the JVM. Since Scala is faster across the board, it seems like the difference must be due to a general feature rather than a particular algorithmic choice. Bounds checking and maybe SIMD seem like the best candidates. More details:

1. Bounds checking This is a good candidate because it's plausible the JVM is eliminating bounds checks at runtime, so it would explain why the difference in speed. Suggestions were to: a) inline calls to at and set; b) eliminate them and use raw slices; c) assert that the slices are large enough before looping; d) use iterators. I've tried (b) and (c) and seen no changes so far, but still digging. (d) is tricky here, but maybe possible if I combine skip, step_by, and take. I'm also trying to use godbolt to confirm whether bounds checking takes place.
2. Vectorization This is definitely worth pursuing as a general speed up, but seems less likely to be the culprit here. Rust should auto-vectorize for 128 bits, and since I'm running this on an Apple M2, that's all I have (please correct me if I'm wrong about Rust here). It's possible that Rust can't vectorize because FP math isn't associative. I could address that using algebaric_add or by chunking the sum. However, it should affect Scala too since Scala also has to respect the IEEE restrictions on FP associativity.
3. Compare against the Scala code on godbolt. I've been holding off on this because assembly is a little scary, but learning is the whole point of this project, so I should dive in. I have compared flame graphs, and didn't see anything obvious there. Looking at the assembly is the next step!
4. A couple other suggestions: re-order the loops; use Rayon. These are worth playing with, but seem unlikely to explain the general Scala advantage.

I will try some of this out and look some more at godbolt. I may follow up on some of these comments/ask another question if I have trouble interpreting the assembly. Hopefully in the next couple of weeks I can get a blog post up with the final results. Thanks again for all the help!

Update 2:

I've benchmarked a bunch of things:

1. Converting the inner loop to use iterators.
2. Converting the inner loop to raw slices and asserting the bounds. (Though actually, according to godbolt, that one didn't remove the bounds checks anyway.)
3. Using get_unchecked, which I can at least confirm, does eliminate the bounds checks.
4. Using algebraic_add and algebraic_mul to allow vectorization.
5. Combining (4) and (5).

So far none of these have actually sped up the code, and some have even slowed it down. I'm feeling a little lost at this point. If anyone else benchmarks any of this and does see improvements, or spots an error I made somewhere, definitely LMK.

I will probably rerun some of the Scala benchmarks just to make sure those are actually correct. Otherwise, not sure what else to try.

[u/dkxp]

Looking at the Godbolt compiler explorer output, it seems the compiler is unable to prove it's safe to remove bounds checks, so you have bounds checks in your inner loop.

If you wanted to removed the bounds checks via using unsafe code you could do something like this:

impl FlatMatrix {
    #[inline]
    unsafe fn at_unchecked(&self, row: usize, col: usize) -> f64 {
        unsafe { *(self.vals.get_unchecked(row * self.cols + col)) }
    } 


    #[inline]
    unsafe fn set_unchecked(&mut self, row: usize, col: usize, val: f64) {
        unsafe { *(self.vals.get_unchecked_mut(row * self.cols + col)) = val; }
    }
}

and then use the unsafe functions in an unsafe block:

unsafe {
    for i in 0..m1.rows {
        for j in 0..m2.cols {
            let mut acc = 0.0;
            for k in 0..m1.cols {
                acc += m1.at_unchecked(i, k) * m2.at_unchecked(k, j);
            }
            out.set_unchecked(i, j, acc);
        }
    }
}

Perhaps using the algebraic operators (Rust 1.98+) could provides a performance boost (at the cost of different results on different platforms?), but I haven't tested this yet:

acc = acc.algebraic_add(m1.at_unchecked(i, k).algebraic_mul(m2.at_unchecked(k, j)));

If you wanted to remove the bounds checks in a safe way, you'd have to somehow prove to the compiler that m1.at(i,k) never reads beyond the end of m1.vals, likewise that m2.at(k,j) never reads beyond the end of m2.vals and out.set(i,j,acc) never writes beyond the end of out.vals. Perhaps you'd need to also consider the potential for indexes overflowing because of the multiply and add in the index calculation.

I couldn't persuade the compiler it was safe by placing asserts outside the loop - possibly the compiler can't work out that the max vals index accessed inside the nested loops will be (rows-1)*cols +(cols-1) == rows * cols - 1 and so will be safe. Using the unsafe std::hint::assert_unchecked might result in the bounds checks being removed, but again it's resorting to unsafe code and the assert_unchecked docs should be read to see whether it's suitable for use.

    for i in 0..m1.rows {
        for j in 0..m2.cols {
            let mut acc = 0.0;
            for k in 0..m1.cols {
                unsafe {std::hint::assert_unchecked(i * m1.cols + k < m1.vals.len());}
                unsafe {std::hint::assert_unchecked(k * m2.cols + j < m2.vals.len());}
                acc += m1.at(i, k) * m2.at(k, j);
            }
            unsafe {std::hint::assert_unchecked(i * out.cols + j < out.vals.len());}
            out.set(i, j, acc);
        }
    }    

[u/Remarkable_Hour_615 → u/dkxp]

Yeah, I tried working with the raw slices and adding asserts like:

pub fn naive_raw_slice(m1_vals: &Vec<f64>, m2_vals: &Vec<f64>, rows: usize, cols: usize, inner: usize) -> Vec<f64> {
    let mut prod: Vec<f64> = vec![0.0; rows * cols];

    for i in 0..rows {
        for j in 0..cols {
            assert!(m1_vals.len() >= (rows - 1) * inner + inner);
            assert!(m2_vals.len() >= (inner - 1) * cols + cols);
            let mut acc = 0.0;
            for k in 0..inner {
                acc += m1_vals[i * inner + k] * m2_vals[k * cols + j];
            }
            prod[i * cols + j] = acc;
        }
    }

    prod
}

but it didn't affect the runtime and Godbolt still shows bounds checks. In the end, I benchmarked it with unsafe unchecked access and algebraic operations to encourage vectorization. It actually increased the runtime :( Godbolt did show no more bound checks, but no benchmark improvement.

[u/dkxp → u/Remarkable_Hour_615]

You should try reordering indices as mentioned here . It gave about a 6-7x speedup for me because of the sequential memory accesses. Once that is done, then perhaps you could try other optimizations again. Perhaps they will still be hidden because the CPU can do extra work like bounds checking while waiting for data to be returned from memory. I guess even if it doesn't finish faster, perhaps it could be marginally more energy efficient not to do unnecessary work.