diff --git a/crates/j2k-jpeg-metal/src/buffers.rs b/crates/j2k-jpeg-metal/src/buffers.rs index 63102a60..a0c926f8 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 150a2a57..161790a4 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 0fbf9601..4cf6f772 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 2159eb16..e8024683 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/scratch.rs b/crates/j2k-jpeg-metal/src/compute/single_decode/scratch.rs new file mode 100644 index 00000000..9178a9f1 --- /dev/null +++ b/crates/j2k-jpeg-metal/src/compute/single_decode/scratch.rs @@ -0,0 +1,75 @@ +// SPDX-License-Identifier: MIT OR Apache-2.0 + +use crate::{ + buffers::{checked_copy_bytes_to_buffer_at, 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) 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, + }) +} 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 8838cee2..c716a0f1 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); @@ -68,12 +79,17 @@ fn try_decode_fast_subsampled_to_surface( 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>, 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 +101,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 +236,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 +248,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 +278,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 +407,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 +507,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 42309cd2..4513ddd9 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 5da3ca17..14f78b1e 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 db219774..89fe4355 100644 --- a/crates/j2k-jpeg-metal/src/tests/single_scratch.rs +++ b/crates/j2k-jpeg-metal/src/tests/single_scratch.rs @@ -1,10 +1,20 @@ // 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"] fn metal_single_session_buffer_allocation_report() { + single_session_buffer_allocations(false); +} + +#[test] +fn warm_subsampled_session_reuses_temporaries_with_444_control() { + single_session_buffer_allocations(true); +} + +fn single_session_buffer_allocations(require_reuse: bool) { if !should_run_metal_runtime() { return; } @@ -32,6 +42,20 @@ 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 { + // 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, shared), + expected, + "warm {sampling}/{fmt:?} allocations including fresh output" + ); + } 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 +65,149 @@ 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; + } + 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); + }); + } + }); +}