From babeb1172c676159c831267d13af67e2dea761e0 Mon Sep 17 00:00:00 2001 From: GF Date: Tue, 8 Sep 2026 05:45:08 -0400 Subject: [PATCH 1/4] perf(jpeg-metal): reuse synchronous single-image scratch --- crates/j2k-jpeg-metal/src/buffers.rs | 13 +- .../src/compute/pack_dispatch/subsampled.rs | 122 ++++++++++--- .../src/compute/scratch_pool.rs | 84 ++++++--- .../src/compute/single_decode.rs | 1 + .../src/compute/single_decode/fast444.rs | 121 ++++++++----- .../src/compute/single_decode/scratch.rs | 119 +++++++++++++ .../src/compute/single_decode/subsampled.rs | 102 +++++++++-- crates/j2k-jpeg-metal/src/compute/status.rs | 19 +- crates/j2k-jpeg-metal/src/compute/tests.rs | 47 +++++ .../src/tests/single_scratch.rs | 166 ++++++++++++++++++ 10 files changed, 679 insertions(+), 115 deletions(-) create mode 100644 crates/j2k-jpeg-metal/src/compute/single_decode/scratch.rs diff --git a/crates/j2k-jpeg-metal/src/buffers.rs b/crates/j2k-jpeg-metal/src/buffers.rs index 63102a603..a0c926f8b 100644 --- a/crates/j2k-jpeg-metal/src/buffers.rs +++ b/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 { + 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, diff --git a/crates/j2k-jpeg-metal/src/compute/pack_dispatch/subsampled.rs b/crates/j2k-jpeg-metal/src/compute/pack_dispatch/subsampled.rs index 150a2a57a..7e981f461 100644 --- a/crates/j2k-jpeg-metal/src/compute/pack_dispatch/subsampled.rs +++ b/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, ) -> Result { let roi = core_rect_to_jpeg(roi); let source_window = fast_subsampled_full_mcu_window::

(packet.dimensions(), roi); @@ -77,19 +82,49 @@ pub(in crate::compute) fn encode_fast_subsampled_region_batch_item, ) -> Result { 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 { 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 { diff --git a/crates/j2k-jpeg-metal/src/compute/scratch_pool.rs b/crates/j2k-jpeg-metal/src/compute/scratch_pool.rs index 0fbf96016..4cf6f7723 100644 --- a/crates/j2k-jpeg-metal/src/compute/scratch_pool.rs +++ b/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; 2], - availability: Mutex<()>, + slots: [Mutex>; 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, 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, 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,22 +65,26 @@ 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>, + scratch: Option, + slot: &'a Mutex>, 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") } } @@ -78,22 +92,29 @@ impl Deref for BatchScratchLease<'_> { 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(); diff --git a/crates/j2k-jpeg-metal/src/compute/single_decode.rs b/crates/j2k-jpeg-metal/src/compute/single_decode.rs index 2159eb161..e8024683a 100644 --- a/crates/j2k-jpeg-metal/src/compute/single_decode.rs +++ b/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; diff --git a/crates/j2k-jpeg-metal/src/compute/single_decode/fast444.rs b/crates/j2k-jpeg-metal/src/compute/single_decode/fast444.rs index e14ce3c31..47d4ae3eb 100644 --- a/crates/j2k-jpeg-metal/src/compute/single_decode/fast444.rs +++ b/crates/j2k-jpeg-metal/src/compute/single_decode/fast444.rs @@ -1,5 +1,6 @@ // SPDX-License-Identifier: MIT OR Apache-2.0 +use super::scratch::decode_buffers as scratch_decode_buffers; use crate::metal_types::prelude::*; use super::super::{ @@ -40,23 +41,29 @@ pub(in crate::compute) fn try_decode_fast444_to_surface( let params = fast444_params(packet)?; let mode = fast444_plane_mode(decoder); let plane_len = params.width as usize * params.height as usize; - let y_plane = new_decode_plane_buffer( - &runtime.device, - plane_len, - fmt == PixelFormat::Gray8 && mode != PlaneMode::Rgb, - )?; - let chroma_blue_plane = new_private_buffer(&runtime.device, plane_len)?; - let chroma_red_plane = new_private_buffer(&runtime.device, plane_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 mut scratch = runtime.batch_scratch()?; + let buffers = scratch_decode_buffers( + &mut scratch, + &runtime.device, + plane_len, + fmt == PixelFormat::Gray8 && mode != PlaneMode::Rgb, + decode_threads, + &packet.entropy_bytes, + &packet.restart_offsets, + &packet.entropy_checkpoints, + )?; + let y_plane = buffers.plane0; + let chroma_blue_plane = buffers.plane1; + let chroma_red_plane = buffers.plane2; + let status_buffer = buffers.status; + let entropy_buffer = buffers.packet.entropy; + let restart_offsets_buffer = buffers.packet.restart_offsets; + let entropy_checkpoints_buffer = buffers.packet.checkpoints; let (dc_tables, ac_tables) = fast_packet_huffman_tables(packet); @@ -205,23 +212,29 @@ pub(in crate::compute) fn try_decode_fast444_region_to_surface( )?; let mode = fast444_plane_mode(decoder); let plane_len = params.width as usize * params.height as usize; - let y_plane = new_decode_plane_buffer( - &runtime.device, - plane_len, - fmt == PixelFormat::Gray8 && mode != PlaneMode::Rgb, - )?; - let chroma_blue_plane = new_private_buffer(&runtime.device, plane_len)?; - let chroma_red_plane = new_private_buffer(&runtime.device, plane_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 mut scratch = runtime.batch_scratch()?; + let buffers = scratch_decode_buffers( + &mut scratch, + &runtime.device, + plane_len, + fmt == PixelFormat::Gray8 && mode != PlaneMode::Rgb, + decode_threads, + &packet.entropy_bytes, + restart_offsets, + &packet.entropy_checkpoints, + )?; + let y_plane = buffers.plane0; + let chroma_blue_plane = buffers.plane1; + let chroma_red_plane = buffers.plane2; + let status_buffer = buffers.status; + let entropy_buffer = buffers.packet.entropy; + let restart_offsets_buffer = buffers.packet.restart_offsets; + let entropy_checkpoints_buffer = buffers.packet.checkpoints; let (dc_tables, ac_tables) = fast_packet_huffman_tables(packet); @@ -286,23 +299,29 @@ pub(in crate::compute) fn try_decode_fast444_scaled_to_surface( let mode = fast444_plane_mode(decoder); let plane_len = params.scaled_width as usize * params.scaled_height as usize; - let y_plane = new_decode_plane_buffer( - &runtime.device, - plane_len, - fmt == PixelFormat::Gray8 && mode != PlaneMode::Rgb, - )?; - let chroma_blue_plane = new_private_buffer(&runtime.device, plane_len)?; - let chroma_red_plane = new_private_buffer(&runtime.device, plane_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 mut scratch = runtime.batch_scratch()?; + let buffers = scratch_decode_buffers( + &mut scratch, + &runtime.device, + plane_len, + fmt == PixelFormat::Gray8 && mode != PlaneMode::Rgb, + decode_threads, + &packet.entropy_bytes, + &packet.restart_offsets, + &packet.entropy_checkpoints, + )?; + let y_plane = buffers.plane0; + let chroma_blue_plane = buffers.plane1; + let chroma_red_plane = buffers.plane2; + let status_buffer = buffers.status; + let entropy_buffer = buffers.packet.entropy; + let restart_offsets_buffer = buffers.packet.restart_offsets; + let entropy_checkpoints_buffer = buffers.packet.checkpoints; let (dc_tables, ac_tables) = fast_packet_huffman_tables(packet); @@ -411,23 +430,29 @@ pub(in crate::compute) fn try_decode_fast444_scaled_region_to_surface_with_mode_ )?; let plane_len = params.scaled_width as usize * params.scaled_height as usize; - let y_plane = new_decode_plane_buffer( - &runtime.device, - plane_len, - fmt == PixelFormat::Gray8 && mode != PlaneMode::Rgb, - )?; - let chroma_blue_plane = new_private_buffer(&runtime.device, plane_len)?; - let chroma_red_plane = new_private_buffer(&runtime.device, plane_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 mut scratch = runtime.batch_scratch()?; + let buffers = scratch_decode_buffers( + &mut scratch, + &runtime.device, + plane_len, + fmt == PixelFormat::Gray8 && mode != PlaneMode::Rgb, + decode_threads, + &packet.entropy_bytes, + restart_offsets, + &packet.entropy_checkpoints, + )?; + let y_plane = buffers.plane0; + let chroma_blue_plane = buffers.plane1; + let chroma_red_plane = buffers.plane2; + let status_buffer = buffers.status; + let entropy_buffer = buffers.packet.entropy; + let restart_offsets_buffer = buffers.packet.restart_offsets; + let entropy_checkpoints_buffer = buffers.packet.checkpoints; let (dc_tables, ac_tables) = fast_packet_huffman_tables(packet); diff --git a/crates/j2k-jpeg-metal/src/compute/single_decode/scratch.rs b/crates/j2k-jpeg-metal/src/compute/single_decode/scratch.rs new file mode 100644 index 000000000..c7c0cd2a7 --- /dev/null +++ b/crates/j2k-jpeg-metal/src/compute/single_decode/scratch.rs @@ -0,0 +1,119 @@ +// SPDX-License-Identifier: MIT OR Apache-2.0 + +use crate::{ + buffers::{checked_copy_bytes_to_buffer_at, new_decode_plane_buffer, MetalBatchScratch}, + metal_types::{Buffer, DeviceRef}, + Error, +}; +use j2k_core::accelerator::GpuAbi; +use j2k_jpeg::adapter::JpegEntropyCheckpointV1; + +use super::super::{JpegDecodeStatus, JpegEntropyCheckpointHost}; + +pub(in crate::compute) struct ScratchPacketBuffers { + pub(in crate::compute) entropy: Buffer, + pub(in crate::compute) restart_offsets: Buffer, + pub(in crate::compute) checkpoints: Buffer, +} + +pub(in crate::compute) struct ScratchDecodeBuffers { + pub(in crate::compute) plane0: Buffer, + pub(in crate::compute) plane1: Buffer, + pub(in crate::compute) plane2: Buffer, + pub(in crate::compute) status: Buffer, + pub(in crate::compute) packet: ScratchPacketBuffers, +} + +pub(in crate::compute) fn status_buffer( + scratch: &mut MetalBatchScratch, + device: &DeviceRef, + count: u32, +) -> Result { + let bytes = crate::batch_allocation::checked_count_product( + count as usize, + core::mem::size_of::(), + "JPEG Metal single decode status bytes", + )?; + scratch.shared_zeroed_buffer(device, "single_decode_status", bytes) +} + +pub(in crate::compute) fn packet_buffers( + scratch: &mut MetalBatchScratch, + device: &DeviceRef, + entropy: &[u8], + restart_offsets: &[u32], + checkpoints: &[JpegEntropyCheckpointV1], +) -> Result { + if restart_offsets.is_empty() { + return Err(Error::MetalKernel { + message: "JPEG Metal restart offsets must contain at least one entry".to_string(), + }); + } + if checkpoints.is_empty() { + return Err(Error::MetalKernel { + message: "JPEG Metal entropy checkpoints must contain at least one entry".to_string(), + }); + } + let checkpoint_bytes = core::mem::size_of::(); + let total_checkpoint_bytes = crate::batch_allocation::checked_count_product( + checkpoints.len(), + checkpoint_bytes, + "JPEG Metal entropy checkpoint upload bytes", + )?; + let checkpoint_buffer = + scratch.shared_buffer(device, "single_decode_checkpoints", total_checkpoint_bytes)?; + for (index, checkpoint) in checkpoints.iter().copied().enumerate() { + let checkpoint = JpegEntropyCheckpointHost::from(checkpoint); + checked_copy_bytes_to_buffer_at( + &checkpoint_buffer, + index * checkpoint_bytes, + JpegEntropyCheckpointHost::as_bytes(&checkpoint), + "upload JPEG Metal entropy checkpoint", + )?; + } + Ok(ScratchPacketBuffers { + entropy: scratch.shared_buffer_with_bytes(device, "single_decode_entropy", entropy)?, + restart_offsets: scratch.shared_buffer_with_slice( + device, + "single_decode_restart_offsets", + restart_offsets, + )?, + checkpoints: checkpoint_buffer, + }) +} + +#[expect( + clippy::too_many_arguments, + reason = "the scratch preparation mirrors the three plane and packet inputs bound together" +)] +pub(in crate::compute) fn decode_buffers( + scratch: &mut MetalBatchScratch, + device: &DeviceRef, + plane_len: usize, + plane0_returned: bool, + decode_threads: u32, + entropy: &[u8], + restart_offsets: &[u32], + checkpoints: &[JpegEntropyCheckpointV1], +) -> Result { + // The fast decode kernels cover the active plane rectangles through the + // bounded deposit helpers in shaders_decode_helpers.metal. Callers reject + // any failed status before publishing output, so successful reuse does not + // require preserving or clearing earlier plane contents. + let plane0 = if plane0_returned { + new_decode_plane_buffer(device, plane_len, true)? + } else { + scratch.private_buffer(device, "single_decode_y", plane_len)? + }; + let plane1 = scratch.private_buffer(device, "single_decode_cb", plane_len)?; + let plane2 = scratch.private_buffer(device, "single_decode_cr", plane_len)?; + let packet = packet_buffers(scratch, device, entropy, restart_offsets, checkpoints)?; + let status = status_buffer(scratch, device, decode_threads)?; + Ok(ScratchDecodeBuffers { + plane0, + plane1, + plane2, + status, + packet, + }) +} diff --git a/crates/j2k-jpeg-metal/src/compute/single_decode/subsampled.rs b/crates/j2k-jpeg-metal/src/compute/single_decode/subsampled.rs index 8838cee26..6dc0a10ab 100644 --- a/crates/j2k-jpeg-metal/src/compute/single_decode/subsampled.rs +++ b/crates/j2k-jpeg-metal/src/compute/single_decode/subsampled.rs @@ -2,6 +2,9 @@ use crate::metal_types::prelude::*; +use super::scratch::{ + packet_buffers as scratch_packet_buffers, status_buffer as scratch_status_buffer, +}; use crate::buffers::new_shared_buffer; use super::super::{ @@ -36,7 +39,14 @@ pub(in crate::compute) fn decode_fast422_to_rgb_buffer( fmt: PixelFormat, output_storage: MTLResourceOptions, ) -> Result, Error> { - decode_fast_subsampled_to_rgb_buffer(runtime, packet, fmt, output_storage, fast422_status_error) + decode_fast_subsampled_to_rgb_buffer( + runtime, + packet, + fmt, + output_storage, + fast422_status_error, + false, + ) } #[cfg(target_os = "macos")] @@ -52,6 +62,7 @@ fn try_decode_fast_subsampled_to_surface( fmt, MTLResourceOptions::StorageModeShared, map_status, + true, )? else { return Ok(None); @@ -74,6 +85,7 @@ fn decode_fast_subsampled_to_rgb_buffer( fmt: PixelFormat, output_storage: MTLResourceOptions, map_status: impl Fn(JpegDecodeStatus) -> Error, + reuse_temporaries: bool, ) -> Result, Error> { let Some(packet) = packet else { return Ok(None); @@ -85,19 +97,54 @@ fn decode_fast_subsampled_to_rgb_buffer( let params = fast_subsampled_params(packet, fmt)?; let y_len = params.width as usize * params.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 mut scratch = reuse_temporaries + .then(|| runtime.batch_scratch()) + .transpose()?; + 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.as_deref_mut() { + 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); @@ -185,6 +232,7 @@ fn try_decode_fast_subsampled_region_to_surface( }; let command_buffer = new_command_buffer(&runtime.queue)?; + let mut scratch = runtime.batch_scratch()?; let item = encode_fast_subsampled_region_batch_item( runtime, &command_buffer, @@ -196,6 +244,7 @@ fn try_decode_fast_subsampled_region_to_surface( w: roi.w, h: roi.h, }, + Some(&mut scratch), )?; commit_and_wait_jpeg(&command_buffer)?; @@ -225,8 +274,15 @@ fn try_decode_fast_subsampled_scaled_to_surface( } let command_buffer = new_command_buffer(&runtime.queue)?; - let item = - encode_fast_subsampled_scaled_batch_item(runtime, &command_buffer, packet, fmt, scale)?; + let mut scratch = runtime.batch_scratch()?; + let item = encode_fast_subsampled_scaled_batch_item( + runtime, + &command_buffer, + packet, + fmt, + scale, + Some(&mut scratch), + )?; commit_and_wait_jpeg(&command_buffer)?; if let Some(status) = first_decode_error_status(&item.status_buffer, item.decode_threads)? { @@ -347,19 +403,26 @@ fn try_decode_fast_subsampled_scaled_region_to_surface( 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 mut scratch = runtime.batch_scratch()?; + let y_plane = scratch.private_buffer(&runtime.device, "single_decode_y", y_len)?; + let cb_plane = scratch.private_buffer(&runtime.device, "single_decode_cb", chroma_len)?; + let cr_plane = scratch.private_buffer(&runtime.device, "single_decode_cr", 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 packet_buffers = scratch_packet_buffers( + &mut scratch, + &runtime.device, + packet.entropy_bytes(), + restart_offsets, + packet.entropy_checkpoints(), + )?; + let status_buffer = scratch_status_buffer(&mut scratch, &runtime.device, decode_threads)?; + let entropy_buffer = packet_buffers.entropy; + let restart_offsets_buffer = packet_buffers.restart_offsets; + let entropy_checkpoints_buffer = packet_buffers.checkpoints; let (dc_tables, ac_tables) = fast_packet_huffman_tables(packet); @@ -440,6 +503,7 @@ pub(in crate::compute) fn decode_fast420_to_rgb_buffer( fmt, output_storage, fast_decode_status_error, + false, ) } diff --git a/crates/j2k-jpeg-metal/src/compute/status.rs b/crates/j2k-jpeg-metal/src/compute/status.rs index 42309cd2b..4513ddd9d 100644 --- a/crates/j2k-jpeg-metal/src/compute/status.rs +++ b/crates/j2k-jpeg-metal/src/compute/status.rs @@ -98,7 +98,7 @@ mod tests { } #[test] - fn batch_status_read_retains_scratch_ownership() { + fn batch_and_single_status_reads_retain_scratch_ownership() { if !j2k_test_support::metal_runtime_gate(module_path!()) { return; } @@ -117,7 +117,7 @@ mod tests { observed_reads.set(observed_reads.get() + 1); })); }); - let decoder = + let mut decoder = crate::Decoder::new(include_bytes!("../../fixtures/jpeg/baseline_420_16x16.jpg")) .expect("decoder"); let requests = [decoder.rgb8_metal_request(crate::batch::BatchOp::Full)]; @@ -131,6 +131,19 @@ mod tests { super::super::batch_entry::decode_full_batch_to_surfaces_with_session(&requests, &session) .expect("RGB batch") .expect("supported batch"); - assert_eq!(reads.get(), 2, "both completion paths must inspect status"); + for fmt in [ + crate::PixelFormat::Gray8, + crate::PixelFormat::Rgb8, + crate::PixelFormat::Rgba8, + ] { + decoder + .decode_to_device_with_session(fmt, &session) + .expect("single surface"); + } + assert_eq!( + reads.get(), + 5, + "batch and single completion paths must inspect status" + ); } } diff --git a/crates/j2k-jpeg-metal/src/compute/tests.rs b/crates/j2k-jpeg-metal/src/compute/tests.rs index 5da3ca173..14f78b1ef 100644 --- a/crates/j2k-jpeg-metal/src/compute/tests.rs +++ b/crates/j2k-jpeg-metal/src/compute/tests.rs @@ -1930,3 +1930,50 @@ impl EntropyBitWriter { self.bit_count = 0; } } + +#[test] +fn single_scratch_recovers_after_gpu_entropy_failure() { + if !should_run_metal_runtime() { + return; + } + let runtime = MetalRuntime::new().expect("runtime"); + let valid = j2k_jpeg::adapter::build_fast422_packet(BASELINE_422).expect("valid packet"); + let mut invalid = + j2k_jpeg::adapter::build_fast422_packet(BASELINE_422).expect("packet to corrupt"); + // All-one codes are invalid for these canonical tables. Reset the saved + // reader reservoir so the device reads the corrupt bytes from the start. + invalid.entropy_bytes.fill(0xff); + invalid.entropy_checkpoints.truncate(1); + invalid.entropy_checkpoints[0].mcu_index = 0; + invalid.entropy_checkpoints[0].entropy_pos = 0; + invalid.entropy_checkpoints[0].bit_acc = 0; + invalid.entropy_checkpoints[0].bit_count = 0; + let (expected, _) = CpuDecoder::new(BASELINE_422) + .unwrap() + .decode_request(DecodeRequest::full(PixelFormat::Rgb8)) + .unwrap(); + for _ in 0..2 { + let error = single_decode::try_decode_fast422_to_surface( + &runtime, + Some(&invalid), + PixelFormat::Rgb8, + ) + .err() + .expect("GPU entropy failure"); + assert!( + matches!(&error, Error::MetalKernel { message } + if message.starts_with("unexpected Metal fast422 failure at entropy byte ")), + "{error}" + ); + assert!( + !runtime.batch_scratch_in_use_for_test(), + "failed call must release scratch after completion" + ); + let surface = + single_decode::try_decode_fast422_to_surface(&runtime, Some(&valid), PixelFormat::Rgb8) + .expect("valid recovery") + .expect("supported packet"); + assert_eq!(surface.as_bytes().unwrap().as_ref(), expected); + assert!(!runtime.batch_scratch_in_use_for_test()); + } +} diff --git a/crates/j2k-jpeg-metal/src/tests/single_scratch.rs b/crates/j2k-jpeg-metal/src/tests/single_scratch.rs index db2197747..881bda92c 100644 --- a/crates/j2k-jpeg-metal/src/tests/single_scratch.rs +++ b/crates/j2k-jpeg-metal/src/tests/single_scratch.rs @@ -5,6 +5,15 @@ use super::*; #[test] #[ignore = "allocation diagnostic; run explicitly with --ignored --nocapture"] fn metal_single_session_buffer_allocation_report() { + single_session_buffer_allocations(false); +} + +#[test] +fn warm_single_session_reuses_temporary_buffers() { + single_session_buffer_allocations(true); +} + +fn single_session_buffer_allocations(require_reuse: bool) { if !should_run_metal_runtime() { return; } @@ -32,6 +41,16 @@ fn metal_single_session_buffer_allocation_report() { let private = compute::jpeg_private_buffer_allocations_for_test(); let shared = compute::jpeg_shared_buffer_allocations_for_test(); assert_eq!(surface.residency(), SurfaceResidency::MetalResidentDecode); + if require_reuse && iteration > 0 { + assert_eq!( + private, 0, + "warm {sampling}/{fmt:?} temporary private buffers" + ); + assert_eq!( + shared, 1, + "warm {sampling}/{fmt:?} retains only fresh shared output allocation" + ); + } let mut actual = vec![0; expected.len()]; let stride = surface.dimensions().0 as usize * fmt.bytes_per_pixel(); surface.download_into(&mut actual, stride).expect("pixels"); @@ -41,3 +60,150 @@ fn metal_single_session_buffer_allocation_report() { } } } + +fn distinct_jpeg(side: u16, variant: u8, sampling: jpeg_encoder::SamplingFactor) -> Vec { + let mut pixels = j2k_test_support::gpu_bench_rgb8(u32::from(side), u32::from(side)); + for sample in &mut pixels { + *sample = sample.wrapping_add(variant.wrapping_mul(37)); + } + let mut bytes = Vec::new(); + let mut encoder = jpeg_encoder::Encoder::new(&mut bytes, 90); + encoder.set_sampling_factor(sampling); + encoder + .encode(&pixels, side, side, jpeg_encoder::ColorType::Rgb) + .unwrap(); + bytes +} + +fn check_retained_outputs(retained: &[(Surface, Vec)]) { + for (surface, expected) in retained { + assert_eq!(surface.residency(), SurfaceResidency::MetalResidentDecode); + // First readback is deliberately after all subsequent decodes. Cached + // host bytes must not hide an overwritten device allocation. + let actual = surface.as_bytes().expect("retained surface pixels"); + assert!( + actual.as_ref() == expected.as_slice(), + "retained device output changed" + ); + } +} + +#[test] +fn single_surfaces_survive_distinct_size_and_sampling_reuse() { + if !should_run_metal_runtime() { + return; + } + use jpeg_encoder::SamplingFactor::{F_1_1, F_2_1, F_2_2}; + let session = MetalBackendSession::system_default().unwrap(); + let mut retained = Vec::new(); + for (side, variant, sampling) in [ + (64, 0, F_2_2), + (64, 1, F_2_2), + (32, 2, F_2_1), + (48, 3, F_1_1), + (16, 4, F_2_2), + (64, 5, F_2_2), + ] { + let bytes = distinct_jpeg(side, variant, sampling); + let mut decoder = Decoder::new(&bytes).unwrap(); + for fmt in [PixelFormat::Gray8, PixelFormat::Rgb8, PixelFormat::Rgba8] { + let (expected, _) = CpuDecoder::new(&bytes) + .unwrap() + .decode_request(DecodeRequest::full(fmt)) + .unwrap(); + let surface = decoder + .decode_to_device_with_session(fmt, &session) + .unwrap(); + retained.push((surface, expected)); + } + } + check_retained_outputs(&retained); +} + +#[test] +fn partial_single_surfaces_survive_operation_and_size_reuse() { + if !should_run_metal_runtime() { + return; + } + let mut retained = Vec::new(); + for (side, variant) in [(64, 1), (32, 2), (48, 3), (64, 4)] { + let bytes = distinct_jpeg(side, variant, jpeg_encoder::SamplingFactor::F_1_1); + let mut decoder = Decoder::new(&bytes).unwrap(); + let roi = Rect { + x: 8, + y: 8, + w: u32::from(side) - 16, + h: u32::from(side) - 16, + }; + let native_roi = j2k_jpeg::Rect { + x: roi.x, + y: roi.y, + w: roi.w, + h: roi.h, + }; + for fmt in [PixelFormat::Gray8, PixelFormat::Rgb8, PixelFormat::Rgba8] { + for (request, native) in [ + ( + MetalDecodeRequest::full(fmt, BackendRequest::Metal), + DecodeRequest::full(fmt), + ), + ( + MetalDecodeRequest::scaled(fmt, Downscale::Half, BackendRequest::Metal), + DecodeRequest::scaled(fmt, Downscale::Half), + ), + ( + MetalDecodeRequest::region(fmt, roi, BackendRequest::Metal), + DecodeRequest::region(fmt, native_roi), + ), + ( + MetalDecodeRequest::region_scaled( + fmt, + roi, + Downscale::Quarter, + BackendRequest::Metal, + ), + DecodeRequest::region_scaled(fmt, native_roi, Downscale::Quarter), + ), + ] { + let (expected, _) = CpuDecoder::new(&bytes) + .unwrap() + .decode_request(native) + .unwrap(); + let surface = decoder.decode_request_to_device(request).unwrap(); + retained.push((surface, expected)); + } + } + } + check_retained_outputs(&retained); +} + +#[test] +fn simultaneous_single_calls_keep_independent_surfaces() { + if !should_run_metal_runtime() { + return; + } + let session = MetalBackendSession::system_default().unwrap(); + session.runtime_result().as_ref().expect("runtime"); + let barrier = std::sync::Barrier::new(2); + std::thread::scope(|scope| { + for variant in [1, 2] { + let session = &session; + let barrier = &barrier; + scope.spawn(move || { + let bytes = distinct_jpeg(64, variant, jpeg_encoder::SamplingFactor::F_2_2); + let mut decoder = Decoder::new(&bytes).unwrap(); + let mut retained = Vec::new(); + barrier.wait(); + for fmt in [PixelFormat::Rgb8, PixelFormat::Gray8, PixelFormat::Rgba8] { + let (expected, _) = CpuDecoder::new(&bytes) + .unwrap() + .decode_request(DecodeRequest::full(fmt)) + .unwrap(); + let surface = decoder.decode_to_device_with_session(fmt, session).unwrap(); + retained.push((surface, expected)); + } + check_retained_outputs(&retained); + }); + } + }); +} From f0c5398d48f7615c0992bef1898fa7ac64138346 Mon Sep 17 00:00:00 2001 From: GF Date: Tue, 8 Sep 2026 05:47:43 -0400 Subject: [PATCH 2/4] refactor(jpeg-metal): clarify scratch command ownership --- .../src/compute/pack_dispatch/subsampled.rs | 4 ++-- .../j2k-jpeg-metal/src/compute/single_decode/fast444.rs | 9 ++++----- .../src/compute/single_decode/subsampled.rs | 4 ++++ 3 files changed, 10 insertions(+), 7 deletions(-) diff --git a/crates/j2k-jpeg-metal/src/compute/pack_dispatch/subsampled.rs b/crates/j2k-jpeg-metal/src/compute/pack_dispatch/subsampled.rs index 7e981f461..161790a4f 100644 --- a/crates/j2k-jpeg-metal/src/compute/pack_dispatch/subsampled.rs +++ b/crates/j2k-jpeg-metal/src/compute/pack_dispatch/subsampled.rs @@ -103,7 +103,7 @@ pub(in crate::compute) fn encode_fast_subsampled_region_batch_item( clippy::similar_names, reason = "Cb and Cr are normative JPEG component names" )] +#[expect( + clippy::too_many_lines, + reason = "the full decode keeps scratch acquisition, ordered command encoding, completion, and status publication together" +)] fn decode_fast_subsampled_to_rgb_buffer( runtime: &MetalRuntime, packet: Option<&P>, From c4feebd9f13806aea0a54fa8ac9578258d9bf57d Mon Sep 17 00:00:00 2001 From: GF Date: Tue, 8 Sep 2026 05:48:42 -0400 Subject: [PATCH 3/4] test(jpeg-metal): organize scratch fixture imports --- crates/j2k-jpeg-metal/src/tests/single_scratch.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/crates/j2k-jpeg-metal/src/tests/single_scratch.rs b/crates/j2k-jpeg-metal/src/tests/single_scratch.rs index 881bda92c..96c9f8220 100644 --- a/crates/j2k-jpeg-metal/src/tests/single_scratch.rs +++ b/crates/j2k-jpeg-metal/src/tests/single_scratch.rs @@ -1,6 +1,7 @@ // SPDX-License-Identifier: MIT OR Apache-2.0 use super::*; +use jpeg_encoder::SamplingFactor::{F_1_1, F_2_1, F_2_2}; #[test] #[ignore = "allocation diagnostic; run explicitly with --ignored --nocapture"] @@ -93,7 +94,6 @@ fn single_surfaces_survive_distinct_size_and_sampling_reuse() { if !should_run_metal_runtime() { return; } - use jpeg_encoder::SamplingFactor::{F_1_1, F_2_1, F_2_2}; let session = MetalBackendSession::system_default().unwrap(); let mut retained = Vec::new(); for (side, variant, sampling) in [ From 56aea9460f681341d43684a3fb70ef897bcc66d3 Mon Sep 17 00:00:00 2001 From: GF Date: Tue, 8 Sep 2026 06:02:30 -0400 Subject: [PATCH 4/4] perf(jpeg-metal): retain fresh 444 buffers after measured regression --- .../src/compute/single_decode/fast444.rs | 130 +++++++----------- .../src/compute/single_decode/scratch.rs | 46 +------ .../src/tests/single_scratch.rs | 18 ++- 3 files changed, 65 insertions(+), 129 deletions(-) diff --git a/crates/j2k-jpeg-metal/src/compute/single_decode/fast444.rs b/crates/j2k-jpeg-metal/src/compute/single_decode/fast444.rs index 7db47d918..e14ce3c31 100644 --- a/crates/j2k-jpeg-metal/src/compute/single_decode/fast444.rs +++ b/crates/j2k-jpeg-metal/src/compute/single_decode/fast444.rs @@ -1,6 +1,5 @@ // SPDX-License-Identifier: MIT OR Apache-2.0 -use super::scratch::decode_buffers as scratch_decode_buffers; use crate::metal_types::prelude::*; use super::super::{ @@ -9,10 +8,11 @@ use super::super::{ entropy_decode_thread_count, fast444_params, fast444_region_params, fast444_scaled_params, fast444_scaled_region_params, fast_decode_status_error, fast_packet_huffman_tables, first_decode_error_status, mcu_range_for_rect, new_command_buffer, new_compute_command_encoder, - new_private_buffer, new_shared_buffer_with_data, pixel_format_to_out_format, - restart_offsets_buffer, restart_work_for_mcu_range, CpuDecoder, Error, FastDecodeEntropyInputs, - JpegColorSpace, JpegDecodeStatus, JpegFast444PacketV1, JpegFast444Params, - JpegFast444ScaledParams, MetalRuntime, PixelFormat, PlaneMode, PlaneStage, Surface, + new_decode_plane_buffer, new_private_buffer, new_shared_buffer_with_data, + pixel_format_to_out_format, restart_offsets_buffer, restart_work_for_mcu_range, CpuDecoder, + Error, FastDecodeEntropyInputs, JpegColorSpace, JpegDecodeStatus, JpegFast444PacketV1, + JpegFast444Params, JpegFast444ScaledParams, MetalRuntime, PixelFormat, PlaneMode, PlaneStage, + Surface, }; #[cfg(target_os = "macos")] @@ -40,29 +40,23 @@ pub(in crate::compute) fn try_decode_fast444_to_surface( let params = fast444_params(packet)?; let mode = fast444_plane_mode(decoder); let plane_len = params.width as usize * params.height as usize; + let y_plane = new_decode_plane_buffer( + &runtime.device, + plane_len, + fmt == PixelFormat::Gray8 && mode != PlaneMode::Rgb, + )?; + let chroma_blue_plane = new_private_buffer(&runtime.device, plane_len)?; + let chroma_red_plane = new_private_buffer(&runtime.device, plane_len)?; let decode_threads = entropy_decode_thread_count( packet.restart_interval_mcus, packet.restart_offsets.len(), packet.entropy_checkpoints.len(), ); - let mut scratch = runtime.batch_scratch()?; - let buffers = scratch_decode_buffers( - &mut scratch, - &runtime.device, - plane_len, - fmt == PixelFormat::Gray8 && mode != PlaneMode::Rgb, - decode_threads, - &packet.entropy_bytes, - &packet.restart_offsets, - &packet.entropy_checkpoints, - )?; - let y_plane = buffers.plane0; - let chroma_blue_plane = buffers.plane1; - let chroma_red_plane = buffers.plane2; - let status_buffer = buffers.status; - let entropy_buffer = buffers.packet.entropy; - let restart_offsets_buffer = buffers.packet.restart_offsets; - let entropy_checkpoints_buffer = buffers.packet.checkpoints; + 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 (dc_tables, ac_tables) = fast_packet_huffman_tables(packet); @@ -211,29 +205,23 @@ pub(in crate::compute) fn try_decode_fast444_region_to_surface( )?; let mode = fast444_plane_mode(decoder); let plane_len = params.width as usize * params.height as usize; + let y_plane = new_decode_plane_buffer( + &runtime.device, + plane_len, + fmt == PixelFormat::Gray8 && mode != PlaneMode::Rgb, + )?; + let chroma_blue_plane = new_private_buffer(&runtime.device, plane_len)?; + let chroma_red_plane = new_private_buffer(&runtime.device, plane_len)?; let decode_threads = entropy_decode_thread_count( packet.restart_interval_mcus, restart_offsets.len(), packet.entropy_checkpoints.len(), ); - let mut scratch = runtime.batch_scratch()?; - let buffers = scratch_decode_buffers( - &mut scratch, - &runtime.device, - plane_len, - fmt == PixelFormat::Gray8 && mode != PlaneMode::Rgb, - decode_threads, - &packet.entropy_bytes, - restart_offsets, - &packet.entropy_checkpoints, - )?; - let y_plane = buffers.plane0; - let chroma_blue_plane = buffers.plane1; - let chroma_red_plane = buffers.plane2; - let status_buffer = buffers.status; - let entropy_buffer = buffers.packet.entropy; - let restart_offsets_buffer = buffers.packet.restart_offsets; - let entropy_checkpoints_buffer = buffers.packet.checkpoints; + 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 (dc_tables, ac_tables) = fast_packet_huffman_tables(packet); @@ -298,29 +286,23 @@ pub(in crate::compute) fn try_decode_fast444_scaled_to_surface( let mode = fast444_plane_mode(decoder); let plane_len = params.scaled_width as usize * params.scaled_height as usize; + let y_plane = new_decode_plane_buffer( + &runtime.device, + plane_len, + fmt == PixelFormat::Gray8 && mode != PlaneMode::Rgb, + )?; + let chroma_blue_plane = new_private_buffer(&runtime.device, plane_len)?; + let chroma_red_plane = new_private_buffer(&runtime.device, plane_len)?; let decode_threads = entropy_decode_thread_count( packet.restart_interval_mcus, packet.restart_offsets.len(), packet.entropy_checkpoints.len(), ); - let mut scratch = runtime.batch_scratch()?; - let buffers = scratch_decode_buffers( - &mut scratch, - &runtime.device, - plane_len, - fmt == PixelFormat::Gray8 && mode != PlaneMode::Rgb, - decode_threads, - &packet.entropy_bytes, - &packet.restart_offsets, - &packet.entropy_checkpoints, - )?; - let y_plane = buffers.plane0; - let chroma_blue_plane = buffers.plane1; - let chroma_red_plane = buffers.plane2; - let status_buffer = buffers.status; - let entropy_buffer = buffers.packet.entropy; - let restart_offsets_buffer = buffers.packet.restart_offsets; - let entropy_checkpoints_buffer = buffers.packet.checkpoints; + 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 (dc_tables, ac_tables) = fast_packet_huffman_tables(packet); @@ -429,29 +411,23 @@ pub(in crate::compute) fn try_decode_fast444_scaled_region_to_surface_with_mode_ )?; let plane_len = params.scaled_width as usize * params.scaled_height as usize; + let y_plane = new_decode_plane_buffer( + &runtime.device, + plane_len, + fmt == PixelFormat::Gray8 && mode != PlaneMode::Rgb, + )?; + let chroma_blue_plane = new_private_buffer(&runtime.device, plane_len)?; + let chroma_red_plane = new_private_buffer(&runtime.device, plane_len)?; let decode_threads = entropy_decode_thread_count( packet.restart_interval_mcus, restart_offsets.len(), packet.entropy_checkpoints.len(), ); - let mut scratch = runtime.batch_scratch()?; - let buffers = scratch_decode_buffers( - &mut scratch, - &runtime.device, - plane_len, - fmt == PixelFormat::Gray8 && mode != PlaneMode::Rgb, - decode_threads, - &packet.entropy_bytes, - restart_offsets, - &packet.entropy_checkpoints, - )?; - let y_plane = buffers.plane0; - let chroma_blue_plane = buffers.plane1; - let chroma_red_plane = buffers.plane2; - let status_buffer = buffers.status; - let entropy_buffer = buffers.packet.entropy; - let restart_offsets_buffer = buffers.packet.restart_offsets; - let entropy_checkpoints_buffer = buffers.packet.checkpoints; + 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 (dc_tables, ac_tables) = fast_packet_huffman_tables(packet); diff --git a/crates/j2k-jpeg-metal/src/compute/single_decode/scratch.rs b/crates/j2k-jpeg-metal/src/compute/single_decode/scratch.rs index c7c0cd2a7..9178a9f1c 100644 --- a/crates/j2k-jpeg-metal/src/compute/single_decode/scratch.rs +++ b/crates/j2k-jpeg-metal/src/compute/single_decode/scratch.rs @@ -1,7 +1,7 @@ // SPDX-License-Identifier: MIT OR Apache-2.0 use crate::{ - buffers::{checked_copy_bytes_to_buffer_at, new_decode_plane_buffer, MetalBatchScratch}, + buffers::{checked_copy_bytes_to_buffer_at, MetalBatchScratch}, metal_types::{Buffer, DeviceRef}, Error, }; @@ -16,14 +16,6 @@ pub(in crate::compute) struct ScratchPacketBuffers { pub(in crate::compute) checkpoints: Buffer, } -pub(in crate::compute) struct ScratchDecodeBuffers { - pub(in crate::compute) plane0: Buffer, - pub(in crate::compute) plane1: Buffer, - pub(in crate::compute) plane2: Buffer, - pub(in crate::compute) status: Buffer, - pub(in crate::compute) packet: ScratchPacketBuffers, -} - pub(in crate::compute) fn status_buffer( scratch: &mut MetalBatchScratch, device: &DeviceRef, @@ -81,39 +73,3 @@ pub(in crate::compute) fn packet_buffers( checkpoints: checkpoint_buffer, }) } - -#[expect( - clippy::too_many_arguments, - reason = "the scratch preparation mirrors the three plane and packet inputs bound together" -)] -pub(in crate::compute) fn decode_buffers( - scratch: &mut MetalBatchScratch, - device: &DeviceRef, - plane_len: usize, - plane0_returned: bool, - decode_threads: u32, - entropy: &[u8], - restart_offsets: &[u32], - checkpoints: &[JpegEntropyCheckpointV1], -) -> Result { - // The fast decode kernels cover the active plane rectangles through the - // bounded deposit helpers in shaders_decode_helpers.metal. Callers reject - // any failed status before publishing output, so successful reuse does not - // require preserving or clearing earlier plane contents. - let plane0 = if plane0_returned { - new_decode_plane_buffer(device, plane_len, true)? - } else { - scratch.private_buffer(device, "single_decode_y", plane_len)? - }; - let plane1 = scratch.private_buffer(device, "single_decode_cb", plane_len)?; - let plane2 = scratch.private_buffer(device, "single_decode_cr", plane_len)?; - let packet = packet_buffers(scratch, device, entropy, restart_offsets, checkpoints)?; - let status = status_buffer(scratch, device, decode_threads)?; - Ok(ScratchDecodeBuffers { - plane0, - plane1, - plane2, - status, - packet, - }) -} diff --git a/crates/j2k-jpeg-metal/src/tests/single_scratch.rs b/crates/j2k-jpeg-metal/src/tests/single_scratch.rs index 96c9f8220..89fe43555 100644 --- a/crates/j2k-jpeg-metal/src/tests/single_scratch.rs +++ b/crates/j2k-jpeg-metal/src/tests/single_scratch.rs @@ -10,7 +10,7 @@ fn metal_single_session_buffer_allocation_report() { } #[test] -fn warm_single_session_reuses_temporary_buffers() { +fn warm_subsampled_session_reuses_temporaries_with_444_control() { single_session_buffer_allocations(true); } @@ -43,13 +43,17 @@ fn single_session_buffer_allocations(require_reuse: bool) { let shared = compute::jpeg_shared_buffer_allocations_for_test(); assert_eq!(surface.residency(), SurfaceResidency::MetalResidentDecode); if require_reuse && iteration > 0 { + // 4:4:4 pooling was rejected by the latency gate; retain + // its fresh-allocation control while 4:2:0/4:2:2 reuse. + let expected = if sampling == "444" { + (if fmt == PixelFormat::Gray8 { 2 } else { 3 }, 5) + } else { + (0, 1) + }; assert_eq!( - private, 0, - "warm {sampling}/{fmt:?} temporary private buffers" - ); - assert_eq!( - shared, 1, - "warm {sampling}/{fmt:?} retains only fresh shared output allocation" + (private, shared), + expected, + "warm {sampling}/{fmt:?} allocations including fresh output" ); } let mut actual = vec![0; expected.len()];