Skip to main content

j2k_transcode_metal/
lib.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! Metal acceleration for coefficient-domain JPEG to HTJ2K transcode stages.
4//!
5//! The supported targets are direct DCT-grid to one-level 5/3 and 9/7 wavelet
6//! projections used by `j2k-transcode`'s HTJ2K paths. CPU scalar code
7//! remains the oracle and fallback.
8
9#[cfg(target_os = "macos")]
10mod metal;
11
12#[doc(hidden)]
13pub mod weights;
14
15#[cfg(target_os = "macos")]
16pub use metal::MetalTranscodeSession;
17
18use core::fmt;
19
20use j2k_core::{BackendKind, BackendRequest};
21use j2k_transcode::accelerator::{
22    DctGridToDwt53Job, DctGridToDwt97Job, DctGridToHtj2k97CodeBlockJob,
23    DctGridToReversibleDwt53Job, DctToWaveletStageAccelerator, Dwt97BatchStageTimings,
24    Htj2k97CodeBlockOptions, PrequantizedHtj2k97Component, ReversibleDwt53FirstLevel,
25    TranscodeStageError,
26};
27use j2k_transcode::dct53_2d::Dwt53TwoDimensional;
28use j2k_transcode::dct97_2d::Dwt97TwoDimensional;
29use j2k_transcode::{
30    BatchTranscodeReport, EncodedTranscode, EncodedTranscodeBatch, JpegTileBatchInput,
31    JpegToHtj2kError, JpegToHtj2kOptions, JpegToHtj2kTranscoder, TranscodePipelineMap,
32    TranscodeTimingReport,
33};
34#[cfg(target_os = "macos")]
35use j2k_transcode::{ResidentBufferRef, ResidentCodestreamBuffer, ResidentHandoffError};
36
37/// Stable message returned when Metal is unavailable.
38pub const METAL_UNAVAILABLE: &str = "Metal is unavailable on this host";
39
40const DEFAULT_AUTO_MIN_SAMPLES: usize = 224 * 224;
41const DEFAULT_AUTO_DWT97_MIN_SAMPLES: usize = usize::MAX;
42const DEFAULT_AUTO_REVERSIBLE_MIN_SAMPLES: usize = usize::MAX;
43const DEFAULT_AUTO_REVERSIBLE_BATCH_MIN_JOBS: usize = 32;
44const DEFAULT_AUTO_REVERSIBLE_BATCH_MIN_SAMPLES: usize = 224 * 224 * 32;
45const DEFAULT_AUTO_DWT97_BATCH_MIN_JOBS: usize = 32;
46const DEFAULT_AUTO_DWT97_BATCH_MIN_SAMPLES: usize = 224 * 224 * 32;
47const MAX_AUTO_DWT97_STAGED_BATCH_AXIS: usize = 1024;
48
49/// Error returned by the Metal transcode accelerator.
50#[derive(Debug, Clone, Copy, PartialEq, Eq)]
51pub enum MetalTranscodeError {
52    /// Metal is unavailable on this host or target.
53    MetalUnavailable,
54    /// The request is outside the current Metal implementation.
55    UnsupportedJob(&'static str),
56    /// Metal runtime creation or device setup failed.
57    Runtime(&'static str),
58    /// Metal runtime or kernel execution failed.
59    Kernel(&'static str),
60}
61
62impl fmt::Display for MetalTranscodeError {
63    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
64        match self {
65            Self::MetalUnavailable => f.write_str(METAL_UNAVAILABLE),
66            Self::UnsupportedJob(reason) | Self::Runtime(reason) | Self::Kernel(reason) => {
67                f.write_str(reason)
68            }
69        }
70    }
71}
72
73impl From<MetalTranscodeError> for TranscodeStageError {
74    fn from(error: MetalTranscodeError) -> Self {
75        match error {
76            MetalTranscodeError::MetalUnavailable => Self::DeviceUnavailable,
77            MetalTranscodeError::UnsupportedJob(reason) => Self::Unsupported(reason),
78            MetalTranscodeError::Runtime(reason) | MetalTranscodeError::Kernel(reason) => {
79                Self::Backend(reason.to_string())
80            }
81        }
82    }
83}
84
85impl std::error::Error for MetalTranscodeError {}
86
87const CUDA_REQUESTED_THROUGH_METAL_ADAPTER: &str = "CUDA transcode requested through Metal adapter";
88const STRICT_METAL_TRANSCODE_NO_DISPATCH: &str =
89    "strict Metal transcode produced no Metal dispatch";
90
91/// Structured CPU fallback reason for the Metal JPEG-to-HTJ2K route facade.
92#[derive(Debug, Clone, Copy, PartialEq, Eq)]
93pub enum MetalTranscodeFallbackReason {
94    /// The caller requested CPU explicitly.
95    CpuRequested,
96    /// Auto found no transform stage eligible for Metal.
97    AutoNoEligibleTransformStage,
98    /// Auto offered transform jobs to Metal, but all jobs used CPU fallback.
99    AutoAllTransformJobsFellBackToCpu,
100    /// Auto used Metal for some transform jobs and CPU fallback for others.
101    AutoPartialTransformFallback,
102}
103
104impl MetalTranscodeFallbackReason {
105    /// Stable reason label for logs, examples, and benchmark output.
106    #[must_use]
107    pub const fn as_str(self) -> &'static str {
108        match self {
109            Self::CpuRequested => "cpu_requested",
110            Self::AutoNoEligibleTransformStage => "auto_no_eligible_transform_stage",
111            Self::AutoAllTransformJobsFellBackToCpu => "auto_all_transform_jobs_fell_back_to_cpu",
112            Self::AutoPartialTransformFallback => "auto_partial_transform_fallback",
113        }
114    }
115}
116
117impl fmt::Display for MetalTranscodeFallbackReason {
118    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
119        f.write_str(self.as_str())
120    }
121}
122
123/// Public route report for Metal-adapted JPEG-to-HTJ2K transcode calls.
124#[derive(Debug, Clone, PartialEq, Eq)]
125pub struct MetalTranscodeRouteReport {
126    /// Caller backend request.
127    pub request: BackendRequest,
128    /// Backend that handled transform stages.
129    pub selected_transform_backend: BackendKind,
130    /// Backend that produced the public codestream byte vector.
131    pub output_backend: BackendKind,
132    /// Structured CPU fallback reason when the route did not fully use Metal.
133    pub fallback_reason: Option<MetalTranscodeFallbackReason>,
134    /// Stage residency map derived from the existing transcode timing counters.
135    pub pipeline_map: TranscodePipelineMap,
136}
137
138/// JPEG-to-HTJ2K transcode output plus Metal route report.
139pub struct MetalEncodedTranscode {
140    /// Encoded HTJ2K codestream and native transcode report.
141    pub encoded: EncodedTranscode,
142    /// Route and residency report for this call.
143    pub route: MetalTranscodeRouteReport,
144}
145
146/// Batch JPEG-to-HTJ2K transcode output plus Metal route report.
147pub struct MetalEncodedTranscodeBatch {
148    /// Per-tile outputs and aggregate native transcode report.
149    pub batch: EncodedTranscodeBatch,
150    /// Route and residency report for this batch call.
151    pub route: MetalTranscodeRouteReport,
152}
153
154/// Build a backend-neutral resident codestream descriptor from a Metal encode output.
155#[cfg(target_os = "macos")]
156pub fn resident_codestream_buffer_from_metal_encoded_j2k(
157    encoded: &j2k_metal::MetalEncodedJ2k,
158) -> Result<ResidentCodestreamBuffer<'_>, ResidentHandoffError> {
159    let memory = encoded
160        .codestream_memory_range()
161        .ok_or(ResidentHandoffError::OffsetOverflow)?;
162    let allocation_len = encoded
163        .codestream_allocation_len()
164        .ok_or(ResidentHandoffError::RangeExceedsAllocation)?;
165    let buffer = ResidentBufferRef::with_allocation_len(memory, allocation_len)?;
166    ResidentCodestreamBuffer::new(buffer, encoded.byte_len, encoded.capacity)?
167        .require_backend(BackendKind::Metal)
168}
169
170/// Transcode JPEG to HTJ2K using CPU, Auto Metal, or strict Metal routing.
171///
172/// `BackendRequest::Metal` uses the explicit Metal accelerator and returns an
173/// error if Metal is unavailable or the required transform stage is unsupported.
174/// `BackendRequest::Auto` may return CPU output with a structured fallback
175/// reason.
176pub fn jpeg_to_htj2k_with_metal_route(
177    bytes: &[u8],
178    options: &JpegToHtj2kOptions,
179    request: BackendRequest,
180) -> Result<MetalEncodedTranscode, JpegToHtj2kError> {
181    let mut transcoder = JpegToHtj2kTranscoder::default();
182    let encoded = match request {
183        BackendRequest::Cpu => transcoder.transcode(bytes, options)?,
184        BackendRequest::Auto => {
185            let mut accelerator = MetalDctToWaveletStageAccelerator::for_auto();
186            transcoder.transcode_with_accelerator(bytes, options, &mut accelerator)?
187        }
188        BackendRequest::Metal => {
189            let mut accelerator = MetalDctToWaveletStageAccelerator::new_explicit();
190            let encoded =
191                transcoder.transcode_with_accelerator(bytes, options, &mut accelerator)?;
192            ensure_strict_metal_dispatched(&encoded.report.timings)?;
193            encoded
194        }
195        BackendRequest::Cuda => {
196            return Err(JpegToHtj2kError::Unsupported(
197                CUDA_REQUESTED_THROUGH_METAL_ADAPTER,
198            ));
199        }
200    };
201    let route = route_report(request, &encoded.report.timings);
202    Ok(MetalEncodedTranscode { encoded, route })
203}
204
205/// Batch transcode JPEG tiles to HTJ2K using CPU, Auto Metal, or strict Metal routing.
206pub fn jpeg_to_htj2k_batch_with_metal_route(
207    tiles: &[JpegTileBatchInput<'_>],
208    options: &JpegToHtj2kOptions,
209    request: BackendRequest,
210) -> Result<MetalEncodedTranscodeBatch, JpegToHtj2kError> {
211    let mut transcoder = JpegToHtj2kTranscoder::default();
212    let batch = match request {
213        BackendRequest::Cpu => transcoder.transcode_batch(tiles, options)?,
214        BackendRequest::Auto => {
215            let mut accelerator = MetalDctToWaveletStageAccelerator::for_auto();
216            transcoder.transcode_batch_with_accelerator(tiles, options, &mut accelerator)?
217        }
218        BackendRequest::Metal => {
219            let mut accelerator = MetalDctToWaveletStageAccelerator::new_explicit();
220            let batch =
221                transcoder.transcode_batch_with_accelerator(tiles, options, &mut accelerator)?;
222            ensure_strict_metal_batch_dispatched(&batch.report)?;
223            batch
224        }
225        BackendRequest::Cuda => {
226            return Err(JpegToHtj2kError::Unsupported(
227                CUDA_REQUESTED_THROUGH_METAL_ADAPTER,
228            ));
229        }
230    };
231    let route = route_report(request, &batch.report.timings);
232    Ok(MetalEncodedTranscodeBatch { batch, route })
233}
234
235fn route_report(
236    request: BackendRequest,
237    timings: &TranscodeTimingReport,
238) -> MetalTranscodeRouteReport {
239    let selected_transform_backend = if timings_use_metal(timings) {
240        BackendKind::Metal
241    } else {
242        BackendKind::Cpu
243    };
244    MetalTranscodeRouteReport {
245        request,
246        selected_transform_backend,
247        output_backend: BackendKind::Cpu,
248        fallback_reason: fallback_reason(request, selected_transform_backend, timings),
249        pipeline_map: timings.pipeline_map(),
250    }
251}
252
253fn fallback_reason(
254    request: BackendRequest,
255    selected_transform_backend: BackendKind,
256    timings: &TranscodeTimingReport,
257) -> Option<MetalTranscodeFallbackReason> {
258    match request {
259        BackendRequest::Cpu => Some(MetalTranscodeFallbackReason::CpuRequested),
260        BackendRequest::Auto
261            if selected_transform_backend == BackendKind::Cpu
262                && timings.accelerator_attempts == 0 =>
263        {
264            Some(MetalTranscodeFallbackReason::AutoNoEligibleTransformStage)
265        }
266        BackendRequest::Auto if selected_transform_backend == BackendKind::Cpu => {
267            Some(MetalTranscodeFallbackReason::AutoAllTransformJobsFellBackToCpu)
268        }
269        BackendRequest::Auto if timings.cpu_fallback_jobs > 0 => {
270            Some(MetalTranscodeFallbackReason::AutoPartialTransformFallback)
271        }
272        BackendRequest::Metal | BackendRequest::Cuda | BackendRequest::Auto => None,
273    }
274}
275
276fn ensure_strict_metal_dispatched(timings: &TranscodeTimingReport) -> Result<(), JpegToHtj2kError> {
277    if timings_use_metal(timings) {
278        Ok(())
279    } else {
280        Err(JpegToHtj2kError::Accelerator(
281            TranscodeStageError::Unsupported(STRICT_METAL_TRANSCODE_NO_DISPATCH),
282        ))
283    }
284}
285
286fn ensure_strict_metal_batch_dispatched(
287    report: &BatchTranscodeReport,
288) -> Result<(), JpegToHtj2kError> {
289    if report.successful_tiles == 0 || timings_use_metal(&report.timings) {
290        Ok(())
291    } else {
292        Err(JpegToHtj2kError::Accelerator(
293            TranscodeStageError::Unsupported(STRICT_METAL_TRANSCODE_NO_DISPATCH),
294        ))
295    }
296}
297
298fn timings_use_metal(timings: &TranscodeTimingReport) -> bool {
299    timings.accelerator_dispatches > 0
300        || timings.dwt97_batch_idct_row_lift_us > 0
301        || timings.dwt97_batch_column_lift_us > 0
302        || timings.dwt97_batch_quantize_codeblock_us > 0
303        || timings.dwt97_batch_ht_encode_us > 0
304        || timings.dwt97_batch_ht_kernel_us > 0
305        || timings.dwt97_batch_ht_compact_us > 0
306        || timings.dwt97_batch_ht_codeblock_dispatches > 0
307}
308
309#[cfg(not(target_os = "macos"))]
310#[derive(Clone, Copy, Debug, Default)]
311/// Placeholder Metal transcode session for hosts without Metal support.
312pub struct MetalTranscodeSession {
313    _private: (),
314}
315
316#[cfg(not(target_os = "macos"))]
317impl MetalTranscodeSession {
318    /// Return `MetalUnavailable` on hosts without Metal support.
319    pub const fn system_default() -> Result<Self, MetalTranscodeError> {
320        Err(MetalTranscodeError::MetalUnavailable)
321    }
322}
323
324/// Optional Metal accelerator for `j2k-transcode` transform stages.
325#[derive(Debug, Clone)]
326pub struct MetalDctToWaveletStageAccelerator {
327    mode: MetalDispatchMode,
328    min_auto_samples: usize,
329    min_auto_dwt97_samples: usize,
330    min_auto_reversible_samples: usize,
331    min_auto_reversible_batch_jobs: usize,
332    min_auto_reversible_batch_samples: usize,
333    reversible_dwt53_attempts: usize,
334    reversible_dwt53_dispatches: usize,
335    reversible_dwt53_batch_attempts: usize,
336    reversible_dwt53_batch_dispatches: usize,
337    dwt53_attempts: usize,
338    dwt53_dispatches: usize,
339    dwt97_attempts: usize,
340    dwt97_dispatches: usize,
341    dwt97_batch_attempts: usize,
342    dwt97_batch_dispatches: usize,
343    htj2k97_codeblock_batch_attempts: usize,
344    htj2k97_codeblock_batch_dispatches: usize,
345    last_dwt97_batch_stage_timings: Option<Dwt97BatchStageTimings>,
346    min_auto_dwt97_batch_jobs: usize,
347    min_auto_dwt97_batch_samples: usize,
348    #[cfg(target_os = "macos")]
349    session: Option<MetalTranscodeSession>,
350}
351
352#[derive(Debug, Clone, Copy, PartialEq, Eq)]
353enum MetalDispatchMode {
354    Explicit,
355    Auto,
356}
357
358impl MetalDctToWaveletStageAccelerator {
359    /// Create an accelerator that treats unsupported Metal dispatch as an
360    /// error.
361    #[must_use]
362    pub const fn new_explicit() -> Self {
363        Self {
364            mode: MetalDispatchMode::Explicit,
365            min_auto_samples: 0,
366            min_auto_dwt97_samples: 0,
367            min_auto_reversible_samples: 0,
368            min_auto_reversible_batch_jobs: 0,
369            min_auto_reversible_batch_samples: 0,
370            reversible_dwt53_attempts: 0,
371            reversible_dwt53_dispatches: 0,
372            reversible_dwt53_batch_attempts: 0,
373            reversible_dwt53_batch_dispatches: 0,
374            dwt53_attempts: 0,
375            dwt53_dispatches: 0,
376            dwt97_attempts: 0,
377            dwt97_dispatches: 0,
378            dwt97_batch_attempts: 0,
379            dwt97_batch_dispatches: 0,
380            htj2k97_codeblock_batch_attempts: 0,
381            htj2k97_codeblock_batch_dispatches: 0,
382            last_dwt97_batch_stage_timings: None,
383            min_auto_dwt97_batch_jobs: 0,
384            min_auto_dwt97_batch_samples: 0,
385            #[cfg(target_os = "macos")]
386            session: None,
387        }
388    }
389
390    /// Create an accelerator that falls back to scalar CPU for small or
391    /// unsupported jobs.
392    #[must_use]
393    pub const fn for_auto() -> Self {
394        Self {
395            mode: MetalDispatchMode::Auto,
396            min_auto_samples: DEFAULT_AUTO_MIN_SAMPLES,
397            min_auto_dwt97_samples: DEFAULT_AUTO_DWT97_MIN_SAMPLES,
398            min_auto_reversible_samples: DEFAULT_AUTO_REVERSIBLE_MIN_SAMPLES,
399            min_auto_reversible_batch_jobs: DEFAULT_AUTO_REVERSIBLE_BATCH_MIN_JOBS,
400            min_auto_reversible_batch_samples: DEFAULT_AUTO_REVERSIBLE_BATCH_MIN_SAMPLES,
401            reversible_dwt53_attempts: 0,
402            reversible_dwt53_dispatches: 0,
403            reversible_dwt53_batch_attempts: 0,
404            reversible_dwt53_batch_dispatches: 0,
405            dwt53_attempts: 0,
406            dwt53_dispatches: 0,
407            dwt97_attempts: 0,
408            dwt97_dispatches: 0,
409            dwt97_batch_attempts: 0,
410            dwt97_batch_dispatches: 0,
411            htj2k97_codeblock_batch_attempts: 0,
412            htj2k97_codeblock_batch_dispatches: 0,
413            last_dwt97_batch_stage_timings: None,
414            min_auto_dwt97_batch_jobs: DEFAULT_AUTO_DWT97_BATCH_MIN_JOBS,
415            min_auto_dwt97_batch_samples: DEFAULT_AUTO_DWT97_BATCH_MIN_SAMPLES,
416            #[cfg(target_os = "macos")]
417            session: None,
418        }
419    }
420
421    /// Create an explicit-dispatch accelerator bound to a caller-owned Metal session.
422    #[cfg(target_os = "macos")]
423    #[must_use]
424    pub fn new_explicit_with_session(session: MetalTranscodeSession) -> Self {
425        Self::new_explicit().with_session(session)
426    }
427
428    /// Create an Auto-mode accelerator bound to a caller-owned Metal session.
429    #[cfg(target_os = "macos")]
430    #[must_use]
431    pub fn for_auto_with_session(session: MetalTranscodeSession) -> Self {
432        Self::for_auto().with_session(session)
433    }
434
435    /// Create an explicit-dispatch accelerator bound to an existing Metal device.
436    #[cfg(target_os = "macos")]
437    #[must_use]
438    pub fn new_explicit_with_device(device: ::metal::Device) -> Self {
439        Self::new_explicit_with_session(MetalTranscodeSession::new(device))
440    }
441
442    /// Create an Auto-mode accelerator bound to an existing Metal device.
443    #[cfg(target_os = "macos")]
444    #[must_use]
445    pub fn for_auto_with_device(device: ::metal::Device) -> Self {
446        Self::for_auto_with_session(MetalTranscodeSession::new(device))
447    }
448
449    /// Bind this accelerator to a caller-owned Metal session.
450    #[cfg(target_os = "macos")]
451    #[must_use]
452    pub fn with_session(mut self, session: MetalTranscodeSession) -> Self {
453        self.session = Some(session);
454        self
455    }
456
457    /// Bind this accelerator to an existing Metal device.
458    #[cfg(target_os = "macos")]
459    #[must_use]
460    pub fn with_device(self, device: ::metal::Device) -> Self {
461        self.with_session(MetalTranscodeSession::new(device))
462    }
463
464    #[cfg(target_os = "macos")]
465    fn metal_session(&mut self) -> &mut MetalTranscodeSession {
466        self.session
467            .get_or_insert_with(MetalTranscodeSession::default)
468    }
469
470    /// Override the minimum component sample count used before Auto mode
471    /// dispatches non-reversible projection jobs to Metal.
472    #[must_use]
473    pub const fn with_auto_min_samples(mut self, min_samples: usize) -> Self {
474        self.min_auto_samples = min_samples;
475        self.min_auto_dwt97_samples = min_samples;
476        self
477    }
478
479    /// Override the minimum component sample count used before Auto mode
480    /// dispatches 9/7 transform jobs to Metal.
481    #[must_use]
482    pub const fn with_auto_dwt97_min_samples(mut self, min_samples: usize) -> Self {
483        self.min_auto_dwt97_samples = min_samples;
484        self
485    }
486
487    /// Override the 9/7 batch thresholds used before Auto mode dispatches a
488    /// same-geometry batch to Metal.
489    #[must_use]
490    pub const fn with_auto_dwt97_batch_thresholds(
491        mut self,
492        min_jobs: usize,
493        min_samples: usize,
494    ) -> Self {
495        self.min_auto_dwt97_batch_jobs = min_jobs;
496        self.min_auto_dwt97_batch_samples = min_samples;
497        self
498    }
499
500    /// Override the minimum component sample count used before Auto mode
501    /// dispatches single reversible 5/3 jobs to Metal.
502    #[must_use]
503    pub const fn with_auto_reversible_min_samples(mut self, min_samples: usize) -> Self {
504        self.min_auto_reversible_samples = min_samples;
505        self
506    }
507
508    /// Override the reversible 5/3 batch thresholds used before Auto mode
509    /// dispatches a same-geometry batch to Metal.
510    #[must_use]
511    pub const fn with_auto_reversible_batch_thresholds(
512        mut self,
513        min_jobs: usize,
514        min_samples: usize,
515    ) -> Self {
516        self.min_auto_reversible_batch_jobs = min_jobs;
517        self.min_auto_reversible_batch_samples = min_samples;
518        self
519    }
520
521    /// Number of reversible integer 5/3 jobs offered to this accelerator.
522    #[must_use]
523    pub const fn reversible_dwt53_attempts(&self) -> usize {
524        self.reversible_dwt53_attempts
525    }
526
527    /// Number of reversible integer 5/3 jobs handled by Metal.
528    #[must_use]
529    pub const fn reversible_dwt53_dispatches(&self) -> usize {
530        self.reversible_dwt53_dispatches
531    }
532
533    /// Number of reversible integer 5/3 batches offered to this accelerator.
534    #[must_use]
535    pub const fn reversible_dwt53_batch_attempts(&self) -> usize {
536        self.reversible_dwt53_batch_attempts
537    }
538
539    /// Number of reversible integer 5/3 batches handled by Metal.
540    #[must_use]
541    pub const fn reversible_dwt53_batch_dispatches(&self) -> usize {
542        self.reversible_dwt53_batch_dispatches
543    }
544
545    /// Number of 5/3 projection jobs offered to this accelerator.
546    #[must_use]
547    pub const fn dwt53_attempts(&self) -> usize {
548        self.dwt53_attempts
549    }
550
551    /// Number of 5/3 projection jobs handled by Metal.
552    #[must_use]
553    pub const fn dwt53_dispatches(&self) -> usize {
554        self.dwt53_dispatches
555    }
556
557    /// Number of 9/7 transform jobs offered to this accelerator.
558    #[must_use]
559    pub const fn dwt97_attempts(&self) -> usize {
560        self.dwt97_attempts
561    }
562
563    /// Number of 9/7 transform jobs handled by Metal.
564    #[must_use]
565    pub const fn dwt97_dispatches(&self) -> usize {
566        self.dwt97_dispatches
567    }
568
569    /// Number of 9/7 transform batches offered to this accelerator.
570    #[must_use]
571    pub const fn dwt97_batch_attempts(&self) -> usize {
572        self.dwt97_batch_attempts
573    }
574
575    /// Number of 9/7 transform batches handled by Metal.
576    #[must_use]
577    pub const fn dwt97_batch_dispatches(&self) -> usize {
578        self.dwt97_batch_dispatches
579    }
580
581    /// Number of 9/7 code-block-ready batches offered to this accelerator.
582    #[must_use]
583    pub const fn htj2k97_codeblock_batch_attempts(&self) -> usize {
584        self.htj2k97_codeblock_batch_attempts
585    }
586
587    /// Number of 9/7 code-block-ready batches handled by Metal.
588    #[must_use]
589    pub const fn htj2k97_codeblock_batch_dispatches(&self) -> usize {
590        self.htj2k97_codeblock_batch_dispatches
591    }
592
593    /// Backend stage timings for the most recent 9/7 batch dispatch.
594    #[must_use]
595    pub const fn last_dwt97_batch_stage_timings(&self) -> Option<Dwt97BatchStageTimings> {
596        self.last_dwt97_batch_stage_timings
597    }
598
599    /// Dispatch a same-geometry batch of reversible integer 5/3 DCT-grid
600    /// projection jobs. This is an experimental Metal-specific extension used
601    /// for WSI tile-component queues; scalar/Rayon remains the portable oracle.
602    pub fn dct_grid_to_reversible_dwt53_batch(
603        &mut self,
604        jobs: &[DctGridToReversibleDwt53Job<'_>],
605    ) -> Result<Option<Vec<ReversibleDwt53FirstLevel>>, TranscodeStageError> {
606        self.dispatch_reversible_dwt53_batch(jobs)
607    }
608
609    fn dispatch_reversible_dwt53_batch(
610        &mut self,
611        jobs: &[DctGridToReversibleDwt53Job<'_>],
612    ) -> Result<Option<Vec<ReversibleDwt53FirstLevel>>, TranscodeStageError> {
613        self.reversible_dwt53_batch_attempts =
614            self.reversible_dwt53_batch_attempts.saturating_add(1);
615
616        if jobs.is_empty() {
617            return Ok(Some(Vec::new()));
618        }
619
620        let total_samples = jobs.iter().fold(0usize, |total, job| {
621            total.saturating_add(job.width.saturating_mul(job.height))
622        });
623        // Auto declines with `Ok(None)` (small batches, unavailable Metal,
624        // unsupported jobs) so the caller runs its scalar fallback — the same
625        // contract as the float 5/3 path and the CUDA accelerator, instead of
626        // silently computing a Rayon fallback inside the backend.
627        if self.mode == MetalDispatchMode::Auto
628            && (jobs.len() < self.min_auto_reversible_batch_jobs
629                || total_samples < self.min_auto_reversible_batch_samples)
630        {
631            return Ok(None);
632        }
633
634        #[cfg(not(target_os = "macos"))]
635        {
636            match self.mode {
637                MetalDispatchMode::Explicit => Err(TranscodeStageError::DeviceUnavailable),
638                MetalDispatchMode::Auto => Ok(None),
639            }
640        }
641
642        #[cfg(target_os = "macos")]
643        {
644            match metal::dispatch_dct_grid_to_reversible_dwt53_batch(self.metal_session(), jobs) {
645                Ok(output) => {
646                    self.reversible_dwt53_batch_dispatches =
647                        self.reversible_dwt53_batch_dispatches.saturating_add(1);
648                    Ok(Some(output))
649                }
650                Err(
651                    MetalTranscodeError::MetalUnavailable | MetalTranscodeError::UnsupportedJob(_),
652                ) if self.mode == MetalDispatchMode::Auto => Ok(None),
653                Err(error) => Err(error.into()),
654            }
655        }
656    }
657}
658
659impl Default for MetalDctToWaveletStageAccelerator {
660    fn default() -> Self {
661        Self::for_auto()
662    }
663}
664
665impl DctToWaveletStageAccelerator for MetalDctToWaveletStageAccelerator {
666    fn supports_dwt97_batch(&self) -> bool {
667        true
668    }
669
670    fn supports_htj2k97_codeblock_batch(&self) -> bool {
671        true
672    }
673
674    fn dct_grid_to_reversible_dwt53(
675        &mut self,
676        job: DctGridToReversibleDwt53Job<'_>,
677    ) -> Result<Option<ReversibleDwt53FirstLevel>, TranscodeStageError> {
678        self.reversible_dwt53_attempts = self.reversible_dwt53_attempts.saturating_add(1);
679
680        // Auto declines with `Ok(None)`; see dispatch_reversible_dwt53_batch.
681        if self.mode == MetalDispatchMode::Auto
682            && job.width.saturating_mul(job.height) < self.min_auto_reversible_samples
683        {
684            return Ok(None);
685        }
686
687        #[cfg(not(target_os = "macos"))]
688        {
689            match self.mode {
690                MetalDispatchMode::Explicit => Err(TranscodeStageError::DeviceUnavailable),
691                MetalDispatchMode::Auto => Ok(None),
692            }
693        }
694
695        #[cfg(target_os = "macos")]
696        {
697            match metal::dispatch_dct_grid_to_reversible_dwt53(self.metal_session(), job) {
698                Ok(output) => {
699                    self.reversible_dwt53_dispatches =
700                        self.reversible_dwt53_dispatches.saturating_add(1);
701                    Ok(Some(output))
702                }
703                Err(
704                    MetalTranscodeError::MetalUnavailable | MetalTranscodeError::UnsupportedJob(_),
705                ) if self.mode == MetalDispatchMode::Auto => Ok(None),
706                Err(error) => Err(error.into()),
707            }
708        }
709    }
710
711    fn dct_grid_to_reversible_dwt53_batch(
712        &mut self,
713        jobs: &[DctGridToReversibleDwt53Job<'_>],
714    ) -> Result<Option<Vec<ReversibleDwt53FirstLevel>>, TranscodeStageError> {
715        self.dispatch_reversible_dwt53_batch(jobs)
716    }
717
718    fn dct_grid_to_dwt53(
719        &mut self,
720        job: DctGridToDwt53Job<'_>,
721    ) -> Result<Option<Dwt53TwoDimensional<f64>>, TranscodeStageError> {
722        self.dwt53_attempts = self.dwt53_attempts.saturating_add(1);
723
724        if self.mode == MetalDispatchMode::Auto
725            && job.width.saturating_mul(job.height) < self.min_auto_samples
726        {
727            return Ok(None);
728        }
729
730        #[cfg(not(target_os = "macos"))]
731        {
732            let _ = job;
733            match self.mode {
734                MetalDispatchMode::Explicit => Err(TranscodeStageError::DeviceUnavailable),
735                MetalDispatchMode::Auto => Ok(None),
736            }
737        }
738
739        #[cfg(target_os = "macos")]
740        {
741            match metal::dispatch_dct_grid_to_dwt53(self.metal_session(), job) {
742                Ok(output) => {
743                    self.dwt53_dispatches = self.dwt53_dispatches.saturating_add(1);
744                    Ok(Some(output))
745                }
746                Err(
747                    MetalTranscodeError::MetalUnavailable | MetalTranscodeError::UnsupportedJob(_),
748                ) if self.mode == MetalDispatchMode::Auto => Ok(None),
749                Err(error) => Err(error.into()),
750            }
751        }
752    }
753
754    fn dct_grid_to_dwt97(
755        &mut self,
756        job: DctGridToDwt97Job<'_>,
757    ) -> Result<Option<Dwt97TwoDimensional<f64>>, TranscodeStageError> {
758        self.dwt97_attempts = self.dwt97_attempts.saturating_add(1);
759
760        if self.mode == MetalDispatchMode::Auto
761            && job.width.saturating_mul(job.height) < self.min_auto_dwt97_samples
762        {
763            return Ok(None);
764        }
765
766        #[cfg(not(target_os = "macos"))]
767        {
768            let _ = job;
769            match self.mode {
770                MetalDispatchMode::Explicit => Err(TranscodeStageError::DeviceUnavailable),
771                MetalDispatchMode::Auto => Ok(None),
772            }
773        }
774
775        #[cfg(target_os = "macos")]
776        {
777            match metal::dispatch_dct_grid_to_dwt97(self.metal_session(), job) {
778                Ok(output) => {
779                    self.dwt97_dispatches = self.dwt97_dispatches.saturating_add(1);
780                    Ok(Some(output))
781                }
782                Err(
783                    MetalTranscodeError::MetalUnavailable | MetalTranscodeError::UnsupportedJob(_),
784                ) if self.mode == MetalDispatchMode::Auto => Ok(None),
785                Err(error) => Err(error.into()),
786            }
787        }
788    }
789
790    fn dct_grid_to_dwt97_batch(
791        &mut self,
792        jobs: &[DctGridToDwt97Job<'_>],
793    ) -> Result<Option<Vec<Dwt97TwoDimensional<f64>>>, TranscodeStageError> {
794        self.dwt97_batch_attempts = self.dwt97_batch_attempts.saturating_add(1);
795        self.last_dwt97_batch_stage_timings = None;
796
797        if jobs.is_empty() {
798            return Ok(Some(Vec::new()));
799        }
800
801        let total_samples = jobs.iter().fold(0usize, |total, job| {
802            total.saturating_add(job.width.saturating_mul(job.height))
803        });
804        if self.mode == MetalDispatchMode::Auto
805            && (jobs.len() < self.min_auto_dwt97_batch_jobs
806                || total_samples < self.min_auto_dwt97_batch_samples)
807        {
808            return Ok(None);
809        }
810        if self.mode == MetalDispatchMode::Auto
811            && jobs.iter().any(|job| {
812                job.width > MAX_AUTO_DWT97_STAGED_BATCH_AXIS
813                    || job.height > MAX_AUTO_DWT97_STAGED_BATCH_AXIS
814            })
815        {
816            return Ok(None);
817        }
818
819        #[cfg(not(target_os = "macos"))]
820        {
821            let _ = jobs;
822            match self.mode {
823                MetalDispatchMode::Explicit => Err(TranscodeStageError::DeviceUnavailable),
824                MetalDispatchMode::Auto => Ok(None),
825            }
826        }
827
828        #[cfg(target_os = "macos")]
829        {
830            match metal::dispatch_dct_grid_to_dwt97_batch(self.metal_session(), jobs) {
831                Ok((output, timings)) => {
832                    self.dwt97_batch_dispatches = self.dwt97_batch_dispatches.saturating_add(1);
833                    self.last_dwt97_batch_stage_timings = Some(timings);
834                    Ok(Some(output))
835                }
836                Err(
837                    MetalTranscodeError::MetalUnavailable | MetalTranscodeError::UnsupportedJob(_),
838                ) if self.mode == MetalDispatchMode::Auto => Ok(None),
839                Err(error) => Err(error.into()),
840            }
841        }
842    }
843
844    fn dct_grid_to_htj2k97_codeblock_batch(
845        &mut self,
846        jobs: &[DctGridToHtj2k97CodeBlockJob<'_>],
847        options: Htj2k97CodeBlockOptions,
848    ) -> Result<Option<Vec<PrequantizedHtj2k97Component>>, TranscodeStageError> {
849        self.dwt97_batch_attempts = self.dwt97_batch_attempts.saturating_add(1);
850        self.htj2k97_codeblock_batch_attempts =
851            self.htj2k97_codeblock_batch_attempts.saturating_add(1);
852        self.last_dwt97_batch_stage_timings = None;
853
854        if jobs.is_empty() {
855            return Ok(Some(Vec::new()));
856        }
857
858        let total_samples = jobs.iter().fold(0usize, |total, job| {
859            total.saturating_add(job.width.saturating_mul(job.height))
860        });
861        if self.mode == MetalDispatchMode::Auto
862            && (jobs.len() < self.min_auto_dwt97_batch_jobs
863                || total_samples < self.min_auto_dwt97_batch_samples)
864        {
865            return Ok(None);
866        }
867        if self.mode == MetalDispatchMode::Auto
868            && jobs.iter().any(|job| {
869                job.width > MAX_AUTO_DWT97_STAGED_BATCH_AXIS
870                    || job.height > MAX_AUTO_DWT97_STAGED_BATCH_AXIS
871            })
872        {
873            return Ok(None);
874        }
875
876        #[cfg(not(target_os = "macos"))]
877        {
878            let _ = (jobs, options);
879            match self.mode {
880                MetalDispatchMode::Explicit => Err(TranscodeStageError::DeviceUnavailable),
881                MetalDispatchMode::Auto => Ok(None),
882            }
883        }
884
885        #[cfg(target_os = "macos")]
886        {
887            match metal::dispatch_dct_grid_to_htj2k97_codeblock_batch(
888                self.metal_session(),
889                jobs,
890                options,
891            ) {
892                Ok((output, timings)) => {
893                    self.dwt97_batch_dispatches = self.dwt97_batch_dispatches.saturating_add(1);
894                    self.htj2k97_codeblock_batch_dispatches =
895                        self.htj2k97_codeblock_batch_dispatches.saturating_add(1);
896                    self.last_dwt97_batch_stage_timings = Some(timings);
897                    Ok(Some(output))
898                }
899                Err(
900                    MetalTranscodeError::MetalUnavailable | MetalTranscodeError::UnsupportedJob(_),
901                ) if self.mode == MetalDispatchMode::Auto => Ok(None),
902                Err(error) => Err(error.into()),
903            }
904        }
905    }
906
907    fn last_dwt97_batch_stage_timings(&self) -> Option<Dwt97BatchStageTimings> {
908        self.last_dwt97_batch_stage_timings
909    }
910}