Accelerate baseline JPEG Metal texture batches and fix scratch lifetime

crates/j2k-jpeg-metal/docs/decode-architecture.md

@@ -0,0 +1,141 @@
+# JPEG Metal decode architecture and measurements
+
+Measured on an Apple M4 Pro, macOS 26.5.2, Rust 1.96.0, using the workspace
+`release-bench` profile, on September 5–6, 2026. These changes concern baseline
+JPEG; JPEG 2000 implementations are unchanged.
+
+## Execution and ownership
+
+Supported 4:2:0 and 4:2:2 full texture batches now decode entropy and perform IDCT
+in one grid spanning the compatible tiles. The grid writes three reusable private
+component planes. A subsequent GPU pass performs chroma reconstruction, color
+conversion, and writes the caller's existing private RGBA textures. Coefficients
+are not materialized in a full-image intermediate. The direct 4:4:4 texture path
+is retained: it has no subsampled chroma repair and its plane decoder uses a
+different ABI.
+
+The prior direct subsampled texture shader ran a separate encoder per tile,
+retained substantial per-thread pixel state, and required additional horizontal,
+vertical, and corner repair passes. Avoiding component-plane storage did not
+compensate for those costs in the measured workloads. The component-plane route
+requires 1.5 MiB of private intermediate storage for sixteen 256×256 4:2:0 tiles,
+reused on subsequent calls.
+
+Packet entropy is copied directly into reusable shared Metal storage. Only the
+offset, length, and checkpoint metadata is materialized in host vectors. Retained
+packet owners and metadata remain budgeted; shared Metal allocations retain the
+existing device and repository size limits. Shared storage is appropriate for
+CPU-populated data, while GPU-only planes and textures stay private, following
+[Apple's storage-mode guidance](https://developer.apple.com/documentation/metal/choosing-a-resource-storage-mode-for-apple-gpus).
+
+Each runtime has two independently locked scratch slots. Concurrent callers can
+stage separate batches without overwriting each other's inputs or statuses;
+further callers block instead of allocating additional slots. Each lease lasts
+through command-buffer completion **and CPU status consumption**. An availability
+gate wakes callers when either slot is returned, including when an owner
+unwinds; poisoned state remains a typed error. The public
+synchronous and deferred-submit contracts remain unchanged; this does not add an
+asynchronous API. Reusing the same output still requires its existing access gate.
+Two concurrent callers can retain two sets of scratch capacity.
+
+Successful immutable pipeline registries are retained per Metal device registry
+ID for the process lifetime. Independently created sessions share those pipelines
+but own separate queues and mutable runtime state. Initialization is serialized
+per device, failures are not published, and a poisoned cache or failed metadata
+reservation falls back to an uncached load while preserving loader errors.
+
+## Preparation and compatibility
+
+Checkpoint traversal now uses the existing entropy-only `skip_block` operation
+instead of constructing dequantized coefficient blocks it discards. Differential
+tests preserve bit-reader snapshots, DC prediction, and exact errors.
+
+Full entropy validation remains eager. Deferring packet construction or deriving
+restart work only from marker offsets would move malformed-Huffman errors out of
+the current constructor and is not part of this change. Small/one-shot `Auto`
+routing and explicit Metal rejection behavior are unchanged.
+
+## Controlled implementation comparison
+
+The ignored `texture_scheduling_experiment` test compares the direct and
+component-plane implementations in the same optimized executable, with identical
+distinct JPEG tiles, reusable sessions/textures, CPU output parity checks for
+every tile, five warm iterations, and fifteen measured iterations per variant.
+The table shows median batch wall time using the existing 256-thread cap.
+
+| 4:2:0 workload, 16 distinct tiles | Direct textures | Component planes |
+|---|---:|---:|
+| 256×256, no restart | 6.723 ms | 0.517 ms |
+| 256×256, restart interval 2 | 12.127 ms | 0.755 ms |
+| 512×512, no restart | 6.993 ms | 1.200 ms |
+| 512×512, restart interval 4 | 26.683 ms | 1.644 ms |
+
+The public-API Criterion benchmark also confirms the improvement. The original
+run used ten samples, one second of warmup and two seconds of measurement. The
+final run used ten samples, two seconds of warmup and five seconds of measurement:
+
+| Warm 256×256 RGBA texture batch, 16 tiles | Original estimate | Final estimate (95% interval) |
+|---|---:|---:|
+| No restart markers | 6.635 ms | 0.677 ms (0.625–0.730 ms) |
+| Restart interval 2 | 12.029 ms | 0.914 ms (0.907–0.918 ms) |
+
+An intervening run was substantially noisier, including a slower CPU-only
+restart control. Repeating that control without source changes measured
+0.314 ms (0.312–0.317 ms), close to the original 0.311 ms; the apparent CPU
+regression was not reproduced. These are throughput comparisons for warm
+resident batches, not latency claims for all JPEG requests.
+
+Additional comparisons covered batch sizes 1, 2, and 4, 16×16 and 64×64 inputs,
+and 4:2:2 at 256×256 and 512×512. Component planes also won those comparisons,
+although the small-workload measurements were noisier. Separate correctness
+coverage includes distinct odd-sized edge tiles, restart streams, mixed tables,
+reused output, and concurrent calls sharing a backend session.
+
+Threadgroup widths of 32, 64, 128, and 256 were compared. Smaller groups modestly
+helped the direct shader, but there was no consistent additional benefit once
+using component planes. Production therefore retains the existing threadgroup
+policy. This follows Apple's advice to measure work distribution rather than
+assuming a particular group size is universally best:
+[Scale compute workloads across Apple GPUs](https://developer.apple.com/videos/play/wwdc2022/10159/).
+
+The earlier P19 experiment tested a different design: GPU entropy decoding into
+full-image coefficient scratch followed by a separate IDCT pass. Its rejection
+does not apply to this component-plane intermediate.
+
+The existing fast-packet planning benchmark measured about 301 µs without
+restarts and 312 µs with restart interval 2 after the entropy-only traversal
+change. Neither comparison showed a statistically supported improvement; no
+planning speedup is claimed. No isolated speedup is claimed for direct shared
+staging or the two-slot pool either.
+
+These measurements are local diagnostics, not a cross-device performance
+guarantee. Variant order was fixed, the desktop was not isolated, and no GPU
+occupancy or spill counters were captured. Warm measurements exclude packet and
+pipeline initialization. The original unmodified warm baseline varied roughly
+6.6–7.2 ms across separate runs.
+
+## Reproduction and checks
+
+```sh
+# Same-input implementation and threadgroup comparison (also asserts parity).
+J2K_REQUIRE_METAL_RUNTIME=1 cargo test --locked -p j2k-jpeg-metal --lib --profile release-bench texture_scheduling_experiment -- --ignored --test-threads=1 --nocapture
+
+# Small batches and 4:2:2 controls.
+J2K_REQUIRE_METAL_RUNTIME=1 J2K_JPEG_TEXTURE_SMALL_BATCHES=1 cargo test --locked -p j2k-jpeg-metal --lib --profile release-bench texture_scheduling_experiment -- --ignored --test-threads=1 --nocapture
+
+# Public API benchmark: fresh CPU/Metal RGB and warm resident batch-16 textures.
+J2K_REQUIRE_METAL_BENCH=1 cargo bench --locked -p j2k-jpeg-metal --bench compare --profile release-bench -- '^(decode_rgb/generated/fast420(_restart2)?_256x256/(cpu|metal)|wsi_tile_batch_rgba_textures/batch16/warm_session_reused_textures/generated/fast420(_restart2)?_256x256)$' --sample-size 10 --warm-up-time 1 --measurement-time 2 --noplot
+
+# Final resident-output run with longer sampling.
+J2K_REQUIRE_METAL_BENCH=1 cargo bench --locked -p j2k-jpeg-metal --bench compare --profile release-bench -- '^wsi_tile_batch_rgba_textures/batch16/warm_session_reused_textures/generated/fast420(_restart2)?_256x256$' --sample-size 10 --warm-up-time 2 --measurement-time 5 --noplot
+
+# Runtime tests require real Metal instead of accepting an unavailable device.
+J2K_REQUIRE_METAL_RUNTIME=1 cargo test --locked -p j2k-jpeg-metal -- --test-threads=1
+cargo test -p j2k-jpeg --lib
+cargo fmt --check -p j2k-jpeg-metal -p j2k-jpeg
+cargo clippy --locked -p j2k-jpeg-metal --lib -- -D warnings
+cargo clippy --locked -p j2k-jpeg-metal --tests --benches --bins -- -D warnings -A clippy::disallowed_methods -A clippy::disallowed_macros
+```
+
+The non-library Clippy allowances match the repository's existing Metal check
+policy. They do not apply to production library checks.

crates/j2k-jpeg-metal/docs/routing-benchmarks.md

@@ -1,5 +1,8 @@
 # JPEG Metal Routing and Benchmarks
 
+The current batch execution design, ownership rules, measured comparisons, and
+reproduction commands are in [Decode architecture](decode-architecture.md).
+
 JPEG Metal decode should stay selective. The current Metal paths are for fast
 baseline/checkpointed packets, coalesced batches, and resident outputs. They are
 not a claim of full JPEG entropy decode coverage.

crates/j2k-jpeg-metal/src/buffers.rs

@@ -154,9 +154,13 @@ pub(crate) fn new_decode_plane_buffer(
 
 #[cfg(test)]
 mod tests {
-    use j2k_metal_support::MetalSupportError;
+    use j2k_metal_support::{system_default_device, MetalSupportError};
 
-    use super::{buffer_access_error, buffer_readback_error};
+    use super::{
+        buffer_access_error, buffer_readback_error, checked_buffer_slice,
+        jpeg_shared_buffer_allocations_for_test, reset_jpeg_shared_buffer_allocations_for_test,
+        MetalBatchScratch,
+    };
     use crate::Error;
 
     #[test]
@@ -191,6 +195,101 @@ mod tests {
         assert!(error.to_string().contains(&usize::MAX.to_string()));
         assert!(std::error::Error::source(&error).is_some());
     }
+
+    #[test]
+    fn shared_scratch_stages_slices_directly_and_reuses_capacity() {
+        if !j2k_test_support::metal_runtime_gate(module_path!()) {
+            return;
+        }
+
+        let device = system_default_device().expect("Metal device");
+        let mut scratch = MetalBatchScratch::default();
+        reset_jpeg_shared_buffer_allocations_for_test();
+        let first = scratch
+            .shared_buffer_with_byte_slices(
+                &device,
+                "direct entropy staging test",
+                5,
+                [b"ab".as_slice(), b"cde".as_slice()],
+            )
+            .expect("first staging");
+        assert_eq!(
+            checked_buffer_slice::<u8>(&first, 5, "direct entropy staging test")
+                .expect("staged bytes"),
+            b"abcde"
+        );
+        let allocations = jpeg_shared_buffer_allocations_for_test();
+
+        let second = scratch
+            .shared_buffer_with_byte_slices(
+                &device,
+                "direct entropy staging test",
+                4,
+                [b"wxyz".as_slice()],
+            )
+            .expect("reused staging");
+
+        assert!(core::ptr::eq(
+            objc2::rc::Retained::as_ptr(&first),
+            objc2::rc::Retained::as_ptr(&second)
+        ));
+        assert_eq!(jpeg_shared_buffer_allocations_for_test(), allocations);
+        assert_eq!(
+            checked_buffer_slice::<u8>(&second, 4, "reused entropy staging test")
+                .expect("restaged bytes"),
+            b"wxyz"
+        );
+    }
+
+    #[test]
+    fn shared_scratch_rejects_direct_staging_length_mismatch() {
+        if !j2k_test_support::metal_runtime_gate(module_path!()) {
+            return;
+        }
+
+        let device = system_default_device().expect("Metal device");
+        let mut scratch = MetalBatchScratch::default();
+        let error = scratch
+            .shared_buffer_with_byte_slices(
+                &device,
+                "direct entropy staging mismatch test",
+                3,
+                [b"ab".as_slice()],
+            )
+            .expect_err("short staging must fail");
+
+        assert!(matches!(error, Error::MetalKernel { .. }));
+    }
+
+    #[test]
+    fn shared_scratch_direct_staging_preserves_typed_buffer_limit() {
+        if !j2k_test_support::metal_runtime_gate(module_path!()) {
+            return;
+        }
+
+        let device = system_default_device().expect("Metal device");
+        let mut scratch = MetalBatchScratch::default();
+        let requested = j2k_core::DEFAULT_MAX_HOST_ALLOCATION_BYTES + 1;
+        let error = scratch
+            .shared_buffer_with_byte_slices(
+                &device,
+                "direct entropy staging limit test",
+                requested,
+                std::iter::empty(),
+            )
+            .expect_err("oversized staging must fail");
+
+        assert!(matches!(
+            error,
+            Error::MetalSupport {
+                source: MetalSupportError::BufferAllocationTooLarge {
+                    requested: actual,
+                    ..
+                },
+                ..
+            } if actual == requested
+        ));
+    }
 }
 
 struct ReusablePrivateBuffer {
@@ -249,13 +348,13 @@ impl MetalBatchScratch {
         Ok(buffer)
     }
 
-    pub(crate) fn shared_buffer_with_bytes(
+    fn shared_buffer(
         &mut self,
         device: &DeviceRef,
         key: &'static str,
-        bytes: &[u8],
+        bytes: usize,
     ) -> Result<Buffer, Error> {
-        let capacity = bytes.len().max(1);
+        let capacity = bytes.max(1);
         let buffer = if let Some(entry) = self
             .shared_buffers
             .iter()
@@ -285,6 +384,17 @@ impl MetalBatchScratch {
             buffer
         };
 
+        Ok(buffer)
+    }
+
+    pub(crate) fn shared_buffer_with_bytes(
+        &mut self,
+        device: &DeviceRef,
+        key: &'static str,
+        bytes: &[u8],
+    ) -> Result<Buffer, Error> {
+        let buffer = self.shared_buffer(device, key, bytes.len())?;
+
         if !bytes.is_empty() {
             // SAFETY: This scratch buffer is exclusively leased during CPU
             // initialization and has not yet been submitted to Metal.
@@ -294,6 +404,43 @@ impl MetalBatchScratch {
         Ok(buffer)
     }
 
+    pub(crate) fn shared_buffer_with_byte_slices<'a>(
+        &mut self,
+        device: &DeviceRef,
+        key: &'static str,
+        total_bytes: usize,
+        slices: impl IntoIterator<Item = &'a [u8]>,
+    ) -> Result<Buffer, Error> {
+        let buffer = self.shared_buffer(device, key, total_bytes)?;
+        let mut offset = 0_usize;
+        for bytes in slices {
+            let end = offset
+                .checked_add(bytes.len())
+                .ok_or_else(|| Error::MetalKernel {
+                    message: "JPEG Metal shared scratch staging length overflowed".to_string(),
+                })?;
+            if end > total_bytes {
+                return Err(Error::MetalKernel {
+                    message: "JPEG Metal shared scratch staging length mismatch".to_string(),
+                });
+            }
+            if !bytes.is_empty() {
+                // SAFETY: This scratch buffer is exclusively leased during CPU
+                // initialization and has not yet been submitted to Metal. Each
+                // checked range is disjoint and lies below `total_bytes`.
+                unsafe { checked_buffer_write::<u8>(&buffer, offset, bytes) }
+                    .map_err(|error| buffer_access_error("shared scratch upload", error))?;
+            }
+            offset = end;
+        }
+        if offset != total_bytes {
+            return Err(Error::MetalKernel {
+                message: "JPEG Metal shared scratch staging length mismatch".to_string(),
+            });
+        }
+        Ok(buffer)
+    }
+
     pub(crate) fn shared_buffer_with_slice<T: GpuAbi>(
         &mut self,
         device: &DeviceRef,