#[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();
}
}