slain-num 0.4.0

pub type SadFn<P> = unsafe fn(*const P, usize, *const P, usize) -> u32;
pub type PairFn<P> = unsafe fn(*const P, usize, *const P, usize) -> u64;

mod sealed {
    pub trait Sealed {}
    impl Sealed for u8 {}
    impl Sealed for u16 {}
}

pub trait Px: Copy + Into<u32> + sealed::Sealed + 'static {
    #[doc(hidden)]
    fn sads() -> &'static [SadFn<Self>; 25];
    #[doc(hidden)]
    fn pair() -> PairFn<Self>;
}

macro_rules! sizes {
    ($($f:ident)::+) => {
        [
            $($f)::+::<4, 4>, $($f)::+::<4, 8>, $($f)::+::<4, 16>, $($f)::+::<4, 32>, $($f)::+::<4, 64>,
            $($f)::+::<8, 4>, $($f)::+::<8, 8>, $($f)::+::<8, 16>, $($f)::+::<8, 32>, $($f)::+::<8, 64>,
            $($f)::+::<16, 4>, $($f)::+::<16, 8>, $($f)::+::<16, 16>, $($f)::+::<16, 32>, $($f)::+::<16, 64>,
            $($f)::+::<32, 4>, $($f)::+::<32, 8>, $($f)::+::<32, 16>, $($f)::+::<32, 32>, $($f)::+::<32, 64>,
            $($f)::+::<64, 4>, $($f)::+::<64, 8>, $($f)::+::<64, 16>, $($f)::+::<64, 32>, $($f)::+::<64, 64>,
        ]
    };
}

#[cfg(target_arch = "x86_64")]
#[path = "block_x86.rs"]
pub mod x86;

#[cfg(not(any(target_arch = "aarch64", target_arch = "wasm32")))]
mod scalar {
    use super::SadFn;

    pub(super) unsafe fn sad<P: Copy + Into<u32>, const W: usize, const H: usize>(mut a: *const P, sa: usize, mut b: *const P, sb: usize) -> u32 {
        let mut s = 0u32;
        for _ in 0..H {
            for c in 0..W {
                s += (*a.add(c)).into().abs_diff((*b.add(c)).into());
            }
            (a, b) = (a.add(sa), b.add(sb));
        }
        s
    }

    pub(super) unsafe fn pair<P: Copy + Into<u32>>(a: *const P, sa: usize, b: *const P, sb: usize) -> u64 {
        u64::from(sad::<P, 8, 8>(a, sa, b, sb)) | u64::from(sad::<P, 8, 8>(a.add(8), sa, b.add(8), sb)) << 32
    }

    #[cfg(not(target_arch = "x86_64"))]
    pub(super) unsafe fn pair_u16(a: *const u16, sa: usize, b: *const u16, sb: usize) -> u64 {
        pair::<u16>(a, sa, b, sb)
    }

    #[cfg(not(target_arch = "x86_64"))]
    pub(super) unsafe fn pair_u8(a: *const u8, sa: usize, b: *const u8, sb: usize) -> u64 {
        pair::<u8>(a, sa, b, sb)
    }

    unsafe fn sad16<const W: usize, const H: usize>(a: *const u16, sa: usize, b: *const u16, sb: usize) -> u32 {
        sad::<u16, W, H>(a, sa, b, sb)
    }

    unsafe fn sad8<const W: usize, const H: usize>(a: *const u8, sa: usize, b: *const u8, sb: usize) -> u32 {
        sad::<u8, W, H>(a, sa, b, sb)
    }

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

#[cfg(target_arch = "aarch64")]
#[path = "block_neon.rs"]
mod arch;

#[cfg(target_arch = "wasm32")]
#[path = "block_wasm.rs"]
mod arch;

#[cfg(target_arch = "x86_64")]
use x86 as arch;

#[cfg(not(any(target_arch = "x86_64", target_arch = "aarch64", target_arch = "wasm32")))]
use scalar as arch;

impl Px for u16 {
    #[inline]
    fn sads() -> &'static [SadFn<u16>; 25] {
        #[cfg(target_arch = "x86_64")]
        if !slain_rt::cpu::avx2() {
            return &scalar::SAD_U16;
        }
        &arch::SAD_U16
    }

    #[inline]
    fn pair() -> PairFn<u16> {
        #[cfg(target_arch = "x86_64")]
        if !slain_rt::cpu::avx2() {
            return scalar::pair::<u16>;
        }
        arch::pair_u16
    }
}

impl Px for u8 {
    #[inline]
    fn sads() -> &'static [SadFn<u8>; 25] {
        #[cfg(target_arch = "x86_64")]
        if !slain_rt::cpu::avx2() {
            return &scalar::SAD_U8;
        }
        &arch::SAD_U8
    }

    #[inline]
    fn pair() -> PairFn<u8> {
        #[cfg(target_arch = "x86_64")]
        if !slain_rt::cpu::avx2() {
            return scalar::pair::<u8>;
        }
        arch::pair_u8
    }
}

#[inline]
fn side(v: usize) -> Option<usize> {
    (v.is_power_of_two() && (4..=64).contains(&v)).then(|| v.trailing_zeros() as usize - 2)
}

#[derive(Clone, Copy)]
pub struct Sad<P: Px> {
    f: SadFn<P>,
    w: usize,
    h: usize,
}

impl<P: Px> Sad<P> {
    #[inline]
    #[must_use]
    pub fn new(w: usize, h: usize) -> Option<Self> {
        let i = side(w)? * 5 + side(h)?;
        Some(Self { f: P::sads()[i], w, h })
    }

    #[inline]
    #[must_use]
    pub fn run(&self, a: &[P], sa: usize, b: &[P], sb: usize) -> u32 {
        assert!(super::fits(a.len(), self.w, self.h, sa) && super::fits(b.len(), self.w, self.h, sb));
        unsafe { (self.f)(a.as_ptr(), sa, b.as_ptr(), sb) }
    }

    #[inline]
    #[must_use]
    pub fn leaf(&self) -> SadFn<P> {
        self.f
    }
}

pub struct Cells {
    pub out: [u32; 64],
}

impl Default for Cells {
    fn default() -> Self {
        Self { out: [0; 64] }
    }
}

impl Cells {
    #[inline]
    pub fn run<P: Px>(&mut self, a: &[P], sa: usize, b: &[P], sb: usize, w: usize, h: usize) -> &[u32] {
        assert!(matches!(w, 8 | 16 | 32 | 64) && h.is_multiple_of(8) && (8..=64).contains(&h) && super::fits(a.len(), w, h, sa) && super::fits(b.len(), w, h, sb));
        let (cw, o) = (w / 8, self.out.as_mut_ptr());
        unsafe {
            if w == 8 {
                let f = P::sads()[6];
                for cy in 0..h / 8 {
                    *o.add(cy) = f(a.as_ptr().add(cy * 8 * sa), sa, b.as_ptr().add(cy * 8 * sb), sb);
                }
            } else {
                let f = P::pair();
                for cy in 0..h / 8 {
                    let (ra, rb) = (a.as_ptr().add(cy * 8 * sa), b.as_ptr().add(cy * 8 * sb));
                    for cx in (0..cw).step_by(2) {
                        let v = f(ra.add(cx * 8), sa, rb.add(cx * 8), sb);
                        *o.add(cy * cw + cx) = v as u32;
                        *o.add(cy * cw + cx + 1) = (v >> 32) as u32;
                    }
                }
            }
        }
        &self.out[..cw * (h / 8)]
    }
}

#[cfg(target_arch = "x86_64")]
pub mod probes {
    #[inline(never)]
    #[target_feature(enable = "avx2")]
    pub unsafe fn sad8x8_u16(a: &[u16; 64], b: &[u16; 64]) -> u32 {
        super::x86::sad_u16::<8, 8>(a.as_ptr(), 8, b.as_ptr(), 8)
    }

    #[inline(never)]
    #[target_feature(enable = "avx2")]
    pub unsafe fn sad8x8_u8(a: &[u8; 64], b: &[u8; 64]) -> u32 {
        super::x86::sad_u8::<8, 8>(a.as_ptr(), 8, b.as_ptr(), 8)
    }

    #[inline(never)]
    #[target_feature(enable = "avx2")]
    pub unsafe fn pair16x8_u8(a: &[u8; 128], b: &[u8; 128]) -> u64 {
        super::x86::pair_u8(a.as_ptr(), 16, b.as_ptr(), 16)
    }

    #[inline(never)]
    #[target_feature(enable = "avx2")]
    pub unsafe fn pair16x8_u16(a: &[u16; 128], b: &[u16; 128]) -> u64 {
        super::x86::pair_u16(a.as_ptr(), 16, b.as_ptr(), 16)
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn next(seed: &mut u64) -> u32 {
        *seed ^= *seed << 13;
        *seed ^= *seed >> 7;
        *seed ^= *seed << 17;
        (*seed >> 32) as u32
    }

    fn want<P: Px>(a: &[P], sa: usize, b: &[P], sb: usize, w: usize, h: usize) -> u32 {
        (0..h).flat_map(|r| (0..w).map(move |c| (r, c))).map(|(r, c)| a[r * sa + c].into().abs_diff(b[r * sb + c].into())).sum()
    }

    #[test]
    fn sad_every_size_and_stride_matches_scalar() {
        let mut seed = 0x2545_f491_9e37_79b9u64;
        for w in [4usize, 8, 16, 32, 64] {
            for h in [4usize, 8, 16, 32, 64] {
                for (pa, pb) in [(0, 0), (11, 5), (1, 64)] {
                    let (sa, sb) = (w + pa, w + pb);
                    for max in [255u32, 1023, 4095, 32767] {
                        let a: Vec<u16> = (0..sa * h).map(|_| (next(&mut seed) % (max + 1)) as u16).collect();
                        let b: Vec<u16> = (0..sb * h).map(|_| (next(&mut seed) % (max + 1)) as u16).collect();
                        let k = Sad::<u16>::new(w, h).unwrap();
                        assert_eq!(k.run(&a, sa, &b, sb), want(&a, sa, &b, sb, w, h), "u16 {w}x{h} max {max}");
                        if max == 255 {
                            let (a8, b8): (Vec<u8>, Vec<u8>) = (a.iter().map(|&v| v as u8).collect(), b.iter().map(|&v| v as u8).collect());
                            assert_eq!(Sad::<u8>::new(w, h).unwrap().run(&a8, sa, &b8, sb), want(&a8, sa, &b8, sb, w, h), "u8 {w}x{h}");
                        }
                    }
                }
            }
        }
        let (z, f) = (vec![0u8; 64 * 64], vec![255u8; 64 * 64]);
        assert_eq!(Sad::<u8>::new(64, 64).unwrap().run(&z, 64, &f, 64), 64 * 64 * 255);
        let (z, f) = (vec![0u16; 64 * 64], vec![32767u16; 64 * 64]);
        assert_eq!(Sad::<u16>::new(64, 64).unwrap().run(&z, 64, &f, 64), 64 * 64 * 32767);
        assert!(Sad::<u16>::new(24, 8).is_none() && Sad::<u8>::new(8, 2).is_none() && Sad::<u8>::new(128, 8).is_none());
    }

    #[test]
    fn cells_every_width_and_band_count_match_scalar() {
        let mut seed = 0x0bad_cafe_1234_5678u64;
        let mut cells = Cells::default();
        for w in [8usize, 16, 32, 64] {
            for h in [8usize, 16, 24, 32, 64] {
                for max in [255u32, 1023, 4095, 32767] {
                    let (sa, sb) = (w + 24, w);
                    let a: Vec<u16> = (0..sa * h).map(|_| (next(&mut seed) % (max + 1)) as u16).collect();
                    let b: Vec<u16> = (0..sb * h).map(|_| (next(&mut seed) % (max + 1)) as u16).collect();
                    let cell = |cx: usize, cy: usize| want(&a[cy * 8 * sa + cx * 8..], sa, &b[cy * 8 * sb + cx * 8..], sb, 8, 8);
                    let expect: Vec<u32> = (0..h / 8).flat_map(|cy| (0..w / 8).map(move |cx| (cx, cy))).map(|(cx, cy)| cell(cx, cy)).collect();
                    assert_eq!(cells.run(&a, sa, &b, sb, w, h), expect, "u16 {w}x{h} max {max}");
                    if max == 255 {
                        let (a8, b8): (Vec<u8>, Vec<u8>) = (a.iter().map(|&v| v as u8).collect(), b.iter().map(|&v| v as u8).collect());
                        assert_eq!(cells.run(&a8, sa, &b8, sb, w, h), expect, "u8 {w}x{h}");
                    }
                }
            }
        }
    }

    #[cfg(target_arch = "x86_64")]
    fn floor_run(symbol: &str, px: usize, mut call: impl FnMut()) {
        if !slain_rt::cpu::avx2() {
            return;
        }
        let calls = 20_000u32;
        let mut best = 0f64;
        for _ in 0..15 {
            let at = std::time::Instant::now();
            for _ in 0..calls {
                call();
            }
            best = best.max(f64::from(calls) / at.elapsed().as_secs_f64());
        }
        println!("FLOOR_SYMBOL slain_num::nd::block::probes::{symbol}");
        println!("FLOOR_PX {px}");
        println!("FLOOR_CALLS_PER_S {best:.9}");
        println!("FLOOR_CORE_HZ {:.9}", slain_rt::cpu::core_hz());
    }

    #[cfg(target_arch = "x86_64")]
    fn floor_fill<const N: usize>(seed: u32) -> [u16; N] {
        let mut s = seed;
        core::array::from_fn(|_| {
            s ^= s << 13;
            s ^= s >> 17;
            s ^= s << 5;
            (s & 0xff) as u16
        })
    }

    #[cfg(target_arch = "x86_64")]
    #[test]
    #[ignore = "floor: just floor sad8x8_u16 slain-num"]
    fn floor_sad8x8_u16() {
        use core::hint::black_box as slain;
        let (a, b) = (floor_fill::<64>(0x1357_9bdf), floor_fill::<64>(0x2468_ace1));
        floor_run("sad8x8_u16", 64, || {
            slain(unsafe { probes::sad8x8_u16(slain(&a), slain(&b)) });
        });
    }

    #[cfg(target_arch = "x86_64")]
    #[test]
    #[ignore = "floor: just floor sad8x8_u8 slain-num"]
    fn floor_sad8x8_u8() {
        use core::hint::black_box as slain;
        let (a, b) = (floor_fill::<64>(0x1357_9bdf).map(|v| v as u8), floor_fill::<64>(0x2468_ace1).map(|v| v as u8));
        floor_run("sad8x8_u8", 64, || {
            slain(unsafe { probes::sad8x8_u8(slain(&a), slain(&b)) });
        });
    }

    #[cfg(target_arch = "x86_64")]
    #[test]
    #[ignore = "floor: just floor pair16x8_u8 slain-num"]
    fn floor_pair16x8_u8() {
        use core::hint::black_box as slain;
        let (a, b) = (floor_fill::<128>(0x1357_9bdf).map(|v| v as u8), floor_fill::<128>(0x2468_ace1).map(|v| v as u8));
        floor_run("pair16x8_u8", 128, || {
            slain(unsafe { probes::pair16x8_u8(slain(&a), slain(&b)) });
        });
    }

    #[cfg(target_arch = "x86_64")]
    #[test]
    #[ignore = "floor: just floor pair16x8_u16 slain-num"]
    fn floor_pair16x8_u16() {
        use core::hint::black_box as slain;
        let (a, b) = (floor_fill::<128>(0x1357_9bdf), floor_fill::<128>(0x2468_ace1));
        floor_run("pair16x8_u16", 128, || {
            slain(unsafe { probes::pair16x8_u16(slain(&a), slain(&b)) });
        });
    }
}