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
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146

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

use super::SadFn;

#[inline(always)]
unsafe fn hsum(v: __m256i) -> u32 {
    let s = _mm_add_epi32(_mm256_castsi256_si128(v), _mm256_extracti128_si256::<1>(v));
    let s = _mm_add_epi32(s, _mm_shuffle_epi32::<0b01_00_11_10>(s));
    let s = _mm_add_epi32(s, _mm_shuffle_epi32::<0b10_11_00_01>(s));
    _mm_cvtsi128_si32(s) as u32
}

#[inline(always)]
unsafe fn d(p: *const u8) -> i32 {
    p.cast::<i32>().read_unaligned()
}

#[inline(always)]
unsafe fn rows2x8b(p: *const u8, s: usize) -> __m128i {
    _mm_castpd_si128(_mm_loadh_pd(_mm_castsi128_pd(_mm_loadl_epi64(p.cast())), p.add(s).cast()))
}

#[inline(always)]
unsafe fn rows4x8b(p: *const u8, s: usize) -> __m256i {
    _mm256_inserti128_si256::<1>(_mm256_castsi128_si256(rows2x8b(p, s)), rows2x8b(p.add(2 * s), s))
}

#[inline(always)]
unsafe fn rows4x4b(p: *const u8, s: usize) -> __m128i {
    _mm_setr_epi32(d(p), d(p.add(s)), d(p.add(2 * s)), d(p.add(3 * s)))
}

#[inline(always)]
unsafe fn rows2(p: *const u8, s: usize) -> __m256i {
    _mm256_loadu2_m128i(p.add(s).cast(), p.cast())
}

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

static ONE: usize = 1;

#[inline(always)]
unsafe fn trips(n: usize) -> usize {
    core::ptr::addr_of!(ONE).read_volatile() * n
}

#[inline(always)]
unsafe fn d16(x: __m256i, y: __m256i) -> __m256i {
    _mm256_madd_epi16(_mm256_abs_epi16(_mm256_sub_epi16(x, y)), _mm256_set1_epi16(1))
}

#[target_feature(enable = "avx2")]
pub unsafe fn sad_u16<const W: usize, const H: usize>(a: *const u16, sa: usize, b: *const u16, sb: usize) -> u32 {
    let (mut a, mut b, sa, sb) = (a.cast::<u8>(), b.cast::<u8>(), 2 * sa, 2 * sb);
    let mut acc = _mm256_setzero_si256();
    if W == 4 && H <= 8 {
        for r in (0..H).step_by(4) {
            acc = _mm256_add_epi32(acc, d16(rows4x8b(a.add(r * sa), sa), rows4x8b(b.add(r * sb), sb)));
        }
    } else if W == 4 {
        let mut x = _mm_setzero_si128();
        for _ in 0..trips(H / 2) {
            let d = _mm_abs_epi16(_mm_sub_epi16(rows2x8b(a, sa), rows2x8b(b, sb)));
            x = _mm_add_epi32(x, _mm_madd_epi16(d, _mm_set1_epi16(1)));
            (a, b) = (a.add(2 * sa), b.add(2 * sb));
        }
        acc = _mm256_zextsi128_si256(x);
    } else if W == 8 {
        for _ in 0..H / 2 {
            acc = _mm256_add_epi32(acc, d16(rows2(a, sa), rows2(b, sb)));
            (a, b) = (a.add(2 * sa), b.add(2 * sb));
        }
    } else {
        for _ in 0..if W == 64 && H < 16 { trips(H) } else { H } {
            for c in (0..2 * W).step_by(32) {
                acc = _mm256_add_epi32(acc, d16(ld(a.add(c)), ld(b.add(c))));
            }
            (a, b) = (a.add(sa), b.add(sb));
        }
    }
    hsum(acc)
}

#[target_feature(enable = "avx2")]
pub unsafe fn sad_u8<const W: usize, const H: usize>(mut a: *const u8, sa: usize, mut b: *const u8, sb: usize) -> u32 {
    let mut acc = _mm256_setzero_si256();
    if W == 4 {
        for _ in 0..H / 4 {
            acc = _mm256_add_epi64(acc, _mm256_zextsi128_si256(_mm_sad_epu8(rows4x4b(a, sa), rows4x4b(b, sb))));
            (a, b) = (a.add(4 * sa), b.add(4 * sb));
        }
    } else if W == 8 && H <= 8 {
        let mut x = _mm_setzero_si128();
        for r in (0..H).step_by(2) {
            x = _mm_add_epi64(x, _mm_sad_epu8(rows2x8b(a.add(r * sa), sa), rows2x8b(b.add(r * sb), sb)));
        }
        acc = _mm256_zextsi128_si256(x);
    } else if W == 8 {
        let mut x = _mm_setzero_si128();
        for _ in 0..trips(H / 2) {
            x = _mm_add_epi64(x, _mm_sad_epu8(rows2x8b(a, sa), rows2x8b(b, sb)));
            (a, b) = (a.add(2 * sa), b.add(2 * sb));
        }
        acc = _mm256_zextsi128_si256(x);
    } else if W == 16 {
        for _ in 0..H / 2 {
            acc = _mm256_add_epi64(acc, _mm256_sad_epu8(rows2(a, sa), rows2(b, sb)));
            (a, b) = (a.add(2 * sa), b.add(2 * sb));
        }
    } else {
        for _ in 0..H {
            for c in (0..W).step_by(32) {
                acc = _mm256_add_epi64(acc, _mm256_sad_epu8(ld(a.add(c)), ld(b.add(c))));
            }
            (a, b) = (a.add(sa), b.add(sb));
        }
    }
    hsum(acc)
}

#[target_feature(enable = "avx2")]
pub unsafe fn pair_u16(a: *const u16, sa: usize, b: *const u16, sb: usize) -> u64 {
    let (a, b, sa, sb) = (a.cast::<u8>(), b.cast::<u8>(), 2 * sa, 2 * sb);
    let mut acc = _mm256_setzero_si256();
    for r in 0..8 {
        acc = _mm256_add_epi32(acc, d16(ld(a.add(r * sa)), ld(b.add(r * sb))));
    }
    let s = _mm_hadd_epi32(_mm256_castsi256_si128(acc), _mm256_extracti128_si256::<1>(acc));
    _mm_cvtsi128_si64(_mm_hadd_epi32(s, s)) as u64
}

#[target_feature(enable = "avx2")]
pub unsafe fn pair_u8(a: *const u8, sa: usize, b: *const u8, sb: usize) -> u64 {
    let mut acc = _mm256_setzero_si256();
    for r in (0..8).step_by(2) {
        acc = _mm256_add_epi64(acc, _mm256_sad_epu8(rows2(a.add(r * sa), sa), rows2(b.add(r * sb), sb)));
    }
    let s = _mm_add_epi64(_mm256_castsi256_si128(acc), _mm256_extracti128_si256::<1>(acc));
    _mm_cvtsi128_si64(_mm_shuffle_epi32::<0b10_00_10_00>(s)) as u64
}

pub(super) static SAD_U16: [SadFn<u16>; 25] = sizes!(sad_u16);
pub(super) static SAD_U8: [SadFn<u8>; 25] = sizes!(sad_u8);