mirror of
https://github.com/storytold/photocraft.git
synced 2026-11-02 18:13:39 +00:00
perf(stack): cache bounded entropy terms across pixels (#9622)
* perf(stack): cache bounded entropy terms across pixels * test(stack): iterate entropy terms directly for Clippy * ci: rerun checks
This commit is contained in:
Khawar Zamman Wani
committed by
GitHub
GitHub
parent
66ac154d68
commit
6c6cfc243a
@@ -0,0 +1,104 @@
|
||||
//! Stack Mode Entropy timings and exact output hashes, including lookup preparation.
|
||||
//! Run the same example on both revisions: `cargo run --release -p photocraft-algo --example
|
||||
//! bench_stack_entropy -- [--quick] [--reps N] [--threads N] [--json path]`.
|
||||
//! Full cases use 24/36 MP; quick cases use 1/16 as many pixels. The deep-stack fallback is
|
||||
//! deliberately small. Input generation, output hashing and deallocation are outside timings.
|
||||
|
||||
use std::path::PathBuf;
|
||||
use std::time::Instant;
|
||||
|
||||
use photocraft_algo::stack::{Stat, combine};
|
||||
use photocraft_testkit::perf::{self, Summary};
|
||||
use serde_json::json;
|
||||
|
||||
struct Options {
|
||||
quick: bool,
|
||||
reps: usize,
|
||||
threads: Option<usize>,
|
||||
json: Option<PathBuf>,
|
||||
}
|
||||
|
||||
fn options() -> Result<Options, String> {
|
||||
let mut o = Options { quick: false, reps: 5, threads: None, json: None };
|
||||
let mut args = std::env::args().skip(1);
|
||||
while let Some(arg) = args.next() {
|
||||
match arg.as_str() {
|
||||
"--quick" => o.quick = true,
|
||||
"--reps" => o.reps = args.next().ok_or("--reps needs an integer")?.parse().map_err(|_| "invalid --reps")?,
|
||||
"--threads" => o.threads = Some(args.next().ok_or("--threads needs an integer")?.parse().map_err(|_| "invalid --threads")?),
|
||||
"--json" => o.json = Some(args.next().ok_or("--json needs a path")?.into()),
|
||||
_ => return Err(format!("unknown option: {arg}")),
|
||||
}
|
||||
}
|
||||
if !(1..=25).contains(&o.reps) || o.threads.is_some_and(|n| !(1..=64).contains(&n)) {
|
||||
return Err("expected 1..=25 repetitions and 1..=64 threads".into());
|
||||
}
|
||||
Ok(o)
|
||||
}
|
||||
|
||||
fn benchmark(o: &Options) -> Result<(), String> {
|
||||
let mut rows = Vec::new();
|
||||
let cases = [
|
||||
("3 layers opaque 24 MP", 3, 24_000_000, false),
|
||||
("7 layers varying alpha 24 MP", 7, 24_000_000, true),
|
||||
("3 layers opaque 36 MP", 3, 36_000_000, false),
|
||||
("257 layers varying alpha", 257, 8192, true),
|
||||
];
|
||||
for (name, count, pixels, masked) in cases {
|
||||
let pixels = if o.quick && pixels > 8192 { pixels / 16 } else { pixels };
|
||||
let layers: Vec<Vec<[f32; 4]>> = (0..count)
|
||||
.map(|k| {
|
||||
(0..pixels)
|
||||
.map(|i| {
|
||||
[
|
||||
((i * 31 + k * 17) % 256) as f32 / 255.0,
|
||||
((i / 7 + k * 29) % 32) as f32 / 31.0,
|
||||
((i + k * 13) % 1024) as f32 / 255.0 - 0.5,
|
||||
if masked && k > i % (count + 1) { 0.0 } else { 1.0 },
|
||||
]
|
||||
})
|
||||
.collect()
|
||||
})
|
||||
.collect();
|
||||
// Warm the worker pool and input pages before collecting samples.
|
||||
std::hint::black_box(combine(&layers, Stat::Entropy));
|
||||
let mut samples = Vec::with_capacity(o.reps);
|
||||
let mut hash = 0u64;
|
||||
for rep in 0..o.reps {
|
||||
let start = Instant::now();
|
||||
let out = std::hint::black_box(combine(&layers, Stat::Entropy));
|
||||
samples.push(start.elapsed().as_secs_f64() * 1000.0);
|
||||
if rep == 0 {
|
||||
hash = out.iter().flatten().fold(0xcbf29ce484222325, |h, v| (h ^ u64::from(v.to_bits())).wrapping_mul(0x100000001b3));
|
||||
}
|
||||
}
|
||||
let p50 = Summary::of(&samples).map_or(0.0, |s| s.p50);
|
||||
println!("{name}: {pixels} pixels, {p50:.3} ms p50, hash {hash:016x}, samples_ms={samples:?}");
|
||||
let mut row = perf::row(name, &samples, perf::process_peak_rss_bytes(), None);
|
||||
row["output_hash"] = json!(format!("{hash:016x}"));
|
||||
row["pixels"] = json!(pixels);
|
||||
rows.push(row);
|
||||
}
|
||||
if let Some(path) = &o.json {
|
||||
let report = perf::report("bench_stack_entropy", json!({"quick": o.quick, "threads": o.threads}), rows, None);
|
||||
std::fs::write(path, serde_json::to_vec_pretty(&report).map_err(|e| e.to_string())?).map_err(|e| format!("write {}: {e}", path.display()))?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn run() -> Result<(), String> {
|
||||
let o = options()?;
|
||||
#[cfg(not(target_arch = "wasm32"))]
|
||||
if let Some(threads) = o.threads {
|
||||
let pool = rayon::ThreadPoolBuilder::new().num_threads(threads).build().map_err(|e| e.to_string())?;
|
||||
return pool.install(|| benchmark(&o));
|
||||
}
|
||||
benchmark(&o)
|
||||
}
|
||||
|
||||
fn main() {
|
||||
if let Err(e) = run() {
|
||||
eprintln!("{e}");
|
||||
std::process::exit(1);
|
||||
}
|
||||
}
|
||||
@@ -37,8 +37,54 @@ impl EntropyScratch {
|
||||
}
|
||||
}
|
||||
|
||||
/// One shared table per combine, never per pixel or worker. Rows cover the contributing sample
|
||||
/// count, which can change at every pixel because of transparency or unequal layer lengths.
|
||||
/// The triangular layout stores log2(n) followed by -p*log2(p) for counts 1..=n. Keeping the
|
||||
/// division after the sum (rather than normalising each term) preserves the old output bits.
|
||||
struct EntropyTerms {
|
||||
values: Vec<f32>,
|
||||
}
|
||||
|
||||
impl EntropyTerms {
|
||||
// At most 130 KiB, independent of image size and total layer count. Deeper stacks retain
|
||||
// the original calculation; pixels with fewer contributors can still use a cached row.
|
||||
const MAX_SAMPLES: usize = 256;
|
||||
|
||||
fn new(samples: usize, pixels: usize) -> Self {
|
||||
let mut values = Vec::new();
|
||||
// Tiny images do not amortise preparation, especially with hundreds of layers.
|
||||
if samples < 2 || pixels < 256 {
|
||||
return Self { values };
|
||||
}
|
||||
let max = samples.min(Self::MAX_SAMPLES);
|
||||
// max is capped at 256 before any size arithmetic or allocation.
|
||||
let len = (max + 1) * (max + 2) / 2;
|
||||
if values.try_reserve_exact(len).is_err() {
|
||||
return Self { values };
|
||||
}
|
||||
values.extend_from_slice(&[0.0; 3]); // Unused rows for zero and one sample.
|
||||
for samples in 2..=max {
|
||||
let n = samples as f32;
|
||||
values.push(n.log2());
|
||||
for count in 1..=samples {
|
||||
let p = count as f32 / n;
|
||||
values.push(-p * p.log2());
|
||||
}
|
||||
}
|
||||
Self { values }
|
||||
}
|
||||
|
||||
fn row(&self, samples: usize) -> Option<&[f32]> {
|
||||
if !(2..=Self::MAX_SAMPLES).contains(&samples) {
|
||||
return None;
|
||||
}
|
||||
let start = samples * (samples + 1) / 2;
|
||||
self.values.get(start..start + samples + 1)
|
||||
}
|
||||
}
|
||||
|
||||
/// Shannon entropy (bits) of the 8-bit values in `v`, over log2(N). `v.len() >= 2`.
|
||||
fn entropy(v: &[f32], scratch: &mut EntropyScratch) -> f32 {
|
||||
fn entropy(v: &[f32], scratch: &mut EntropyScratch, terms: Option<&[f32]>) -> f32 {
|
||||
let n = v.len() as f32;
|
||||
for x in v {
|
||||
let bin = (x.clamp(0.0, 1.0) * 255.0).round() as u8;
|
||||
@@ -57,8 +103,11 @@ fn entropy(v: &[f32], scratch: &mut EntropyScratch) -> f32 {
|
||||
.touched
|
||||
.iter()
|
||||
.map(|&bin| {
|
||||
let p = scratch.counts[usize::from(bin)] as f32 / n;
|
||||
-p * p.log2()
|
||||
let count = scratch.counts[usize::from(bin)];
|
||||
terms.and_then(|row| row.get(count as usize)).copied().unwrap_or_else(|| {
|
||||
let p = count as f32 / n;
|
||||
-p * p.log2()
|
||||
})
|
||||
})
|
||||
.sum();
|
||||
// Reset only the bins used, so the 256-entry array is never cleared wholesale.
|
||||
@@ -66,7 +115,7 @@ fn entropy(v: &[f32], scratch: &mut EntropyScratch) -> f32 {
|
||||
scratch.counts[usize::from(bin)] = 0;
|
||||
}
|
||||
scratch.touched.clear();
|
||||
h / n.log2()
|
||||
h / terms.and_then(|row| row.first()).copied().unwrap_or_else(|| n.log2())
|
||||
}
|
||||
|
||||
/// `x` to the power `k` by repeated squaring. Every step is one IEEE 754 multiplication, so the
|
||||
@@ -115,7 +164,7 @@ fn stat(v: &mut [f32], s: Stat, scratch: &mut EntropyScratch) -> f32 {
|
||||
if v.len() < 2 {
|
||||
return 0.0;
|
||||
}
|
||||
return entropy(v, scratch);
|
||||
return entropy(v, scratch, None);
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
@@ -153,7 +202,11 @@ pub fn combine(layers: &[Vec<[f32; 4]>], s: Stat) -> Vec<[f32; 4]> {
|
||||
// serial loop (Mean, Variance, StdDev, Skewness, Kurtosis); this alone cuts Mean's
|
||||
// instructions by a fifth.
|
||||
match s {
|
||||
Stat::Entropy => combine_with(layers, |v, sc| stat(v, Stat::Entropy, sc)),
|
||||
Stat::Entropy => {
|
||||
let pixels = layers.iter().map(Vec::len).max().unwrap_or(0);
|
||||
let terms = EntropyTerms::new(layers.len(), pixels);
|
||||
combine_with(layers, |v, sc| if v.len() < 2 { 0.0 } else { entropy(v, sc, terms.row(v.len())) })
|
||||
}
|
||||
Stat::Kurtosis => combine_with(layers, |v, sc| stat(v, Stat::Kurtosis, sc)),
|
||||
Stat::Maximum => combine_with(layers, |v, sc| stat(v, Stat::Maximum, sc)),
|
||||
Stat::Mean => combine_with(layers, |v, sc| stat(v, Stat::Mean, sc)),
|
||||
|
||||
@@ -323,6 +323,102 @@ fn entropy_is_deterministic() {
|
||||
}
|
||||
}
|
||||
|
||||
/// The pre-lookup calculation, summing the 8-bit bins in ascending order. This deliberately
|
||||
/// does not call `entropy` or use its lookup/scratch, so it catches changes in rounding too.
|
||||
fn entropy_reference(v: &[f32]) -> f32 {
|
||||
if v.len() < 2 {
|
||||
return 0.0;
|
||||
}
|
||||
let mut counts = [0u32; 256];
|
||||
for x in v {
|
||||
counts[(x.clamp(0.0, 1.0) * 255.0).round() as usize] += 1;
|
||||
}
|
||||
let n = v.len() as f32;
|
||||
let h: f32 = counts
|
||||
.iter()
|
||||
.filter(|&&count| count > 0)
|
||||
.map(|&count| {
|
||||
let p = count as f32 / n;
|
||||
-p * p.log2()
|
||||
})
|
||||
.sum();
|
||||
h / n.log2()
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn entropy_terms_preserve_each_operation_and_bound_storage() {
|
||||
let terms = EntropyTerms::new(usize::MAX, usize::MAX);
|
||||
let max = EntropyTerms::MAX_SAMPLES;
|
||||
assert_eq!(terms.values.len(), (max + 1) * (max + 2) / 2);
|
||||
assert!(terms.values.len() * size_of::<f32>() < 130 * 1024);
|
||||
for samples in 2..=max {
|
||||
let row = terms.row(samples).unwrap();
|
||||
let n = samples as f32;
|
||||
assert_eq!(row[0].to_bits(), n.log2().to_bits());
|
||||
for (count, &term) in row.iter().enumerate().skip(1) {
|
||||
let p = count as f32 / n;
|
||||
assert_eq!(term.to_bits(), (-p * p.log2()).to_bits(), "{count}/{samples}");
|
||||
}
|
||||
}
|
||||
for samples in [0, 1, max + 1, usize::MAX] {
|
||||
assert!(terms.row(samples).is_none());
|
||||
}
|
||||
assert!(EntropyTerms::new(1, usize::MAX).values.is_empty());
|
||||
assert!(EntropyTerms::new(usize::MAX, 255).values.is_empty());
|
||||
assert!(EntropyTerms { values: Vec::new() }.row(2).is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn entropy_lookup_and_fallback_match_reference_bits() {
|
||||
let terms = EntropyTerms::new(EntropyTerms::MAX_SAMPLES, CHUNK);
|
||||
let mut scratch = EntropyScratch::new();
|
||||
let special = [f32::NAN, f32::NEG_INFINITY, f32::INFINITY, -3.0, -0.0, 0.0, 0.5 / 255.0, 0.5, 1.0, 4.0];
|
||||
// Move repeatedly between counts and distributions to exercise scratch reset, duplicate
|
||||
// bins, every cached denominator and the uncached side of the table boundary.
|
||||
for samples in (2..=EntropyTerms::MAX_SAMPLES + 2).chain([513, 3, 256, 2]) {
|
||||
for pattern in 0..4 {
|
||||
let v: Vec<f32> = (0..samples)
|
||||
.map(|i| match pattern {
|
||||
0 => 0.5,
|
||||
1 => (i % 256) as f32 / 255.0,
|
||||
2 => ((i * 137 + samples * 29) % 65536) as f32 / 65535.0,
|
||||
_ => special[i % special.len()],
|
||||
})
|
||||
.collect();
|
||||
let expected = entropy_reference(&v).to_bits();
|
||||
assert_eq!(entropy(&v, &mut scratch, terms.row(samples)).to_bits(), expected, "cached: n={samples}, pattern={pattern}");
|
||||
assert_eq!(entropy(&v, &mut scratch, None).to_bits(), expected, "fallback: n={samples}, pattern={pattern}");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn entropy_combine_matches_reference_with_changing_contributor_counts() {
|
||||
for (count, len) in [(19, 2 * CHUNK + 17), (EntropyTerms::MAX_SAMPLES + 1, 513)] {
|
||||
let layers: Vec<Vec<[f32; 4]>> = (0..count)
|
||||
.map(|k| {
|
||||
(0..len.saturating_sub(k % 13))
|
||||
.map(|i| {
|
||||
[
|
||||
((i * 31 + k * 7) % 256) as f32 / 255.0,
|
||||
((i * 17 + k) % 65536) as f32 / 65535.0,
|
||||
if (i + k).is_multiple_of(23) { f32::NAN } else { (i % 7) as f32 - 2.0 },
|
||||
if k > i % (count + 1) { 0.0 } else { 0.5 },
|
||||
]
|
||||
})
|
||||
.collect()
|
||||
})
|
||||
.collect();
|
||||
let expected = bits(&combine_with(&layers, |v, _| entropy_reference(v)));
|
||||
assert_eq!(bits(&combine(&layers, Stat::Entropy)), expected);
|
||||
#[cfg(not(target_arch = "wasm32"))]
|
||||
for threads in [1, 3] {
|
||||
let pool = rayon::ThreadPoolBuilder::new().num_threads(threads).build().unwrap();
|
||||
assert_eq!(pool.install(|| bits(&combine(&layers, Stat::Entropy))), expected, "{threads} threads");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Chunking is invisible: a multi-chunk image (with a ragged last chunk and layers of unequal
|
||||
/// length) gives, pixel by pixel, what combining that pixel on its own gives.
|
||||
#[test]
|
||||
|
||||
Reference in New Issue
Block a user