Skip to main content

j2k_cuda_runtime/
context.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3use std::sync::Arc;
4
5use crate::{
6    error::CudaError,
7    execution::{CudaExecutionStats, CudaLaunchMode},
8    htj2k_decode::{
9        htj2k_decode_needs_zero_fill, CudaHtj2kCodeBlockJob, CudaHtj2kDecodeOutput,
10        CudaHtj2kDecodeStageTimings, CudaQueuedHtj2kCleanup,
11    },
12    memory::pooled_device_buffer,
13};
14
15mod band_transfer;
16mod compact;
17mod creation;
18mod device;
19mod host_budget;
20mod inner;
21mod kernel_cache;
22mod kernel_dispatch;
23mod lifecycle;
24mod operations;
25mod pinned_host;
26mod pointer;
27mod resource_creation;
28#[cfg(test)]
29mod test_kernels;
30
31pub use self::compact::{CudaHtj2kCompactEncodedCodeBlock, CudaHtj2kCompactEncodedCodeBlocks};
32#[doc(hidden)]
33pub use self::host_budget::{CudaExternalHostOwner, CudaExternalHostReservation};
34#[cfg(test)]
35pub(crate) use self::pinned_host::validate_non_null_pinned_host_allocation;
36#[cfg(test)]
37pub(crate) use self::test_kernels::{CudaKernelModule, CudaKernelName};
38pub(crate) use self::{
39    band_transfer::cuda_idwt_trace_enabled,
40    compact::HTJ2K_UVLC_ENCODE_TABLE_BYTES,
41    inner::ContextInner,
42    kernel_cache::{CompiledKernel, CompiledKernelKey},
43    lifecycle::ContextResourceLifecycle,
44    operations::ensure_context_ownership,
45    pinned_host::PinnedUploadStaging,
46    resource_creation::{validate_device_allocation, validate_resource_handle},
47};
48
49/// CUDA driver context shared by J2K CUDA adapter crates.
50#[derive(Clone)]
51pub struct CudaContext {
52    pub(crate) inner: Arc<ContextInner>,
53}
54
55impl CudaContext {
56    /// Returns whether both handles own the same CUDA driver context.
57    #[doc(hidden)]
58    #[must_use]
59    pub fn is_same_context(&self, other: &Self) -> bool {
60        self.inner.context == other.inner.context
61    }
62
63    /// Device ordinal associated with this context.
64    #[doc(hidden)]
65    #[must_use]
66    pub fn device_ordinal(&self) -> usize {
67        self.inner.device_ordinal
68    }
69
70    /// Dequantize HTJ2K cleanup outputs using the metadata buffer already held
71    /// live by a queued cleanup launch.
72    #[doc(hidden)]
73    pub fn j2k_dequantize_queued_htj2k_cleanup_with_pool(
74        &self,
75        cleanup: &CudaQueuedHtj2kCleanup,
76    ) -> Result<CudaExecutionStats, CudaError> {
77        if !self.is_same_context(&cleanup.context) {
78            return Err(CudaError::InvalidArgument {
79                message: "queued HTJ2K cleanup belongs to a different CUDA context".to_string(),
80            });
81        }
82        self.inner.set_current()?;
83        if cleanup.status_count == 0 {
84            return Ok(CudaExecutionStats::default());
85        }
86        let Some(jobs_buffer) = cleanup.resources.first() else {
87            return Err(CudaError::InvalidArgument {
88                message: "queued HTJ2K cleanup has no metadata buffer".to_string(),
89            });
90        };
91        self.launch_j2k_dequantize_htj2k_cleanup_jobs_multi(
92            pooled_device_buffer(jobs_buffer)?,
93            cleanup.status_count,
94            CudaLaunchMode::Sync,
95        )?;
96        Ok(CudaExecutionStats {
97            kernel_dispatches: 1,
98            copy_kernel_dispatches: 0,
99            decode_kernel_dispatches: 1,
100            hardware_decode: false,
101        })
102    }
103
104    pub(crate) fn decode_empty_htj2k_codeblocks(
105        &self,
106        jobs: &[CudaHtj2kCodeBlockJob],
107        output_words: usize,
108    ) -> Result<CudaHtj2kDecodeOutput, CudaError> {
109        self.inner.set_current()?;
110        let output_bytes = output_words
111            .checked_mul(std::mem::size_of::<f32>())
112            .ok_or(CudaError::LengthTooLarge { len: output_words })?;
113        let coefficients = self.allocate(output_bytes)?;
114        if htj2k_decode_needs_zero_fill(jobs, output_words)? {
115            self.memset_d32(&coefficients, 0, output_words)?;
116            self.synchronize()?;
117        }
118        Ok(CudaHtj2kDecodeOutput {
119            coefficients,
120            execution: CudaExecutionStats::default(),
121            statuses: Vec::new(),
122            stage_timings: CudaHtj2kDecodeStageTimings::default(),
123        })
124    }
125}
126
127impl std::fmt::Debug for CudaContext {
128    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
129        f.debug_struct("CudaContext").finish_non_exhaustive()
130    }
131}
132
133#[cfg(test)]
134mod structure_tests;