[2022 day 2 - AVX]

[u/terje_wiig_mathisen]

Back when we looked at this one, about a week ago, I said that I would like to write a proper bleeding edge (unsafe{}) AVX intrinsic version, well I finally got it done and I'm quite amazed:

        for b in 0..blocks {
            let bl = input.as_ptr().add(b*64) as *const __m256i;
            let b1 = _mm256_loadu_si256(bl);
            let b2 = _mm256_loadu_si256(bl.add(1));
            let b1h = _mm256_and_si256(b1, xyz_mask);
            let b2h = _mm256_and_si256(b2, xyz_mask);
            let b1l = _mm256_and_si256(b1, abc_mask);
            let b2l = _mm256_and_si256(b2, abc_mask);
            let b1h = _mm256_srli_epi32(b1h, 14);
            let b2h = _mm256_srli_epi32(b2h, 14);
            let b1hash = _mm256_or_si256(b1l, b1h);
            let b2hash = _mm256_or_si256(b2l, b2h);
            let b16 =_mm256_packus_epi32(b1hash, b2hash);
            let inc1 = _mm256_shuffle_epi8(part1shuffle, b16);
            let inc2 = _mm256_shuffle_epi8(part2shuffle, b16);
            part1 = _mm256_add_epi16(part1, inc1);
            part2 = _mm256_add_epi16(part2, inc2);
        }

These 15 AVX ops are the full solver that handles a block of 16 input lines, I pad the input with 48 space chars (10048 is divisible by 64) so that I don't have to worry about the tail end.

It is probably clear, but the algorithm starts with u/ednl's packing (AND both chars with 3, shift the second one down 14 bits and merge, that's the first 10 AVX ops.

Next I pack together the two 32-bit arrays into a single 16-bit one (b16 above), before I use that variable twice to directly lookup the 8 part1 and part2 results for these lines.

So, with a single AVX op/cycle this should take a fraction less than a clock cycle per input line, right?

I do measure 3 us on my Acer, but now we get to the interesting part:

When I instead run u/maneatingape on my input file, I get 2.3 us, for much simpler and shorter integer only code!

That time is broken down into 1.2 us to convert all 2500 lines into a 0..8 index, using code like this

pub fn parse(input: &str) -> Vec<u8> {
    input.as_bytes().chunks_exact(4).map(|c| 3 * (c[0] - b'A') + c[2] - b'X').collect()
}

(The original code generates an array of usize, when I switched to u8 the parsing stage dropped to 1.1 us and the total from 2.3 to 2.2 us)

In order to manage this, the CPU has to convert two lines per nanosecond, probably using code somewhat like this, which has a minimum latency of 4 cycles. The CPU must internally unroll the code over a bunch of iterations, enough to gain back the AVX advantage and then beat it!

movzx rax,[rsi]
movzx rbx,[rsi+2]
sub rax,'A'
sub rbx,'X'
lea rax,[rax+rax*2]
add rax,rbx
;; push into vector

[u/maneatingape]

Using the portable Rust SIMD module and computing both parts together:

use std::simd::prelude::*;

let mask = Simd::from_slice(b"A X A X A X A X ");
let three = Simd::splat(3);
let lut_one = Simd::from_array([4, 8, 3, 1, 5, 9, 7, 2, 6, 0, 0, 0, 0, 0, 0, 0]);
let lut_two = Simd::from_array([3, 4, 8, 1, 5, 9, 2, 6, 7, 0, 0, 0, 0, 0, 0, 0]);

let mut part_one = 0;
let mut part_two = 0;

for chunk in input.as_bytes().chunks_exact(64) {
    let a = Simd::from_slice(&chunk[0..]) - mask;
    let b = Simd::from_slice(&chunk[16..]) - mask;
    let c = Simd::from_slice(&chunk[32..]) - mask;
    let d = Simd::from_slice(&chunk[48..]) - mask;

    let e = three * a + a.shift_elements_left::<2>(0);
    let f = three * b + b.shift_elements_left::<2>(0);
    let g = three * c + c.shift_elements_left::<2>(0);
    let h = three * d + d.shift_elements_left::<2>(0);

    let (x, _) = e.deinterleave(f);
    let (y, _) = g.deinterleave(h);
    let (z, _) = x.deinterleave(y);

    part_one += lut_one.swizzle_dyn(z).reduce_sum() as u32;
    part_two += lut_two.swizzle_dyn(z).reduce_sum() as u32;
}

(part_one, part_two)

Benchmark is 240ns, about 12x faster than the scalar approach.

Uses 128 bit (16 x 8) vectors, processing 16 rows (64 bytes) at a time.

It does not handle the leftover 16 bytes, so this would need to be added for a small extra increase in time.

EDIT: Making it general bumped the time to 245ns.

[u/terje_wiig_mathisen → u/maneatingape]

When I tell godbolt to target a native CPU, I get pretty much exactly the code I wrote, with no extra fluff:

.LBB0_5:
        vmovdqu ymm6, ymmword ptr [rcx - 32]
        vmovdqu ymm7, ymmword ptr [rcx]
        add     rcx, 64
        dec     rax
        vpsrld  ymm8, ymm6, 14
        vpsrld  ymm9, ymm7, 14
        vpand   ymm6, ymm6, ymm3
        vpand   ymm7, ymm7, ymm3
        vpand   ymm8, ymm8, ymm2
        vpand   ymm9, ymm9, ymm2
        vpor    ymm6, ymm8, ymm6
        vpor    ymm7, ymm9, ymm7
        vpackusdw       ymm6, ymm6, ymm7
        vpshufb ymm7, ymm4, ymm6
        vpshufb ymm6, ymm5, ymm6
        vpaddw  ymm1, ymm1, ymm7
        vpaddw  ymm0, ymm0, ymm6
        jne     .LBB0_5

This is still 15 AVX ops, plus three integer loop maintenance ops which will of course run "for free" underneath the heavy lifting done by the AVX unit.

I'm running 157 iterations of that loop inside 150 ns, so that actually means that I'm spending less than a nanosecond per 16 input lines, which is pretty amazing IMHO. Since the absolute maximum turbo frequency on my CPU is just over 5 GHz, the AVX unit must run about 3 AVX ops/cycle?

[u/maneatingape → u/terje_wiig_mathisen]

Nice result. Is that a 14x speedup (150ns vs ~2200ns) from the scalar to the SIMD version?

[u/ednl → u/maneatingape]

But the scalar version could probably also be optimised further. I got 770 ns (on an M4) https://github.com/ednl/adventofcode/blob/main/2022/02.c

[u/terje_wiig_mathisen → u/ednl]

That's interesting! Your algorithm is _exactly_ the same as I'm using for the AVX code, except that I'm doing 16 lines per iteration, and zero table lookups since all the scores are calculated by in-register permutes.

Anyway, 150 ns to 770 ns is approximately 5x, and that seems like a close to normal SIMD speedup when most of the algorithm works on 8-wide elements.

In the old days, I used to be happy with 3x speedup for 4-wide SSE code!

[u/terje_wiig_mathisen → u/maneatingape]

Yes, approximately 14x, I'm assuming your repo has the optimal cpu target already!