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

use core::arch::aarch64::*;

use super::SadFn;

#[inline(always)]
unsafe fn w2(p: *const u8, s: usize) -> uint8x8_t {
    vcreate_u8(u64::from(p.cast::<u32>().read_unaligned()) | u64::from(p.add(s).cast::<u32>().read_unaligned()) << 32)
}

#[inline(always)]
unsafe fn h2(p: *const u16, s: usize) -> uint16x8_t {
    vcombine_u16(vld1_u16(p), vld1_u16(p.add(s)))
}

pub unsafe fn sad_u16<const W: usize, const H: usize>(mut a: *const u16, sa: usize, mut b: *const u16, sb: usize) -> u32 {
    let (mut s0, mut s1) = (vdupq_n_u32(0), vdupq_n_u32(0));
    if W == 4 {
        for _ in 0..H / 4 {
            s0 = vpadalq_u16(s0, vabdq_u16(h2(a, sa), h2(b, sb)));
            s1 = vpadalq_u16(s1, vabdq_u16(h2(a.add(2 * sa), sa), h2(b.add(2 * sb), sb)));
            (a, b) = (a.add(4 * sa), b.add(4 * sb));
        }
    } else {
        for _ in 0..H / 2 {
            for c in (0..W).step_by(8) {
                s0 = vpadalq_u16(s0, vabdq_u16(vld1q_u16(a.add(c)), vld1q_u16(b.add(c))));
                s1 = vpadalq_u16(s1, vabdq_u16(vld1q_u16(a.add(sa + c)), vld1q_u16(b.add(sb + c))));
            }
            (a, b) = (a.add(2 * sa), b.add(2 * sb));
        }
    }
    vaddvq_u32(vaddq_u32(s0, s1))
}

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 s0, mut s1) = (vdupq_n_u16(0), vdupq_n_u16(0));
    if W == 4 {
        for _ in 0..H / 4 {
            s0 = vabal_u8(s0, w2(a, sa), w2(b, sb));
            s1 = vabal_u8(s1, w2(a.add(2 * sa), sa), w2(b.add(2 * sb), sb));
            (a, b) = (a.add(4 * sa), b.add(4 * sb));
        }
    } else if W == 8 {
        for _ in 0..H / 2 {
            s0 = vabal_u8(s0, vld1_u8(a), vld1_u8(b));
            s1 = vabal_u8(s1, vld1_u8(a.add(sa)), vld1_u8(b.add(sb)));
            (a, b) = (a.add(2 * sa), b.add(2 * sb));
        }
    } else {
        for _ in 0..H {
            for c in (0..W).step_by(16) {
                let (x, y) = (vld1q_u8(a.add(c)), vld1q_u8(b.add(c)));
                s0 = vabal_u8(s0, vget_low_u8(x), vget_low_u8(y));
                s1 = vabal_high_u8(s1, x, y);
            }
            (a, b) = (a.add(sa), b.add(sb));
        }
    }
    vaddlvq_u16(s0) + vaddlvq_u16(s1)
}

pub unsafe fn pair_u16(mut a: *const u16, sa: usize, mut b: *const u16, sb: usize) -> u64 {
    let (mut l, mut r) = (vdupq_n_u32(0), vdupq_n_u32(0));
    for _ in 0..8 {
        l = vpadalq_u16(l, vabdq_u16(vld1q_u16(a), vld1q_u16(b)));
        r = vpadalq_u16(r, vabdq_u16(vld1q_u16(a.add(8)), vld1q_u16(b.add(8))));
        (a, b) = (a.add(sa), b.add(sb));
    }
    u64::from(vaddvq_u32(l)) | u64::from(vaddvq_u32(r)) << 32
}

pub unsafe fn pair_u8(mut a: *const u8, sa: usize, mut b: *const u8, sb: usize) -> u64 {
    let (mut l, mut r) = (vdupq_n_u16(0), vdupq_n_u16(0));
    for _ in 0..8 {
        let (x, y) = (vld1q_u8(a), vld1q_u8(b));
        l = vabal_u8(l, vget_low_u8(x), vget_low_u8(y));
        r = vabal_high_u8(r, x, y);
        (a, b) = (a.add(sa), b.add(sb));
    }
    u64::from(vaddlvq_u16(l)) | u64::from(vaddlvq_u16(r)) << 32
}

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