1#[cfg(target_arch = "x86_64")]
2use std::arch::x86_64::*;
3
4pub fn sad(a: &[u8], sa: usize, b: &[u8], sb: usize, w: usize, h: usize) -> u32 {
5 check(a, sa, b, sb, w, h);
6 #[cfg(target_arch = "x86_64")]
7 if crate::media::isa::avx2() {
8 return unsafe { sad_avx2(a, sa, b, sb, w, h) };
9 }
10 sad_scalar(a, sa, b, sb, w, h)
11}
12
13pub fn satd(a: &[u8], sa: usize, b: &[u8], sb: usize, w: usize, h: usize) -> u32 {
14 check(a, sa, b, sb, w, h);
15 if w < 8 || h < 8 {
16 #[cfg(target_arch = "x86_64")]
17 if crate::media::isa::avx2() {
18 return unsafe { satd4_avx2(a, sa, b, sb, w, h) };
19 }
20 return satd4_scalar(a, sa, b, sb, w, h);
21 }
22 #[cfg(target_arch = "x86_64")]
23 if crate::media::isa::avx2() {
24 return unsafe { satd8_avx2(a, sa, b, sb, w, h) };
25 }
26 satd8_scalar(a, sa, b, sb, w, h)
27}
28
29pub fn sse(a: &[u8], sa: usize, b: &[u8], sb: usize, w: usize, h: usize) -> u32 {
30 check(a, sa, b, sb, w, h);
31 #[cfg(target_arch = "x86_64")]
32 if crate::media::isa::avx2() && w >= 16 {
33 return unsafe { sse_avx2(a, sa, b, sb, w, h) };
34 }
35 sse_scalar(a, sa, b, sb, w, h)
36}
37
38pub fn variance(a: &[u8], sa: usize, w: usize, h: usize) -> u32 {
39 let (mut s, mut q) = (0u64, 0u64);
40 for i in 0..h {
41 for &v in &a[i * sa..i * sa + w] {
42 s += v as u64;
43 q += (v as u64) * (v as u64);
44 }
45 }
46 let n = (w * h) as u64;
47 (q - s * s / n) as u32
48}
49
50#[inline(always)]
51fn check(a: &[u8], sa: usize, b: &[u8], sb: usize, w: usize, h: usize) {
52 assert!(h > 0 && w > 0);
53 assert!(a.len() >= (h - 1) * sa + w && b.len() >= (h - 1) * sb + w);
54}
55
56pub fn sad_scalar(a: &[u8], sa: usize, b: &[u8], sb: usize, w: usize, h: usize) -> u32 {
57 let mut s = 0u32;
58 for i in 0..h {
59 for (x, y) in a[i * sa..i * sa + w].iter().zip(&b[i * sb..i * sb + w]) {
60 s += x.abs_diff(*y) as u32;
61 }
62 }
63 s
64}
65
66pub fn sse_scalar(a: &[u8], sa: usize, b: &[u8], sb: usize, w: usize, h: usize) -> u32 {
67 let mut s = 0u32;
68 for i in 0..h {
69 for (x, y) in a[i * sa..i * sa + w].iter().zip(&b[i * sb..i * sb + w]) {
70 let d = *x as i32 - *y as i32;
71 s += (d * d) as u32;
72 }
73 }
74 s
75}
76
77fn hadamard_sum<const N: usize>(d: &mut [[i32; N]; N]) -> u32 {
78 for r in d.iter_mut() {
79 let mut step = 1;
80 while step < N {
81 for i in 0..N {
82 if i & step == 0 {
83 let (x, y) = (r[i], r[i + step]);
84 r[i] = x + y;
85 r[i + step] = x - y;
86 }
87 }
88 step <<= 1;
89 }
90 }
91 let mut s = 0u32;
92 for c in 0..N {
93 let mut col = [0i32; N];
94 for i in 0..N {
95 col[i] = d[i][c];
96 }
97 let mut step = 1;
98 while step < N {
99 for i in 0..N {
100 if i & step == 0 {
101 let (x, y) = (col[i], col[i + step]);
102 col[i] = x + y;
103 col[i + step] = x - y;
104 }
105 }
106 step <<= 1;
107 }
108 s += col.iter().map(|v| v.unsigned_abs()).sum::<u32>();
109 }
110 s
111}
112
113fn satd_tiles<const N: usize>(a: &[u8], sa: usize, b: &[u8], sb: usize, w: usize, h: usize) -> u32 {
114 let mut s = 0;
115 for by in (0..h).step_by(N) {
116 for bx in (0..w).step_by(N) {
117 let mut d = [[0i32; N]; N];
118 for i in 0..N {
119 for j in 0..N {
120 d[i][j] = a[(by + i) * sa + bx + j] as i32 - b[(by + i) * sb + bx + j] as i32;
121 }
122 }
123 s += hadamard_sum(&mut d);
124 }
125 }
126 s
127}
128
129pub fn satd4_scalar(a: &[u8], sa: usize, b: &[u8], sb: usize, w: usize, h: usize) -> u32 {
130 satd_tiles::<4>(a, sa, b, sb, w, h)
131}
132
133pub fn satd8_scalar(a: &[u8], sa: usize, b: &[u8], sb: usize, w: usize, h: usize) -> u32 {
134 satd_tiles::<8>(a, sa, b, sb, w, h)
135}
136
137#[cfg(target_arch = "x86_64")]
138#[target_feature(enable = "avx2")]
139fn hsum_epi64(v: __m256i) -> u32 {
140 let s = _mm_add_epi64(_mm256_castsi256_si128(v), _mm256_extracti128_si256::<1>(v));
141 (_mm_cvtsi128_si64(s) + _mm_extract_epi64::<1>(s)) as u32
142}
143
144#[cfg(target_arch = "x86_64")]
145#[target_feature(enable = "avx2")]
146fn hsum_epi32(v: __m256i) -> u32 {
147 let s = _mm_add_epi32(_mm256_castsi256_si128(v), _mm256_extracti128_si256::<1>(v));
148 let s = _mm_add_epi32(s, _mm_shuffle_epi32::<0b01_00_11_10>(s));
149 let s = _mm_add_epi32(s, _mm_shuffle_epi32::<0b10_11_00_01>(s));
150 _mm_cvtsi128_si32(s) as u32
151}
152
153#[cfg(target_arch = "x86_64")]
154#[target_feature(enable = "avx2")]
155unsafe fn sad_avx2(a: &[u8], sa: usize, b: &[u8], sb: usize, w: usize, h: usize) -> u32 {
156 unsafe {
157 let (pa, pb) = (a.as_ptr(), b.as_ptr());
158 let mut acc = _mm256_setzero_si256();
159 match w {
160 64 => {
161 for i in 0..h {
162 let (ra, rb) = (pa.add(i * sa), pb.add(i * sb));
163 let x0 = _mm256_loadu_si256(ra as *const __m256i);
164 let y0 = _mm256_loadu_si256(rb as *const __m256i);
165 let x1 = _mm256_loadu_si256(ra.add(32) as *const __m256i);
166 let y1 = _mm256_loadu_si256(rb.add(32) as *const __m256i);
167 acc = _mm256_add_epi64(acc, _mm256_sad_epu8(x0, y0));
168 acc = _mm256_add_epi64(acc, _mm256_sad_epu8(x1, y1));
169 }
170 }
171 32 => {
172 for i in 0..h {
173 let x = _mm256_loadu_si256(pa.add(i * sa) as *const __m256i);
174 let y = _mm256_loadu_si256(pb.add(i * sb) as *const __m256i);
175 acc = _mm256_add_epi64(acc, _mm256_sad_epu8(x, y));
176 }
177 }
178 16 => {
179 let mut i = 0;
180 while i + 1 < h {
181 let x = _mm256_loadu2_m128i(
182 pa.add((i + 1) * sa) as *const __m128i,
183 pa.add(i * sa) as *const __m128i,
184 );
185 let y = _mm256_loadu2_m128i(
186 pb.add((i + 1) * sb) as *const __m128i,
187 pb.add(i * sb) as *const __m128i,
188 );
189 acc = _mm256_add_epi64(acc, _mm256_sad_epu8(x, y));
190 i += 2;
191 }
192 if i < h {
193 return hsum_epi64(acc)
194 + sad_scalar(&a[i * sa..], sa, &b[i * sb..], sb, w, h - i);
195 }
196 }
197 8 => {
198 let mut i = 0;
199 while i + 3 < h {
200 let x = _mm256_set_epi64x(
201 (pa.add((i + 3) * sa) as *const i64).read_unaligned(),
202 (pa.add((i + 2) * sa) as *const i64).read_unaligned(),
203 (pa.add((i + 1) * sa) as *const i64).read_unaligned(),
204 (pa.add(i * sa) as *const i64).read_unaligned(),
205 );
206 let y = _mm256_set_epi64x(
207 (pb.add((i + 3) * sb) as *const i64).read_unaligned(),
208 (pb.add((i + 2) * sb) as *const i64).read_unaligned(),
209 (pb.add((i + 1) * sb) as *const i64).read_unaligned(),
210 (pb.add(i * sb) as *const i64).read_unaligned(),
211 );
212 acc = _mm256_add_epi64(acc, _mm256_sad_epu8(x, y));
213 i += 4;
214 }
215 if i < h {
216 return hsum_epi64(acc)
217 + sad_scalar(&a[i * sa..], sa, &b[i * sb..], sb, w, h - i);
218 }
219 }
220 4 if h.is_multiple_of(4) => {
221 let mut acc = _mm_setzero_si128();
222 for i in (0..h).step_by(4) {
223 let x = _mm_set_epi32(
224 (pa.add((i + 3) * sa) as *const i32).read_unaligned(),
225 (pa.add((i + 2) * sa) as *const i32).read_unaligned(),
226 (pa.add((i + 1) * sa) as *const i32).read_unaligned(),
227 (pa.add(i * sa) as *const i32).read_unaligned(),
228 );
229 let y = _mm_set_epi32(
230 (pb.add((i + 3) * sb) as *const i32).read_unaligned(),
231 (pb.add((i + 2) * sb) as *const i32).read_unaligned(),
232 (pb.add((i + 1) * sb) as *const i32).read_unaligned(),
233 (pb.add(i * sb) as *const i32).read_unaligned(),
234 );
235 acc = _mm_add_epi64(acc, _mm_sad_epu8(x, y));
236 }
237 return (_mm_cvtsi128_si64(acc) + _mm_extract_epi64::<1>(acc)) as u32;
238 }
239 _ => return sad_scalar(a, sa, b, sb, w, h),
240 }
241 hsum_epi64(acc)
242 }
243}
244
245#[cfg(target_arch = "x86_64")]
246#[target_feature(enable = "avx2")]
247unsafe fn sse_avx2(a: &[u8], sa: usize, b: &[u8], sb: usize, w: usize, h: usize) -> u32 {
248 unsafe {
249 let mut acc = _mm256_setzero_si256();
250 for i in 0..h {
251 let (ra, rb) = (a.as_ptr().add(i * sa), b.as_ptr().add(i * sb));
252 for x in (0..w).step_by(16) {
253 let p = _mm256_cvtepu8_epi16(_mm_loadu_si128(ra.add(x) as *const __m128i));
254 let q = _mm256_cvtepu8_epi16(_mm_loadu_si128(rb.add(x) as *const __m128i));
255 let d = _mm256_sub_epi16(p, q);
256 acc = _mm256_add_epi32(acc, _mm256_madd_epi16(d, d));
257 }
258 }
259 hsum_epi32(acc)
260 }
261}
262
263#[cfg(target_arch = "x86_64")]
264#[target_feature(enable = "avx2")]
265fn butterfly(r: &mut [__m256i; 8], step: usize) {
266 for i in 0..8 {
267 if i & step == 0 {
268 let (x, y) = (r[i], r[i + step]);
269 r[i] = _mm256_add_epi16(x, y);
270 r[i + step] = _mm256_sub_epi16(x, y);
271 }
272 }
273}
274
275#[cfg(target_arch = "x86_64")]
276#[target_feature(enable = "avx2")]
277fn transpose8(r: &mut [__m256i; 8]) {
278 let t0 = _mm256_unpacklo_epi16(r[0], r[1]);
279 let t1 = _mm256_unpackhi_epi16(r[0], r[1]);
280 let t2 = _mm256_unpacklo_epi16(r[2], r[3]);
281 let t3 = _mm256_unpackhi_epi16(r[2], r[3]);
282 let t4 = _mm256_unpacklo_epi16(r[4], r[5]);
283 let t5 = _mm256_unpackhi_epi16(r[4], r[5]);
284 let t6 = _mm256_unpacklo_epi16(r[6], r[7]);
285 let t7 = _mm256_unpackhi_epi16(r[6], r[7]);
286 let u0 = _mm256_unpacklo_epi32(t0, t2);
287 let u1 = _mm256_unpackhi_epi32(t0, t2);
288 let u2 = _mm256_unpacklo_epi32(t1, t3);
289 let u3 = _mm256_unpackhi_epi32(t1, t3);
290 let u4 = _mm256_unpacklo_epi32(t4, t6);
291 let u5 = _mm256_unpackhi_epi32(t4, t6);
292 let u6 = _mm256_unpacklo_epi32(t5, t7);
293 let u7 = _mm256_unpackhi_epi32(t5, t7);
294 r[0] = _mm256_unpacklo_epi64(u0, u4);
295 r[1] = _mm256_unpackhi_epi64(u0, u4);
296 r[2] = _mm256_unpacklo_epi64(u1, u5);
297 r[3] = _mm256_unpackhi_epi64(u1, u5);
298 r[4] = _mm256_unpacklo_epi64(u2, u6);
299 r[5] = _mm256_unpackhi_epi64(u2, u6);
300 r[6] = _mm256_unpacklo_epi64(u3, u7);
301 r[7] = _mm256_unpackhi_epi64(u3, u7);
302}
303
304#[cfg(target_arch = "x86_64")]
305#[target_feature(enable = "avx2")]
306fn hadamard_abs(r: &mut [__m256i; 8]) -> __m256i {
307 butterfly(r, 1);
308 butterfly(r, 2);
309 butterfly(r, 4);
310 transpose8(r);
311 butterfly(r, 1);
312 butterfly(r, 2);
313 butterfly(r, 4);
314 let ones = _mm256_set1_epi16(1);
315 let mut acc = _mm256_setzero_si256();
316 for k in (0..8).step_by(2) {
317 let s = _mm256_add_epi16(_mm256_abs_epi16(r[k]), _mm256_abs_epi16(r[k + 1]));
318 acc = _mm256_add_epi32(acc, _mm256_madd_epi16(s, ones));
319 }
320 acc
321}
322
323#[cfg(target_arch = "x86_64")]
324#[target_feature(enable = "avx2")]
325unsafe fn satd8_avx2(a: &[u8], sa: usize, b: &[u8], sb: usize, w: usize, h: usize) -> u32 {
326 unsafe {
327 let (pa, pb) = (a.as_ptr(), b.as_ptr());
328 let mut acc = _mm256_setzero_si256();
329 let mut r = [_mm256_setzero_si256(); 8];
330 for by in (0..h).step_by(8) {
331 let mut bx = 0;
332 while bx + 16 <= w {
333 for i in 0..8 {
334 let p = _mm256_cvtepu8_epi16(_mm_loadu_si128(
335 pa.add((by + i) * sa + bx) as *const __m128i
336 ));
337 let q = _mm256_cvtepu8_epi16(_mm_loadu_si128(
338 pb.add((by + i) * sb + bx) as *const __m128i
339 ));
340 r[i] = _mm256_sub_epi16(p, q);
341 }
342 acc = _mm256_add_epi32(acc, hadamard_abs(&mut r));
343 bx += 16;
344 }
345 if bx < w {
346 let mut i = 0;
347 while i < 8 {
348 let p = _mm256_cvtepu8_epi16(_mm_set_epi64x(
349 (pa.add((by + i + 1) * sa + bx) as *const i64).read_unaligned(),
350 (pa.add((by + i) * sa + bx) as *const i64).read_unaligned(),
351 ));
352 let q = _mm256_cvtepu8_epi16(_mm_set_epi64x(
353 (pb.add((by + i + 1) * sb + bx) as *const i64).read_unaligned(),
354 (pb.add((by + i) * sb + bx) as *const i64).read_unaligned(),
355 ));
356 let d = _mm256_sub_epi16(p, q);
357 r[i] = _mm256_permute4x64_epi64::<0b01_00_01_00>(d);
358 r[i + 1] = _mm256_permute4x64_epi64::<0b11_10_11_10>(d);
359 i += 2;
360 }
361 let half = hadamard_abs(&mut r);
362 acc = _mm256_add_epi32(acc, _mm256_permute2x128_si256::<0x80>(half, half));
363 }
364 }
365 hsum_epi32(acc)
366 }
367}
368
369#[cfg(target_arch = "x86_64")]
370#[target_feature(enable = "avx2")]
371unsafe fn satd4_avx2(a: &[u8], sa: usize, b: &[u8], sb: usize, w: usize, h: usize) -> u32 {
372 unsafe {
373 let (pa, pb) = (a.as_ptr(), b.as_ptr());
374 let ones = _mm256_set1_epi16(1);
375 let mut acc = _mm256_setzero_si256();
376 let rows = |p: *const u8, s: usize| {
377 _mm256_cvtepu8_epi16(_mm_set_epi32(
378 (p.add(3 * s) as *const i32).read_unaligned(),
379 (p.add(2 * s) as *const i32).read_unaligned(),
380 (p.add(s) as *const i32).read_unaligned(),
381 (p as *const i32).read_unaligned(),
382 ))
383 };
384 for by in (0..h).step_by(4) {
385 for bx in (0..w).step_by(4) {
386 let v = _mm256_sub_epi16(
387 rows(pa.add(by * sa + bx), sa),
388 rows(pb.add(by * sb + bx), sb),
389 );
390 let p = _mm256_permute2x128_si256::<0x01>(v, v);
391 let x = _mm256_permute2x128_si256::<0x20>(
392 _mm256_add_epi16(v, p),
393 _mm256_sub_epi16(v, p),
394 );
395 let q = _mm256_shuffle_epi32::<0b01_00_11_10>(x);
396 let y = _mm256_unpacklo_epi64(_mm256_add_epi16(x, q), _mm256_sub_epi16(x, q));
397 let q = _mm256_shuffle_epi32::<0b10_11_00_01>(y);
398 let z = _mm256_blend_epi32::<0b1010_1010>(
399 _mm256_add_epi16(y, q),
400 _mm256_sub_epi16(y, q),
401 );
402 let q = _mm256_shufflehi_epi16::<0b10_11_00_01>(_mm256_shufflelo_epi16::<
403 0b10_11_00_01,
404 >(z));
405 let f = _mm256_blend_epi16::<0b1010_1010>(
406 _mm256_add_epi16(z, q),
407 _mm256_sub_epi16(z, q),
408 );
409 acc = _mm256_add_epi32(acc, _mm256_madd_epi16(_mm256_abs_epi16(f), ones));
410 }
411 }
412 hsum_epi32(acc)
413 }
414}