Revision control
Copy as Markdown
Other Tools
// Copyright (c) the JPEG XL Project Authors. All rights reserved.
//
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
use super::*;
use jxl_simd::{test_all_instruction_sets, ScalarDescriptor, SimdDescriptor};
use rand::Rng;
use rand::SeedableRng;
use rand_chacha::ChaCha12Rng;
use test_log::test;
use std::f64::consts::FRAC_1_SQRT_2;
use std::f64::consts::PI;
use std::f64::consts::SQRT_2;
#[inline(always)]
fn alpha(u: usize) -> f64 {
if u == 0 {
FRAC_1_SQRT_2
} else {
1.0
}
}
pub fn dct1d(input_matrix: &[Vec<f64>]) -> Vec<Vec<f64>> {
let num_rows = input_matrix.len();
if num_rows == 0 {
return Vec::new();
}
let num_cols = input_matrix[0].len();
let mut output_matrix = vec![vec![0.0f64; num_cols]; num_rows];
let scale: f64 = SQRT_2;
// Precompute the DCT matrix (size: n_rows x n_rows)
let mut dct_coeff_matrix = vec![vec![0.0f64; num_rows]; num_rows];
for (u_freq, row) in dct_coeff_matrix.iter_mut().enumerate() {
let alpha_u_val = alpha(u_freq);
for (y_spatial, coeff) in row.iter_mut().enumerate() {
*coeff = alpha_u_val
* ((y_spatial as f64 + 0.5) * u_freq as f64 * PI / num_rows as f64).cos()
* scale;
}
}
// Perform the DCT calculation column by column
for x_col_idx in 0..num_cols {
for u_freq_idx in 0..num_rows {
let mut sum = 0.0;
for (y_spatial_idx, col) in input_matrix.iter().enumerate() {
// This access `input_matrix[y_spatial_idx][x_col_idx]` assumes the input_matrix
// is rectangular. If not, it might panic here.
sum += dct_coeff_matrix[u_freq_idx][y_spatial_idx] * col[x_col_idx];
}
output_matrix[u_freq_idx][x_col_idx] = sum;
}
}
output_matrix
}
pub fn idct1d(input_matrix: &[Vec<f64>]) -> Vec<Vec<f64>> {
let num_rows = input_matrix.len();
if num_rows == 0 {
return Vec::new();
}
let num_cols = input_matrix[0].len();
let mut output_matrix = vec![vec![0.0f64; num_cols]; num_rows];
let scale: f64 = SQRT_2;
// Precompute the DCT matrix (size: num_rows x num_rows)
let mut dct_coeff_matrix = vec![vec![0.0f64; num_rows]; num_rows];
for (u_freq, row) in dct_coeff_matrix.iter_mut().enumerate() {
let alpha_u_val = alpha(u_freq);
for (y_def_idx, coeff) in row.iter_mut().enumerate() {
*coeff = alpha_u_val
* ((y_def_idx as f64 + 0.5) * u_freq as f64 * PI / num_rows as f64).cos()
* scale;
}
}
// Perform the IDCT calculation column by column
for x_col_idx in 0..num_cols {
for (y_row_idx, row) in output_matrix.iter_mut().enumerate() {
let mut sum = 0.0;
for (u_freq_idx, col) in input_matrix.iter().enumerate() {
// This access input_coeffs_matrix[u_freq_idx][x_col_idx] assumes input_coeffs_matrix
// is rectangular. If not, it might panic here.
sum += dct_coeff_matrix[u_freq_idx][y_row_idx] * col[x_col_idx];
}
row[x_col_idx] = sum;
}
}
output_matrix
}
fn transpose_f64(matrix: &[Vec<f64>]) -> Vec<Vec<f64>> {
if matrix.is_empty() {
return Vec::new();
}
let num_rows = matrix.len();
let num_cols = matrix[0].len();
let mut transposed = vec![vec![0.0; num_rows]; num_cols];
for i in 0..num_rows {
for j in 0..num_cols {
transposed[j][i] = matrix[i][j];
}
}
transposed
}
pub fn slow_idct2d(input: &[Vec<f64>]) -> Vec<Vec<f64>> {
let rows = input.len();
let cols = input[0].len();
let idct1 = if rows < cols {
let transposed = transpose_f64(input);
idct1d(&transposed)
} else {
let input: Vec<_> = input.iter().flat_map(|x| x.iter()).copied().collect();
let input: Vec<_> = input.chunks_exact(rows).map(|x| x.to_vec()).collect();
idct1d(&input)
};
let transposed1 = transpose_f64(&idct1);
idct1d(&transposed1)
}
pub fn scales(n: usize) -> Vec<f64> {
(0..n)
.map(|i| {
(i as f64 / (16 * n) as f64 * PI).cos()
* (i as f64 / (8 * n) as f64 * PI).cos()
* (i as f64 / (4 * n) as f64 * PI).cos()
* n as f64
})
.collect()
}
pub fn slow_reinterpreting_dct2d(input: &[Vec<f64>]) -> Vec<Vec<f64>> {
let rows = input.len();
let cols = input[0].len();
let dct1 = dct1d(input);
let tdct1 = transpose_f64(&dct1);
let dct2 = dct1d(&tdct1);
let mut res = if rows < cols {
transpose_f64(&dct2)
} else {
dct2
};
let row_scales = scales(rows);
let col_scales = scales(cols);
if rows < cols {
for y in 0..rows {
for x in 0..cols {
res[y][x] /= row_scales[y] * col_scales[x];
}
}
} else {
for y in 0..cols {
for x in 0..rows {
res[y][x] /= row_scales[x] * col_scales[y];
}
}
}
res
}
#[track_caller]
fn check_close(a: f64, b: f64, max_err: f64) {
let abs = (a - b).abs();
let rel = abs / a.abs().max(b.abs());
assert!(
abs < max_err || rel < max_err,
"a: {a} b: {b} abs diff: {abs:?} rel diff: {rel:?}"
);
}
macro_rules! test_reinterpreting_dct1d_eq_slow_n {
($test_name:ident, $n_val:expr, $do_idct_fun:path, $tolerance:expr) => {
#[test]
fn $test_name() {
const N: usize = $n_val;
let input_matrix_for_ref = random_matrix(N, 1);
let output_matrix_slow: Vec<Vec<f64>> = dct1d(&input_matrix_for_ref);
let mut output: Vec<_> = input_matrix_for_ref.iter().map(|x| x[0] as f32).collect();
let d = ScalarDescriptor {};
let (output_chunks, remainder) = output.as_chunks_mut::<1>();
assert!(remainder.is_empty());
$do_idct_fun(d, output_chunks, 1);
let scales = scales(N);
for i in 0..N {
check_close(
output[i] as f64,
output_matrix_slow[i][0] / scales[i],
$tolerance,
);
}
}
};
}
test_reinterpreting_dct1d_eq_slow_n!(
test_reinterpreting_dct1d_2_eq_slow,
2,
do_reinterpreting_dct_2,
1e-6
);
test_reinterpreting_dct1d_eq_slow_n!(
test_reinterpreting_dct1d_4_eq_slow,
4,
do_reinterpreting_dct_4,
1e-6
);
test_reinterpreting_dct1d_eq_slow_n!(
test_reinterpreting_dct1d_8_eq_slow,
8,
do_reinterpreting_dct_8,
1e-6
);
test_reinterpreting_dct1d_eq_slow_n!(
test_reinterpreting_dct1d_16_eq_slow,
16,
do_reinterpreting_dct_16,
5e-6
);
test_reinterpreting_dct1d_eq_slow_n!(
test_reinterpreting_dct1d_32_eq_slow,
32,
do_reinterpreting_dct_32,
5e-6
);
fn random_matrix(n: usize, m: usize) -> Vec<Vec<f64>> {
let mut rng = ChaCha12Rng::seed_from_u64(0);
let mut data = vec![vec![0.0; m]; n];
data.iter_mut()
.flat_map(|x| x.iter_mut())
.for_each(|x| *x = rng.random_range(-1.0..1.0));
data
}
macro_rules! test_idct1d_eq_slow_n {
($test_name:ident, $n_val:expr, $do_idct_fun:path, $tolerance:expr) => {
#[test]
fn $test_name() {
const N: usize = $n_val;
let input_matrix_for_ref = random_matrix(N, 1);
let output_matrix_slow: Vec<Vec<f64>> = idct1d(&input_matrix_for_ref);
let mut output: Vec<_> = input_matrix_for_ref.iter().map(|x| x[0] as f32).collect();
let d = ScalarDescriptor {};
let (output_chunks, remainder) = output.as_chunks_mut::<1>();
assert!(remainder.is_empty());
$do_idct_fun(d, output_chunks, 1);
for i in 0..N {
check_close(output[i] as f64, output_matrix_slow[i][0], $tolerance);
}
}
};
}
test_idct1d_eq_slow_n!(test_idct1d_2_eq_slow, 2, do_idct_2, 1e-6);
test_idct1d_eq_slow_n!(test_idct1d_4_eq_slow, 4, do_idct_4, 1e-6);
test_idct1d_eq_slow_n!(test_idct1d_8_eq_slow, 8, do_idct_8, 1e-6);
test_idct1d_eq_slow_n!(test_idct1d_16_eq_slow, 16, do_idct_16, 1e-6);
test_idct1d_eq_slow_n!(test_idct1d_32_eq_slow, 32, do_idct_32, 5e-6);
test_idct1d_eq_slow_n!(test_idct1d_64_eq_slow, 64, do_idct_64, 5e-6);
test_idct1d_eq_slow_n!(test_idct1d_128_eq_slow, 128, do_idct_128, 5e-5);
test_idct1d_eq_slow_n!(test_idct1d_256_eq_slow, 256, do_idct_256, 5e-5);
macro_rules! test_idct2d_eq_slow {
($test_name:ident, $rows:expr, $cols:expr, $fast_idct:ident, $tol:expr) => {
fn $test_name<D: SimdDescriptor>(d: D) {
const N: usize = $rows;
const M: usize = $cols;
let slow_input = random_matrix(N, M);
let slow_output = slow_idct2d(&slow_input);
let mut fast_input: Vec<_> = slow_input
.iter()
.flat_map(|x| x.iter())
.map(|x| *x as f32)
.collect();
$fast_idct(d, &mut fast_input);
for r in 0..N {
for c in 0..M {
check_close(fast_input[r * M + c] as f64, slow_output[r][c], $tol);
}
}
}
test_all_instruction_sets!($test_name);
};
}
test_idct2d_eq_slow!(test_idct2d_2_2_eq_slow, 2, 2, idct2d_2_2, 1e-6);
test_idct2d_eq_slow!(test_idct2d_4_4_eq_slow, 4, 4, idct2d_4_4, 1e-6);
test_idct2d_eq_slow!(test_idct2d_4_8_eq_slow, 4, 8, idct2d_4_8, 1e-6);
test_idct2d_eq_slow!(test_idct2d_8_4_eq_slow, 8, 4, idct2d_8_4, 1e-6);
test_idct2d_eq_slow!(test_idct2d_8_8_eq_slow, 8, 8, idct2d_8_8, 5e-6);
test_idct2d_eq_slow!(test_idct2d_16_8_eq_slow, 16, 8, idct2d_16_8, 5e-6);
test_idct2d_eq_slow!(test_idct2d_8_16_eq_slow, 8, 16, idct2d_8_16, 5e-6);
test_idct2d_eq_slow!(test_idct2d_16_16_eq_slow, 16, 16, idct2d_16_16, 1e-5);
test_idct2d_eq_slow!(test_idct2d_32_8_eq_slow, 32, 8, idct2d_32_8, 5e-6);
test_idct2d_eq_slow!(test_idct2d_8_32_eq_slow, 8, 32, idct2d_8_32, 5e-6);
test_idct2d_eq_slow!(test_idct2d_32_16_eq_slow, 32, 16, idct2d_32_16, 1e-5);
test_idct2d_eq_slow!(test_idct2d_16_32_eq_slow, 16, 32, idct2d_16_32, 1e-5);
test_idct2d_eq_slow!(test_idct2d_32_32_eq_slow, 32, 32, idct2d_32_32, 5e-5);
test_idct2d_eq_slow!(test_idct2d_64_32_eq_slow, 64, 32, idct2d_64_32, 1e-4);
test_idct2d_eq_slow!(test_idct2d_32_64_eq_slow, 32, 64, idct2d_32_64, 1e-4);
test_idct2d_eq_slow!(test_idct2d_64_64_eq_slow, 64, 64, idct2d_64_64, 1e-4);
test_idct2d_eq_slow!(test_idct2d_128_64_eq_slow, 128, 64, idct2d_128_64, 5e-4);
test_idct2d_eq_slow!(test_idct2d_64_128_eq_slow, 64, 128, idct2d_64_128, 5e-4);
test_idct2d_eq_slow!(test_idct2d_128_128_eq_slow, 128, 128, idct2d_128_128, 5e-4);
test_idct2d_eq_slow!(test_idct2d_256_128_eq_slow, 256, 128, idct2d_256_128, 1e-3);
test_idct2d_eq_slow!(test_idct2d_128_256_eq_slow, 128, 256, idct2d_128_256, 1e-3);
test_idct2d_eq_slow!(test_idct2d_256_256_eq_slow, 256, 256, idct2d_256_256, 5e-3);
macro_rules! test_reinterpreting_dct_eq_slow {
($test_name:ident, $fun: ident, $rows:expr, $cols:expr, $tol:expr) => {
fn $test_name<D: SimdDescriptor>(d: D) {
const N: usize = $rows;
const M: usize = $cols;
let slow_input = random_matrix(N, M);
let slow_output = slow_reinterpreting_dct2d(&slow_input);
let mut fast_input: Vec<_> = slow_input
.iter()
.flat_map(|x| x.iter())
.map(|x| *x as f32)
.collect();
let mut output = [0.0; $rows * $cols * 64];
$fun(d, &mut fast_input, &mut output);
let on = slow_output.len();
let om = slow_output[0].len();
for r in 0..on {
for c in 0..om {
check_close(output[r * om * 8 + c] as f64, slow_output[r][c], $tol);
}
}
}
test_all_instruction_sets!($test_name);
};
}
test_reinterpreting_dct_eq_slow!(
test_reinterpreting_dct_1x2_eq_slow,
reinterpreting_dct2d_1_2,
1,
2,
1e-6
);
test_reinterpreting_dct_eq_slow!(
test_reinterpreting_dct_2x1_eq_slow,
reinterpreting_dct2d_2_1,
2,
1,
1e-6
);
test_reinterpreting_dct_eq_slow!(
test_reinterpreting_dct_2x2_eq_slow,
reinterpreting_dct2d_2_2,
2,
2,
1e-6
);
test_reinterpreting_dct_eq_slow!(
test_reinterpreting_dct_1x4_eq_slow,
reinterpreting_dct2d_1_4,
1,
4,
1e-6
);
test_reinterpreting_dct_eq_slow!(
test_reinterpreting_dct_4x1_eq_slow,
reinterpreting_dct2d_4_1,
4,
1,
1e-6
);
test_reinterpreting_dct_eq_slow!(
test_reinterpreting_dct_2x4_eq_slow,
reinterpreting_dct2d_2_4,
2,
4,
1e-6
);
test_reinterpreting_dct_eq_slow!(
test_reinterpreting_dct_4x2_eq_slow,
reinterpreting_dct2d_4_2,
4,
2,
1e-6
);
test_reinterpreting_dct_eq_slow!(
test_reinterpreting_dct_4x4_eq_slow,
reinterpreting_dct2d_4_4,
4,
4,
1e-6
);
test_reinterpreting_dct_eq_slow!(
test_reinterpreting_dct_8x4_eq_slow,
reinterpreting_dct2d_8_4,
8,
4,
1e-6
);
test_reinterpreting_dct_eq_slow!(
test_reinterpreting_dct_4x8_eq_slow,
reinterpreting_dct2d_4_8,
4,
8,
1e-6
);
test_reinterpreting_dct_eq_slow!(
test_reinterpreting_dct_8x8_eq_slow,
reinterpreting_dct2d_8_8,
8,
8,
1e-6
);
test_reinterpreting_dct_eq_slow!(
test_reinterpreting_dct_8x16_eq_slow,
reinterpreting_dct2d_8_16,
8,
16,
5e-6
);
test_reinterpreting_dct_eq_slow!(
test_reinterpreting_dct_16x8_eq_slow,
reinterpreting_dct2d_16_8,
16,
8,
5e-6
);
test_reinterpreting_dct_eq_slow!(
test_reinterpreting_dct_16x16_eq_slow,
reinterpreting_dct2d_16_16,
16,
16,
5e-6
);
test_reinterpreting_dct_eq_slow!(
test_reinterpreting_dct_32x16_eq_slow,
reinterpreting_dct2d_32_16,
32,
16,
5e-6
);
test_reinterpreting_dct_eq_slow!(
test_reinterpreting_dct_16x32_eq_slow,
reinterpreting_dct2d_16_32,
16,
32,
5e-6
);
test_reinterpreting_dct_eq_slow!(
test_reinterpreting_dct_32x32_eq_slow,
reinterpreting_dct2d_32_32,
32,
32,
5e-6
);