1pub mod block;
2
3#[derive(Clone, Copy, Debug)]
4pub struct View<'a, T> {
5 data: &'a [T],
6 w: usize,
7 h: usize,
8 stride: usize,
9}
10
11#[derive(Debug)]
12pub struct ViewMut<'a, T> {
13 data: &'a mut [T],
14 w: usize,
15 h: usize,
16 stride: usize,
17}
18
19#[inline]
20fn fits(len: usize, w: usize, h: usize, stride: usize) -> bool {
21 stride >= w && (h == 0 || w == 0 || (h - 1).checked_mul(stride).and_then(|n| n.checked_add(w)).is_some_and(|n| n <= len))
22}
23
24impl<'a, T: Copy> View<'a, T> {
25 #[must_use]
26 pub fn new(data: &'a [T], w: usize, h: usize, stride: usize) -> Option<Self> {
27 fits(data.len(), w, h, stride).then_some(Self { data, w, h, stride })
28 }
29
30 #[must_use]
31 pub fn w(&self) -> usize {
32 self.w
33 }
34
35 #[must_use]
36 pub fn h(&self) -> usize {
37 self.h
38 }
39
40 #[must_use]
41 pub fn stride(&self) -> usize {
42 self.stride
43 }
44
45 #[inline]
46 #[must_use]
47 pub fn row(&self, y: usize) -> &'a [T] {
48 assert!(y < self.h);
49 &self.data[y * self.stride..y * self.stride + self.w]
50 }
51
52 pub fn rows(&self) -> impl Iterator<Item = &'a [T]> + '_ {
53 (0..self.h).map(|y| self.row(y))
54 }
55
56 #[must_use]
57 pub fn get(&self, x: usize, y: usize) -> T {
58 self.row(y)[x]
59 }
60
61 #[must_use]
62 pub fn sub(&self, x: usize, y: usize, w: usize, h: usize) -> Option<Self> {
63 if x.checked_add(w)? > self.w || y.checked_add(h)? > self.h {
64 return None;
65 }
66 if w == 0 || h == 0 {
67 return Some(Self { data: &self.data[..0], w, h, stride: self.stride });
68 }
69 Some(Self { data: &self.data[y * self.stride + x..], w, h, stride: self.stride })
70 }
71}
72
73impl<'a, T: Copy> ViewMut<'a, T> {
74 #[must_use]
75 pub fn new(data: &'a mut [T], w: usize, h: usize, stride: usize) -> Option<Self> {
76 fits(data.len(), w, h, stride).then_some(Self { data, w, h, stride })
77 }
78
79 #[must_use]
80 pub fn w(&self) -> usize {
81 self.w
82 }
83
84 #[must_use]
85 pub fn h(&self) -> usize {
86 self.h
87 }
88
89 #[must_use]
90 pub fn stride(&self) -> usize {
91 self.stride
92 }
93
94 #[must_use]
95 pub fn view(&self) -> View<'_, T> {
96 View { data: self.data, w: self.w, h: self.h, stride: self.stride }
97 }
98
99 #[inline]
100 pub fn row_mut(&mut self, y: usize) -> &mut [T] {
101 assert!(y < self.h);
102 &mut self.data[y * self.stride..y * self.stride + self.w]
103 }
104
105 pub fn sub_mut(&mut self, x: usize, y: usize, w: usize, h: usize) -> Option<ViewMut<'_, T>> {
106 if x.checked_add(w)? > self.w || y.checked_add(h)? > self.h {
107 return None;
108 }
109 if w == 0 || h == 0 {
110 return Some(ViewMut { data: &mut self.data[..0], w, h, stride: self.stride });
111 }
112 Some(ViewMut { data: &mut self.data[y * self.stride + x..], w, h, stride: self.stride })
113 }
114}
115
116#[cfg(target_arch = "x86_64")]
117#[inline]
118fn fast() -> bool {
119 slain_rt::cpu::avx2()
120}
121
122fn same<A: Copy, B: Copy>(a: &View<'_, A>, b: &View<'_, B>) {
123 assert!(a.w == b.w && a.h == b.h, "shape {}x{} vs {}x{}", a.w, a.h, b.w, b.h);
124}
125
126pub fn fill(d: &mut ViewMut<'_, u16>, v: u16) {
127 for y in 0..d.h {
128 d.row_mut(y).fill(v);
129 }
130}
131
132pub fn copy(d: &mut ViewMut<'_, u16>, s: &View<'_, u16>) {
133 same(&d.view(), s);
134 for y in 0..d.h {
135 d.row_mut(y).copy_from_slice(s.row(y));
136 }
137}
138
139pub fn clip(d: &mut ViewMut<'_, u16>, lo: u16, hi: u16) {
140 #[cfg(target_arch = "x86_64")]
141 let f = fast();
142 for y in 0..d.h {
143 let r = d.row_mut(y);
144 #[cfg(target_arch = "x86_64")]
145 if f {
146 unsafe { x86::clip(r, lo, hi) };
147 continue;
148 }
149 scalar::clip(r, lo, hi);
150 }
151}
152
153pub fn avg(d: &mut ViewMut<'_, u16>, a: &View<'_, u16>, b: &View<'_, u16>) {
154 same(&d.view(), a);
155 same(a, b);
156 #[cfg(target_arch = "x86_64")]
157 let f = fast();
158 for y in 0..d.h {
159 let r = d.row_mut(y);
160 #[cfg(target_arch = "x86_64")]
161 if f {
162 unsafe { x86::avg(r, a.row(y), b.row(y)) };
163 continue;
164 }
165 scalar::avg(r, a.row(y), b.row(y));
166 }
167}
168
169pub fn add_sat(d: &mut ViewMut<'_, u16>, a: &View<'_, u16>) {
170 same(&d.view(), a);
171 #[cfg(target_arch = "x86_64")]
172 let f = fast();
173 for y in 0..d.h {
174 let r = d.row_mut(y);
175 #[cfg(target_arch = "x86_64")]
176 if f {
177 unsafe { x86::add_sat(r, a.row(y)) };
178 continue;
179 }
180 scalar::add_sat(r, a.row(y));
181 }
182}
183
184pub fn add_residual(d: &mut ViewMut<'_, u16>, res: &View<'_, i16>, max: u16) {
185 assert!(max <= 0x7fff);
186 same(&d.view(), res);
187 #[cfg(target_arch = "x86_64")]
188 let f = fast();
189 for y in 0..d.h {
190 let r = d.row_mut(y);
191 #[cfg(target_arch = "x86_64")]
192 if f {
193 unsafe { x86::add_residual(r, res.row(y), max) };
194 continue;
195 }
196 scalar::add_residual(r, res.row(y), max);
197 }
198}
199
200#[must_use]
201pub fn sum(a: &View<'_, u16>) -> u64 {
202 #[cfg(target_arch = "x86_64")]
203 let f = fast();
204 a.rows()
205 .map(|r| {
206 #[cfg(target_arch = "x86_64")]
207 if f {
208 return unsafe { x86::sum(r) };
209 }
210 scalar::sum(r)
211 })
212 .sum()
213}
214
215#[must_use]
216pub fn sad(a: &View<'_, u16>, b: &View<'_, u16>) -> u64 {
217 same(a, b);
218 #[cfg(target_arch = "x86_64")]
219 let f = fast();
220 (0..a.h)
221 .map(|y| {
222 #[cfg(target_arch = "x86_64")]
223 if f {
224 return unsafe { x86::sad(a.row(y), b.row(y)) };
225 }
226 scalar::sad(a.row(y), b.row(y))
227 })
228 .sum()
229}
230
231#[must_use]
232pub fn ssd(a: &View<'_, u16>, b: &View<'_, u16>) -> u64 {
233 same(a, b);
234 #[cfg(target_arch = "x86_64")]
235 let f = fast();
236 (0..a.h)
237 .map(|y| {
238 #[cfg(target_arch = "x86_64")]
239 if f {
240 return unsafe { x86::ssd(a.row(y), b.row(y)) };
241 }
242 scalar::ssd(a.row(y), b.row(y))
243 })
244 .sum()
245}
246
247#[must_use]
248pub fn min_max(a: &View<'_, u16>) -> Option<(u16, u16)> {
249 if a.w == 0 || a.h == 0 {
250 return None;
251 }
252 #[cfg(target_arch = "x86_64")]
253 let f = fast();
254 Some(a.rows().fold((u16::MAX, 0), |(lo, hi), r| {
255 #[cfg(target_arch = "x86_64")]
256 if f {
257 let (l, h) = unsafe { x86::min_max(r) };
258 return (lo.min(l), hi.max(h));
259 }
260 let (l, h) = scalar::min_max(r);
261 (lo.min(l), hi.max(h))
262 }))
263}
264
265pub fn absdiff(d: &mut ViewMut<'_, u16>, a: &View<'_, u16>, b: &View<'_, u16>) {
266 same(&d.view(), a);
267 same(a, b);
268 #[cfg(target_arch = "x86_64")]
269 let f = fast();
270 for y in 0..d.h {
271 let r = d.row_mut(y);
272 #[cfg(target_arch = "x86_64")]
273 if f {
274 unsafe { x86::absdiff(r, a.row(y), b.row(y)) };
275 continue;
276 }
277 scalar::absdiff(r, a.row(y), b.row(y));
278 }
279}
280
281#[must_use]
282pub fn count_range(a: &View<'_, u16>, lo: u16, hi: u16) -> u64 {
283 #[cfg(target_arch = "x86_64")]
284 let f = fast();
285 a.rows()
286 .map(|r| {
287 #[cfg(target_arch = "x86_64")]
288 if f {
289 return unsafe { x86::count_range(r, lo, hi) };
290 }
291 scalar::count_range(r, lo, hi)
292 })
293 .sum()
294}
295
296pub mod scalar {
297 pub fn absdiff(d: &mut [u16], a: &[u16], b: &[u16]) {
298 for ((o, &x), &y) in d.iter_mut().zip(a).zip(b) {
299 *o = x.abs_diff(y);
300 }
301 }
302
303 #[must_use]
304 pub fn count_range(a: &[u16], lo: u16, hi: u16) -> u64 {
305 a.iter().filter(|&&v| lo <= v && v <= hi).count() as u64
306 }
307
308 pub fn clip(d: &mut [u16], lo: u16, hi: u16) {
309 for v in d {
310 *v = (*v).max(lo).min(hi);
311 }
312 }
313
314 pub fn avg(d: &mut [u16], a: &[u16], b: &[u16]) {
315 for ((o, &x), &y) in d.iter_mut().zip(a).zip(b) {
316 *o = ((u32::from(x) + u32::from(y) + 1) >> 1) as u16;
317 }
318 }
319
320 pub fn add_sat(d: &mut [u16], a: &[u16]) {
321 for (o, &x) in d.iter_mut().zip(a) {
322 *o = o.saturating_add(x);
323 }
324 }
325
326 pub fn add_residual(d: &mut [u16], r: &[i16], max: u16) {
327 for (o, &x) in d.iter_mut().zip(r) {
328 *o = (i32::from(*o) + i32::from(x)).clamp(0, i32::from(max)) as u16;
329 }
330 }
331
332 #[must_use]
333 pub fn sum(a: &[u16]) -> u64 {
334 a.iter().map(|&v| u64::from(v)).sum()
335 }
336
337 #[must_use]
338 pub fn sad(a: &[u16], b: &[u16]) -> u64 {
339 a.iter().zip(b).map(|(&x, &y)| u64::from(x.abs_diff(y))).sum()
340 }
341
342 #[must_use]
343 pub fn ssd(a: &[u16], b: &[u16]) -> u64 {
344 a.iter().zip(b).map(|(&x, &y)| u64::from(x.abs_diff(y)).pow(2)).sum()
345 }
346
347 #[must_use]
348 pub fn min_max(a: &[u16]) -> (u16, u16) {
349 a.iter().fold((u16::MAX, 0), |(lo, hi), &v| (lo.min(v), hi.max(v)))
350 }
351}
352
353#[cfg(target_arch = "x86_64")]
354pub mod x86 {
355 use core::arch::x86_64::*;
356
357 #[inline(always)]
358 unsafe fn ld(p: *const u16) -> __m256i {
359 _mm256_loadu_si256(p.cast())
360 }
361
362 #[inline(always)]
363 unsafe fn st(p: *mut u16, v: __m256i) {
364 _mm256_storeu_si256(p.cast(), v)
365 }
366
367 #[inline(always)]
368 unsafe fn widen_add(acc: __m256i, v: __m256i) -> __m256i {
369 let z = _mm256_setzero_si256();
370 _mm256_add_epi32(acc, _mm256_add_epi32(_mm256_unpacklo_epi16(v, z), _mm256_unpackhi_epi16(v, z)))
371 }
372
373 #[inline(always)]
374 unsafe fn hsum32(v: __m256i) -> u64 {
375 let z = _mm256_setzero_si256();
376 let q = _mm256_add_epi64(_mm256_unpacklo_epi32(v, z), _mm256_unpackhi_epi32(v, z));
377 let mut l = [0u64; 4];
378 _mm256_storeu_si256(l.as_mut_ptr().cast(), q);
379 l[0] + l[1] + l[2] + l[3]
380 }
381
382 #[target_feature(enable = "avx2")]
383 pub unsafe fn clip(d: &mut [u16], lo: u16, hi: u16) {
384 let (vl, vh) = (_mm256_set1_epi16(lo as i16), _mm256_set1_epi16(hi as i16));
385 let n = d.len();
386 let p = d.as_mut_ptr();
387 let mut i = 0;
388 while i + 16 <= n {
389 st(p.add(i), _mm256_min_epu16(_mm256_max_epu16(ld(p.add(i)), vl), vh));
390 i += 16;
391 }
392 super::scalar::clip(&mut d[i..], lo, hi);
393 }
394
395 #[target_feature(enable = "avx2")]
396 pub unsafe fn avg(d: &mut [u16], a: &[u16], b: &[u16]) {
397 let n = d.len().min(a.len()).min(b.len());
398 let (p, pa, pb) = (d.as_mut_ptr(), a.as_ptr(), b.as_ptr());
399 let mut i = 0;
400 while i + 16 <= n {
401 st(p.add(i), _mm256_avg_epu16(ld(pa.add(i)), ld(pb.add(i))));
402 i += 16;
403 }
404 super::scalar::avg(&mut d[i..n], &a[i..n], &b[i..n]);
405 }
406
407 #[target_feature(enable = "avx2")]
408 pub unsafe fn add_sat(d: &mut [u16], a: &[u16]) {
409 let n = d.len().min(a.len());
410 let (p, pa) = (d.as_mut_ptr(), a.as_ptr());
411 let mut i = 0;
412 while i + 16 <= n {
413 st(p.add(i), _mm256_adds_epu16(ld(p.add(i)), ld(pa.add(i))));
414 i += 16;
415 }
416 super::scalar::add_sat(&mut d[i..n], &a[i..n]);
417 }
418
419 #[target_feature(enable = "avx2")]
420 pub unsafe fn add_residual(d: &mut [u16], r: &[i16], max: u16) {
421 let n = d.len().min(r.len());
422 let (p, pr) = (d.as_mut_ptr(), r.as_ptr());
423 let (z, m) = (_mm256_setzero_si256(), _mm256_set1_epi16(max as i16));
424 let mut i = 0;
425 while i + 16 <= n {
426 let s = _mm256_adds_epi16(ld(p.add(i)), _mm256_loadu_si256(pr.add(i).cast()));
427 st(p.add(i), _mm256_min_epi16(_mm256_max_epi16(s, z), m));
428 i += 16;
429 }
430 super::scalar::add_residual(&mut d[i..n], &r[i..n], max);
431 }
432
433 #[inline(always)]
434 unsafe fn pairs(v: __m256i) -> __m256i {
435 _mm256_madd_epi16(_mm256_xor_si256(v, _mm256_set1_epi16(i16::MIN)), _mm256_set1_epi16(1))
436 }
437
438 #[inline(always)]
439 unsafe fn hsum_i32(v: __m256i) -> i64 {
440 let q = _mm256_add_epi64(_mm256_cvtepi32_epi64(_mm256_castsi256_si128(v)), _mm256_cvtepi32_epi64(_mm256_extracti128_si256::<1>(v)));
441 let mut l = [0i64; 4];
442 _mm256_storeu_si256(l.as_mut_ptr().cast(), q);
443 l[0] + l[1] + l[2] + l[3]
444 }
445
446 #[inline(always)]
447 unsafe fn sum16(n: usize, f: impl Fn(usize) -> __m256i) -> (u64, usize) {
448 let m = n - n % 16;
449 let z = _mm256_setzero_si256();
450 let mut s = 0i64;
451 let mut i = 0;
452 while i + 64 <= m {
453 let end = i + ((m - i) / 64).min(4096) * 64;
454 let (mut a0, mut a1, mut a2, mut a3) = (z, z, z, z);
455 while i < end {
456 a0 = _mm256_add_epi32(a0, pairs(f(i)));
457 a1 = _mm256_add_epi32(a1, pairs(f(i + 16)));
458 a2 = _mm256_add_epi32(a2, pairs(f(i + 32)));
459 a3 = _mm256_add_epi32(a3, pairs(f(i + 48)));
460 i += 64;
461 }
462 s += hsum_i32(_mm256_add_epi32(_mm256_add_epi32(a0, a1), _mm256_add_epi32(a2, a3)));
463 }
464 let mut a0 = z;
465 while i < m {
466 a0 = _mm256_add_epi32(a0, pairs(f(i)));
467 i += 16;
468 }
469 s += hsum_i32(a0);
470 ((s + 32768 * i as i64) as u64, i)
471 }
472
473 #[target_feature(enable = "avx2")]
474 pub unsafe fn sum(a: &[u16]) -> u64 {
475 let p = a.as_ptr();
476 let (t, i) = sum16(a.len(), |i| ld(p.add(i)));
477 t + super::scalar::sum(&a[i..])
478 }
479
480 #[target_feature(enable = "avx2")]
481 pub unsafe fn sad(a: &[u16], b: &[u16]) -> u64 {
482 let n = a.len().min(b.len());
483 let (pa, pb) = (a.as_ptr(), b.as_ptr());
484 let (t, i) = sum16(n, |i| {
485 let (x, y) = (ld(pa.add(i)), ld(pb.add(i)));
486 _mm256_sub_epi16(_mm256_max_epu16(x, y), _mm256_min_epu16(x, y))
487 });
488 t + super::scalar::sad(&a[i..n], &b[i..n])
489 }
490
491 #[target_feature(enable = "avx2")]
492 pub unsafe fn ssd(a: &[u16], b: &[u16]) -> u64 {
493 let n = a.len().min(b.len());
494 let (pa, pb) = (a.as_ptr(), b.as_ptr());
495 let z = _mm256_setzero_si256();
496 let (mut s0, mut s1) = (z, z);
497 let mut i = 0;
498 while i + 16 <= n {
499 let (x, y) = (ld(pa.add(i)), ld(pb.add(i)));
500 let d = _mm256_sub_epi16(_mm256_max_epu16(x, y), _mm256_min_epu16(x, y));
501 let (lo, hi) = (_mm256_mullo_epi16(d, d), _mm256_mulhi_epu16(d, d));
502 let (p0, p1) = (_mm256_unpacklo_epi16(lo, hi), _mm256_unpackhi_epi16(lo, hi));
503 s0 = _mm256_add_epi64(s0, _mm256_add_epi64(_mm256_unpacklo_epi32(p0, z), _mm256_unpackhi_epi32(p0, z)));
504 s1 = _mm256_add_epi64(s1, _mm256_add_epi64(_mm256_unpacklo_epi32(p1, z), _mm256_unpackhi_epi32(p1, z)));
505 i += 16;
506 }
507 let mut l = [0u64; 4];
508 _mm256_storeu_si256(l.as_mut_ptr().cast(), _mm256_add_epi64(s0, s1));
509 l[0] + l[1] + l[2] + l[3] + super::scalar::ssd(&a[i..n], &b[i..n])
510 }
511
512 #[target_feature(enable = "avx2")]
513 pub unsafe fn absdiff(d: &mut [u16], a: &[u16], b: &[u16]) {
514 let n = d.len().min(a.len()).min(b.len());
515 let (p, pa, pb) = (d.as_mut_ptr(), a.as_ptr(), b.as_ptr());
516 let mut i = 0;
517 while i + 16 <= n {
518 let (x, y) = (ld(pa.add(i)), ld(pb.add(i)));
519 st(p.add(i), _mm256_sub_epi16(_mm256_max_epu16(x, y), _mm256_min_epu16(x, y)));
520 i += 16;
521 }
522 super::scalar::absdiff(&mut d[i..n], &a[i..n], &b[i..n]);
523 }
524
525 #[target_feature(enable = "avx2")]
526 pub unsafe fn count_range(a: &[u16], lo: u16, hi: u16) -> u64 {
527 if lo > hi {
528 return 0;
529 }
530 let n = a.len();
531 let p = a.as_ptr();
532 let (vl, vh, z) = (_mm256_set1_epi16(lo as i16), _mm256_set1_epi16(hi as i16), _mm256_setzero_si256());
533 let hit = |v: __m256i| _mm256_cmpeq_epi16(_mm256_min_epu16(_mm256_max_epu16(v, vl), vh), v);
534 let m = n - n % 16;
535 let mut total = 0u64;
536 let mut i = 0;
537 while i + 64 <= m {
538 let end = i + ((m - i) / 64).min(32768) * 64;
539 let (mut c0, mut c1, mut c2, mut c3) = (z, z, z, z);
540 while i < end {
541 c0 = _mm256_sub_epi16(c0, hit(ld(p.add(i))));
542 c1 = _mm256_sub_epi16(c1, hit(ld(p.add(i + 16))));
543 c2 = _mm256_sub_epi16(c2, hit(ld(p.add(i + 32))));
544 c3 = _mm256_sub_epi16(c3, hit(ld(p.add(i + 48))));
545 i += 64;
546 }
547 total += hsum32(widen_add(widen_add(widen_add(widen_add(z, c0), c1), c2), c3));
548 }
549 let mut c0 = z;
550 while i < m {
551 c0 = _mm256_sub_epi16(c0, hit(ld(p.add(i))));
552 i += 16;
553 }
554 total += hsum32(widen_add(z, c0));
555 total + super::scalar::count_range(&a[i..], lo, hi)
556 }
557
558 #[target_feature(enable = "avx2")]
559 pub unsafe fn min_max(a: &[u16]) -> (u16, u16) {
560 let n = a.len();
561 if n < 16 {
562 return super::scalar::min_max(a);
563 }
564 let p = a.as_ptr();
565 let (f, z) = (_mm256_set1_epi16(-1), _mm256_setzero_si256());
566 let (mut l0, mut l1, mut l2, mut l3) = (f, f, f, f);
567 let (mut h0, mut h1, mut h2, mut h3) = (z, z, z, z);
568 let mut i = 0;
569 while i + 64 <= n {
570 let (v0, v1, v2, v3) = (ld(p.add(i)), ld(p.add(i + 16)), ld(p.add(i + 32)), ld(p.add(i + 48)));
571 l0 = _mm256_min_epu16(l0, v0);
572 l1 = _mm256_min_epu16(l1, v1);
573 l2 = _mm256_min_epu16(l2, v2);
574 l3 = _mm256_min_epu16(l3, v3);
575 h0 = _mm256_max_epu16(h0, v0);
576 h1 = _mm256_max_epu16(h1, v1);
577 h2 = _mm256_max_epu16(h2, v2);
578 h3 = _mm256_max_epu16(h3, v3);
579 i += 64;
580 }
581 while i + 16 <= n {
582 let v = ld(p.add(i));
583 l0 = _mm256_min_epu16(l0, v);
584 h0 = _mm256_max_epu16(h0, v);
585 i += 16;
586 }
587 let v = ld(p.add(n - 16));
588 l1 = _mm256_min_epu16(l1, v);
589 h1 = _mm256_max_epu16(h1, v);
590 let lo = _mm256_min_epu16(_mm256_min_epu16(l0, l1), _mm256_min_epu16(l2, l3));
591 let hi = _mm256_max_epu16(_mm256_max_epu16(h0, h1), _mm256_max_epu16(h2, h3));
592 let lo = _mm_min_epu16(_mm256_castsi256_si128(lo), _mm256_extracti128_si256::<1>(lo));
593 let hi = _mm_max_epu16(_mm256_castsi256_si128(hi), _mm256_extracti128_si256::<1>(hi));
594 let all = _mm_set1_epi16(-1);
595 let vl = _mm_extract_epi16::<0>(_mm_minpos_epu16(lo)) as u16;
596 let vh = !(_mm_extract_epi16::<0>(_mm_minpos_epu16(_mm_xor_si128(hi, all))) as u16);
597 (vl, vh)
598 }
599}
600
601#[cfg(test)]
602mod tests {
603 use super::*;
604
605 fn rng(seed: u64) -> impl FnMut() -> u64 {
606 let mut s = seed;
607 move || {
608 s ^= s << 13;
609 s ^= s >> 7;
610 s ^= s << 17;
611 s
612 }
613 }
614
615 fn plane(_w: usize, h: usize, stride: usize, max: u16, seed: u64) -> Vec<u16> {
616 let mut r = rng(seed);
617 (0..stride * h).map(|_| (r() % (u64::from(max) + 1)) as u16).collect()
618 }
619
620 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)];
621
622 #[test]
623 fn view_bounds() {
624 let d = vec![0u16; 100];
625 assert!(View::new(&d, 10, 10, 10).is_some());
626 assert!(View::new(&d, 10, 11, 10).is_none());
627 assert!(View::new(&d, 11, 1, 10).is_none());
628 assert!(View::new(&d, 4, 9, 11).is_some());
629 assert!(View::new(&d, 4, 10, 11).is_none());
630 let v = View::new(&d, 10, 10, 10).unwrap();
631 assert!(v.sub(5, 5, 5, 5).is_some());
632 assert!(v.sub(6, 5, 5, 5).is_none());
633 assert_eq!(v.sub(2, 3, 4, 5).unwrap().row(4).len(), 4);
634 }
635
636 #[cfg(target_arch = "x86_64")]
637 #[test]
638 fn x86_matches_scalar_every_shape_and_limit() {
639 if !slain_rt::cpu::avx2() {
640 return;
641 }
642 for (k, &(w, h, s)) in SHAPES.iter().enumerate() {
643 for max in [1u16, 255, 1023, 4095, 0x7fff, u16::MAX] {
644 let a = plane(w, h, s, max, 11 + k as u64);
645 let b = plane(w, h, s, max, 97 + k as u64);
646 let va = View::new(&a, w, h, s).unwrap();
647 let vb = View::new(&b, w, h, s).unwrap();
648 for y in 0..h {
649 let (ra, rb) = (va.row(y), vb.row(y));
650 unsafe {
651 assert_eq!(x86::sum(ra), scalar::sum(ra));
652 assert_eq!(x86::sad(ra, rb), scalar::sad(ra, rb));
653 assert_eq!(x86::min_max(ra), scalar::min_max(ra));
654 let (mut p, mut q) = (ra.to_vec(), ra.to_vec());
655 x86::clip(&mut p, max / 4, max / 2);
656 scalar::clip(&mut q, max / 4, max / 2);
657 assert_eq!(p, q);
658 let (mut p, mut q) = (vec![0; w], vec![0; w]);
659 x86::avg(&mut p, ra, rb);
660 scalar::avg(&mut q, ra, rb);
661 assert_eq!(p, q);
662 let (mut p, mut q) = (ra.to_vec(), ra.to_vec());
663 x86::add_sat(&mut p, rb);
664 scalar::add_sat(&mut q, rb);
665 assert_eq!(p, q);
666 let (mut p, mut q) = (vec![0; w], vec![0; w]);
667 x86::absdiff(&mut p, ra, rb);
668 scalar::absdiff(&mut q, ra, rb);
669 assert_eq!(p, q);
670 for (lo, hi) in [(0, 0), (1, u16::MAX), (1, 1), (max / 4, max / 2), (max, max), (5, 3), (0, u16::MAX)] {
671 assert_eq!(x86::count_range(ra, lo, hi), scalar::count_range(ra, lo, hi));
672 }
673 }
674 }
675 if max <= 0x7fff {
676 let mut r = rng(5 + k as u64);
677 let res: Vec<i16> = (0..w).map(|_| r() as i16).collect();
678 let lim: Vec<i16> = (0..w).map(|i| if i % 2 == 0 { i16::MAX } else { i16::MIN }).collect();
679 for rr in [&res, &lim] {
680 let (mut p, mut q) = (va.row(0).to_vec(), va.row(0).to_vec());
681 unsafe { x86::add_residual(&mut p, rr, max) };
682 scalar::add_residual(&mut q, rr, max);
683 assert_eq!(p, q);
684 }
685 }
686 }
687 }
688 }
689
690 #[cfg(target_arch = "x86_64")]
691 #[test]
692 fn count_range_lane_counters_do_not_wrap() {
693 if !slain_rt::cpu::avx2() {
694 return;
695 }
696 let a = vec![9u16; 16 * 32768 * 2 + 16 * 70000 + 5];
697 unsafe {
698 assert_eq!(x86::count_range(&a, 9, 9), a.len() as u64);
699 assert_eq!(x86::count_range(&a, 0, 8), 0);
700 }
701 assert_eq!(count_range(&View::new(&a, a.len(), 1, a.len()).unwrap(), 1, u16::MAX), a.len() as u64);
702 }
703
704 #[test]
705 fn ssd_matches_scalar_and_strided() {
706 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() };
707 for (k, &(w, h, s)) in SHAPES.iter().enumerate() {
708 for max in [1u16, 255, 1023, 0x7fff, u16::MAX] {
709 let a = plane(w, h, s, max, 21 + k as u64);
710 let b = plane(w, h, s, max, 87 + k as u64);
711 let (va, vb) = (View::new(&a, w, h, s).unwrap(), View::new(&b, w, h, s).unwrap());
712 assert_eq!(ssd(&va, &vb), want(&va, &vb));
713 #[cfg(target_arch = "x86_64")]
714 if slain_rt::cpu::avx2() {
715 for y in 0..h {
716 assert_eq!(unsafe { x86::ssd(va.row(y), vb.row(y)) }, scalar::ssd(va.row(y), vb.row(y)));
717 }
718 }
719 }
720 }
721 let (z, m) = (vec![0u16; 4099], vec![u16::MAX; 4099]);
722 assert_eq!(ssd(&View::new(&z, 4099, 1, 4099).unwrap(), &View::new(&m, 4099, 1, 4099).unwrap()), 4099 * 65535 * 65535);
723 let a = plane(67, 13, 96, 255, 3);
724 let b = plane(67, 13, 96, 255, 4);
725 let va = View::new(&a, 67, 13, 96).unwrap().sub(3, 2, 50, 9).unwrap();
726 let vb = View::new(&b, 67, 13, 96).unwrap().sub(3, 2, 50, 9).unwrap();
727 assert_eq!(ssd(&va, &vb), want(&va, &vb));
728 }
729
730 #[test]
731 fn ufuncs_on_strided_subviews() {
732 let (w, h, s) = (67, 13, 96);
733 let a = plane(w, h, s, 1023, 3);
734 let b = plane(w, h, s, 1023, 4);
735 let va = View::new(&a, w, h, s).unwrap().sub(3, 2, 50, 9).unwrap();
736 let vb = View::new(&b, w, h, s).unwrap().sub(3, 2, 50, 9).unwrap();
737 let want_sum: u64 = va.rows().flatten().map(|&v| u64::from(v)).sum();
738 let want_sad: u64 = va.rows().flatten().zip(vb.rows().flatten()).map(|(&x, &y)| u64::from(x.abs_diff(y))).sum();
739 assert_eq!(sum(&va), want_sum);
740 assert_eq!(sad(&va, &vb), want_sad);
741 let mut out = vec![7u16; s * h];
742 let mut vo = ViewMut::new(&mut out, w, h, s).unwrap();
743 let mut so = vo.sub_mut(3, 2, 50, 9).unwrap();
744 avg(&mut so, &va, &vb);
745 for y in 0..9 {
746 for x in 0..50 {
747 assert_eq!(out[(y + 2) * s + x + 3], ((u32::from(va.get(x, y)) + u32::from(vb.get(x, y)) + 1) >> 1) as u16);
748 }
749 assert_eq!(out[(y + 2) * s + 2], 7);
750 assert_eq!(out[(y + 2) * s + 53], 7);
751 }
752 let mut c = a.clone();
753 let mut vc = ViewMut::new(&mut c, w, h, s).unwrap();
754 clip(&mut vc, 100, 900);
755 assert!(View::new(&c, w, h, s).unwrap().rows().flatten().all(|&v| (100..=900).contains(&v)));
756 assert_eq!(min_max(&View::new(&c, w, h, s).unwrap()), Some((100, 900)));
757 }
758
759 #[cfg(target_arch = "x86_64")]
760 #[test]
761 #[ignore]
762 fn bench_ufuncs() {
763 use core::hint::black_box as slain;
764 if !slain_rt::cpu::avx2() {
765 return;
766 }
767 let (w, h) = (1920usize, 1080usize);
768 let a = plane(w, h, w, 1023, 1);
769 let b = plane(w, h, w, 1023, 2);
770 let mut d = vec![0u16; w * h];
771 let time = |f: &mut dyn FnMut()| {
772 let mut best = u128::MAX;
773 for _ in 0..15 {
774 let t = std::time::Instant::now();
775 f();
776 best = best.min(t.elapsed().as_nanos());
777 }
778 best
779 };
780 let rows: Vec<(u128, u128, &str)> = vec![
781 (
782 time(&mut || {
783 slain(scalar::sum(slain(&a)));
784 }),
785 time(&mut || {
786 slain(unsafe { x86::sum(slain(&a)) });
787 }),
788 "sum",
789 ),
790 (
791 time(&mut || {
792 slain(scalar::sad(slain(&a), slain(&b)));
793 }),
794 time(&mut || {
795 slain(unsafe { x86::sad(slain(&a), slain(&b)) });
796 }),
797 "sad",
798 ),
799 (
800 time(&mut || {
801 slain(scalar::min_max(slain(&a)));
802 }),
803 time(&mut || {
804 slain(unsafe { x86::min_max(slain(&a)) });
805 }),
806 "min_max",
807 ),
808 (time(&mut || scalar::avg(slain(&mut d), slain(&a), slain(&b))), time(&mut || unsafe { x86::avg(slain(&mut d), slain(&a), slain(&b)) }), "avg"),
809 (time(&mut || scalar::clip(slain(&mut d), 64, 940)), time(&mut || unsafe { x86::clip(slain(&mut d), 64, 940) }), "clip"),
810 ];
811 for (s, v, name) in rows {
812 println!("{name:8} 1920x1080 scalar {:>9} ns avx2 {:>9} ns {:.2}x", s, v, s as f64 / v as f64);
813 }
814 }
815}
816
817#[cfg(all(test, target_arch = "x86_64"))]
818#[path = "nd_prev.rs"]
819mod prev;
820
821#[cfg(all(test, target_arch = "x86_64"))]
822mod reduce_tests {
823 use super::*;
824
825 fn data(n: usize, max: u32, seed: u64) -> Vec<u16> {
826 let mut s = seed | 1;
827 (0..n)
828 .map(|_| {
829 s ^= s << 13;
830 s ^= s >> 7;
831 s ^= s << 17;
832 (s % (u64::from(max) + 1)) as u16
833 })
834 .collect()
835 }
836
837 #[test]
838 fn register_mapped_reductions_match_prev_and_scalar() {
839 if !slain_rt::cpu::avx2() {
840 return;
841 }
842 for n in [0usize, 1, 15, 16, 17, 63, 64, 65, 127, 4096, 4096 * 64 + 37, 32768 * 64 + 5, 2_000_003] {
843 for (k, max) in [1u32, 1023, 4095, 65535].into_iter().enumerate() {
844 let a = data(n, max, 3 + k as u64);
845 let b = data(n, max, 9 + k as u64);
846 let full = vec![u16::MAX; n];
847 for v in [&a, &full] {
848 unsafe {
849 assert_eq!(x86::sum(v), scalar::sum(v));
850 assert_eq!(x86::sum(v), prev::sum(v));
851 assert_eq!(x86::min_max(v), scalar::min_max(v));
852 assert_eq!(x86::count_range(v, 3, 900), scalar::count_range(v, 3, 900));
853 assert_eq!(x86::count_range(v, 0, u16::MAX), n as u64);
854 }
855 }
856 let zero = vec![0u16; n];
857 unsafe {
858 assert_eq!(x86::sad(&a, &b), scalar::sad(&a, &b));
859 assert_eq!(x86::sad(&a, &b), prev::sad(&a, &b));
860 assert_eq!(x86::sad(&full, &zero), u64::from(u16::MAX) * n as u64);
861 }
862 }
863 }
864 }
865
866 #[test]
867 #[ignore]
868 fn bench_reductions() {
869 use core::hint::black_box as slain;
870 if !slain_rt::cpu::avx2() {
871 return;
872 }
873 for (label, n) in [("L1 4096", 4096usize), ("L2 65536", 65536), ("DRAM 1920x1080", 1920 * 1080)] {
874 let a = data(n, 1023, 1);
875 let b = data(n, 1023, 2);
876 let reps = (1 << 24) / n.max(1);
877 let arms: [(&str, &dyn Fn()); 8] = [
878 ("sum prev", &|| {
879 slain(unsafe { prev::sum(slain(&a)) });
880 }),
881 ("sum new", &|| {
882 slain(unsafe { x86::sum(slain(&a)) });
883 }),
884 ("sad prev", &|| {
885 slain(unsafe { prev::sad(slain(&a), slain(&b)) });
886 }),
887 ("sad new", &|| {
888 slain(unsafe { x86::sad(slain(&a), slain(&b)) });
889 }),
890 ("min_max prev", &|| {
891 slain(unsafe { prev::min_max(slain(&a)) });
892 }),
893 ("min_max new", &|| {
894 slain(unsafe { x86::min_max(slain(&a)) });
895 }),
896 ("count_range prev", &|| {
897 slain(unsafe { prev::count_range(slain(&a), 64, 940) });
898 }),
899 ("count_range new", &|| {
900 slain(unsafe { x86::count_range(slain(&a), 64, 940) });
901 }),
902 ];
903 let mut ns = vec![Vec::new(); arms.len()];
904 for round in 0..21 {
905 for k in 0..arms.len() {
906 let j = (k + round) % arms.len();
907 let t = std::time::Instant::now();
908 for _ in 0..reps {
909 (arms[j].1)();
910 }
911 ns[j].push(t.elapsed().as_nanos() as f64 / reps as f64);
912 }
913 }
914 for v in &mut ns {
915 v.sort_by(f64::total_cmp);
916 }
917 println!("{label}");
918 for pair in 0..arms.len() / 2 {
919 let (o, w) = (ns[2 * pair][10], ns[2 * pair + 1][10]);
920 println!(" {:12} prev {:>10.1} ns new {:>10.1} ns {:.2}x", arms[2 * pair].0.trim_end_matches(" prev"), o, w, o / w);
921 }
922 }
923 }
924}