perf(jpeg-metal): reuse subsampled single-image scratch

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

@@ -348,7 +348,7 @@ impl MetalBatchScratch {
         Ok(buffer)
     }
 
-    fn shared_buffer(
+    pub(crate) fn shared_buffer(
         &mut self,
         device: &DeviceRef,
         key: &'static str,
@@ -387,6 +387,17 @@ impl MetalBatchScratch {
         Ok(buffer)
     }
 
+    pub(crate) fn shared_zeroed_buffer(
+        &mut self,
+        device: &DeviceRef,
+        key: &'static str,
+        bytes: usize,
+    ) -> Result<Buffer, Error> {
+        let buffer = self.shared_buffer(device, key, bytes)?;
+        checked_fill_buffer_u8(&buffer, bytes, 0, "initialize shared scratch")?;
+        Ok(buffer)
+    }
+
     pub(crate) fn shared_buffer_with_bytes(
         &mut self,
         device: &DeviceRef,

crates/j2k-jpeg-metal/src/compute/pack_dispatch/subsampled.rs

@@ -3,6 +3,10 @@
 use crate::metal_types::prelude::*;
 
 use crate::buffers::new_shared_buffer;
+use crate::buffers::MetalBatchScratch;
+use crate::compute::single_decode::scratch::{
+    packet_buffers as scratch_packet_buffers, status_buffer as scratch_status_buffer,
+};
 
 use super::super::{
     batch, bind_fast_decode_entropy_inputs, bind_three_plane_pack, checked_entropy_segment_count,
@@ -37,6 +41,7 @@ pub(in crate::compute) fn encode_fast_subsampled_region_batch_item<P: FastSubsam
     packet: &P,
     fmt: PixelFormat,
     roi: Rect,
+    mut scratch: Option<&mut MetalBatchScratch>,
 ) -> Result<BatchedDecodeItem, Error> {
     let roi = core_rect_to_jpeg(roi);
     let source_window = fast_subsampled_full_mcu_window::<P>(packet.dimensions(), roi);
@@ -77,19 +82,49 @@ pub(in crate::compute) fn encode_fast_subsampled_region_batch_item<P: FastSubsam
     let y_len = source_window.w as usize * source_window.h as usize;
     let chroma_len =
         source_window.w.div_ceil(2) as usize * P::chroma_height(source_window.h) as usize;
-    let y_plane = new_decode_plane_buffer(&runtime.device, y_len, false)?;
-    let cb_plane = new_private_buffer(&runtime.device, chroma_len)?;
-    let cr_plane = new_private_buffer(&runtime.device, chroma_len)?;
+    let y_plane = if let Some(scratch) = scratch.as_deref_mut() {
+        scratch.private_buffer(&runtime.device, "single_decode_y", y_len)?
+    } else {
+        new_decode_plane_buffer(&runtime.device, y_len, false)?
+    };
+    let cb_plane = if let Some(scratch) = scratch.as_deref_mut() {
+        scratch.private_buffer(&runtime.device, "single_decode_cb", chroma_len)?
+    } else {
+        new_private_buffer(&runtime.device, chroma_len)?
+    };
+    let cr_plane = if let Some(scratch) = scratch.as_deref_mut() {
+        scratch.private_buffer(&runtime.device, "single_decode_cr", chroma_len)?
+    } else {
+        new_private_buffer(&runtime.device, chroma_len)?
+    };
     let decode_threads = entropy_decode_thread_count(
         packet.restart_interval_mcus(),
         restart_offsets.len(),
         packet.entropy_checkpoints().len(),
     );
-    let status_buffer = decode_status_buffer(&runtime.device, decode_threads)?;
-    let entropy_buffer = new_shared_buffer_with_data(&runtime.device, packet.entropy_bytes())?;
-    let restart_offsets_buffer = restart_offsets_buffer(&runtime.device, restart_offsets)?;
-    let entropy_checkpoints_buffer =
-        entropy_checkpoints_buffer(&runtime.device, packet.entropy_checkpoints())?;
+    let (status_buffer, entropy_buffer, restart_offsets_buffer, entropy_checkpoints_buffer) =
+        if let Some(scratch) = scratch {
+            let packet_buffers = scratch_packet_buffers(
+                scratch,
+                &runtime.device,
+                packet.entropy_bytes(),
+                restart_offsets,
+                packet.entropy_checkpoints(),
+            )?;
+            (
+                scratch_status_buffer(scratch, &runtime.device, decode_threads)?,
+                packet_buffers.entropy,
+                packet_buffers.restart_offsets,
+                packet_buffers.checkpoints,
+            )
+        } else {
+            (
+                decode_status_buffer(&runtime.device, decode_threads)?,
+                new_shared_buffer_with_data(&runtime.device, packet.entropy_bytes())?,
+                restart_offsets_buffer(&runtime.device, restart_offsets)?,
+                entropy_checkpoints_buffer(&runtime.device, packet.entropy_checkpoints())?,
+            )
+        };
 
     let (dc_tables, ac_tables) = fast_packet_huffman_tables(packet);
 
@@ -161,6 +196,7 @@ pub(in crate::compute) fn encode_fast_subsampled_scaled_batch_item<P: FastSubsam
     packet: &P,
     fmt: PixelFormat,
     scale: j2k_core::Downscale,
+    mut scratch: Option<&mut MetalBatchScratch>,
 ) -> Result<BatchedDecodeItem, Error> {
     let Some(params) = fast_subsampled_scaled_params(packet, scale) else {
         return Err(Error::MetalKernel {
@@ -170,19 +206,51 @@ pub(in crate::compute) fn encode_fast_subsampled_scaled_batch_item<P: FastSubsam
 
     let y_len = params.scaled_width as usize * params.scaled_height as usize;
     let chroma_len = params.chroma_width as usize * params.chroma_height as usize;
-    let y_plane = new_decode_plane_buffer(&runtime.device, y_len, fmt == PixelFormat::Gray8)?;
-    let cb_plane = new_private_buffer(&runtime.device, chroma_len)?;
-    let cr_plane = new_private_buffer(&runtime.device, chroma_len)?;
+    let y_plane = if fmt == PixelFormat::Gray8 {
+        new_decode_plane_buffer(&runtime.device, y_len, true)?
+    } else if let Some(scratch) = scratch.as_deref_mut() {
+        scratch.private_buffer(&runtime.device, "single_decode_y", y_len)?
+    } else {
+        new_private_buffer(&runtime.device, y_len)?
+    };
+    let cb_plane = if let Some(scratch) = scratch.as_deref_mut() {
+        scratch.private_buffer(&runtime.device, "single_decode_cb", chroma_len)?
+    } else {
+        new_private_buffer(&runtime.device, chroma_len)?
+    };
+    let cr_plane = if let Some(scratch) = scratch.as_deref_mut() {
+        scratch.private_buffer(&runtime.device, "single_decode_cr", chroma_len)?
+    } else {
+        new_private_buffer(&runtime.device, chroma_len)?
+    };
     let decode_threads = entropy_decode_thread_count(
         packet.restart_interval_mcus(),
         packet.restart_offsets().len(),
         packet.entropy_checkpoints().len(),
     );
-    let status_buffer = decode_status_buffer(&runtime.device, decode_threads)?;
-    let entropy_buffer = new_shared_buffer_with_data(&runtime.device, packet.entropy_bytes())?;
-    let restart_offsets_buffer = restart_offsets_buffer(&runtime.device, packet.restart_offsets())?;
-    let entropy_checkpoints_buffer =
-        entropy_checkpoints_buffer(&runtime.device, packet.entropy_checkpoints())?;
+    let (status_buffer, entropy_buffer, restart_offsets_buffer, entropy_checkpoints_buffer) =
+        if let Some(scratch) = scratch {
+            let packet_buffers = scratch_packet_buffers(
+                scratch,
+                &runtime.device,
+                packet.entropy_bytes(),
+                packet.restart_offsets(),
+                packet.entropy_checkpoints(),
+            )?;
+            (
+                scratch_status_buffer(scratch, &runtime.device, decode_threads)?,
+                packet_buffers.entropy,
+                packet_buffers.restart_offsets,
+                packet_buffers.checkpoints,
+            )
+        } else {
+            (
+                decode_status_buffer(&runtime.device, decode_threads)?,
+                new_shared_buffer_with_data(&runtime.device, packet.entropy_bytes())?,
+                restart_offsets_buffer(&runtime.device, packet.restart_offsets())?,
+                entropy_checkpoints_buffer(&runtime.device, packet.entropy_checkpoints())?,
+            )
+        };
 
     let (dc_tables, ac_tables) = fast_packet_huffman_tables(packet);
 
@@ -567,12 +635,22 @@ pub(in crate::compute) fn encode_fast_subsampled_op_batch_item<P: FastSubsampled
         batch::BatchOp::Full => {
             encode_fast_subsampled_batch_item(runtime, command_buffer, packet, fmt)
         }
-        batch::BatchOp::Region(roi) => {
-            encode_fast_subsampled_region_batch_item(runtime, command_buffer, packet, fmt, roi)
-        }
-        batch::BatchOp::Scaled(scale) => {
-            encode_fast_subsampled_scaled_batch_item(runtime, command_buffer, packet, fmt, scale)
-        }
+        batch::BatchOp::Region(roi) => encode_fast_subsampled_region_batch_item(
+            runtime,
+            command_buffer,
+            packet,
+            fmt,
+            roi,
+            None,
+        ),
+        batch::BatchOp::Scaled(scale) => encode_fast_subsampled_scaled_batch_item(
+            runtime,
+            command_buffer,
+            packet,
+            fmt,
+            scale,
+            None,
+        ),
         batch::BatchOp::RegionScaled { roi, scale } => {
             encode_fast_subsampled_scaled_region_batch_item(
                 FastSubsampledScaledRegionBatchItemRequest {

crates/j2k-jpeg-metal/src/compute/scratch_pool.rs

@@ -5,20 +5,22 @@
 use crate::{buffers::MetalBatchScratch, Error};
 use std::{
     ops::{Deref, DerefMut},
-    sync::{Condvar, Mutex, MutexGuard, TryLockError},
+    sync::{Condvar, Mutex, TryLockError},
 };
 
 pub(super) struct BatchScratchPool {
-    slots: [Mutex<MetalBatchScratch>; 2],
-    availability: Mutex<()>,
+    slots: [Mutex<Option<MetalBatchScratch>>; 2],
+    // A panic in an owned lease must remain fail-closed after moving the
+    // scratch value out of its mutex. Access this flag under the same gate.
+    availability: Mutex<bool>,
     released: Condvar,
 }
 
 impl Default for BatchScratchPool {
     fn default() -> Self {
         Self {
-            slots: std::array::from_fn(|_| Mutex::new(MetalBatchScratch::default())),
-            availability: Mutex::new(()),
+            slots: std::array::from_fn(|_| Mutex::new(Some(MetalBatchScratch::default()))),
+            availability: Mutex::new(false),
             released: Condvar::new(),
         }
     }
@@ -34,14 +36,22 @@ impl BatchScratchPool {
     pub(super) fn acquire(&self) -> Result<BatchScratchLease<'_>, Error> {
         let mut availability = self.availability.lock().map_err(|_| poisoned())?;
         loop {
+            if *availability {
+                return Err(poisoned());
+            }
             for slot in &self.slots {
                 match slot.try_lock() {
-                    Ok(scratch) => {
-                        drop(availability);
-                        return Ok(BatchScratchLease {
-                            scratch: Some(scratch),
-                            pool: self,
-                        });
+                    Ok(mut entry) => {
+                        if let Some(scratch) = entry.take() {
+                            drop(entry);
+                            drop(availability);
+                            return Ok(BatchScratchLease {
+                                scratch: Some(scratch),
+                                slot,
+                                pool: self,
+                                panicking_on_acquire: std::thread::panicking(),
+                            });
+                        }
                     }
                     Err(TryLockError::WouldBlock) => {}
                     Err(TryLockError::Poisoned(_)) => return Err(poisoned()),
@@ -55,45 +65,56 @@ impl BatchScratchPool {
 
     #[cfg(test)]
     pub(super) fn in_use(&self) -> bool {
-        self.slots
-            .iter()
-            .any(|slot| matches!(slot.try_lock(), Err(TryLockError::WouldBlock)))
+        self.slots.iter().any(|slot| match slot.try_lock() {
+            Ok(entry) => entry.is_none(),
+            Err(TryLockError::WouldBlock) => true,
+            Err(TryLockError::Poisoned(_)) => false,
+        })
     }
 }
 
 pub(in crate::compute) struct BatchScratchLease<'a> {
-    scratch: Option<MutexGuard<'a, MetalBatchScratch>>,
+    scratch: Option<MetalBatchScratch>,
+    slot: &'a Mutex<Option<MetalBatchScratch>>,
     pool: &'a BatchScratchPool,
+    panicking_on_acquire: bool,
 }
 
 impl Deref for BatchScratchLease<'_> {
     type Target = MetalBatchScratch;
     fn deref(&self) -> &Self::Target {
         self.scratch
-            .as_deref()
+            .as_ref()
             .expect("scratch is present until lease drop")
     }
 }
 
 impl DerefMut for BatchScratchLease<'_> {
     fn deref_mut(&mut self) -> &mut Self::Target {
         self.scratch
-            .as_deref_mut()
+            .as_mut()
             .expect("scratch is present until lease drop")
     }
 }
 
 impl Drop for BatchScratchLease<'_> {
     fn drop(&mut self) {
-        // Release under the same gate used by acquire before notifying. Recover
-        // the gate only to finish cleanup during unwinding; future acquisition
-        // still reports poison, and a panicking owner poisons its scratch slot.
-        let _availability = self
+        // Restore under the acquisition gate before notifying, without holding
+        // any pool lock during GPU work. Preserve MutexGuard's panic contract.
+        let mut availability = self
             .pool
             .availability
             .lock()
             .unwrap_or_else(std::sync::PoisonError::into_inner);
-        drop(self.scratch.take());
+        if std::thread::panicking() && !self.panicking_on_acquire {
+            *availability = true;
+        }
+        let mut entry = self
+            .slot
+            .lock()
+            .unwrap_or_else(std::sync::PoisonError::into_inner);
+        *entry = self.scratch.take();
+        drop(entry);
         // Wake every waiter so a poisoned owner cannot leave other callers
         // asleep indefinitely after the first waiter reports that failure.
         self.pool.released.notify_all();
@@ -104,6 +125,25 @@ impl Drop for BatchScratchLease<'_> {
 mod tests {
     use super::*;
 
+    #[test]
+    fn scratch_leases_do_not_hold_pool_mutexes() {
+        let pool = BatchScratchPool::default();
+        let first = pool.acquire().expect("first lease");
+        let second = pool.acquire().expect("second lease");
+        assert!(pool.in_use());
+        assert!(
+            pool.availability.try_lock().is_ok(),
+            "availability gate held by lease"
+        );
+        assert!(
+            pool.slots.iter().all(|slot| slot.try_lock().is_ok()),
+            "slot mutex held by lease"
+        );
+        drop(first);
+        drop(second);
+        assert!(!pool.in_use());
+    }
+
     #[test]
     fn scratch_slots_fail_closed_after_panicking_owner() {
         let pool = BatchScratchPool::default();

crates/j2k-jpeg-metal/src/compute/single_decode.rs

@@ -2,6 +2,7 @@
 
 mod fast444;
 mod routing;
+pub(in crate::compute) mod scratch;
 mod subsampled;
 
 pub(super) use fast444::try_decode_fast444_scaled_region_to_surface_with_mode_and_status;