1use super::color::Siting;
2use super::plane::{Sample, Simd, check, plane_dims, reshape, samples, samples_mut, simd, sized};
3use crate::media::{Error, Frame, PixelFormat, Result};
4use rayon::prelude::*;
5
6const BITS: i32 = 14;
7const ONE: f64 = (1 << BITS) as f64;
8
9#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
10pub enum Kernel {
11 Bilinear,
12 Bicubic,
13 Lanczos,
14}
15
16impl Kernel {
17 fn radius(self) -> f64 {
18 match self {
19 Kernel::Bilinear => 1.0,
20 Kernel::Bicubic => 2.0,
21 Kernel::Lanczos => 3.0,
22 }
23 }
24
25 fn weight(self, x: f64) -> f64 {
26 let x = x.abs();
27 match self {
28 Kernel::Bilinear => (1.0 - x).max(0.0),
29 Kernel::Bicubic => {
30 if x < 1.0 {
31 (1.5 * x - 2.5) * x * x + 1.0
32 } else if x < 2.0 {
33 ((-0.5 * x + 2.5) * x - 4.0) * x + 2.0
34 } else {
35 0.0
36 }
37 }
38 Kernel::Lanczos => {
39 if x >= 3.0 {
40 0.0
41 } else {
42 sinc(x) * sinc(x / 3.0)
43 }
44 }
45 }
46 }
47}
48
49fn sinc(x: f64) -> f64 {
50 if x < 1e-9 {
51 1.0
52 } else {
53 let p = std::f64::consts::PI * x;
54 p.sin() / p
55 }
56}
57
58#[derive(Debug, Clone)]
59struct Taps {
60 n: usize,
61 start: Vec<usize>,
62 coef: Vec<i16>,
63 blocks: usize,
64 reach: usize,
65 packed: Vec<i16>,
66 #[cfg_attr(not(target_arch = "x86_64"), allow(dead_code))]
67 quad: Vec<i16>,
68}
69
70impl Taps {
71 fn new(src: usize, dst: usize, ratio: f64, sub: f64, offset: f64, kernel: Kernel) -> Self {
72 let widen = ratio.max(1.0);
73 let support = kernel.radius() * widen;
74 let n = ((2.0 * support).ceil() as usize).clamp(1, src);
75 let mut start = Vec::with_capacity(dst);
76 let mut coef = Vec::with_capacity(dst * n);
77 let mut w = vec![0.0f64; n];
78 for j in 0..dst {
79 let c = ((j as f64 * sub + offset + 0.5) * ratio - 0.5 - offset) / sub;
80 let first = (c - support).floor() as i64 + 1;
81 let s = first.clamp(0, (src - n) as i64);
82 w.iter_mut().for_each(|v| *v = 0.0);
83 for k in 0..n as i64 {
84 let i = first + k;
85 let at = i.clamp(0, src as i64 - 1);
86 w[(at - s) as usize] += kernel.weight((i as f64 - c) / widen);
87 }
88 let sum: f64 = w.iter().sum();
89 let mut q: Vec<i16> = w.iter().map(|v| (v / sum * ONE).round() as i16).collect();
90 let total: i32 = q.iter().map(|&v| v as i32).sum();
91 let big = (0..n).max_by_key(|&k| q[k]).unwrap_or(0);
92 q[big] = (q[big] as i32 + (1 << BITS) - total) as i16;
93 start.push(s as usize);
94 coef.extend_from_slice(&q);
95 }
96 let blocks = n.div_ceil(8);
97 let pairs = dst.div_ceil(2);
98 let mut packed = vec![0i16; pairs * blocks * 16];
99 for j in 0..dst {
100 let (pair, side) = (j / 2, j % 2);
101 for t in 0..n {
102 let (b, k) = (t / 8, t % 8);
103 packed[(pair * blocks + b) * 16 + side * 8 + k] = coef[j * n + t];
104 }
105 }
106 let reach = start.iter().max().copied().unwrap_or(0) + blocks * 8;
107 let mut quad = Vec::new();
108 if n <= 4 {
109 quad = vec![0i16; dst.div_ceil(4) * 16];
110 for j in 0..dst {
111 for t in 0..n {
112 quad[j / 4 * 16 + j % 4 * 4 + t] = coef[j * n + t];
113 }
114 }
115 }
116 Taps {
117 n,
118 start,
119 coef,
120 blocks,
121 reach,
122 packed,
123 quad,
124 }
125 }
126}
127
128#[derive(Debug, Clone)]
129struct PlaneTaps {
130 src: (usize, usize),
131 dst: (usize, usize),
132 h: Taps,
133 v: Taps,
134}
135
136impl PlaneTaps {
137 fn new(src: (u32, u32), dst: (u32, u32), i: usize, offset: f64, kernel: Kernel) -> Self {
138 let rx = src.0 as f64 / dst.0 as f64;
139 let ry = src.1 as f64 / dst.1 as f64;
140 let (sub, oy) = if i == 0 { (1.0, 0.0) } else { (2.0, 0.5) };
141 let s = plane_dims(src.0, src.1, i);
142 let d = plane_dims(dst.0, dst.1, i);
143 PlaneTaps {
144 src: s,
145 dst: d,
146 h: Taps::new(s.0, d.0, rx, sub, offset, kernel),
147 v: Taps::new(s.1, d.1, ry, sub, oy, kernel),
148 }
149 }
150}
151
152#[derive(Debug, Clone)]
153pub struct Scaler {
154 format: PixelFormat,
155 kernel: Kernel,
156 src: (u32, u32),
157 dst: (u32, u32),
158 planes: [PlaneTaps; 2],
159 simd: Simd,
160}
161
162impl Scaler {
163 pub fn new(
164 src_width: u32,
165 src_height: u32,
166 dst_width: u32,
167 dst_height: u32,
168 format: PixelFormat,
169 kernel: Kernel,
170 ) -> Result<Self> {
171 if src_width == 0 || src_height == 0 || dst_width == 0 || dst_height == 0 {
172 return Err(Error::Invalid("scale needs non-zero sizes".into()));
173 }
174 sized(src_width, src_height)?;
175 sized(dst_width, dst_height)?;
176 let src = (src_width, src_height);
177 let dst = (dst_width, dst_height);
178 Ok(Scaler {
179 format,
180 kernel,
181 src,
182 dst,
183 planes: [
184 PlaneTaps::new(src, dst, 0, 0.0, kernel),
185 PlaneTaps::new(src, dst, 1, 0.0, kernel),
186 ],
187 simd: simd(),
188 })
189 }
190
191 pub fn siting(mut self, siting: Siting) -> Self {
192 let offset = match siting {
193 Siting::Left => 0.0,
194 Siting::Centre => 0.5,
195 };
196 self.planes[1] = PlaneTaps::new(self.src, self.dst, 1, offset, self.kernel);
197 self
198 }
199
200 #[cfg(test)]
201 pub(crate) fn with_simd(mut self, simd: Simd) -> Self {
202 self.simd = simd;
203 self
204 }
205
206 pub fn run(&self, src: &Frame, dst: &mut Frame) -> Result<()> {
207 if src.format != self.format || (src.width, src.height) != self.src {
208 return Err(Error::Invalid(format!(
209 "scaler expects {:?} {}x{}, got {:?} {}x{}",
210 self.format, self.src.0, self.src.1, src.format, src.width, src.height
211 )));
212 }
213 check(src)?;
214 reshape(
215 dst,
216 self.format,
217 self.dst.0,
218 self.dst.1,
219 src.pts,
220 src.time_base,
221 );
222 let depth = self.format.bit_depth() as i32;
223 for i in 0..3 {
224 let taps = &self.planes[i.min(1)];
225 let (s, d) = (&src.planes[i], &mut dst.planes[i]);
226 let (ss, ds) = (s.stride, d.stride);
227 match self.format {
228 PixelFormat::Yuv420p8 => {
229 run_plane::<u8>(samples(s)?, ss, samples_mut(d)?, ds, taps, depth, self.simd)
230 }
231 PixelFormat::Yuv420p10 => run_plane::<u16>(
232 samples(s)?,
233 ss / 2,
234 samples_mut(d)?,
235 ds / 2,
236 taps,
237 depth,
238 self.simd,
239 ),
240 }
241 }
242 Ok(())
243 }
244}
245
246pub fn scale(frame: &Frame, width: u32, height: u32, kernel: Kernel) -> Result<Frame> {
247 let scaler = Scaler::new(
248 frame.width,
249 frame.height,
250 width,
251 height,
252 frame.format,
253 kernel,
254 )?;
255 let mut out = Frame::new(frame.format, width, height, frame.pts, frame.time_base);
256 scaler.run(frame, &mut out)?;
257 Ok(out)
258}
259
260fn run_plane<S: Sample>(
261 src: &[S],
262 src_stride: usize,
263 dst: &mut [S],
264 dst_stride: usize,
265 taps: &PlaneTaps,
266 depth: i32,
267 simd: Simd,
268) {
269 let (sw, sh) = taps.src;
270 let (dw, dh) = taps.dst;
271 if dw == 0 || dh == 0 {
272 return;
273 }
274 if (sw, sh) == (dw, dh) {
275 dst.par_chunks_mut(dst_stride)
276 .zip(src.par_chunks(src_stride))
277 .for_each(|(d, s)| d[..dw].copy_from_slice(&s[..dw]));
278 return;
279 }
280 let rows = rows_per_job(dh, sw + dw);
281 let pad = taps.h.blocks * 8 + 16;
282 dst.par_chunks_mut(dst_stride * rows)
283 .enumerate()
284 .for_each_init(
285 || vec![0i16; sw + pad],
286 |tmp, (job, chunk)| {
287 for (r, out) in chunk.chunks_mut(dst_stride).enumerate() {
288 let y = job * rows + r;
289 let n = taps.v.n;
290 let coef = &taps.v.coef[y * n..(y + 1) * n];
291 let base = &src[taps.v.start[y] * src_stride..];
292 vertical(base, src_stride, coef, sw, tmp, depth, simd);
293 horizontal(tmp, &taps.h, &mut out[..dw], depth, simd);
294 }
295 },
296 );
297}
298
299fn rows_per_job(height: usize, width: usize) -> usize {
300 let threads = rayon::current_num_threads().max(1);
301 let by_threads = height.div_ceil(threads * 4);
302 let by_work = (16384 / width.max(1)).max(1);
303 by_threads.max(by_work).clamp(1, height)
304}
305
306fn vertical<S: Sample>(
307 src: &[S],
308 stride: usize,
309 coef: &[i16],
310 width: usize,
311 tmp: &mut [i16],
312 depth: i32,
313 simd: Simd,
314) {
315 let n = coef.len();
316 assert!(src.len() >= (n - 1) * stride + width && tmp.len() >= width);
317 let from = match simd {
318 #[cfg(target_arch = "x86_64")]
319 Simd::Avx2 => unsafe { x86::vertical(src, stride, coef, width, tmp, depth) },
320 #[cfg(target_arch = "aarch64")]
321 Simd::Neon => neon::vertical(src, stride, coef, width, tmp, depth),
322 _ => 0,
323 };
324 vertical_scalar(src, stride, coef, from, width, tmp, depth);
325}
326
327fn vertical_scalar<S: Sample>(
328 src: &[S],
329 stride: usize,
330 coef: &[i16],
331 from: usize,
332 width: usize,
333 tmp: &mut [i16],
334 depth: i32,
335) {
336 let round = 1 << (depth - 1);
337 let mut acc = [0i32; 64];
338 let mut x = from;
339 while x < width {
340 let len = (width - x).min(64);
341 acc[..len].iter_mut().for_each(|a| *a = round);
342 for (t, &c) in coef.iter().enumerate() {
343 let row = &src[t * stride + x..t * stride + x + len];
344 for (a, &s) in acc[..len].iter_mut().zip(row) {
345 *a += s.get() * c as i32;
346 }
347 }
348 for (o, &a) in tmp[x..x + len].iter_mut().zip(&acc[..len]) {
349 *o = (a >> depth).clamp(i16::MIN as i32, i16::MAX as i32) as i16;
350 }
351 x += len;
352 }
353}
354
355fn horizontal<S: Sample>(tmp: &[i16], taps: &Taps, out: &mut [S], depth: i32, simd: Simd) {
356 assert!(
357 tmp.len() >= taps.reach && taps.packed.len() >= out.len().div_ceil(2) * taps.blocks * 16
358 );
359 let from = match simd {
360 #[cfg(target_arch = "x86_64")]
361 Simd::Avx2 => unsafe { x86::horizontal(tmp, taps, out, depth) },
362 #[cfg(target_arch = "aarch64")]
363 Simd::Neon => neon::horizontal(tmp, taps, out, depth),
364 _ => 0,
365 };
366 horizontal_scalar(tmp, taps, out, from, depth);
367}
368
369fn horizontal_scalar<S: Sample>(tmp: &[i16], taps: &Taps, out: &mut [S], from: usize, depth: i32) {
370 let shift = 2 * BITS - depth;
371 let round = 1 << (shift - 1);
372 let max = (1 << depth) - 1;
373 let n = taps.n;
374 for (x, o) in out.iter_mut().enumerate().skip(from) {
375 let s = taps.start[x];
376 let c = &taps.coef[x * n..(x + 1) * n];
377 let acc: i32 = tmp[s..s + n]
378 .iter()
379 .zip(c)
380 .map(|(&a, &b)| a as i32 * b as i32)
381 .sum();
382 *o = S::put(((acc + round) >> shift).clamp(0, max));
383 }
384}
385
386#[cfg(target_arch = "x86_64")]
387mod x86 {
388 use super::{BITS, Sample, Taps};
389 use std::arch::x86_64::*;
390
391 #[target_feature(enable = "avx2")]
392 unsafe fn load16<S: Sample>(p: *const S) -> __m256i {
393 unsafe {
394 if S::BYTES == 1 {
395 _mm256_cvtepu8_epi16(_mm_loadu_si128(p as *const __m128i))
396 } else {
397 _mm256_loadu_si256(p as *const __m256i)
398 }
399 }
400 }
401
402 #[target_feature(enable = "avx2")]
403 pub(super) unsafe fn vertical<S: Sample>(
404 src: &[S],
405 stride: usize,
406 coef: &[i16],
407 width: usize,
408 tmp: &mut [i16],
409 depth: i32,
410 ) -> usize {
411 let n = coef.len();
412 let base = src.as_ptr();
413 let round = _mm256_set1_epi32(1 << (depth - 1));
414 let shift = _mm_cvtsi32_si128(depth);
415 let zero = _mm256_setzero_si256();
416 let mut x = 0;
417 while x + 16 <= width {
418 let mut lo = round;
419 let mut hi = round;
420 let mut t = 0;
421 unsafe {
422 while t + 2 <= n {
423 let a = load16(base.add(t * stride + x));
424 let b = load16(base.add((t + 1) * stride + x));
425 let c =
426 _mm256_set1_epi32((coef[t] as u16 as i32) | ((coef[t + 1] as i32) << 16));
427 lo = _mm256_add_epi32(lo, _mm256_madd_epi16(_mm256_unpacklo_epi16(a, b), c));
428 hi = _mm256_add_epi32(hi, _mm256_madd_epi16(_mm256_unpackhi_epi16(a, b), c));
429 t += 2;
430 }
431 if t < n {
432 let a = load16(base.add(t * stride + x));
433 let c = _mm256_set1_epi32(coef[t] as u16 as i32);
434 lo = _mm256_add_epi32(lo, _mm256_madd_epi16(_mm256_unpacklo_epi16(a, zero), c));
435 hi = _mm256_add_epi32(hi, _mm256_madd_epi16(_mm256_unpackhi_epi16(a, zero), c));
436 }
437 lo = _mm256_sra_epi32(lo, shift);
438 hi = _mm256_sra_epi32(hi, shift);
439 _mm256_storeu_si256(
440 tmp.as_mut_ptr().add(x) as *mut __m256i,
441 _mm256_packs_epi32(lo, hi),
442 );
443 }
444 x += 16;
445 }
446 x
447 }
448
449 #[target_feature(enable = "avx2")]
450 pub(super) unsafe fn horizontal<S: Sample>(
451 tmp: &[i16],
452 taps: &Taps,
453 out: &mut [S],
454 depth: i32,
455 ) -> usize {
456 let width = out.len();
457 let blocks = taps.blocks;
458 let shift = 2 * BITS - depth;
459 let round = _mm256_set1_epi32(1 << (shift - 1));
460 let count = _mm_cvtsi32_si128(shift);
461 let max = _mm256_set1_epi32((1 << depth) - 1);
462 let zero = _mm256_setzero_si256();
463 let order = _mm256_setr_epi32(0, 4, 1, 5, 2, 6, 3, 7);
464 let t = tmp.as_ptr();
465 let mut x = 0;
466 if taps.n <= 4 && taps.quad.len() >= width.div_ceil(4) * 16 {
467 let order = _mm256_setr_epi32(0, 1, 4, 5, 2, 3, 6, 7);
468 let q = taps.quad.as_ptr();
469 let at = |p: usize| unsafe { (t.add(taps.start[p]) as *const i64).read_unaligned() };
470 while x + 8 <= width {
471 let a = _mm256_setr_epi64x(at(x), at(x + 1), at(x + 2), at(x + 3));
472 let b = _mm256_setr_epi64x(at(x + 4), at(x + 5), at(x + 6), at(x + 7));
473 let (wa, wb) = unsafe {
474 (
475 _mm256_loadu_si256(q.add(x / 4 * 16) as *const __m256i),
476 _mm256_loadu_si256(q.add(x / 4 * 16 + 16) as *const __m256i),
477 )
478 };
479 let h = _mm256_hadd_epi32(_mm256_madd_epi16(a, wa), _mm256_madd_epi16(b, wb));
480 let h = _mm256_permutevar8x32_epi32(h, order);
481 let v = _mm256_sra_epi32(_mm256_add_epi32(h, round), count);
482 let v = _mm256_min_epi32(_mm256_max_epi32(v, zero), max);
483 let v = _mm256_castsi256_si128(_mm256_permute4x64_epi64::<8>(_mm256_packus_epi32(
484 v, v,
485 )));
486 unsafe {
487 let o = out.as_mut_ptr().add(x);
488 if S::BYTES == 1 {
489 _mm_storel_epi64(o as *mut __m128i, _mm_packus_epi16(v, v));
490 } else {
491 _mm_storeu_si128(o as *mut __m128i, v);
492 }
493 }
494 x += 8;
495 }
496 return x;
497 }
498 while x + 8 <= width {
499 let mut r = [zero; 4];
500 for (k, acc) in r.iter_mut().enumerate() {
501 let p = x + 2 * k;
502 let s0 = taps.start[p];
503 let s1 = taps.start[p + 1];
504 let c = unsafe { taps.packed.as_ptr().add(p / 2 * blocks * 16) };
505 for b in 0..blocks {
506 unsafe {
507 let a = _mm256_inserti128_si256::<1>(
508 _mm256_castsi128_si256(_mm_loadu_si128(
509 t.add(s0 + 8 * b) as *const __m128i
510 )),
511 _mm_loadu_si128(t.add(s1 + 8 * b) as *const __m128i),
512 );
513 let w = _mm256_loadu_si256(c.add(16 * b) as *const __m256i);
514 *acc = _mm256_add_epi32(*acc, _mm256_madd_epi16(a, w));
515 }
516 }
517 }
518 let h01 = _mm256_hadd_epi32(r[0], r[1]);
519 let h23 = _mm256_hadd_epi32(r[2], r[3]);
520 let h = _mm256_permutevar8x32_epi32(_mm256_hadd_epi32(h01, h23), order);
521 let v = _mm256_sra_epi32(_mm256_add_epi32(h, round), count);
522 let v = _mm256_min_epi32(_mm256_max_epi32(v, zero), max);
523 let v =
524 _mm256_castsi256_si128(_mm256_permute4x64_epi64::<8>(_mm256_packus_epi32(v, v)));
525 unsafe {
526 let o = out.as_mut_ptr().add(x);
527 if S::BYTES == 1 {
528 _mm_storel_epi64(o as *mut __m128i, _mm_packus_epi16(v, v));
529 } else {
530 _mm_storeu_si128(o as *mut __m128i, v);
531 }
532 }
533 x += 8;
534 }
535 x
536 }
537}
538
539#[cfg(target_arch = "aarch64")]
540mod neon {
541 use super::{BITS, Sample, Taps};
542 use std::arch::aarch64::*;
543
544 unsafe fn load8<S: Sample>(p: *const S) -> int16x8_t {
545 unsafe {
546 if S::BYTES == 1 {
547 vreinterpretq_s16_u16(vmovl_u8(vld1_u8(p as *const u8)))
548 } else {
549 vreinterpretq_s16_u16(vld1q_u16(p as *const u16))
550 }
551 }
552 }
553
554 pub(super) fn vertical<S: Sample>(
555 src: &[S],
556 stride: usize,
557 coef: &[i16],
558 width: usize,
559 tmp: &mut [i16],
560 depth: i32,
561 ) -> usize {
562 let base = src.as_ptr();
563 let mut x = 0;
564 while x + 8 <= width {
565 unsafe {
566 let mut lo = vdupq_n_s32(1 << (depth - 1));
567 let mut hi = lo;
568 for (t, &c) in coef.iter().enumerate() {
569 let row = load8(base.add(t * stride + x));
570 lo = vmlal_n_s16(lo, vget_low_s16(row), c);
571 hi = vmlal_high_n_s16(hi, row, c);
572 }
573 let shift = vdupq_n_s32(-depth);
574 let v = vcombine_s16(
575 vqmovn_s32(vshlq_s32(lo, shift)),
576 vqmovn_s32(vshlq_s32(hi, shift)),
577 );
578 vst1q_s16(tmp.as_mut_ptr().add(x), v);
579 }
580 x += 8;
581 }
582 x
583 }
584
585 pub(super) fn horizontal<S: Sample>(
586 tmp: &[i16],
587 taps: &Taps,
588 out: &mut [S],
589 depth: i32,
590 ) -> usize {
591 let shift = 2 * BITS - depth;
592 let round = 1 << (shift - 1);
593 let max = (1 << depth) - 1;
594 let blocks = taps.blocks;
595 for (p, o) in out.iter_mut().enumerate() {
596 let s = taps.start[p];
597 let c = &taps.packed[(p / 2 * blocks) * 16 + (p % 2) * 8..];
598 let sum = unsafe {
599 let mut acc = vdupq_n_s32(0);
600 for b in 0..blocks {
601 let v = vld1q_s16(tmp.as_ptr().add(s + 8 * b));
602 let w = vld1q_s16(c.as_ptr().add(16 * b));
603 acc = vmlal_s16(acc, vget_low_s16(v), vget_low_s16(w));
604 acc = vmlal_high_s16(acc, v, w);
605 }
606 vaddvq_s32(acc)
607 };
608 *o = S::put(((sum + round) >> shift).clamp(0, max));
609 }
610 out.len()
611 }
612}