j2k_cuda/encode/
stage.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3use j2k::{
4    EncodedHtJ2kCodeBlock, EncodedJ2kCodeBlock, J2kDeinterleaveToF32Job, J2kEncodeDispatchReport,
5    J2kEncodeStageAccelerator, J2kEncodeStageError, J2kForwardDwt53Job, J2kForwardDwt53Output,
6    J2kForwardDwt97Job, J2kForwardDwt97Output, J2kForwardIctJob, J2kForwardRctJob,
7    J2kHtCodeBlockEncodeJob, J2kHtSubbandEncodeJob, J2kHtj2kTileEncodeJob,
8    J2kPacketizationEncodeJob, J2kQuantizeSubbandJob, J2kTier1CodeBlockEncodeJob,
9};
10#[cfg(feature = "cuda-runtime")]
11use j2k_cuda_runtime::{CudaContext, CudaError, CudaHtj2kEncodeResources, CudaJ2kQuantizeJob};
12#[cfg(feature = "cuda-runtime")]
13use std::sync::Arc;
14
15#[cfg(feature = "cuda-runtime")]
16use crate::allocation::HostPhaseBudget;
17use crate::profile;
18
19#[cfg(feature = "cuda-runtime")]
20use super::cuda_component_count_u8;
21#[cfg(feature = "cuda-runtime")]
22use super::htj2k::{
23    cuda_encode_ht_code_block, cuda_encode_ht_code_blocks, cuda_encode_ht_subband,
24    cuda_encode_htj2k_tile_body, cuda_htj2k_encode_tables, encoded_ht_code_blocks_from_cuda,
25};
26#[cfg(feature = "cuda-runtime")]
27use super::packetization::{
28    cuda_packetization_blocks, cuda_packetization_packets, cuda_packetization_subbands,
29    cuda_packetization_tag_nodes, cuda_packetization_tag_states,
30};
31use super::packetization::{
32    flatten_cuda_htj2k_packetization_job_classified, CudaHtj2kPacketizationPlanError,
33};
34#[cfg(feature = "cuda-runtime")]
35use super::stage_error::internal_invariant;
36#[cfg(feature = "cuda-runtime")]
37use super::stage_error::runtime_error;
38use super::stage_error::{adapter_error, arithmetic_overflow, CudaStageResult};
39
40#[cfg(feature = "cuda-runtime")]
41mod dwt_output;
42#[cfg(feature = "cuda-runtime")]
43pub(super) use self::dwt_output::{cuda_dwt53_output_to_j2k, cuda_dwt97_output_to_j2k};
44
45macro_rules! emit_cuda_encode_route {
46    ($(($key:expr, $value:expr)),+ $(,)?) => {{
47        crate::profile::emit_optional_gpu_route_fields(
48            "j2k_cuda_encode_route_fields",
49            || Ok([$(j2k_profile::ProfileField::label($key, $value)?),+]),
50            |fields| j2k_profile::emit_gpu_route_fields("j2k", "cuda", &fields),
51        );
52    }};
53}
54
55pub(super) fn cuda_packetization_plan_fallback_reason(
56    error: CudaHtj2kPacketizationPlanError,
57) -> CudaStageResult<&'static str> {
58    match error {
59        CudaHtj2kPacketizationPlanError::Invalid(reason) => Ok(reason),
60        CudaHtj2kPacketizationPlanError::ArithmeticOverflow(what) => Err(arithmetic_overflow(what)),
61        CudaHtj2kPacketizationPlanError::MemoryCapExceeded {
62            what,
63            requested,
64            cap,
65        } => Err(J2kEncodeStageError::memory_cap_exceeded(
66            what, requested, cap,
67        )),
68        CudaHtj2kPacketizationPlanError::HostAllocation { what, bytes } => {
69            Err(J2kEncodeStageError::host_allocation_failed(what, bytes))
70        }
71        CudaHtj2kPacketizationPlanError::Adapter(source) => Err(adapter_error(
72            "prepare CUDA HTJ2K packetization plan",
73            source,
74        )),
75    }
76}
77
78/// CUDA implementation of selected JPEG 2000 encode stages.
79#[derive(Debug, Default, Clone)]
80#[expect(
81    clippy::struct_excessive_bools,
82    reason = "independent route switches mirror distinct accelerator-stage policies"
83)]
84pub struct CudaEncodeStageAccelerator {
85    #[cfg(feature = "cuda-runtime")]
86    context: Option<CudaContext>,
87    #[cfg(feature = "cuda-runtime")]
88    encode_resources: Option<Arc<CudaHtj2kEncodeResources>>,
89    #[cfg(feature = "cuda-runtime")]
90    device_unavailable_observed: bool,
91    #[cfg_attr(
92        not(feature = "cuda-runtime"),
93        expect(dead_code, reason = "profiling state is used only by the CUDA runtime")
94    )]
95    collect_profile: bool,
96    deinterleave_attempts: usize,
97    forward_rct_attempts: usize,
98    forward_ict_attempts: usize,
99    forward_dwt53_attempts: usize,
100    forward_dwt97_attempts: usize,
101    htj2k_tile_attempts: usize,
102    quantize_subband_attempts: usize,
103    ht_subband_attempts: usize,
104    tier1_code_block_attempts: usize,
105    ht_code_block_attempts: usize,
106    packetization_attempts: usize,
107    prefer_cpu_forward_rct: bool,
108    prefer_cpu_ht_subband: bool,
109    prefer_cpu_quantize_subband: bool,
110    prefer_cpu_packetization: bool,
111    deinterleave_dispatches: usize,
112    forward_rct_dispatches: usize,
113    forward_ict_dispatches: usize,
114    forward_dwt53_dispatches: usize,
115    forward_dwt97_dispatches: usize,
116    #[cfg(feature = "cuda-runtime")]
117    htj2k_tile_dispatches: usize,
118    quantize_subband_dispatches: usize,
119    #[cfg(feature = "cuda-runtime")]
120    ht_subband_dispatches: usize,
121    tier1_code_block_dispatches: usize,
122    ht_code_block_dispatches: usize,
123    packetization_dispatches: usize,
124    deinterleave_us: u128,
125    mct_us: u128,
126    dwt_us: u128,
127    quantize_us: u128,
128    ht_encode_us: u128,
129    packetize_us: u128,
130}
131
132impl CudaEncodeStageAccelerator {
133    pub(super) fn begin_encode_attempt(&mut self) {
134        #[cfg(feature = "cuda-runtime")]
135        {
136            self.device_unavailable_observed = false;
137        }
138        #[cfg(not(feature = "cuda-runtime"))]
139        {
140            let _ = self;
141        }
142    }
143
144    pub(super) const fn device_unavailable_observed(&self) -> bool {
145        #[cfg(feature = "cuda-runtime")]
146        {
147            self.device_unavailable_observed
148        }
149        #[cfg(not(feature = "cuda-runtime"))]
150        {
151            let _ = self;
152            true
153        }
154    }
155
156    /// Create an encode-stage accelerator with optional CUDA stage timing collection.
157    #[must_use]
158    #[doc(hidden)]
159    pub fn with_profile_collection(collect_profile: bool) -> Self {
160        Self {
161            collect_profile,
162            ..Self::default()
163        }
164    }
165
166    /// Create the measured Auto route for host-output HTJ2K encode.
167    ///
168    /// CUDA keeps the DWT and HT code-block stages, while forward RCT and
169    /// Tier-2 packetization stay on the CPU for the current host-pixel path.
170    #[must_use]
171    pub fn for_auto_host_output() -> Self {
172        Self::default()
173            .prefer_cpu_forward_rct(true)
174            .prefer_cpu_packetization(true)
175    }
176
177    /// Prefer scalar CPU forward RCT while keeping later CUDA stages enabled.
178    #[must_use]
179    pub fn prefer_cpu_forward_rct(mut self, prefer_cpu_forward_rct: bool) -> Self {
180        self.prefer_cpu_forward_rct = prefer_cpu_forward_rct;
181        self
182    }
183
184    /// Prefer scalar CPU Tier-2 packetization while keeping CUDA Tier-1/HT block coding enabled.
185    ///
186    /// This is useful for batches of many small tiles where launching a CUDA
187    /// packetization kernel and copying several tiny descriptor buffers per tile
188    /// costs more than forming the packet body on the host.
189    #[must_use]
190    pub fn prefer_cpu_packetization(mut self, prefer_cpu_packetization: bool) -> Self {
191        self.prefer_cpu_packetization = prefer_cpu_packetization;
192        self
193    }
194
195    /// Prefer host sub-band quantization while keeping batched CUDA HT code-block encode enabled.
196    ///
197    /// This avoids launching one CUDA quantize/subband path for every prepared
198    /// subband in multi-resolution precomputed transcode outputs, where the
199    /// many tiny launches cost more than CPU quantization.
200    #[must_use]
201    pub fn prefer_cpu_ht_subband(mut self, prefer_cpu_ht_subband: bool) -> Self {
202        self.prefer_cpu_ht_subband = prefer_cpu_ht_subband;
203        self
204    }
205
206    /// Prefer host sub-band quantization while keeping CUDA HT code-block encode enabled.
207    ///
208    /// Multi-resolution transcode workloads can contain thousands of small
209    /// subbands; for those, CPU quantization plus one batched HT code-block
210    /// encode per tile is currently faster than launching CUDA quantization for
211    /// every subband.
212    #[must_use]
213    pub fn prefer_cpu_quantize_subband(mut self, prefer_cpu_quantize_subband: bool) -> Self {
214        self.prefer_cpu_quantize_subband = prefer_cpu_quantize_subband;
215        self
216    }
217
218    /// Return cumulative CUDA encode stage timings collected by this accelerator.
219    #[must_use]
220    pub const fn collected_stage_timings(&self) -> CudaEncodeStageTimings {
221        CudaEncodeStageTimings {
222            deinterleave_us: self.deinterleave_us,
223            mct_us: self.mct_us,
224            dwt_us: self.dwt_us,
225            quantize_us: self.quantize_us,
226            ht_encode_us: self.ht_encode_us,
227            packetize_us: self.packetize_us,
228        }
229    }
230
231    /// Clear cumulative CUDA encode stage timings without changing dispatch counters.
232    pub fn reset_collected_stage_timings(&mut self) {
233        self.deinterleave_us = 0;
234        self.mct_us = 0;
235        self.dwt_us = 0;
236        self.quantize_us = 0;
237        self.ht_encode_us = 0;
238        self.packetize_us = 0;
239    }
240
241    #[cfg(feature = "cuda-runtime")]
242    fn cuda_context(&mut self) -> CudaStageResult<Option<CudaContext>> {
243        if self.context.is_none() {
244            match CudaContext::system_default() {
245                Ok(context) => self.context = Some(context),
246                Err(error) if !cuda_runtime_required() && error.is_unavailable() => {
247                    self.device_unavailable_observed = true;
248                    return Ok(None);
249                }
250                Err(error) => {
251                    return Err(runtime_error("initialize CUDA encode context", error));
252                }
253            }
254        }
255        Ok(self.context.clone())
256    }
257
258    #[cfg(feature = "cuda-runtime")]
259    fn cuda_encode_resources(
260        &mut self,
261        context: &CudaContext,
262    ) -> CudaStageResult<Arc<CudaHtj2kEncodeResources>> {
263        if self.encode_resources.is_none() {
264            let resources = context
265                .upload_htj2k_encode_resources(cuda_htj2k_encode_tables())
266                .map_err(|error| runtime_error("upload CUDA HTJ2K encode resources", error))?;
267            self.encode_resources = Some(Arc::new(resources));
268        }
269        self.encode_resources
270            .clone()
271            .ok_or_else(|| internal_invariant("CUDA HTJ2K encode resources unavailable"))
272    }
273
274    pub(super) fn encode_profile_report(
275        &self,
276        encoded: &j2k::EncodedJ2k,
277        input_bytes: usize,
278        total_us: u128,
279    ) -> profile::CudaHtj2kEncodeProfileReport {
280        profile::CudaHtj2kEncodeProfileReport {
281            deinterleave_us: self.deinterleave_us,
282            mct_us: self.mct_us,
283            dwt_us: self.dwt_us,
284            quantize_us: self.quantize_us,
285            ht_encode_us: self.ht_encode_us,
286            packetize_us: self.packetize_us,
287            total_us,
288            input_bytes,
289            codestream_bytes: encoded.codestream.len(),
290            block_count: self.ht_code_block_attempts,
291            dispatch_count: self.dispatch_report().total(),
292            backend: encoded.backend,
293        }
294    }
295
296    /// Number of forward RCT attempts observed by crate-local diagnostics.
297    #[cfg(test)]
298    pub(crate) fn forward_rct_attempts(&self) -> usize {
299        self.forward_rct_attempts
300    }
301
302    /// Number of forward ICT attempts observed by crate-local diagnostics.
303    #[cfg(all(test, feature = "cuda-runtime"))]
304    pub(crate) fn forward_ict_attempts(&self) -> usize {
305        self.forward_ict_attempts
306    }
307
308    /// Number of forward 5/3 DWT attempts observed by crate-local diagnostics.
309    #[cfg(test)]
310    pub(crate) fn forward_dwt53_attempts(&self) -> usize {
311        self.forward_dwt53_attempts
312    }
313
314    /// Number of forward 9/7 DWT attempts observed by crate-local diagnostics.
315    #[cfg(all(test, feature = "cuda-runtime"))]
316    pub(crate) fn forward_dwt97_attempts(&self) -> usize {
317        self.forward_dwt97_attempts
318    }
319
320    /// Number of resident HTJ2K tile-body attempts observed by crate-local diagnostics.
321    #[cfg(all(test, feature = "cuda-runtime"))]
322    pub(crate) fn htj2k_tile_attempts(&self) -> usize {
323        self.htj2k_tile_attempts
324    }
325
326    /// Number of sub-band quantization attempts observed by crate-local diagnostics.
327    #[cfg(test)]
328    pub(crate) fn quantize_subband_attempts(&self) -> usize {
329        self.quantize_subband_attempts
330    }
331
332    /// Number of classic Tier-1 code-block attempts observed by crate-local diagnostics.
333    #[cfg(test)]
334    pub(crate) fn tier1_code_block_attempts(&self) -> usize {
335        self.tier1_code_block_attempts
336    }
337
338    /// Number of HT code-block attempts observed by crate-local diagnostics.
339    #[cfg(test)]
340    pub(crate) fn ht_code_block_attempts(&self) -> usize {
341        self.ht_code_block_attempts
342    }
343
344    /// Number of HT sub-band attempts observed by crate-local diagnostics.
345    #[cfg(test)]
346    pub(crate) fn ht_subband_attempts(&self) -> usize {
347        self.ht_subband_attempts
348    }
349
350    /// Number of packetization attempts observed by crate-local diagnostics.
351    #[cfg(test)]
352    pub(crate) fn packetization_attempts(&self) -> usize {
353        self.packetization_attempts
354    }
355
356    /// Number of deinterleave CUDA dispatches observed by crate-local diagnostics.
357    #[cfg(all(test, feature = "cuda-runtime"))]
358    pub(crate) fn deinterleave_dispatches(&self) -> usize {
359        self.deinterleave_dispatches
360    }
361
362    /// Number of forward RCT CUDA dispatches observed by crate-local diagnostics.
363    #[cfg(all(test, feature = "cuda-runtime"))]
364    pub(crate) fn forward_rct_dispatches(&self) -> usize {
365        self.forward_rct_dispatches
366    }
367
368    /// Number of forward ICT CUDA dispatches observed by crate-local diagnostics.
369    #[cfg(all(test, feature = "cuda-runtime"))]
370    pub(crate) fn forward_ict_dispatches(&self) -> usize {
371        self.forward_ict_dispatches
372    }
373
374    /// Number of forward 5/3 DWT CUDA dispatches observed by crate-local diagnostics.
375    #[cfg(all(test, feature = "cuda-runtime"))]
376    pub(crate) fn forward_dwt53_dispatches(&self) -> usize {
377        self.forward_dwt53_dispatches
378    }
379
380    /// Number of forward 9/7 DWT CUDA dispatches observed by crate-local diagnostics.
381    #[cfg(all(test, feature = "cuda-runtime"))]
382    pub(crate) fn forward_dwt97_dispatches(&self) -> usize {
383        self.forward_dwt97_dispatches
384    }
385
386    /// Number of resident HTJ2K tile-body CUDA dispatches observed by crate-local diagnostics.
387    #[cfg(all(test, feature = "cuda-runtime"))]
388    pub(crate) fn htj2k_tile_dispatches(&self) -> usize {
389        self.htj2k_tile_dispatches
390    }
391
392    /// Number of sub-band quantization CUDA dispatches observed by crate-local diagnostics.
393    #[cfg(all(test, feature = "cuda-runtime"))]
394    pub(crate) fn quantize_subband_dispatches(&self) -> usize {
395        self.quantize_subband_dispatches
396    }
397
398    /// Number of HT code-block CUDA dispatches observed by crate-local diagnostics.
399    #[cfg(all(test, feature = "cuda-runtime"))]
400    pub(crate) fn ht_code_block_dispatches(&self) -> usize {
401        self.ht_code_block_dispatches
402    }
403
404    /// Number of HT sub-band CUDA dispatches observed by crate-local diagnostics.
405    #[cfg(all(test, feature = "cuda-runtime"))]
406    pub(crate) fn ht_subband_dispatches(&self) -> usize {
407        self.ht_subband_dispatches
408    }
409
410    /// Number of packetization CUDA dispatches observed by crate-local diagnostics.
411    #[cfg(test)]
412    pub(crate) fn packetization_dispatches(&self) -> usize {
413        self.packetization_dispatches
414    }
415}
416
417#[cfg(feature = "cuda-runtime")]
418fn cuda_runtime_required() -> bool {
419    std::env::var_os("J2K_REQUIRE_CUDA_RUNTIME").is_some()
420}
421
422#[cfg(feature = "cuda-runtime")]
423pub(super) fn time_cuda_stage<T>(
424    name: &'static str,
425    context: &CudaContext,
426    collect_profile: bool,
427    work: impl FnMut() -> core::result::Result<T, CudaError>,
428) -> core::result::Result<(T, u128), CudaError> {
429    context.time_default_stream_named_us_if(collect_profile, name, work)
430}
431
432/// Cumulative CUDA encode-stage timings collected by `CudaEncodeStageAccelerator`.
433#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
434pub struct CudaEncodeStageTimings {
435    /// Pixel deinterleave and level-shift CUDA stage time.
436    pub deinterleave_us: u128,
437    /// Forward MCT CUDA stage time.
438    pub mct_us: u128,
439    /// Forward DWT CUDA stage time.
440    pub dwt_us: u128,
441    /// Quantization CUDA stage time.
442    pub quantize_us: u128,
443    /// HT code-block encode CUDA stage time.
444    pub ht_encode_us: u128,
445    /// HTJ2K packetization CUDA stage time.
446    pub packetize_us: u128,
447}
448
449impl CudaEncodeStageTimings {
450    /// Return field-wise saturating timing sums.
451    #[must_use]
452    pub const fn saturating_add(self, other: Self) -> Self {
453        Self {
454            deinterleave_us: self.deinterleave_us.saturating_add(other.deinterleave_us),
455            mct_us: self.mct_us.saturating_add(other.mct_us),
456            dwt_us: self.dwt_us.saturating_add(other.dwt_us),
457            quantize_us: self.quantize_us.saturating_add(other.quantize_us),
458            ht_encode_us: self.ht_encode_us.saturating_add(other.ht_encode_us),
459            packetize_us: self.packetize_us.saturating_add(other.packetize_us),
460        }
461    }
462
463    /// Total collected CUDA encode-stage time.
464    #[must_use]
465    pub const fn total_us(self) -> u128 {
466        self.deinterleave_us
467            .saturating_add(self.mct_us)
468            .saturating_add(self.dwt_us)
469            .saturating_add(self.quantize_us)
470            .saturating_add(self.ht_encode_us)
471            .saturating_add(self.packetize_us)
472    }
473}
474
475fn ht_subband_code_block_count(job: J2kHtSubbandEncodeJob<'_>) -> CudaStageResult<usize> {
476    if job.code_block_width == 0 || job.code_block_height == 0 {
477        return Err(J2kEncodeStageError::invalid_request(
478            "CUDA HTJ2K subband encode job has invalid code-block dimensions",
479        ));
480    }
481    let num_cbs_x = job.width.div_ceil(job.code_block_width);
482    let num_cbs_y = job.height.div_ceil(job.code_block_height);
483    (num_cbs_x as usize)
484        .checked_mul(num_cbs_y as usize)
485        .ok_or_else(|| arithmetic_overflow("CUDA HTJ2K subband code-block count overflow"))
486}
487
488#[doc(hidden)]
489impl J2kEncodeStageAccelerator for CudaEncodeStageAccelerator {
490    fn dispatch_report(&self) -> J2kEncodeDispatchReport {
491        J2kEncodeDispatchReport {
492            deinterleave: self.deinterleave_dispatches,
493            forward_rct: self.forward_rct_dispatches,
494            forward_ict: self.forward_ict_dispatches,
495            forward_dwt53: self.forward_dwt53_dispatches,
496            forward_dwt97: self.forward_dwt97_dispatches,
497            quantize_subband: self.quantize_subband_dispatches,
498            tier1_code_block: self.tier1_code_block_dispatches,
499            ht_code_block: self.ht_code_block_dispatches,
500            packetization: self.packetization_dispatches,
501        }
502    }
503
504    fn encode_deinterleave(
505        &mut self,
506        job: J2kDeinterleaveToF32Job<'_>,
507    ) -> CudaStageResult<Option<Vec<Vec<f32>>>> {
508        self.deinterleave_attempts = self.deinterleave_attempts.saturating_add(1);
509        if job.num_components > 4 {
510            emit_cuda_encode_route!(
511                ("op", "encode_deinterleave"),
512                ("decision", "cpu_fallback"),
513                ("reason", "component_count_unsupported"),
514                ("components", job.num_components),
515            );
516            return Ok(None);
517        }
518        #[cfg(feature = "cuda-runtime")]
519        if let Some(context) = self.cuda_context()? {
520            let num_components = cuda_component_count_u8(
521                job.num_components,
522                "CUDA deinterleave encode supports at most 255 components",
523            )?;
524            let (output, elapsed_us) = time_cuda_stage(
525                "j2k.j2k.cuda.encode.deinterleave",
526                &context,
527                self.collect_profile,
528                || {
529                    context.j2k_deinterleave_to_f32(
530                        job.pixels,
531                        job.num_pixels,
532                        num_components,
533                        job.bit_depth,
534                        job.signed,
535                    )
536                },
537            )
538            .map_err(|error| runtime_error("deinterleave encode pixels", error))?;
539            let dispatches = output.execution().kernel_dispatches();
540            self.deinterleave_dispatches = self.deinterleave_dispatches.saturating_add(dispatches);
541            self.deinterleave_us = self.deinterleave_us.saturating_add(elapsed_us);
542            emit_cuda_encode_route!(
543                ("op", "encode_deinterleave"),
544                ("decision", "cuda_dispatch"),
545                ("pixels", job.num_pixels),
546                ("components", job.num_components),
547                ("dispatches", dispatches),
548            );
549            return Ok(Some(output.into_components()));
550        }
551        #[cfg(not(feature = "cuda-runtime"))]
552        let _ = job;
553        emit_cuda_encode_route!(
554            ("op", "encode_deinterleave"),
555            ("decision", "cpu_fallback"),
556            ("reason", "cuda_unavailable"),
557        );
558        Ok(None)
559    }
560
561    fn encode_forward_rct(&mut self, job: J2kForwardRctJob<'_>) -> CudaStageResult<bool> {
562        self.forward_rct_attempts = self.forward_rct_attempts.saturating_add(1);
563        if self.prefer_cpu_forward_rct {
564            emit_cuda_encode_route!(
565                ("op", "encode_forward_rct"),
566                ("decision", "cpu_fallback"),
567                ("reason", "prefer_cpu_forward_rct"),
568            );
569            let _ = job;
570            return Ok(false);
571        }
572        #[cfg(feature = "cuda-runtime")]
573        if let Some(context) = self.cuda_context()? {
574            let (execution, elapsed_us) = time_cuda_stage(
575                "j2k.j2k.cuda.encode.rct",
576                &context,
577                self.collect_profile,
578                || context.j2k_forward_rct(job.plane0, job.plane1, job.plane2),
579            )
580            .map_err(|error| runtime_error("apply forward RCT", error))?;
581            self.forward_rct_dispatches = self
582                .forward_rct_dispatches
583                .saturating_add(execution.kernel_dispatches());
584            self.mct_us = self.mct_us.saturating_add(elapsed_us);
585            emit_cuda_encode_route!(
586                ("op", "encode_forward_rct"),
587                ("decision", "cuda_dispatch"),
588                ("dispatches", 1),
589            );
590            return Ok(true);
591        }
592        #[cfg(not(feature = "cuda-runtime"))]
593        let _ = job;
594        emit_cuda_encode_route!(
595            ("op", "encode_forward_rct"),
596            ("decision", "cpu_fallback"),
597            ("reason", "cuda_unavailable"),
598        );
599        Ok(false)
600    }
601
602    fn encode_forward_ict(&mut self, job: J2kForwardIctJob<'_>) -> CudaStageResult<bool> {
603        self.forward_ict_attempts = self.forward_ict_attempts.saturating_add(1);
604        #[cfg(feature = "cuda-runtime")]
605        if let Some(context) = self.cuda_context()? {
606            let (execution, elapsed_us) = time_cuda_stage(
607                "j2k.j2k.cuda.encode.ict",
608                &context,
609                self.collect_profile,
610                || context.j2k_forward_ict(job.plane0, job.plane1, job.plane2),
611            )
612            .map_err(|error| runtime_error("apply forward ICT", error))?;
613            self.forward_ict_dispatches = self
614                .forward_ict_dispatches
615                .saturating_add(execution.kernel_dispatches());
616            self.mct_us = self.mct_us.saturating_add(elapsed_us);
617            emit_cuda_encode_route!(
618                ("op", "encode_forward_ict"),
619                ("decision", "cuda_dispatch"),
620                ("dispatches", 1),
621            );
622            return Ok(true);
623        }
624        #[cfg(not(feature = "cuda-runtime"))]
625        let _ = job;
626        emit_cuda_encode_route!(
627            ("op", "encode_forward_ict"),
628            ("decision", "cpu_fallback"),
629            ("reason", "cuda_unavailable"),
630        );
631        Ok(false)
632    }
633
634    fn encode_forward_dwt53(
635        &mut self,
636        job: J2kForwardDwt53Job<'_>,
637    ) -> CudaStageResult<Option<J2kForwardDwt53Output>> {
638        self.forward_dwt53_attempts = self.forward_dwt53_attempts.saturating_add(1);
639        if job.num_levels == 0 {
640            emit_cuda_encode_route!(
641                ("op", "encode_forward_dwt53"),
642                ("decision", "cpu_fallback"),
643                ("reason", "zero_levels"),
644            );
645            return Ok(None);
646        }
647        #[cfg(feature = "cuda-runtime")]
648        if let Some(context) = self.cuda_context()? {
649            let (output, elapsed_us) = time_cuda_stage(
650                "j2k.j2k.cuda.encode.dwt53",
651                &context,
652                self.collect_profile,
653                || context.j2k_forward_dwt53(job.samples, job.width, job.height, job.num_levels),
654            )
655            .map_err(|error| runtime_error("apply forward 5/3 DWT", error))?;
656            let dispatches = output.execution().kernel_dispatches();
657            self.forward_dwt53_dispatches =
658                self.forward_dwt53_dispatches.saturating_add(dispatches);
659            self.dwt_us = self.dwt_us.saturating_add(elapsed_us);
660            emit_cuda_encode_route!(
661                ("op", "encode_forward_dwt53"),
662                ("decision", "cuda_dispatch"),
663                ("width", job.width),
664                ("height", job.height),
665                ("levels", job.num_levels),
666                ("dispatches", dispatches),
667            );
668            return Ok(Some(cuda_dwt53_output_to_j2k(&output)?));
669        }
670        #[cfg(not(feature = "cuda-runtime"))]
671        let _ = job;
672        emit_cuda_encode_route!(
673            ("op", "encode_forward_dwt53"),
674            ("decision", "cpu_fallback"),
675            ("reason", "cuda_unavailable"),
676        );
677        Ok(None)
678    }
679
680    fn encode_forward_dwt97(
681        &mut self,
682        job: J2kForwardDwt97Job<'_>,
683    ) -> CudaStageResult<Option<J2kForwardDwt97Output>> {
684        self.forward_dwt97_attempts = self.forward_dwt97_attempts.saturating_add(1);
685        if job.num_levels == 0 {
686            emit_cuda_encode_route!(
687                ("op", "encode_forward_dwt97"),
688                ("decision", "cpu_fallback"),
689                ("reason", "zero_levels"),
690            );
691            return Ok(None);
692        }
693        #[cfg(feature = "cuda-runtime")]
694        if let Some(context) = self.cuda_context()? {
695            let (output, elapsed_us) = time_cuda_stage(
696                "j2k.j2k.cuda.encode.dwt97",
697                &context,
698                self.collect_profile,
699                || context.j2k_forward_dwt97(job.samples, job.width, job.height, job.num_levels),
700            )
701            .map_err(|error| runtime_error("apply forward 9/7 DWT", error))?;
702            let dispatches = output.execution().kernel_dispatches();
703            self.forward_dwt97_dispatches =
704                self.forward_dwt97_dispatches.saturating_add(dispatches);
705            self.dwt_us = self.dwt_us.saturating_add(elapsed_us);
706            emit_cuda_encode_route!(
707                ("op", "encode_forward_dwt97"),
708                ("decision", "cuda_dispatch"),
709                ("width", job.width),
710                ("height", job.height),
711                ("levels", job.num_levels),
712                ("dispatches", dispatches),
713            );
714            return Ok(Some(cuda_dwt97_output_to_j2k(&output)?));
715        }
716        #[cfg(not(feature = "cuda-runtime"))]
717        let _ = job;
718        emit_cuda_encode_route!(
719            ("op", "encode_forward_dwt97"),
720            ("decision", "cpu_fallback"),
721            ("reason", "cuda_unavailable"),
722        );
723        Ok(None)
724    }
725
726    fn encode_quantize_subband(
727        &mut self,
728        job: J2kQuantizeSubbandJob<'_>,
729    ) -> CudaStageResult<Option<Vec<i32>>> {
730        self.quantize_subband_attempts = self.quantize_subband_attempts.saturating_add(1);
731        if self.prefer_cpu_quantize_subband {
732            emit_cuda_encode_route!(
733                ("op", "encode_quantize_subband"),
734                ("decision", "cpu_fallback"),
735                ("reason", "prefer_cpu_quantize_subband"),
736            );
737            let _ = job;
738            return Ok(None);
739        }
740        #[cfg(feature = "cuda-runtime")]
741        if let Some(context) = self.cuda_context()? {
742            let (output, elapsed_us) = time_cuda_stage(
743                "j2k.j2k.cuda.encode.quantize",
744                &context,
745                self.collect_profile,
746                || {
747                    context.j2k_quantize_subband(
748                        job.coefficients,
749                        CudaJ2kQuantizeJob {
750                            step_exponent: job.step_exponent,
751                            step_mantissa: job.step_mantissa,
752                            range_bits: job.range_bits,
753                            reversible: job.reversible,
754                        },
755                    )
756                },
757            )
758            .map_err(|error| runtime_error("quantize encode subband", error))?;
759            let dispatches = output.execution().kernel_dispatches();
760            self.quantize_subband_dispatches =
761                self.quantize_subband_dispatches.saturating_add(dispatches);
762            self.quantize_us = self.quantize_us.saturating_add(elapsed_us);
763            emit_cuda_encode_route!(
764                ("op", "encode_quantize_subband"),
765                ("decision", "cuda_dispatch"),
766                ("samples", job.coefficients.len()),
767                ("dispatches", dispatches),
768            );
769            return Ok(Some(output.into_coefficients()));
770        }
771        #[cfg(not(feature = "cuda-runtime"))]
772        let _ = job;
773        emit_cuda_encode_route!(
774            ("op", "encode_quantize_subband"),
775            ("decision", "cpu_fallback"),
776            ("reason", "cuda_unavailable"),
777        );
778        Ok(None)
779    }
780
781    fn encode_tier1_code_block(
782        &mut self,
783        _job: J2kTier1CodeBlockEncodeJob<'_>,
784    ) -> CudaStageResult<Option<EncodedJ2kCodeBlock>> {
785        self.tier1_code_block_attempts = self.tier1_code_block_attempts.saturating_add(1);
786        emit_cuda_encode_route!(
787            ("op", "encode_tier1_code_block"),
788            ("decision", "cpu_fallback"),
789            ("reason", "unsupported_stage"),
790        );
791        Ok(None)
792    }
793
794    fn encode_ht_code_block(
795        &mut self,
796        job: J2kHtCodeBlockEncodeJob<'_>,
797    ) -> CudaStageResult<Option<EncodedHtJ2kCodeBlock>> {
798        self.ht_code_block_attempts = self.ht_code_block_attempts.saturating_add(1);
799        #[cfg(feature = "cuda-runtime")]
800        if let Some(context) = self.cuda_context()? {
801            let resources = self.cuda_encode_resources(&context)?;
802            let encoded = cuda_encode_ht_code_block(&context, resources.as_ref(), job)?;
803            let dispatches = encoded.execution().kernel_dispatches();
804            let ht_encode_us = encoded.stage_timings().ht_encode_us;
805            let mut outputs = encoded_ht_code_blocks_from_cuda(encoded)?;
806            let output = outputs.pop().ok_or_else(|| {
807                internal_invariant("CUDA HTJ2K code-block encode returned no output")
808            })?;
809            self.ht_code_block_dispatches =
810                self.ht_code_block_dispatches.saturating_add(dispatches);
811            if self.collect_profile {
812                self.ht_encode_us = self.ht_encode_us.saturating_add(ht_encode_us);
813            }
814            emit_cuda_encode_route!(
815                ("op", "encode_ht_code_block"),
816                ("decision", "cuda_dispatch"),
817                ("width", job.width),
818                ("height", job.height),
819                ("dispatches", dispatches),
820            );
821            return Ok(Some(output));
822        }
823        #[cfg(not(feature = "cuda-runtime"))]
824        let _ = job;
825        emit_cuda_encode_route!(
826            ("op", "encode_ht_code_block"),
827            ("decision", "cpu_fallback"),
828            ("reason", "unsupported_stage"),
829        );
830        Ok(None)
831    }
832
833    fn encode_ht_code_blocks(
834        &mut self,
835        jobs: &[J2kHtCodeBlockEncodeJob<'_>],
836    ) -> CudaStageResult<Option<Vec<EncodedHtJ2kCodeBlock>>> {
837        self.ht_code_block_attempts = self.ht_code_block_attempts.saturating_add(jobs.len());
838        #[cfg(feature = "cuda-runtime")]
839        if let Some(context) = self.cuda_context()? {
840            let resources = self.cuda_encode_resources(&context)?;
841            let encoded = cuda_encode_ht_code_blocks(&context, resources.as_ref(), jobs)?;
842            let dispatches = encoded.execution().kernel_dispatches();
843            let ht_encode_us = encoded.stage_timings().ht_encode_us;
844            let outputs = encoded_ht_code_blocks_from_cuda(encoded)?;
845            self.ht_code_block_dispatches =
846                self.ht_code_block_dispatches.saturating_add(dispatches);
847            if self.collect_profile {
848                self.ht_encode_us = self.ht_encode_us.saturating_add(ht_encode_us);
849            }
850            emit_cuda_encode_route!(
851                ("op", "encode_ht_code_blocks"),
852                ("decision", "cuda_dispatch"),
853                ("jobs", jobs.len()),
854                ("dispatches", dispatches),
855            );
856            return Ok(Some(outputs));
857        }
858        #[cfg(not(feature = "cuda-runtime"))]
859        let _ = jobs;
860        emit_cuda_encode_route!(
861            ("op", "encode_ht_code_blocks"),
862            ("decision", "cpu_fallback"),
863            ("reason", "cuda_unavailable"),
864        );
865        Ok(None)
866    }
867
868    #[expect(
869        clippy::too_many_lines,
870        reason = "accelerator route preserves CUDA stage attempts, fallbacks, and counters"
871    )]
872    fn encode_htj2k_tile(
873        &mut self,
874        job: J2kHtj2kTileEncodeJob<'_>,
875    ) -> CudaStageResult<Option<Vec<u8>>> {
876        self.htj2k_tile_attempts = self.htj2k_tile_attempts.saturating_add(1);
877        if self.prefer_cpu_forward_rct || self.prefer_cpu_packetization {
878            emit_cuda_encode_route!(
879                ("op", "encode_htj2k_tile"),
880                ("decision", "cpu_fallback"),
881                ("reason", "prefer_stage_hybrid"),
882            );
883            let _ = job;
884            return Ok(None);
885        }
886        #[cfg(feature = "cuda-runtime")]
887        if let Some(context) = self.cuda_context()? {
888            let resources = self.cuda_encode_resources(&context)?;
889            let Some(encoded) = cuda_encode_htj2k_tile_body(
890                &context,
891                resources.as_ref(),
892                job,
893                self.collect_profile,
894            )?
895            else {
896                return Ok(None);
897            };
898            self.htj2k_tile_dispatches = self.htj2k_tile_dispatches.saturating_add(1);
899            self.deinterleave_attempts = self.deinterleave_attempts.saturating_add(1);
900            self.deinterleave_dispatches = self
901                .deinterleave_dispatches
902                .saturating_add(encoded.deinterleave_dispatches);
903            if job.use_mct {
904                if job.reversible {
905                    self.forward_rct_attempts = self.forward_rct_attempts.saturating_add(1);
906                } else {
907                    self.forward_ict_attempts = self.forward_ict_attempts.saturating_add(1);
908                }
909            }
910            self.forward_rct_dispatches = self
911                .forward_rct_dispatches
912                .saturating_add(encoded.forward_rct_dispatches);
913            self.forward_ict_dispatches = self
914                .forward_ict_dispatches
915                .saturating_add(encoded.forward_ict_dispatches);
916            if job.num_decomposition_levels > 0 {
917                if job.reversible {
918                    self.forward_dwt53_attempts = self
919                        .forward_dwt53_attempts
920                        .saturating_add(usize::from(job.num_components));
921                } else {
922                    self.forward_dwt97_attempts = self
923                        .forward_dwt97_attempts
924                        .saturating_add(usize::from(job.num_components));
925                }
926            }
927            self.forward_dwt53_dispatches = self
928                .forward_dwt53_dispatches
929                .saturating_add(encoded.forward_dwt53_dispatches);
930            self.forward_dwt97_dispatches = self
931                .forward_dwt97_dispatches
932                .saturating_add(encoded.forward_dwt97_dispatches);
933            self.quantize_subband_attempts = self
934                .quantize_subband_attempts
935                .saturating_add(encoded.quantize_jobs);
936            self.quantize_subband_dispatches = self
937                .quantize_subband_dispatches
938                .saturating_add(encoded.quantize_dispatches);
939            self.ht_code_block_attempts = self
940                .ht_code_block_attempts
941                .saturating_add(encoded.ht_code_block_jobs);
942            self.ht_code_block_dispatches = self
943                .ht_code_block_dispatches
944                .saturating_add(encoded.ht_code_block_dispatches);
945            self.packetization_attempts = self.packetization_attempts.saturating_add(1);
946            self.packetization_dispatches = self
947                .packetization_dispatches
948                .saturating_add(encoded.packetization_dispatches);
949            if self.collect_profile {
950                self.deinterleave_us = self
951                    .deinterleave_us
952                    .saturating_add(encoded.timings.deinterleave_us);
953                self.mct_us = self.mct_us.saturating_add(encoded.timings.mct_us);
954                self.dwt_us = self.dwt_us.saturating_add(encoded.timings.dwt_us);
955                self.quantize_us = self.quantize_us.saturating_add(encoded.timings.quantize_us);
956                self.ht_encode_us = self
957                    .ht_encode_us
958                    .saturating_add(encoded.timings.ht_encode_us);
959                self.packetize_us = self
960                    .packetize_us
961                    .saturating_add(encoded.timings.packetize_us);
962            }
963            emit_cuda_encode_route!(
964                ("op", "encode_htj2k_tile"),
965                ("decision", "cuda_dispatch"),
966                ("components", job.num_components),
967                ("blocks", encoded.ht_code_block_jobs),
968            );
969            return Ok(Some(encoded.tile_data));
970        }
971        #[cfg(not(feature = "cuda-runtime"))]
972        let _ = job;
973        emit_cuda_encode_route!(
974            ("op", "encode_htj2k_tile"),
975            ("decision", "cpu_fallback"),
976            ("reason", "cuda_unavailable"),
977        );
978        Ok(None)
979    }
980
981    fn encode_ht_subband(
982        &mut self,
983        job: J2kHtSubbandEncodeJob<'_>,
984    ) -> CudaStageResult<Option<Vec<EncodedHtJ2kCodeBlock>>> {
985        let code_block_count = ht_subband_code_block_count(job)?;
986        self.ht_subband_attempts = self.ht_subband_attempts.saturating_add(1);
987        self.quantize_subband_attempts = self.quantize_subband_attempts.saturating_add(1);
988        self.ht_code_block_attempts = self.ht_code_block_attempts.saturating_add(code_block_count);
989        if self.prefer_cpu_ht_subband {
990            emit_cuda_encode_route!(
991                ("op", "encode_ht_subband"),
992                ("decision", "cpu_fallback"),
993                ("reason", "prefer_cpu_ht_subband"),
994            );
995            return Ok(None);
996        }
997        #[cfg(feature = "cuda-runtime")]
998        if let Some(context) = self.cuda_context()? {
999            let resources = self.cuda_encode_resources(&context)?;
1000            let encoded =
1001                cuda_encode_ht_subband(&context, resources.as_ref(), job, self.collect_profile)?;
1002            let quantize_dispatches = encoded.quantize_dispatches;
1003            let encode_dispatches = encoded.encode.execution().kernel_dispatches();
1004            let timings = encoded.timings;
1005            let outputs = encoded_ht_code_blocks_from_cuda(encoded.encode)?;
1006            self.ht_subband_dispatches = self.ht_subband_dispatches.saturating_add(1);
1007            self.quantize_subband_dispatches = self
1008                .quantize_subband_dispatches
1009                .saturating_add(quantize_dispatches);
1010            self.ht_code_block_dispatches = self
1011                .ht_code_block_dispatches
1012                .saturating_add(encode_dispatches);
1013            if self.collect_profile {
1014                self.quantize_us = self.quantize_us.saturating_add(timings.quantize_us);
1015                self.ht_encode_us = self.ht_encode_us.saturating_add(timings.ht_encode_us);
1016            }
1017            emit_cuda_encode_route!(
1018                ("op", "encode_ht_subband"),
1019                ("decision", "cuda_dispatch"),
1020                ("width", job.width),
1021                ("height", job.height),
1022                ("blocks", code_block_count),
1023                ("quantize_dispatches", quantize_dispatches),
1024                ("encode_dispatches", encode_dispatches),
1025            );
1026            return Ok(Some(outputs));
1027        }
1028        #[cfg(not(feature = "cuda-runtime"))]
1029        let _ = job;
1030        emit_cuda_encode_route!(
1031            ("op", "encode_ht_subband"),
1032            ("decision", "cpu_fallback"),
1033            ("reason", "cuda_unavailable"),
1034        );
1035        Ok(None)
1036    }
1037
1038    fn encode_packetization(
1039        &mut self,
1040        job: J2kPacketizationEncodeJob<'_>,
1041    ) -> CudaStageResult<Option<Vec<u8>>> {
1042        self.packetization_attempts = self.packetization_attempts.saturating_add(1);
1043        if self.prefer_cpu_packetization {
1044            emit_cuda_encode_route!(
1045                ("op", "encode_packetization"),
1046                ("decision", "cpu_fallback"),
1047                ("reason", "prefer_cpu_packetization"),
1048            );
1049            let _ = job;
1050            return Ok(None);
1051        }
1052        let plan = match flatten_cuda_htj2k_packetization_job_classified(job) {
1053            Ok(plan) => plan,
1054            Err(error) => {
1055                let reason = cuda_packetization_plan_fallback_reason(error)?;
1056                emit_cuda_encode_route!(
1057                    ("op", "encode_packetization"),
1058                    ("decision", "cpu_fallback"),
1059                    ("reason", reason),
1060                );
1061                return Ok(None);
1062            }
1063        };
1064        #[cfg(feature = "cuda-runtime")]
1065        if let Some(context) = self.cuda_context()? {
1066            let mut host_budget = HostPhaseBudget::new("j2k CUDA HTJ2K staged packetization");
1067            host_budget
1068                .account_vec(&plan.payload)
1069                .map_err(|error| adapter_error("retain CUDA packetization payload", error))?;
1070            host_budget
1071                .account_vec(&plan.packets)
1072                .map_err(|error| adapter_error("retain CUDA packet descriptors", error))?;
1073            host_budget
1074                .account_vec(&plan.subbands)
1075                .map_err(|error| adapter_error("retain CUDA packet subbands", error))?;
1076            host_budget
1077                .account_vec(&plan.blocks)
1078                .map_err(|error| adapter_error("retain CUDA packet blocks", error))?;
1079            host_budget
1080                .account_vec(&plan.tag_states)
1081                .map_err(|error| adapter_error("retain CUDA packet tag states", error))?;
1082            host_budget
1083                .account_vec(&plan.tag_nodes)
1084                .map_err(|error| adapter_error("retain CUDA packet tag nodes", error))?;
1085            let packets = cuda_packetization_packets(&plan, &mut host_budget)?;
1086            let subbands = cuda_packetization_subbands(&plan, &mut host_budget)?;
1087            let blocks = cuda_packetization_blocks(&plan, &mut host_budget)?;
1088            let tag_states = cuda_packetization_tag_states(&plan, &mut host_budget)?;
1089            let tag_nodes = cuda_packetization_tag_nodes(&plan, &mut host_budget)?;
1090            let packetized = context
1091                .packetize_htj2k_cleanup_packets_with_tag_state_and_live_host_bytes(
1092                    &plan.payload,
1093                    &packets,
1094                    &subbands,
1095                    &blocks,
1096                    &tag_states,
1097                    &tag_nodes,
1098                    host_budget.live_bytes(),
1099                )
1100                .map_err(|error| runtime_error("packetize HTJ2K cleanup packets", error))?;
1101            let dispatches = packetized.execution().kernel_dispatches();
1102            let packetize_us = packetized.stage_timings().packetize_us;
1103            self.packetization_dispatches =
1104                self.packetization_dispatches.saturating_add(dispatches);
1105            if self.collect_profile {
1106                self.packetize_us = self.packetize_us.saturating_add(packetize_us);
1107            }
1108            emit_cuda_encode_route!(
1109                ("op", "encode_packetization"),
1110                ("decision", "cuda_dispatch"),
1111                ("packets", packets.len()),
1112                ("dispatches", dispatches),
1113            );
1114            return Ok(Some(packetized.into_data()));
1115        }
1116        #[cfg(not(feature = "cuda-runtime"))]
1117        let _ = plan;
1118        emit_cuda_encode_route!(
1119            ("op", "encode_packetization"),
1120            ("decision", "cpu_fallback"),
1121            ("reason", "unsupported_stage"),
1122        );
1123        Ok(None)
1124    }
1125}
j2k_cuda/encode/stage.rs

j2k_cuda/encode/
stage.rs
