use crate::dct::{INV_WC4, INV_WC8, INV_WC16, INV_WC32, WC4, WC8, WC16, WC32};
use std::arch::x86_64::*;
#[inline]
#[target_feature(enable = "avx2")]
fn transpose_8x8(c: &mut [__m256; 8]) {
let t0 = _mm256_unpacklo_ps(c[0], c[1]);
let t1 = _mm256_unpackhi_ps(c[0], c[1]);
let t2 = _mm256_unpacklo_ps(c[2], c[3]);
let t3 = _mm256_unpackhi_ps(c[2], c[3]);
let t4 = _mm256_unpacklo_ps(c[4], c[5]);
let t5 = _mm256_unpackhi_ps(c[4], c[5]);
let t6 = _mm256_unpacklo_ps(c[6], c[7]);
let t7 = _mm256_unpackhi_ps(c[6], c[7]);
let s0 = _mm256_castpd_ps(_mm256_unpacklo_pd(
_mm256_castps_pd(t0),
_mm256_castps_pd(t2),
));
let s1 = _mm256_castpd_ps(_mm256_unpackhi_pd(
_mm256_castps_pd(t0),
_mm256_castps_pd(t2),
));
let s2 = _mm256_castpd_ps(_mm256_unpacklo_pd(
_mm256_castps_pd(t1),
_mm256_castps_pd(t3),
));
let s3 = _mm256_castpd_ps(_mm256_unpackhi_pd(
_mm256_castps_pd(t1),
_mm256_castps_pd(t3),
));
let s4 = _mm256_castpd_ps(_mm256_unpacklo_pd(
_mm256_castps_pd(t4),
_mm256_castps_pd(t6),
));
let s5 = _mm256_castpd_ps(_mm256_unpackhi_pd(
_mm256_castps_pd(t4),
_mm256_castps_pd(t6),
));
let s6 = _mm256_castpd_ps(_mm256_unpacklo_pd(
_mm256_castps_pd(t5),
_mm256_castps_pd(t7),
));
let s7 = _mm256_castpd_ps(_mm256_unpackhi_pd(
_mm256_castps_pd(t5),
_mm256_castps_pd(t7),
));
c[0] = _mm256_permute2f128_ps(s0, s4, 0x20);
c[1] = _mm256_permute2f128_ps(s1, s5, 0x20);
c[2] = _mm256_permute2f128_ps(s2, s6, 0x20);
c[3] = _mm256_permute2f128_ps(s3, s7, 0x20);
c[4] = _mm256_permute2f128_ps(s0, s4, 0x31);
c[5] = _mm256_permute2f128_ps(s1, s5, 0x31);
c[6] = _mm256_permute2f128_ps(s2, s6, 0x31);
c[7] = _mm256_permute2f128_ps(s3, s7, 0x31);
}
#[inline]
#[target_feature(enable = "avx2")]
fn load(input: &[f32; 64]) -> [__m256; 8] {
unsafe { std::array::from_fn(|i| _mm256_loadu_ps(input[i * 8..].as_ptr())) }
}
#[inline]
#[target_feature(enable = "avx2,fma")]
fn dct1d_8_flat(c: &mut [__m256; 8]) {
let e0 = _mm256_add_ps(c[0], c[7]);
let e1 = _mm256_add_ps(c[1], c[6]);
let e2 = _mm256_add_ps(c[2], c[5]);
let e3 = _mm256_add_ps(c[3], c[4]);
let o0 = _mm256_mul_ps(_mm256_sub_ps(c[0], c[7]), _mm256_set1_ps(WC8[0]));
let o1 = _mm256_mul_ps(_mm256_sub_ps(c[1], c[6]), _mm256_set1_ps(WC8[1]));
let o2 = _mm256_mul_ps(_mm256_sub_ps(c[2], c[5]), _mm256_set1_ps(WC8[2]));
let o3 = _mm256_mul_ps(_mm256_sub_ps(c[3], c[4]), _mm256_set1_ps(WC8[3]));
let et0 = _mm256_add_ps(e0, e3);
let et1 = _mm256_add_ps(e1, e2);
let esum = _mm256_add_ps(et0, et1);
let ediff = _mm256_sub_ps(et0, et1);
let et2 = _mm256_mul_ps(_mm256_sub_ps(e0, e3), _mm256_set1_ps(WC4[0]));
let et3 = _mm256_mul_ps(_mm256_sub_ps(e1, e2), _mm256_set1_ps(WC4[1]));
let et2p = _mm256_add_ps(et2, et3);
let et3p = _mm256_sub_ps(et2, et3);
let et2pp = _mm256_fmadd_ps(et2p, _mm256_set1_ps(std::f32::consts::SQRT_2), et3p);
let evens = [esum, et2pp, ediff, et3p];
let ot0 = _mm256_add_ps(o0, o3);
let ot1 = _mm256_add_ps(o1, o2);
let osum = _mm256_add_ps(ot0, ot1);
let odiff = _mm256_sub_ps(ot0, ot1);
let ot2 = _mm256_mul_ps(_mm256_sub_ps(o0, o3), _mm256_set1_ps(WC4[0]));
let ot3 = _mm256_mul_ps(_mm256_sub_ps(o1, o2), _mm256_set1_ps(WC4[1]));
let ot2p = _mm256_add_ps(ot2, ot3);
let ot3p = _mm256_sub_ps(ot2, ot3);
let ot2pp = _mm256_fmadd_ps(ot2p, _mm256_set1_ps(std::f32::consts::SQRT_2), ot3p);
let mut odds = [osum, ot2pp, odiff, ot3p];
odds[0] = _mm256_fmadd_ps(odds[0], _mm256_set1_ps(std::f32::consts::SQRT_2), odds[1]);
odds[1] = _mm256_add_ps(odds[1], odds[2]);
odds[2] = _mm256_add_ps(odds[2], odds[3]);
c[0] = evens[0];
c[1] = odds[0];
c[2] = evens[1];
c[3] = odds[1];
c[4] = evens[2];
c[5] = odds[2];
c[6] = evens[3];
c[7] = odds[3];
}
#[target_feature(enable = "avx2,fma")]
pub(crate) fn dct8x8_avx2(input: &[f32; 64], output: &mut [f32; 64]) {
let mut rows = load(input);
dct1d_8_flat(&mut rows);
transpose_8x8(&mut rows);
dct1d_8_flat(&mut rows);
let scale = _mm256_set1_ps(1.0 / 64.0);
for (k, row) in rows.iter().enumerate() {
unsafe {
_mm256_storeu_ps(output[k * 8..].as_mut_ptr(), _mm256_mul_ps(*row, scale));
}
}
}
#[inline]
#[target_feature(enable = "avx2,fma")]
fn dct1d_16_flat(c: &mut [__m256; 16]) {
let mut evens = [
_mm256_add_ps(c[0], c[15]),
_mm256_add_ps(c[1], c[14]),
_mm256_add_ps(c[2], c[13]),
_mm256_add_ps(c[3], c[12]),
_mm256_add_ps(c[4], c[11]),
_mm256_add_ps(c[5], c[10]),
_mm256_add_ps(c[6], c[9]),
_mm256_add_ps(c[7], c[8]),
];
let mut odds = [
_mm256_mul_ps(_mm256_sub_ps(c[0], c[15]), _mm256_set1_ps(WC16[0])),
_mm256_mul_ps(_mm256_sub_ps(c[1], c[14]), _mm256_set1_ps(WC16[1])),
_mm256_mul_ps(_mm256_sub_ps(c[2], c[13]), _mm256_set1_ps(WC16[2])),
_mm256_mul_ps(_mm256_sub_ps(c[3], c[12]), _mm256_set1_ps(WC16[3])),
_mm256_mul_ps(_mm256_sub_ps(c[4], c[11]), _mm256_set1_ps(WC16[4])),
_mm256_mul_ps(_mm256_sub_ps(c[5], c[10]), _mm256_set1_ps(WC16[5])),
_mm256_mul_ps(_mm256_sub_ps(c[6], c[9]), _mm256_set1_ps(WC16[6])),
_mm256_mul_ps(_mm256_sub_ps(c[7], c[8]), _mm256_set1_ps(WC16[7])),
];
dct1d_8_flat(&mut evens);
dct1d_8_flat(&mut odds);
odds[0] = _mm256_fmadd_ps(odds[0], _mm256_set1_ps(std::f32::consts::SQRT_2), odds[1]);
odds[1] = _mm256_add_ps(odds[1], odds[2]);
odds[2] = _mm256_add_ps(odds[2], odds[3]);
odds[3] = _mm256_add_ps(odds[3], odds[4]);
odds[4] = _mm256_add_ps(odds[4], odds[5]);
odds[5] = _mm256_add_ps(odds[5], odds[6]);
odds[6] = _mm256_add_ps(odds[6], odds[7]);
c[0] = evens[0];
c[1] = odds[0];
c[2] = evens[1];
c[3] = odds[1];
c[4] = evens[2];
c[5] = odds[2];
c[6] = evens[3];
c[7] = odds[3];
c[8] = evens[4];
c[9] = odds[4];
c[10] = evens[5];
c[11] = odds[5];
c[12] = evens[6];
c[13] = odds[6];
c[14] = evens[7];
c[15] = odds[7];
}
#[target_feature(enable = "avx2,fma")]
pub(crate) fn dct8x16_avx2(input: &[f32; 128], output: &mut [f32; 128]) {
let mut rows_lo: [__m256; 8] =
std::array::from_fn(|k| unsafe { _mm256_loadu_ps(input[k * 16..].as_ptr()) });
let mut rows_hi: [__m256; 8] =
std::array::from_fn(|k| unsafe { _mm256_loadu_ps(input[k * 16 + 8..].as_ptr()) });
transpose_8x8(&mut rows_lo);
transpose_8x8(&mut rows_hi);
let mut c = [_mm256_undefined_ps(); 16];
c[0..8].copy_from_slice(&rows_lo);
c[8..16].copy_from_slice(&rows_hi);
dct1d_16_flat(&mut c);
let mut cl: [__m256; 8] = c[0..8].try_into().unwrap();
let mut cr: [__m256; 8] = c[8..16].try_into().unwrap();
transpose_8x8(&mut cl);
transpose_8x8(&mut cr);
dct1d_8_flat(&mut cl);
dct1d_8_flat(&mut cr);
let scale = _mm256_set1_ps(1.0 / 128.0);
for m in 0..8 {
let base = &mut output[m * 16..];
unsafe {
_mm256_storeu_ps(base.as_mut_ptr(), _mm256_mul_ps(cl[m], scale));
_mm256_storeu_ps(base[8..].as_mut_ptr(), _mm256_mul_ps(cr[m], scale));
}
}
}
#[target_feature(enable = "avx2,fma")]
pub(crate) fn dct16x8_avx2(input: &[f32; 128], output: &mut [f32; 128]) {
let mut c: [__m256; 16] =
std::array::from_fn(|v| unsafe { _mm256_loadu_ps(input[v * 8..].as_ptr()) });
dct1d_16_flat(&mut c);
let mut top: [__m256; 8] = c[0..8].try_into().unwrap();
let mut bot: [__m256; 8] = c[8..16].try_into().unwrap();
transpose_8x8(&mut top);
transpose_8x8(&mut bot);
dct1d_8_flat(&mut top);
dct1d_8_flat(&mut bot);
let scale = _mm256_set1_ps(1.0 / 128.0);
for m in 0..8 {
let base = &mut output[m * 16..];
unsafe {
_mm256_storeu_ps(base.as_mut_ptr(), _mm256_mul_ps(top[m], scale));
_mm256_storeu_ps(base[8..].as_mut_ptr(), _mm256_mul_ps(bot[m], scale));
}
}
}
#[target_feature(enable = "avx2,fma")]
pub(crate) fn dct16x16_avx2(input: &[f32; 256], output: &mut [f32; 256]) {
let mut scratch = [0.0f32; 256];
for g in 0..2 {
let mut c: [__m256; 16] =
std::array::from_fn(|r| unsafe { _mm256_loadu_ps(input[r * 16 + g * 8..].as_ptr()) });
dct1d_16_flat(&mut c);
for t in 0..2 {
let mut tile: [__m256; 8] = c[t * 8..t * 8 + 8].try_into().unwrap();
transpose_8x8(&mut tile);
for (j, v) in tile.iter().enumerate() {
unsafe { _mm256_storeu_ps(scratch[(g * 8 + j) * 16 + t * 8..].as_mut_ptr(), *v) };
}
}
}
let scale = _mm256_set1_ps(1.0 / 256.0);
for g in 0..2 {
let mut c: [__m256; 16] = std::array::from_fn(|col| unsafe {
_mm256_loadu_ps(scratch[col * 16 + g * 8..].as_ptr())
});
dct1d_16_flat(&mut c);
for u in 0..16 {
unsafe {
_mm256_storeu_ps(
output[u * 16 + g * 8..].as_mut_ptr(),
_mm256_mul_ps(c[u], scale),
);
}
}
}
}
#[target_feature(enable = "avx2,fma")]
fn dct1d_32_flat(c: &mut [__m256; 32]) {
let mut evens = [_mm256_undefined_ps(); 16];
let mut odds = [_mm256_undefined_ps(); 16];
for i in 0..16 {
evens[i] = _mm256_add_ps(c[i], c[31 - i]);
odds[i] = _mm256_mul_ps(_mm256_sub_ps(c[i], c[31 - i]), _mm256_set1_ps(WC32[i]));
}
dct1d_16_flat(&mut evens);
dct1d_16_flat(&mut odds);
odds[0] = _mm256_fmadd_ps(odds[0], _mm256_set1_ps(std::f32::consts::SQRT_2), odds[1]);
for i in 1..15 {
odds[i] = _mm256_add_ps(odds[i], odds[i + 1]);
}
for i in 0..16 {
c[2 * i] = evens[i];
c[2 * i + 1] = odds[i];
}
}
#[target_feature(enable = "avx2,fma")]
pub(crate) fn dct32x32_avx2(input: &[f32; 1024], output: &mut [f32; 1024]) {
let mut colt = [0.0f32; 1024];
for g in 0..4 {
let mut c: [__m256; 32] =
std::array::from_fn(|r| unsafe { _mm256_loadu_ps(input[r * 32 + g * 8..].as_ptr()) });
dct1d_32_flat(&mut c);
for t in 0..4 {
let mut tile: [__m256; 8] = c[t * 8..t * 8 + 8].try_into().unwrap();
transpose_8x8(&mut tile);
for (j, v) in tile.iter().enumerate() {
unsafe { _mm256_storeu_ps(colt[(g * 8 + j) * 32 + t * 8..].as_mut_ptr(), *v) };
}
}
}
let scale = _mm256_set1_ps(1.0 / 1024.0);
for g in 0..4 {
let mut c: [__m256; 32] = std::array::from_fn(|col| unsafe {
_mm256_loadu_ps(colt[col * 32 + g * 8..].as_ptr())
});
dct1d_32_flat(&mut c);
for u in 0..32 {
unsafe {
_mm256_storeu_ps(
output[u * 32 + g * 8..].as_mut_ptr(),
_mm256_mul_ps(c[u], scale),
);
}
}
}
}
const IS2: f32 = std::f32::consts::FRAC_1_SQRT_2;
#[inline]
#[target_feature(enable = "avx2")]
fn idct2_flat(a: __m256, b: __m256) -> (__m256, __m256) {
let half = _mm256_set1_ps(0.5);
(
_mm256_mul_ps(_mm256_add_ps(a, b), half),
_mm256_mul_ps(_mm256_sub_ps(a, b), half),
)
}
#[inline]
#[target_feature(enable = "avx2")]
fn inv_dct1d_4_flat(c: &mut [__m256; 4]) {
let (t0, t1, mut t2, t3) = (c[0], c[2], c[1], c[3]);
t2 = _mm256_mul_ps(_mm256_sub_ps(t2, t3), _mm256_set1_ps(IS2));
let (n2, n3) = idct2_flat(t2, t3);
let a2 = _mm256_mul_ps(n2, _mm256_set1_ps(INV_WC4[0]));
let a3 = _mm256_mul_ps(n3, _mm256_set1_ps(INV_WC4[1]));
let (m0, m1) = idct2_flat(t0, t1);
let half = _mm256_set1_ps(0.5);
c[0] = _mm256_mul_ps(_mm256_add_ps(m0, a2), half);
c[3] = _mm256_mul_ps(_mm256_sub_ps(m0, a2), half);
c[1] = _mm256_mul_ps(_mm256_add_ps(m1, a3), half);
c[2] = _mm256_mul_ps(_mm256_sub_ps(m1, a3), half);
}
#[inline]
#[target_feature(enable = "avx2")]
fn inv_dct1d_8_flat(c: &mut [__m256; 8]) {
let mut t = [c[0]; 8];
for i in 0..4 {
t[i] = c[2 * i];
t[4 + i] = c[2 * i + 1];
}
t[6] = _mm256_sub_ps(t[6], t[7]);
t[5] = _mm256_sub_ps(t[5], t[6]);
t[4] = _mm256_mul_ps(_mm256_sub_ps(t[4], t[5]), _mm256_set1_ps(IS2));
let mut o = [t[4], t[5], t[6], t[7]];
inv_dct1d_4_flat(&mut o);
for i in 0..4 {
o[i] = _mm256_mul_ps(o[i], _mm256_set1_ps(INV_WC8[i]));
}
let mut e = [t[0], t[1], t[2], t[3]];
inv_dct1d_4_flat(&mut e);
let half = _mm256_set1_ps(0.5);
for i in 0..4 {
c[i] = _mm256_mul_ps(_mm256_add_ps(e[i], o[i]), half);
c[7 - i] = _mm256_mul_ps(_mm256_sub_ps(e[i], o[i]), half);
}
}
#[inline]
#[target_feature(enable = "avx2")]
fn inv_dct1d_16_flat(c: &mut [__m256; 16]) {
let mut t = [c[0]; 16];
for i in 0..8 {
t[i] = c[2 * i];
t[8 + i] = c[2 * i + 1];
}
for i in (9..=14).rev() {
t[i] = _mm256_sub_ps(t[i], t[i + 1]);
}
t[8] = _mm256_mul_ps(_mm256_sub_ps(t[8], t[9]), _mm256_set1_ps(IS2));
let mut o: [__m256; 8] = std::array::from_fn(|i| t[8 + i]);
inv_dct1d_8_flat(&mut o);
for i in 0..8 {
o[i] = _mm256_mul_ps(o[i], _mm256_set1_ps(INV_WC16[i]));
}
let mut e: [__m256; 8] = std::array::from_fn(|i| t[i]);
inv_dct1d_8_flat(&mut e);
let half = _mm256_set1_ps(0.5);
for i in 0..8 {
c[i] = _mm256_mul_ps(_mm256_add_ps(e[i], o[i]), half);
c[15 - i] = _mm256_mul_ps(_mm256_sub_ps(e[i], o[i]), half);
}
}
#[inline]
#[target_feature(enable = "avx2")]
fn inv_dct1d_32_flat(c: &mut [__m256; 32]) {
let mut t = [c[0]; 32];
for i in 0..16 {
t[i] = c[2 * i];
t[16 + i] = c[2 * i + 1];
}
for i in (17..=30).rev() {
t[i] = _mm256_sub_ps(t[i], t[i + 1]);
}
t[16] = _mm256_mul_ps(_mm256_sub_ps(t[16], t[17]), _mm256_set1_ps(IS2));
let mut o: [__m256; 16] = std::array::from_fn(|i| t[16 + i]);
inv_dct1d_16_flat(&mut o);
for i in 0..16 {
o[i] = _mm256_mul_ps(o[i], _mm256_set1_ps(INV_WC32[i]));
}
let mut e: [__m256; 16] = std::array::from_fn(|i| t[i]);
inv_dct1d_16_flat(&mut e);
let half = _mm256_set1_ps(0.5);
for i in 0..16 {
c[i] = _mm256_mul_ps(_mm256_add_ps(e[i], o[i]), half);
c[31 - i] = _mm256_mul_ps(_mm256_sub_ps(e[i], o[i]), half);
}
}
macro_rules! horizontal_idct_pass_avx2 {
($tmp:expr, $out:expr, $h:literal, $w:literal, $inv_h:path) => {{
for y8 in (0..$h).step_by(8) {
let mut c = [_mm256_undefined_ps(); $w];
for u8 in (0..$w).step_by(8) {
let mut tile: [__m256; 8] = std::array::from_fn(|j| unsafe {
_mm256_loadu_ps($tmp.as_ptr().add((y8 + j) * $w + u8))
});
transpose_8x8(&mut tile);
c[u8..u8 + 8].copy_from_slice(&tile);
}
$inv_h(&mut c);
for x8 in (0..$w).step_by(8) {
let mut tile: [__m256; 8] = c[x8..x8 + 8].try_into().unwrap();
transpose_8x8(&mut tile);
for (j, row) in tile.iter().enumerate() {
unsafe { _mm256_storeu_ps($out.as_mut_ptr().add((y8 + j) * $w + x8), *row) };
}
}
}
}};
}
macro_rules! inv_dct_natural_avx2 {
($name:ident, $n:literal, $h:literal, $w:literal, $inv_v:path, $inv_h:path) => {
#[target_feature(enable = "avx2")]
pub(crate) fn $name(coeff: &[f32; $n], out: &mut [f32; $n]) {
let mut tmp = [0.0f32; $n];
let scale = _mm256_set1_ps($n as f32);
for u8 in (0..$w).step_by(8) {
let mut c: [__m256; $h] = std::array::from_fn(|v| unsafe {
_mm256_mul_ps(_mm256_loadu_ps(coeff.as_ptr().add(v * $w + u8)), scale)
});
$inv_v(&mut c);
for y in 0..$h {
unsafe { _mm256_storeu_ps(tmp.as_mut_ptr().add(y * $w + u8), c[y]) };
}
}
horizontal_idct_pass_avx2!(tmp, out, $h, $w, $inv_h);
}
};
}
macro_rules! inv_dct_transposed_avx2 {
($name:ident, $n:literal, $h:literal, $w:literal, $inv_v:path, $inv_h:path) => {
#[target_feature(enable = "avx2")]
pub(crate) fn $name(coeff: &[f32; $n], out: &mut [f32; $n]) {
let mut tmp = [0.0f32; $n];
let scale = _mm256_set1_ps($n as f32);
for u8 in (0..$w).step_by(8) {
let mut c = [_mm256_undefined_ps(); $h];
for v8 in (0..$h).step_by(8) {
let mut tile: [__m256; 8] = std::array::from_fn(|j| unsafe {
_mm256_loadu_ps(coeff.as_ptr().add((u8 + j) * $h + v8))
});
transpose_8x8(&mut tile);
for (j, v) in tile.iter().enumerate() {
c[v8 + j] = _mm256_mul_ps(*v, scale);
}
}
$inv_v(&mut c);
for y in 0..$h {
unsafe { _mm256_storeu_ps(tmp.as_mut_ptr().add(y * $w + u8), c[y]) };
}
}
horizontal_idct_pass_avx2!(tmp, out, $h, $w, $inv_h);
}
};
}
inv_dct_transposed_avx2!(
inv_dct8x8_avx2,
64,
8,
8,
inv_dct1d_8_flat,
inv_dct1d_8_flat
);
inv_dct_natural_avx2!(
inv_dct8x16_avx2,
128,
8,
16,
inv_dct1d_8_flat,
inv_dct1d_16_flat
);
inv_dct_transposed_avx2!(
inv_dct16x8_avx2,
128,
16,
8,
inv_dct1d_16_flat,
inv_dct1d_8_flat
);
inv_dct_transposed_avx2!(
inv_dct16x16_avx2,
256,
16,
16,
inv_dct1d_16_flat,
inv_dct1d_16_flat
);
inv_dct_natural_avx2!(
inv_dct16x32_avx2,
512,
16,
32,
inv_dct1d_16_flat,
inv_dct1d_32_flat
);
inv_dct_transposed_avx2!(
inv_dct32x16_avx2,
512,
32,
16,
inv_dct1d_32_flat,
inv_dct1d_16_flat
);
inv_dct_transposed_avx2!(
inv_dct32x32_avx2,
1024,
32,
32,
inv_dct1d_32_flat,
inv_dct1d_32_flat
);
#[target_feature(enable = "avx2,fma")]
fn dct1d_4_m128(c: &mut [__m128; 4]) {
let t0 = _mm_add_ps(c[0], c[3]);
let t1 = _mm_add_ps(c[1], c[2]);
let e0 = _mm_add_ps(t0, t1);
let e1 = _mm_sub_ps(t0, t1);
let t2 = _mm_mul_ps(_mm_sub_ps(c[0], c[3]), _mm_set1_ps(WC4[0]));
let t3 = _mm_mul_ps(_mm_sub_ps(c[1], c[2]), _mm_set1_ps(WC4[1]));
let o0 = _mm_add_ps(t2, t3);
let o1 = _mm_sub_ps(t2, t3);
let m = _mm_fmadd_ps(o0, _mm_set1_ps(std::f32::consts::SQRT_2), o1);
c[0] = e0;
c[2] = e1;
c[1] = m;
c[3] = o1;
}
#[target_feature(enable = "avx2,fma")]
pub(crate) fn dct4x4_avx2(input: &[f32; 64], output: &mut [f32; 64]) {
let mut q = [_mm_undefined_ps(); 16];
for r in 0..4 {
for col in 0..4 {
q[r * 4 + col] = _mm_set_ps(
input[(4 + r) * 8 + (4 + col)], input[(4 + r) * 8 + col], input[r * 8 + (4 + col)], input[r * 8 + col], );
}
}
for r in 0..4 {
let mut row = [q[r * 4], q[r * 4 + 1], q[r * 4 + 2], q[r * 4 + 3]];
dct1d_4_m128(&mut row);
for col in 0..4 {
q[r * 4 + col] = row[col];
}
}
let inv16 = _mm_set1_ps(1.0 / 16.0);
let mut d = [[0.0f32; 4]; 16];
for col in 0..4 {
let mut cc = [q[col], q[4 + col], q[8 + col], q[12 + col]];
dct1d_4_m128(&mut cc);
for i in 0..4 {
unsafe { _mm_storeu_ps(d[col * 4 + i].as_mut_ptr(), _mm_mul_ps(cc[i], inv16)) };
}
}
for iy in 0..4 {
for ix in 0..4 {
let dd = &d[iy * 4 + ix];
for k in 0..4 {
let qy = k / 2;
let qx = k % 2;
output[(qy + iy * 2) * 8 + (qx + ix * 2)] = dd[k];
}
}
}
let b00 = output[0];
let b01 = output[1];
let b10 = output[8];
let b11 = output[9];
output[0] = (b00 + b01 + b10 + b11) * 0.25;
output[1] = (b00 + b01 - b10 - b11) * 0.25;
output[8] = (b00 - b01 + b10 - b11) * 0.25;
output[9] = (b00 - b01 - b10 + b11) * 0.25;
}
#[inline]
#[target_feature(enable = "avx2,fma")]
fn dct1d_4_flat(c: &mut [__m256; 4]) {
let s0 = _mm256_add_ps(c[0], c[3]);
let s1 = _mm256_add_ps(c[1], c[2]);
let e0 = _mm256_add_ps(s0, s1);
let e1 = _mm256_sub_ps(s0, s1);
let d2 = _mm256_mul_ps(_mm256_sub_ps(c[0], c[3]), _mm256_set1_ps(WC4[0]));
let d3 = _mm256_mul_ps(_mm256_sub_ps(c[1], c[2]), _mm256_set1_ps(WC4[1]));
let osum = _mm256_add_ps(d2, d3);
let odiff = _mm256_sub_ps(d2, d3);
let o0 = _mm256_fmadd_ps(osum, _mm256_set1_ps(std::f32::consts::SQRT_2), odiff);
c[0] = e0;
c[1] = o0;
c[2] = e1;
c[3] = odiff;
}
#[target_feature(enable = "avx2,fma")]
pub(crate) fn dct4x8_avx2(input: &[f32; 64], output: &mut [f32; 64]) {
let rows = load(input);
let mut top: [__m256; 4] = [rows[0], rows[1], rows[2], rows[3]];
let mut bot: [__m256; 4] = [rows[4], rows[5], rows[6], rows[7]];
dct1d_4_flat(&mut top);
dct1d_4_flat(&mut bot);
let mut r: [__m256; 8] = [
top[0], top[1], top[2], top[3], bot[0], bot[1], bot[2], bot[3],
];
transpose_8x8(&mut r);
dct1d_8_flat(&mut r);
transpose_8x8(&mut r);
let scale = _mm256_set1_ps(1.0 / 32.0);
let mut buf = [0.0f32; 64];
for k in 0..8 {
unsafe {
_mm256_storeu_ps(buf[k * 8..].as_mut_ptr(), _mm256_mul_ps(r[k], scale));
}
}
for k in 0..8 {
let vf = k % 4;
let half = k / 4;
for hf in 0..8 {
output[(half + vf * 2) * 8 + hf] = buf[k * 8 + hf];
}
}
let b0 = output[0];
let b1 = output[8];
output[0] = (b0 + b1) * 0.5;
output[8] = (b0 - b1) * 0.5;
}
#[target_feature(enable = "avx2,fma")]
pub(crate) fn dct8x4_avx2(input: &[f32; 64], output: &mut [f32; 64]) {
let mut rows = load(input);
dct1d_8_flat(&mut rows);
transpose_8x8(&mut rows); let mut left: [__m256; 4] = [rows[0], rows[1], rows[2], rows[3]];
let mut right: [__m256; 4] = [rows[4], rows[5], rows[6], rows[7]];
dct1d_4_flat(&mut left);
dct1d_4_flat(&mut right);
let scale = _mm256_set1_ps(1.0 / 32.0);
let mut lb = [0.0f32; 32];
let mut rb = [0.0f32; 32];
for hf in 0..4 {
unsafe {
_mm256_storeu_ps(lb[hf * 8..].as_mut_ptr(), _mm256_mul_ps(left[hf], scale));
_mm256_storeu_ps(rb[hf * 8..].as_mut_ptr(), _mm256_mul_ps(right[hf], scale));
}
}
for hf in 0..4 {
for vf in 0..8 {
output[(hf * 2) * 8 + vf] = lb[hf * 8 + vf];
output[(1 + hf * 2) * 8 + vf] = rb[hf * 8 + vf];
}
}
let b0 = output[0];
let b1 = output[8];
output[0] = (b0 + b1) * 0.5;
output[8] = (b0 - b1) * 0.5;
}
#[target_feature(enable = "avx2,fma")]
pub(crate) fn dct32x16_avx2(input: &[f32; 512], output: &mut [f32; 512]) {
let mut cols = [[_mm256_undefined_ps(); 16]; 4];
for g in 0..2 {
let mut c: [__m256; 32] =
std::array::from_fn(|r| unsafe { _mm256_loadu_ps(input[r * 16 + g * 8..].as_ptr()) });
dct1d_32_flat(&mut c);
for t in 0..4 {
let mut tile: [__m256; 8] = c[t * 8..t * 8 + 8].try_into().unwrap();
transpose_8x8(&mut tile);
cols[t][g * 8..g * 8 + 8].copy_from_slice(&tile);
}
}
let scale = _mm256_set1_ps(1.0 / 512.0);
for (t, tile) in cols.iter_mut().enumerate() {
dct1d_16_flat(tile);
for u in 0..16 {
unsafe {
_mm256_storeu_ps(
output[u * 32 + t * 8..].as_mut_ptr(),
_mm256_mul_ps(tile[u], scale),
);
}
}
}
}
#[target_feature(enable = "avx2,fma")]
pub(crate) fn dct16x32_avx2(input: &[f32; 512], output: &mut [f32; 512]) {
let mut scratch = [0.0f32; 512];
for g in 0..2 {
let mut c = [_mm256_undefined_ps(); 32];
for ct in 0..4 {
let mut tile: [__m256; 8] = std::array::from_fn(|j| unsafe {
_mm256_loadu_ps(input[(g * 8 + j) * 32 + ct * 8..].as_ptr())
});
transpose_8x8(&mut tile);
c[ct * 8..ct * 8 + 8].copy_from_slice(&tile);
}
dct1d_32_flat(&mut c);
for u in 0..32 {
unsafe { _mm256_storeu_ps(scratch[u * 16 + g * 8..].as_mut_ptr(), c[u]) };
}
}
let scale = _mm256_set1_ps(1.0 / 512.0);
for g in 0..4 {
let mut c = [_mm256_undefined_ps(); 16];
for rt in 0..2 {
let mut tile: [__m256; 8] = std::array::from_fn(|j| unsafe {
_mm256_loadu_ps(scratch[(g * 8 + j) * 16 + rt * 8..].as_ptr())
});
transpose_8x8(&mut tile);
c[rt * 8..rt * 8 + 8].copy_from_slice(&tile);
}
dct1d_16_flat(&mut c);
for v in 0..16 {
unsafe {
_mm256_storeu_ps(
output[v * 32 + g * 8..].as_mut_ptr(),
_mm256_mul_ps(c[v], scale),
);
}
}
}
}
#[cfg(test)]
mod tests {
const ATOL: f32 = 1e-4;
fn assert_close(neon: &[f32], scalar: &[f32], label: &str) {
assert_eq!(neon.len(), scalar.len(), "{label}: length mismatch");
let mut max_err: f32 = 0.0;
let mut worst = 0usize;
for (i, (n, s)) in neon.iter().zip(scalar.iter()).enumerate() {
let e = (n - s).abs();
if e > max_err {
max_err = e;
worst = i;
}
}
assert!(
max_err < ATOL,
"{label}: max error {max_err:.2e} at index {worst} \
(neon={:.6}, scalar={:.6})",
neon[worst],
scalar[worst]
);
}
fn rng_f32(seed: u64) -> f32 {
let mut x = seed.wrapping_add(0x9e3779b97f4a7c15);
x = (x ^ (x >> 30)).wrapping_mul(0xbf58476d1ce4e5b9);
x = (x ^ (x >> 27)).wrapping_mul(0x94d049bb133111eb);
x ^= x >> 31;
let u = (x >> 41) as f32; u / (1u32 << 23) as f32 * 2.0 - 1.0
}
fn fill<const N: usize>(seed: u64) -> [f32; N] {
let mut buf = [0.0f32; N];
for (i, v) in buf.iter_mut().enumerate() {
*v = rng_f32(seed.wrapping_add((i as u64).wrapping_mul(6364136223846793005)));
}
buf
}
fn assert_inverse_matches_scalar<const N: usize>(
scalar: fn(&[f32; N], &mut [f32; N]),
avx2: unsafe fn(&[f32; N], &mut [f32; N]),
label: &str,
) {
if !is_x86_feature_detected!("avx2") {
return;
}
let mut cases = Vec::with_capacity(35);
cases.push([0.0f32; N]);
let mut dc = [0.0f32; N];
dc[0] = 0.75;
cases.push(dc);
let mut alternating = [0.0f32; N];
for (i, v) in alternating.iter_mut().enumerate() {
*v = if i.is_multiple_of(2) { 1.0 } else { -1.0 };
}
cases.push(alternating);
for seed in 0..32 {
cases.push(fill(0xa7_c2_0000 + seed));
}
for (case, input) in cases.iter().enumerate() {
let mut got = [0.0f32; N];
let mut want = [0.0f32; N];
unsafe { avx2(input, &mut got) };
scalar(input, &mut want);
assert_close(&got, &want, &format!("{label} case={case}"));
}
}
#[test]
fn test_dct16x16_avx2_vs_scalar_random() {
if !is_x86_feature_detected!("avx2") || !is_x86_feature_detected!("fma") {
return;
}
use crate::avx::dct16x16_avx2;
use crate::dct::dct16x16_scalar;
for seed in 0u64..32 {
let input: [f32; 256] = fill(seed.wrapping_add(0xf00d));
let mut got = [0.0f32; 256];
let mut want = [0.0f32; 256];
unsafe { dct16x16_avx2(&input, &mut got) };
dct16x16_scalar(&input, &mut want);
assert_close(&got, &want, &format!("dct16x16 seed={seed}"));
}
}
#[test]
fn test_dct32x32_avx2_vs_scalar_random() {
if !is_x86_feature_detected!("avx2") || !is_x86_feature_detected!("fma") {
return;
}
use crate::avx::dct32x32_avx2;
use crate::dct::dct32x32_scalar;
for seed in 0u64..16 {
let input: [f32; 1024] = fill(seed.wrapping_add(0x3232));
let mut got = [0.0f32; 1024];
let mut want = [0.0f32; 1024];
unsafe { dct32x32_avx2(&input, &mut got) };
dct32x32_scalar(&input, &mut want);
assert_close(&got, &want, &format!("dct32x32 seed={seed}"));
}
}
#[test]
fn test_dct32x16_avx2_vs_scalar_random() {
if !is_x86_feature_detected!("avx2") || !is_x86_feature_detected!("fma") {
return;
}
use crate::avx::dct32x16_avx2;
use crate::dct::dct32x16_scalar;
for seed in 0u64..16 {
let input: [f32; 512] = fill(seed.wrapping_add(0x3216));
let mut got = [0.0f32; 512];
let mut want = [0.0f32; 512];
unsafe { dct32x16_avx2(&input, &mut got) };
dct32x16_scalar(&input, &mut want);
assert_close(&got, &want, &format!("dct32x16 seed={seed}"));
}
}
#[test]
fn test_dct16x32_avx2_vs_scalar_random() {
if !is_x86_feature_detected!("avx2") || !is_x86_feature_detected!("fma") {
return;
}
use crate::avx::dct16x32_avx2;
use crate::dct::dct16x32_scalar;
for seed in 0u64..16 {
let input: [f32; 512] = fill(seed.wrapping_add(0x1632));
let mut got = [0.0f32; 512];
let mut want = [0.0f32; 512];
unsafe { dct16x32_avx2(&input, &mut got) };
dct16x32_scalar(&input, &mut want);
assert_close(&got, &want, &format!("dct16x32 seed={seed}"));
}
}
#[test]
fn test_dct4x4_avx2_vs_scalar_random() {
if !is_x86_feature_detected!("avx2") || !is_x86_feature_detected!("fma") {
return;
}
use crate::avx::dct4x4_avx2;
use crate::dct::dct4x4_scalar;
for seed in 0u64..32 {
let input: [f32; 64] = fill(seed.wrapping_add(0x4a4));
let mut got = [0.0f32; 64];
let mut want = [0.0f32; 64];
unsafe { dct4x4_avx2(&input, &mut got) };
dct4x4_scalar(&input, &mut want);
assert_close(&got, &want, &format!("dct4x4 seed={seed}"));
}
}
#[test]
fn test_dct4x8_avx2_vs_scalar_random() {
if !is_x86_feature_detected!("avx2") || !is_x86_feature_detected!("fma") {
return;
}
use crate::avx::dct4x8_avx2;
use crate::dct::dct4x8_scalar;
for seed in 0u64..32 {
let input: [f32; 64] = fill(seed.wrapping_add(0x4a8));
let mut got = [0.0f32; 64];
let mut want = [0.0f32; 64];
unsafe { dct4x8_avx2(&input, &mut got) };
dct4x8_scalar(&input, &mut want);
assert_close(&got, &want, &format!("dct4x8 seed={seed}"));
}
}
#[test]
fn test_dct8x4_avx2_vs_scalar_random() {
if !is_x86_feature_detected!("avx2") || !is_x86_feature_detected!("fma") {
return;
}
use crate::avx::dct8x4_avx2;
use crate::dct::dct8x4_scalar;
for seed in 0u64..32 {
let input: [f32; 64] = fill(seed.wrapping_add(0x8a4));
let mut got = [0.0f32; 64];
let mut want = [0.0f32; 64];
unsafe { dct8x4_avx2(&input, &mut got) };
dct8x4_scalar(&input, &mut want);
assert_close(&got, &want, &format!("dct8x4 seed={seed}"));
}
}
#[test]
fn test_dct16x16_avx2_dc_only() {
if !is_x86_feature_detected!("avx2") || !is_x86_feature_detected!("fma") {
return;
}
use crate::avx::dct16x16_avx2;
use crate::dct::dct16x16_scalar;
let input = [0.5f32; 256];
let mut got = [0.0f32; 256];
let mut want = [0.0f32; 256];
unsafe { dct16x16_avx2(&input, &mut got) };
dct16x16_scalar(&input, &mut want);
assert_close(&got, &want, "dct16x16 dc-only");
}
#[test]
fn test_dct16x16_avx2_zero() {
if !is_x86_feature_detected!("avx2") || !is_x86_feature_detected!("fma") {
return;
}
use crate::avx::dct16x16_avx2;
use crate::dct::dct16x16_scalar;
let input = [0.0f32; 256];
let mut got = [0.0f32; 256];
let mut want = [0.0f32; 256];
unsafe { dct16x16_avx2(&input, &mut got) };
dct16x16_scalar(&input, &mut want);
assert_close(&got, &want, "dct16x16 zero");
}
#[test]
fn test_dct16x16_avx2_basis_vectors() {
if !is_x86_feature_detected!("avx2") || !is_x86_feature_detected!("fma") {
return;
}
use crate::avx::dct16x16_avx2;
use crate::dct::dct16x16_scalar;
for k in 0..256 {
let mut input = [0.0f32; 256];
input[k] = 1.0;
let mut got = [0.0f32; 256];
let mut want = [0.0f32; 256];
unsafe { dct16x16_avx2(&input, &mut got) };
dct16x16_scalar(&input, &mut want);
assert_close(&got, &want, &format!("dct16x16 basis[{k}]"));
}
}
#[test]
fn test_dct16x16_avx2_linearity() {
if !is_x86_feature_detected!("avx2") || !is_x86_feature_detected!("fma") {
return;
}
use crate::avx::dct16x16_avx2;
let a: [f32; 256] = fill(700);
let b: [f32; 256] = fill(800);
let mut sum = [0.0f32; 256];
for i in 0..256 {
sum[i] = a[i] + b[i];
}
let mut da = [0.0f32; 256];
let mut db = [0.0f32; 256];
let mut dsum = [0.0f32; 256];
unsafe {
dct16x16_avx2(&a, &mut da);
dct16x16_avx2(&b, &mut db);
dct16x16_avx2(&sum, &mut dsum);
}
let expected: Vec<f32> = (0..256).map(|i| da[i] + db[i]).collect();
assert_close(&dsum, &expected, "dct16x16 linearity");
}
#[test]
fn test_dct16x16_avx2_extreme_values() {
if !is_x86_feature_detected!("avx2") || !is_x86_feature_detected!("fma") {
return;
}
use crate::avx::dct16x16_avx2;
use crate::dct::dct16x16_scalar;
let mut input = [0.0f32; 256];
for i in 0..256 {
input[i] = if i % 2 == 0 { 1.0 } else { -1.0 };
}
let mut got = [0.0f32; 256];
let mut want = [0.0f32; 256];
unsafe { dct16x16_avx2(&input, &mut got) };
dct16x16_scalar(&input, &mut want);
assert_close(&got, &want, "dct16x16 alternating +-1");
}
#[test]
fn test_idct8x8_avx2_matches_scalar() {
assert_inverse_matches_scalar(
crate::dct::inv_dct8x8,
crate::avx::inv_dct8x8_avx2,
"idct8x8",
);
}
#[test]
fn test_idct8x16_avx2_matches_scalar() {
assert_inverse_matches_scalar(
crate::dct::inv_dct8x16,
crate::avx::inv_dct8x16_avx2,
"idct8x16",
);
}
#[test]
fn test_idct16x8_avx2_matches_scalar() {
assert_inverse_matches_scalar(
crate::dct::inv_dct16x8,
crate::avx::inv_dct16x8_avx2,
"idct16x8",
);
}
#[test]
fn test_idct16x16_avx2_matches_scalar() {
assert_inverse_matches_scalar(
crate::dct::inv_dct16x16,
crate::avx::inv_dct16x16_avx2,
"idct16x16",
);
}
#[test]
fn test_idct16x32_avx2_matches_scalar() {
assert_inverse_matches_scalar(
crate::dct::inv_dct16x32,
crate::avx::inv_dct16x32_avx2,
"idct16x32",
);
}
#[test]
fn test_idct32x16_avx2_matches_scalar() {
assert_inverse_matches_scalar(
crate::dct::inv_dct32x16,
crate::avx::inv_dct32x16_avx2,
"idct32x16",
);
}
#[test]
fn test_idct32x32_avx2_matches_scalar() {
assert_inverse_matches_scalar(
crate::dct::inv_dct32x32,
crate::avx::inv_dct32x32_avx2,
"idct32x32",
);
}
}