use core::arch::aarch64::*;
use crate::transform::{DST4, inv_transform_into_scalar};
#[inline]
#[target_feature(enable = "neon")]
fn load_s32x4(src: &[i32]) -> int32x4_t {
debug_assert!(src.len() >= 4);
unsafe { vld1q_s32(src.as_ptr()) }
}
#[inline]
#[target_feature(enable = "neon")]
fn load_rows4_s32x4(src: &[i32], stride: usize, x: usize) -> int32x4_t {
debug_assert!(src.len() > 3 * stride + x);
let lanes = [
src[x],
src[stride + x],
src[2 * stride + x],
src[3 * stride + x],
];
unsafe { vld1q_s32(lanes.as_ptr()) }
}
#[inline]
#[target_feature(enable = "neon")]
fn store_s32x4(dst: &mut [i32], v: int32x4_t) {
debug_assert!(dst.len() >= 4);
unsafe { vst1q_s32(dst.as_mut_ptr(), v) }
}
#[inline]
#[target_feature(enable = "neon")]
fn store_rows4_s32x4(dst: &mut [i32], stride: usize, x: usize, v: int32x4_t) {
debug_assert!(dst.len() > 3 * stride + x);
let mut lanes = [0i32; 4];
unsafe { vst1q_s32(lanes.as_mut_ptr(), v) };
dst[x] = lanes[0];
dst[stride + x] = lanes[1];
dst[2 * stride + x] = lanes[2];
dst[3 * stride + x] = lanes[3];
}
#[inline]
#[target_feature(enable = "neon")]
fn zero() -> int32x4_t {
vdupq_n_s32(0)
}
#[inline]
#[target_feature(enable = "neon")]
fn add(a: int32x4_t, b: int32x4_t) -> int32x4_t {
vaddq_s32(a, b)
}
#[inline]
#[target_feature(enable = "neon")]
fn sub(a: int32x4_t, b: int32x4_t) -> int32x4_t {
vsubq_s32(a, b)
}
#[inline]
#[target_feature(enable = "neon")]
fn mul_const(v: int32x4_t, c: i32) -> int32x4_t {
vmulq_s32(v, vdupq_n_s32(c))
}
#[inline]
#[target_feature(enable = "neon")]
fn madd_const(acc: int32x4_t, v: int32x4_t, c: i32) -> int32x4_t {
add(acc, mul_const(v, c))
}
macro_rules! lin_s32x4 {
($v0:expr, $k0:expr $(, $v:expr, $k:expr)+ $(,)?) => {{
let mut acc = mul_const($v0, $k0);
$(
acc = madd_const(acc, $v, $k);
)+
acc
}};
}
#[inline]
#[target_feature(enable = "neon")]
fn round_shift_s32x4(v: int32x4_t, add: i32, shift: i32) -> int32x4_t {
vshlq_s32(vaddq_s32(v, vdupq_n_s32(add)), vdupq_n_s32(-shift))
}
#[inline]
#[target_feature(enable = "neon")]
fn round_shift_clip_i16_s32x4(v: int32x4_t, add: i32, shift: i32) -> int32x4_t {
let v = round_shift_s32x4(v, add, shift);
vmaxq_s32(vminq_s32(v, vdupq_n_s32(32767)), vdupq_n_s32(-32768))
}
#[inline]
#[target_feature(enable = "neon")]
fn idct_odd_16_s32x4(c: [int32x4_t; 16]) -> [int32x4_t; 8] {
let c1 = c[1];
let c3 = c[3];
let c5 = c[5];
let c7 = c[7];
let c9 = c[9];
let c11 = c[11];
let c13 = c[13];
let c15 = c[15];
[
lin_s32x4!(
c1, 90, c3, 87, c5, 80, c7, 70, c9, 57, c11, 43, c13, 25, c15, 9
),
lin_s32x4!(
c1, 87, c3, 57, c5, 9, c7, -43, c9, -80, c11, -90, c13, -70, c15, -25
),
lin_s32x4!(
c1, 80, c3, 9, c5, -70, c7, -87, c9, -25, c11, 57, c13, 90, c15, 43
),
lin_s32x4!(
c1, 70, c3, -43, c5, -87, c7, 9, c9, 90, c11, 25, c13, -80, c15, -57
),
lin_s32x4!(
c1, 57, c3, -80, c5, -25, c7, 90, c9, -9, c11, -87, c13, 43, c15, 70
),
lin_s32x4!(
c1, 43, c3, -90, c5, 57, c7, 25, c9, -87, c11, 70, c13, 9, c15, -80
),
lin_s32x4!(
c1, 25, c3, -70, c5, 90, c7, -80, c9, 43, c11, 9, c13, -57, c15, 87
),
lin_s32x4!(
c1, 9, c3, -25, c5, 43, c7, -57, c9, 70, c11, -80, c13, 87, c15, -90
),
]
}
#[inline]
#[target_feature(enable = "neon")]
fn idct_odd_32_s32x4(c: [int32x4_t; 32]) -> [int32x4_t; 16] {
let c1 = c[1];
let c3 = c[3];
let c5 = c[5];
let c7 = c[7];
let c9 = c[9];
let c11 = c[11];
let c13 = c[13];
let c15 = c[15];
let c17 = c[17];
let c19 = c[19];
let c21 = c[21];
let c23 = c[23];
let c25 = c[25];
let c27 = c[27];
let c29 = c[29];
let c31 = c[31];
[
lin_s32x4!(
c1, 90, c3, 90, c5, 88, c7, 85, c9, 82, c11, 78, c13, 73, c15, 67, c17, 61, c19, 54,
c21, 46, c23, 38, c25, 31, c27, 22, c29, 13, c31, 4
),
lin_s32x4!(
c1, 90, c3, 82, c5, 67, c7, 46, c9, 22, c11, -4, c13, -31, c15, -54, c17, -73, c19,
-85, c21, -90, c23, -88, c25, -78, c27, -61, c29, -38, c31, -13
),
lin_s32x4!(
c1, 88, c3, 67, c5, 31, c7, -13, c9, -54, c11, -82, c13, -90, c15, -78, c17, -46, c19,
-4, c21, 38, c23, 73, c25, 90, c27, 85, c29, 61, c31, 22
),
lin_s32x4!(
c1, 85, c3, 46, c5, -13, c7, -67, c9, -90, c11, -73, c13, -22, c15, 38, c17, 82, c19,
88, c21, 54, c23, -4, c25, -61, c27, -90, c29, -78, c31, -31
),
lin_s32x4!(
c1, 82, c3, 22, c5, -54, c7, -90, c9, -61, c11, 13, c13, 78, c15, 85, c17, 31, c19,
-46, c21, -90, c23, -67, c25, 4, c27, 73, c29, 88, c31, 38
),
lin_s32x4!(
c1, 78, c3, -4, c5, -82, c7, -73, c9, 13, c11, 85, c13, 67, c15, -22, c17, -88, c19,
-61, c21, 31, c23, 90, c25, 54, c27, -38, c29, -90, c31, -46
),
lin_s32x4!(
c1, 73, c3, -31, c5, -90, c7, -22, c9, 78, c11, 67, c13, -38, c15, -90, c17, -13, c19,
82, c21, 61, c23, -46, c25, -88, c27, -4, c29, 85, c31, 54
),
lin_s32x4!(
c1, 67, c3, -54, c5, -78, c7, 38, c9, 85, c11, -22, c13, -90, c15, 4, c17, 90, c19, 13,
c21, -88, c23, -31, c25, 82, c27, 46, c29, -73, c31, -61
),
lin_s32x4!(
c1, 61, c3, -73, c5, -46, c7, 82, c9, 31, c11, -88, c13, -13, c15, 90, c17, -4, c19,
-90, c21, 22, c23, 85, c25, -38, c27, -78, c29, 54, c31, 67
),
lin_s32x4!(
c1, 54, c3, -85, c5, -4, c7, 88, c9, -46, c11, -61, c13, 82, c15, 13, c17, -90, c19,
38, c21, 67, c23, -78, c25, -22, c27, 90, c29, -31, c31, -73
),
lin_s32x4!(
c1, 46, c3, -90, c5, 38, c7, 54, c9, -90, c11, 31, c13, 61, c15, -88, c17, 22, c19, 67,
c21, -85, c23, 13, c25, 73, c27, -82, c29, 4, c31, 78
),
lin_s32x4!(
c1, 38, c3, -88, c5, 73, c7, -4, c9, -67, c11, 90, c13, -46, c15, -31, c17, 85, c19,
-78, c21, 13, c23, 61, c25, -90, c27, 54, c29, 22, c31, -82
),
lin_s32x4!(
c1, 31, c3, -78, c5, 90, c7, -61, c9, 4, c11, 54, c13, -88, c15, 82, c17, -38, c19,
-22, c21, 73, c23, -90, c25, 67, c27, -13, c29, -46, c31, 85
),
lin_s32x4!(
c1, 22, c3, -61, c5, 85, c7, -90, c9, 73, c11, -38, c13, -4, c15, 46, c17, -78, c19,
90, c21, -82, c23, 54, c25, -13, c27, -31, c29, 67, c31, -88
),
lin_s32x4!(
c1, 13, c3, -38, c5, 61, c7, -78, c9, 88, c11, -90, c13, 85, c15, -73, c17, 54, c19,
-31, c21, 4, c23, 22, c25, -46, c27, 67, c29, -82, c31, 90
),
lin_s32x4!(
c1, 4, c3, -13, c5, 22, c7, -31, c9, 38, c11, -46, c13, 54, c15, -61, c17, 67, c19,
-73, c21, 78, c23, -82, c25, 85, c27, -88, c29, 90, c31, -90
),
]
}
#[inline]
#[target_feature(enable = "neon")]
fn idct_raw_4_s32x4(c: [int32x4_t; 4]) -> [int32x4_t; 4] {
let e0 = mul_const(add(c[0], c[2]), 64);
let e1 = mul_const(sub(c[0], c[2]), 64);
let o0 = add(mul_const(c[1], 83), mul_const(c[3], 36));
let o1 = sub(mul_const(c[1], 36), mul_const(c[3], 83));
[add(e0, o0), add(e1, o1), sub(e1, o1), sub(e0, o0)]
}
#[inline]
#[target_feature(enable = "neon")]
fn idct_raw_8_s32x4(c: [int32x4_t; 8]) -> [int32x4_t; 8] {
let ee = idct_raw_4_s32x4([c[0], c[2], c[4], c[6]]);
let c1 = c[1];
let c3 = c[3];
let c5 = c[5];
let c7 = c[7];
let o0 = add(
add(mul_const(c1, 89), mul_const(c3, 75)),
add(mul_const(c5, 50), mul_const(c7, 18)),
);
let o1 = sub(
sub(mul_const(c1, 75), mul_const(c3, 18)),
add(mul_const(c5, 89), mul_const(c7, 50)),
);
let o2 = add(
sub(mul_const(c1, 50), mul_const(c3, 89)),
add(mul_const(c5, 18), mul_const(c7, 75)),
);
let o3 = sub(
add(mul_const(c1, 18), mul_const(c5, 75)),
add(mul_const(c3, 50), mul_const(c7, 89)),
);
let oo = [o0, o1, o2, o3];
let mut out = [zero(); 8];
for k in 0..4 {
out[k] = add(ee[k], oo[k]);
out[7 - k] = sub(ee[k], oo[k]);
}
out
}
#[inline]
#[target_feature(enable = "neon")]
fn idct_raw_16_s32x4(c: [int32x4_t; 16]) -> [int32x4_t; 16] {
let ee = idct_raw_8_s32x4(std::array::from_fn(|j| c[2 * j]));
let oo = idct_odd_16_s32x4(c);
let mut out = [zero(); 16];
for k in 0..8 {
out[k] = add(ee[k], oo[k]);
out[15 - k] = sub(ee[k], oo[k]);
}
out
}
#[inline]
#[target_feature(enable = "neon")]
fn idct_raw_32_s32x4(c: [int32x4_t; 32]) -> [int32x4_t; 32] {
let ee = idct_raw_16_s32x4(std::array::from_fn(|j| c[2 * j]));
let oo = idct_odd_32_s32x4(c);
let mut out = [zero(); 32];
for k in 0..16 {
out[k] = add(ee[k], oo[k]);
out[31 - k] = sub(ee[k], oo[k]);
}
out
}
#[inline]
#[target_feature(enable = "neon")]
fn idct_raw_s32x4<const N: usize>(c: [int32x4_t; N]) -> [int32x4_t; N] {
debug_assert!(N == 4 || N == 8 || N == 16 || N == 32);
match N {
4 => {
let src = [c[0], c[1], c[2], c[3]];
let r = idct_raw_4_s32x4(src);
std::array::from_fn(|i| r[i])
}
8 => {
let src = [c[0], c[1], c[2], c[3], c[4], c[5], c[6], c[7]];
let r = idct_raw_8_s32x4(src);
std::array::from_fn(|i| r[i])
}
16 => {
let src = [
c[0], c[1], c[2], c[3], c[4], c[5], c[6], c[7], c[8], c[9], c[10], c[11], c[12],
c[13], c[14], c[15],
];
let r = idct_raw_16_s32x4(src);
std::array::from_fn(|i| r[i])
}
32 => {
let src = [
c[0], c[1], c[2], c[3], c[4], c[5], c[6], c[7], c[8], c[9], c[10], c[11], c[12],
c[13], c[14], c[15], c[16], c[17], c[18], c[19], c[20], c[21], c[22], c[23], c[24],
c[25], c[26], c[27], c[28], c[29], c[30], c[31],
];
let r = idct_raw_32_s32x4(src);
std::array::from_fn(|i| r[i])
}
_ => unreachable!(),
}
}
#[target_feature(enable = "neon")]
fn inv_dct_n_into_neon<const N: usize>(coeff: &[i32], bit_depth: u8, out: &mut [i32]) {
debug_assert!(N == 4 || N == 8 || N == 16 || N == 32);
debug_assert!(coeff.len() >= N * N);
debug_assert!(out.len() >= N * N);
let shift1 = 7i32;
let add1 = 1i32 << (shift1 - 1);
let shift2 = 20 - bit_depth as i32;
let add2 = 1i32 << (shift2 - 1);
let mut tmp = [0i32; 32 * 32];
for c in (0..N).step_by(4) {
let src = std::array::from_fn(|k| load_s32x4(&coeff[k * N + c..]));
let raw = idct_raw_s32x4::<N>(src);
for (m, raw) in raw.iter().copied().enumerate() {
let v = round_shift_clip_i16_s32x4(raw, add1, shift1);
store_s32x4(&mut tmp[m * N + c..], v);
}
}
for r in (0..N).step_by(4) {
let src = std::array::from_fn(|k| load_rows4_s32x4(&tmp[r * N..], N, k));
let raw = idct_raw_s32x4::<N>(src);
for (x, raw) in raw.iter().copied().enumerate() {
let v = round_shift_s32x4(raw, add2, shift2);
store_rows4_s32x4(&mut out[r * N..], N, x, v);
}
}
}
#[target_feature(enable = "neon")]
fn inv_transform_dst4_into_neon(coeff: &[i32], bit_depth: u8, out: &mut [i32]) {
debug_assert!(coeff.len() >= 16);
debug_assert!(out.len() >= 16);
let shift1 = 7i32;
let add1 = 1i32 << (shift1 - 1);
let shift2 = 20 - bit_depth as i32;
let add2 = 1i32 << (shift2 - 1);
let mut tmp = [0i32; 16];
let mut acc = [zero(); 4];
for k in 0..4 {
let ck = load_s32x4(&coeff[k * 4..]);
for (acc, &tm) in acc.iter_mut().zip(DST4[k].iter()) {
*acc = madd_const(*acc, ck, tm);
}
}
for (m, acc) in acc.iter().copied().enumerate() {
let v = round_shift_clip_i16_s32x4(acc, add1, shift1);
store_s32x4(&mut tmp[m * 4..], v);
}
let src: [int32x4_t; 4] = std::array::from_fn(|k| load_rows4_s32x4(&tmp, 4, k));
for x in 0..4 {
let mut acc = zero();
for (rk, trow) in src.iter().copied().zip(DST4.iter()) {
acc = madd_const(acc, rk, trow[x]);
}
let v = round_shift_s32x4(acc, add2, shift2);
store_rows4_s32x4(out, 4, x, v);
}
}
pub(crate) fn inv_transform_into_neon(coeff: &[i32], n: usize, bit_depth: u8, out: &mut [i32]) {
unsafe {
match n {
4 => inv_dct_n_into_neon::<4>(coeff, bit_depth, out),
8 => inv_dct_n_into_neon::<8>(coeff, bit_depth, out),
16 => inv_dct_n_into_neon::<16>(coeff, bit_depth, out),
32 => inv_dct_n_into_neon::<32>(coeff, bit_depth, out),
_ => inv_transform_into_scalar(coeff, n, bit_depth, out),
}
}
}
pub(crate) fn inv_transform_dst_into_neon(coeff: &[i32], bit_depth: u8, out: &mut [i32]) {
unsafe { inv_transform_dst4_into_neon(coeff, bit_depth, out) }
}