1use super::color::{ColorSpec, FRAC, Siting, fixed, rows_job};
2use super::plane::{Sample, Simd, check, reshape, samples, samples_mut, simd, sized, span};
3use crate::media::Rgba;
4use crate::media::{Error, Frame, PixelFormat, Result};
5use rayon::prelude::*;
6
7#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Default)]
8pub enum Order {
9 #[default]
10 Rgba,
11 Bgra,
12 Rgbx,
13 Bgrx,
14}
15
16impl Order {
17 fn swapped(self) -> bool {
18 matches!(self, Order::Bgra | Order::Bgrx)
19 }
20}
21
22#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Default)]
23pub enum Chroma {
24 #[default]
25 Nearest,
26 Bilinear,
27}
28
29#[derive(Debug, Clone, Copy)]
30pub struct Packed<'a> {
31 pub width: u32,
32 pub height: u32,
33 pub stride: usize,
34 pub order: Order,
35 pub data: &'a [u8],
36}
37
38#[derive(Debug)]
39pub struct PackedMut<'a> {
40 pub width: u32,
41 pub height: u32,
42 pub stride: usize,
43 pub order: Order,
44 pub data: &'a mut [u8],
45}
46
47impl<'a> Packed<'a> {
48 pub fn new(width: u32, height: u32, order: Order, data: &'a [u8]) -> Self {
49 Packed {
50 width,
51 height,
52 stride: width as usize * 4,
53 order,
54 data,
55 }
56 }
57}
58
59impl<'a> PackedMut<'a> {
60 pub fn new(width: u32, height: u32, order: Order, data: &'a mut [u8]) -> Self {
61 PackedMut {
62 width,
63 height,
64 stride: width as usize * 4,
65 order,
66 data,
67 }
68 }
69}
70
71impl Rgba {
72 pub fn packed(&self) -> Packed<'_> {
73 Packed::new(self.width, self.height, Order::Rgba, &self.data)
74 }
75
76 pub fn packed_mut(&mut self) -> PackedMut<'_> {
77 PackedMut::new(self.width, self.height, Order::Rgba, &mut self.data)
78 }
79}
80
81fn fits(width: u32, height: u32, stride: usize, len: usize) -> Result<()> {
82 sized(width, height)?;
83 let (w, h) = (width as usize, height as usize);
84 if stride < w * 4 || span(stride, h, w * 4).is_none_or(|n| len < n) {
85 return Err(Error::Invalid(format!(
86 "packed {width}x{height} with stride {stride} needs more than {len} bytes"
87 )));
88 }
89 Ok(())
90}
91
92#[derive(Clone, Copy)]
93pub(crate) struct ToPacked {
94 cy: i32,
95 ky: i32,
96 c_off: i32,
97 rv: i32,
98 gu: i32,
99 gv: i32,
100 bu: i32,
101 swap: bool,
102}
103
104#[derive(Default)]
105struct Terms {
106 r: Vec<i32>,
107 g: Vec<i32>,
108 b: Vec<i32>,
109 u: Vec<i32>,
110 v: Vec<i32>,
111 near: [Vec<i32>; 4],
112}
113
114impl Terms {
115 fn nearest<S: Sample>(&mut self, u: &[S], v: &[S], k: ToPacked, simd: Simd) {
116 match simd {
117 #[cfg(target_arch = "x86_64")]
118 Simd::Avx2 => unsafe { self.nearest_avx2(u, v, k) },
119 _ => self.nearest_plain(u, v, k),
120 }
121 }
122
123 #[cfg(target_arch = "x86_64")]
124 #[target_feature(enable = "avx2")]
125 unsafe fn nearest_avx2<S: Sample>(&mut self, u: &[S], v: &[S], k: ToPacked) {
126 self.nearest_plain(u, v, k);
127 }
128
129 #[inline(always)]
130 fn nearest_plain<S: Sample>(&mut self, u: &[S], v: &[S], k: ToPacked) {
131 let n = u.len();
132 for t in [&mut self.r, &mut self.g, &mut self.b] {
133 t.clear();
134 t.resize(n, 0);
135 }
136 for ((((r, g), b), &u), &v) in self
137 .r
138 .iter_mut()
139 .zip(self.g.iter_mut())
140 .zip(self.b.iter_mut())
141 .zip(u)
142 .zip(v)
143 {
144 let (u, v) = (u.get() - k.c_off, v.get() - k.c_off);
145 *r = k.rv * v;
146 *g = k.gu * u + k.gv * v;
147 *b = k.bu * u;
148 }
149 }
150}
151
152fn crow<S: Sample>(p: &[S], s: usize, r: usize, cw: usize, out: &mut Vec<i32>) {
153 out.clear();
154 out.extend(p[r * s..r * s + cw].iter().map(|v| v.get()));
155}
156
157pub fn yuv_to_rgba(frame: &Frame, spec: ColorSpec, out: &mut Rgba) -> Result<()> {
158 check(frame)?;
159 out.width = frame.width;
160 out.height = frame.height;
161 out.data
162 .resize(frame.width as usize * frame.height as usize * 4, 0);
163 yuv_to_packed(frame, spec, Chroma::Nearest, out.packed_mut())
164}
165
166pub fn yuv_to_packed(frame: &Frame, spec: ColorSpec, chroma: Chroma, out: PackedMut) -> Result<()> {
167 to_packed(frame, spec, chroma, out, simd())
168}
169
170pub(crate) fn to_packed(
171 frame: &Frame,
172 spec: ColorSpec,
173 chroma: Chroma,
174 out: PackedMut,
175 simd: Simd,
176) -> Result<()> {
177 check(frame)?;
178 if (out.width, out.height) != (frame.width, frame.height) {
179 return Err(Error::Invalid(format!(
180 "packed output {}x{} for a {}x{} frame",
181 out.width, out.height, frame.width, frame.height
182 )));
183 }
184 fits(out.width, out.height, out.stride, out.data.len())?;
185 let l = spec.levels(frame.format.bit_depth());
186 let m = spec.matrix.inverse();
187 let cy = fixed(255.0 / l.y_range);
188 let k = ToPacked {
189 cy,
190 ky: (1 << (FRAC - 1)) - cy * l.y_off as i32,
191 c_off: l.c_off as i32,
192 rv: fixed(255.0 * m[0][2] / l.c_range),
193 gu: fixed(255.0 * m[1][1] / l.c_range),
194 gv: fixed(255.0 * m[1][2] / l.c_range),
195 bu: fixed(255.0 * m[2][1] / l.c_range),
196 swap: out.order.swapped(),
197 };
198 match frame.format {
199 PixelFormat::Yuv420p8 => to_packed_with::<u8>(frame, k, chroma, spec.siting, out, simd),
200 PixelFormat::Yuv420p10 => to_packed_with::<u16>(frame, k, chroma, spec.siting, out, simd),
201 }
202}
203
204fn to_packed_with<S: Sample>(
205 frame: &Frame,
206 k: ToPacked,
207 chroma: Chroma,
208 siting: Siting,
209 out: PackedMut,
210 simd: Simd,
211) -> Result<()> {
212 let w = frame.width as usize;
213 let h = frame.height as usize;
214 if w == 0 || h == 0 {
215 return Ok(());
216 }
217 let b = S::BYTES;
218 let (ys, us, vs) = (
219 frame.planes[0].stride / b,
220 frame.planes[1].stride / b,
221 frame.planes[2].stride / b,
222 );
223 let (yp, up, vp) = (
224 samples::<S>(&frame.planes[0])?,
225 samples::<S>(&frame.planes[1])?,
226 samples::<S>(&frame.planes[2])?,
227 );
228 let stride = out.stride;
229 let cw = w.div_ceil(2);
230 let ch = h.div_ceil(2);
231 out.data[..stride * (h - 1) + w * 4]
232 .par_chunks_mut(stride * 2)
233 .enumerate()
234 .with_min_len(rows_job(ch))
235 .for_each_init(Terms::default, |t, (cy, pair)| match chroma {
236 Chroma::Nearest => {
237 let u = &up[cy * us..cy * us + cw];
238 let v = &vp[cy * vs..cy * vs + cw];
239 t.nearest(u, v, k, simd);
240 for (r, row) in pair.chunks_mut(stride).enumerate() {
241 let y = cy * 2 + r;
242 let luma = &yp[y * ys..y * ys + w];
243 pixels::<S, true>(luma, &t.r, &t.g, &t.b, &mut row[..w * 4], k, simd);
244 }
245 }
246 Chroma::Bilinear => {
247 for (r, row) in pair.chunks_mut(stride).enumerate() {
248 let y = cy * 2 + r;
249 let near = if r == 0 {
250 cy.saturating_sub(1)
251 } else {
252 (cy + 1).min(ch - 1)
253 };
254 let [ua, ub, va, vb] = &mut t.near;
255 crow(up, us, cy, cw, ua);
256 crow(up, us, near, cw, ub);
257 crow(vp, vs, cy, cw, va);
258 crow(vp, vs, near, cw, vb);
259 upsample(ua, ub, w, siting, &mut t.u);
260 upsample(va, vb, w, siting, &mut t.v);
261 t.r.clear();
262 t.g.clear();
263 t.b.clear();
264 let off = 16 * k.c_off;
265 for (&u, &v) in t.u.iter().zip(&t.v) {
266 let (u, v) = (u - off, v - off);
267 t.r.push((k.rv * v + 8) >> 4);
268 t.g.push((k.gu * u + k.gv * v + 8) >> 4);
269 t.b.push((k.bu * u + 8) >> 4);
270 }
271 let luma = &yp[y * ys..y * ys + w];
272 pixels::<S, false>(luma, &t.r, &t.g, &t.b, &mut row[..w * 4], k, simd);
273 }
274 }
275 });
276 Ok(())
277}
278
279fn upsample(a: &[i32], b: &[i32], w: usize, siting: Siting, out: &mut Vec<i32>) {
280 let cw = a.len();
281 out.clear();
282 out.resize(w, 0);
283 let v = |i: usize| 3 * a[i] + b[i];
284 for (x, o) in out.iter_mut().enumerate() {
285 let i = x / 2;
286 *o = match (siting, x % 2) {
287 (Siting::Left, 0) => 4 * v(i),
288 (Siting::Left, _) => 2 * (v(i) + v((i + 1).min(cw - 1))),
289 (Siting::Centre, 0) => 3 * v(i) + v(i.saturating_sub(1)),
290 (Siting::Centre, _) => 3 * v(i) + v((i + 1).min(cw - 1)),
291 };
292 }
293}
294
295#[inline(always)]
296fn pixel<S: Sample>(y: S, tr: i32, tg: i32, tb: i32, k: ToPacked, o: &mut [u8]) {
297 let y = k.cy * y.get() + k.ky;
298 let r = ((y + tr) >> FRAC).clamp(0, 255) as u8;
299 let g = ((y + tg) >> FRAC).clamp(0, 255) as u8;
300 let b = ((y + tb) >> FRAC).clamp(0, 255) as u8;
301 o.copy_from_slice(&if k.swap {
302 [b, g, r, 255]
303 } else {
304 [r, g, b, 255]
305 });
306}
307
308fn pixels<S: Sample, const DUP: bool>(
309 luma: &[S],
310 tr: &[i32],
311 tg: &[i32],
312 tb: &[i32],
313 out: &mut [u8],
314 k: ToPacked,
315 simd: Simd,
316) {
317 let w = luma.len();
318 let need = if DUP { w.div_ceil(2) } else { w };
319 assert!(tr.len() >= need && tg.len() >= need && tb.len() >= need && out.len() >= w * 4);
320 let from = match simd {
321 #[cfg(target_arch = "x86_64")]
322 Simd::Avx2 => unsafe { x86::pixels::<S, DUP>(luma, tr, tg, tb, out, k) },
323 #[cfg(target_arch = "aarch64")]
324 Simd::Neon => neon::pixels::<S, DUP>(luma, tr, tg, tb, out, k),
325 _ => 0,
326 };
327 for x in from..w {
328 let i = if DUP { x / 2 } else { x };
329 pixel(luma[x], tr[i], tg[i], tb[i], k, &mut out[x * 4..x * 4 + 4]);
330 }
331}
332
333#[derive(Clone, Copy)]
334pub(crate) struct ToYuv {
335 y: [i32; 3],
336 u: [i32; 3],
337 v: [i32; 3],
338 y_off: i32,
339 c_off: i32,
340 max: i32,
341 swap: bool,
342}
343
344pub fn rgba_to_yuv(
345 src: &Rgba,
346 spec: ColorSpec,
347 format: PixelFormat,
348 out: &mut Frame,
349) -> Result<()> {
350 packed_to_yuv(src.packed(), spec, format, out)
351}
352
353pub fn packed_to_yuv(
354 src: Packed,
355 spec: ColorSpec,
356 format: PixelFormat,
357 out: &mut Frame,
358) -> Result<()> {
359 from_packed(src, spec, format, out, simd())
360}
361
362pub(crate) fn from_packed(
363 src: Packed,
364 spec: ColorSpec,
365 format: PixelFormat,
366 out: &mut Frame,
367 simd: Simd,
368) -> Result<()> {
369 fits(src.width, src.height, src.stride, src.data.len())?;
370 let depth = format.bit_depth();
371 let l = spec.levels(depth);
372 let m = spec.matrix.forward();
373 let row = |i: usize, range: f64| {
374 [
375 fixed(m[i][0] * range / 255.0),
376 fixed(m[i][1] * range / 255.0),
377 fixed(m[i][2] * range / 255.0),
378 ]
379 };
380 let k = ToYuv {
381 y: row(0, l.y_range),
382 u: row(1, l.c_range),
383 v: row(2, l.c_range),
384 y_off: ((l.y_off as i32) << FRAC) + (1 << (FRAC - 1)),
385 c_off: ((l.c_off as i32) << FRAC) + (1 << (FRAC - 1)),
386 max: (1 << depth) - 1,
387 swap: src.order.swapped(),
388 };
389 reshape(out, format, src.width, src.height, out.pts, out.time_base);
390 if src.width == 0 || src.height == 0 {
391 return Ok(());
392 }
393 match format {
394 PixelFormat::Yuv420p8 => from_packed_with::<u8>(src, k, spec.siting, out, simd),
395 PixelFormat::Yuv420p10 => from_packed_with::<u16>(src, k, spec.siting, out, simd),
396 }
397}
398
399fn from_packed_with<S: Sample>(
400 src: Packed,
401 k: ToYuv,
402 siting: Siting,
403 out: &mut Frame,
404 simd: Simd,
405) -> Result<()> {
406 let w = src.width as usize;
407 let rows = src.height as usize;
408 let stride = src.stride;
409 let cw = w.div_ceil(2);
410 let [py, pu, pv] = &mut out.planes;
411 let (ys, us, vs) = (
412 py.stride / S::BYTES,
413 pu.stride / S::BYTES,
414 pv.stride / S::BYTES,
415 );
416 let yp = samples_mut::<S>(py)?;
417 let up = samples_mut::<S>(pu)?;
418 let vp = samples_mut::<S>(pv)?;
419 let line = |r: usize| &src.data[r * stride..r * stride + w * 4];
420 yp.par_chunks_mut(ys * 2)
421 .zip(up.par_chunks_mut(us))
422 .zip(vp.par_chunks_mut(vs))
423 .enumerate()
424 .with_min_len(rows_job(rows.div_ceil(2)))
425 .for_each_init(
426 || (vec![0i32; w], vec![0i32; w]),
427 |(su, sv), (cy, ((yrows, urow), vrow))| {
428 let top = line(cy * 2);
429 let bottom = if cy * 2 + 1 < rows {
430 line(cy * 2 + 1)
431 } else {
432 top
433 };
434 let (yt, yb) = yrows.split_at_mut(ys.min(yrows.len()));
435 luma(top, &mut yt[..w], k, simd);
436 if cy * 2 + 1 < rows {
437 luma(bottom, &mut yb[..w], k, simd);
438 }
439 sums(top, bottom, su, sv, k, simd);
440 let last = w - 1;
441 for x in 0..cw {
442 let (a, b) = (2 * x, (2 * x + 1).min(last));
443 let (u, v) = match siting {
444 Siting::Centre => (
445 (su[a] + su[b] + 4 * k.c_off) >> (FRAC + 2),
446 (sv[a] + sv[b] + 4 * k.c_off) >> (FRAC + 2),
447 ),
448 Siting::Left => {
449 let p = a.saturating_sub(1);
450 (
451 (su[p] + 2 * su[a] + su[b] + 8 * k.c_off) >> (FRAC + 3),
452 (sv[p] + 2 * sv[a] + sv[b] + 8 * k.c_off) >> (FRAC + 3),
453 )
454 }
455 };
456 urow[x] = S::put(u.clamp(0, k.max));
457 vrow[x] = S::put(v.clamp(0, k.max));
458 }
459 },
460 );
461 Ok(())
462}
463
464#[inline(always)]
465fn rgb(p: &[u8], swap: bool) -> [i32; 3] {
466 if swap {
467 [p[2] as i32, p[1] as i32, p[0] as i32]
468 } else {
469 [p[0] as i32, p[1] as i32, p[2] as i32]
470 }
471}
472
473fn luma<S: Sample>(src: &[u8], out: &mut [S], k: ToYuv, simd: Simd) {
474 let w = out.len();
475 assert!(src.len() >= w * 4);
476 let from = match simd {
477 #[cfg(target_arch = "x86_64")]
478 Simd::Avx2 => unsafe { x86::luma(src, out, k) },
479 #[cfg(target_arch = "aarch64")]
480 Simd::Neon => neon::luma(src, out, k),
481 _ => 0,
482 };
483 for x in from..w {
484 let [r, g, b] = rgb(&src[x * 4..], k.swap);
485 let y = (k.y[0] * r + k.y[1] * g + k.y[2] * b + k.y_off) >> FRAC;
486 out[x] = S::put(y.clamp(0, k.max));
487 }
488}
489
490fn sums(top: &[u8], bottom: &[u8], su: &mut [i32], sv: &mut [i32], k: ToYuv, simd: Simd) {
491 let w = su.len();
492 assert!(top.len() >= w * 4 && bottom.len() >= w * 4 && sv.len() >= w);
493 let from = match simd {
494 #[cfg(target_arch = "x86_64")]
495 Simd::Avx2 => unsafe { x86::sums(top, bottom, su, sv, k) },
496 #[cfg(target_arch = "aarch64")]
497 Simd::Neon => neon::sums(top, bottom, su, sv, k),
498 _ => 0,
499 };
500 for x in from..w {
501 let a = rgb(&top[x * 4..], k.swap);
502 let b = rgb(&bottom[x * 4..], k.swap);
503 let s = [a[0] + b[0], a[1] + b[1], a[2] + b[2]];
504 su[x] = k.u[0] * s[0] + k.u[1] * s[1] + k.u[2] * s[2];
505 sv[x] = k.v[0] * s[0] + k.v[1] * s[1] + k.v[2] * s[2];
506 }
507}
508
509#[cfg(target_arch = "x86_64")]
510mod x86 {
511 use super::{FRAC, Sample, ToPacked, ToYuv};
512 use std::arch::x86_64::*;
513
514 #[target_feature(enable = "avx2")]
515 unsafe fn luma8<S: Sample>(p: *const S) -> __m256i {
516 unsafe {
517 if S::BYTES == 1 {
518 _mm256_cvtepu8_epi32(_mm_loadl_epi64(p as *const __m128i))
519 } else {
520 _mm256_cvtepu16_epi32(_mm_loadu_si128(p as *const __m128i))
521 }
522 }
523 }
524
525 #[target_feature(enable = "avx2")]
526 fn ordered(a: __m256i, b: __m256i) -> __m256i {
527 _mm256_permute4x64_epi64::<0xD8>(_mm256_packs_epi32(a, b))
528 }
529
530 #[target_feature(enable = "avx2")]
531 pub(super) unsafe fn pixels<S: Sample, const DUP: bool>(
532 luma: &[S],
533 tr: &[i32],
534 tg: &[i32],
535 tb: &[i32],
536 out: &mut [u8],
537 k: ToPacked,
538 ) -> usize {
539 let w = luma.len();
540 let cy = _mm256_set1_epi32(k.cy);
541 let ky = _mm256_set1_epi32(k.ky);
542 let lo_idx = _mm256_setr_epi32(0, 0, 1, 1, 2, 2, 3, 3);
543 let hi_idx = _mm256_setr_epi32(4, 4, 5, 5, 6, 6, 7, 7);
544 let alpha = _mm256_set1_epi16(255);
545 let mut x = 0;
546 while x + 16 <= w {
547 unsafe {
548 let y0 = _mm256_add_epi32(_mm256_mullo_epi32(luma8(luma.as_ptr().add(x)), cy), ky);
549 let y1 =
550 _mm256_add_epi32(_mm256_mullo_epi32(luma8(luma.as_ptr().add(x + 8)), cy), ky);
551 let term = |t: &[i32]| -> (__m256i, __m256i) {
552 if DUP {
553 let v = _mm256_loadu_si256(t.as_ptr().add(x / 2) as *const __m256i);
554 (
555 _mm256_permutevar8x32_epi32(v, lo_idx),
556 _mm256_permutevar8x32_epi32(v, hi_idx),
557 )
558 } else {
559 (
560 _mm256_loadu_si256(t.as_ptr().add(x) as *const __m256i),
561 _mm256_loadu_si256(t.as_ptr().add(x + 8) as *const __m256i),
562 )
563 }
564 };
565 let channel = |t: &[i32]| {
566 let (a, b) = term(t);
567 ordered(
568 _mm256_srai_epi32::<FRAC>(_mm256_add_epi32(y0, a)),
569 _mm256_srai_epi32::<FRAC>(_mm256_add_epi32(y1, b)),
570 )
571 };
572 let (r, g, b) = (channel(tr), channel(tg), channel(tb));
573 let (first, third) = if k.swap { (b, r) } else { (r, b) };
574 let p01 = _mm256_packus_epi16(first, g);
575 let p23 = _mm256_packus_epi16(third, alpha);
576 let fg = _mm256_unpacklo_epi8(p01, _mm256_srli_si256::<8>(p01));
577 let ba = _mm256_unpacklo_epi8(p23, _mm256_srli_si256::<8>(p23));
578 let lo = _mm256_unpacklo_epi16(fg, ba);
579 let hi = _mm256_unpackhi_epi16(fg, ba);
580 let o = out.as_mut_ptr().add(x * 4) as *mut __m256i;
581 _mm256_storeu_si256(o, _mm256_permute2x128_si256::<0x20>(lo, hi));
582 _mm256_storeu_si256(o.add(1), _mm256_permute2x128_si256::<0x31>(lo, hi));
583 }
584 x += 16;
585 }
586 x
587 }
588
589 #[target_feature(enable = "avx2")]
590 fn channels(v: __m256i, swap: bool) -> [__m256i; 3] {
591 let m = _mm256_set1_epi32(0xff);
592 let c0 = _mm256_and_si256(v, m);
593 let c1 = _mm256_and_si256(_mm256_srli_epi32::<8>(v), m);
594 let c2 = _mm256_and_si256(_mm256_srli_epi32::<16>(v), m);
595 if swap { [c2, c1, c0] } else { [c0, c1, c2] }
596 }
597
598 #[target_feature(enable = "avx2")]
599 fn dot(c: [__m256i; 3], k: [i32; 3]) -> __m256i {
600 _mm256_add_epi32(
601 _mm256_add_epi32(
602 _mm256_mullo_epi32(c[0], _mm256_set1_epi32(k[0])),
603 _mm256_mullo_epi32(c[1], _mm256_set1_epi32(k[1])),
604 ),
605 _mm256_mullo_epi32(c[2], _mm256_set1_epi32(k[2])),
606 )
607 }
608
609 #[target_feature(enable = "avx2")]
610 pub(super) unsafe fn luma<S: Sample>(src: &[u8], out: &mut [S], k: ToYuv) -> usize {
611 let w = out.len();
612 let off = _mm256_set1_epi32(k.y_off);
613 let max = _mm_set1_epi16(k.max as i16);
614 let zero = _mm_setzero_si128();
615 let mut x = 0;
616 while x + 8 <= w {
617 unsafe {
618 let v = _mm256_loadu_si256(src.as_ptr().add(x * 4) as *const __m256i);
619 let y =
620 _mm256_srai_epi32::<FRAC>(_mm256_add_epi32(dot(channels(v, k.swap), k.y), off));
621 let p =
622 _mm256_castsi256_si128(_mm256_permute4x64_epi64::<8>(_mm256_packs_epi32(y, y)));
623 if S::BYTES == 1 {
624 _mm_storel_epi64(
625 out.as_mut_ptr().add(x) as *mut __m128i,
626 _mm_packus_epi16(p, p),
627 );
628 } else {
629 let p = _mm_min_epi16(_mm_max_epi16(p, zero), max);
630 _mm_storeu_si128(out.as_mut_ptr().add(x) as *mut __m128i, p);
631 }
632 }
633 x += 8;
634 }
635 x
636 }
637
638 #[target_feature(enable = "avx2")]
639 pub(super) unsafe fn sums(
640 top: &[u8],
641 bottom: &[u8],
642 su: &mut [i32],
643 sv: &mut [i32],
644 k: ToYuv,
645 ) -> usize {
646 let w = su.len();
647 let mut x = 0;
648 while x + 8 <= w {
649 unsafe {
650 let a = channels(
651 _mm256_loadu_si256(top.as_ptr().add(x * 4) as *const __m256i),
652 k.swap,
653 );
654 let b = channels(
655 _mm256_loadu_si256(bottom.as_ptr().add(x * 4) as *const __m256i),
656 k.swap,
657 );
658 let s = [
659 _mm256_add_epi32(a[0], b[0]),
660 _mm256_add_epi32(a[1], b[1]),
661 _mm256_add_epi32(a[2], b[2]),
662 ];
663 _mm256_storeu_si256(su.as_mut_ptr().add(x) as *mut __m256i, dot(s, k.u));
664 _mm256_storeu_si256(sv.as_mut_ptr().add(x) as *mut __m256i, dot(s, k.v));
665 }
666 x += 8;
667 }
668 x
669 }
670}
671
672#[cfg(target_arch = "aarch64")]
673mod neon {
674 use super::{FRAC, Sample, ToPacked, ToYuv};
675 use std::arch::aarch64::*;
676
677 unsafe fn luma16<S: Sample>(p: *const S) -> [int32x4_t; 4] {
678 unsafe {
679 let (lo, hi) = if S::BYTES == 1 {
680 let v = vld1q_u8(p as *const u8);
681 (vmovl_u8(vget_low_u8(v)), vmovl_high_u8(v))
682 } else {
683 (
684 vld1q_u16(p as *const u16),
685 vld1q_u16((p as *const u16).add(8)),
686 )
687 };
688 [
689 vreinterpretq_s32_u32(vmovl_u16(vget_low_u16(lo))),
690 vreinterpretq_s32_u32(vmovl_high_u16(lo)),
691 vreinterpretq_s32_u32(vmovl_u16(vget_low_u16(hi))),
692 vreinterpretq_s32_u32(vmovl_high_u16(hi)),
693 ]
694 }
695 }
696
697 fn narrow(v: [int32x4_t; 4]) -> uint8x16_t {
698 unsafe {
699 let a = vcombine_s16(vqmovn_s32(v[0]), vqmovn_s32(v[1]));
700 let b = vcombine_s16(vqmovn_s32(v[2]), vqmovn_s32(v[3]));
701 vcombine_u8(vqmovun_s16(a), vqmovun_s16(b))
702 }
703 }
704
705 pub(super) fn pixels<S: Sample, const DUP: bool>(
706 luma: &[S],
707 tr: &[i32],
708 tg: &[i32],
709 tb: &[i32],
710 out: &mut [u8],
711 k: ToPacked,
712 ) -> usize {
713 let w = luma.len();
714 let mut x = 0;
715 while x + 16 <= w {
716 unsafe {
717 let ky = vdupq_n_s32(k.ky);
718 let y = luma16(luma.as_ptr().add(x)).map(|v| vmlaq_n_s32(ky, v, k.cy));
719 let term = |t: &[i32]| -> [int32x4_t; 4] {
720 if DUP {
721 let a = vld1q_s32(t.as_ptr().add(x / 2));
722 let b = vld1q_s32(t.as_ptr().add(x / 2 + 4));
723 [
724 vzip1q_s32(a, a),
725 vzip2q_s32(a, a),
726 vzip1q_s32(b, b),
727 vzip2q_s32(b, b),
728 ]
729 } else {
730 [0, 4, 8, 12].map(|o| vld1q_s32(t.as_ptr().add(x + o)))
731 }
732 };
733 let channel = |t: &[i32]| {
734 let t = term(t);
735 narrow([0, 1, 2, 3].map(|i| vshrq_n_s32::<FRAC>(vaddq_s32(y[i], t[i]))))
736 };
737 let (r, g, b) = (channel(tr), channel(tg), channel(tb));
738 let (first, third) = if k.swap { (b, r) } else { (r, b) };
739 vst4q_u8(
740 out.as_mut_ptr().add(x * 4),
741 uint8x16x4_t(first, g, third, vdupq_n_u8(255)),
742 );
743 }
744 x += 16;
745 }
746 x
747 }
748
749 fn widen(v: uint8x16_t) -> [int32x4_t; 4] {
750 unsafe {
751 let lo = vmovl_u8(vget_low_u8(v));
752 let hi = vmovl_high_u8(v);
753 [
754 vreinterpretq_s32_u32(vmovl_u16(vget_low_u16(lo))),
755 vreinterpretq_s32_u32(vmovl_high_u16(lo)),
756 vreinterpretq_s32_u32(vmovl_u16(vget_low_u16(hi))),
757 vreinterpretq_s32_u32(vmovl_high_u16(hi)),
758 ]
759 }
760 }
761
762 fn dot(c: [[int32x4_t; 4]; 3], k: [i32; 3], i: usize, base: int32x4_t) -> int32x4_t {
763 unsafe {
764 let s = vmlaq_n_s32(base, c[0][i], k[0]);
765 let s = vmlaq_n_s32(s, c[1][i], k[1]);
766 vmlaq_n_s32(s, c[2][i], k[2])
767 }
768 }
769
770 pub(super) fn luma<S: Sample>(src: &[u8], out: &mut [S], k: ToYuv) -> usize {
771 let w = out.len();
772 let mut x = 0;
773 while x + 16 <= w {
774 unsafe {
775 let p = vld4q_u8(src.as_ptr().add(x * 4));
776 let (r, b) = if k.swap { (p.2, p.0) } else { (p.0, p.2) };
777 let c = [widen(r), widen(p.1), widen(b)];
778 let off = vdupq_n_s32(k.y_off);
779 let y = [0, 1, 2, 3].map(|i| vshrq_n_s32::<FRAC>(dot(c, k.y, i, off)));
780 if S::BYTES == 1 {
781 vst1q_u8(out.as_mut_ptr().add(x) as *mut u8, narrow(y));
782 } else {
783 let max = vdupq_n_u16(k.max as u16);
784 let lo = vminq_u16(vcombine_u16(vqmovun_s32(y[0]), vqmovun_s32(y[1])), max);
785 let hi = vminq_u16(vcombine_u16(vqmovun_s32(y[2]), vqmovun_s32(y[3])), max);
786 let o = out.as_mut_ptr().add(x) as *mut u16;
787 vst1q_u16(o, lo);
788 vst1q_u16(o.add(8), hi);
789 }
790 }
791 x += 16;
792 }
793 x
794 }
795
796 pub(super) fn sums(
797 top: &[u8],
798 bottom: &[u8],
799 su: &mut [i32],
800 sv: &mut [i32],
801 k: ToYuv,
802 ) -> usize {
803 let w = su.len();
804 let mut x = 0;
805 while x + 16 <= w {
806 unsafe {
807 let a = vld4q_u8(top.as_ptr().add(x * 4));
808 let b = vld4q_u8(bottom.as_ptr().add(x * 4));
809 let add = |p: uint8x16_t, q: uint8x16_t| -> [int32x4_t; 4] {
810 let lo = vaddl_u8(vget_low_u8(p), vget_low_u8(q));
811 let hi = vaddl_high_u8(p, q);
812 [
813 vreinterpretq_s32_u32(vmovl_u16(vget_low_u16(lo))),
814 vreinterpretq_s32_u32(vmovl_high_u16(lo)),
815 vreinterpretq_s32_u32(vmovl_u16(vget_low_u16(hi))),
816 vreinterpretq_s32_u32(vmovl_high_u16(hi)),
817 ]
818 };
819 let (ra, ba) = if k.swap { (a.2, a.0) } else { (a.0, a.2) };
820 let (rb, bb) = if k.swap { (b.2, b.0) } else { (b.0, b.2) };
821 let c = [add(ra, rb), add(a.1, b.1), add(ba, bb)];
822 let zero = vdupq_n_s32(0);
823 for i in 0..4 {
824 vst1q_s32(su.as_mut_ptr().add(x + 4 * i), dot(c, k.u, i, zero));
825 vst1q_s32(sv.as_mut_ptr().add(x + 4 * i), dot(c, k.v, i, zero));
826 }
827 }
828 x += 16;
829 }
830 x
831 }
832}