perf(j2k): stream large decodes and bounded rows through pool bands

Cargo.toml

@@ -103,6 +103,7 @@ objc2-metal-performance-shaders-graph = { version = "=0.3.2", default-features =
 ] }
 block2 = { version = "=0.6.2", default-features = false, features = ["std"] }
 rayon  = "1"
+bytemuck = { version = "1.25", default-features = false }
 cc     = "1.2"
 libc   = "0.2"
 windows-sys = "0.61"

crates/j2k-alloc-probe/src/lib.rs

@@ -8,9 +8,13 @@
 //! delayed allocator call from reporting into a later measurement, and each
 //! measurement waits for already-admitted reporters before taking its
 //! snapshot.
+//!
+//! The allocator also keeps the bytes live at any moment and their high
+//! water, always on and independent of measurements, for probes that bound
+//! working memory rather than allocation calls (see [`reset_peak`]).
 
 use std::alloc::{GlobalAlloc, Layout, System};
-use std::sync::atomic::{AtomicU64, AtomicU8, Ordering};
+use std::sync::atomic::{AtomicU64, AtomicU8, AtomicUsize, Ordering};
 
 #[cfg(test)]
 use std::sync::atomic::AtomicBool;
@@ -33,6 +37,8 @@ static ALLOCATIONS: AtomicU64 = AtomicU64::new(0);
 static REALLOCATIONS: AtomicU64 = AtomicU64::new(0);
 static DEALLOCATIONS: AtomicU64 = AtomicU64::new(0);
 static REQUESTED_BYTES: AtomicU64 = AtomicU64::new(0);
+static LIVE_BYTES: AtomicUsize = AtomicUsize::new(0);
+static PEAK_BYTES: AtomicUsize = AtomicUsize::new(0);
 
 #[cfg(test)]
 static TEST_HOOK: AtomicU8 = AtomicU8::new(TEST_HOOK_IDLE);
@@ -60,6 +66,7 @@ unsafe impl GlobalAlloc for CountingAllocator {
         let pointer = unsafe { System.alloc(layout) };
         if !pointer.is_null() {
             reporter.record_allocation(layout.size());
+            grow_live(layout.size());
         }
         pointer
     }
@@ -70,13 +77,15 @@ unsafe impl GlobalAlloc for CountingAllocator {
         let pointer = unsafe { System.alloc_zeroed(layout) };
         if !pointer.is_null() {
             reporter.record_allocation(layout.size());
+            grow_live(layout.size());
         }
         pointer
     }
 
     unsafe fn dealloc(&self, pointer: *mut u8, layout: Layout) {
         let reporter = Reporter::begin();
         reporter.record_deallocation();
+        shrink_live(layout.size());
         // SAFETY: forwards the pointer and its original layout unchanged.
         unsafe { System.dealloc(pointer, layout) };
     }
@@ -88,6 +97,11 @@ unsafe impl GlobalAlloc for CountingAllocator {
         let replacement = unsafe { System.realloc(pointer, layout, new_size) };
         if !replacement.is_null() {
             reporter.record_reallocation(new_size);
+            if new_size >= layout.size() {
+                grow_live(new_size - layout.size());
+            } else {
+                shrink_live(layout.size() - new_size);
+            }
         }
         replacement
     }
@@ -191,6 +205,44 @@ pub fn assert_allocations<R>(label: &str, budget: Budget, operation: impl FnOnce
     result
 }
 
+/// Bytes held through the global allocator right now, by every thread.
+#[must_use]
+pub fn live_bytes() -> usize {
+    LIVE_BYTES.load(Ordering::Relaxed)
+}
+
+/// Start a new high-water window at the bytes live now and return them.
+///
+/// The window is process-wide like a measurement: allocations on every
+/// thread raise [`peak_bytes`], so unrelated threads must stay idle while it
+/// is open. Windows need no guard and never block the allocator.
+pub fn reset_peak() -> usize {
+    let live = LIVE_BYTES.load(Ordering::Relaxed);
+    PEAK_BYTES.store(live, Ordering::Relaxed);
+    // An allocation between the load and the store raised the live bytes
+    // past `live`; keep it in the new window.
+    PEAK_BYTES.fetch_max(LIVE_BYTES.load(Ordering::Relaxed), Ordering::Relaxed);
+    live
+}
+
+/// The most bytes live at once since the last [`reset_peak`] (or since the
+/// process started).
+#[must_use]
+pub fn peak_bytes() -> usize {
+    PEAK_BYTES.load(Ordering::Relaxed)
+}
+
+fn grow_live(bytes: usize) {
+    let live = LIVE_BYTES
+        .fetch_add(bytes, Ordering::Relaxed)
+        .wrapping_add(bytes);
+    PEAK_BYTES.fetch_max(live, Ordering::Relaxed);
+}
+
+fn shrink_live(bytes: usize) {
+    LIVE_BYTES.fetch_sub(bytes, Ordering::Relaxed);
+}
+
 struct MeasurementGuard {
     finished: bool,
 }

crates/j2k-alloc-probe/tests/main.rs

@@ -4,7 +4,7 @@ use std::hint::black_box;
 use std::panic::{catch_unwind, AssertUnwindSafe};
 use std::sync::{mpsc, OnceLock};
 
-use j2k_alloc_probe::{assert_allocations, measure, Budget};
+use j2k_alloc_probe::{assert_allocations, live_bytes, measure, peak_bytes, reset_peak, Budget};
 use j2k_native::{
     decode_ht_code_block_scalar_with_workspace, encode,
     encode_ht_code_block_scalar_with_passes_and_workspace,
@@ -32,6 +32,10 @@ fn main() {
             harness_counts_allocation_and_deallocation,
         ),
         ("harness_counts_reallocation", harness_counts_reallocation),
+        (
+            "harness_tracks_live_bytes_and_their_peak",
+            harness_tracks_live_bytes_and_their_peak,
+        ),
         (
             "preexisting_deallocation_cannot_credit_budget",
             preexisting_deallocation_cannot_credit_budget,
@@ -131,6 +135,44 @@ fn harness_counts_reallocation() {
     assert!(stats.requested_bytes >= 4096);
 }
 
+fn harness_tracks_live_bytes_and_their_peak() {
+    const BLOCK: usize = 64 * 1024;
+    let base = reset_peak();
+    assert_eq!(peak_bytes(), base, "a new window starts at the live bytes");
+
+    let first = black_box(Box::<[u8]>::new_uninit_slice(BLOCK));
+    let second = black_box(Box::<[u8]>::new_uninit_slice(BLOCK / 2));
+    assert_eq!(live_bytes(), base + BLOCK + BLOCK / 2);
+    drop(first);
+    assert_eq!(
+        live_bytes(),
+        base + BLOCK / 2,
+        "a free lowers the live bytes"
+    );
+    assert_eq!(
+        peak_bytes(),
+        base + BLOCK + BLOCK / 2,
+        "a free never lowers the peak"
+    );
+
+    let mut grown: Vec<u8> = Vec::new();
+    grown
+        .try_reserve_exact(BLOCK / 4)
+        .expect("small allocation must succeed");
+    grown
+        .try_reserve_exact(2 * BLOCK)
+        .expect("growth allocation must succeed");
+    assert_eq!(live_bytes(), base + BLOCK / 2 + grown.capacity());
+    assert_eq!(peak_bytes(), base + BLOCK / 2 + grown.capacity());
+    // Nothing was pushed: shrinking the empty vector frees its buffer.
+    grown.shrink_to_fit();
+    drop(second);
+    assert_eq!(live_bytes(), base, "every byte allocated here was freed");
+
+    assert_eq!(reset_peak(), base);
+    assert_eq!(peak_bytes(), base, "a reset forgets the earlier peak");
+}
+
 fn preexisting_deallocation_cannot_credit_budget() {
     let preexisting = black_box(Box::new([0_u8; 8192]));
     let violation = without_panic_output(|| {

crates/j2k-native/Cargo.toml

@@ -26,6 +26,7 @@ name = "j2k_native"
 path = "src/lib.rs"
 
 [dependencies]
+bytemuck = { workspace = true, features = ["extern_crate_alloc"] }
 fearless_simd = { workspace = true, optional = true }
 libm = { workspace = true }
 log = { workspace = true, optional = true }
@@ -36,6 +37,8 @@ j2k-types = { path = "../j2k-types", version = "=0.13.2" }
 
 [dev-dependencies]
 criterion = { workspace = true }
+# Counts the bytes the band executor allocates, against what it charges.
+j2k-alloc-probe = { path = "../j2k-alloc-probe" }
 # Lets SIMD kernel tests also run on the portable fallback level.
 fearless_simd = { workspace = true, features = ["force_support_fallback"] }
 

crates/j2k-native/src/color/mod.rs

@@ -23,6 +23,7 @@ pub use output_planes::{
     DecodedComponentsParts, DecodedNativeComponents, DecodedNativeComponentsParts,
     NativeComponentPlane, NativeComponentPlaneParts,
 };
+pub use packing::pack_rgb8;
 #[cfg(all(feature = "simd", target_arch = "x86_64"))]
 pub(crate) use packing::store_rgb8_x8;
 pub(crate) use packing::{

crates/j2k-native/src/color/output_planes.rs

@@ -208,6 +208,17 @@ impl<'a> ComponentPlane<'a> {
 pub trait ComponentRowSink {
     /// Receive the samples of one rectangle of the region.
     fn write_rows(&mut self, x: u32, y: u32, planes: &[ComponentPlane<'_>]);
+
+    /// The sink's destination for a `width` x `height` region of three 8-bit
+    /// components, when it is interleaved RGB8 rows of `3 * width` bytes back
+    /// to back that the sink would fill from such planes by rounding and
+    /// saturating each sample. A decode may then pack its rows straight into
+    /// this buffer instead of calling [`Self::write_rows`]. The default,
+    /// `None`, always uses `write_rows`.
+    fn rgb8_region(&mut self, width: u32, height: u32) -> Option<&mut [u8]> {
+        let _ = (width, height);
+        None
+    }
 }
 
 /// Borrowed decoded component planes for an image.

crates/j2k-native/src/color/packing.rs

@@ -301,7 +301,7 @@ pub(crate) fn interleave_and_convert(
 }
 
 /// Pack 8-bit RGB planes as interleaved bytes, rounded and saturated.
-fn pack_rgb8(planes: [&[f32]; 3], output: &mut [u8]) {
+pub fn pack_rgb8(planes: [&[f32]; 3], output: &mut [u8]) {
     #[cfg(all(feature = "simd", target_arch = "x86_64"))]
     {
         use fearless_simd::{dispatch, Level};

crates/j2k-native/src/direct_cpu.rs

@@ -95,12 +95,21 @@ pub struct J2kDirectCpuScratch {
     ht_task_workspaces: Vec<HtCodeBlockDecodeWorkspace>,
     /// Extra classic workspaces for parallel referenced classic execution,
     /// one per additional Rayon worker.
-    classic_task_workspaces: Vec<J2kCodeBlockDecodeWorkspace>,
+    classic_task_workspaces: Vec<ClassicTaskWorkspace>,
     /// Whether referenced execution may spread work over the Rayon pool.
     parallel: bool,
     staged_state: Option<StagedDirectState>,
 }
 
+/// One extra Rayon worker's state for parallel referenced classic execution:
+/// its code-block workspace and the buffer it joins a code block's payload
+/// fragments into when a multi-layer codestream splits them.
+#[derive(Debug, Default)]
+struct ClassicTaskWorkspace {
+    decode: J2kCodeBlockDecodeWorkspace,
+    payload: Vec<u8>,
+}
+
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
 enum StagedDirectRoute {
     Classic,
@@ -170,12 +179,21 @@ impl J2kDirectCpuScratch {
             .iter()
             .try_fold(
                 self.classic_workspace.allocated_bytes().ok(),
-                |bytes, workspace| Some(bytes?.checked_add(workspace.allocated_bytes().ok()?)),
+                |bytes, task| Some(bytes?.checked_add(task.decode.allocated_bytes().ok()?)),
             )
             .flatten()
             .unwrap_or(usize::MAX)
     }
 
+    /// Length and capacity of each extra classic task's payload buffer.
+    #[cfg(test)]
+    fn classic_task_payloads(&self) -> Vec<(usize, usize)> {
+        self.classic_task_workspaces
+            .iter()
+            .map(|task| (task.payload.len(), task.payload.capacity()))
+            .collect()
+    }
+
     /// Number of coefficient/IDWT buffer owners retained by this scratch.
     ///
     /// Referenced multi-tile execution retains the maximum live tile shape,

crates/j2k-native/src/direct_cpu/allocation.rs

@@ -452,8 +452,9 @@ fn parallel_ht_task_workspace_count(
 }
 
 /// Largest code block and most code-block rows of any sub-band a parallel
-/// referenced classic execution decodes. One dispatch decodes at most one
-/// row per Rayon worker, so it never needs more workspaces than rows.
+/// referenced classic execution decodes, and its largest payload split over
+/// several input fragments. One dispatch decodes at most one row per Rayon
+/// worker, so it never needs more workspaces than rows.
 #[derive(Debug, Clone, Copy)]
 struct ClassicTaskShape {
     dimensions: (u32, u32),
@@ -462,6 +463,9 @@ struct ClassicTaskShape {
         expect(dead_code, reason = "rows bound parallel dispatch only")
     )]
     rows: usize,
+    /// Bytes each task joins one fragmented payload into; zero when every
+    /// payload is one input slice.
+    fragmented_payload_bytes: usize,
 }
 
 /// Extra classic workspaces beside the main one: one per additional Rayon
@@ -478,36 +482,44 @@ fn classic_task_workspace_count(parallel: bool, classic_tasks: Option<ClassicTas
 }
 
 /// Charge the extra classic workspaces of parallel referenced classic
-/// execution. `classic_tasks` is `None` for every other route.
+/// execution and their payload buffers. `classic_tasks` is `None` for every
+/// other route.
 fn include_classic_task_workspaces(
     budget: &mut DirectAllocationBudget,
     classic_tasks: Option<ClassicTaskShape>,
     scratch: &J2kDirectCpuScratch,
 ) -> Result<()> {
     let target_count = classic_task_workspace_count(scratch.parallel, classic_tasks);
-    budget.include_capacity::<crate::J2kCodeBlockDecodeWorkspace>(
+    budget.include_capacity::<super::ClassicTaskWorkspace>(
         scratch.classic_task_workspaces.capacity().max(target_count),
     )?;
-    let target_bytes = classic_tasks.map_or(Ok(0), |shape| {
+    let target_workspace_bytes = classic_tasks.map_or(Ok(0), |shape| {
         bitplane::classic_decode_workspace_bytes(shape.dimensions.0, shape.dimensions.1)
     })?;
+    let target_payload_bytes = classic_tasks.map_or(0, |shape| shape.fragmented_payload_bytes);
     for index in 0..scratch.classic_task_workspaces.len().max(target_count) {
-        let retained = scratch
+        let (retained_workspace, retained_payload) = scratch
             .classic_task_workspaces
             .get(index)
-            .map_or(Ok(0), crate::J2kCodeBlockDecodeWorkspace::allocated_bytes)?;
-        let target = if index < target_count {
-            target_bytes
+            .map_or(Ok((0, 0)), |task| {
+                task.decode
+                    .allocated_bytes()
+                    .map(|bytes| (bytes, task.payload.capacity()))
+            })?;
+        let (target_workspace, target_payload) = if index < target_count {
+            (target_workspace_bytes, target_payload_bytes)
         } else {
-            0
+            (0, 0)
         };
-        budget.include_bytes(retained.max(target))?;
+        budget.include_bytes(retained_workspace.max(target_workspace))?;
+        budget.include_bytes(retained_payload.max(target_payload))?;
     }
     Ok(())
 }
 
 /// Keep exactly the extra classic workspaces parallel referenced classic
-/// execution uses, each prepared for the plan's largest code block.
+/// execution uses, each prepared for the plan's largest code block and
+/// holding a buffer for its largest fragmented payload.
 fn reserve_classic_task_workspaces(
     classic_tasks: Option<ClassicTaskShape>,
     scratch: &mut J2kDirectCpuScratch,
@@ -518,15 +530,23 @@ fn reserve_classic_task_workspaces(
     while scratch.classic_task_workspaces.len() < count {
         scratch
             .classic_task_workspaces
-            .push(crate::J2kCodeBlockDecodeWorkspace::default());
+            .push(super::ClassicTaskWorkspace::default());
     }
     if let Some(ClassicTaskShape {
         dimensions: (width, height),
+        fragmented_payload_bytes,
         ..
     }) = classic_tasks
     {
-        for workspace in &mut scratch.classic_task_workspaces {
-            workspace.prepare(width, height)?;
+        for task in &mut scratch.classic_task_workspaces {
+            task.decode.prepare(width, height)?;
+            task.payload.clear();
+            if task.payload.capacity() < fragmented_payload_bytes {
+                // Release before growing so the old and new buffers never
+                // coexist beyond the charged maximum.
+                task.payload = Vec::new();
+            }
+            try_reserve_decode_elements(&mut task.payload, fragmented_payload_bytes)?;
         }
     }
     Ok(())