slain_num/

simd.rs

#[cfg(target_arch = "x86_64")]
#[inline]
fn fast() -> bool {
    slain_rt::cpu::avx2()
}

#[must_use]
pub fn dot(a: &[f64], b: &[f64]) -> f64 {
    let n = a.len().min(b.len());
    #[cfg(target_arch = "x86_64")]
    if fast() {
        return unsafe { x86::dot(&a[..n], &b[..n]) };
    }
    lanes::dot(&a[..n], &b[..n])
}

pub fn axpy(k: f64, x: &[f64], y: &mut [f64]) {
    let n = x.len().min(y.len());
    #[cfg(target_arch = "x86_64")]
    if fast() {
        unsafe { x86::axpy(k, &x[..n], &mut y[..n]) };
        return;
    }
    lanes::axpy(k, &x[..n], &mut y[..n]);
}

#[must_use]
pub fn sum(a: &[f64]) -> f64 {
    #[cfg(target_arch = "x86_64")]
    if fast() {
        return unsafe { x86::sum(a) };
    }
    lanes::sum(a)
}

#[must_use]
pub fn sq_dev(a: &[f64], m: f64) -> f64 {
    #[cfg(target_arch = "x86_64")]
    if fast() {
        return unsafe { x86::sq_dev(a, m) };
    }
    lanes::sq_dev(a, m)
}

#[must_use]
pub fn sq_diff_u16(a: &[u16], b: &[u16]) -> u64 {
    let n = a.len().min(b.len());
    #[cfg(target_arch = "x86_64")]
    if fast() {
        return unsafe { x86::sq_diff_u16(&a[..n], &b[..n]) };
    }
    a[..n].iter().zip(&b[..n]).map(|(&x, &y)| (i64::from(x) - i64::from(y)).pow(2) as u64).sum()
}

pub fn gram_rhs(cols: &[&[f64]], t: &[f64], gram: &mut [f64], rhs: &mut [f64]) {
    let k = cols.len();
    assert!(gram.len() >= k * k && rhs.len() >= k, "gram k*k and rhs k");
    let n = cols.iter().map(|c| c.len()).chain([t.len()]).min().unwrap_or(0);
    if k <= 4 {
        let mut c = [&[][..]; 4];
        for (d, s) in c.iter_mut().zip(cols) {
            *d = &s[..n];
        }
        let mut g = [0.0; 10];
        let mut r = [0.0; 4];
        let t = &t[..n];
        #[cfg(target_arch = "x86_64")]
        let done = fast()
            && unsafe {
                match k {
                    1 => x86::gram_rhs4::<1>(&c, t, &mut g, &mut r),
                    2 => x86::gram_rhs4::<2>(&c, t, &mut g, &mut r),
                    3 => x86::gram_rhs4::<3>(&c, t, &mut g, &mut r),
                    4 => x86::gram_rhs4::<4>(&c, t, &mut g, &mut r),
                    _ => {}
                }
                true
            };
        #[cfg(not(target_arch = "x86_64"))]
        let done = false;
        if !done {
            match k {
                1 => lanes::gram_rhs4::<1>(&c, t, &mut g, &mut r),
                2 => lanes::gram_rhs4::<2>(&c, t, &mut g, &mut r),
                3 => lanes::gram_rhs4::<3>(&c, t, &mut g, &mut r),
                4 => lanes::gram_rhs4::<4>(&c, t, &mut g, &mut r),
                _ => {}
            }
        }
        let mut idx = 0;
        for i in 0..k {
            for j in i..k {
                gram[i * k + j] = g[idx];
                gram[j * k + i] = g[idx];
                idx += 1;
            }
            rhs[i] = r[i];
        }
        return;
    }
    for i in 0..k {
        for j in i..k {
            let v = dot(&cols[i][..n], &cols[j][..n]);
            gram[i * k + j] = v;
            gram[j * k + i] = v;
        }
        rhs[i] = dot(&cols[i][..n], &t[..n]);
    }
}

mod lanes {
    type L = [f64; 4];
    const Z: L = [0.0; 4];

    #[inline(always)]
    fn ld(x: &[f64], i: usize) -> L {
        [x[i], x[i + 1], x[i + 2], x[i + 3]]
    }

    #[inline(always)]
    fn add(a: L, b: L) -> L {
        [a[0] + b[0], a[1] + b[1], a[2] + b[2], a[3] + b[3]]
    }

    #[inline(always)]
    fn fma(a: L, b: L, c: L) -> L {
        [a[0].mul_add(b[0], c[0]), a[1].mul_add(b[1], c[1]), a[2].mul_add(b[2], c[2]), a[3].mul_add(b[3], c[3])]
    }

    #[inline(always)]
    fn sub(a: L, m: f64) -> L {
        [a[0] - m, a[1] - m, a[2] - m, a[3] - m]
    }

    #[inline(always)]
    fn hsum(v: L) -> f64 {
        (v[0] + v[2]) + (v[1] + v[3])
    }

    #[inline(always)]
    fn hsum8(s: [L; 8]) -> f64 {
        let a = add(add(s[0], s[1]), add(s[2], s[3]));
        let b = add(add(s[4], s[5]), add(s[6], s[7]));
        hsum(add(a, b))
    }

    #[inline(always)]
    fn reduce(a: &[f64], mut s: [L; 8], step: impl Fn(L, usize, L) -> L, tail: impl Fn(usize) -> f64) -> f64 {
        let n = a.len();
        let mut i = 0;
        while i + 32 <= n {
            for (j, acc) in s.iter_mut().enumerate() {
                *acc = step(*acc, i + 4 * j, ld(a, i + 4 * j));
            }
            i += 32;
        }
        let mut j = 0;
        while i + 4 <= n {
            s[j] = step(s[j], i, ld(a, i));
            j += 1;
            i += 4;
        }
        let mut t = hsum8(s);
        while i < n {
            t += tail(i);
            i += 1;
        }
        t
    }

    pub(super) fn dot(a: &[f64], b: &[f64]) -> f64 {
        reduce(a, [Z; 8], |acc, i, x| fma(x, ld(b, i), acc), |i| a[i] * b[i])
    }

    pub(super) fn sum(a: &[f64]) -> f64 {
        reduce(a, [Z; 8], |acc, _, x| add(acc, x), |i| a[i])
    }

    pub(super) fn sq_dev(a: &[f64], m: f64) -> f64 {
        reduce(
            a,
            [Z; 8],
            |acc, _, x| {
                let d = sub(x, m);
                fma(d, d, acc)
            },
            |i| (a[i] - m) * (a[i] - m),
        )
    }

    pub(super) fn axpy(k: f64, x: &[f64], y: &mut [f64]) {
        let n = x.len();
        let vk = [k; 4];
        let mut i = 0;
        while i + 4 <= n {
            let r = fma(vk, ld(x, i), ld(y, i));
            y[i..i + 4].copy_from_slice(&r);
            i += 4;
        }
        while i < n {
            y[i] += k * x[i];
            i += 1;
        }
    }

    pub(super) fn gram_rhs4<const K: usize>(cols: &[&[f64]; 4], t: &[f64], g: &mut [f64; 10], r: &mut [f64; 4]) {
        let n = t.len();
        let mut acc = [Z; 10];
        let mut rh = [Z; 4];
        let mut i = 0;
        while i + 4 <= n {
            let tv = ld(t, i);
            let mut v = [Z; 4];
            for c in 0..K {
                v[c] = ld(cols[c], i);
            }
            let mut idx = 0;
            for a in 0..K {
                for b in a..K {
                    acc[idx] = fma(v[a], v[b], acc[idx]);
                    idx += 1;
                }
                rh[a] = fma(v[a], tv, rh[a]);
            }
            i += 4;
        }
        let mut idx = 0;
        for a in 0..K {
            for b in a..K {
                let mut s = hsum(acc[idx]);
                for row in i..n {
                    s += cols[a][row] * cols[b][row];
                }
                g[idx] = s;
                idx += 1;
            }
            let mut s = hsum(rh[a]);
            for row in i..n {
                s += cols[a][row] * t[row];
            }
            r[a] = s;
        }
    }
}

#[cfg(target_arch = "x86_64")]
mod x86 {
    use core::arch::x86_64::*;

    #[inline(always)]
    unsafe fn hsum(v: __m256d) -> f64 {
        let s = _mm_add_pd(_mm256_castpd256_pd128(v), _mm256_extractf128_pd::<1>(v));
        _mm_cvtsd_f64(_mm_add_sd(s, _mm_unpackhi_pd(s, s)))
    }

    #[inline(always)]
    unsafe fn hsum8(s: [__m256d; 8]) -> f64 {
        let a = _mm256_add_pd(_mm256_add_pd(s[0], s[1]), _mm256_add_pd(s[2], s[3]));
        let b = _mm256_add_pd(_mm256_add_pd(s[4], s[5]), _mm256_add_pd(s[6], s[7]));
        hsum(_mm256_add_pd(a, b))
    }

    #[target_feature(enable = "avx2,fma")]
    pub(super) unsafe fn dot(a: &[f64], b: &[f64]) -> f64 {
        let n = a.len();
        let (pa, pb) = (a.as_ptr(), b.as_ptr());
        let mut s = [_mm256_setzero_pd(); 8];
        let mut i = 0;
        while i + 32 <= n {
            s[0] = _mm256_fmadd_pd(_mm256_loadu_pd(pa.add(i)), _mm256_loadu_pd(pb.add(i)), s[0]);
            s[1] = _mm256_fmadd_pd(_mm256_loadu_pd(pa.add(i + 4)), _mm256_loadu_pd(pb.add(i + 4)), s[1]);
            s[2] = _mm256_fmadd_pd(_mm256_loadu_pd(pa.add(i + 8)), _mm256_loadu_pd(pb.add(i + 8)), s[2]);
            s[3] = _mm256_fmadd_pd(_mm256_loadu_pd(pa.add(i + 12)), _mm256_loadu_pd(pb.add(i + 12)), s[3]);
            s[4] = _mm256_fmadd_pd(_mm256_loadu_pd(pa.add(i + 16)), _mm256_loadu_pd(pb.add(i + 16)), s[4]);
            s[5] = _mm256_fmadd_pd(_mm256_loadu_pd(pa.add(i + 20)), _mm256_loadu_pd(pb.add(i + 20)), s[5]);
            s[6] = _mm256_fmadd_pd(_mm256_loadu_pd(pa.add(i + 24)), _mm256_loadu_pd(pb.add(i + 24)), s[6]);
            s[7] = _mm256_fmadd_pd(_mm256_loadu_pd(pa.add(i + 28)), _mm256_loadu_pd(pb.add(i + 28)), s[7]);
            i += 32;
        }
        let mut j = 0;
        while i + 4 <= n {
            s[j] = _mm256_fmadd_pd(_mm256_loadu_pd(pa.add(i)), _mm256_loadu_pd(pb.add(i)), s[j]);
            j += 1;
            i += 4;
        }
        let mut t = hsum8(s);
        while i < n {
            t += a[i] * b[i];
            i += 1;
        }
        t
    }

    #[target_feature(enable = "avx2,fma")]
    pub(super) unsafe fn axpy(k: f64, x: &[f64], y: &mut [f64]) {
        let n = x.len();
        let vk = _mm256_set1_pd(k);
        let mut i = 0;
        while i + 4 <= n {
            let r = _mm256_fmadd_pd(vk, _mm256_loadu_pd(x.as_ptr().add(i)), _mm256_loadu_pd(y.as_ptr().add(i)));
            _mm256_storeu_pd(y.as_mut_ptr().add(i), r);
            i += 4;
        }
        while i < n {
            y[i] += k * x[i];
            i += 1;
        }
    }

    #[target_feature(enable = "avx2")]
    pub(super) unsafe fn sum(a: &[f64]) -> f64 {
        let n = a.len();
        let p = a.as_ptr();
        let mut s = [_mm256_setzero_pd(); 8];
        let mut i = 0;
        while i + 32 <= n {
            for (j, acc) in s.iter_mut().enumerate() {
                *acc = _mm256_add_pd(*acc, _mm256_loadu_pd(p.add(i + 4 * j)));
            }
            i += 32;
        }
        let mut j = 0;
        while i + 4 <= n {
            s[j] = _mm256_add_pd(s[j], _mm256_loadu_pd(p.add(i)));
            j += 1;
            i += 4;
        }
        let mut t = hsum8(s);
        while i < n {
            t += a[i];
            i += 1;
        }
        t
    }

    #[target_feature(enable = "avx2,fma")]
    pub(super) unsafe fn sq_dev(a: &[f64], m: f64) -> f64 {
        let n = a.len();
        let p = a.as_ptr();
        let vm = _mm256_set1_pd(m);
        let mut s = [_mm256_setzero_pd(); 8];
        let mut i = 0;
        while i + 32 <= n {
            for (j, acc) in s.iter_mut().enumerate() {
                let d = _mm256_sub_pd(_mm256_loadu_pd(p.add(i + 4 * j)), vm);
                *acc = _mm256_fmadd_pd(d, d, *acc);
            }
            i += 32;
        }
        let mut j = 0;
        while i + 4 <= n {
            let d = _mm256_sub_pd(_mm256_loadu_pd(p.add(i)), vm);
            s[j] = _mm256_fmadd_pd(d, d, s[j]);
            j += 1;
            i += 4;
        }
        let mut t = hsum8(s);
        while i < n {
            t += (a[i] - m) * (a[i] - m);
            i += 1;
        }
        t
    }

    #[target_feature(enable = "avx2")]
    pub(super) unsafe fn sq_diff_u16(a: &[u16], b: &[u16]) -> u64 {
        let n = a.len();
        let zero = _mm256_setzero_si256();
        let (mut acc0, mut acc1) = (zero, zero);
        let mut i = 0;
        while i + 16 <= n {
            let pa = _mm256_loadu_si256(a.as_ptr().add(i).cast());
            let pb = _mm256_loadu_si256(b.as_ptr().add(i).cast());
            if _mm256_movemask_epi8(_mm256_or_si256(pa, pb)) & 0xAAAA_AAAAu32 as i32 == 0 {
                let d = _mm256_sub_epi16(pa, pb);
                let sq = _mm256_madd_epi16(d, d);
                acc0 = _mm256_add_epi64(acc0, _mm256_unpacklo_epi32(sq, zero));
                acc1 = _mm256_add_epi64(acc1, _mm256_unpackhi_epi32(sq, zero));
            } else {
                let lo = _mm256_sub_epi32(_mm256_unpacklo_epi16(pa, zero), _mm256_unpacklo_epi16(pb, zero));
                let hi = _mm256_sub_epi32(_mm256_unpackhi_epi16(pa, zero), _mm256_unpackhi_epi16(pb, zero));
                let odd = |v: __m256i| _mm256_srli_epi64::<32>(v);
                acc0 = _mm256_add_epi64(acc0, _mm256_add_epi64(_mm256_mul_epi32(lo, lo), _mm256_mul_epi32(odd(lo), odd(lo))));
                acc1 = _mm256_add_epi64(acc1, _mm256_add_epi64(_mm256_mul_epi32(hi, hi), _mm256_mul_epi32(odd(hi), odd(hi))));
            }
            i += 16;
        }
        let mut lanes = [0u64; 4];
        _mm256_storeu_si256(lanes.as_mut_ptr().cast(), _mm256_add_epi64(acc0, acc1));
        let mut total: u64 = lanes.iter().sum();
        while i < n {
            total += (i64::from(a[i]) - i64::from(b[i])).pow(2) as u64;
            i += 1;
        }
        total
    }

    #[target_feature(enable = "avx2,fma")]
    pub(super) unsafe fn gram_rhs4<const K: usize>(cols: &[&[f64]; 4], t: &[f64], g: &mut [f64; 10], r: &mut [f64; 4]) {
        let n = t.len();
        let z = _mm256_setzero_pd();
        let mut acc = [z; 10];
        let mut rh = [z; 4];
        let p: [*const f64; 4] = [cols[0].as_ptr(), cols[1].as_ptr(), cols[2].as_ptr(), cols[3].as_ptr()];
        let mut i = 0;
        while i + 4 <= n {
            let tv = _mm256_loadu_pd(t.as_ptr().add(i));
            let mut v = [z; 4];
            for c in 0..K {
                v[c] = _mm256_loadu_pd(p[c].add(i));
            }
            let mut idx = 0;
            for a in 0..K {
                for b in a..K {
                    acc[idx] = _mm256_fmadd_pd(v[a], v[b], acc[idx]);
                    idx += 1;
                }
                rh[a] = _mm256_fmadd_pd(v[a], tv, rh[a]);
            }
            i += 4;
        }
        let k = K;
        let mut idx = 0;
        for a in 0..k {
            for b in a..k {
                let mut s = hsum(acc[idx]);
                for row in i..n {
                    s += cols[a][row] * cols[b][row];
                }
                g[idx] = s;
                idx += 1;
            }
            let mut s = hsum(rh[a]);
            for row in i..n {
                s += cols[a][row] * t[row];
            }
            r[a] = s;
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn close(a: f64, b: f64) -> bool {
        (a - b).abs() <= 1e-9 * (1.0 + a.abs().max(b.abs()))
    }

    #[test]
    fn kernels_match_scalar() {
        for len in [0usize, 1, 3, 4, 7, 31, 32, 33, 100, 1003] {
            let a: Vec<f64> = (0..len).map(|i| (i as f64 * 0.37).sin()).collect();
            let b: Vec<f64> = (0..len).map(|i| (i as f64 * 0.11).cos()).collect();
            let want: f64 = a.iter().zip(&b).map(|(x, y)| x * y).sum();
            assert!(close(dot(&a, &b), want), "dot len {len}");
            assert!(close(sum(&a), a.iter().sum::<f64>()), "sum len {len}");
            let m = 0.25;
            assert!(close(sq_dev(&a, m), a.iter().map(|v| (v - m) * (v - m)).sum()), "sq_dev len {len}");
            let mut y = b.clone();
            axpy(2.5, &a, &mut y);
            assert!(y.iter().zip(a.iter().zip(&b)).all(|(o, (x, z))| (o - (z + 2.5 * x)).abs() < 1e-12));
        }
        for (len, top) in [(1037usize, 65536u32), (1037, 32768), (1037, 4096), (15, 1024), (16, 32768), (40, 65536)] {
            let p: Vec<u16> = (0..len).map(|i| (i as u32 * 7919 % top) as u16).collect();
            let q: Vec<u16> = (0..len).map(|i| (i as u32 * 104_729 % top) as u16).collect();
            let want: u64 = p.iter().zip(&q).map(|(&x, &y)| (i64::from(x) - i64::from(y)).pow(2) as u64).sum();
            assert_eq!(sq_diff_u16(&p, &q), want, "sq_diff len {len} range {top}");
            assert_eq!(sq_diff_u16(&p, &q), crate::simd_prev::sq_diff_u16(&p, &q));
        }
    }

    #[cfg(target_arch = "x86_64")]
    #[test]
    fn lanes_equal_the_x86_kernels_bitwise() {
        if !fast() {
            return;
        }
        for len in [0usize, 1, 3, 4, 7, 31, 32, 33, 100, 1003] {
            let a: Vec<f64> = (0..len).map(|i| (i as f64 * 0.37).sin()).collect();
            let b: Vec<f64> = (0..len).map(|i| (i as f64 * 0.11).cos()).collect();
            unsafe {
                assert_eq!(lanes::dot(&a, &b).to_bits(), x86::dot(&a, &b).to_bits(), "dot len {len}");
                assert_eq!(lanes::sum(&a).to_bits(), x86::sum(&a).to_bits(), "sum len {len}");
                assert_eq!(lanes::sq_dev(&a, 0.25).to_bits(), x86::sq_dev(&a, 0.25).to_bits(), "sq_dev len {len}");
                let (mut y0, mut y1) = (b.clone(), b.clone());
                lanes::axpy(2.5, &a, &mut y0);
                x86::axpy(2.5, &a, &mut y1);
                assert!(y0.iter().zip(&y1).all(|(p, q)| p.to_bits() == q.to_bits()), "axpy len {len}");
                let c = [&a[..], &b[..], &a[..], &b[..]];
                let (mut g0, mut r0, mut g1, mut r1) = ([0.0; 10], [0.0; 4], [0.0; 10], [0.0; 4]);
                lanes::gram_rhs4::<3>(&c, &b, &mut g0, &mut r0);
                x86::gram_rhs4::<3>(&c, &b, &mut g1, &mut r1);
                assert_eq!(g0.map(f64::to_bits), g1.map(f64::to_bits), "gram len {len}");
                assert_eq!(r0.map(f64::to_bits), r1.map(f64::to_bits), "rhs len {len}");
            }
        }
    }

    #[test]
    fn gram_rhs_matches_dots() {
        for k in 0..=6usize {
            for len in [0usize, 3, 4, 17, 4096] {
                let cols: Vec<Vec<f64>> = (0..k).map(|c| (0..len).map(|i| ((i * (c + 3)) as f64 * 0.013).sin()).collect()).collect();
                let t: Vec<f64> = (0..len).map(|i| (i as f64 * 0.007).cos()).collect();
                let refs: Vec<&[f64]> = cols.iter().map(Vec::as_slice).collect();
                let (mut g, mut r) = (vec![0.0; k * k], vec![0.0; k]);
                gram_rhs(&refs, &t, &mut g, &mut r);
                for i in 0..k {
                    for j in 0..k {
                        let want: f64 = cols[i].iter().zip(&cols[j]).map(|(x, y)| x * y).sum();
                        assert!(close(g[i * k + j], want), "k {k} len {len} g[{i}][{j}]");
                    }
                    let want: f64 = cols[i].iter().zip(&t).map(|(x, y)| x * y).sum();
                    assert!(close(r[i], want), "k {k} len {len} rhs[{i}]");
                }
            }
        }
    }

    #[test]
    #[ignore = "bench"]
    fn simd_bench() {
        use crate::bench::{slain, Harness};
        use crate::simd_prev as prev;
        let mut h = Harness::default();
        let n = 4096;
        let a: Vec<f64> = (0..n).map(|i| (i as f64 * 0.37).sin()).collect();
        let b: Vec<f64> = (0..n).map(|i| (i as f64 * 0.11).cos()).collect();
        let p: Vec<u16> = (0..n).map(|i| (i * 7919 % 1024) as u16).collect();
        let q: Vec<u16> = (0..n).map(|i| (i * 104_729 % 1024) as u16).collect();
        let cols: Vec<Vec<f64>> = (0..3).map(|c| (0..n).map(|i| ((i * (c + 3)) as f64 * 0.013).sin()).collect()).collect();
        let refs: Vec<&[f64]> = cols.iter().map(Vec::as_slice).collect();
        let mut g = h.benchmark_group("dot 4096 f64");
        g.bench_function("previous", |bn| bn.iter(|| prev::dot(slain(&a), slain(&b))));
        g.bench_function("register-mapped", |bn| bn.iter(|| dot(slain(&a), slain(&b))));
        g.finish();
        let mut g = h.benchmark_group("sum 4096 f64");
        g.bench_function("previous", |bn| bn.iter(|| prev::sum(slain(&a))));
        g.bench_function("register-mapped", |bn| bn.iter(|| sum(slain(&a))));
        g.finish();
        let mut g = h.benchmark_group("sq_diff_u16 4096 (10-bit)");
        g.bench_function("previous", |bn| bn.iter(|| prev::sq_diff_u16(slain(&p), slain(&q))));
        g.bench_function("register-mapped", |bn| bn.iter(|| sq_diff_u16(slain(&p), slain(&q))));
        g.finish();
        let mut g = h.benchmark_group("gram + rhs, 3 columns x 4096 (Wiener fit)");
        g.bench_function("previous (9 dots)", |bn| {
            bn.iter(|| {
                let mut o = [0.0; 12];
                for i in 0..3 {
                    for j in i..3 {
                        o[i * 3 + j] = prev::dot(slain(refs[i]), slain(refs[j]));
                    }
                    o[9 + i] = prev::dot(slain(refs[i]), slain(&a));
                }
                o
            })
        });
        g.bench_function("register-mapped (one pass)", |bn| {
            bn.iter(|| {
                let (mut gm, mut r) = ([0.0; 9], [0.0; 3]);
                gram_rhs(slain(&refs), slain(&a), &mut gm, &mut r);
                (gm, r)
            })
        });
        g.finish();
    }
}