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 std::sync::OnceLock;
use jxl_simd::{F32SimdVec, I32SimdVec, SimdDescriptor, simd_function};
use super::step::TiledChannelView;
use crate::frame::modular::ModularStorage;
use crate::image::{ImageRectMut, OwnedRawImage, Rect};
use crate::util::{DITHER_TABLE, fast_jinc_windowed_sq_simd};
fn compute_jinc_subkernel(delta_x: f32, delta_y: f32) -> [f32; 25] {
let mut w = [0.0f32; 25];
let mut sum = 0.0f32;
for ty in 0..5 {
let py = (ty as f32 - 2.0) - delta_y;
for tx in 0..5 {
let px = (tx as f32 - 2.0) - delta_x;
let r2 = px * px + py * py;
let weight = crate::util::fast_jinc_windowed_sq(r2);
w[ty * 5 + tx] = weight;
sum += weight;
}
}
let inv_sum = 1.0 / sum;
for weight in &mut w {
*weight *= inv_sum;
}
w
}
#[inline(always)]
fn compute_row_kernel<D: SimdDescriptor>(
d: D,
fx: usize,
delta_y: f32,
num_weights: usize,
kernel_storage: &mut [f32],
) {
const IOTA: [f32; 16] = [
0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0, 11.0, 12.0, 13.0, 14.0, 15.0,
];
const { assert!(D::F32Vec::LEN <= 16) };
let lanes = D::F32Vec::LEN;
let inv_fx = 1.0 / (fx as f32);
let num_chunks = num_weights.div_ceil(lanes);
let idx_vec = D::F32Vec::load(d, &IOTA[..lanes]);
for chunk_idx in 0..num_chunks {
let ox_base = chunk_idx * lanes;
let ox_vec = D::F32Vec::splat(d, ox_base as f32) + idx_vec;
let delta_x_vec = (ox_vec + D::F32Vec::splat(d, 0.5)) * D::F32Vec::splat(d, inv_fx)
- D::F32Vec::splat(d, 0.5);
let mut subk_weights = [D::F32Vec::zero(d); 25];
let mut sum_vec = D::F32Vec::zero(d);
for ty in 0..5 {
let py = (ty as f32 - 2.0) - delta_y;
let py2_vec = D::F32Vec::splat(d, py * py);
for tx in 0..5 {
let px_vec = D::F32Vec::splat(d, tx as f32 - 2.0) - delta_x_vec;
let r2_vec = px_vec.mul_add(px_vec, py2_vec);
let w_vec = fast_jinc_windowed_sq_simd(d, r2_vec);
subk_weights[ty * 5 + tx] = w_vec;
sum_vec += w_vec;
}
}
let inv_sum_vec = D::F32Vec::splat(d, 1.0) / sum_vec;
let chunk_offset = chunk_idx * 25 * lanes;
for (tap, w_vec) in subk_weights.iter().enumerate() {
let norm_w = *w_vec * inv_sum_vec;
norm_w.store(&mut kernel_storage[chunk_offset + tap * lanes..]);
}
}
}
fn compute_jinc_squeeze_kernel(shift_diff: (usize, usize)) -> Vec<[f32; 25]> {
let (dx, dy) = shift_diff;
let fx = 1usize << dx;
let fy = 1usize << dy;
let mut kernel = vec![[0.0f32; 25]; fx * fy];
for oy in 0..fy {
let delta_y = (oy as f32 + 0.5) / (fy as f32) - 0.5;
for ox in 0..fx {
let delta_x = (ox as f32 + 0.5) / (fx as f32) - 0.5;
kernel[oy * fx + ox] = compute_jinc_subkernel(delta_x, delta_y);
}
}
kernel
}
const SMALL_SHIFT_LIMIT: usize = 4;
static SMALL_KERNEL_CACHE: [[OnceLock<Vec<[f32; 25]>>; SMALL_SHIFT_LIMIT]; SMALL_SHIFT_LIMIT] =
[const { [const { OnceLock::new() }; SMALL_SHIFT_LIMIT] }; SMALL_SHIFT_LIMIT];
fn get_small_squeeze_kernel(shift_diff: (usize, usize)) -> &'static [[f32; 25]] {
let (dx, dy) = shift_diff;
debug_assert!(dx < SMALL_SHIFT_LIMIT && dy < SMALL_SHIFT_LIMIT);
SMALL_KERNEL_CACHE[dx][dy].get_or_init(|| compute_jinc_squeeze_kernel(shift_diff))
}
#[derive(Debug, Default)]
pub(crate) struct SmoothUpsampleScratch {
buffer: [Vec<f32>; 5],
ibuf: Vec<i32>,
out_buf: Vec<i32>,
kernel_storage: Vec<f32>,
row_float: Vec<f32>,
}
impl SmoothUpsampleScratch {
fn init(&mut self, in_len: usize, out_len: usize, kernel_len: usize) {
for b in &mut self.buffer {
b.resize(in_len, 0.0);
}
self.ibuf.resize(in_len, 0);
self.out_buf.resize(out_len, 0);
self.row_float.resize(out_len, 0.0);
self.kernel_storage.resize(kernel_len, 0.0);
}
}
fn make_float<D: SimdDescriptor>(d: D, inp: &[i32], out: &mut [f32]) {
for (i, o) in inp
.chunks_exact(D::I32Vec::LEN)
.zip(out.chunks_exact_mut(D::F32Vec::LEN))
{
D::I32Vec::load(d, i).as_f32().store(o);
}
}
#[inline(always)]
fn store_interleaved_f32<D: SimdDescriptor, const FX: usize>(
out: &[D::F32Vec; FX],
dest: &mut [f32],
) {
match FX {
1 => out[0].store(dest),
2 => D::F32Vec::store_interleaved_2(out[0], out[1], dest),
4 => D::F32Vec::store_interleaved_4(out[0], out[1], out[2], out[3], dest),
8 => D::F32Vec::store_interleaved_8(
out[0], out[1], out[2], out[3], out[4], out[5], out[6], out[7], dest,
),
_ => unreachable!(),
}
}
#[inline(always)]
fn process_row<D: SimdDescriptor, const FX: usize>(
d: D,
weights: &[[f32; 25]],
in_xs: usize,
buffer: &[Vec<f32>; 5],
row_float: &mut [f32],
) {
let lanes = D::F32Vec::LEN;
let num_chunks = in_xs.div_ceil(lanes);
for chunk_idx in 0..num_chunks {
let ix = chunk_idx * lanes;
let mut out = [D::F32Vec::zero(d); FX];
for ty in 0..5 {
let row = &buffer[ty][ix..];
for tx in 0..5 {
let v = D::F32Vec::load(d, &row[tx..]);
for (ox, out_vec) in out.iter_mut().enumerate() {
let w = weights[ox][ty * 5 + tx];
*out_vec = v.mul_add(D::F32Vec::splat(d, w), *out_vec);
}
}
}
store_interleaved_f32::<D, FX>(&out, &mut row_float[FX * ix..]);
}
}
#[inline(always)]
fn process_row_large<D: SimdDescriptor>(
d: D,
kernel_storage: &[f32],
fx: usize,
xs: usize,
buffer: &[Vec<f32>; 5],
row_float: &mut [f32],
) {
let lanes = D::F32Vec::LEN;
let in_xs = xs.div_ceil(fx);
for ix in 0..in_xs {
let base_x = ix * fx;
let mut c = [D::F32Vec::zero(d); 25];
for ty in 0..5 {
let row = &buffer[ty][ix..];
for tx in 0..5 {
c[ty * 5 + tx] = D::F32Vec::splat(d, row[tx]);
}
}
let num_ox = (xs - base_x).min(fx);
let num_chunks = num_ox.div_ceil(lanes);
for ox_chunk in 0..num_chunks {
let ox_base = ox_chunk * lanes;
let chunk_offset = ox_chunk * 25 * lanes;
let mut out_vec = D::F32Vec::zero(d);
for tap in 0..25 {
let w_vec = D::F32Vec::load(d, &kernel_storage[chunk_offset + tap * lanes..]);
out_vec = c[tap].mul_add(w_vec, out_vec);
}
out_vec.store(&mut row_float[base_x + ox_base..]);
}
}
}
#[inline(always)]
fn dither_round_and_store<D: SimdDescriptor>(
d: D,
dither: bool,
dither_y: usize,
x0: usize,
xs: usize,
row_float: &[f32],
output_row: &mut [i32],
) {
const { assert!(D::F32Vec::LEN <= 16) };
let lanes = D::F32Vec::LEN;
let half = D::F32Vec::splat(d, 0.5);
let num_chunks = xs.div_ceil(lanes);
for chunk_idx in 0..num_chunks {
let x = chunk_idx * lanes;
let cur_len = (xs - x).min(lanes);
let val = D::F32Vec::load(d, &row_float[x..]);
let dither_val = if dither {
let dither_x = (x0 + x) % 32;
D::F32Vec::load(d, &DITHER_TABLE[dither_y][dither_x..])
} else {
D::F32Vec::zero(d)
};
let dithered = val + dither_val;
let rounded = (dithered + half.copysign(dithered)).as_i32();
if cur_len == lanes && output_row.len() >= x + lanes {
rounded.store(&mut output_row[x..]);
} else {
let mut temp = [0i32; 16];
rounded.store(&mut temp[..lanes]);
let to_copy = cur_len.min(output_row.len() - x);
output_row[x..x + to_copy].copy_from_slice(&temp[..to_copy]);
}
}
}
#[allow(clippy::too_many_arguments)]
#[inline(always)]
fn smooth_upsample_simd_impl<D: SimdDescriptor>(
d: D,
input: &TiledChannelView<'_>,
shift_diff: (usize, usize),
dither: bool,
rect: Rect,
output: &mut OwnedRawImage,
storage: ModularStorage,
scratch: &mut SmoothUpsampleScratch,
) {
let (dx, dy) = shift_diff;
let (fx, fy) = (1usize << dx, 1usize << dy);
let (x0, y0) = (rect.origin.0, rect.origin.1);
let (xs, ys) = (rect.size.0, rect.size.1);
let (in_xs, in_ys) = (xs.div_ceil(fx), ys.div_ceil(fy));
let (col_offset, row_offset) = (x0 / fx, y0 / fy);
let lanes = D::I32Vec::LEN;
if in_xs == 0 || in_ys == 0 {
return;
}
let num_chunks = in_xs.div_ceil(lanes);
let in_len = (num_chunks + 1) * lanes + 8;
let out_len = if fx <= 8 {
(num_chunks + 1) * lanes * fx
} else {
xs.next_multiple_of(lanes)
};
let num_weights = xs.min(fx);
let kernel_len = if fx > 8 {
num_weights.div_ceil(lanes) * 25 * lanes
} else {
0
};
scratch.init(in_len, out_len, kernel_len);
for (dy_idx, buf) in scratch.buffer.iter_mut().enumerate().take(4) {
let yg = (row_offset + dy_idx) as isize - 2;
if storage == ModularStorage::I16 {
input.load_row_to_scratch::<i16>(yg, col_offset, in_xs + 4, &mut scratch.ibuf);
} else {
input.load_row_to_scratch::<i32>(yg, col_offset, in_xs + 4, &mut scratch.ibuf);
}
make_float(d, &scratch.ibuf, buf);
}
let is_small = dx < SMALL_SHIFT_LIMIT && dy < SMALL_SHIFT_LIMIT;
let small_kernel = if is_small {
Some(get_small_squeeze_kernel(shift_diff))
} else {
None
};
let mut small_dynamic_weights = [[0.0f32; 25]; 8];
for iy_center in 0..in_ys {
let yg = (row_offset + iy_center) as isize + 2;
if storage == ModularStorage::I16 {
input.load_row_to_scratch::<i16>(yg, col_offset, in_xs + 4, &mut scratch.ibuf);
} else {
input.load_row_to_scratch::<i32>(yg, col_offset, in_xs + 4, &mut scratch.ibuf);
}
make_float(d, &scratch.ibuf, &mut scratch.buffer[4]);
for oy in 0..fy {
let yout = fy * iy_center + oy;
if yout >= ys {
continue;
}
let dither_y = (y0 + yout) % 32;
let delta_y = (oy as f32 + 0.5) / (fy as f32) - 0.5;
if fx <= 8 {
let row_weights: &[[f32; 25]] = if let Some(k) = small_kernel {
&k[oy * fx..(oy + 1) * fx]
} else {
for (ox, w) in small_dynamic_weights.iter_mut().enumerate().take(fx) {
let delta_x = (ox as f32 + 0.5) / (fx as f32) - 0.5;
*w = compute_jinc_subkernel(delta_x, delta_y);
}
&small_dynamic_weights[..fx]
};
match fx {
1 => process_row::<D, 1>(
d,
row_weights,
in_xs,
&scratch.buffer,
&mut scratch.row_float,
),
2 => process_row::<D, 2>(
d,
row_weights,
in_xs,
&scratch.buffer,
&mut scratch.row_float,
),
4 => process_row::<D, 4>(
d,
row_weights,
in_xs,
&scratch.buffer,
&mut scratch.row_float,
),
8 => process_row::<D, 8>(
d,
row_weights,
in_xs,
&scratch.buffer,
&mut scratch.row_float,
),
_ => unreachable!(),
}
} else {
compute_row_kernel(d, fx, delta_y, num_weights, &mut scratch.kernel_storage);
process_row_large(
d,
&scratch.kernel_storage,
fx,
xs,
&scratch.buffer,
&mut scratch.row_float,
);
}
let mut img;
let out_row = if storage == ModularStorage::I16 {
&mut scratch.out_buf
} else {
img = ImageRectMut::<i32>::from_raw(output.as_rect_mut());
img.row(yout)
};
dither_round_and_store(d, dither, dither_y, x0, xs, &scratch.row_float, out_row);
let mut img;
if storage == ModularStorage::I16 {
img = ImageRectMut::<i16>::from_raw(output.as_rect_mut());
for (dst, &src) in img.row(yout).iter_mut().zip(&scratch.out_buf) {
*dst = src as i16;
}
}
}
scratch.buffer.rotate_left(1);
}
}
simd_function!(
smooth_upsample,
d: D,
#[allow(clippy::too_many_arguments)]
pub(super) fn smooth_upsample_dispatch(
input: &TiledChannelView<'_>,
shift_diff: (usize, usize),
dither: bool,
rect: Rect,
output: &mut OwnedRawImage,
storage: ModularStorage,
scratch: &mut SmoothUpsampleScratch,
) {
smooth_upsample_simd_impl(
d, input, shift_diff, dither, rect, output, storage, scratch,
);
}
);
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_compute_jinc_squeeze_kernel_normalized() {
for &(dx, dy) in &[
(0, 0),
(1, 0),
(0, 1),
(1, 1),
(2, 2),
(3, 3),
(1, 2),
(2, 1),
(4, 4),
] {
let kernel = compute_jinc_squeeze_kernel((dx, dy));
let fx = 1 << dx;
let fy = 1 << dy;
assert_eq!(kernel.len(), fx * fy);
for w in &kernel {
let sum: f32 = w.iter().sum();
assert!(
(sum - 1.0).abs() < 1e-5,
"Kernel normalization failed for shift ({dx}, {dy}): sum={sum}"
);
}
}
}
#[test]
fn test_compute_jinc_subkernel_symmetry() {
let k_center = compute_jinc_subkernel(0.0, 0.0);
// Top-left vs bottom-right symmetry at (0, 0)
for ty in 0..5 {
for tx in 0..5 {
let w1 = k_center[ty * 5 + tx];
let w2 = k_center[(4 - ty) * 5 + (4 - tx)];
assert!(
(w1 - w2).abs() < 1e-6,
"Symmetry mismatch at ({tx}, {ty}): {w1} vs {w2}"
);
}
}
// Horizontal flip with +/- delta_x
let k_left = compute_jinc_subkernel(-0.25, 0.0);
let k_right = compute_jinc_subkernel(0.25, 0.0);
for ty in 0..5 {
for tx in 0..5 {
let w_l = k_left[ty * 5 + tx];
let w_r = k_right[ty * 5 + (4 - tx)];
assert!(
(w_l - w_r).abs() < 1e-6,
"Horizontal symmetry mismatch at ({tx}, {ty}): {w_l} vs {w_r}"
);
}
}
}
fn test_compute_row_kernel_simd_equivalent<D: SimdDescriptor>(d: D) {
let lanes = D::F32Vec::LEN;
for fx in [8, 15, 16, 32, 64] {
for num_weights in [1usize, 3, 7, 8, 15, 16, fx] {
if num_weights > fx {
continue;
}
for delta_y in [-0.5, -0.25, 0.0, 0.123, 0.456] {
let num_chunks = num_weights.div_ceil(lanes);
let mut simd_storage = vec![0.0f32; num_chunks * 25 * lanes];
compute_row_kernel(d, fx, delta_y, num_weights, &mut simd_storage);
for ox in 0..num_weights {
let delta_x = (ox as f32 + 0.5) / (fx as f32) - 0.5;
let scalar_subk = compute_jinc_subkernel(delta_x, delta_y);
let ox_chunk = ox / lanes;
let ox_lane = ox % lanes;
let chunk_offset = ox_chunk * 25 * lanes;
for tap in 0..25 {
let simd_val = simd_storage[chunk_offset + tap * lanes + ox_lane];
let abs_err = (simd_val - scalar_subk[tap]).abs();
assert!(
abs_err < 1e-6,
"Mismatch at fx={fx}, delta_y={delta_y}, ox={ox}, tap={tap}: simd={}, scalar={}, err={}",
simd_val,
scalar_subk[tap],
abs_err
);
}
}
}
}
}
}
jxl_simd::test_all_instruction_sets!(test_compute_row_kernel_simd_equivalent);
fn test_process_row_large_constant<D: SimdDescriptor>(d: D) {
let lanes = D::F32Vec::LEN;
for fx in [16, 32, 64] {
for xs in [1usize, 3, 7, 15, 16, 20, 32, 50, 64, 100] {
let in_xs = xs.div_ceil(fx);
let buffer = [
vec![42.0f32; in_xs + 8],
vec![42.0f32; in_xs + 8],
vec![42.0f32; in_xs + 8],
vec![42.0f32; in_xs + 8],
vec![42.0f32; in_xs + 8],
];
let mut row_float = vec![0.0f32; xs.next_multiple_of(lanes)];
let num_weights = xs.min(fx);
let num_chunks = num_weights.div_ceil(lanes);
let mut kernel_storage = vec![0.0f32; num_chunks * 25 * lanes];
compute_row_kernel(d, fx, 0.0, num_weights, &mut kernel_storage);
process_row_large(d, &kernel_storage, fx, xs, &buffer, &mut row_float);
for &val in &row_float[..xs] {
assert!(
(val - 42.0).abs() < 1e-4,
"Mismatch at fx={fx}, xs={xs}: val={val}"
);
}
}
}
}
jxl_simd::test_all_instruction_sets!(test_process_row_large_constant);
}