Rust SIMD — a tutorial

SIMD in Rust

Single instruction, multiple data (SIMD) is a type of parallel computing that computes on multiple values in a single CPU instruction. Special CPU registers allow for this possibility. Examples on an x86_64 platform are xmm0 — xmm15, 128-bit wide, and ymm0 — ymm15, which are 256-bit wide. 256-bit wide means that four 64-bit or eight 32-bit (etc…) values can be computed on simultaneously. For those interested, the ymm-based registers were introduced as part of the AVX (Advanced Vector Extensions) instruction set. However, they are not the widest registers — the AVX-512 instruction set introduced the zmm0 — zmm15 registers, 512 bits wide, as the name suggests, although they are currently not very widely used.

Due to the popularity of the AVX instruction set, this article will primarily focus on the AVX instruction set, that is, the 256-bit wide registers, in Rust.

Structure of the SIMD register

When using SIMD instructions, you will typically structure your data in terms of lanes. A lane represents the bits of a particular value that is being computed. Drawn below is an illustration of the lanes for a 256-bit wide register containing 32-bit sized values. There are 8 lanes, each representing a single 32-bit value.

Manipulating a SIMD register is done on a lane-by-lane basis. This is why SIMD calculations are often referred to as vectorization, since they are computationally equivalent to calculations in linear algebra.

SIMD crates

It’s worth briefly referencing some crates that attempt to abstract SIMD from the underlying architecture. The most popular of these is packed_simd. This package works around a Simd<[T;N]> type in order to perform the vector math. There is also the crate packed_simd_2, which points to the same repo. Furthermore, with the nightly Rust compiler, one can use portable-simd, which supersedes packedsimd. The idea behind this project is to include cross-platform SIMD functionality from within the Rust standard library. Below is some example code on how to use it (taken from the GitHub repo):

#![feature(portable_simd)]
use std::simd::f32x4;
fn main() {
    let a = f32x4::splat(10.0);
    let b = f32x4::from_array([1.0, 2.0, 3.0, 4.0]);
    println!("{:?}", a + b);
}

If you want to use a cross-platform SIMD library that does not require the nightly version of the compiler, the crate wide might fit your needs.

SIMD in practice

The purpose of the code below is not to observe the performance gains of SIMD but instead to explore how to write vectorized code. Below is the code snippet that will be vectorized using architecture-specific primitives provided by the Rust standard library.

pub fn occurrence_non_simd(needle: u8, bytes: &[u8]) -> usize {
    let mut count = 0;
    for byte in bytes {
        if needle == *byte {
            count += 1;
        }
    }
    count
}

The above code simply returns the occurrences of a needle in a byte slice. As previously mentioned, we are not attempting to beat the performance of this code (which will be challenging since the Rust compiler can aggressively optimize simple linear loops in ways that go beyond SIMD).

The code below is the SIMD example that calculates the same output. Note that I am using x86_64 primitives.

#[cfg(target_arch = "x86_64")]
pub fn occurrence_u8_simd(needle: u8, bytes: &[u8]) -> usize {
    use std::arch::x86_64::{
        __m256i, _mm256_cmpeq_epi8, _mm256_loadu_si256, _mm256_movemask_epi8, _mm256_set1_epi8,
    };

    unsafe {
        let needles = _mm256_set1_epi8(needle as i8);
        let mut count = 0;
        let num_iter = bytes.len() / 32;
        (0..num_iter).for_each(|i| {
            let simd_ = _mm256_loadu_si256(bytes.as_ptr().add(i * 32) as *const __m256i);
            let r = _mm256_cmpeq_epi8(needles, simd_);
            let r_moved = _mm256_movemask_epi8(r).count_ones() as usize;
            count += r_moved;
        });

        count
    }
}

For the observant, there is a slight subtle difference in the output of the SIMD code. Namely, that byte slices of lengths that are not divisible by 32 could yield inaccurate results. We will ignore that point since the byte slice length provided in the testing was a factor of 32. For those interested in how this restriction can be avoided, the code below can be added to the end of the unsafe block (although this will not work in this case since it is using portable-simd representation, but the idea remains).

let remainder = bytes.len() % 32;
if remainder > 0 {
   let mut tail = [0u8; 32];
   tail[..remainder].copy_from_slice(&bytes[offset..]);
   // The code below is from using 
   let result = u8x32::from(tail).cmp_eq(needles);
   count -= result;
 }
 count.to_array().iter().sum::<u64>()

Back to the SIMD equivalent code. The first step is to initialize a 256-bit register filled with our needle value, which will act as the comparitor, the width of each lane is 8 bits.

let needles = _mm256_set1_epi8(needle as i8);

Since we have 256 bits available and are using an 8 bit number, we can compare 32 bytes with our needle in a single instruction call. First we load 32 bytes at a time into a 256-bit register

let simd_ = _mm256_loadu_si256(bytes.as_ptr().add(i * 32) as *const __m256i);

Then we perform a comparison with the needle register on a lane by lane basis.

let r = _mm256_cmpeq_epi8(needles, simd);

We then extract the most significant bit of each 8-bit element in r in order to find out which lanes contained the same value as our needle

let r_moved = _mm256_movemask_epi8(r).count_ones() as usize;

countones() returns the number of ones in the binary representation of the result of the mask (which is an i32). It is not an O(n) operation since countones() is translated into a single hardware instruction called POPCNT (population count).

Get BWinter’s stories in your inbox

Join Medium for free to get updates from this writer.

Remember me for faster sign in

Then we simply increment the count.

Performance

Although I mentioned that the purpose of the SIMD code is not to beat the simple linear one it is still interesting to compare the performance. Using criterion the linear code is 23% faster.

It would be interesting to see if we can get it a bit closer. One simple technique is loop unrolling, essentially reducing the number of iterations needed to loop.

#[cfg(target_arch = "x86_64")]
pub fn occurrence_u8_simd_optimized(needle: u8, bytes: &[u8]) -> usize {
    use std::arch::x86_64::{
        __m256i, _mm256_cmpeq_epi8, _mm256_loadu_si256, _mm256_movemask_epi8, _mm256_set1_epi8,
    };

    unsafe {
        let needles = _mm256_set1_epi8(needle as i8);
        let mut count = 0;
        let num_iter = bytes.len() / 32;
        (0..num_iter).step_by(4).for_each(|i| {
            let simd_ = _mm256_loadu_si256(bytes.as_ptr().add(i * 32) as *const __m256i);
            let r = _mm256_cmpeq_epi8(needles, simd_);
            let r_moved = _mm256_movemask_epi8(r).count_ones() as usize;
            count += r_moved;

            let simd_ = _mm256_loadu_si256(bytes.as_ptr().add(i * 32 + 32) as *const __m256i);
            let r = _mm256_cmpeq_epi8(needles, simd_);
            let r_moved = _mm256_movemask_epi8(r).count_ones() as usize;
            count += r_moved;

            let simd_ = _mm256_loadu_si256(bytes.as_ptr().add(i * 32 + 32 * 2) as *const __m256i);
            let r = _mm256_cmpeq_epi8(needles, simd_);
            let r_moved = _mm256_movemask_epi8(r).count_ones() as usize;
            count += r_moved;

            let simd_ = _mm256_loadu_si256(bytes.as_ptr().add(i * 32 + 32 * 3) as *const __m256i);
            let r = _mm256_cmpeq_epi8(needles, simd_);
            let r_moved = _mm256_movemask_epi8(r).count_ones() as usize;
            count += r_moved;
        });
        count
    }
}

With this optimization, the SIMD code is 11% slower than the linear code. We are getting closer.

Another way we can try to improve is to make the assumption that the bytes are 32-bit aligned. This should be faster because the SIMD register would load data in a single instruction, avoiding extra CPU cycles handling misaligned memory.

#[cfg(target_arch = "x86_64")]
pub fn occurrence_u8_simd_optimized(needle: u8, bytes: &[u8]) -> usize {
    use std::arch::x86_64::{
        __m256i, _mm256_cmpeq_epi8, _mm256_load_si256, _mm256_movemask_epi8, _mm256_set1_epi8,
    };

    unsafe {
        let needles = _mm256_set1_epi8(needle as i8);
        let mut count = 0;
        let num_iter = bytes.len() / 32;
        (0..num_iter).step_by(4).for_each(|i| {
            let simd_ = _mm256_load_si256(bytes.as_ptr().add(i * 32) as *const __m256i);
            let r = _mm256_cmpeq_epi8(needles, simd_);
            let r_moved = _mm256_movemask_epi8(r).count_ones() as usize;
            count += r_moved;

            let simd_ = _mm256_load_si256(bytes.as_ptr().add(i * 32 + 32) as *const __m256i);
            let r = _mm256_cmpeq_epi8(needles, simd_);
            let r_moved = _mm256_movemask_epi8(r).count_ones() as usize;
            count += r_moved;

            let simd_ = _mm256_load_si256(bytes.as_ptr().add(i * 32 + 32 * 2) as *const __m256i);
            let r = _mm256_cmpeq_epi8(needles, simd_);
            let r_moved = _mm256_movemask_epi8(r).count_ones() as usize;
            count += r_moved;

            let simd_ = _mm256_load_si256(bytes.as_ptr().add(i * 32 + 32 * 3) as *const __m256i);
            let r = _mm256_cmpeq_epi8(needles, simd_);
            let r_moved = _mm256_movemask_epi8(r).count_ones() as usize;
            count += r_moved;
        });
        count
    }
}

The only difference here is the load, which changed from _mm256_loadu_si256 to _mm256_loadsi256. The benchmark shows that this code is 10.8% faster, but within the margin of error of the previous benchmark, so no real difference.

When to use SIMD?

The above code, although less performant than its linear counterpart, was written to explore how to perform the Rust equivalent of “hand-rolled” SIMD. So if hand-rolled is generally slower how can we help the compiler use SIMD. The simplest way is to ensure contiguous memory since loads are predictable and cache misses are rare.

Array of Structures vs Structure of Arrays

Say you have the following structure:

struct MyStruct {
    value_a: u8,
    value_b: u8
}

let myarr = [MyStruct { value_a: 1, value_b: 2 }; 1000]

If you wanted to perform calculations independently on value_a, for example, e.g., find the sum, this structure would not be properly optimized for SIMD, simply because value_a is not contiguous—each value_a is separated by a value_b.

A more suitable SIMD implementation would be Structure of Arrays:

struct MyStruct {
    values_a: [1; 1000],
    values_b: [1; 1000]
}

The contiguous nature of values_a ensures that loading into SIMD registers is fast and can be further optimized if aligned.

Ensuring alignment

Below is the code (without comments) used to generate a 32-bit aligned vector.

#[repr(align(32))]
pub struct AlignedTo32Bytes([u8; 32]);
pub fn aligned_vec(n_bytes: usize) -> Vec<u8> {
    assert!(n_bytes % 32 == 0, "n_bytes must be a multiple of 32");
    let num_elements = n_bytes;
    let aligned: Vec<AlignedTo32Bytes> = vec![AlignedTo32Bytes([0u8; 32]); num_elements];
    let ptr = aligned.as_ptr();
    assert!((ptr as usize) % 32 == 0);
    // Avoid double free
    std::mem::forget(aligned);
    unsafe { Vec::from_raw_parts(ptr as *mut u8, n_bytes, n_bytes) }
}