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