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::{collections::HashMap, sync::atomic::AtomicUsize};
use num_traits::FromPrimitive;
use crate::{
error::{Error, Result},
frame::modular::{
ChannelInfo, ModularBuffer, ModularBufferInfo, ModularGridKind, Predictor,
transforms::step::{TransformStep, TransformStepChunk},
},
headers::{self, frame_header::FrameHeader, modular::TransformId},
image::Rect,
util::{ShiftRightCeil, tracing_wrappers::*},
};
use super::{RctOp, RctPermutation};
#[instrument(level = "trace", err)]
fn check_equal_channels(
channels: &[(usize, ChannelInfo)],
first_chan: usize,
num: usize,
) -> Result<()> {
if first_chan + num > channels.len() {
return Err(Error::InvalidChannelRange(
first_chan,
first_chan + num,
channels.len(),
));
}
for inc in 1..num {
if !channels[first_chan]
.1
.is_equivalent(&channels[first_chan + inc].1)
{
return Err(Error::MixingDifferentChannels);
}
}
Ok(())
}
pub(super) fn meta_apply_single_transform(
transform: &headers::modular::Transform,
header: &headers::modular::GroupHeader,
channels: &mut Vec<(usize, ChannelInfo)>,
transform_steps: &mut Vec<TransformStep>,
mut add_transform_buffer: impl FnMut(ChannelInfo, String) -> usize,
) -> Result<()> {
match transform.id {
TransformId::Rct => {
let begin_channel = transform.begin_channel as usize;
let op = RctOp::from_u32(transform.rct_type % 7).unwrap();
let perm = RctPermutation::from_u32(transform.rct_type / 7)
.expect("header decoding should ensure rct_type < 42");
check_equal_channels(channels, begin_channel, 3)?;
let mut buf_in = [0; 3];
let buf_out = [
channels[begin_channel].0,
channels[begin_channel + 1].0,
channels[begin_channel + 2].0,
];
for i in 0..3 {
let c = &mut channels[begin_channel + i];
let mut info = c.1;
info.output_channel_idx = None;
c.0 = add_transform_buffer(
info,
format!(
"RCT (op {op:?} perm {perm:?}) starting at channel {begin_channel}, \
input {i}"
),
);
buf_in[i] = c.0;
}
transform_steps.push(TransformStep::Rct {
buf_out,
buf_in,
op,
perm,
});
trace!("applied RCT: {channels:?}");
}
TransformId::Squeeze => {
let steps = if transform.squeezes.is_empty() {
super::squeeze::default_squeeze(channels)
} else {
transform.squeezes.clone()
};
let mut transform_step_for_buf = HashMap::new();
for step in steps {
super::squeeze::check_squeeze_params(channels, &step)?;
let in_place = step.in_place;
let horizontal = step.horizontal;
let begin_channel = step.begin_channel as usize;
let num_channels = step.num_channels as usize;
let end_channel = begin_channel + num_channels;
let new_chan_offset = if in_place {
end_channel
} else {
channels.len()
};
for ic in 0..num_channels {
let chan = &channels[begin_channel + ic].1;
let new_shift = if let Some(shift) = chan.shift {
if shift.0 > 30 || shift.1 > 30 {
return Err(Error::TooManySqueezes);
}
if horizontal {
Some((shift.0 + 1, shift.1))
} else {
Some((shift.0, shift.1 + 1))
}
} else {
None
};
let w = chan.size.0;
let h = chan.size.1;
let (new_size_0, new_size_1) = if horizontal {
((w.div_ceil(2), h), (w - w.div_ceil(2), h))
} else {
((w, h.div_ceil(2)), (w, h - h.div_ceil(2)))
};
let new_0 = ChannelInfo {
output_channel_idx: None,
shift: new_shift,
size: new_size_0,
bit_depth: chan.bit_depth,
};
let buf_0 = add_transform_buffer(
new_0,
format!("Squeezed channel, original channel {}", begin_channel + ic),
);
let new_1 = ChannelInfo {
output_channel_idx: None,
shift: new_shift,
size: new_size_1,
bit_depth: chan.bit_depth,
};
let buf_1 = add_transform_buffer(
new_1,
format!("Squeeze residual, original channel {}", begin_channel + ic),
);
transform_step_for_buf.insert(buf_0, (transform_steps.len(), horizontal));
let buf_out = channels[begin_channel + ic].0;
if let Some(t) = transform_step_for_buf.get(&buf_out)
&& t.1 != horizontal
{
let (TransformStep::VSqueeze { buf_in_avg, .. }
| TransformStep::HSqueeze { buf_in_avg, .. }) = &mut transform_steps[t.0]
else {
unreachable!()
};
*buf_in_avg = Some([buf_0, buf_1]);
}
if horizontal {
transform_steps.push(TransformStep::HSqueeze {
buf_in: [buf_0, buf_1],
buf_out,
buf_in_avg: None,
});
} else {
transform_steps.push(TransformStep::VSqueeze {
buf_in: [buf_0, buf_1],
buf_out,
buf_in_avg: None,
});
}
channels[begin_channel + ic] = (buf_0, new_0);
channels.insert(new_chan_offset + ic, (buf_1, new_1));
trace!("applied squeeze: {channels:?}");
}
}
}
TransformId::Palette => {
let begin_channel = transform.begin_channel as usize;
let num_channels = transform.num_channels as usize;
let num_colors = transform.num_colors as usize;
let num_deltas = transform.num_deltas as usize;
let pred = Predictor::from_u32(transform.predictor_id)
.expect("header decoding should ensure a valid predictor");
check_equal_channels(channels, begin_channel, num_channels)?;
// We already checked the bit_depth for all channels from `begin_channel` is
// equal in the line above.
let bit_depth = channels[begin_channel].1.bit_depth;
let pchan_info = ChannelInfo {
output_channel_idx: None,
shift: None,
size: (num_colors + num_deltas, num_channels),
bit_depth,
};
let pchan = add_transform_buffer(
pchan_info,
format!(
"Palette for palette transform starting at channel {begin_channel} with \
{num_channels} channels"
),
);
let mut inchan_info = channels[begin_channel].1;
inchan_info.output_channel_idx = None;
let inchan = add_transform_buffer(
inchan_info,
format!(
"Pixel data for palette transform starting at channel {begin_channel} with \
{num_channels} channels",
),
);
transform_steps.push(TransformStep::Palette {
buf_in: inchan,
buf_pal: pchan,
buf_out: channels[begin_channel..(begin_channel + num_channels)]
.iter()
.map(|x| x.0)
.collect(),
num_colors,
num_deltas,
predictor: pred,
wp_header: header.wp_header.clone(),
});
channels.drain(begin_channel + 1..begin_channel + num_channels);
channels[begin_channel].0 = inchan;
channels.insert(0, (pchan, pchan_info));
trace!("applied palette: {channels:?}");
}
TransformId::Invalid => {
unreachable!("header decoding for invalid transforms should fail");
}
}
Ok(())
}
#[instrument(level = "trace", ret)]
pub fn meta_apply_transforms(
channels: &[ChannelInfo],
header: &headers::modular::GroupHeader,
) -> Result<(Vec<ModularBufferInfo>, Vec<TransformStep>)> {
let mut buffer_info = vec![];
let mut transform_steps = vec![];
// (buffer id, channel info)
let mut channels: Vec<_> = channels.iter().cloned().enumerate().collect();
// First, add all the pre-transform channels to the buffer list.
for chan in channels.iter() {
buffer_info.push(ModularBufferInfo {
info: chan.1,
coded_channel_id: -1,
description: format!(
"Input channel {}, size {}x{}",
chan.0, chan.1.size.0, chan.1.size.1
),
// To be filled by make_grids.
grid_kind: ModularGridKind::None,
grid_shape: (0, 0),
buffer_grid: vec![],
});
}
let mut add_transform_buffer = |info, description| {
buffer_info.push(ModularBufferInfo {
info,
coded_channel_id: -1,
description,
// To be filled by make_grids.
grid_kind: ModularGridKind::None,
grid_shape: (0, 0),
buffer_grid: vec![],
});
buffer_info.len() - 1
};
// Apply transforms to the channel list.
for transform in &header.transforms {
meta_apply_single_transform(
transform,
header,
&mut channels,
&mut transform_steps,
&mut add_transform_buffer,
)?;
}
// All the channels left over at the end of applying transforms are the channels that are
// actually coded.
for (chid, chan) in channels.iter().enumerate() {
buffer_info[chan.0].coded_channel_id = chid as isize;
}
#[cfg(feature = "tracing")]
for (i, transform) in transform_steps.iter().enumerate() {
trace!("Transform step {i}: {transform:?}");
}
Ok((buffer_info, transform_steps))
}
#[instrument(level = "trace", skip_all, ret)]
pub fn make_grids(
frame_header: &FrameHeader,
transform_steps: Vec<TransformStep>,
section_buffer_indices: &[Vec<usize>],
buffer_info: &mut Vec<ModularBufferInfo>,
color_channels: usize,
) -> Vec<TransformStepChunk> {
// Initialize grid sizes, starting from coded channels.
for i in section_buffer_indices[1].iter() {
buffer_info[*i].grid_kind = ModularGridKind::Lf;
}
for buffer_indices in section_buffer_indices.iter().skip(2) {
for i in buffer_indices.iter() {
buffer_info[*i].grid_kind = ModularGridKind::Hf;
}
}
trace!(?buffer_info, "post set grid kind for coded channels");
// Transforms can be un-applied in the opposite order they appear with in the array,
// so we can use that information to propagate grid kinds.
for step in transform_steps.iter().rev() {
match step {
TransformStep::Rct {
buf_in, buf_out, ..
} => {
let grid_in = buffer_info[buf_in[0]].grid_kind;
for i in 0..3 {
assert_eq!(grid_in, buffer_info[buf_in[i]].grid_kind);
}
for i in 0..3 {
buffer_info[buf_out[i]].grid_kind = grid_in;
}
}
TransformStep::Palette {
buf_in, buf_out, ..
} => {
for buf in buf_out.iter() {
buffer_info[*buf].grid_kind = buffer_info[*buf_in].grid_kind;
}
}
TransformStep::HSqueeze {
buf_in, buf_out, ..
}
| TransformStep::VSqueeze {
buf_in, buf_out, ..
} => {
let mut grid_kind = buffer_info[buf_in[0]]
.grid_kind
.max(buffer_info[buf_in[1]].grid_kind);
if grid_kind == ModularGridKind::None
&& !buffer_info[*buf_out]
.info
.is_meta_or_small(frame_header.group_dim())
{
grid_kind = ModularGridKind::Hf;
}
buffer_info[*buf_out].grid_kind = grid_kind;
}
TransformStep::Output { .. } => {
unreachable!()
}
}
}
// Set grid shapes.
for buf in buffer_info.iter_mut() {
buf.grid_shape = buf.grid_kind.grid_shape(frame_header);
}
trace!(?buffer_info, "post propagate grid kind");
let get_grid_indices = |shape: (usize, usize)| {
(0..shape.1).flat_map(move |y| (0..shape.0).map(move |x| (x as isize, y as isize)))
};
// Create grids.
for g in buffer_info.iter_mut() {
g.buffer_grid = get_grid_indices(g.grid_shape)
.map(|(x, y)| {
ModularBuffer::new(
g.get_grid_rect(frame_header, g.grid_kind, (x as usize, y as usize))
.size,
)
})
.collect();
}
trace!(?buffer_info, "with grids");
let add_transform_step =
|transform: &TransformStep,
grid_pos: (isize, isize),
grid_transform_steps: &mut Vec<TransformStepChunk>| {
let ts = grid_transform_steps.len();
grid_transform_steps.push(TransformStepChunk {
step: transform.clone(),
grid_pos: (grid_pos.0 as usize, grid_pos.1 as usize),
missing_final_deps: 0,
missing_deps: AtomicUsize::new(0),
});
ts
};
let add_grid_use = |ts: usize,
input_buffer_idx: usize,
output_grid_kind: ModularGridKind,
output_grid_shape: (usize, usize),
output_grid_pos: (isize, isize),
grid_transform_steps: &mut Vec<TransformStepChunk>,
buffer_info: &mut Vec<ModularBufferInfo>| {
let output_grid_size = (output_grid_shape.0 as isize, output_grid_shape.1 as isize);
if output_grid_pos.0 < 0
|| output_grid_pos.0 >= output_grid_size.0
|| output_grid_pos.1 < 0
|| output_grid_pos.1 >= output_grid_size.1
{
// Skip adding uses of non-existent grid positions.
return;
}
let output_grid_pos = (output_grid_pos.0 as usize, output_grid_pos.1 as usize);
let input_grid_pos =
buffer_info[input_buffer_idx].get_grid_idx(output_grid_kind, output_grid_pos);
let grid = &mut buffer_info[input_buffer_idx].buffer_grid[input_grid_pos];
if !grid.used_by_transforms_final.contains(&ts) {
grid_transform_steps[ts].missing_final_deps += 1;
grid.used_by_transforms_final.push(ts);
}
};
// Add grid-ed transforms.
let mut grid_transform_steps = vec![];
for transform in transform_steps {
match &transform {
TransformStep::Rct {
buf_in, buf_out, ..
} => {
// Easy case: we just depend on the 3 input buffers in the same location.
let out_kind = buffer_info[buf_out[0]].grid_kind;
let out_shape = buffer_info[buf_out[0]].grid_shape;
for (x, y) in get_grid_indices(out_shape) {
let ts = add_transform_step(&transform, (x, y), &mut grid_transform_steps);
for bin in buf_in {
add_grid_use(
ts,
*bin,
out_kind,
out_shape,
(x, y),
&mut grid_transform_steps,
buffer_info,
);
}
}
}
TransformStep::Palette {
buf_in,
buf_pal,
buf_out,
predictor,
..
} if predictor.requires_full_row() => {
// Delta palette with AverageAll or Weighted. Those are special, because we can
// only make progress one full image row at a time (since we need decoded values
// from the previous row or two rows).
let out_kind = buffer_info[buf_out[0]].grid_kind;
let out_shape = buffer_info[buf_out[0]].grid_shape;
let mut ts = 0;
for (x, y) in get_grid_indices(out_shape) {
if x == 0 {
ts = add_transform_step(&transform, (x, y), &mut grid_transform_steps);
add_grid_use(
ts,
*buf_pal,
out_kind,
out_shape,
(x, y),
&mut grid_transform_steps,
buffer_info,
);
}
add_grid_use(
ts,
*buf_in,
out_kind,
out_shape,
(x, y),
&mut grid_transform_steps,
buffer_info,
);
for out in buf_out.iter() {
add_grid_use(
ts,
*out,
out_kind,
out_shape,
(x, y - 1),
&mut grid_transform_steps,
buffer_info,
);
}
}
}
TransformStep::Palette {
buf_in,
buf_pal,
buf_out,
predictor,
..
} => {
// Maybe-delta palette: we depend on the palette and the input buffer in the same
// location. We may also depend on other grid positions in the output buffer,
// according to the used predictor.
let out_kind = buffer_info[buf_out[0]].grid_kind;
let out_shape = buffer_info[buf_out[0]].grid_shape;
for (x, y) in get_grid_indices(out_shape) {
let ts = add_transform_step(&transform, (x, y), &mut grid_transform_steps);
add_grid_use(
ts,
*buf_pal,
out_kind,
out_shape,
(x, y),
&mut grid_transform_steps,
buffer_info,
);
add_grid_use(
ts,
*buf_in,
out_kind,
out_shape,
(x, y),
&mut grid_transform_steps,
buffer_info,
);
let offsets = match predictor {
Predictor::Zero => [].as_slice(),
_ => &[(0, -1), (-1, 0), (-1, -1)],
};
for (dx, dy) in offsets {
for out in buf_out.iter() {
add_grid_use(
ts,
*out,
out_kind,
out_shape,
(x + dx, y + dy),
&mut grid_transform_steps,
buffer_info,
);
}
}
}
}
TransformStep::HSqueeze {
buf_in, buf_out, ..
} => {
let out_kind = buffer_info[*buf_out].grid_kind;
let out_shape = buffer_info[*buf_out].grid_shape;
for (x, y) in get_grid_indices(out_shape) {
let ts = add_transform_step(&transform, (x, y), &mut grid_transform_steps);
// Average and residuals from the same position
for bin in buf_in {
add_grid_use(
ts,
*bin,
out_kind,
out_shape,
(x, y),
&mut grid_transform_steps,
buffer_info,
);
}
// Next average
add_grid_use(
ts,
buf_in[0],
out_kind,
out_shape,
(x + 1, y),
&mut grid_transform_steps,
buffer_info,
);
// Previous decoded
add_grid_use(
ts,
*buf_out,
out_kind,
out_shape,
(x - 1, y),
&mut grid_transform_steps,
buffer_info,
);
}
}
TransformStep::VSqueeze {
buf_in, buf_out, ..
} => {
let out_kind = buffer_info[*buf_out].grid_kind;
let out_shape = buffer_info[*buf_out].grid_shape;
for (x, y) in get_grid_indices(out_shape) {
let ts = add_transform_step(&transform, (x, y), &mut grid_transform_steps);
// Average and residuals from the same position
for bin in buf_in {
add_grid_use(
ts,
*bin,
out_kind,
out_shape,
(x, y),
&mut grid_transform_steps,
buffer_info,
);
}
// Next average
add_grid_use(
ts,
buf_in[0],
out_kind,
out_shape,
(x, y + 1),
&mut grid_transform_steps,
buffer_info,
);
// Previous decoded
add_grid_use(
ts,
*buf_out,
out_kind,
out_shape,
(x, y - 1),
&mut grid_transform_steps,
buffer_info,
);
}
}
TransformStep::Output { .. } => {
unreachable!()
}
}
}
// Write on transform outputs which step produces them.
for (ts, step) in grid_transform_steps.iter().enumerate() {
for &(buf, grid) in step.outputs(buffer_info).iter() {
buffer_info[buf].buffer_grid[grid].produced_by_step = Some(ts);
}
}
// Add "output" transform steps.
// TODO(veluca): filter this by pipeline_used_channels.
for (idx, bi) in buffer_info.iter_mut().enumerate() {
if let Some(chan) = bi.info.output_channel_idx {
let (shift_x, shift_y) = bi.info.shift.unwrap_or((0, 0));
let gsx = frame_header.group_dim().shrc(shift_x);
let gsy = frame_header.group_dim().shrc(shift_y);
// If the channel has a grid size of Hf groups, then we can just output
// the entire buffer. Otherwise, we need to copy the region of the image
// that corresponds to a HF group.
let grid_dim = match bi.grid_kind {
ModularGridKind::None => frame_header.size().0.max(frame_header.size().1),
ModularGridKind::Hf => frame_header.group_dim(),
ModularGridKind::Lf => frame_header.lf_group_dim(),
};
let gtx = grid_dim.shrc(shift_x);
let gty = grid_dim.shrc(shift_y);
let group_shape = frame_header.size_groups();
for y in 0..group_shape.1 {
let gy = y * gsy / gty;
let gyo = y * gsy - gy * gty;
for x in 0..group_shape.0 {
let gx = x * gsx / gtx;
let gxo = x * gsx - gx * gtx;
let grid = gy * bi.grid_shape.0 + gx;
let rect = if gxo == 0 && gyo == 0 && gtx == gsx && gty == gsy {
None
} else {
let rect = Rect {
origin: (gxo, gyo),
size: (gsx, gsy),
};
Some(rect.clip(bi.buffer_grid[grid].size))
};
let chans = if chan == 0 && color_channels == 1 {
0..=2
} else {
chan..=chan
};
for c in chans {
bi.buffer_grid[grid]
.used_by_transforms_final
.push(grid_transform_steps.len());
grid_transform_steps.push(TransformStepChunk {
step: TransformStep::Output {
buf_in: idx,
rect,
group: y * frame_header.size_groups().0 + x,
channel: c,
},
grid_pos: (gx, gy),
missing_final_deps: 1,
missing_deps: AtomicUsize::new(0),
});
}
}
}
}
}
trace!(?grid_transform_steps, ?buffer_info);
grid_transform_steps
}