SIMD for Collision

32 points by mikejsavage


mitchellh

SIMD is surprisingly easy to reason about for tight tasks like this and more software people should be utilizing it. I think people get hung up on the crazy (and genuinely impressive) SIMD projects like simdutf and simdjson, which are head-spinning in complexity.

But basic "process N bytes at a time" SIMD that could speed up many ordinary loops are fairly simple! They all follow the same pattern:

  1. Broadcast any constants you need and initialize vector accumulators, if any.
  2. Loop over input one vector-width chunk at a time
  3. Perform the comparison or arithmetic across all lanes in parallel.
  4. Reduce or store the vector result as needed (into number 1)
  5. Handle the remaining elements with a scalar tail. A scalar tail is just your normal loop you had before vectorizing, but its only processing the remainder that doesn't fit into the vector-width.

Here is a simple example. I have a byte string that I want to keep consuming until I see a value less than 0xF (a control character limit). The naive loop:

// Loop while we see codepoints above 0xF
while (end < cps.len and cps[end] > 0xF) {
    end += 1;
}

And here is the generic vector version (no CPU-specific intrinsics):

// If the target CPU doesn't have a vector size for u32 just fallback to
// scalar. This has some CPU detection and hardcoding so we pick
// the proper vector width for the hardware.
if (simd.lanes(u32)) |lanes| {
    // A vector containing `lanes` independent `u32` values. Basically a chunk of u32!
    const V = @Vector(lanes, u32);

    // Broadcast 0xF into every lane so one vector comparison checks every value.
    const threshold: V = @splat(0xF);

    // Only enter the SIMD loop when a complete vector remains in the input.
    while (end + lanes <= cps.len) {
        // Load the next `lanes` codepoints into a single SIMD vector.
        const values: V = cps[end..][0..lanes].*;

        // Compare every codepoint against 0xF in parallel, producing one boolean per lane.
        const greater_than_threshold = values > threshold;

        // When every lane is greater than 0xF, the entire vector can be skipped.
        if (@reduce(.And, greater_than_threshold)) {
            end += lanes;
            continue;
        }

        // Convert the per-lane booleans into a bit mask: 1 means the value was > 0xF.
        const mask: std.meta.Int(.unsigned, lanes) =
            @bitCast(greater_than_threshold);

        // Invert the mask so failing lanes become 1, then find the first failing lane.
        end += @ctz(~mask);

        // We found the first value <= 0xF, so the scan is complete.
        break :scan;
    }
}

// Finish the partial vector-sized tail, or do the entire scan when SIMD is unavailable.
// It's just our normal scalar loop!
while (end < cps.len and cps[end] > 0xF) {
    end += 1;
}

This gets you straight throughput speedup of anywhere from 2x to 16x depending on the hardware. This is real, because its from my actual project.

I realize that it looks like a LOT more code, but its code following a really specific pattern that you get used to quickly (the steps outlined above). And, the comments and whitespace add a lot of noise. :)

Why can’t the compiler just do this through auto-vectorization? Sometimes it can! Compilers are quite good at vectorizing simple, regular arithmetic loops. But you still can't reliably expect them to discover more complicated transformations like this one, particularly when the loop contains an early exit and the desired implementation involves a comparison mask, a reduction, and locating the first failing lane. More details here: https://llvm.org/docs/Vectorizers.html (always know your compilers!)

You could say “compilers aren't good enough for now,” but auto-vectorization has been an active compiler research problem for decades, and recent research still begins from the observation that production compilers frequently miss vectorization opportunities. (Ref from 2024: https://arxiv.org/abs/2406.04693)

Once you understand the basic SIMD pattern, code like this becomes nearly as natural as writing the scalar loop... more engineers should know how to do this and more languages should provide the tools to do it.

EDIT: I expanded this comment into a full more detailed blog post: https://mitchellh.com/writing/everyone-should-know-simd