Skip to main content

j2k_cuda/batch/
decoder.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! Persistent CUDA batch decoder facade.
4
5#[cfg(not(feature = "cuda-runtime"))]
6use super::IndexedBatchError;
7#[cfg(feature = "cuda-runtime")]
8use super::{
9    decode_warnings, group_pixel_format, native_color_inputs, native_decode_settings,
10    native_referenced_classic_plan, native_referenced_htj2k_plan, validate_layout,
11    CudaExternalBatchGroup, PixelFormat, PreparedBatchGroup, SubmittedCudaCodecBatch,
12    SubmittedCudaExternalBatch,
13};
14use super::{
15    prepare_batch, prepare_batch_from_images, BatchDecodeOptions, BatchDecoder,
16    BatchInfrastructureError, CudaBatchDecodeResult, CudaBatchError, CudaSession, EncodedImage,
17    Error, PreparedBatch, PreparedImage,
18};
19
20/// Persistent CUDA batch decoder that reuses one [`CudaSession`].
21#[derive(Clone, Debug, Default)]
22pub struct CudaBatchDecoder {
23    pub(super) session: CudaSession,
24    pub(super) options: BatchDecodeOptions,
25}
26
27impl CudaBatchDecoder {
28    /// Create a decoder with a lazily initialized CUDA session and strict
29    /// shared batch options.
30    #[must_use]
31    pub fn new() -> Self {
32        Self::default()
33    }
34
35    /// Create a decoder with explicit shared preparation options.
36    #[must_use]
37    pub fn with_options(options: BatchDecodeOptions) -> Self {
38        Self {
39            session: CudaSession::default(),
40            options,
41        }
42    }
43
44    /// Create a decoder around an existing CUDA session.
45    #[must_use]
46    pub fn with_session(session: CudaSession) -> Self {
47        Self {
48            session,
49            options: BatchDecodeOptions::default(),
50        }
51    }
52
53    /// Create a decoder around an existing session and preparation policy.
54    #[must_use]
55    pub const fn with_session_and_options(
56        session: CudaSession,
57        options: BatchDecodeOptions,
58    ) -> Self {
59        Self { session, options }
60    }
61
62    /// Shared preparation options used by [`Self::prepare`] and
63    /// [`Self::decode_batch`].
64    #[must_use]
65    pub const fn options(&self) -> BatchDecodeOptions {
66        self.options
67    }
68
69    /// Borrow the persistent session for diagnostics.
70    #[must_use]
71    pub const fn session(&self) -> &CudaSession {
72        &self.session
73    }
74
75    /// Snapshot the persistent session's two private decode buffer pools.
76    #[cfg(feature = "cuda-runtime")]
77    pub fn decode_pool_diagnostics(&self) -> Result<crate::CudaDecodePoolDiagnostics, Error> {
78        self.session.decode_pool_diagnostics()
79    }
80
81    /// Snapshot CUDA transfer/event counters and retained decode-pool memory.
82    #[cfg(feature = "cuda-runtime")]
83    pub fn diagnostics(&self) -> Result<crate::CudaSessionDiagnostics, Error> {
84        self.session.diagnostics()
85    }
86
87    /// Mutably borrow the persistent session for advanced configuration.
88    #[must_use]
89    pub fn session_mut(&mut self) -> &mut CudaSession {
90        &mut self.session
91    }
92
93    /// Inspect and group owned inputs without copying their compressed bytes.
94    pub fn prepare(
95        &self,
96        inputs: Vec<EncodedImage>,
97    ) -> Result<PreparedBatch, BatchInfrastructureError> {
98        prepare_batch(inputs, self.options)
99    }
100
101    /// Regroup caller-supplied prepared images without reparsing codestream bytes.
102    ///
103    /// Returned source indices are positions in `images`; each image retains
104    /// its original [`PreparedImage::source_index`] for provenance.
105    pub fn prepare_prepared_images(
106        &self,
107        images: Vec<PreparedImage>,
108    ) -> Result<PreparedBatch, BatchInfrastructureError> {
109        prepare_batch_from_images(images, self.options)
110    }
111
112    /// Prepare and strictly decode one owned batch to CUDA-resident groups.
113    pub fn decode_batch(
114        &mut self,
115        inputs: Vec<EncodedImage>,
116    ) -> Result<CudaBatchDecodeResult, CudaBatchError> {
117        let prepared = self.prepare(inputs)?;
118        self.decode_prepared(&prepared)
119    }
120
121    /// Regroup and decode prepared images without reparsing their encoded bytes.
122    pub fn decode_prepared_images(
123        &mut self,
124        images: Vec<PreparedImage>,
125    ) -> Result<CudaBatchDecodeResult, CudaBatchError> {
126        let prepared = self.prepare_prepared_images(images)?;
127        self.decode_prepared(&prepared)
128    }
129
130    /// Strictly decode a reusable shared codec batch to CUDA-resident groups.
131    ///
132    /// Explicit CUDA execution never falls back to CPU-decoded pixels. A CUDA
133    /// execution error discards the entire affected dense group.
134    pub fn decode_prepared(
135        &mut self,
136        prepared: &PreparedBatch,
137    ) -> Result<CudaBatchDecodeResult, CudaBatchError> {
138        #[cfg(feature = "cuda-runtime")]
139        {
140            self.submit_prepared(prepared)?.wait()
141        }
142        #[cfg(not(feature = "cuda-runtime"))]
143        {
144            if let Some(group) = prepared.groups().first() {
145                return Err(CudaBatchError::group(group, Error::CudaUnavailable));
146            }
147            let mut errors = Vec::new();
148            errors
149                .try_reserve_exact(prepared.errors().len())
150                .map_err(|_| BatchInfrastructureError::HostAllocationFailed {
151                    what: "CUDA stub indexed errors",
152                    bytes: prepared
153                        .errors()
154                        .len()
155                        .saturating_mul(core::mem::size_of::<IndexedBatchError>()),
156                })?;
157            errors.extend_from_slice(prepared.errors());
158            Ok(CudaBatchDecodeResult {
159                groups: Vec::new(),
160                errors,
161                group_errors: Vec::new(),
162            })
163        }
164    }
165
166    /// Decode one prepared exact-native Gray/RGB/RGBA group directly into a
167    /// validated caller-owned CUDA allocation range.
168    ///
169    /// The destination must belong to this decoder's CUDA context and cover
170    /// the tightly concatenated group output. Decoded pixels are never staged
171    /// through host memory or copied through an intermediate device output.
172    ///
173    /// # Safety
174    ///
175    /// The destination allocation must remain live until this method returns
176    /// success. If CUDA completion cannot be proven and an error is returned,
177    /// the caller must quarantine the allocation rather than free or reuse it.
178    #[cfg(feature = "cuda-runtime")]
179    pub unsafe fn decode_batch_into(
180        &mut self,
181        group: &PreparedBatchGroup,
182        destination: &mut j2k_cuda_runtime::CudaExternalDeviceBufferViewMut<'_>,
183    ) -> Result<CudaExternalBatchGroup, CudaBatchError> {
184        // SAFETY: this synchronous convenience immediately waits while the
185        // caller's external view and exclusive managed-owner borrow are live.
186        unsafe { self.submit_batch_into(group, destination) }?.wait()
187    }
188
189    /// Submit one prepared exact-native Gray/RGB/RGBA group into CUDA storage
190    /// without a host completion wait.
191    ///
192    /// Callers integrating another CUDA runtime must use
193    /// [`j2k_cuda_runtime::CudaContext::with_primary_stream_ordering`] around
194    /// this submission, then retain the returned value alongside the tensor so
195    /// codec-internal resources outlive the ordered final store.
196    ///
197    /// # Safety
198    ///
199    /// The destination allocation must remain live and may only be consumed on
200    /// a CUDA stream ordered after codec completion until this value is waited
201    /// or dropped. No unordered host or device access may overlap the decode.
202    /// Stream completion alone does not validate entropy status: callers must
203    /// not expose the destination as decoded pixels until
204    /// [`SubmittedCudaExternalBatch::wait`] succeeds. If CUDA completion
205    /// cannot be proven, the external allocation must be quarantined rather
206    /// than freed or reused.
207    #[cfg(feature = "cuda-runtime")]
208    #[expect(
209        clippy::too_many_lines,
210        reason = "external submission keeps routing, safety-critical destination ownership, and fallible metadata capture in one boundary"
211    )]
212    pub unsafe fn submit_batch_into(
213        &mut self,
214        group: &PreparedBatchGroup,
215        destination: &mut j2k_cuda_runtime::CudaExternalDeviceBufferViewMut<'_>,
216    ) -> Result<SubmittedCudaExternalBatch, CudaBatchError> {
217        let fmt = group_pixel_format(group.info())
218            .and_then(|fmt| {
219                validate_layout(group.info())?;
220                Ok(fmt)
221            })
222            .map_err(|source| CudaBatchError::group(group, source))?;
223        let pending = if matches!(
224            fmt,
225            PixelFormat::Rgb8
226                | PixelFormat::Rgb16
227                | PixelFormat::RgbI16
228                | PixelFormat::Rgba8
229                | PixelFormat::Rgba16
230                | PixelFormat::RgbaI16
231        ) {
232            let inputs = native_color_inputs(group)
233                .map_err(|source| CudaBatchError::group(group, source))?;
234            SubmittedCudaCodecBatch::Color(
235                crate::decoder::submit_native_color_resident_prepared_batch_into(
236                    &inputs,
237                    &mut self.session,
238                    fmt,
239                    group.info().layout,
240                    destination,
241                )
242                .map_err(|source| CudaBatchError::group(group, source))?,
243            )
244        } else if matches!(
245            fmt,
246            PixelFormat::Gray8 | PixelFormat::Gray16 | PixelFormat::GrayI16
247        ) {
248            let inputs = group
249                .images()
250                .iter()
251                .zip(group.source_indices().iter().copied())
252                .map(|(image, source_index)| {
253                    let referenced_plan = image
254                        .htj2k_plan()
255                        .map(native_referenced_htj2k_plan)
256                        .transpose()?;
257                    let referenced_classic_plan = image
258                        .classic_plan()
259                        .map(native_referenced_classic_plan)
260                        .transpose()?;
261                    Ok(crate::decoder::grayscale_batch::GrayscaleBatchInput {
262                        source_index,
263                        bytes: image.bytes().as_ref(),
264                        device_plan: Some(image.plan()),
265                        referenced_plan,
266                        referenced_classic_plan,
267                    })
268                })
269                .collect::<Result<Vec<_>, Error>>()
270                .map_err(|source| CudaBatchError::group(group, source))?;
271            SubmittedCudaCodecBatch::Grayscale(
272                crate::decoder::grayscale_batch::submit_grayscale_cuda_resident_prepared_batch_into(
273                    &inputs,
274                    native_decode_settings(group.options().settings),
275                    &mut self.session,
276                    fmt,
277                    destination,
278                )
279                .map_err(|source| CudaBatchError::group(group, source))?,
280            )
281        } else {
282            return Err(CudaBatchError::group(
283                group,
284                Error::UnsupportedCudaRequest {
285                    reason:
286                        "direct external CUDA batch decode requires exact Gray, RGB, or RGBA output",
287                },
288            ));
289        };
290        let mut metadata_budget =
291            crate::allocation::HostPhaseBudget::new("CUDA external batch result metadata");
292        let source_indices = metadata_budget
293            .try_clone_slice(group.source_indices())
294            .map_err(|source| CudaBatchError::group(group, source))?;
295        let decoded_rects = metadata_budget
296            .try_collect_exact(
297                group
298                    .images()
299                    .iter()
300                    .map(|image| image.plan().output_rect()),
301            )
302            .map_err(|source| CudaBatchError::group(group, source))?;
303        let warnings = decode_warnings(group.images())
304            .map_err(|source| CudaBatchError::group(group, source))?;
305        let ranges = metadata_budget
306            .try_clone_slice(pending.ranges())
307            .map_err(|source| CudaBatchError::group(group, source))?;
308        let external_group = CudaExternalBatchGroup {
309            info: group.info().clone(),
310            source_indices,
311            decoded_rects,
312            warnings,
313            ranges,
314        };
315        Ok(SubmittedCudaExternalBatch {
316            group: external_group,
317            pending,
318        })
319    }
320}
321
322impl BatchDecoder for CudaBatchDecoder {
323    type Output = CudaBatchDecodeResult;
324    type Error = CudaBatchError;
325
326    fn decode_prepared(&mut self, prepared: &PreparedBatch) -> Result<Self::Output, Self::Error> {
327        CudaBatchDecoder::decode_prepared(self, prepared)
328    }
329
330    fn options(&self) -> BatchDecodeOptions {
331        self.options
332    }
333}