Improve memory usage of evaluator (#688)

This adds a new `final_round` parameter to the evaluator, which allows
us to determine when to keep the `filtered` data and when to keep the
compressed `idat_data`. By only keeping one of these we can
significantly improve memory usage.

Ref #683
This commit is contained in:
andrews05 2025-03-09 00:01:02 +13:00 committed by GitHub
parent 6cf42fbefa
commit 7650953faf
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 53 additions and 43 deletions

View file

@ -21,23 +21,18 @@ use crate::{atomicmin::AtomicMin, deflate, filters::RowFilter, png::PngImage, De
pub(crate) struct Candidate {
pub image: Arc<PngImage>,
pub idat_data: Vec<u8>,
pub filtered: Vec<u8>,
pub data: Vec<u8>,
pub data_is_compressed: bool,
pub estimated_output_size: usize,
pub filter: RowFilter,
// For determining tie-breaker
nth: usize,
}
impl Candidate {
/// Return an estimate of the output size which can help with evaluation of very small data
#[must_use]
pub fn estimated_output_size(&self) -> usize {
self.idat_data.len() + self.image.key_chunks_size()
}
fn cmp_key(&self) -> impl Ord {
(
self.estimated_output_size(),
self.estimated_output_size,
self.image.data.len(),
self.filter,
// Prefer the later image added (e.g. baseline, which is always added last)
@ -52,6 +47,7 @@ pub(crate) struct Evaluator {
filters: IndexSet<RowFilter>,
deflater: Deflaters,
optimize_alpha: bool,
final_round: bool,
nth: AtomicUsize,
executed: Arc<AtomicUsize>,
best_candidate_size: Arc<AtomicMin>,
@ -69,6 +65,7 @@ impl Evaluator {
filters: IndexSet<RowFilter>,
deflater: Deflaters,
optimize_alpha: bool,
final_round: bool,
) -> Self {
#[cfg(feature = "parallel")]
let eval_channel = unbounded();
@ -77,6 +74,7 @@ impl Evaluator {
filters,
deflater,
optimize_alpha,
final_round,
nth: AtomicUsize::new(0),
executed: Arc::new(AtomicUsize::new(0)),
best_candidate_size: Arc::new(AtomicMin::new(None)),
@ -127,6 +125,7 @@ impl Evaluator {
let filters = self.filters.clone();
let deflater = self.deflater;
let optimize_alpha = self.optimize_alpha;
let final_round = self.final_round;
let executed = self.executed.clone();
let best_candidate_size = self.best_candidate_size.clone();
let description = description.to_string();
@ -149,21 +148,23 @@ impl Evaluator {
let filtered = image.filter_image(filter, optimize_alpha);
let idat_data = deflater.deflate(&filtered, best_candidate_size.get());
if let Ok(idat_data) = idat_data {
let estimated_output_size = image.estimated_output_size(&idat_data);
// For the final round we need the IDAT data, otherwise the filtered data
let new = Candidate {
image: image.clone(),
idat_data,
filtered,
data: if final_round { idat_data } else { filtered },
data_is_compressed: final_round,
estimated_output_size,
filter,
nth,
};
let size = new.estimated_output_size();
best_candidate_size.set_min(size);
best_candidate_size.set_min(estimated_output_size);
trace!(
"Eval: {}-bit {:23} {:8} {} bytes",
image.ihdr.bit_depth,
description,
filter,
size
estimated_output_size
);
#[cfg(feature = "parallel")]

View file

@ -170,7 +170,7 @@ impl RawImage {
let mut png = PngData {
raw: result.image,
idat_data: result.idat_data,
idat_data: result.data,
aux_chunks,
frames: Vec::new(),
};
@ -355,11 +355,11 @@ fn optimize_png(
let max_size = if opts.force {
None
} else {
Some(png.estimated_output_size())
Some(png.raw.estimated_output_size(&png.idat_data))
};
if let Some(result) = optimize_raw(raw.clone(), &opts, deadline.clone(), max_size) {
png.raw = result.image;
png.idat_data = result.idat_data;
png.idat_data = result.data;
recompress_frames(png, &opts, deadline, result.filter)?;
postprocess_chunks(&mut png.aux_chunks, &png.raw.ihdr, &raw.ihdr);
}
@ -433,7 +433,13 @@ fn optimize_raw(
indexset! {RowFilter::None, RowFilter::Bigrams}
};
// This will collect all versions of images and pick one that compresses best
let eval = Evaluator::new(deadline.clone(), eval_filters.clone(), eval_deflater, false);
let eval = Evaluator::new(
deadline.clone(),
eval_filters.clone(),
eval_deflater,
false,
opts.deflate == eval_deflater,
);
let mut new_image = perform_reductions(image.clone(), opts, &deadline, &eval);
let eval_result = eval.get_best_candidate();
if let Some(ref result) = eval_result {
@ -464,7 +470,9 @@ fn optimize_raw(
(eval_result?, eval_deflater)
};
if max_size.map_or(true, |max_size| result.estimated_output_size() < max_size) {
if result.data_is_compressed
&& max_size.map_or(true, |max_size| result.estimated_output_size < max_size)
{
debug!("Found better result:");
debug!(" {}, f = {}", deflater, result.filter);
return Some(result);
@ -499,35 +507,36 @@ fn perform_trials(
filters,
eval_deflater,
opts.optimize_alpha,
opts.deflate == eval_deflater,
);
if let Some(result) = &eval_result {
eval.set_best_size(result.estimated_output_size());
eval.set_best_size(result.estimated_output_size);
}
eval.try_image(image.clone());
if let Some(result) = eval.get_best_candidate() {
eval_result = Some(result);
}
}
if opts.deflate == eval_deflater {
// No further compression required
return eval_result;
}
// We should have a result here - fail if not (e.g. deadline passed)
let mut result = eval_result?;
// Recompress with the main deflater
debug!("Trying filter {} with {}", result.filter, opts.deflate);
match opts.deflate.deflate(&result.filtered, max_size) {
Ok(idat_data) => {
result.idat_data = idat_data;
trace!("{} bytes", result.estimated_output_size());
}
Err(PngError::DeflatedDataTooLong(bytes)) => {
trace!(">{bytes} bytes");
}
Err(_) => (),
};
if !result.data_is_compressed {
// Compress with the main deflater
debug!("Trying filter {} with {}", result.filter, opts.deflate);
match opts.deflate.deflate(&result.data, max_size) {
Ok(idat_data) => {
result.estimated_output_size = result.image.estimated_output_size(&idat_data);
result.data = idat_data;
result.data_is_compressed = true;
trace!("{} bytes", result.estimated_output_size);
}
Err(PngError::DeflatedDataTooLong(bytes)) => {
trace!(">{bytes} bytes");
}
Err(_) => (),
};
}
return Some(result);
}
@ -545,7 +554,7 @@ fn perform_trials(
}
debug!("Trying {} filters with {}", filters.len(), opts.deflate);
let eval = Evaluator::new(deadline, filters, opts.deflate, opts.optimize_alpha);
let eval = Evaluator::new(deadline, filters, opts.deflate, opts.optimize_alpha, true);
if let Some(max_size) = max_size {
eval.set_best_size(max_size);
}

View file

@ -185,12 +185,6 @@ impl PngData {
})
}
/// Return an estimate of the output size which can help with evaluation of very small data
#[must_use]
pub fn estimated_output_size(&self) -> usize {
self.idat_data.len() + self.raw.key_chunks_size()
}
/// Format the `PngData` struct into a valid PNG bytestream
#[must_use]
pub fn output(&self) -> Vec<u8> {
@ -355,6 +349,12 @@ impl PngImage {
}
}
/// Return an estimate of the output size which can help with evaluation of very small data
#[must_use]
pub fn estimated_output_size(&self, idat_data: &[u8]) -> usize {
idat_data.len() + self.key_chunks_size()
}
/// Return an iterator over the scanlines of the image
#[inline]
#[must_use]