perf(io): Interleave PSD pixels in the final buffer

Remove intermediate band buffers while retaining byte-for-byte equivalence with the previous path at every sample depth.
This commit is contained in:
Rene Leonhardt
2026-10-19 11:31:59 +02:00
parent d95fdf3181
commit 76a6405f9b
2 changed files with 139 additions and 41 deletions
+101 -41
View File
@@ -43,40 +43,43 @@ pub fn interleave(planes: &[Option<&[u8]>], fill: &[Vec<u8>], n: usize, s: Sampl
let ch = planes.len();
let row = ch * bps;
let mut out = vec![0u8; n * row];
// One band of rows per core, written straight into its slice of the shared output: the
// per-pixel channel loop stays, but a band holds ~1/cores of a 20 MP image's iterations
// instead of all of them.
if out.is_empty() {
return out;
}
let bs = bands(n);
let parts: Vec<Vec<u8>> = par_map(bs.clone(), |band| {
let mut part = vec![0u8; band.len() * row];
for (i, p) in band.clone().enumerate() {
for (c, plane) in planes.iter().enumerate() {
let dst = &mut part[(i * ch + c) * bps..(i * ch + c + 1) * bps];
match plane {
Some(pl) if pl.len() >= (p + 1) * bps => {
let src = &pl[p * bps..(p + 1) * bps];
for k in 0..bps {
dst[k] = src[bps - 1 - k];
}
if cfg!(target_endian = "big") {
dst.copy_from_slice(src);
}
if invert[c] {
invert_sample(dst, s);
}
}
_ => dst.copy_from_slice(&fill[c]),
}
let step = bs.first().map_or(n, |b| b.len()) * row;
par_map(bs.into_iter().zip(out.chunks_mut(step)).collect(), |(band, part)| {
for (c, plane) in planes.iter().enumerate() {
let src = plane.and_then(|p| p.get(band.start * bps..)).unwrap_or_default();
let inv = src.len() >= bps && invert[c];
let fill = fill.get(c).map(Vec::as_slice).unwrap_or_default();
match s {
SampleType::U8 => write_channel(part, ch, c, src, fill, |[v]| [if inv { 255 - v } else { v }]),
SampleType::U16 => write_channel(part, ch, c, src, fill, |v| {
let v = u16::from_be_bytes(v);
(if inv { 65535 - v } else { v }).to_ne_bytes()
}),
SampleType::F32 => write_channel(part, ch, c, src, fill, |v| {
// Preserve float bits when no inversion is requested.
if inv { (1.0 - f32::from_be_bytes(v)).to_ne_bytes() } else { u32::from_be_bytes(v).to_ne_bytes() }
}),
}
}
part
});
for (band, part) in bs.into_iter().zip(parts) {
out[band.start * row..band.end * row].copy_from_slice(&part);
}
out
}
fn write_channel<const N: usize>(out: &mut [u8], ch: usize, c: usize, src: &[u8], fill: &[u8], decode: impl Fn([u8; N]) -> [u8; N]) {
let mut samples = src.as_chunks::<N>().0.iter();
for pixel in out.chunks_exact_mut(ch * N) {
let dst = &mut pixel[c * N..(c + 1) * N];
match samples.next() {
Some(sample) => dst.copy_from_slice(&decode(*sample)),
None => dst.copy_from_slice(fill),
}
}
}
/// Splits interleaved native-endian bytes into big-endian planes.
pub fn deinterleave(bytes: &[u8], ch: usize, s: SampleType, invert: &[bool]) -> Vec<Vec<u8>> {
let bps = s.bytes();
@@ -155,6 +158,7 @@ pub fn zero_sample(s: SampleType) -> Vec<u8> {
}
/// Quantizes a normalized value to big-endian bytes.
#[inline]
pub fn encode_be(v: f32, s: SampleType, out: &mut Vec<u8>) {
match s {
SampleType::U8 => out.push((v.clamp(0.0, 1.0) * 255.0 + 0.5) as u8),
@@ -217,21 +221,29 @@ pub fn bands(n: usize) -> Vec<std::ops::Range<usize>> {
mod tests {
use super::*;
#[test]
fn banded_shuffles_match_a_sequential_reference() {
fn seq_interleave(planes: &[Option<&[u8]>], fill: &[Vec<u8>], n: usize, s: SampleType, invert: &[bool]) -> Vec<u8> {
let bps = s.bytes();
let ch = planes.len();
let mut out = vec![0u8; n * ch * bps];
for (c, plane) in planes.iter().enumerate() {
for i in 0..n {
let dst = &mut out[(i * ch + c) * bps..(i * ch + c + 1) * bps];
pub(super) fn previous_interleave(planes: &[Option<&[u8]>], fill: &[Vec<u8>], n: usize, s: SampleType, invert: &[bool]) -> Vec<u8> {
let bps = s.bytes();
let ch = planes.len();
let row = ch * bps;
let mut out = vec![0u8; n * row];
// One band of rows per core, written straight into its slice of the shared output: the
// per-pixel channel loop stays, but a band holds ~1/cores of a 20 MP image's iterations
// instead of all of them.
let bs = bands(n);
let parts: Vec<Vec<u8>> = par_map(bs.clone(), |band| {
let mut part = vec![0u8; band.len() * row];
for (i, p) in band.clone().enumerate() {
for (c, plane) in planes.iter().enumerate() {
let dst = &mut part[(i * ch + c) * bps..(i * ch + c + 1) * bps];
match plane {
Some(p) if p.len() >= (i + 1) * bps => {
let src = &p[i * bps..(i + 1) * bps];
Some(pl) if pl.len() >= (p + 1) * bps => {
let src = &pl[p * bps..(p + 1) * bps];
for k in 0..bps {
dst[k] = src[bps - 1 - k];
}
if cfg!(target_endian = "big") {
dst.copy_from_slice(src);
}
if invert[c] {
invert_sample(dst, s);
}
@@ -240,8 +252,16 @@ mod tests {
}
}
}
out
part
});
for (band, part) in bs.into_iter().zip(parts) {
out[band.start * row..band.end * row].copy_from_slice(&part);
}
out
}
#[test]
fn banded_shuffles_match_a_sequential_reference() {
fn seq_deinterleave(bytes: &[u8], ch: usize, s: SampleType, invert: &[bool]) -> Vec<Vec<u8>> {
let bps = s.bytes();
let n = bytes.len() / (ch * bps).max(1);
@@ -260,16 +280,17 @@ mod tests {
planes
}
// Odd pixel counts cross band boundaries; a missing plane exercises the fill path.
for n in [1usize, 7, 64, 1000] {
for n in [0usize, 1, 7, 64, 1000] {
for s in [SampleType::U8, SampleType::U16, SampleType::F32] {
let bps = s.bytes();
let planes: Vec<Vec<u8>> = (0..4).map(|c| (0..n * bps).map(|i| (i * 13 + c * 5) as u8).collect()).collect();
let mut refs: Vec<Option<&[u8]>> = planes.iter().map(|p| Some(&p[..])).collect();
refs[2] = None;
refs[3] = Some(&planes[3][..(n * bps).saturating_sub(1)]);
let fill: Vec<Vec<u8>> = (0..4).map(|c| vec![c as u8 * 40; bps]).collect();
let invert = [false, true, false, true];
let got = interleave(&refs, &fill, n, s, &invert);
let want = seq_interleave(&refs, &fill, n, s, &invert);
let want = previous_interleave(&refs, &fill, n, s, &invert);
assert_eq!(got, want, "interleave n={n}");
let planes_back = deinterleave(&got, 4, s, &invert);
let want_back = seq_deinterleave(&got, 4, s, &invert);
@@ -287,6 +308,7 @@ mod tests {
let fill = vec![zero_sample(s); 3];
let inv = [false, true, false];
let il = interleave(&refs, &fill, 4, s, &inv);
assert_eq!(interleave(&refs, &[], 4, s, &inv), il);
let back = deinterleave(&il, 3, s, &inv);
if s == SampleType::F32 {
// 1 - (1 - v) may not be exact for floats; only check non-inverted planes.
@@ -302,6 +324,12 @@ mod tests {
fn missing_planes_use_fill() {
let il = interleave(&[None, Some(&[5, 6])], &[vec![9], vec![0]], 2, SampleType::U8, &[false, false]);
assert_eq!(il, vec![9, 5, 9, 6]);
for s in [SampleType::U8, SampleType::U16, SampleType::F32] {
for n in [1, 1000] {
let fill = vec![vec![9; s.bytes()]; 2];
assert_eq!(interleave(&[None, Some(&[1, 2, 3][..s.bytes() - 1])], &fill, n, s, &[]), vec![9; n * 2 * s.bytes()]);
}
}
}
#[test]
@@ -343,3 +371,35 @@ mod tests {
assert_eq!(decode_be(&f, 0, SampleType::F32), 2.5);
}
}
#[cfg(test)]
mod measurement {
use super::*;
#[test]
#[ignore = "manual 24 MP release benchmark"]
fn benchmark_interleave() {
use std::{hint::black_box, time::Instant};
const N: usize = 24_000_000;
for sample in [SampleType::U8, SampleType::U16, SampleType::F32] {
let bps = sample.bytes();
let planes: Vec<Vec<u8>> = (0..4).map(|c| (0..N * bps).map(|i| (i * 13 + c * 5) as u8).collect()).collect();
let refs: Vec<_> = planes.iter().map(|p| Some(p.as_slice())).collect();
let fill = vec![vec![0; bps]; 4];
let invert = [false, true, false, true];
assert_eq!(tests::previous_interleave(&refs, &fill, N, sample, &invert), interleave(&refs, &fill, N, sample, &invert));
for trial in 0..6 {
let mut paths = [("main", tests::previous_interleave as fn(_, _, _, _, _) -> _), ("split", interleave)];
if trial % 2 == 1 {
paths.reverse();
}
for (name, path) in paths {
let t = Instant::now();
let out = black_box(path)(&refs, &fill, N, sample, &invert);
let elapsed = t.elapsed().as_secs_f64() * 1000.0;
black_box(&out);
println!("interleave,{sample:?},{trial},{name},{elapsed:.6}");
}
}
}
}
}
+38
View File
@@ -0,0 +1,38 @@
# PSD pixel buffer allocation
Only PSD pixel shuffling and sample-codec inlining change. Interleave writes into the final buffer rather than allocating an intermediate buffer per band. The kernel processes 24 million pixels with four full planes at 8-bit, 16-bit and float depth.
Measured on 2026-10-19 against `d95fdf3181c01ad41e95456091a394b105a9afb2` on an Apple M4 Max
(16 CPUs, 128 GiB), macOS arm64, rustc 1.99.1. Each branch starts directly at that
base; none of the other performance splits is applied. Release uses optimization
level 3, LTO off, 16 codegen units, debug info off and incremental compilation off.
Builds ran at nice 10 with APFS-cloned warm caches; measurements ran at normal
priority with no concurrent builds. Timings are medians of five samples after one
warm-up, except the compositor's ten Criterion samples. These synthetic workloads
are reproducible checks of this change, not a representative image corpus.
| Workload | Main ms | This change ms | Time change |
|---|---:|---:|---:|
| U8 | 20.493 | 5.097 | -75.1% |
| U16 | 31.827 | 6.312 | -80.2% |
| F32 | 41.915 | 11.305 | -73.0% |
These are interleave-kernel timings, not whole PSD import/export timings. Direct RLE output is already on main and is not part of this PR.
The committed equivalence test uses the literal previous main-branch interleave implementation. It compares raw bytes across U8/U16/F32, empty and partial bands, missing/truncated planes, fill pixels and inversion. Existing deinterleave reference and big-endian sample-codec round trips cover the companion paths.
Reproduce with the same release profile environment:
```sh
CARGO_PROFILE_RELEASE_LTO=false CARGO_PROFILE_RELEASE_CODEGEN_UNITS=16 CARGO_PROFILE_RELEASE_DEBUG=0 CARGO_INCREMENTAL=0 nice -n 10 cargo test --release -p photocraft-io --lib --no-run
# Run the printed executable at normal priority:
<printed-executable> benchmark_interleave --ignored --nocapture --test-threads=1
```
## Local validation
- `cargo test -p photocraft-io --locked`: passed (617 passed, 6 ignored).
- `cargo clippy -p photocraft-io --locked --all-targets -- -D warnings`: passed.
- `cargo xtask layers`: passed.
- `cargo check -p photocraft-io --locked --target wasm32-unknown-unknown`: passed.
- `cargo xtask test-corpus -p io`: passed (646 passed, 6 ignored).