(Not) fast dot products via SIMD

Lately I've been tinkering with optimizing Sasha's llama2.rs, a fast Rust port of Karpathy's llama2.c. It takes advantage of Rust nightly's portable_simd, which allows you to emit AVX2 or AVX512 instructions using a relatively clean set of abstractions, and also run inference on a quantized llama2, so it's pretty fast.

One of the core loops looks like this:

let mask = (1 << BITS) - 1; // BITS is 4
let elems_per_i32 = 32 / BITS;
let ipg: usize = GROUPSIZE / 32 * BITS; // GROUPSIZE is 128, so this is 16 (left-associative)
let mask_4bits = i32x8::splat(mask);
let shift_right = i32x8::from_array([0, 4, 8, 12, 16, 20, 24, 28]);
...
// Do K at a time
let zero = f32x8::splat(0.0);
let qzeros = &self.qzeros[oi / elems_per_i32]; // self is the neural net block this is part of
let out_elem = oi % elems_per_i32;
let qweight = self.qweight[oi].chunks_exact(ipg);

let collect = self.scales[oi]
    .into_iter()
    .zip(qweight)
    .enumerate()
    .map(|(group, (scale, qweight))| {
        let qz = ((qzeros[group] >> (BITS * out_elem)) & mask) + 1;
        let scale_simd = f32x8::splat(scale);
        let zero_simd = i32x8::splat(qz);
        let in_pos = group * GROUPSIZE;
        let xs = x[in_pos..in_pos + GROUPSIZE].chunks_exact(8);
        qweight
            .iter()
            .zip(xs)
            .map(|(v, x)| {
                //Extract v into 8 chunks
                let x = f32x8::from_slice(x);
                let num_simd = i32x8::splat(*v);
                let qw: i32x8 = (num_simd >> shift_right) & mask_4bits;
                let combine: f32x8 = (qw - zero_simd).cast::<f32>();
                let weight: f32x8 = scale_simd * combine;
                weight * x
            })
            .fold(zero, |x, y| x + y)
    })
    .fold(zero, |x, y| x + y);
*o = collect.reduce_sum(); // output reference to write to

Here, the f32x8 and i32x8 primitives are aliases for the std::simd types that denote a SIMD vector with 8 lanes of f32 and i32 respectively (256 bits each). I have a 5700X, a Zen 3 CPU with AVX2 support but not AVX512.

When I profiled this, I saw that a lot of the time was being spent in the fold on the inside (i.e. the sum over results from mapping over qweight).

Ok, so this is basically a SIMD dot-product between x[in_pos..in_pos + GROUPSIZE] and whatever the weight becomes. My limited SIMD experience is based on writing CUDA kernels, so naturally I thought that this might be better if I rewrote this to use the fma instructions that have been around for the past few years.

I rewrote the interior slightly:

// Do K at a time
let zero = f32x8::splat(0.0);
let qzeros = &self.qzeros[oi / elems_per_i32];
let out_elem = oi % elems_per_i32;
let qweight = self.qweight[oi].chunks_exact(ipg);

let collect = self.scales[oi]
    .into_iter()
    .zip(qweight)
    .enumerate()
    .map(|(group, (scale, qweight))| {
        let qz = ((qzeros[group] >> (BITS * out_elem)) & mask) + 1;
        let scale_simd = f32x8::splat(scale);
        let qzero_simd = i32x8::splat(qz);
        let in_pos = group * GROUPSIZE;
        let xs = x[in_pos..in_pos + GROUPSIZE]
            .chunks_exact(8)
            .map(f32x8::from_slice);
        let q_op = qweight
            .iter()
            .map(|v| {
                //Extract v into 8 chunks
                let num_simd = i32x8::splat(*v);
                let qw: i32x8 = (num_simd >> shift_right) & mask_4bits;
                (qw - qzero_simd).cast::<f32>()
            });
        xs.zip(q_op).fold(zero, |acc, (x, y)| {
            x.mul_add(y, acc)
        }) * scale_simd
    })
    .fold(zero, |x, y| x + y);
*o = collect.reduce_sum();

However, this tanked performance. I'm working on getting a proper hyperfine benchmark up and running¹, but it's safe to say that the performance is at least 50% worse. I thought issuing 1 instructon (fma) per element of the vector (well, N/8 since I have 8 lanes) would be faster than issuing 2 instructions (mul and add).

The reason, as best as I can understand it (not SIMD expert, yet!) is that in the above code, we're bottlenecked on instruction latency, not throughput. Apparently the computation

result = weight[0] * x[0] + weight[1] * x[1] + ... + x[n] * weight[n]

can be run highly out-of-order because it doesn't matter which multiplication happens first, or technically in which order the results are added². That's why I see the add in the profiling output, because that's the instruction that ends up being where the CPU is waiting for the result of all the individual multiplications to finish.

In contrast:

result = fma(fma(x[2], weight[2], fma(x[1], weight[1], fma(x[0], weight[0]), 0)))

is highly order dependent, so the execution is much slower (also, I think in this case fma latency is 4 cycles vs 3 cycles for add and mul, which makes things much slower). Compilers are pretty smart! And modern out-of-order execution is pretty impressive.

Here's the people who helped me figure this out:

Is vfmadd132pd slow on AMD Zen 3 architecture?

Latency and throughput