Everyone Should Know SIMD (Mitchell Hashimoto)

Everyone Should Know SIMD

32 Likes

I always think it’s a shame to hard code the vector body-scalar tail like he proposes here.

if (simd.lanes(u32)) |lanes| {
    const V = @Vector(lanes, u32);
    const threshold: V = @splat(0xF);
    while (end + lanes <= cps.len) : (end += lanes) {
        const values: V = cps[end..][0..lanes].*;
        const greater_than_threshold = values > threshold;
        if (@reduce(.And, greater_than_threshold)) continue;
        const mask: std.meta.Int(.unsigned, lanes) = @bitCast(greater_than_threshold);
        end += @ctz(~mask);
        break;
    }
}

while (end < cps.len and cps[end] > 0xF) end += 1;

My day job is RISC-V, and if we look at the code generated from that we can see the two loops. Body loop .LBB1_1. Tail loop .LBB1_6.

(Interestingly, if you flip the Zig version to trunk, the code calculating the mask after continue becomes a lot bigger)

If I write an assembly loop in RISC-V to do this, I can write it so:

  • The number of elements done each iteration is up to the hardware.
  • The tail is handled by the same loop as the body.
  • I can use multiple vector registers per iteration, like a loop unroll, by setting LMul > 1.

(Forgive any errors. I haven’t actually tried this.)

# Example: Load 32-bit values, count until you find one less than 15
  # On entry:
  #   a0 holds the address of the source array
  #   a1 holds the total number of elements to process
  # During execution:
  #   v1 is used for mask calculations
  #   v4-v11 hold the input data
  #   t0 holds the number of elements this iteration
  #   t2 holds the total count of items with values over 15
  #   t3 is the number of items this iteration with values over 15
count_cp:
    li t2,0                           # `end` = 0
.loop:
    vsetvli t0, a1, e32, m8, tu, mu   # vtype = u32, LMul = 8
                                      # Also set t0 to number of elements for this loop
                                      # VLEN=128 => max(t0) = 4*8 = 32
                                      # VLEN=256 => max(t0) = 8*8 = 64
    vle32.v v4, (a0)                  # Load v4-v11
    sh2add a0, a0, t0                 # Bump pointer

    vmsleu.vi v1, v4, 15              # Calc mask of elements =< 15
    vmsbf.m v1, v1                    # Set bits up to first bit set, clear others
    vcpop.m t3, v1                    # Count bits set 
    add t2, t2, t3                    # Increment `end` by count
    bne t0, t3, .exit                 # Exit if count wasn't maximum number
    sub a1, a1, t3                    # Decrement count
    bnez a1, .loop                    # Any more?
.exit:
    mv a0, t3                         # Return `end`
    ret

That’s special to RISC-V as far as I’m aware.

I can’t figure out how I can tell a compiler to write that code (in any language). I feel like the one line version of the loop should be enough, maybe with a hint to say I expect it to loop many times, and it’s therefore worth vectorizing.

10 Likes

Finally, an explanation of SIMD I can understand.

Now I need a reference to the supported SIMD operations in Zig so I have some idea if SIMD can even be used for a given calculation/processing. Thinking mainly about DSP, but text processing applications (including parsing) are interesting.

1 Like

I did an extensive shoot-out comparing solutions for cross-platform SIMD for my job, including Google Highway, C++, Rust and some other libraries. Zig won by a mile both in reasonable coding effort computational performance and cross-platform code/build ergonomics and became the language of choice for DSP.

Zig SIMD is lovely, codegen seems great and very rarely have I been forced to write in intrinsics though that does still occasionally happen for some tricky twisty stuff.

Currently I’m building a rather fancy FFT pipeline library which uses comptime to select the optimal pipeline from profiling run. A bit like FFTW but planning decision made at comptime. Pipelines beat FFTW and MKL based ones by 15% for our uses, I’m a very happy camper!

12 Likes

ARM also has variable length simd these days; they call it Scalable Vector Extension.

They even want to include variable sized matrix operations in the CPU, called Scalable Matrix Extension.

4 Likes

If possible also extend the data if you will use SIMD to multiples of the Vector size. E.g. in this blog post, have cps be multiple of lanes and pad with 0. Of course, you must be able to control the length of the buffer and have sensible padding bytes. Then you can skip the one-line-loop, because there will be no trailing bytes left.

1 Like

I found this article to be quite nice - It’s intuitive, straight-to-the-point, and generalizes to the types of problems I actually want to use SIMD for. It would be nice to see a followup article for cool SWAR tricks (SIMD within a register) too, though of course not too many people care about that nowadays. :‍)

2 Likes

I hate always having to implement the tail for scalar leftovers, though I have view it as a necessary evil. In my own experience with the data types I work with, it is quite rare for the input to not be divisible by the vector lanes.

My Assembly-foo is moron-tier, so unfortunately implementing it that way is out of reach for me, but how much difference does this achieve? My initial (naive) thought is that the extra cycles used for masking on every pass is about on-par with the (maybe) few scalar ops? I could see how this might depend on the operations actually being done, though.

Well, the vector loop generated from Mitchell’s Zig code is 12 instructions long (7 integer + 5 vector). Mine is 10 instructions long (5 integer + 5 vector), and I don’t have the tail loop. The speed of each will depend on the specific CPU you run it on, but I think the loops are basically equal. You’ll just do 1 extra iteration of mine for any tail, whereas you have to iterate over every item of the tail when you have a separate tail-loop.

Calculating number of instructions for a few block sizes, where there’s no early exit, gives me this table. I’ve assumed 128-bit vector (i.e. 4 items per vector) and I did an LMul = 1 and LMul = 8 version. That’s just like unrolling the vector loop, but without needing more instructions. Basically it cuts the integer operations down by a factor of 8 on big blocks. Where there’s a tail I’ve assumed a worst-case of 3 items left over.

Block Size Body/Tail Loop One Loop (Lmul = 1) % One Loop (LMul = 8) %
12 42 (15v + 27i) 33 (15v + 18i) 78% 23 (15v + 8i) 54%
15 59 (15v + 44i) 43 (20v + 23i) 72% 28 (20v + 8i) 47%
256 776 (320v + 456i) 643 (320v + 323i) 82% 363 (320v + 43i) 46%
259 796 (320v + 476i) 653 (325v + 328i) 82% 373 (325v + 48i) 46%
5 Likes

This might be caused by the changes to @bitCast in https://codeberg.org/ziglang/zig/pulls/35711

I have a dumb question. Why do we need te scalar tail? Why not just pad the final vector with a value that doesn’t skew the results? Would that not be more efficient, at least in some cases?

2 Likes

I never used SIMD, but what I like about the scalar tail is: it shows the intention of the SIMD instructions in plain Zig. This may be worth a few CPU cycles.

1 Like

And to add to this (for the others reading this): Over CPUs with bigger vector register sizes can take better advantage of the same machine code since the machine code automatically adapts to whatever the vector size is.

Here’s the vector register size 128bit, but there are also RISC-V CPUs with a register size of 254bit or even 1024bit (in both cases the SpacemiT K3 is an example) and I heard that some manufacturers even want to go as big as 2048bit, but the spec itself allows I think up to 65536bits.

1 Like

I think the article covers this, that the scaler tail also covers the fallback situation when the target architecture doesnt support SIMD. I like that.

6 Likes

I havent used SIMD myself, and found the article really enlightening, and would have enjoyed experimenting with this back when I was a performance tester on some very big systems. Not sure my lil PinePhone Zig OS project will benefit much, if it even supports SIMD, but gonna have a play anyway :slight_smile:

1 Like

It strongly reminds me of my days working on NEC Vector cards. There was an automatic vectorizer in the compiler that was extremely clever in finding vectorization and SIMD opportunities , but to get the full potential I would occasionally have to address the vector registers by hand.
The structure was basically identical to what is shown here.
Ah yes, good old times

3 Likes

Yes… the software model of the RISC-V spec supports 64k vectors and the first version of the code absolutely killed the compiler :wink:

PinePhone has an ARM A53 which does support SIMD

2 Likes

Hmm I was just checking the architectural reference manual, and I was just thinking that V0 seems familiar syntax… I found an old reference to Arm Neon in my global assembly. It seems I already tried a very basic loop to clear the screen faster without making the connection that this was SIMD hehe. It seems that I ended up going with a simple assembly loop, so I assume Zig compiler was doing something cleverer that was even faster :slight_smile: assuming it optimizes GA?

Instead of the separate tail loop, you could also ensure the input array comes padded with some “dummy” bytes that will produce equivalent output.

Might be worth it if you’re doing a lot of vector operations on your arrays, but won’t work for every situation.

1 Like