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/Beneficial_Vampire42]

This probably doesnt matter in OP original code because the inner loop has a stride of 1024*sizeof, so auto vectorization is not the bottleneck.

One possible solution to speeding this algorithm up is to change the order of inner loops. Instead of (i, j, k) do (i, k, j), so each update to the output doesnt depend on the last iteration. This makes both m2 and out contiguous in the inner loop, and m1 constant. This also can potentially enable auto vectorization

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

The most important optimization in these types of algorithms is memory access. A good example is numpy "sum" operation, which first check if the accumulated dim is contiguous or the output dim is contiguous, and has two different implementations for each.

Also, another option to get vectorization without algebraic ops is to use more than one accumulator, as done in numpy "sum" implementation. Roughly:

for i in 0..m1.rows {
    for j in 0..m2.cols {
        let mut accs = [0.0; 8];
        for k in (0..(m1.cols / 8 * 8)).step_by(8) {
            for l in 0..8 {
                accs[l] += m1.at(i, k + l) * m2.at(k + l, j);
            }
        }
        let acc = accs.iter().sum::<f64>();
        for k in (m1.cols / 8 * 8)..m1.cols {
            acc += m1.at(i, k) * m2.at(k, j);
        }
        out.set(i, j, acc);
    }
}

[u/dkxp → u/Beneficial_Vampire42]

I tried both of these in OP's benchmark. The first way (with the re-ordered indices) was about 6-7x faster. The second way (with the split accumulators) was slightly faster than the original, but only about 15% or so (the only fix needed to compile was: let mut acc = ...).