From 33826b73c0689fa84e4dad265d685ad0f33bcd29 Mon Sep 17 00:00:00 2001 From: GF Date: Tue, 8 Sep 2026 06:45:00 -0400 Subject: [PATCH] perf(metal): copy completed IDWT results directly into host slices --- .../src/engine/decode_dispatch/idwt.rs | 44 ++++++++++----- .../decode_dispatch/idwt/irreversible.rs | 24 +++------ crates/j2k-metal/src/engine/direct_buffers.rs | 54 +++++++++++++++++-- crates/j2k-metal/src/engine/test_counters.rs | 8 --- crates/j2k-metal/src/idwt_host_slice_tests.rs | 9 +--- 5 files changed, 90 insertions(+), 49 deletions(-) diff --git a/crates/j2k-metal/src/engine/decode_dispatch/idwt.rs b/crates/j2k-metal/src/engine/decode_dispatch/idwt.rs index d608fb15..0176b94c 100644 --- a/crates/j2k-metal/src/engine/decode_dispatch/idwt.rs +++ b/crates/j2k-metal/src/engine/decode_dispatch/idwt.rs @@ -6,11 +6,35 @@ use crate::metal_types::prelude::*; use super::{ checked_buffer_copy_into, commit_and_wait_metal, copied_slice_buffer, dispatch_2d_pipeline, dispatch_3d_pipeline, hybrid_stage_signpost, label_compute_encoder, new_command_buffer, - new_compute_command_encoder, with_runtime, Buffer, CommandBufferRef, ComputeCommandEncoderRef, - DirectIdwtCommandBuffers, Error, J2kIdwtSingleDecompositionParams, + new_compute_command_encoder, new_shared_buffer, with_runtime, Buffer, CommandBufferRef, + ComputeCommandEncoderRef, DirectIdwtCommandBuffers, Error, J2kIdwtSingleDecompositionParams, J2kRepeatedIdwtSingleDecompositionParams, J2kSingleDecompositionIdwtJob, SIGNPOST_DECODE_HYBRID_IDWT_COMMAND_ENCODE, }; + +#[cfg(target_os = "macos")] +fn checked_host_output_layout( + width: u32, + height: u32, + output_len: usize, +) -> Result<(usize, usize), Error> { + let required_len = (width as usize) + .checked_mul(height as usize) + .ok_or_else(|| Error::MetalKernel { + message: "J2K Metal IDWT output length overflow".to_string(), + })?; + let required_bytes = required_len + .checked_mul(std::mem::size_of::()) + .ok_or_else(|| Error::MetalKernel { + message: "J2K Metal IDWT output byte length overflow".to_string(), + })?; + if output_len < required_len { + return Err(Error::MetalKernel { + message: "J2K Metal IDWT output slice is too small".to_string(), + }); + } + Ok((required_len, required_bytes)) +} #[cfg(target_os = "macos")] mod batched_irreversible; #[cfg(target_os = "macos")] @@ -39,12 +63,8 @@ pub(crate) fn decode_reversible53_single_decomposition_idwt( output: &mut [f32], ) -> Result<(), Error> { with_runtime(|runtime| { - let required_len = job.rect.width() as usize * job.rect.height() as usize; - if output.len() < required_len { - return Err(Error::MetalKernel { - message: "J2K Metal IDWT output slice is too small".to_string(), - }); - } + let (required_len, required_bytes) = + checked_host_output_layout(job.rect.width(), job.rect.height(), output.len())?; let params = J2kIdwtSingleDecompositionParams { x0: job.rect.x0, @@ -71,15 +91,11 @@ pub(crate) fn decode_reversible53_single_decomposition_idwt( hh_height: job.hh.rect.height(), }; + let decoded = new_shared_buffer(&runtime.device, required_bytes)?; let ll = copied_slice_buffer(&runtime.device, job.ll.coefficients)?; let hl = copied_slice_buffer(&runtime.device, job.hl.coefficients)?; let lh = copied_slice_buffer(&runtime.device, job.lh.coefficients)?; let hh = copied_slice_buffer(&runtime.device, job.hh.coefficients)?; - let decoded = copied_slice_buffer(&runtime.device, output)?; - #[cfg(test)] - crate::engine::test_counters::record_idwt_host_overwritten_output_upload( - std::mem::size_of_val(output), - ); let command_buffer = new_command_buffer(&runtime.queue)?; @@ -128,7 +144,7 @@ pub(crate) fn decode_reversible53_single_decomposition_idwt( ); encoder.endEncoding(); commit_and_wait_metal(&command_buffer)?; - checked_buffer_copy_into(&decoded, 0, output, "IDWT output")?; + checked_buffer_copy_into(&decoded, 0, &mut output[..required_len], "IDWT output")?; Ok(()) }) } diff --git a/crates/j2k-metal/src/engine/decode_dispatch/idwt/irreversible.rs b/crates/j2k-metal/src/engine/decode_dispatch/idwt/irreversible.rs index 7608efe9..4d3d3b7b 100644 --- a/crates/j2k-metal/src/engine/decode_dispatch/idwt/irreversible.rs +++ b/crates/j2k-metal/src/engine/decode_dispatch/idwt/irreversible.rs @@ -6,11 +6,11 @@ use crate::metal_types::prelude::*; use super::super::{ checked_buffer_copy_into, commit_and_wait_metal, copied_slice_buffer, dispatch_2d_pipeline, dispatch_3d_pipeline, hybrid_stage_signpost, label_compute_encoder, new_command_buffer, - new_compute_command_encoder, with_runtime, Buffer, CommandBufferRef, ComputeCommandEncoderRef, - Error, J2kIdwt97StepParams, J2kIdwtSingleDecompositionParams, J2kSingleDecompositionIdwtJob, - SIGNPOST_DECODE_HYBRID_IDWT_COMMAND_ENCODE, + new_compute_command_encoder, new_shared_buffer, with_runtime, Buffer, CommandBufferRef, + ComputeCommandEncoderRef, Error, J2kIdwt97StepParams, J2kIdwtSingleDecompositionParams, + J2kSingleDecompositionIdwtJob, SIGNPOST_DECODE_HYBRID_IDWT_COMMAND_ENCODE, }; -use super::{IdwtSubBandBuffers, SingleIdwtDispatch}; +use super::{checked_host_output_layout, IdwtSubBandBuffers, SingleIdwtDispatch}; use j2k_codec_math::dwt; #[cfg(test)] @@ -51,12 +51,8 @@ fn decode_irreversible97_staged_single_decomposition_idwt_with_high_pass( high_pass: f32, ) -> Result<(), Error> { with_runtime(|runtime| { - let required_len = job.rect.width() as usize * job.rect.height() as usize; - if output.len() < required_len { - return Err(Error::MetalKernel { - message: "J2K Metal IDWT output slice is too small".to_string(), - }); - } + let (required_len, required_bytes) = + checked_host_output_layout(job.rect.width(), job.rect.height(), output.len())?; let params = J2kIdwtSingleDecompositionParams { x0: job.rect.x0, @@ -83,15 +79,11 @@ fn decode_irreversible97_staged_single_decomposition_idwt_with_high_pass( hh_height: job.hh.rect.height(), }; + let decoded = new_shared_buffer(&runtime.device, required_bytes)?; let ll = copied_slice_buffer(&runtime.device, job.ll.coefficients)?; let hl = copied_slice_buffer(&runtime.device, job.hl.coefficients)?; let lh = copied_slice_buffer(&runtime.device, job.lh.coefficients)?; let hh = copied_slice_buffer(&runtime.device, job.hh.coefficients)?; - let decoded = copied_slice_buffer(&runtime.device, output)?; - #[cfg(test)] - crate::engine::test_counters::record_idwt_host_overwritten_output_upload( - std::mem::size_of_val(output), - ); let command_buffer = new_command_buffer(&runtime.queue)?; let encoder = new_compute_command_encoder(&command_buffer)?; dispatch_irreversible97_single_decomposition_buffers_in_encoder_with_high_pass( @@ -117,7 +109,7 @@ fn decode_irreversible97_staged_single_decomposition_idwt_with_high_pass( encoder.endEncoding(); commit_and_wait_metal(&command_buffer)?; - checked_buffer_copy_into(&decoded, 0, output, "IDWT output")?; + checked_buffer_copy_into(&decoded, 0, &mut output[..required_len], "IDWT output")?; Ok(()) }) } diff --git a/crates/j2k-metal/src/engine/direct_buffers.rs b/crates/j2k-metal/src/engine/direct_buffers.rs index 160b54bf..c9cdad2e 100644 --- a/crates/j2k-metal/src/engine/direct_buffers.rs +++ b/crates/j2k-metal/src/engine/direct_buffers.rs @@ -92,9 +92,57 @@ pub(crate) fn checked_buffer_copy_into( output: &mut [T], context: &str, ) -> Result<(), Error> { - let values = checked_buffer_slice_at(buffer, byte_offset, output.len(), context)?; - output.copy_from_slice(&values); - Ok(()) + let result = (|| { + let buffer_len = buffer.length(); + let element_size = size_of::(); + if element_size == 0 { + return Err(MetalSupportError::BufferZeroSizedType { abi_name: T::NAME }); + } + let byte_len = + output + .len() + .checked_mul(element_size) + .ok_or(MetalSupportError::BufferBounds { + offset_bytes: byte_offset, + byte_len: usize::MAX, + buffer_len, + })?; + let end = byte_offset + .checked_add(byte_len) + .ok_or(MetalSupportError::BufferBounds { + offset_bytes: byte_offset, + byte_len, + buffer_len, + })?; + if end > buffer_len { + return Err(MetalSupportError::BufferBounds { + offset_bytes: byte_offset, + byte_len, + buffer_len, + }); + } + let align = core::mem::align_of::(); + if !byte_offset.is_multiple_of(align) { + return Err(MetalSupportError::BufferAlignment { + offset_bytes: byte_offset, + align, + }); + } + if output.is_empty() { + return Ok(()); + } + // SAFETY: These private readback callers use completed/synchronized, + // immutable Metal storage and a separate, exclusively borrowed host + // destination. Reuse the audited support read for storage and actual + // pointer alignment, then the existing checked byte-borrow boundary. + let bytes = unsafe { + support_checked_buffer_read::(buffer, byte_offset)?; + completed_metal_buffer_bytes(buffer, byte_offset, byte_len)? + }; + T::slice_as_bytes_mut(output).copy_from_slice(bytes); + Ok(()) + })(); + result.map_err(|error| buffer_access_error(context, error)) } #[cfg(target_os = "macos")] diff --git a/crates/j2k-metal/src/engine/test_counters.rs b/crates/j2k-metal/src/engine/test_counters.rs index 2670c83a..766567e8 100644 --- a/crates/j2k-metal/src/engine/test_counters.rs +++ b/crates/j2k-metal/src/engine/test_counters.rs @@ -295,14 +295,6 @@ pub(crate) fn idwt_host_transfer_counters_for_test() -> (usize, usize, usize) { ) } -pub(crate) fn record_idwt_host_overwritten_output_upload(bytes: usize) { - IDWT_HOST_OVERWRITTEN_OUTPUT_UPLOAD_BYTES.set( - IDWT_HOST_OVERWRITTEN_OUTPUT_UPLOAD_BYTES - .get() - .saturating_add(bytes), - ); -} - pub(crate) fn record_idwt_host_temporary_readback_vec(bytes: usize) { IDWT_HOST_TEMPORARY_READBACK_VEC_ALLOCATIONS.set( IDWT_HOST_TEMPORARY_READBACK_VEC_ALLOCATIONS diff --git a/crates/j2k-metal/src/idwt_host_slice_tests.rs b/crates/j2k-metal/src/idwt_host_slice_tests.rs index 5daac64c..40558365 100644 --- a/crates/j2k-metal/src/idwt_host_slice_tests.rs +++ b/crates/j2k-metal/src/idwt_host_slice_tests.rs @@ -142,14 +142,7 @@ fn host_slice_idwt_output_transfer_accounting() { assert!(decoder .decode_single_decomposition_idwt(fixture.job(transform), &mut output) .unwrap()); - assert_eq!( - idwt_host_transfer_counters_for_test(), - ( - output.len() * size_of::(), - 1, - output.len() * size_of::() - ) - ); + assert_eq!(idwt_host_transfer_counters_for_test(), (0, 0, 0)); } }