Deep Dive: Optimizing SIMD loops in Rust with std::arch

When you care about throughput and latency at the same time, the compiler's auto-vectorization usually gets you 80% of the way there. LLVM is very good at noticing that a tight loop over &[u8] or &[f32] is really a packed compare, a horizontal add, or a masked store. Most production code should stop there.

The remaining 20% lives in hot paths: kernels that run on every packet, every audio window, or every compiler pass. Those loops have just enough control flow, aliasing, or alignment weirdness that the optimizer backs off. At that point you stop hoping for a miracle and write the vector width yourself.

This article walks through that jump. We start from ordinary Rust, look at why Result and .unwrap() still show up next to unsafe, and then build a small AVX2 kernel that counts newlines. The goal is not to memorize intrinsics. It is to keep the unsafe island small, testable, and honest about what the CPU is doing.

Why .unwrap()?

std::fs::read_to_string returns an std::io::Result<String>, which is an alias for std::result::Result<String, std::io::Error>. That type is a sum type: Ok(T) when the read succeeds, Err(E) when it does not. Nothing about SIMD changes that contract. The file still might not exist. The buffer still might be empty. The kernel still has to decide what a failure means.

In toy examples you will often see something like this, and it is fine as a teaching shortcut:

let content = std::fs::read_to_string("input.txt").unwrap();

Here we are saying: if reading the file fails, panic. That is a developer error, not a runtime policy. For a blog-post kernel that loads a fixture, that is acceptable. For a production SIMD path that lives inside a server, you propagate the error instead of crashing the process.

fn load_data(path: &Path) -> color_eyre::Result<String> {
    let content = std::fs::read_to_string(path)?;
    Ok(content)
}

If reading the file fails, just panic and crash — this is a developer error.

That quote is the intent of .unwrap(). Once the bytes are in memory, the SIMD loop itself should be free of I/O. Load, validate, then vectorize. Mixing those layers is how you end up with an unsafe block that also opens files.

When the compiler already does the work

Auto-vectorization works well when the loop has no complex control flow, there are no hidden dependencies between iterations, and the compiler has good aliasing information. The moment you mix in branchy logic or raw pointer arithmetic, LLVM gets conservative. That is when std::arch is the right escape hatch.

A good rule of thumb: write the scalar version first, measure it, and only then open std::arch::x86_64. If the auto-vectorized loop is already within 20% of the hand-written one, keep the safe Rust.

A tiny SIMD example

Assume we want to count how many bytes in a slice are equal to b'\n'. The naive version is one iterator chain and it is already a useful baseline:

pub fn count_newlines_naive(buf: &[u8]) -> usize {
    buf.iter().filter(|&&b| b == b'\n').count()
}

An AVX2 version (simplified, error handling omitted) processes 32 bytes at a time, then finishes the tail in scalar code. The remainder loop is deliberately boring. That is the point: keep the weird part in one function.

use std::arch::x86_64::*;

pub unsafe fn count_newlines_avx2(buf: &[u8]) -> usize {
    let mut cnt = 0usize;
    let mut i = 0;
    let needle = _mm256_set1_epi8(b'\n' as i8);

    while i + 32 <= buf.len() {
        let chunk = _mm256_loadu_si256(buf[i..].as_ptr() as *const __m256i);
        let eq = _mm256_cmpeq_epi8(chunk, needle);
        let mask = _mm256_movemask_epi8(eq) as u32;
        cnt += mask.count_ones() as usize;
        i += 32;
    }

    cnt += buf[i..].iter().filter(|&&b| b == b'\n').count();
    cnt
}

The shape we want is always the same:

• Load a vector, compare, reduce, advance by the lane width.
• Leave the scalar remainder obvious and easy to audit.
• Isolate unsafe in a small, testable function instead of sprinkling it through the caller.

If you later need SSE4.1 or NEON, the public wrapper stays safe and the target-specific kernels stay behind #[target_feature].

Performance metrics

On a 4-core laptop (release build, best-of-5 runs) the numbers look like this. They are rounded on purpose: they show the shape of the improvement, not a number you should quote in a paper.

The interesting row is usually the middle one. Auto-vectorization already bought 2.6x. Manual AVX2 is worth it only if this kernel sits on the critical path and you have tests that lock the result against the scalar version.

Note. These numbers assume aligned-enough loads (_mm256_loadu_si256) and a buffer large enough that the 32-byte loop dominates. Tiny payloads will not look like this.

Reader suggestion: std::cmp::Reverse

Sometimes we want to sort by “largest first”, so we reach for tricks like subtracting from u64::MAX. That works, but it hides the intent:

.sorted_by_key(|&v| u64::MAX - v)

A clearer version uses std::cmp::Reverse. Same order, no clever overflow story:

.sorted_by_key(|&v| std::cmp::Reverse(v))

Reader suggestion: Reverse. I prefer Reverse because it encodes intent instead of clever math with u64::MAX. Thanks to freax13 on Reddit for the idea.

Combining it with k_smallest

We can go one step further and combine k_smallest from itertools, Reverse<T>, and a batching pattern with fold_while. The result is denser, but each combinator still tells a story:

use itertools::{FoldWhile, Itertools};
use std::cmp::Reverse;

fn main() -> color_eyre::Result<()> {
    color_eyre::install()?;

    let answer = include_str!("input.txt")
        .lines()
        .map(|v| v.parse::<u64>().ok())
        .batching(|it| {
            it.fold_while(None, |acc: Option<u64>, v| match v {
                Some(v) => FoldWhile::Continue(Some(acc.unwrap_or_default() + v)),
                None => FoldWhile::Done(acc),
            })
            .into_inner()
        })
        .map(Reverse)
        .k_smallest(3)
        .map(|x| x.0)
        .sum::<u64>();

    println!("{answer:?}");
    Ok(())
}

The important part is not this exact puzzle. It is that the combinators line up: parse, batch, reverse, take k, strip the wrapper, sum. Once that story is readable, you can replace any step with a SIMD kernel without rewriting the rest.

Verifying the output

Finally we sanity-check the binary against the provided test vectors. If this does not match the scalar implementation, the AVX2 loop is wrong — no amount of throughput saves you.

cargo run --release --bin day01
Compiling aoc2024 v0.1.0 (/home/jorge/dev/aoc2024)
Finished release [optimized] target(s) in 0.45s
Running `target/release/day01`
24000

Once this passes on both sample and real input, start micro-benchmarking the kernel itself with criterion or cargo bench. Measure the scalar version, the auto-vectorized version, and the std::arch version on the same buffers. Keep the winner. Throw the rest away.