j2k_ml/cuda/
batch.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3use burn_core::tensor::backend::Backend;
4use burn_cuda::{Cuda, CudaDevice};
5use j2k::{BatchDecodeOptions, EncodedImage, PreparedBatch, PreparedImage};
6use j2k_cuda::{
7    CudaBatchDecodeResult, CudaBatchDecoder as CodecDecoder, SubmittedCudaResidentBatch,
8};
9
10use crate::{BurnBatchDecode, BurnBatchGroupError, BurnDecodeError};
11
12/// Pending CUDA codec decode whose completed pixels will be staged through
13/// host memory and uploaded with Burn's ordinary tensor API.
14#[must_use = "submitted CUDA upload batches must be waited or dropped"]
15pub struct SubmittedCudaBurnBatch {
16    pending: SubmittedCudaResidentBatch,
17    device: CudaDevice,
18}
19
20impl core::fmt::Debug for SubmittedCudaBurnBatch {
21    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
22        f.debug_struct("SubmittedCudaBurnBatch")
23            .field("pending", &self.pending)
24            .field("device", &self.device)
25            .finish()
26    }
27}
28
29impl SubmittedCudaBurnBatch {
30    /// Number of successfully submitted homogeneous codec groups.
31    #[must_use]
32    pub fn len(&self) -> usize {
33        self.pending.pending_group_count()
34    }
35
36    /// Whether no accelerator codec work was submitted.
37    #[must_use]
38    pub fn is_empty(&self) -> bool {
39        self.len() == 0
40    }
41
42    /// Wait for CUDA codec completion, copy each completed group to host
43    /// staging, and upload it into an ordinary Burn CUDA tensor.
44    pub fn wait(self) -> Result<BurnBatchDecode<Cuda>, BurnDecodeError> {
45        materialize(self.pending.wait()?, &self.device)
46    }
47}
48
49/// Persistent CUDA codec session followed by an explicit staged Burn upload.
50///
51/// JPEG 2000 decoding executes on CUDA. Completed codec-owned device output is
52/// copied to host staging and then uploaded through [`burn_core::tensor::Tensor::from_data`].
53/// This type does not provide direct-destination or zero-copy behavior.
54#[derive(Debug)]
55pub struct CudaBurnDecoder {
56    codec: CodecDecoder,
57    device: CudaDevice,
58}
59
60impl CudaBurnDecoder {
61    /// Create a staged decoder for one Burn CUDA device.
62    #[must_use]
63    pub fn new(device: CudaDevice, options: BatchDecodeOptions) -> Self {
64        Self {
65            codec: CodecDecoder::with_options(options),
66            device,
67        }
68    }
69
70    /// Borrow the persistent CUDA codec session.
71    #[must_use]
72    pub const fn codec(&self) -> &CodecDecoder {
73        &self.codec
74    }
75
76    /// Borrow the Burn CUDA device receiving the staged upload.
77    #[must_use]
78    pub const fn device(&self) -> &CudaDevice {
79        &self.device
80    }
81
82    /// Parse and group owned inputs once for repeated decode calls.
83    pub fn prepare(&self, inputs: Vec<EncodedImage>) -> Result<PreparedBatch, BurnDecodeError> {
84        Ok(self.codec.prepare(inputs)?)
85    }
86
87    /// Regroup caller-supplied prepared images without reparsing codestream bytes.
88    pub fn prepare_prepared_images(
89        &self,
90        images: Vec<PreparedImage>,
91    ) -> Result<PreparedBatch, BurnDecodeError> {
92        Ok(self.codec.prepare_prepared_images(images)?)
93    }
94
95    /// Decode on CUDA, stage completed pixels through host memory, and upload to Burn.
96    pub fn decode(
97        &mut self,
98        inputs: Vec<EncodedImage>,
99    ) -> Result<BurnBatchDecode<Cuda>, BurnDecodeError> {
100        self.submit(inputs)?.wait()
101    }
102
103    /// Regroup prepared images, decode on CUDA, and perform the staged upload.
104    pub fn decode_prepared_images(
105        &mut self,
106        images: Vec<PreparedImage>,
107    ) -> Result<BurnBatchDecode<Cuda>, BurnDecodeError> {
108        let prepared = self.prepare_prepared_images(images)?;
109        self.decode_prepared(&prepared)
110    }
111
112    /// Decode a reusable preparation on CUDA and perform the staged upload.
113    pub fn decode_prepared(
114        &mut self,
115        prepared: &PreparedBatch,
116    ) -> Result<BurnBatchDecode<Cuda>, BurnDecodeError> {
117        self.submit_prepared(prepared)?.wait()
118    }
119
120    /// Submit CUDA codec work. The later [`SubmittedCudaBurnBatch::wait`]
121    /// performs a synchronous device-to-host copy and ordinary Burn upload.
122    pub fn submit(
123        &mut self,
124        inputs: Vec<EncodedImage>,
125    ) -> Result<SubmittedCudaBurnBatch, BurnDecodeError> {
126        let prepared = self.prepare(inputs)?;
127        self.submit_prepared(&prepared)
128    }
129
130    /// Submit a reusable preparation to codec-owned CUDA output.
131    pub fn submit_prepared(
132        &mut self,
133        prepared: &PreparedBatch,
134    ) -> Result<SubmittedCudaBurnBatch, BurnDecodeError> {
135        ensure_dtypes::<Cuda>(prepared, &self.device)?;
136        Ok(SubmittedCudaBurnBatch {
137            pending: self.codec.submit_prepared(prepared)?,
138            device: self.device.clone(),
139        })
140    }
141}
142
143fn ensure_dtypes<B: Backend>(
144    prepared: &PreparedBatch,
145    device: &B::Device,
146) -> Result<(), BurnDecodeError> {
147    for group in prepared.groups() {
148        let dtype = crate::batch_contract::dtype(group.info().sample_type)?;
149        if !B::supports_dtype(device, dtype) {
150            return Err(BurnDecodeError::UnsupportedDType { dtype });
151        }
152    }
153    Ok(())
154}
155
156fn materialize(
157    decoded: CudaBatchDecodeResult,
158    device: &CudaDevice,
159) -> Result<BurnBatchDecode<Cuda>, BurnDecodeError> {
160    let (codec_groups, errors, codec_group_errors) = decoded.into_parts();
161    let mut groups = Vec::new();
162    groups
163        .try_reserve_exact(codec_groups.len())
164        .map_err(|_| BurnDecodeError::SizeOverflow)?;
165    for group in codec_groups {
166        let (info, source_indices, decoded_rects, warnings, _surfaces, dense) = group.into_parts();
167        let mut bytes = vec![0; crate::staging::byte_len(source_indices.len(), &info)?];
168        dense.buffer().copy_to_host(&mut bytes).map_err(|source| {
169            BurnDecodeError::AcceleratorInterop {
170                backend: "CUDA staged readback",
171                message: source.to_string(),
172            }
173        })?;
174        groups.push(crate::staging::materialize(
175            info,
176            source_indices,
177            decoded_rects,
178            warnings,
179            bytes,
180            device,
181        )?);
182    }
183    let group_errors = codec_group_errors
184        .into_iter()
185        .map(|error| {
186            let (source_indices, source) = error.into_parts();
187            cuda_group_error(source_indices, source)
188        })
189        .collect();
190    Ok(BurnBatchDecode {
191        groups,
192        errors,
193        group_errors,
194    })
195}
196
197fn cuda_group_error(source_indices: Vec<usize>, source: j2k_cuda::Error) -> BurnBatchGroupError {
198    let adapter_source = BurnDecodeError::Cuda(j2k_cuda::CudaBatchError::GroupExecution {
199        source_indices: source_indices.clone(),
200        source: Box::new(source),
201    });
202    BurnBatchGroupError::new(source_indices, adapter_source)
203}
204
205#[cfg(test)]
206mod tests {
207    use super::cuda_group_error;
208    use crate::BurnDecodeError;
209
210    #[test]
211    fn cuda_group_failure_preserves_the_existing_typed_error_chain() {
212        let error = cuda_group_error(
213            vec![2, 4],
214            j2k_cuda::Error::UnsupportedCudaRequest {
215                reason: "test rejection",
216            },
217        );
218
219        assert_eq!(error.source_indices(), [2, 4]);
220        let BurnDecodeError::Cuda(j2k_cuda::CudaBatchError::GroupExecution {
221            source_indices,
222            source,
223            ..
224        }) = error.source()
225        else {
226            panic!("CUDA codec failure must use the patch-compatible batch error variant");
227        };
228        assert_eq!(source_indices, &[2, 4]);
229        assert!(matches!(
230            source.as_ref(),
231            j2k_cuda::Error::UnsupportedCudaRequest {
232                reason: "test rejection"
233            }
234        ));
235    }
236}
j2k_ml/cuda/batch.rs

j2k_ml/cuda/
batch.rs
