pub mod block;
#[derive(Clone, Copy, Debug)]
pub struct View<'a, T> {
data: &'a [T],
w: usize,
h: usize,
stride: usize,
}
#[derive(Debug)]
pub struct ViewMut<'a, T> {
data: &'a mut [T],
w: usize,
h: usize,
stride: usize,
}
#[inline]
fn fits(len: usize, w: usize, h: usize, stride: usize) -> bool {
stride >= w && (h == 0 || w == 0 || (h - 1).checked_mul(stride).and_then(|n| n.checked_add(w)).is_some_and(|n| n <= len))
}
impl<'a, T: Copy> View<'a, T> {
#[must_use]
pub fn new(data: &'a [T], w: usize, h: usize, stride: usize) -> Option<Self> {
fits(data.len(), w, h, stride).then_some(Self { data, w, h, stride })
}
#[must_use]
pub fn w(&self) -> usize {
self.w
}
#[must_use]
pub fn h(&self) -> usize {
self.h
}
#[must_use]
pub fn stride(&self) -> usize {
self.stride
}
#[inline]
#[must_use]
pub fn row(&self, y: usize) -> &'a [T] {
assert!(y < self.h);
&self.data[y * self.stride..y * self.stride + self.w]
}
pub fn rows(&self) -> impl Iterator<Item = &'a [T]> + '_ {
(0..self.h).map(|y| self.row(y))
}
#[must_use]
pub fn get(&self, x: usize, y: usize) -> T {
self.row(y)[x]
}
#[must_use]
pub fn sub(&self, x: usize, y: usize, w: usize, h: usize) -> Option<Self> {
if x.checked_add(w)? > self.w || y.checked_add(h)? > self.h {
return None;
}
if w == 0 || h == 0 {
return Some(Self { data: &self.data[..0], w, h, stride: self.stride });
}
Some(Self { data: &self.data[y * self.stride + x..], w, h, stride: self.stride })
}
}
impl<'a, T: Copy> ViewMut<'a, T> {
#[must_use]
pub fn new(data: &'a mut [T], w: usize, h: usize, stride: usize) -> Option<Self> {
fits(data.len(), w, h, stride).then_some(Self { data, w, h, stride })
}
#[must_use]
pub fn w(&self) -> usize {
self.w
}
#[must_use]
pub fn h(&self) -> usize {
self.h
}
#[must_use]
pub fn stride(&self) -> usize {
self.stride
}
#[must_use]
pub fn view(&self) -> View<'_, T> {
View { data: self.data, w: self.w, h: self.h, stride: self.stride }
}
#[inline]
pub fn row_mut(&mut self, y: usize) -> &mut [T] {
assert!(y < self.h);
&mut self.data[y * self.stride..y * self.stride + self.w]
}
pub fn sub_mut(&mut self, x: usize, y: usize, w: usize, h: usize) -> Option<ViewMut<'_, T>> {
if x.checked_add(w)? > self.w || y.checked_add(h)? > self.h {
return None;
}
if w == 0 || h == 0 {
return Some(ViewMut { data: &mut self.data[..0], w, h, stride: self.stride });
}
Some(ViewMut { data: &mut self.data[y * self.stride + x..], w, h, stride: self.stride })
}
}
#[cfg(target_arch = "x86_64")]
#[inline]
fn fast() -> bool {
slain_rt::cpu::avx2()
}
fn same<A: Copy, B: Copy>(a: &View<'_, A>, b: &View<'_, B>) {
assert!(a.w == b.w && a.h == b.h, "shape {}x{} vs {}x{}", a.w, a.h, b.w, b.h);
}
pub fn fill(d: &mut ViewMut<'_, u16>, v: u16) {
for y in 0..d.h {
d.row_mut(y).fill(v);
}
}
pub fn copy(d: &mut ViewMut<'_, u16>, s: &View<'_, u16>) {
same(&d.view(), s);
for y in 0..d.h {
d.row_mut(y).copy_from_slice(s.row(y));
}
}
pub fn clip(d: &mut ViewMut<'_, u16>, lo: u16, hi: u16) {
#[cfg(target_arch = "x86_64")]
let f = fast();
for y in 0..d.h {
let r = d.row_mut(y);
#[cfg(target_arch = "x86_64")]
if f {
unsafe { x86::clip(r, lo, hi) };
continue;
}
scalar::clip(r, lo, hi);
}
}
pub fn avg(d: &mut ViewMut<'_, u16>, a: &View<'_, u16>, b: &View<'_, u16>) {
same(&d.view(), a);
same(a, b);
#[cfg(target_arch = "x86_64")]
let f = fast();
for y in 0..d.h {
let r = d.row_mut(y);
#[cfg(target_arch = "x86_64")]
if f {
unsafe { x86::avg(r, a.row(y), b.row(y)) };
continue;
}
scalar::avg(r, a.row(y), b.row(y));
}
}
pub fn add_sat(d: &mut ViewMut<'_, u16>, a: &View<'_, u16>) {
same(&d.view(), a);
#[cfg(target_arch = "x86_64")]
let f = fast();
for y in 0..d.h {
let r = d.row_mut(y);
#[cfg(target_arch = "x86_64")]
if f {
unsafe { x86::add_sat(r, a.row(y)) };
continue;
}
scalar::add_sat(r, a.row(y));
}
}
pub fn add_residual(d: &mut ViewMut<'_, u16>, res: &View<'_, i16>, max: u16) {
assert!(max <= 0x7fff);
same(&d.view(), res);
#[cfg(target_arch = "x86_64")]
let f = fast();
for y in 0..d.h {
let r = d.row_mut(y);
#[cfg(target_arch = "x86_64")]
if f {
unsafe { x86::add_residual(r, res.row(y), max) };
continue;
}
scalar::add_residual(r, res.row(y), max);
}
}
#[must_use]
pub fn sum(a: &View<'_, u16>) -> u64 {
#[cfg(target_arch = "x86_64")]
let f = fast();
a.rows()
.map(|r| {
#[cfg(target_arch = "x86_64")]
if f {
return unsafe { x86::sum(r) };
}
scalar::sum(r)
})
.sum()
}
#[must_use]
pub fn sad(a: &View<'_, u16>, b: &View<'_, u16>) -> u64 {
same(a, b);
#[cfg(target_arch = "x86_64")]
let f = fast();
(0..a.h)
.map(|y| {
#[cfg(target_arch = "x86_64")]
if f {
return unsafe { x86::sad(a.row(y), b.row(y)) };
}
scalar::sad(a.row(y), b.row(y))
})
.sum()
}
#[must_use]
pub fn ssd(a: &View<'_, u16>, b: &View<'_, u16>) -> u64 {
same(a, b);
#[cfg(target_arch = "x86_64")]
let f = fast();
(0..a.h)
.map(|y| {
#[cfg(target_arch = "x86_64")]
if f {
return unsafe { x86::ssd(a.row(y), b.row(y)) };
}
scalar::ssd(a.row(y), b.row(y))
})
.sum()
}
#[must_use]
pub fn min_max(a: &View<'_, u16>) -> Option<(u16, u16)> {
if a.w == 0 || a.h == 0 {
return None;
}
#[cfg(target_arch = "x86_64")]
let f = fast();
Some(a.rows().fold((u16::MAX, 0), |(lo, hi), r| {
#[cfg(target_arch = "x86_64")]
if f {
let (l, h) = unsafe { x86::min_max(r) };
return (lo.min(l), hi.max(h));
}
let (l, h) = scalar::min_max(r);
(lo.min(l), hi.max(h))
}))
}
pub fn absdiff(d: &mut ViewMut<'_, u16>, a: &View<'_, u16>, b: &View<'_, u16>) {
same(&d.view(), a);
same(a, b);
#[cfg(target_arch = "x86_64")]
let f = fast();
for y in 0..d.h {
let r = d.row_mut(y);
#[cfg(target_arch = "x86_64")]
if f {
unsafe { x86::absdiff(r, a.row(y), b.row(y)) };
continue;
}
scalar::absdiff(r, a.row(y), b.row(y));
}
}
#[must_use]
pub fn count_range(a: &View<'_, u16>, lo: u16, hi: u16) -> u64 {
#[cfg(target_arch = "x86_64")]
let f = fast();
a.rows()
.map(|r| {
#[cfg(target_arch = "x86_64")]
if f {
return unsafe { x86::count_range(r, lo, hi) };
}
scalar::count_range(r, lo, hi)
})
.sum()
}
pub mod scalar {
pub fn absdiff(d: &mut [u16], a: &[u16], b: &[u16]) {
for ((o, &x), &y) in d.iter_mut().zip(a).zip(b) {
*o = x.abs_diff(y);
}
}
#[must_use]
pub fn count_range(a: &[u16], lo: u16, hi: u16) -> u64 {
a.iter().filter(|&&v| lo <= v && v <= hi).count() as u64
}
pub fn clip(d: &mut [u16], lo: u16, hi: u16) {
for v in d {
*v = (*v).max(lo).min(hi);
}
}
pub fn avg(d: &mut [u16], a: &[u16], b: &[u16]) {
for ((o, &x), &y) in d.iter_mut().zip(a).zip(b) {
*o = ((u32::from(x) + u32::from(y) + 1) >> 1) as u16;
}
}
pub fn add_sat(d: &mut [u16], a: &[u16]) {
for (o, &x) in d.iter_mut().zip(a) {
*o = o.saturating_add(x);
}
}
pub fn add_residual(d: &mut [u16], r: &[i16], max: u16) {
for (o, &x) in d.iter_mut().zip(r) {
*o = (i32::from(*o) + i32::from(x)).clamp(0, i32::from(max)) as u16;
}
}
#[must_use]
pub fn sum(a: &[u16]) -> u64 {
a.iter().map(|&v| u64::from(v)).sum()
}
#[must_use]
pub fn sad(a: &[u16], b: &[u16]) -> u64 {
a.iter().zip(b).map(|(&x, &y)| u64::from(x.abs_diff(y))).sum()
}
#[must_use]
pub fn ssd(a: &[u16], b: &[u16]) -> u64 {
a.iter().zip(b).map(|(&x, &y)| u64::from(x.abs_diff(y)).pow(2)).sum()
}
#[must_use]
pub fn min_max(a: &[u16]) -> (u16, u16) {
a.iter().fold((u16::MAX, 0), |(lo, hi), &v| (lo.min(v), hi.max(v)))
}
}
#[cfg(target_arch = "x86_64")]
pub mod x86 {
use core::arch::x86_64::*;
#[inline(always)]
unsafe fn ld(p: *const u16) -> __m256i {
_mm256_loadu_si256(p.cast())
}
#[inline(always)]
unsafe fn st(p: *mut u16, v: __m256i) {
_mm256_storeu_si256(p.cast(), v)
}
#[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 clip(d: &mut [u16], lo: u16, hi: u16) {
let (vl, vh) = (_mm256_set1_epi16(lo as i16), _mm256_set1_epi16(hi as i16));
let n = d.len();
let p = d.as_mut_ptr();
let mut i = 0;
while i + 16 <= n {
st(p.add(i), _mm256_min_epu16(_mm256_max_epu16(ld(p.add(i)), vl), vh));
i += 16;
}
super::scalar::clip(&mut d[i..], lo, hi);
}
#[target_feature(enable = "avx2")]
pub unsafe fn avg(d: &mut [u16], a: &[u16], b: &[u16]) {
let n = d.len().min(a.len()).min(b.len());
let (p, pa, pb) = (d.as_mut_ptr(), a.as_ptr(), b.as_ptr());
let mut i = 0;
while i + 16 <= n {
st(p.add(i), _mm256_avg_epu16(ld(pa.add(i)), ld(pb.add(i))));
i += 16;
}
super::scalar::avg(&mut d[i..n], &a[i..n], &b[i..n]);
}
#[target_feature(enable = "avx2")]
pub unsafe fn add_sat(d: &mut [u16], a: &[u16]) {
let n = d.len().min(a.len());
let (p, pa) = (d.as_mut_ptr(), a.as_ptr());
let mut i = 0;
while i + 16 <= n {
st(p.add(i), _mm256_adds_epu16(ld(p.add(i)), ld(pa.add(i))));
i += 16;
}
super::scalar::add_sat(&mut d[i..n], &a[i..n]);
}
#[target_feature(enable = "avx2")]
pub unsafe fn add_residual(d: &mut [u16], r: &[i16], max: u16) {
let n = d.len().min(r.len());
let (p, pr) = (d.as_mut_ptr(), r.as_ptr());
let (z, m) = (_mm256_setzero_si256(), _mm256_set1_epi16(max as i16));
let mut i = 0;
while i + 16 <= n {
let s = _mm256_adds_epi16(ld(p.add(i)), _mm256_loadu_si256(pr.add(i).cast()));
st(p.add(i), _mm256_min_epi16(_mm256_max_epi16(s, z), m));
i += 16;
}
super::scalar::add_residual(&mut d[i..n], &r[i..n], max);
}
#[inline(always)]
unsafe fn pairs(v: __m256i) -> __m256i {
_mm256_madd_epi16(_mm256_xor_si256(v, _mm256_set1_epi16(i16::MIN)), _mm256_set1_epi16(1))
}
#[inline(always)]
unsafe fn hsum_i32(v: __m256i) -> i64 {
let q = _mm256_add_epi64(_mm256_cvtepi32_epi64(_mm256_castsi256_si128(v)), _mm256_cvtepi32_epi64(_mm256_extracti128_si256::<1>(v)));
let mut l = [0i64; 4];
_mm256_storeu_si256(l.as_mut_ptr().cast(), q);
l[0] + l[1] + l[2] + l[3]
}
#[inline(always)]
unsafe fn sum16(n: usize, f: impl Fn(usize) -> __m256i) -> (u64, usize) {
let m = n - n % 16;
let z = _mm256_setzero_si256();
let mut s = 0i64;
let mut i = 0;
while i + 64 <= m {
let end = i + ((m - i) / 64).min(4096) * 64;
let (mut a0, mut a1, mut a2, mut a3) = (z, z, z, z);
while i < end {
a0 = _mm256_add_epi32(a0, pairs(f(i)));
a1 = _mm256_add_epi32(a1, pairs(f(i + 16)));
a2 = _mm256_add_epi32(a2, pairs(f(i + 32)));
a3 = _mm256_add_epi32(a3, pairs(f(i + 48)));
i += 64;
}
s += hsum_i32(_mm256_add_epi32(_mm256_add_epi32(a0, a1), _mm256_add_epi32(a2, a3)));
}
let mut a0 = z;
while i < m {
a0 = _mm256_add_epi32(a0, pairs(f(i)));
i += 16;
}
s += hsum_i32(a0);
((s + 32768 * i as i64) as u64, i)
}
#[target_feature(enable = "avx2")]
pub unsafe fn sum(a: &[u16]) -> u64 {
let p = a.as_ptr();
let (t, i) = sum16(a.len(), |i| ld(p.add(i)));
t + super::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 (t, i) = sum16(n, |i| {
let (x, y) = (ld(pa.add(i)), ld(pb.add(i)));
_mm256_sub_epi16(_mm256_max_epu16(x, y), _mm256_min_epu16(x, y))
});
t + super::scalar::sad(&a[i..n], &b[i..n])
}
#[target_feature(enable = "avx2")]
pub unsafe fn ssd(a: &[u16], b: &[u16]) -> u64 {
let n = a.len().min(b.len());
let (pa, pb) = (a.as_ptr(), b.as_ptr());
let z = _mm256_setzero_si256();
let (mut s0, mut s1) = (z, z);
let mut i = 0;
while i + 16 <= n {
let (x, y) = (ld(pa.add(i)), ld(pb.add(i)));
let d = _mm256_sub_epi16(_mm256_max_epu16(x, y), _mm256_min_epu16(x, y));
let (lo, hi) = (_mm256_mullo_epi16(d, d), _mm256_mulhi_epu16(d, d));
let (p0, p1) = (_mm256_unpacklo_epi16(lo, hi), _mm256_unpackhi_epi16(lo, hi));
s0 = _mm256_add_epi64(s0, _mm256_add_epi64(_mm256_unpacklo_epi32(p0, z), _mm256_unpackhi_epi32(p0, z)));
s1 = _mm256_add_epi64(s1, _mm256_add_epi64(_mm256_unpacklo_epi32(p1, z), _mm256_unpackhi_epi32(p1, z)));
i += 16;
}
let mut l = [0u64; 4];
_mm256_storeu_si256(l.as_mut_ptr().cast(), _mm256_add_epi64(s0, s1));
l[0] + l[1] + l[2] + l[3] + super::scalar::ssd(&a[i..n], &b[i..n])
}
#[target_feature(enable = "avx2")]
pub unsafe fn absdiff(d: &mut [u16], a: &[u16], b: &[u16]) {
let n = d.len().min(a.len()).min(b.len());
let (p, pa, pb) = (d.as_mut_ptr(), a.as_ptr(), b.as_ptr());
let mut i = 0;
while i + 16 <= n {
let (x, y) = (ld(pa.add(i)), ld(pb.add(i)));
st(p.add(i), _mm256_sub_epi16(_mm256_max_epu16(x, y), _mm256_min_epu16(x, y)));
i += 16;
}
super::scalar::absdiff(&mut d[i..n], &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 hit = |v: __m256i| _mm256_cmpeq_epi16(_mm256_min_epu16(_mm256_max_epu16(v, vl), vh), v);
let m = n - n % 16;
let mut total = 0u64;
let mut i = 0;
while i + 64 <= m {
let end = i + ((m - i) / 64).min(32768) * 64;
let (mut c0, mut c1, mut c2, mut c3) = (z, z, z, z);
while i < end {
c0 = _mm256_sub_epi16(c0, hit(ld(p.add(i))));
c1 = _mm256_sub_epi16(c1, hit(ld(p.add(i + 16))));
c2 = _mm256_sub_epi16(c2, hit(ld(p.add(i + 32))));
c3 = _mm256_sub_epi16(c3, hit(ld(p.add(i + 48))));
i += 64;
}
total += hsum32(widen_add(widen_add(widen_add(widen_add(z, c0), c1), c2), c3));
}
let mut c0 = z;
while i < m {
c0 = _mm256_sub_epi16(c0, hit(ld(p.add(i))));
i += 16;
}
total += hsum32(widen_add(z, c0));
total + super::scalar::count_range(&a[i..], lo, hi)
}
#[target_feature(enable = "avx2")]
pub unsafe fn min_max(a: &[u16]) -> (u16, u16) {
let n = a.len();
if n < 16 {
return super::scalar::min_max(a);
}
let p = a.as_ptr();
let (f, z) = (_mm256_set1_epi16(-1), _mm256_setzero_si256());
let (mut l0, mut l1, mut l2, mut l3) = (f, f, f, f);
let (mut h0, mut h1, mut h2, mut h3) = (z, z, z, z);
let mut i = 0;
while i + 64 <= n {
let (v0, v1, v2, v3) = (ld(p.add(i)), ld(p.add(i + 16)), ld(p.add(i + 32)), ld(p.add(i + 48)));
l0 = _mm256_min_epu16(l0, v0);
l1 = _mm256_min_epu16(l1, v1);
l2 = _mm256_min_epu16(l2, v2);
l3 = _mm256_min_epu16(l3, v3);
h0 = _mm256_max_epu16(h0, v0);
h1 = _mm256_max_epu16(h1, v1);
h2 = _mm256_max_epu16(h2, v2);
h3 = _mm256_max_epu16(h3, v3);
i += 64;
}
while i + 16 <= n {
let v = ld(p.add(i));
l0 = _mm256_min_epu16(l0, v);
h0 = _mm256_max_epu16(h0, v);
i += 16;
}
let v = ld(p.add(n - 16));
l1 = _mm256_min_epu16(l1, v);
h1 = _mm256_max_epu16(h1, v);
let lo = _mm256_min_epu16(_mm256_min_epu16(l0, l1), _mm256_min_epu16(l2, l3));
let hi = _mm256_max_epu16(_mm256_max_epu16(h0, h1), _mm256_max_epu16(h2, h3));
let lo = _mm_min_epu16(_mm256_castsi256_si128(lo), _mm256_extracti128_si256::<1>(lo));
let hi = _mm_max_epu16(_mm256_castsi256_si128(hi), _mm256_extracti128_si256::<1>(hi));
let all = _mm_set1_epi16(-1);
let vl = _mm_extract_epi16::<0>(_mm_minpos_epu16(lo)) as u16;
let vh = !(_mm_extract_epi16::<0>(_mm_minpos_epu16(_mm_xor_si128(hi, all))) as u16);
(vl, vh)
}
}
#[cfg(test)]
mod tests {
use super::*;
fn rng(seed: u64) -> impl FnMut() -> u64 {
let mut s = seed;
move || {
s ^= s << 13;
s ^= s >> 7;
s ^= s << 17;
s
}
}
fn plane(_w: usize, h: usize, stride: usize, max: u16, seed: u64) -> Vec<u16> {
let mut r = rng(seed);
(0..stride * h).map(|_| (r() % (u64::from(max) + 1)) as u16).collect()
}
const SHAPES: [(usize, usize, usize); 7] = [(1, 1, 1), (15, 3, 17), (16, 4, 16), (17, 5, 40), (64, 64, 80), (127, 9, 128), (1920, 4, 1984)];
#[test]
fn view_bounds() {
let d = vec![0u16; 100];
assert!(View::new(&d, 10, 10, 10).is_some());
assert!(View::new(&d, 10, 11, 10).is_none());
assert!(View::new(&d, 11, 1, 10).is_none());
assert!(View::new(&d, 4, 9, 11).is_some());
assert!(View::new(&d, 4, 10, 11).is_none());
let v = View::new(&d, 10, 10, 10).unwrap();
assert!(v.sub(5, 5, 5, 5).is_some());
assert!(v.sub(6, 5, 5, 5).is_none());
assert_eq!(v.sub(2, 3, 4, 5).unwrap().row(4).len(), 4);
}
#[cfg(target_arch = "x86_64")]
#[test]
fn x86_matches_scalar_every_shape_and_limit() {
if !slain_rt::cpu::avx2() {
return;
}
for (k, &(w, h, s)) in SHAPES.iter().enumerate() {
for max in [1u16, 255, 1023, 4095, 0x7fff, u16::MAX] {
let a = plane(w, h, s, max, 11 + k as u64);
let b = plane(w, h, s, max, 97 + k as u64);
let va = View::new(&a, w, h, s).unwrap();
let vb = View::new(&b, w, h, s).unwrap();
for y in 0..h {
let (ra, rb) = (va.row(y), vb.row(y));
unsafe {
assert_eq!(x86::sum(ra), scalar::sum(ra));
assert_eq!(x86::sad(ra, rb), scalar::sad(ra, rb));
assert_eq!(x86::min_max(ra), scalar::min_max(ra));
let (mut p, mut q) = (ra.to_vec(), ra.to_vec());
x86::clip(&mut p, max / 4, max / 2);
scalar::clip(&mut q, max / 4, max / 2);
assert_eq!(p, q);
let (mut p, mut q) = (vec![0; w], vec![0; w]);
x86::avg(&mut p, ra, rb);
scalar::avg(&mut q, ra, rb);
assert_eq!(p, q);
let (mut p, mut q) = (ra.to_vec(), ra.to_vec());
x86::add_sat(&mut p, rb);
scalar::add_sat(&mut q, rb);
assert_eq!(p, q);
let (mut p, mut q) = (vec![0; w], vec![0; w]);
x86::absdiff(&mut p, ra, rb);
scalar::absdiff(&mut q, ra, rb);
assert_eq!(p, q);
for (lo, hi) in [(0, 0), (1, u16::MAX), (1, 1), (max / 4, max / 2), (max, max), (5, 3), (0, u16::MAX)] {
assert_eq!(x86::count_range(ra, lo, hi), scalar::count_range(ra, lo, hi));
}
}
}
if max <= 0x7fff {
let mut r = rng(5 + k as u64);
let res: Vec<i16> = (0..w).map(|_| r() as i16).collect();
let lim: Vec<i16> = (0..w).map(|i| if i % 2 == 0 { i16::MAX } else { i16::MIN }).collect();
for rr in [&res, &lim] {
let (mut p, mut q) = (va.row(0).to_vec(), va.row(0).to_vec());
unsafe { x86::add_residual(&mut p, rr, max) };
scalar::add_residual(&mut q, rr, max);
assert_eq!(p, q);
}
}
}
}
}
#[cfg(target_arch = "x86_64")]
#[test]
fn count_range_lane_counters_do_not_wrap() {
if !slain_rt::cpu::avx2() {
return;
}
let a = vec![9u16; 16 * 32768 * 2 + 16 * 70000 + 5];
unsafe {
assert_eq!(x86::count_range(&a, 9, 9), a.len() as u64);
assert_eq!(x86::count_range(&a, 0, 8), 0);
}
assert_eq!(count_range(&View::new(&a, a.len(), 1, a.len()).unwrap(), 1, u16::MAX), a.len() as u64);
}
#[test]
fn ssd_matches_scalar_and_strided() {
let want = |a: &View<'_, u16>, b: &View<'_, u16>| -> u64 { a.rows().flatten().zip(b.rows().flatten()).map(|(&x, &y)| u64::from(x.abs_diff(y)).pow(2)).sum() };
for (k, &(w, h, s)) in SHAPES.iter().enumerate() {
for max in [1u16, 255, 1023, 0x7fff, u16::MAX] {
let a = plane(w, h, s, max, 21 + k as u64);
let b = plane(w, h, s, max, 87 + k as u64);
let (va, vb) = (View::new(&a, w, h, s).unwrap(), View::new(&b, w, h, s).unwrap());
assert_eq!(ssd(&va, &vb), want(&va, &vb));
#[cfg(target_arch = "x86_64")]
if slain_rt::cpu::avx2() {
for y in 0..h {
assert_eq!(unsafe { x86::ssd(va.row(y), vb.row(y)) }, scalar::ssd(va.row(y), vb.row(y)));
}
}
}
}
let (z, m) = (vec![0u16; 4099], vec![u16::MAX; 4099]);
assert_eq!(ssd(&View::new(&z, 4099, 1, 4099).unwrap(), &View::new(&m, 4099, 1, 4099).unwrap()), 4099 * 65535 * 65535);
let a = plane(67, 13, 96, 255, 3);
let b = plane(67, 13, 96, 255, 4);
let va = View::new(&a, 67, 13, 96).unwrap().sub(3, 2, 50, 9).unwrap();
let vb = View::new(&b, 67, 13, 96).unwrap().sub(3, 2, 50, 9).unwrap();
assert_eq!(ssd(&va, &vb), want(&va, &vb));
}
#[test]
fn ufuncs_on_strided_subviews() {
let (w, h, s) = (67, 13, 96);
let a = plane(w, h, s, 1023, 3);
let b = plane(w, h, s, 1023, 4);
let va = View::new(&a, w, h, s).unwrap().sub(3, 2, 50, 9).unwrap();
let vb = View::new(&b, w, h, s).unwrap().sub(3, 2, 50, 9).unwrap();
let want_sum: u64 = va.rows().flatten().map(|&v| u64::from(v)).sum();
let want_sad: u64 = va.rows().flatten().zip(vb.rows().flatten()).map(|(&x, &y)| u64::from(x.abs_diff(y))).sum();
assert_eq!(sum(&va), want_sum);
assert_eq!(sad(&va, &vb), want_sad);
let mut out = vec![7u16; s * h];
let mut vo = ViewMut::new(&mut out, w, h, s).unwrap();
let mut so = vo.sub_mut(3, 2, 50, 9).unwrap();
avg(&mut so, &va, &vb);
for y in 0..9 {
for x in 0..50 {
assert_eq!(out[(y + 2) * s + x + 3], ((u32::from(va.get(x, y)) + u32::from(vb.get(x, y)) + 1) >> 1) as u16);
}
assert_eq!(out[(y + 2) * s + 2], 7);
assert_eq!(out[(y + 2) * s + 53], 7);
}
let mut c = a.clone();
let mut vc = ViewMut::new(&mut c, w, h, s).unwrap();
clip(&mut vc, 100, 900);
assert!(View::new(&c, w, h, s).unwrap().rows().flatten().all(|&v| (100..=900).contains(&v)));
assert_eq!(min_max(&View::new(&c, w, h, s).unwrap()), Some((100, 900)));
}
#[cfg(target_arch = "x86_64")]
#[test]
#[ignore]
fn bench_ufuncs() {
use core::hint::black_box as slain;
if !slain_rt::cpu::avx2() {
return;
}
let (w, h) = (1920usize, 1080usize);
let a = plane(w, h, w, 1023, 1);
let b = plane(w, h, w, 1023, 2);
let mut d = vec![0u16; w * h];
let time = |f: &mut dyn FnMut()| {
let mut best = u128::MAX;
for _ in 0..15 {
let t = std::time::Instant::now();
f();
best = best.min(t.elapsed().as_nanos());
}
best
};
let rows: Vec<(u128, u128, &str)> = vec![
(
time(&mut || {
slain(scalar::sum(slain(&a)));
}),
time(&mut || {
slain(unsafe { x86::sum(slain(&a)) });
}),
"sum",
),
(
time(&mut || {
slain(scalar::sad(slain(&a), slain(&b)));
}),
time(&mut || {
slain(unsafe { x86::sad(slain(&a), slain(&b)) });
}),
"sad",
),
(
time(&mut || {
slain(scalar::min_max(slain(&a)));
}),
time(&mut || {
slain(unsafe { x86::min_max(slain(&a)) });
}),
"min_max",
),
(time(&mut || scalar::avg(slain(&mut d), slain(&a), slain(&b))), time(&mut || unsafe { x86::avg(slain(&mut d), slain(&a), slain(&b)) }), "avg"),
(time(&mut || scalar::clip(slain(&mut d), 64, 940)), time(&mut || unsafe { x86::clip(slain(&mut d), 64, 940) }), "clip"),
];
for (s, v, name) in rows {
println!("{name:8} 1920x1080 scalar {:>9} ns avx2 {:>9} ns {:.2}x", s, v, s as f64 / v as f64);
}
}
}
#[cfg(all(test, target_arch = "x86_64"))]
#[path = "nd_prev.rs"]
mod prev;
#[cfg(all(test, target_arch = "x86_64"))]
mod reduce_tests {
use super::*;
fn data(n: usize, max: u32, seed: u64) -> Vec<u16> {
let mut s = seed | 1;
(0..n)
.map(|_| {
s ^= s << 13;
s ^= s >> 7;
s ^= s << 17;
(s % (u64::from(max) + 1)) as u16
})
.collect()
}
#[test]
fn register_mapped_reductions_match_prev_and_scalar() {
if !slain_rt::cpu::avx2() {
return;
}
for n in [0usize, 1, 15, 16, 17, 63, 64, 65, 127, 4096, 4096 * 64 + 37, 32768 * 64 + 5, 2_000_003] {
for (k, max) in [1u32, 1023, 4095, 65535].into_iter().enumerate() {
let a = data(n, max, 3 + k as u64);
let b = data(n, max, 9 + k as u64);
let full = vec![u16::MAX; n];
for v in [&a, &full] {
unsafe {
assert_eq!(x86::sum(v), scalar::sum(v));
assert_eq!(x86::sum(v), prev::sum(v));
assert_eq!(x86::min_max(v), scalar::min_max(v));
assert_eq!(x86::count_range(v, 3, 900), scalar::count_range(v, 3, 900));
assert_eq!(x86::count_range(v, 0, u16::MAX), n as u64);
}
}
let zero = vec![0u16; n];
unsafe {
assert_eq!(x86::sad(&a, &b), scalar::sad(&a, &b));
assert_eq!(x86::sad(&a, &b), prev::sad(&a, &b));
assert_eq!(x86::sad(&full, &zero), u64::from(u16::MAX) * n as u64);
}
}
}
}
#[test]
#[ignore]
fn bench_reductions() {
use core::hint::black_box as slain;
if !slain_rt::cpu::avx2() {
return;
}
for (label, n) in [("L1 4096", 4096usize), ("L2 65536", 65536), ("DRAM 1920x1080", 1920 * 1080)] {
let a = data(n, 1023, 1);
let b = data(n, 1023, 2);
let reps = (1 << 24) / n.max(1);
let arms: [(&str, &dyn Fn()); 8] = [
("sum prev", &|| {
slain(unsafe { prev::sum(slain(&a)) });
}),
("sum new", &|| {
slain(unsafe { x86::sum(slain(&a)) });
}),
("sad prev", &|| {
slain(unsafe { prev::sad(slain(&a), slain(&b)) });
}),
("sad new", &|| {
slain(unsafe { x86::sad(slain(&a), slain(&b)) });
}),
("min_max prev", &|| {
slain(unsafe { prev::min_max(slain(&a)) });
}),
("min_max new", &|| {
slain(unsafe { x86::min_max(slain(&a)) });
}),
("count_range prev", &|| {
slain(unsafe { prev::count_range(slain(&a), 64, 940) });
}),
("count_range new", &|| {
slain(unsafe { x86::count_range(slain(&a), 64, 940) });
}),
];
let mut ns = vec![Vec::new(); arms.len()];
for round in 0..21 {
for k in 0..arms.len() {
let j = (k + round) % arms.len();
let t = std::time::Instant::now();
for _ in 0..reps {
(arms[j].1)();
}
ns[j].push(t.elapsed().as_nanos() as f64 / reps as f64);
}
}
for v in &mut ns {
v.sort_by(f64::total_cmp);
}
println!("{label}");
for pair in 0..arms.len() / 2 {
let (o, w) = (ns[2 * pair][10], ns[2 * pair + 1][10]);
println!(" {:12} prev {:>10.1} ns new {:>10.1} ns {:.2}x", arms[2 * pair].0.trim_end_matches(" prev"), o, w, o / w);
}
}
}
}