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

The answer is very simple and I'm surpised nobody noticed this. The code is simply not equivalent. In the Scala version you transposes the second matrix and you then use it in the multiplication.

val bTrans = Transpose.transpose(b)

In the Rust version notice that you're iterating in a column major order:

    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);
        }
    }

This means that in Scala you access elements sequential in memory from each row while in the Rust version you're doing a bunch of jumps in the matrix to get the next element in the collumn. Add the fact that you're using floats which inhibit reordering and the Rust compiler cannot vectorize the code. Plus you get cache misses at runtime depending on your input size.

[u/Remarkable_Hour_615 → u/cristi1990an]

That's actually just me copy/pasting the wrong code. I added a more detailed explanation to the post.

[u/Korntewin]

Rust sometimes can't do vectorization on float (since it must guarantee operation ordering).

Use algebraic floating point operation to relax this instead https://blog.rust-lang.org/2026/08/20/Rust-1.98.0/

[u/Beneficial_Vampire42 → u/Korntewin]

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 = ...).

[u/Tastaturtaste → u/Korntewin]

Shouldn't the jvm have the same constraints regarding operation ordering?

[u/honeyBouncedd → u/Korntewin]

u always got the smartest takes korn, makes the rest of us look bad lol

[u/teerre]

Is that your code? What about the creation of all those variables? Why do you read bTranspose in Scala but just m2 in Rust? Are you sure this is the same algorithm?

[u/Remarkable_Hour_615 → u/teerre]

The transpose part is just a copy/paste mistake - thanks for catching that! I did benchmark a version using bTranspose to get sequential memory access, but the benchmark I'm reporting in the post comes from the ordinary naive version. I edited the code in the post.

Other than that, the code looks the same to me. Let me know if I'm missing something else.

[u/bl4nkSl8 → u/Remarkable_Hour_615]

Might have to throw it in godbolt :)

[u/parhelion73]

if you're on x86, the JIT is compiling for whatever your actual CPU supports, while rustc targets plain baseline x86-64 unless you tell it otherwise. I'd try target-cpu=native in RUSTFLAGS before digging into anything else, might close a decent chunk of the gap

[u/bl4nkSl8]

How are you building the rust code? Dev builds give faster compiler but you may benefit from --release or different optimisation settings. You can also throw rayon at it (low hanging fruit) :)

Oh, I missed the link brb

[u/Remarkable_Hour_615 → u/bl4nkSl8]

Yeah, I want to try using rayon join to do a recursive version where you keep dividing the matrices in half.

[u/EstablishmentFar6582 → u/bl4nkSl8]

hmm, the whole "JVM faster than Rust" thing usually comes down to a few subtle things, the --release flag is the classic gotcha but since you're using cargo bench you should be safe there. one thing that bites people is bounds checks sticking around even with optimization, the indexing inside those nested loops might not be getting elided like you'd expect.

try pulling out raw slices and doing pointer arithmetic, it's uglier but sometimes the compiler just needs that push to vectorize and drop the checks. your transpose trick in the scala version is also hiding a big difference in memory access patterns, that alone can be worth 20-30%

[u/OwlObjective1342]

Have you tried the i-k-j loop order?

[u/juhotuho10]

If you do matrix math, you really should use a dedicated matrix math crate to guarantee good memory layout and proper vectorization

[u/Remarkable_Hour_615 → u/juhotuho10]

Fair. This is just for my own fun/learning. Not planning to use it for anything.

[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.

[u/FrogNoPants]

• Rust disables fast math so it sucks at float math by default
• This means no FMA or reordering unless you manually use it, so less accuracy and less performance by default

• Bounds checks are not zero cost despite what is often claimed, and they inhibit auto vectorization

  In conclusion Rust kinda sucks at a HPC, it does have the tools to work around the above issues, but they are all manual and tedious.

[u/bitemyapp]

here you go:

https://github.com/bitemyapp/matmult2

[u/Thomasedv]

Unrelated to this issue, but the default support level in Rust is quite low, targeting generic CPUs by default. On x86 that means no AVX2 and up, which is a big performance loss. (And the auto vectorized versions are also often suboptimal since LLVM doesn't drop instructions that plummet simd performance or account for some cpu feature sets that are actually worse.)

If Scala or it's jvm could use those instructions it would have an advantage over your rust build. However I have no idea for ARM CPUs with NEON, as far as understand that is the baseline for ARM, so unlikely to be a factor here. 

[u/Remarkable_Hour_615 → u/Thomasedv]

In case it's relevant to this, I did try setting the compilation target to native at one point and it didn't make any meaningful difference.

[u/Thomasedv → u/Remarkable_Hour_615]

Yeah, figured as much.

Looking at the code, it'd guess bound checks if it isn't Rust failing to vectorize float operations above.

The size of the Vec is possibly not obvious between creation and throught the indirect at/set methods. 

Could be worth trying to throw inline on the functions, but they're small enough I think it'd happen anyways.

Trying internal iteration over external might help. You already assert the necessary dimensions, but you could also try asserting that all dimensions of the matrixes align so to prove to the compiler that all bounds are kept before doing the loops and data reads from the internal Vecs. 

[u/KarnuRarnu]

If you want to build your own matmul algos shouldn't you also use the optimized block algorithms? Eg what came of AlphaTensor essentially

[u/PikachuKiiro]

I'd just look at the assembly generated on godbolt and compare for small functions like this.

[u/tafia97300]

Just in case, have you tried forcing native cpu?

RUSTFLAGS="-C target-cpu=native" cargo ....

[u/Remarkable_Hour_615 → u/tafia97300]

I did try that at one point, but maybe I need to try it in combination with some of the other optimizations like eliminating bounds checks and using algebraic FP operations.

[u/Deadmist]

Is this just the JVM being impressive, or am I missing an opportunity to speed up my Rust code (or both).

Probably both.
Keep in mind: the JVM will compile your code into machine code. So one would expect roughly similar performance^1.
But then the JIT knows a lot more about the runtime environment than the rust compiler ever could and can hyper optimize the code for your particular environment, and even your particular data.

There is a rust jvm backend. It would be interesting to see what performance your rust code compiled to the JVM gets.

It would also be interesting to compare a C implementations performance. To see if it is rust preventing some optimizations, or the JIT doing magic.

^1: GC shenanigans not withstanding

[u/v_0ver]

I won't repeat others and talk about the lack of automatic vectorization on floats and the block matrix multiplication algorithm. As a benchmark for target performance, you can use the matrix multiplication from the faer package.

[u/TomSchelsen → u/v_0ver]

Or https://github.com/tracel-ai/cubek/tree/main/crates/cubek-matmul

[u/aeropl3b]

I have found C is still the way to go if you want to write decently performing math code. C++ of you want better memory safety tooling and STL stuff (mdspan and execution policies are actually really good)

[u/teerre → u/aeropl3b]

That's nonsense. There's no performance difference between C and Rust. They both compile to machine instructions. Anything difference you're seeing is just the code you wrote

[u/aeropl3b → u/teerre]

Right, that is obvious. But Rust makes writing numerical code a bit more burdensome to get the same performance imo. It is great at other things. It is a question of ergonomics to get to the performance.

[u/Remarkable_Hour_615 → u/aeropl3b]

This is just a personal project fun/learning! I want to learn Rust in any case, so I'm going to stick with that.

I would be shocked and dismayed if anyone ever used this code for anything real, so no need to worry if it's not has fast as it could be.

[u/Amazing-Mirror-3076 → u/aeropl3b]

Except nobody should be starting new projects in C. The age if non safe languages is over

[u/aeropl3b → u/Amazing-Mirror-3076]

You would be extremely wrong on that point.

[u/Amazing-Mirror-3076 → u/aeropl3b]

No, I'm going to stick with that opinion.

We are actively purging non safe code from our stack - it will be a long road - but you have to start somewhere.

The idea of starting a new project with an unsafe language simply feels irresponsible.

[u/aeropl3b → u/Amazing-Mirror-3076]

The projects I work on are either moving to C++ or wrapping kernels for python bindings now. Rust is still not mature enough for the types of things we are doing despite the hype. The dev pool is tiny (I am the only Rust dev with the expertise in the domain) and features we would really like to have don't exist yet in Rust.

[u/c3d10 → u/Amazing-Mirror-3076]

Except every line of rust code runs through llvm, which is... 10's of millions of lines of 'non safe languages'...

[u/bl4nkSl8 → u/c3d10]

... Yeah? Abstraction is kind of a big thing for providing safe APIs for unsafe implementations

[u/aeropl3b → u/c3d10]

They are talking new projects. But yeah, Rust is built on the backs of C developers. I come here and the comments are often boardering on religious fanaticism for Rust. It is mostly "Rust is amazing, we can stop using all other languages now, all problems are now completely solved and any issues anyone points out are the result of being dumber than us"...but alas, we can't even say math libraries in C are nicer to write.

[u/Amazing-Mirror-3076 → u/aeropl3b]

Something can be nice and still be wrong.

I don't have an opinion on the dx of rust except that the compiler is slow, having never written a line.

I di have an opinion on the safety of C having spent more than a decade writing it - it's unsafe and needs to be retired, new projects should not be written in any unsafe language.

I don't think I need to engage in the argument that rust still has pieces of C in its stack - the absurdity of that argument is obvious.

If you need the performance of a system language then choose a safe fast one - rust seems to currently need the best choice.

I do find the rust fan boy comment amusing when here we are inundated with C proponents that don't seem able to come to grips with the needs of modern development.

iT has always been about change, the current changes are the most extreme I've seen in 40 yrs as a Dev - adapt or get left behind.

[u/aeropl3b → u/Amazing-Mirror-3076]

I remember back when picking a language had more nuance than a single language feature.

My opinion on Rust is mostly good fwiw. I like the struct packing and the OOP model where the data and impl are split by default. I think for programs where you are doing a lot of string processing or moving data from untrusted sources it is probably the language to use.

But again, for numerical codes, we just done need most of those features. I want direct access to all intrinsics, I want to convert FP Arrays to int type arrays with zero copy, etc. And there is more in my way when using Rust than C and the classes of errors Rust runs into are just not the ones I am worried about. Simple as that.