slain-num 0.4.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

use core::arch::x86_64::*;

#[inline(always)]
unsafe fn ld(p: *const u16) -> __m256i {
    _mm256_loadu_si256(p.cast())
}

#[inline(always)]
unsafe fn widen_add(acc: __m256i, v: __m256i) -> __m256i {
    let z = _mm256_setzero_si256();
    _mm256_add_epi32(acc, _mm256_add_epi32(_mm256_unpacklo_epi16(v, z), _mm256_unpackhi_epi16(v, z)))
}

#[inline(always)]
unsafe fn hsum32(v: __m256i) -> u64 {
    let z = _mm256_setzero_si256();
    let q = _mm256_add_epi64(_mm256_unpacklo_epi32(v, z), _mm256_unpackhi_epi32(v, z));
    let mut l = [0u64; 4];
    _mm256_storeu_si256(l.as_mut_ptr().cast(), q);
    l[0] + l[1] + l[2] + l[3]
}

#[target_feature(enable = "avx2")]
pub unsafe fn sum(a: &[u16]) -> u64 {
    let n = a.len();
    let p = a.as_ptr();
    let mut total = 0u64;
    let mut i = 0;
    while i + 16 <= n {
        let end = (i + 16 * 16384).min(n - n % 16);
        let mut acc = _mm256_setzero_si256();
        while i < end {
            acc = widen_add(acc, ld(p.add(i)));
            i += 16;
        }
        total += hsum32(acc);
    }
    total + crate::nd::scalar::sum(&a[i..])
}

#[target_feature(enable = "avx2")]
pub unsafe fn sad(a: &[u16], b: &[u16]) -> u64 {
    let n = a.len().min(b.len());
    let (pa, pb) = (a.as_ptr(), b.as_ptr());
    let mut total = 0u64;
    let mut i = 0;
    while i + 16 <= n {
        let end = (i + 16 * 16384).min(n - n % 16);
        let mut acc = _mm256_setzero_si256();
        while i < end {
            let (x, y) = (ld(pa.add(i)), ld(pb.add(i)));
            acc = widen_add(acc, _mm256_sub_epi16(_mm256_max_epu16(x, y), _mm256_min_epu16(x, y)));
            i += 16;
        }
        total += hsum32(acc);
    }
    total + crate::nd::scalar::sad(&a[i..n], &b[i..n])
}

#[target_feature(enable = "avx2")]
pub unsafe fn count_range(a: &[u16], lo: u16, hi: u16) -> u64 {
    if lo > hi {
        return 0;
    }
    let n = a.len();
    let p = a.as_ptr();
    let (vl, vh, z) = (_mm256_set1_epi16(lo as i16), _mm256_set1_epi16(hi as i16), _mm256_setzero_si256());
    let mut total = 0u64;
    let mut i = 0;
    while i + 16 <= n {
        let end = (i + 16 * 32768).min(n - n % 16);
        let mut c = z;
        while i < end {
            let v = ld(p.add(i));
            c = _mm256_sub_epi16(c, _mm256_cmpeq_epi16(_mm256_min_epu16(_mm256_max_epu16(v, vl), vh), v));
            i += 16;
        }
        total += hsum32(widen_add(z, c));
    }
    total + crate::nd::scalar::count_range(&a[i..], lo, hi)
}

#[target_feature(enable = "avx2")]
pub unsafe fn min_max(a: &[u16]) -> (u16, u16) {
    let n = a.len();
    let p = a.as_ptr();
    let (mut lo, mut hi) = (_mm256_set1_epi16(-1), _mm256_setzero_si256());
    let mut i = 0;
    while i + 16 <= n {
        let v = ld(p.add(i));
        lo = _mm256_min_epu16(lo, v);
        hi = _mm256_max_epu16(hi, v);
        i += 16;
    }
    let (mut l, mut h) = ([0u16; 16], [0u16; 16]);
    _mm256_storeu_si256(l.as_mut_ptr().cast(), lo);
    _mm256_storeu_si256(h.as_mut_ptr().cast(), hi);
    let (tl, th) = crate::nd::scalar::min_max(&a[i..]);
    (l.iter().copied().fold(tl, u16::min), h.iter().copied().fold(th, u16::max))
}