jxr_cuda/
session.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3use std::sync::Arc;
4
5use jxr_core::{DecodedImage, DecodedSamples, PlaneDescriptor, StorageKind};
6
7use crate::{
8    CudaBatchDestination, CudaBatchDestinationSubmission, CudaBatchSubmission, CudaDecodePlan,
9    CudaDestination, CudaDestinationSubmission, CudaError, CudaResidentBatch,
10    CudaResidentBatchSubmission, CudaSubmission, DenseCudaBatchLayout, ResidentCudaImage,
11    runtime::{BATCH_SCRATCH_BUDGET, CudaRuntime},
12};
13
14/// Reusable CUDA runtime session with cached kernels, uploads, streams, and scratch.
15#[derive(Clone)]
16pub struct CudaDecoderSession {
17    runtime: Arc<CudaRuntime>,
18}
19
20impl core::fmt::Debug for CudaDecoderSession {
21    fn fmt(&self, formatter: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
22        formatter
23            .debug_struct("CudaDecoderSession")
24            .field("device_ordinal", &self.runtime.context.ordinal())
25            .field("stream_count", &self.runtime.streams.len())
26            .finish_non_exhaustive()
27    }
28}
29
30impl CudaDecoderSession {
31    /// Probe for both dynamically loaded CUDA driver and NVRTC libraries.
32    ///
33    /// This performs no CUDA initialization and is safe on CPU-only hosts.
34    #[must_use]
35    pub fn is_available() -> bool {
36        // SAFETY: These cudarc probes only attempt to open candidate shared
37        // libraries and do not retain or dereference foreign handles.
38        unsafe { cudarc::driver::sys::is_culib_present() && cudarc::nvrtc::sys::is_culib_present() }
39    }
40
41    /// Create a reusable session on device ordinal zero.
42    pub fn system_default() -> Result<Self, CudaError> {
43        Self::new(0)
44    }
45
46    /// Create a reusable session on a caller-selected CUDA device ordinal.
47    pub fn new(device_ordinal: usize) -> Result<Self, CudaError> {
48        if !Self::is_available() {
49            return Err(CudaError::Unavailable);
50        }
51        Ok(Self {
52            runtime: CudaRuntime::build(device_ordinal)?,
53        })
54    }
55
56    /// Whether this initialized session can accept reconstruction work.
57    #[must_use]
58    pub const fn is_usable(&self) -> bool {
59        true
60    }
61
62    /// Device ordinal retained by this session.
63    #[must_use]
64    pub fn device_ordinal(&self) -> usize {
65        self.runtime.context.ordinal()
66    }
67
68    /// Number of reusable nonblocking streams available to batch work.
69    #[must_use]
70    pub fn stream_count(&self) -> usize {
71        self.runtime.streams.len()
72    }
73
74    /// Return retained exact-size scratch allocation counters.
75    pub fn buffer_pool_diagnostics(&self) -> Result<crate::CudaBufferPoolDiagnostics, CudaError> {
76        self.runtime.buffer_pool.diagnostics()
77    }
78
79    /// Return immutable coefficient upload-cache counters.
80    pub fn upload_cache_diagnostics(&self) -> Result<crate::CudaUploadCacheDiagnostics, CudaError> {
81        self.runtime.upload_cache.diagnostics()
82    }
83
84    /// Allocate an exact device output owned by the caller.
85    pub fn allocate_destination(
86        &self,
87        layout: jxr_core::SurfaceLayout,
88    ) -> Result<CudaDestination, CudaError> {
89        let buffer = self.runtime.stream(0).alloc_zeros(layout.byte_len)?;
90        CudaDestination::from_device_buffer(buffer, layout)
91    }
92
93    /// Allocate a dense device batch output owned by the caller.
94    pub fn allocate_batch_destination(
95        &self,
96        layout: DenseCudaBatchLayout,
97    ) -> Result<CudaBatchDestination, CudaError> {
98        let buffer = self.runtime.stream(0).alloc_zeros(layout.byte_len())?;
99        CudaBatchDestination::from_device_buffer(buffer, layout)
100    }
101
102    /// Submit one validated plan without waiting for completion.
103    pub fn submit(&self, plan: &CudaDecodePlan) -> Result<CudaSubmission, CudaError> {
104        let encoded = crate::encode::encode_owned(&self.runtime, plan, 0)?;
105        Ok(CudaSubmission::submitted(
106            encoded,
107            plan.decode_report(false),
108        ))
109    }
110
111    /// Submit a bounded batch across the session's reusable streams.
112    pub fn submit_batch(&self, plans: &[CudaDecodePlan]) -> Result<CudaBatchSubmission, CudaError> {
113        check_batch_scratch(plans)?;
114        let mut submissions = Vec::with_capacity(plans.len());
115        for (index, plan) in plans.iter().enumerate() {
116            let encoded = crate::encode::encode_owned(&self.runtime, plan, index)?;
117            submissions.push(CudaSubmission::submitted(
118                encoded,
119                plan.decode_report(false),
120            ));
121        }
122        Ok(CudaBatchSubmission::new(submissions))
123    }
124
125    /// Submit a homogeneous batch into one internally owned device allocation.
126    pub fn submit_dense_batch(
127        &self,
128        plans: &[CudaDecodePlan],
129    ) -> Result<CudaResidentBatchSubmission, CudaError> {
130        let first = plans.first().ok_or(CudaError::InvalidPlan {
131            reason: "dense CUDA batch cannot be empty",
132        })?;
133        if plans.iter().any(|plan| plan.output() != first.output()) {
134            return Err(CudaError::InvalidPlan {
135                reason: "dense CUDA batch plans are not homogeneous",
136            });
137        }
138        check_batch_scratch(plans)?;
139        let layout = DenseCudaBatchLayout::new(first.output().clone(), plans.len())?;
140        let stream = self.runtime.stream(0);
141        let mut output = stream.alloc_zeros::<u8>(layout.byte_len())?;
142        let mut submissions = Vec::with_capacity(plans.len());
143        let mut reports = Vec::with_capacity(plans.len());
144        for (image, plan) in plans.iter().enumerate() {
145            let encoded = crate::encode::encode_into(
146                &self.runtime,
147                plan,
148                &mut output,
149                layout.image_offset(image)?,
150                0,
151            )?;
152            let report = plan.decode_report(false);
153            submissions.push(CudaSubmission::submitted(encoded, report.clone()));
154            reports.push(report);
155        }
156        Ok(CudaResidentBatchSubmission::new(
157            submissions,
158            output,
159            layout,
160            reports,
161        ))
162    }
163
164    /// Submit into an exclusively retained caller-owned device allocation.
165    pub fn submit_into(
166        &self,
167        plan: &CudaDecodePlan,
168        mut destination: CudaDestination,
169    ) -> Result<CudaDestinationSubmission, CudaError> {
170        destination.validate_plan(plan)?;
171        destination.validate_context(&self.runtime.context)?;
172        let encoded =
173            crate::encode::encode_into(&self.runtime, plan, destination.buffer_mut(), 0, 0)?;
174        let report = plan.decode_report(false);
175        Ok(CudaDestinationSubmission::new(
176            CudaSubmission::submitted(encoded, report.clone()),
177            destination,
178            report,
179        ))
180    }
181
182    /// Submit a homogeneous batch into caller-owned dense device storage.
183    pub fn submit_batch_into(
184        &self,
185        plans: &[CudaDecodePlan],
186        mut destination: CudaBatchDestination,
187    ) -> Result<CudaBatchDestinationSubmission, CudaError> {
188        destination.validate_plans(plans)?;
189        destination.validate_context(&self.runtime.context)?;
190        check_batch_scratch(plans)?;
191        let mut submissions = Vec::with_capacity(plans.len());
192        let mut reports = Vec::with_capacity(plans.len());
193        for (image, plan) in plans.iter().enumerate() {
194            let offset = destination.layout().image_offset(image)?;
195            let encoded = crate::encode::encode_into(
196                &self.runtime,
197                plan,
198                destination.buffer_mut(),
199                offset,
200                0,
201            )?;
202            let report = plan.decode_report(false);
203            submissions.push(CudaSubmission::submitted(encoded, report.clone()));
204            reports.push(report);
205        }
206        Ok(CudaBatchDestinationSubmission::new(
207            submissions,
208            destination,
209            reports,
210        ))
211    }
212
213    /// Decode one plan and copy its exact native bytes to typed host storage.
214    pub fn decode_to_host(&self, plan: &CudaDecodePlan) -> Result<DecodedImage, CudaError> {
215        let resident = self.submit(plan)?.wait()?;
216        let bytes = self.readback(&resident)?;
217        decoded_image(plan, bytes)
218    }
219
220    /// Decode a batch to host storage while preserving caller order.
221    pub fn decode_batch_to_host(
222        &self,
223        plans: &[CudaDecodePlan],
224    ) -> Result<Vec<DecodedImage>, CudaError> {
225        self.submit_batch(plans)?
226            .wait()?
227            .into_iter()
228            .zip(plans)
229            .map(|(resident, plan)| {
230                let bytes = self.readback(&resident)?;
231                decoded_image(plan, bytes)
232            })
233            .collect()
234    }
235
236    /// Copy one completed resident image to host bytes.
237    pub fn readback(&self, image: &ResidentCudaImage) -> Result<Vec<u8>, CudaError> {
238        validate_image_context(image, &self.runtime)?;
239        let stream = image.buffer.stream();
240        let bytes = stream.clone_dtoh(&image.buffer)?;
241        stream.synchronize()?;
242        Ok(bytes)
243    }
244
245    /// Copy one image from a completed dense batch to host bytes.
246    pub fn readback_batch_image(
247        &self,
248        batch: &CudaResidentBatch,
249        image: usize,
250    ) -> Result<Vec<u8>, CudaError> {
251        if batch.buffer.context() != &self.runtime.context {
252            return Err(CudaError::InvalidDestination {
253                reason: "resident CUDA batch belongs to a different context",
254            });
255        }
256        let start = batch.layout().image_offset(image)?;
257        let end = start
258            .checked_add(batch.layout().image_stride_bytes())
259            .ok_or(CudaError::InvalidDestination {
260                reason: "resident CUDA batch image range overflows usize",
261            })?;
262        let view = batch.buffer.slice(start..end);
263        let stream = batch.buffer.stream();
264        let bytes = stream.clone_dtoh(&view)?;
265        stream.synchronize()?;
266        Ok(bytes)
267    }
268}
269
270fn validate_image_context(
271    image: &ResidentCudaImage,
272    runtime: &CudaRuntime,
273) -> Result<(), CudaError> {
274    if image.buffer.context() == &runtime.context {
275        Ok(())
276    } else {
277        Err(CudaError::InvalidDestination {
278            reason: "resident CUDA image belongs to a different context",
279        })
280    }
281}
282
283fn check_batch_scratch(plans: &[CudaDecodePlan]) -> Result<(), CudaError> {
284    let total = plans.iter().try_fold(0_usize, |total, plan| {
285        total
286            .checked_add(plan.scratch_bytes()?)
287            .ok_or(CudaError::ResourceLimit {
288                reason: "aggregate CUDA batch scratch overflows usize",
289                requested: usize::MAX,
290                maximum: BATCH_SCRATCH_BUDGET,
291            })
292    })?;
293    if total > BATCH_SCRATCH_BUDGET {
294        return Err(CudaError::ResourceLimit {
295            reason: "aggregate CUDA batch scratch exceeds the bounded session budget",
296            requested: total,
297            maximum: BATCH_SCRATCH_BUDGET,
298        });
299    }
300    Ok(())
301}
302
303fn decoded_image(plan: &CudaDecodePlan, bytes: Vec<u8>) -> Result<DecodedImage, CudaError> {
304    let info = plan.info().ok_or(CudaError::InvalidPlan {
305        reason: "metadata-only plan cannot produce a decoded image",
306    })?;
307    let decoded_region = plan.output_region().ok_or(CudaError::InvalidPlan {
308        reason: "metadata-only plan omits its output region",
309    })?;
310    let samples = decoded_samples(plan.output().format, bytes)?;
311    let decoded = DecodedImage {
312        info: info.clone(),
313        decoded_region,
314        format: plan.output().format,
315        planes: plan
316            .output()
317            .planes
318            .iter()
319            .map(|plane| PlaneDescriptor {
320                byte_offset: plane.byte_offset,
321                row_stride_bytes: plane.row_stride_bytes,
322                width: plane.width,
323                height: plane.height,
324                channels: plane.channels,
325            })
326            .collect(),
327        samples,
328        report: plan.decode_report(true),
329    };
330    decoded
331        .validate_layout()
332        .map_err(|_| CudaError::InvalidPlan {
333            reason: "CUDA host result violates the decoded-image layout",
334        })?;
335    Ok(decoded)
336}
337
338fn decoded_samples(
339    format: jxr_core::PixelFormat,
340    bytes: Vec<u8>,
341) -> Result<DecodedSamples, CudaError> {
342    match format.storage_kind() {
343        StorageKind::BitPacked => Ok(DecodedSamples::BitPacked(bytes)),
344        StorageKind::U8 => Ok(DecodedSamples::U8(bytes)),
345        StorageKind::U16 => Ok(DecodedSamples::U16(read_u16(&bytes)?)),
346        StorageKind::I16 => Ok(DecodedSamples::I16(
347            read_u16(&bytes)?
348                .into_iter()
349                .map(|value| i16::from_ne_bytes(value.to_ne_bytes()))
350                .collect(),
351        )),
352        StorageKind::I32 => Ok(DecodedSamples::I32(
353            read_u32(&bytes)?
354                .into_iter()
355                .map(|value| i32::from_ne_bytes(value.to_ne_bytes()))
356                .collect(),
357        )),
358        StorageKind::F16Bits => Ok(DecodedSamples::F16(read_u16(&bytes)?)),
359        StorageKind::F32 => Ok(DecodedSamples::F32(
360            read_u32(&bytes)?.into_iter().map(f32::from_bits).collect(),
361        )),
362        StorageKind::PackedU16 => match format {
363            jxr_core::PixelFormat::Rgb555 => Ok(DecodedSamples::Rgb555(read_u16(&bytes)?)),
364            jxr_core::PixelFormat::Rgb565 => Ok(DecodedSamples::Rgb565(read_u16(&bytes)?)),
365            _ => Err(CudaError::InvalidPlan {
366                reason: "unknown packed 16-bit output",
367            }),
368        },
369        StorageKind::PackedU32 => match format {
370            jxr_core::PixelFormat::Rgb101010 => Ok(DecodedSamples::Rgb101010(read_u32(&bytes)?)),
371            jxr_core::PixelFormat::Rgbe => Ok(DecodedSamples::Rgbe(read_u32(&bytes)?)),
372            _ => Err(CudaError::InvalidPlan {
373                reason: "unknown packed 32-bit output",
374            }),
375        },
376    }
377}
378
379fn read_u16(bytes: &[u8]) -> Result<Vec<u16>, CudaError> {
380    let (chunks, remainder) = bytes.as_chunks::<2>();
381    if !remainder.is_empty() {
382        return Err(CudaError::InvalidPlan {
383            reason: "16-bit CUDA host readback has a trailing byte",
384        });
385    }
386    Ok(chunks
387        .iter()
388        .map(|chunk| u16::from_ne_bytes(*chunk))
389        .collect())
390}
391
392fn read_u32(bytes: &[u8]) -> Result<Vec<u32>, CudaError> {
393    let (chunks, remainder) = bytes.as_chunks::<4>();
394    if !remainder.is_empty() {
395        return Err(CudaError::InvalidPlan {
396            reason: "32-bit CUDA host readback has trailing bytes",
397        });
398    }
399    Ok(chunks
400        .iter()
401        .map(|chunk| u32::from_ne_bytes(*chunk))
402        .collect())
403}
jxr_cuda/session.rs

jxr_cuda/
session.rs
