Vectorization is a critical optimization step without which numerical computing code will not perform well. Achieving vectorization in Rust, however, can be a little difficult. Rust emphasizes memory safety and as a result many common indexing operations do runtime bounds-checking. This bounds checking can inhibit vectorization when the compiler cannot prove that the checks are redundant, although optimized Rust can and often does eliminate bounds checks. One can make the proof easier by relying on iterators instead. I show here an interesting implementation of matrix-matrix multiplication in Rust using iterators alone and no raw indexing. For this example, the iterator structure lets Rust omit the bounds checking in the inner loop. I show the reference code with raw indexing, its translation to iterators, and I also include an “unsafe” version using raw pointer dereferencing for performance comparison.
In what follows I will compute \[ \mathbf{C} \leftarrow \mathbf{C} + \mathbf{A}\mathbf{B} \] and these symbols will roughly correspond to variables in the Rust code which follows. I map matrices to simple 1D arrays of contiguous floating point numbers as is common in libraries like BLAS and LAPACK. I store the matrices in row-major order so that rows are stored contiguous in memory.
The basic Rust implementation is
pub fn gemm(m : usize, n : usize, p : usize, a : &[f64], lda : usize, b : &[f64], ldb : usize, c : &mut [f64], ldc : usize) -> () {
//m : number of rows of A and C
//n : number of columns of B and C
//p : number of columns of A,rows of B
for i in 0..m{
for k in 0..p{
for j in 0..n{
let aik=a[k+lda*i];
let bkj=b[j+ldb*k];
let ref mut cij=c[j+ldc*i];
*cij += aik*bkj;
}
}
}
}
We can see from Godbolt.org that this code does not vectorize even with optimizations turned on (I set the architecture to Skylake+AVX512 so we are looking for register usage from ymm* and packed vector operations like vmulpd,vaddpd.)
For this loop, the generated code indicates that bounds checking prevented vectorization. Rust can eliminate these checks in many indexed loops, but switching to iterators makes it easy here.
Rust data slices (the types &[f64] in this case) are somewhat similar to pointers in C and C++ but one key difference is that they come with a length. This means many iterator idioms common for more advanced containers in Rust (e.g. BTreeSet, Vec) also work with slices. However, it is not immediately obvious how to translate the above indexed code to iterator code. If instead of iterating directly over floating point numbers in the slices we iterate over rows using chunks and use zip for co-iteration then the indexed code translates to the following:
pub fn gemm(_m : usize, _n : usize, _p : usize, a : &[f64], lda : usize, b : &[f64], ldb : usize, c : &mut [f64], ldc : usize) -> () {
for (ci,ai) in c.chunks_exact_mut(ldc).zip(a.chunks_exact(lda)){
for (aik,bk) in ai.iter().zip(b.chunks_exact(ldb)){
for (cij,bkj) in ci.iter_mut().zip(bk.iter()){
*cij += (*aik) * (*bkj);
}
}
}
}
This compact version assumes contiguous row-major matrices with no padding or submatrices. The iterator structure is the point of the example, not a complete BLAS interface.
Since this is done with iterators Rust can omit the bounds checking. Godbolt shows decent usage of ymm* registers and vmulpd,vaddpd where one should expect them to be:
For comparison I will show an unsafe version of this code and one can compare how they vectorize and whether there is loss of codegen quality because of the use of iterators.
This implementation matches more closely how it would work in C or C++. The unsafe block in Rust in this case lets us dereference raw pointers without Rust inserting a bounds check. Thus this should vectorize about as well as the equivalent would in C or C++ and so will serve as a useful reference to the iterator implementation. The caller must still guarantee that every computed offset is in bounds and that the mutable output does not alias the inputs; otherwise the unsafe block can invoke undefined behavior.
pub fn gemm(m : usize, n : usize, p : usize, a : &[f64], lda : usize, b : &[f64], ldb : usize, c : &mut [f64], ldc : usize) -> () {
//Grab pointers to the slices.
//This is fine to do outside unsafe blocks, but you can't
//realy do anything useful with them except inside unsafe
//blocks. Putting it here makes the inner loop look cleaner.
let ap=a.as_ptr();
let bp=b.as_ptr();
let cp=c.as_mut_ptr();
//Pointer offsets have to be done with
//signed integers. Convert everything to
//signed for later use.
let ldai = lda as isize;
let ldbi = ldb as isize;
let ldci = ldc as isize;
for i in 0..(m as isize){
for k in 0..(p as isize){
for j in 0..(n as isize){
let ik = k+ldai*i;
let kj = j+ldbi*k;
let ij = j+ldci*i;
unsafe{
let aik = *(ap.offset(ik));
let bkj = *(bp.offset(kj));
let cij = cp.offset(ij);
*cij += aik*bkj;
}
}
}
}
}
The binary size is a little smaller for this version than the iterator version, but the actual inner loop vectorization is roughly identical.
Now I show how these implementations compare
Running on my laptop I benchmarked these matrix-matrix multiplications for various matrix sizes by running each implementation 100 times for each size and taking the fastest runtime. Thus these are best-case timings rather than typical timings. I plot below all three implementations together
Both the iterator version and the unsafe version handily defeat basic Rust indexing version because they were able to similarly vectorize. On this scale it is hard however to compare the iterator version to the unsafe version so I removed the “naive” numbers to make it easier to see these together
It appears that on these matrix sizes, with this compiler and these layouts, the iterator implementation is effectively equivalent in performance to the unsafe version.
Vectorization is only one piece of the puzzle of performance for matrix-matrix multiply. Other important optimizations are:
Implementing all of these is a tremendous amount of work so it’s usually best to just use an optimized BLAS wherever possible. I was just curious here whether one could implement the matrix-matrix multiply using pure iterators in Rust and achieve vectorization - it does appear that this is possible.