From 070ef2959796b02dbac18b89f444e48b52d72453 Mon Sep 17 00:00:00 2001 From: GF Date: Tue, 8 Sep 2026 05:26:35 -0400 Subject: [PATCH] test(jpeg-metal): benchmark retained single-image sessions --- crates/j2k-jpeg-metal/benches/compare.rs | 255 +++++++++++++++++- .../src/compute/scratch_pool.rs | 19 ++ crates/j2k-jpeg-metal/src/tests.rs | 3 + .../src/tests/single_scratch.rs | 43 +++ 4 files changed, 318 insertions(+), 2 deletions(-) create mode 100644 crates/j2k-jpeg-metal/src/tests/single_scratch.rs diff --git a/crates/j2k-jpeg-metal/benches/compare.rs b/crates/j2k-jpeg-metal/benches/compare.rs index d4c09ae2f..2cb3f422d 100644 --- a/crates/j2k-jpeg-metal/benches/compare.rs +++ b/crates/j2k-jpeg-metal/benches/compare.rs @@ -18,7 +18,10 @@ use j2k_jpeg_metal::{ MetalSession, ScratchPool, ViewportTile, ViewportWorkload, }; #[cfg(target_os = "macos")] -use j2k_jpeg_metal::{MetalBackendSession, MetalBatchTextureOutput}; +use j2k_jpeg_metal::{ + MetalBackendSession, MetalBatchOutputBuffer, MetalBatchTextureOutput, MetalBufferBatchTarget, + Rgb8MetalBatchOp, Rgb8MetalBatchRequest, Rgb8MetalBatchSource, SurfaceResidency, +}; use jpeg_encoder::{ColorType, Encoder, SamplingFactor}; use std::collections::HashSet; @@ -126,7 +129,23 @@ fn generated_rgb_jpeg( sampling: SamplingFactor, restart_interval: Option, ) -> Vec { - let rgb = j2k_test_support::gpu_bench_rgb8(u32::from(width), u32::from(height)); + generated_rgb_jpeg_variant(width, height, sampling, restart_interval, 0) +} + +fn generated_rgb_jpeg_variant( + width: u16, + height: u16, + sampling: SamplingFactor, + restart_interval: Option, + variant: u8, +) -> Vec { + let mut rgb = j2k_test_support::gpu_bench_rgb8(u32::from(width), u32::from(height)); + if variant != 0 { + for (index, sample) in rgb.iter_mut().enumerate() { + let position = u8::try_from(index % 251).expect("generated JPEG position fits u8"); + *sample = sample.wrapping_add(position.wrapping_mul(variant)); + } + } let mut jpeg = Vec::new(); let mut encoder = Encoder::new(&mut jpeg, 90); @@ -868,6 +887,237 @@ fn bench_full_and_tile_decode_groups(c: &mut Criterion, inputs: &[BenchInput], h wsi_tile_batch_rgb.finish(); } +#[cfg(target_os = "macos")] +fn pixel_format_label(fmt: PixelFormat) -> &'static str { + match fmt { + PixelFormat::Gray8 => "gray8", + PixelFormat::Rgb8 => "rgb8", + PixelFormat::Rgba8 => "rgba8", + _ => "other", + } +} + +#[cfg(target_os = "macos")] +fn native_request_pixels(bytes: &[u8], request: DecodeRequest) -> Vec { + CpuDecoder::new(bytes) + .expect("native retained-session benchmark decoder") + .decode_request(request) + .expect("native retained-session benchmark decode") + .0 +} + +#[cfg(target_os = "macos")] +fn assert_metal_surface_pixels(surface: &j2k_jpeg_metal::Surface, expected: &[u8]) { + assert_eq!(surface.residency(), SurfaceResidency::MetalResidentDecode); + let actual = surface.as_bytes().expect("Metal benchmark surface bytes"); + let actual = actual.as_ref(); + if actual != expected { + let first_mismatch = actual + .iter() + .zip(expected) + .position(|(actual, expected)| actual != expected); + let mismatch_count = actual + .iter() + .zip(expected) + .filter(|(actual, expected)| actual != expected) + .count() + + actual.len().abs_diff(expected.len()); + let max_delta = actual + .iter() + .zip(expected) + .map(|(actual, expected)| actual.abs_diff(*expected)) + .max() + .unwrap_or(0); + panic!( + "Metal benchmark pixels differ: actual_len={} expected_len={} first_mismatch={first_mismatch:?} mismatch_count={mismatch_count} max_delta={max_delta}", + actual.len(), + expected.len(), + ); + } +} + +#[cfg(target_os = "macos")] +fn bench_retained_full_rows(c: &mut Criterion) { + let fixtures = [ + ( + "420_nr", + generated_rgb_jpeg(256, 256, SamplingFactor::F_2_2, None), + "fast420", + ), + ( + "420_r2", + generated_rgb_jpeg(256, 256, SamplingFactor::F_2_2, Some(2)), + "fast420", + ), + ( + "422_nr", + generated_rgb_jpeg(256, 256, SamplingFactor::F_2_1, None), + "fast422", + ), + ( + "444_nr", + generated_rgb_jpeg(256, 256, SamplingFactor::F_1_1, None), + "fast444", + ), + ]; + let mut group = c.benchmark_group("jpeg_metal_single_retained"); + for (fixture_name, bytes, expected_family) in &fixtures { + let packet = fast_packet_plan(bytes).expect("retained-session fast packet"); + assert_eq!(fast_packet_family_label(packet), *expected_family); + for fmt in [PixelFormat::Gray8, PixelFormat::Rgb8, PixelFormat::Rgba8] { + let expected = native_request_pixels(bytes, DecodeRequest::full(fmt)); + let session = MetalBackendSession::system_default().expect("Metal backend session"); + let mut decoder = Decoder::new(bytes).expect("retained Metal decoder"); + let probe = decoder + .decode_to_device_with_session(fmt, &session) + .expect("retained-session Metal probe"); + assert_metal_surface_pixels(&probe, &expected); + drop(probe); + + group.bench_function( + format!("{fixture_name}/full/{}", pixel_format_label(fmt)), + |b| { + b.iter(|| { + let surface = decoder + .decode_to_device_with_session(fmt, &session) + .expect("retained-session Metal decode"); + std::hint::black_box(surface); + }); + }, + ); + } + } + group.finish(); +} + +#[cfg(target_os = "macos")] +fn bench_single_request_controls(c: &mut Criterion) { + let scaled_bytes = generated_rgb_jpeg(256, 256, SamplingFactor::F_2_2, Some(2)); + let region_bytes = generated_rgb_jpeg(256, 256, SamplingFactor::F_1_1, None); + let controls = [ + ( + "420_r2", + "scaled_q4", + scaled_bytes.as_slice(), + MetalDecodeRequest::scaled( + PixelFormat::Rgb8, + Downscale::Quarter, + BackendRequest::Metal, + ), + DecodeRequest::scaled(PixelFormat::Rgb8, Downscale::Quarter), + ), + ( + "444_nr", + "region192_q4", + region_bytes.as_slice(), + MetalDecodeRequest::region_scaled( + PixelFormat::Rgb8, + centered_roi((256, 256), 192), + Downscale::Quarter, + BackendRequest::Metal, + ), + DecodeRequest::region_scaled( + PixelFormat::Rgb8, + to_jpeg_rect(centered_roi((256, 256), 192)), + Downscale::Quarter, + ), + ), + ]; + let mut group = c.benchmark_group("jpeg_metal_single_request_controls"); + for (fixture_name, name, bytes, metal_request, native_request) in controls { + let expected = native_request_pixels(bytes, native_request); + let mut decoder = Decoder::new(bytes).expect("retained request-control decoder"); + let probe = decoder + .decode_request_to_device(metal_request) + .expect("global-runtime request control probe"); + assert_metal_surface_pixels(&probe, &expected); + drop(probe); + group.bench_function( + format!("{fixture_name}/global_runtime_control/{name}"), + |b| { + b.iter(|| { + let surface = decoder + .decode_request_to_device(metal_request) + .expect("global-runtime request control"); + std::hint::black_box(surface); + }); + }, + ); + } + group.finish(); +} + +#[cfg(target_os = "macos")] +fn bench_retained_distinct_batch(c: &mut Criterion) { + let distinct_bytes = (1..=4) + .map(|variant| { + generated_rgb_jpeg_variant(256, 256, SamplingFactor::F_2_2, Some(2), variant) + }) + .collect::>(); + for bytes in &distinct_bytes { + let packet = fast_packet_plan(bytes).expect("distinct batch fast packet"); + assert_eq!(fast_packet_family_label(packet), "fast420"); + } + let expected = distinct_bytes + .iter() + .map(|bytes| native_request_pixels(bytes, DecodeRequest::full(PixelFormat::Rgb8))) + .collect::>(); + for first in 0..distinct_bytes.len() { + for second in first + 1..distinct_bytes.len() { + assert_ne!( + distinct_bytes[first], distinct_bytes[second], + "distinct batch JPEG inputs {first} and {second}" + ); + assert_ne!( + expected[first], expected[second], + "distinct batch pixels {first} and {second}" + ); + } + } + let decoders = distinct_bytes + .iter() + .map(|bytes| Decoder::new(bytes).expect("distinct retained Metal decoder")) + .collect::>(); + let decoder_refs = decoders.iter().collect::>(); + let session = MetalBackendSession::system_default().expect("distinct batch Metal session"); + let output = MetalBatchOutputBuffer::new_rgb8_tiles(&session, (256, 256), decoder_refs.len()) + .expect("distinct batch output"); + let decode_batch = || { + Codec::decode_rgb8_batch_into_buffer_with_session( + Rgb8MetalBatchRequest { + source: Rgb8MetalBatchSource::Decoders(&decoder_refs), + op: Rgb8MetalBatchOp::Full, + }, + MetalBufferBatchTarget::Reusable(&output), + &session, + ) + .expect("distinct retained-session batch") + }; + let probe = decode_batch(); + assert_eq!(probe.len(), expected.len()); + for (surface, expected) in probe.into_iter().zip(&expected) { + assert_metal_surface_pixels(&surface.expect("distinct batch surface"), expected); + } + let mut batch_group = c.benchmark_group("jpeg_metal_single_retained_batch_control"); + batch_group.bench_function("420_r2/distinct4/rgb8", |b| { + b.iter(|| std::hint::black_box(decode_batch())); + }); + batch_group.finish(); +} + +#[cfg(target_os = "macos")] +fn bench_retained_single_decode(c: &mut Criterion, has_metal: bool) { + if !has_metal { + return; + } + bench_retained_full_rows(c); + bench_single_request_controls(c); + bench_retained_distinct_batch(c); +} + +#[cfg(not(target_os = "macos"))] +fn bench_retained_single_decode(_c: &mut Criterion, _has_metal: bool) {} + fn bench_region_and_scaled_decode_groups( c: &mut Criterion, inputs: &[BenchInput], @@ -1374,6 +1624,7 @@ fn bench_compare(c: &mut Criterion) { bench_fast_packet_planning(c, &inputs); bench_full_and_tile_decode_groups(c, &inputs, has_metal); + bench_retained_single_decode(c, has_metal); bench_resident_texture_batches(c, &inputs, has_metal); bench_resident_viewport_outputs(c, &inputs, has_metal); bench_region_and_scaled_decode_groups(c, &inputs, has_metal); diff --git a/crates/j2k-jpeg-metal/src/compute/scratch_pool.rs b/crates/j2k-jpeg-metal/src/compute/scratch_pool.rs index 6125fc166..0fbf96016 100644 --- a/crates/j2k-jpeg-metal/src/compute/scratch_pool.rs +++ b/crates/j2k-jpeg-metal/src/compute/scratch_pool.rs @@ -99,3 +99,22 @@ impl Drop for BatchScratchLease<'_> { self.pool.released.notify_all(); } } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn scratch_slots_fail_closed_after_panicking_owner() { + let pool = BatchScratchPool::default(); + let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + let _lease = pool.acquire().expect("scratch lease"); + panic!("abandon scratch during an owner panic"); + })); + assert!(result.is_err()); + assert!(matches!( + pool.acquire(), + Err(Error::MetalStatePoisoned { .. }) + )); + } +} diff --git a/crates/j2k-jpeg-metal/src/tests.rs b/crates/j2k-jpeg-metal/src/tests.rs index 522d3963d..26bd54157 100644 --- a/crates/j2k-jpeg-metal/src/tests.rs +++ b/crates/j2k-jpeg-metal/src/tests.rs @@ -2155,3 +2155,6 @@ fn warm_session_reuses_shared_upload_buffers_for_reusable_output_batches() { "warm session batch should reuse shared upload/status buffers" ); } + +#[cfg(target_os = "macos")] +mod single_scratch; diff --git a/crates/j2k-jpeg-metal/src/tests/single_scratch.rs b/crates/j2k-jpeg-metal/src/tests/single_scratch.rs new file mode 100644 index 000000000..db2197747 --- /dev/null +++ b/crates/j2k-jpeg-metal/src/tests/single_scratch.rs @@ -0,0 +1,43 @@ +// SPDX-License-Identifier: MIT OR Apache-2.0 + +use super::*; + +#[test] +#[ignore = "allocation diagnostic; run explicitly with --ignored --nocapture"] +fn metal_single_session_buffer_allocation_report() { + if !should_run_metal_runtime() { + return; + } + for (sampling, input) in [ + ("420", BASELINE_420), + ("420-restart", BASELINE_420_RESTART), + ("422", BASELINE_422), + ("444", BASELINE_444), + ] { + for fmt in [PixelFormat::Gray8, PixelFormat::Rgb8, PixelFormat::Rgba8] { + let session = MetalBackendSession::system_default().expect("session"); + // Exclude one-time runtime resources from per-decode allocation counts. + session.runtime_result().as_ref().expect("runtime"); + let mut decoder = Decoder::new(input).expect("decoder"); + let (expected, _) = CpuDecoder::new(input) + .expect("native decoder") + .decode_request(DecodeRequest::full(fmt)) + .expect("native pixels"); + for iteration in 0..3 { + compute::reset_jpeg_private_buffer_allocations_for_test(); + compute::reset_jpeg_shared_buffer_allocations_for_test(); + let surface = decoder + .decode_to_device_with_session(fmt, &session) + .expect("single decode"); + let private = compute::jpeg_private_buffer_allocations_for_test(); + let shared = compute::jpeg_shared_buffer_allocations_for_test(); + assert_eq!(surface.residency(), SurfaceResidency::MetalResidentDecode); + 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"); + assert_eq!(actual, expected); + println!("metal_single_session_allocations sampling={sampling} format={fmt:?} iteration={iteration} private_buffers={private} shared_buffers={shared} returned_surfaces=1"); + } + } + } +}