slain-num 0.3.0

SLAIN's own numeric crate: dense matrices, least squares, polynomials, FFT, convolution, statistics and 1-D/N-D minimisation, with AVX2 kernels.
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110

#[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]) };
    }
    a[..n].iter().zip(&b[..n]).map(|(x, y)| x * y).sum()
}

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

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

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

    #[target_feature(enable = "avx2,fma")]
    pub(super) unsafe fn dot(a: &[f64], b: &[f64]) -> f64 {
        let n = a.len();
        let (mut s0, mut s1) = (_mm256_setzero_pd(), _mm256_setzero_pd());
        let mut i = 0;
        while i + 8 <= n {
            s0 = _mm256_fmadd_pd(_mm256_loadu_pd(a.as_ptr().add(i)), _mm256_loadu_pd(b.as_ptr().add(i)), s0);
            s1 = _mm256_fmadd_pd(_mm256_loadu_pd(a.as_ptr().add(i + 4)), _mm256_loadu_pd(b.as_ptr().add(i + 4)), s1);
            i += 8;
        }
        if i + 4 <= n {
            s0 = _mm256_fmadd_pd(_mm256_loadu_pd(a.as_ptr().add(i)), _mm256_loadu_pd(b.as_ptr().add(i)), s0);
            i += 4;
        }
        let mut s = hsum(_mm256_add_pd(s0, s1));
        while i < n {
            s += a[i] * b[i];
            i += 1;
        }
        s
    }

    #[target_feature(enable = "avx2")]
    pub(super) unsafe fn sum(a: &[f64]) -> f64 {
        let n = a.len();
        let mut s = _mm256_setzero_pd();
        let mut i = 0;
        while i + 4 <= n {
            s = _mm256_add_pd(s, _mm256_loadu_pd(a.as_ptr().add(i)));
            i += 4;
        }
        let mut t = hsum(s);
        while i < n {
            t += a[i];
            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 acc = 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());
            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);
            acc = _mm256_add_epi64(acc, _mm256_add_epi64(_mm256_mul_epi32(lo, lo), _mm256_mul_epi32(odd(lo), odd(lo))));
            acc = _mm256_add_epi64(acc, _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(), acc);
        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
    }
}