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.
#![allow(clippy::needless_range_loop)]
use std::any::Any;
use row_buffers::RowBuffer;
use crate::api::JxlOutputBuffer;
use crate::error::Result;
use crate::image::{Image, ImageDataType, OwnedRawImage, Rect};
use crate::render::MAX_BORDER;
use crate::render::buffer_splitter::{BufferSplitter, SaveStageBufferInfo};
use crate::render::internal::Stage;
use crate::util::{ShiftRightCeil, tracing_wrappers::*};
use super::RenderPipeline;
use super::internal::{RenderPipelineShared, RunInOutStage, RunInPlaceStage};
mod helpers;
mod render_group;
pub(super) mod row_buffers;
mod run_stage;
mod save;
struct InputBuffer {
// One buffer per channel.
data: Vec<Option<OwnedRawImage>>,
completed_passes: usize,
}
pub struct LowMemoryRenderPipeline {
shared: RenderPipelineShared<RowBuffer>,
input_buffers: Vec<InputBuffer>,
row_buffers: Vec<Vec<RowBuffer>>,
save_buffer_info: Vec<Option<SaveStageBufferInfo>>,
// The input buffer that each channel of each stage should use.
// This is indexed both by stage index (0 corresponds to input data, 1 to stage[0], etc) and by
// index *of those channels that are used*.
stage_input_buffer_index: Vec<Vec<(usize, usize)>>,
// Tracks whether we already rendered the padding around the core frame (if any).
padding_was_rendered: bool,
// The amount of pixels that each stage needs to *output* around the current group to
// run future stages correctly.
stage_output_border_pixels: Vec<(usize, usize)>,
// The amount of pixels that we need to read (for every channel) in non-edge groups to run all
// stages correctly.
input_border_pixels: Vec<(usize, usize)>,
has_nontrivial_border: bool,
// For every stage, the downsampling level of *any* channel that the stage uses at that point.
// Note that this must be equal across all the used channels.
downsampling_for_stage: Vec<(usize, usize)>,
// Local states of each stage, if any.
local_states: Vec<Option<Box<dyn Any>>>,
// Pre-filled opaque alpha buffers for stages that need fill_opaque_alpha.
// Indexed by stage index; None if stage doesn't need alpha fill.
opaque_alpha_buffers: Vec<Option<RowBuffer>>,
// Sorted indices to call get_distinct_indices.
sorted_buffer_indices: Vec<Vec<(usize, usize, usize)>>,
// For each channel, buffers that could be reused to store group data for that channel.
scratch_channel_buffers: Vec<Vec<OwnedRawImage>>,
}
impl LowMemoryRenderPipeline {
// TODO(veluca): most of this logic will need to change to ensure better cache utilization and
// lower memory usage.
fn render_with_new_group(
&mut self,
new_group_id: usize,
buffer_splitter: &mut BufferSplitter,
) -> Result<()> {
let (gx, gy) = self.shared.group_position(new_group_id);
// We put groups that are 2 afar here, because even if they could not have become
// renderable, they might have become freeable.
let mut possible_groups = vec![];
for dy in -2..=2 {
let igy = gy as isize + dy;
if igy < 0 || igy >= self.shared.group_count.1 as isize {
continue;
}
for dx in -2..=2 {
let igx = gx as isize + dx;
if igx < 0 || igx >= self.shared.group_count.0 as isize {
continue;
}
possible_groups.push(igy as usize * self.shared.group_count.0 + igx as usize);
}
}
// First, render all groups that have made progress; only check those that *could* have
// made progress.
for g in possible_groups.iter().copied() {
let ready_passes = self.shared.group_chan_ready_passes[g]
.iter()
.copied()
.min()
.unwrap();
if self.input_buffers[g].completed_passes < ready_passes {
let (gx, gy) = self.shared.group_position(g);
let mut fully_ready_passes = ready_passes;
// Here we assume that we never need more than one group worth of border.
if self.has_nontrivial_border {
for dy in -1..=1 {
let igy = gy as isize + dy;
if igy < 0 || igy >= self.shared.group_count.1 as isize {
continue;
}
for dx in -1..=1 {
let igx = gx as isize + dx;
if igx < 0 || igx >= self.shared.group_count.0 as isize {
continue;
}
let ig = (igy as usize) * self.shared.group_count.0 + igx as usize;
let ready_passes = self.shared.group_chan_ready_passes[ig]
.iter()
.copied()
.min()
.unwrap();
fully_ready_passes = fully_ready_passes.min(ready_passes);
}
}
}
if self.input_buffers[g].completed_passes >= fully_ready_passes {
continue;
}
debug!(
"new ready passes for group {gx},{gy} ({} completed, \
{ready_passes} ready, {fully_ready_passes} ready including neighbours)",
self.input_buffers[g].completed_passes
);
// Prepare output buffers for the group.
let (origin, size) = if let Some(e) = self.shared.extend_stage_index {
let Stage::Extend(e) = &self.shared.stages[e] else {
unreachable!("extend stage is not an extend stage");
};
(e.frame_origin, e.image_size)
} else {
((0, 0), self.shared.input_size)
};
let gsz = (
1 << self.shared.log_group_size,
1 << self.shared.log_group_size,
);
let rect_to_render = Rect {
size: gsz,
origin: (gsz.0 * gx, gsz.1 * gy),
};
let mut local_buffers = buffer_splitter.get_local_buffers(
&self.save_buffer_info,
rect_to_render,
false,
self.shared.input_size,
size,
origin,
);
self.render_group((gx, gy), &mut local_buffers)?;
self.input_buffers[g].completed_passes = fully_ready_passes;
}
}
// Clear buffers that will not be used again.
for g in possible_groups.iter().copied() {
let (gx, gy) = self.shared.group_position(g);
let mut neigh_complete_passes = self.input_buffers[g].completed_passes;
if self.has_nontrivial_border {
for dy in -1..=1 {
let igy = gy as isize + dy;
if igy < 0 || igy >= self.shared.group_count.1 as isize {
continue;
}
for dx in -1..=1 {
let igx = gx as isize + dx;
if igx < 0 || igx >= self.shared.group_count.0 as isize {
continue;
}
let ig = (igy as usize) * self.shared.group_count.0 + igx as usize;
neigh_complete_passes = self.input_buffers[ig]
.completed_passes
.min(neigh_complete_passes);
}
}
}
if self.shared.num_passes <= neigh_complete_passes {
for (c, b) in self.input_buffers[g].data.iter_mut().enumerate() {
if let Some(b) = std::mem::take(b) {
self.scratch_channel_buffers[c].push(b);
}
}
}
}
Ok(())
}
}
impl RenderPipeline for LowMemoryRenderPipeline {
type Buffer = RowBuffer;
fn new_from_shared(shared: RenderPipelineShared<Self::Buffer>) -> Result<Self> {
let mut input_buffers = vec![];
for _ in 0..shared.group_chan_ready_passes.len() {
input_buffers.push(InputBuffer {
data: vec![],
completed_passes: 0,
});
for _ in 0..shared.group_chan_ready_passes[0].len() {
input_buffers.last_mut().unwrap().data.push(None);
}
}
let nc = shared.channel_info[0].len();
let mut previous_inout: Vec<_> = (0..nc).map(|x| (0usize, x)).collect();
let mut stage_input_buffer_index = vec![];
let mut next_border_and_cur_downsample = vec![vec![]];
for ci in shared.channel_info[0].iter() {
next_border_and_cur_downsample[0].push((0, ci.downsample));
}
// For each stage, compute in which stage its input was buffered (the previous InOut
// stage). Also, compute for each InOut stage and channel the border with which the stage
// output is used; this will used to allocate buffers of the correct size.
for (i, stage) in shared.stages.iter().enumerate() {
stage_input_buffer_index.push(previous_inout.clone());
next_border_and_cur_downsample.push(vec![]);
if let Stage::InOut(p) = stage {
for (chan, (ps, pc)) in previous_inout.iter_mut().enumerate() {
if !p.uses_channel(chan) {
continue;
}
next_border_and_cur_downsample[*ps][*pc].0 = p.border().1;
*ps = i + 1;
*pc = next_border_and_cur_downsample[i + 1].len();
next_border_and_cur_downsample[i + 1]
.push((0, shared.channel_info[i + 1][chan].downsample));
}
}
}
let mut initial_buffers = vec![];
for chan in 0..nc {
initial_buffers.push(RowBuffer::new(
shared.channel_info[0][chan].ty.unwrap(),
next_border_and_cur_downsample[0][chan].0 as usize,
0,
shared.chunk_size >> shared.channel_info[0][chan].downsample.0,
)?);
}
let mut row_buffers = vec![initial_buffers];
// Allocate buffers.
for (i, stage) in shared.stages.iter().enumerate() {
let mut stage_buffers = vec![];
for (next_y_border, (dsx, _)) in next_border_and_cur_downsample[i + 1].iter() {
stage_buffers.push(RowBuffer::new(
stage.output_type().unwrap(),
*next_y_border as usize,
stage.shift().1 as usize,
shared.chunk_size >> *dsx,
)?);
}
row_buffers.push(stage_buffers);
}
// Compute information to be used to compute sub-rects for "save" stages to operate on
// rects.
let mut save_buffer_info = vec![];
'stage: for (i, (s, ci)) in shared
.stages
.iter()
.zip(shared.channel_info.iter())
.enumerate()
{
let Stage::Save(s) = s else {
continue;
};
for (c, ci) in ci.iter().enumerate() {
if s.uses_channel(c) {
let info = SaveStageBufferInfo {
downsample: ci.downsample,
orientation: s.orientation,
byte_size: s.data_format.bytes_per_sample() * s.output_channels(),
after_extend: shared.extend_stage_index.is_some_and(|e| i > e),
};
while save_buffer_info.len() <= s.output_buffer_index {
save_buffer_info.push(None);
}
save_buffer_info[s.output_buffer_index] = Some(info);
continue 'stage;
}
}
}
// Compute the amount of border pixels needed per channel, per stage.
let mut border_pixels = vec![(0usize, 0usize); nc];
let mut border_pixels_per_stage = vec![];
for s in shared.stages.iter().rev() {
let mut stage_max = (0, 0);
for (c, bp) in border_pixels.iter_mut().enumerate() {
if !s.uses_channel(c) {
continue;
}
stage_max.0 = stage_max.0.max(bp.0);
stage_max.1 = stage_max.1.max(bp.1);
bp.0 = bp.0.shrc(s.shift().0) + s.border().0 as usize;
bp.1 = bp.1.shrc(s.shift().1) + s.border().1 as usize;
}
border_pixels_per_stage.push(stage_max);
}
border_pixels_per_stage.reverse();
assert!(border_pixels_per_stage[0].0 <= MAX_BORDER);
let downsampling_for_stage: Vec<_> = shared
.stages
.iter()
.zip(shared.channel_info.iter())
.map(|(s, ci)| {
let dowsamplings: Vec<_> = (0..nc)
.filter_map(|c| {
if s.uses_channel(c) {
Some(ci[c].downsample)
} else {
None
}
})
.collect();
for &d in dowsamplings.iter() {
assert_eq!(d, dowsamplings[0]);
}
(dowsamplings[0].0 as usize, dowsamplings[0].1 as usize)
})
.collect();
// Create opaque alpha buffers for save stages that need fill_opaque_alpha
let mut opaque_alpha_buffers = vec![];
for (i, stage) in shared.stages.iter().enumerate() {
if let Stage::Save(s) = stage {
if s.fill_opaque_alpha {
let (dx, _dy) = downsampling_for_stage[i];
let row_len = shared.chunk_size >> dx;
let fill_pattern = s.data_format.opaque_alpha_bytes();
let buf =
RowBuffer::new_filled(s.data_format.data_type(), row_len, &fill_pattern)?;
opaque_alpha_buffers.push(Some(buf));
} else {
opaque_alpha_buffers.push(None);
}
} else {
opaque_alpha_buffers.push(None);
}
}
let default_channels: Vec<usize> = (0..nc).collect();
for (s, ibi) in stage_input_buffer_index.iter_mut().enumerate() {
let mut filtered = vec![];
// For SaveStage, use s.channels to get correct output ordering (e.g., BGRA).
let channels = if let Stage::Save(save_stage) = &shared.stages[s] {
save_stage.channels.as_slice()
} else {
default_channels.as_slice()
};
for &c in channels {
if shared.stages[s].uses_channel(c) {
filtered.push(ibi[c]);
}
}
*ibi = filtered;
}
let sorted_buffer_indices = (0..shared.stages.len())
.map(|s| {
let mut v: Vec<_> = stage_input_buffer_index[s]
.iter()
.enumerate()
.map(|(i, (outer, inner))| (*outer, *inner, i))
.collect();
v.sort();
v
})
.collect();
Ok(Self {
input_buffers,
stage_input_buffer_index,
row_buffers,
padding_was_rendered: false,
save_buffer_info,
stage_output_border_pixels: border_pixels_per_stage,
has_nontrivial_border: border_pixels.iter().any(|x| *x != (0, 0)),
input_border_pixels: border_pixels,
local_states: shared
.stages
.iter()
.map(|x| x.init_local_state(0)) // Thread index 0 for single-threaded execution
.collect::<Result<_>>()?,
shared,
downsampling_for_stage,
opaque_alpha_buffers,
sorted_buffer_indices,
scratch_channel_buffers: (0..nc).map(|_| vec![]).collect(),
})
}
#[instrument(skip_all, err)]
fn get_buffer<T: ImageDataType>(&mut self, channel: usize) -> Result<Image<T>> {
if let Some(b) = self.scratch_channel_buffers[channel].pop() {
return Ok(Image::from_raw(b));
}
let sz = self.shared.group_size_for_channel(channel, T::DATA_TYPE_ID);
Image::<T>::new(sz)
}
fn set_buffer_for_group<T: ImageDataType>(
&mut self,
channel: usize,
group_id: usize,
num_passes: usize,
buf: Image<T>,
buffer_splitter: &mut BufferSplitter,
) -> Result<()> {
debug!(
"filling data for group {}, channel {}, using type {:?}",
group_id,
channel,
T::DATA_TYPE_ID,
);
self.input_buffers[group_id].data[channel] = Some(buf.into_raw());
self.shared.group_chan_ready_passes[group_id][channel] += num_passes;
self.render_with_new_group(group_id, buffer_splitter)
}
fn check_buffer_sizes(&self, buffers: &mut [Option<JxlOutputBuffer>]) -> Result<()> {
// Check that buffer sizes are correct.
let mut size = self.shared.input_size;
for (i, s) in self.shared.stages.iter().enumerate() {
match s {
Stage::Extend(e) => size = e.image_size,
Stage::Save(s) => {
let (dx, dy) = self.downsampling_for_stage[i];
s.check_buffer_size(
(size.0 >> dx, size.1 >> dy),
buffers[s.output_buffer_index].as_ref(),
)?
}
_ => {}
}
}
Ok(())
}
fn render_outside_frame(&mut self, buffer_splitter: &mut BufferSplitter) -> Result<()> {
if self.shared.extend_stage_index.is_none() || self.padding_was_rendered {
return Ok(());
}
self.padding_was_rendered = true;
// TODO(veluca): consider pre-computing those strips at pipeline construction and making
// smaller strips.
let e = self.shared.extend_stage_index.unwrap();
let Stage::Extend(e) = &self.shared.stages[e] else {
unreachable!("extend stage is not an extend stage");
};
let frame_end = (
e.frame_origin.0 + self.shared.input_size.0 as isize,
e.frame_origin.1 + self.shared.input_size.1 as isize,
);
// Split the full image area in 4 strips: left and right of the frame, and above and below.
// We divide each part further in strips of width self.shared.chunk_size.
let mut strips = vec![];
// Above (including left and right)
if e.frame_origin.1 > 0 {
let xend = e.image_size.0;
let yend = (e.frame_origin.1 as usize).min(e.image_size.1);
for x in (0..xend).step_by(self.shared.chunk_size) {
let xe = (x + self.shared.chunk_size).min(xend);
strips.push((x..xe, 0..yend));
}
}
// Below
if frame_end.1 < e.image_size.1 as isize {
let ystart = frame_end.1.max(0) as usize;
let yend = e.image_size.1;
let xend = e.image_size.0;
for x in (0..xend).step_by(self.shared.chunk_size) {
let xe = (x + self.shared.chunk_size).min(xend);
strips.push((x..xe, ystart..yend));
}
}
// Left
if e.frame_origin.0 > 0 {
let ystart = e.frame_origin.1.max(0) as usize;
let yend = (frame_end.1 as usize).min(e.image_size.1);
let xend = (e.frame_origin.0 as usize).min(e.image_size.0);
for x in (0..xend).step_by(self.shared.chunk_size) {
let xe = (x + self.shared.chunk_size).min(xend);
strips.push((x..xe, ystart..yend));
}
}
// Right
if frame_end.0 < e.image_size.0 as isize {
let xstart = frame_end.0.max(0) as usize;
let xend = e.image_size.0;
let ystart = e.frame_origin.1.max(0) as usize;
let yend = (frame_end.1 as usize).min(e.image_size.1);
for x in (xstart..xend).step_by(self.shared.chunk_size) {
let xe = (x + self.shared.chunk_size).min(xend);
strips.push((x..xe, ystart..yend));
}
}
let full_image_size = e.image_size;
for (xrange, yrange) in strips {
let rect_to_render = Rect {
origin: (xrange.start, yrange.start),
size: (xrange.clone().count(), yrange.clone().count()),
};
if rect_to_render.size.0 == 0 || rect_to_render.size.1 == 0 {
continue;
}
let mut local_buffers = buffer_splitter.get_local_buffers(
&self.save_buffer_info,
rect_to_render,
true,
full_image_size,
full_image_size,
(0, 0),
);
self.render_outside_frame(xrange, yrange, &mut local_buffers)?;
}
Ok(())
}
fn box_inout_stage<S: super::RenderPipelineInOutStage>(
stage: S,
) -> Box<dyn RunInOutStage<Self::Buffer>> {
Box::new(stage)
}
fn box_inplace_stage<S: super::RenderPipelineInPlaceStage>(
stage: S,
) -> Box<dyn RunInPlaceStage<Self::Buffer>> {
Box::new(stage)
}
}