1pub type SadFn<P> = unsafe fn(*const P, usize, *const P, usize) -> u32;
2pub type PairFn<P> = unsafe fn(*const P, usize, *const P, usize) -> u64;
3
4mod sealed {
5 pub trait Sealed {}
6 impl Sealed for u8 {}
7 impl Sealed for u16 {}
8}
9
10pub trait Px: Copy + Into<u32> + sealed::Sealed + 'static {
11 #[doc(hidden)]
12 fn sads() -> &'static [SadFn<Self>; 25];
13 #[doc(hidden)]
14 fn pair() -> PairFn<Self>;
15}
16
17macro_rules! sizes {
18 ($($f:ident)::+) => {
19 [
20 $($f)::+::<4, 4>, $($f)::+::<4, 8>, $($f)::+::<4, 16>, $($f)::+::<4, 32>, $($f)::+::<4, 64>,
21 $($f)::+::<8, 4>, $($f)::+::<8, 8>, $($f)::+::<8, 16>, $($f)::+::<8, 32>, $($f)::+::<8, 64>,
22 $($f)::+::<16, 4>, $($f)::+::<16, 8>, $($f)::+::<16, 16>, $($f)::+::<16, 32>, $($f)::+::<16, 64>,
23 $($f)::+::<32, 4>, $($f)::+::<32, 8>, $($f)::+::<32, 16>, $($f)::+::<32, 32>, $($f)::+::<32, 64>,
24 $($f)::+::<64, 4>, $($f)::+::<64, 8>, $($f)::+::<64, 16>, $($f)::+::<64, 32>, $($f)::+::<64, 64>,
25 ]
26 };
27}
28
29#[cfg(target_arch = "x86_64")]
30#[path = "block_x86.rs"]
31pub mod x86;
32
33#[cfg(target_arch = "x86_64")]
34mod scalar {
35 use super::SadFn;
36
37 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 {
38 let mut s = 0u32;
39 for _ in 0..H {
40 for c in 0..W {
41 s += (*a.add(c)).into().abs_diff((*b.add(c)).into());
42 }
43 (a, b) = (a.add(sa), b.add(sb));
44 }
45 s
46 }
47
48 pub(super) unsafe fn pair<P: Copy + Into<u32>>(a: *const P, sa: usize, b: *const P, sb: usize) -> u64 {
49 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
50 }
51
52 unsafe fn sad16<const W: usize, const H: usize>(a: *const u16, sa: usize, b: *const u16, sb: usize) -> u32 {
53 sad::<u16, W, H>(a, sa, b, sb)
54 }
55
56 unsafe fn sad8<const W: usize, const H: usize>(a: *const u8, sa: usize, b: *const u8, sb: usize) -> u32 {
57 sad::<u8, W, H>(a, sa, b, sb)
58 }
59
60 pub(super) static SAD_U16: [SadFn<u16>; 25] = sizes!(sad16);
61 pub(super) static SAD_U8: [SadFn<u8>; 25] = sizes!(sad8);
62}
63
64#[cfg(target_arch = "aarch64")]
65#[path = "block_neon.rs"]
66mod arch;
67
68#[cfg(target_arch = "wasm32")]
69#[path = "block_wasm.rs"]
70mod arch;
71
72#[cfg(target_arch = "x86_64")]
73use x86 as arch;
74
75impl Px for u16 {
76 #[inline]
77 fn sads() -> &'static [SadFn<u16>; 25] {
78 #[cfg(target_arch = "x86_64")]
79 if !slain_rt::cpu::avx2() {
80 return &scalar::SAD_U16;
81 }
82 &arch::SAD_U16
83 }
84
85 #[inline]
86 fn pair() -> PairFn<u16> {
87 #[cfg(target_arch = "x86_64")]
88 if !slain_rt::cpu::avx2() {
89 return scalar::pair::<u16>;
90 }
91 arch::pair_u16
92 }
93}
94
95impl Px for u8 {
96 #[inline]
97 fn sads() -> &'static [SadFn<u8>; 25] {
98 #[cfg(target_arch = "x86_64")]
99 if !slain_rt::cpu::avx2() {
100 return &scalar::SAD_U8;
101 }
102 &arch::SAD_U8
103 }
104
105 #[inline]
106 fn pair() -> PairFn<u8> {
107 #[cfg(target_arch = "x86_64")]
108 if !slain_rt::cpu::avx2() {
109 return scalar::pair::<u8>;
110 }
111 arch::pair_u8
112 }
113}
114
115#[inline]
116fn side(v: usize) -> Option<usize> {
117 (v.is_power_of_two() && (4..=64).contains(&v)).then(|| v.trailing_zeros() as usize - 2)
118}
119
120#[derive(Clone, Copy)]
121pub struct Sad<P: Px> {
122 f: SadFn<P>,
123 w: usize,
124 h: usize,
125}
126
127impl<P: Px> Sad<P> {
128 #[inline]
129 #[must_use]
130 pub fn new(w: usize, h: usize) -> Option<Self> {
131 let i = side(w)? * 5 + side(h)?;
132 Some(Self { f: P::sads()[i], w, h })
133 }
134
135 #[inline]
136 #[must_use]
137 pub fn run(&self, a: &[P], sa: usize, b: &[P], sb: usize) -> u32 {
138 assert!(super::fits(a.len(), self.w, self.h, sa) && super::fits(b.len(), self.w, self.h, sb));
139 unsafe { (self.f)(a.as_ptr(), sa, b.as_ptr(), sb) }
140 }
141
142 #[inline]
143 #[must_use]
144 pub fn leaf(&self) -> SadFn<P> {
145 self.f
146 }
147}
148
149pub struct Cells {
150 pub out: [u32; 64],
151}
152
153impl Default for Cells {
154 fn default() -> Self {
155 Self { out: [0; 64] }
156 }
157}
158
159impl Cells {
160 #[inline]
161 pub fn run<P: Px>(&mut self, a: &[P], sa: usize, b: &[P], sb: usize, w: usize, h: usize) -> &[u32] {
162 assert!(matches!(w, 8 | 16 | 32 | 64) && h % 8 == 0 && (8..=64).contains(&h) && super::fits(a.len(), w, h, sa) && super::fits(b.len(), w, h, sb));
163 let (cw, o) = (w / 8, self.out.as_mut_ptr());
164 unsafe {
165 if w == 8 {
166 let f = P::sads()[6];
167 for cy in 0..h / 8 {
168 *o.add(cy) = f(a.as_ptr().add(cy * 8 * sa), sa, b.as_ptr().add(cy * 8 * sb), sb);
169 }
170 } else {
171 let f = P::pair();
172 for cy in 0..h / 8 {
173 let (ra, rb) = (a.as_ptr().add(cy * 8 * sa), b.as_ptr().add(cy * 8 * sb));
174 for cx in (0..cw).step_by(2) {
175 let v = f(ra.add(cx * 8), sa, rb.add(cx * 8), sb);
176 *o.add(cy * cw + cx) = v as u32;
177 *o.add(cy * cw + cx + 1) = (v >> 32) as u32;
178 }
179 }
180 }
181 }
182 &self.out[..cw * (h / 8)]
183 }
184}
185
186#[cfg(target_arch = "x86_64")]
187pub mod probes {
188 #[inline(never)]
189 #[target_feature(enable = "avx2")]
190 pub unsafe fn sad8x8_u16(a: &[u16; 64], b: &[u16; 64]) -> u32 {
191 super::x86::sad_u16::<8, 8>(a.as_ptr(), 8, b.as_ptr(), 8)
192 }
193
194 #[inline(never)]
195 #[target_feature(enable = "avx2")]
196 pub unsafe fn sad8x8_u8(a: &[u8; 64], b: &[u8; 64]) -> u32 {
197 super::x86::sad_u8::<8, 8>(a.as_ptr(), 8, b.as_ptr(), 8)
198 }
199
200 #[inline(never)]
201 #[target_feature(enable = "avx2")]
202 pub unsafe fn pair16x8_u8(a: &[u8; 128], b: &[u8; 128]) -> u64 {
203 super::x86::pair_u8(a.as_ptr(), 16, b.as_ptr(), 16)
204 }
205
206 #[inline(never)]
207 #[target_feature(enable = "avx2")]
208 pub unsafe fn pair16x8_u16(a: &[u16; 128], b: &[u16; 128]) -> u64 {
209 super::x86::pair_u16(a.as_ptr(), 16, b.as_ptr(), 16)
210 }
211}
212
213#[cfg(test)]
214mod tests {
215 use super::*;
216
217 fn next(seed: &mut u64) -> u32 {
218 *seed ^= *seed << 13;
219 *seed ^= *seed >> 7;
220 *seed ^= *seed << 17;
221 (*seed >> 32) as u32
222 }
223
224 fn want<P: Px>(a: &[P], sa: usize, b: &[P], sb: usize, w: usize, h: usize) -> u32 {
225 (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()
226 }
227
228 #[test]
229 fn sad_every_size_and_stride_matches_scalar() {
230 let mut seed = 0x2545_f491_9e37_79b9u64;
231 for w in [4usize, 8, 16, 32, 64] {
232 for h in [4usize, 8, 16, 32, 64] {
233 for (pa, pb) in [(0, 0), (11, 5), (1, 64)] {
234 let (sa, sb) = (w + pa, w + pb);
235 for max in [255u32, 1023, 4095, 32767] {
236 let a: Vec<u16> = (0..sa * h).map(|_| (next(&mut seed) % (max + 1)) as u16).collect();
237 let b: Vec<u16> = (0..sb * h).map(|_| (next(&mut seed) % (max + 1)) as u16).collect();
238 let k = Sad::<u16>::new(w, h).unwrap();
239 assert_eq!(k.run(&a, sa, &b, sb), want(&a, sa, &b, sb, w, h), "u16 {w}x{h} max {max}");
240 if max == 255 {
241 let (a8, b8): (Vec<u8>, Vec<u8>) = (a.iter().map(|&v| v as u8).collect(), b.iter().map(|&v| v as u8).collect());
242 assert_eq!(Sad::<u8>::new(w, h).unwrap().run(&a8, sa, &b8, sb), want(&a8, sa, &b8, sb, w, h), "u8 {w}x{h}");
243 }
244 }
245 }
246 }
247 }
248 let (z, f) = (vec![0u8; 64 * 64], vec![255u8; 64 * 64]);
249 assert_eq!(Sad::<u8>::new(64, 64).unwrap().run(&z, 64, &f, 64), 64 * 64 * 255);
250 let (z, f) = (vec![0u16; 64 * 64], vec![32767u16; 64 * 64]);
251 assert_eq!(Sad::<u16>::new(64, 64).unwrap().run(&z, 64, &f, 64), 64 * 64 * 32767);
252 assert!(Sad::<u16>::new(24, 8).is_none() && Sad::<u8>::new(8, 2).is_none() && Sad::<u8>::new(128, 8).is_none());
253 }
254
255 #[test]
256 fn cells_every_width_and_band_count_match_scalar() {
257 let mut seed = 0x0bad_cafe_1234_5678u64;
258 let mut cells = Cells::default();
259 for w in [8usize, 16, 32, 64] {
260 for h in [8usize, 16, 24, 32, 64] {
261 for max in [255u32, 1023, 4095, 32767] {
262 let (sa, sb) = (w + 24, w);
263 let a: Vec<u16> = (0..sa * h).map(|_| (next(&mut seed) % (max + 1)) as u16).collect();
264 let b: Vec<u16> = (0..sb * h).map(|_| (next(&mut seed) % (max + 1)) as u16).collect();
265 let cell = |cx: usize, cy: usize| want(&a[cy * 8 * sa + cx * 8..], sa, &b[cy * 8 * sb + cx * 8..], sb, 8, 8);
266 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();
267 assert_eq!(cells.run(&a, sa, &b, sb, w, h), expect, "u16 {w}x{h} max {max}");
268 if max == 255 {
269 let (a8, b8): (Vec<u8>, Vec<u8>) = (a.iter().map(|&v| v as u8).collect(), b.iter().map(|&v| v as u8).collect());
270 assert_eq!(cells.run(&a8, sa, &b8, sb, w, h), expect, "u8 {w}x{h}");
271 }
272 }
273 }
274 }
275 }
276
277 #[cfg(target_arch = "x86_64")]
278 fn floor_run(symbol: &str, px: usize, mut call: impl FnMut()) {
279 if !slain_rt::cpu::avx2() {
280 return;
281 }
282 let calls = 20_000u32;
283 let mut best = 0f64;
284 for _ in 0..15 {
285 let at = std::time::Instant::now();
286 for _ in 0..calls {
287 call();
288 }
289 best = best.max(f64::from(calls) / at.elapsed().as_secs_f64());
290 }
291 println!("FLOOR_SYMBOL slain_num::nd::block::probes::{symbol}");
292 println!("FLOOR_PX {px}");
293 println!("FLOOR_CALLS_PER_S {best:.9}");
294 println!("FLOOR_CORE_HZ {:.9}", slain_rt::cpu::core_hz());
295 }
296
297 #[cfg(target_arch = "x86_64")]
298 fn floor_fill<const N: usize>(seed: u32) -> [u16; N] {
299 let mut s = seed;
300 core::array::from_fn(|_| {
301 s ^= s << 13;
302 s ^= s >> 17;
303 s ^= s << 5;
304 (s & 0xff) as u16
305 })
306 }
307
308 #[cfg(target_arch = "x86_64")]
309 #[test]
310 #[ignore = "floor: just floor sad8x8_u16 slain-num"]
311 fn floor_sad8x8_u16() {
312 use core::hint::black_box as slain;
313 let (a, b) = (floor_fill::<64>(0x1357_9bdf), floor_fill::<64>(0x2468_ace1));
314 floor_run("sad8x8_u16", 64, || {
315 slain(unsafe { probes::sad8x8_u16(slain(&a), slain(&b)) });
316 });
317 }
318
319 #[cfg(target_arch = "x86_64")]
320 #[test]
321 #[ignore = "floor: just floor sad8x8_u8 slain-num"]
322 fn floor_sad8x8_u8() {
323 use core::hint::black_box as slain;
324 let (a, b) = (floor_fill::<64>(0x1357_9bdf).map(|v| v as u8), floor_fill::<64>(0x2468_ace1).map(|v| v as u8));
325 floor_run("sad8x8_u8", 64, || {
326 slain(unsafe { probes::sad8x8_u8(slain(&a), slain(&b)) });
327 });
328 }
329
330 #[cfg(target_arch = "x86_64")]
331 #[test]
332 #[ignore = "floor: just floor pair16x8_u8 slain-num"]
333 fn floor_pair16x8_u8() {
334 use core::hint::black_box as slain;
335 let (a, b) = (floor_fill::<128>(0x1357_9bdf).map(|v| v as u8), floor_fill::<128>(0x2468_ace1).map(|v| v as u8));
336 floor_run("pair16x8_u8", 128, || {
337 slain(unsafe { probes::pair16x8_u8(slain(&a), slain(&b)) });
338 });
339 }
340
341 #[cfg(target_arch = "x86_64")]
342 #[test]
343 #[ignore = "floor: just floor pair16x8_u16 slain-num"]
344 fn floor_pair16x8_u16() {
345 use core::hint::black_box as slain;
346 let (a, b) = (floor_fill::<128>(0x1357_9bdf), floor_fill::<128>(0x2468_ace1));
347 floor_run("pair16x8_u16", 128, || {
348 slain(unsafe { probes::pair16x8_u16(slain(&a), slain(&b)) });
349 });
350 }
351}