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

use core::arch::wasm32::*;

use super::SadFn;

#[inline(always)]
unsafe fn d16(x: v128, y: v128) -> v128 {
    i32x4_dot_i16x8(i16x8_abs(i16x8_sub(x, y)), i16x8_splat(1))
}

#[inline(always)]
unsafe fn d8(x: v128, y: v128) -> v128 {
    u16x8_extadd_pairwise_u8x16(v128_or(u8x16_sub_sat(x, y), u8x16_sub_sat(y, x)))
}

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

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

#[inline(always)]
unsafe fn sum32(v: v128) -> u32 {
    u32x4_extract_lane::<0>(v) + u32x4_extract_lane::<1>(v) + u32x4_extract_lane::<2>(v) + u32x4_extract_lane::<3>(v)
}

pub unsafe fn sad_u16<const W: usize, const H: usize>(a: *const u16, sa: usize, b: *const u16, sb: usize) -> u32 {
    let (a, b, sa, sb) = (a.cast::<u8>(), b.cast::<u8>(), 2 * sa, 2 * sb);
    let mut acc = i32x4_splat(0);
    if W == 4 {
        for r in (0..H).step_by(2) {
            let (x, y) = (u64x2(q(a.add(r * sa)), q(a.add((r + 1) * sa))), u64x2(q(b.add(r * sb)), q(b.add((r + 1) * sb))));
            acc = i32x4_add(acc, d16(x, y));
        }
    } else {
        for r in 0..H {
            for c in (0..W).step_by(8) {
                acc = i32x4_add(acc, d16(v128_load(a.add(r * sa + 2 * c).cast()), v128_load(b.add(r * sb + 2 * c).cast())));
            }
        }
    }
    sum32(acc)
}

pub unsafe fn sad_u8<const W: usize, const H: usize>(a: *const u8, sa: usize, b: *const u8, sb: usize) -> u32 {
    let mut acc = i32x4_splat(0);
    for r0 in (0..H).step_by(16) {
        let mut s = i16x8_splat(0);
        if W == 4 {
            for r in (r0..r0 + 16).take_while(|&r| r < H).step_by(4) {
                let row = |p: *const u8, s: usize| u32x4(d(p), d(p.add(s)), d(p.add(2 * s)), d(p.add(3 * s)));
                s = i16x8_add(s, d8(row(a.add(r * sa), sa), row(b.add(r * sb), sb)));
            }
        } else if W == 8 {
            for r in (r0..r0 + 16).take_while(|&r| r < H).step_by(2) {
                let (x, y) = (u64x2(q(a.add(r * sa)), q(a.add((r + 1) * sa))), u64x2(q(b.add(r * sb)), q(b.add((r + 1) * sb))));
                s = i16x8_add(s, d8(x, y));
            }
        } else {
            for r in (r0..r0 + 16).take_while(|&r| r < H) {
                for c in (0..W).step_by(16) {
                    s = i16x8_add(s, d8(v128_load(a.add(r * sa + c).cast()), v128_load(b.add(r * sb + c).cast())));
                }
            }
        }
        acc = i32x4_add(acc, u32x4_extadd_pairwise_u16x8(s));
    }
    sum32(acc)
}

pub unsafe fn pair_u16(a: *const u16, sa: usize, b: *const u16, sb: usize) -> u64 {
    let (mut a, mut b, sa, sb) = (a.cast::<u8>(), b.cast::<u8>(), 2 * sa, 2 * sb);
    let (mut l, mut r) = (i32x4_splat(0), i32x4_splat(0));
    for _ in 0..8 {
        l = i32x4_add(l, d16(v128_load(a.cast()), v128_load(b.cast())));
        r = i32x4_add(r, d16(v128_load(a.add(16).cast()), v128_load(b.add(16).cast())));
        (a, b) = (a.add(sa), b.add(sb));
    }
    u64::from(sum32(l)) | u64::from(sum32(r)) << 32
}

pub unsafe fn pair_u8(mut a: *const u8, sa: usize, mut b: *const u8, sb: usize) -> u64 {
    let mut s = i16x8_splat(0);
    for _ in 0..8 {
        s = i16x8_add(s, d8(v128_load(a.cast()), v128_load(b.cast())));
        (a, b) = (a.add(sa), b.add(sb));
    }
    let v = u32x4_extadd_pairwise_u16x8(s);
    u64::from(u32x4_extract_lane::<0>(v) + u32x4_extract_lane::<1>(v)) | u64::from(u32x4_extract_lane::<2>(v) + u32x4_extract_lane::<3>(v)) << 32
}

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