j2k_cuda/batch/
resident_submission.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3use j2k::{
4    EncodedImage, IndexedBatchError, J2kDecodeWarning, PreparedBatch, PreparedBatchGroup, Rect,
5};
6use j2k_core::{BatchInfrastructureError, PixelFormat};
7
8use super::{
9    decode_warnings, group_pixel_format, native_color_group_storage, native_color_inputs,
10    native_decode_settings, native_referenced_classic_plan, native_referenced_htj2k_plan,
11    validate_layout, CudaBatchDecodeResult, CudaBatchDecoder, CudaBatchError, CudaBatchGroup,
12    CudaBatchGroupError, CudaResidentBatchBuffer, Error, Surface,
13};
14
15struct ResidentGroupMetadata {
16    info: j2k::BatchGroupInfo,
17    source_indices: Vec<usize>,
18    decoded_rects: Vec<Rect>,
19    warnings: Vec<Vec<J2kDecodeWarning>>,
20}
21
22impl ResidentGroupMetadata {
23    fn from_prepared(group: &PreparedBatchGroup) -> Result<Self, Error> {
24        let mut budget =
25            crate::allocation::HostPhaseBudget::new("CUDA resident batch result metadata");
26        Ok(Self {
27            info: group.info().clone(),
28            source_indices: budget.try_clone_slice(group.source_indices())?,
29            decoded_rects: budget.try_collect_exact(
30                group
31                    .images()
32                    .iter()
33                    .map(|image| image.plan().output_rect()),
34            )?,
35            warnings: decode_warnings(group.images())?,
36        })
37    }
38
39    fn finish(
40        self,
41        surfaces: Vec<Surface>,
42        dense_output: CudaResidentBatchBuffer,
43    ) -> CudaBatchGroup {
44        CudaBatchGroup {
45            info: self.info,
46            source_indices: self.source_indices,
47            decoded_rects: self.decoded_rects,
48            warnings: self.warnings,
49            surfaces,
50            dense_output,
51        }
52    }
53}
54
55enum SubmittedResidentCodecGroup {
56    Grayscale(crate::decoder::grayscale_batch::SubmittedGrayscaleResidentBatch),
57    Color(crate::decoder::SubmittedNativeColorResidentBatch),
58}
59
60impl SubmittedResidentCodecGroup {
61    fn is_complete(&self) -> Result<bool, Error> {
62        match self {
63            Self::Grayscale(pending) => pending.is_complete(),
64            Self::Color(pending) => pending.is_complete(),
65        }
66    }
67
68    fn finish(
69        self,
70        info: &j2k::BatchGroupInfo,
71    ) -> Result<(Vec<Surface>, CudaResidentBatchBuffer), Error> {
72        match self {
73            Self::Grayscale(pending) => {
74                let (output, _report) = pending.finish()?;
75                Ok((
76                    output.surfaces,
77                    CudaResidentBatchBuffer {
78                        buffer: output.buffer,
79                        ranges: output.ranges,
80                    },
81                ))
82            }
83            Self::Color(pending) => {
84                let (output, _report) = pending.finish()?;
85                let fmt = group_pixel_format(info)?;
86                Ok(native_color_group_storage(info, fmt, output))
87            }
88        }
89    }
90}
91
92struct SubmittedResidentGroup {
93    metadata: ResidentGroupMetadata,
94    pending: SubmittedResidentCodecGroup,
95}
96
97/// Asynchronously submitted codec-owned CUDA resident batch.
98///
99/// Every homogeneous group owns its final CUDA allocation before submission.
100/// [`Self::wait`] exposes those allocations only after final-store completion
101/// and entropy-status validation. Dropping this guard safely retires all work.
102#[must_use = "submitted CUDA resident decode must be retained or waited"]
103pub struct SubmittedCudaResidentBatch {
104    pending: Vec<SubmittedResidentGroup>,
105    errors: Vec<IndexedBatchError>,
106    group_errors: Vec<CudaBatchGroupError>,
107}
108
109impl core::fmt::Debug for SubmittedCudaResidentBatch {
110    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
111        f.debug_struct("SubmittedCudaResidentBatch")
112            .field("pending_groups", &self.pending.len())
113            .field("errors", &self.errors)
114            .field("group_errors", &self.group_errors)
115            .finish()
116    }
117}
118
119impl SubmittedCudaResidentBatch {
120    /// Number of successfully submitted homogeneous groups still retained.
121    #[must_use]
122    pub fn pending_group_count(&self) -> usize {
123        self.pending.len()
124    }
125
126    /// Query whether every successfully submitted group has completed.
127    pub fn is_complete(&self) -> Result<bool, CudaBatchError> {
128        for group in &self.pending {
129            let source_indices = group.metadata.source_indices.clone();
130            if !group
131                .pending
132                .is_complete()
133                .map_err(|source| CudaBatchError::GroupExecution {
134                    source_indices,
135                    source: Box::new(source),
136                })?
137            {
138                return Ok(false);
139            }
140        }
141        Ok(true)
142    }
143
144    /// Wait once for all group final stores and return codec-owned outputs.
145    pub fn wait(mut self) -> Result<CudaBatchDecodeResult, CudaBatchError> {
146        let mut groups = Vec::new();
147        groups.try_reserve_exact(self.pending.len()).map_err(|_| {
148            BatchInfrastructureError::HostAllocationFailed {
149                what: "completed CUDA resident batch groups",
150                bytes: self
151                    .pending
152                    .len()
153                    .saturating_mul(core::mem::size_of::<CudaBatchGroup>()),
154            }
155        })?;
156        let mut fatal = None;
157        for submitted in self.pending.drain(..) {
158            let SubmittedResidentGroup { metadata, pending } = submitted;
159            let source_indices = metadata.source_indices.clone();
160            match pending.finish(&metadata.info) {
161                Ok((surfaces, dense_output)) => {
162                    groups.push(metadata.finish(surfaces, dense_output));
163                }
164                Err(source) if source.session_is_unusable() => {
165                    if fatal.is_none() {
166                        fatal = Some(CudaBatchError::GroupExecution {
167                            source_indices,
168                            source: Box::new(source),
169                        });
170                    }
171                }
172                Err(source) => self
173                    .group_errors
174                    .push(CudaBatchGroupError::from_parts(source_indices, source)),
175            }
176        }
177        if let Some(error) = fatal {
178            return Err(error);
179        }
180        Ok(CudaBatchDecodeResult {
181            groups,
182            errors: core::mem::take(&mut self.errors),
183            group_errors: core::mem::take(&mut self.group_errors),
184        })
185    }
186}
187
188impl CudaBatchDecoder {
189    /// Prepare and asynchronously submit a codec-owned CUDA resident batch.
190    pub fn submit_batch(
191        &mut self,
192        inputs: Vec<EncodedImage>,
193    ) -> Result<SubmittedCudaResidentBatch, CudaBatchError> {
194        let prepared = self.prepare(inputs)?;
195        self.submit_prepared(&prepared)
196    }
197
198    /// Asynchronously submit reusable prepared inputs to codec-owned CUDA output.
199    ///
200    /// Recoverable execution failures remain group-local and do not suppress
201    /// later groups. No decoded host staging or final device copy is used.
202    pub fn submit_prepared(
203        &mut self,
204        prepared: &PreparedBatch,
205    ) -> Result<SubmittedCudaResidentBatch, CudaBatchError> {
206        let mut pending = Vec::new();
207        pending
208            .try_reserve_exact(prepared.groups().len())
209            .map_err(|_| BatchInfrastructureError::HostAllocationFailed {
210                what: "submitted CUDA resident groups",
211                bytes: prepared
212                    .groups()
213                    .len()
214                    .saturating_mul(core::mem::size_of::<SubmittedResidentGroup>()),
215            })?;
216        let mut group_errors = Vec::new();
217        group_errors
218            .try_reserve_exact(prepared.groups().len())
219            .map_err(|_| BatchInfrastructureError::HostAllocationFailed {
220                what: "submitted CUDA resident group failures",
221                bytes: prepared
222                    .groups()
223                    .len()
224                    .saturating_mul(core::mem::size_of::<CudaBatchGroupError>()),
225            })?;
226        for group in prepared.groups() {
227            match self.submit_resident_group(group) {
228                Ok(submitted) => pending.push(submitted),
229                Err(source) if source.session_is_unusable() => {
230                    return Err(CudaBatchError::group(group, source));
231                }
232                Err(source) => group_errors.push(CudaBatchGroupError::new(group, source)),
233            }
234        }
235        Ok(SubmittedCudaResidentBatch {
236            pending,
237            errors: {
238                let mut errors = Vec::new();
239                errors
240                    .try_reserve_exact(prepared.errors().len())
241                    .map_err(|_| BatchInfrastructureError::HostAllocationFailed {
242                        what: "submitted CUDA resident indexed errors",
243                        bytes: prepared
244                            .errors()
245                            .len()
246                            .saturating_mul(core::mem::size_of::<IndexedBatchError>()),
247                    })?;
248                errors.extend_from_slice(prepared.errors());
249                errors
250            },
251            group_errors,
252        })
253    }
254
255    fn submit_resident_group(
256        &mut self,
257        group: &PreparedBatchGroup,
258    ) -> Result<SubmittedResidentGroup, Error> {
259        let fmt = group_pixel_format(group.info())?;
260        validate_layout(group.info())?;
261        let pending = if matches!(
262            fmt,
263            PixelFormat::Rgb8
264                | PixelFormat::Rgb16
265                | PixelFormat::RgbI16
266                | PixelFormat::Rgba8
267                | PixelFormat::Rgba16
268                | PixelFormat::RgbaI16
269        ) {
270            let inputs = native_color_inputs(group)?;
271            SubmittedResidentCodecGroup::Color(
272                crate::decoder::submit_native_color_resident_prepared_batch(
273                    &inputs,
274                    &mut self.session,
275                    fmt,
276                    group.info().layout,
277                )?,
278            )
279        } else if matches!(
280            fmt,
281            PixelFormat::Gray8 | PixelFormat::Gray16 | PixelFormat::GrayI16
282        ) {
283            let inputs = resident_grayscale_inputs(group)?;
284            SubmittedResidentCodecGroup::Grayscale(
285                crate::decoder::grayscale_batch::submit_grayscale_cuda_resident_prepared_batch(
286                    &inputs,
287                    native_decode_settings(group.options().settings),
288                    &mut self.session,
289                    fmt,
290                )?,
291            )
292        } else {
293            return Err(Error::UnsupportedCudaRequest {
294                reason: "resident CUDA batch submission requires exact Gray, RGB, or RGBA output",
295            });
296        };
297        Ok(SubmittedResidentGroup {
298            metadata: ResidentGroupMetadata::from_prepared(group)?,
299            pending,
300        })
301    }
302}
303
304fn resident_grayscale_inputs(
305    group: &PreparedBatchGroup,
306) -> Result<Vec<crate::decoder::grayscale_batch::GrayscaleBatchInput<'_>>, Error> {
307    group
308        .images()
309        .iter()
310        .zip(group.source_indices().iter().copied())
311        .map(|(image, source_index)| {
312            let referenced_plan = image
313                .htj2k_plan()
314                .map(native_referenced_htj2k_plan)
315                .transpose()?;
316            let referenced_classic_plan = image
317                .classic_plan()
318                .map(native_referenced_classic_plan)
319                .transpose()?;
320            Ok(crate::decoder::grayscale_batch::GrayscaleBatchInput {
321                source_index,
322                bytes: image.bytes().as_ref(),
323                device_plan: Some(image.plan()),
324                referenced_plan,
325                referenced_classic_plan,
326            })
327        })
328        .collect()
329}
j2k_cuda/batch/resident_submission.rs

j2k_cuda/batch/
resident_submission.rs
