j2k_cuda/
surface.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3#[cfg(feature = "cuda-runtime")]
4use std::sync::Arc;
5
6use j2k_core::{
7    copy_tight_pixels_to_strided_output, BackendKind, BufferError, DeviceMemoryRange,
8    DeviceSurface, ExecutionStats, PixelFormat, SurfaceMetadata,
9};
10#[cfg(feature = "cuda-runtime")]
11use j2k_cuda_runtime::CudaDeviceBuffer;
12
13use crate::allocation::try_vec_filled;
14#[cfg(feature = "cuda-runtime")]
15use crate::allocation::try_vec_with_capacity;
16#[cfg(feature = "cuda-runtime")]
17use crate::runtime::cuda_error;
18use crate::Error;
19
20pub use j2k_core::SurfaceResidency;
21
22#[derive(Debug)]
23pub(crate) enum Storage {
24    Host(Vec<u8>),
25    #[cfg(feature = "cuda-runtime")]
26    Cuda(CudaDeviceBuffer),
27    #[cfg(feature = "cuda-runtime")]
28    CudaRange {
29        buffer: Arc<CudaDeviceBuffer>,
30        offset: usize,
31        len: usize,
32    },
33}
34
35/// CUDA surface execution counters.
36#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
37#[doc(hidden)]
38pub struct CudaSurfaceStats {
39    pub(crate) total: usize,
40    pub(crate) copy: usize,
41    pub(crate) decode: usize,
42}
43
44#[doc(hidden)]
45impl CudaSurfaceStats {
46    /// Total CUDA kernel dispatches associated with the surface.
47    pub fn kernel_dispatches(self) -> usize {
48        self.total
49    }
50
51    /// CUDA copy/upload kernel dispatches associated with the surface.
52    pub fn copy_kernel_dispatches(self) -> usize {
53        self.copy
54    }
55
56    /// CUDA codestream decode kernel dispatches associated with the surface.
57    pub fn decode_kernel_dispatches(self) -> usize {
58        self.decode
59    }
60}
61
62/// Borrowed view of a CUDA-resident surface.
63#[derive(Clone, Copy, Debug)]
64pub struct CudaSurface<'a> {
65    device_ptr: u64,
66    _marker: core::marker::PhantomData<&'a ()>,
67    pub(crate) stats: CudaSurfaceStats,
68}
69
70impl CudaSurface<'_> {
71    /// Raw CUDA device pointer value.
72    pub fn device_ptr(&self) -> u64 {
73        self.device_ptr
74    }
75
76    /// Execution counters for this surface.
77    #[doc(hidden)]
78    pub fn stats(&self) -> CudaSurfaceStats {
79        self.stats
80    }
81}
82
83/// Host- or CUDA-backed decoded surface.
84#[derive(Debug)]
85pub struct Surface {
86    pub(crate) backend: BackendKind,
87    pub(crate) residency: SurfaceResidency,
88    pub(crate) dimensions: (u32, u32),
89    pub(crate) fmt: PixelFormat,
90    pub(crate) pitch_bytes: usize,
91    pub(crate) stats: CudaSurfaceStats,
92    pub(crate) storage: Storage,
93}
94
95impl Surface {
96    fn metadata(&self) -> SurfaceMetadata {
97        SurfaceMetadata::new(
98            self.backend,
99            self.residency,
100            self.dimensions,
101            self.fmt,
102            self.pitch_bytes,
103        )
104    }
105
106    /// Return where the surface's pixels currently reside.
107    pub fn residency(&self) -> SurfaceResidency {
108        self.residency
109    }
110
111    /// Row pitch in bytes.
112    pub fn pitch_bytes(&self) -> usize {
113        self.pitch_bytes
114    }
115
116    /// Borrow host bytes when the surface is host-backed.
117    pub fn as_host_bytes(&self) -> Option<&[u8]> {
118        match &self.storage {
119            Storage::Host(bytes) => Some(bytes),
120            #[cfg(feature = "cuda-runtime")]
121            Storage::Cuda(_) | Storage::CudaRange { .. } => None,
122        }
123    }
124
125    /// Download or copy the surface into caller-owned strided output.
126    pub fn download_into(&self, out: &mut [u8], stride: usize) -> Result<(), Error> {
127        match &self.storage {
128            Storage::Host(bytes) => {
129                copy_tight_pixels_to_strided_output(bytes, self.dimensions, self.fmt, out, stride)
130                    .map_err(Error::from)
131            }
132            #[cfg(feature = "cuda-runtime")]
133            Storage::Cuda(buffer) => {
134                let byte_len = self.byte_len();
135                if let Some(len) =
136                    tight_cuda_download_len(byte_len, self.pitch_bytes, stride, out.len())
137                {
138                    return buffer.copy_to_host(&mut out[..len]).map_err(cuda_error);
139                }
140                let mut tight = try_vec_filled(byte_len, 0u8, "j2k CUDA surface download staging")?;
141                buffer.copy_to_host(&mut tight).map_err(cuda_error)?;
142                copy_tight_pixels_to_strided_output(&tight, self.dimensions, self.fmt, out, stride)
143                    .map_err(Error::from)
144            }
145            #[cfg(feature = "cuda-runtime")]
146            Storage::CudaRange {
147                buffer,
148                offset,
149                len,
150            } => {
151                let byte_len = self.byte_len();
152                if *len < byte_len {
153                    return Err(BufferError::InputTooSmall {
154                        required: byte_len,
155                        have: *len,
156                    }
157                    .into());
158                }
159                if let Some(len) =
160                    tight_cuda_download_len(byte_len, self.pitch_bytes, stride, out.len())
161                {
162                    return buffer
163                        .copy_range_to_host(*offset, &mut out[..len])
164                        .map_err(cuda_error);
165                }
166                let mut tight = try_vec_filled(byte_len, 0u8, "j2k CUDA range download staging")?;
167                buffer
168                    .copy_range_to_host(*offset, &mut tight)
169                    .map_err(cuda_error)?;
170                copy_tight_pixels_to_strided_output(&tight, self.dimensions, self.fmt, out, stride)
171                    .map_err(Error::from)
172            }
173        }
174    }
175
176    /// Borrow CUDA metadata when the surface is CUDA-backed.
177    pub fn cuda_surface(&self) -> Option<CudaSurface<'_>> {
178        #[cfg(feature = "cuda-runtime")]
179        match &self.storage {
180            Storage::Cuda(buffer) => Some(CudaSurface {
181                device_ptr: buffer.device_ptr(),
182                _marker: core::marker::PhantomData,
183                stats: self.stats,
184            }),
185            Storage::CudaRange { buffer, offset, .. } => {
186                let offset = u64::try_from(*offset).ok()?;
187                let device_ptr = buffer.device_ptr().checked_add(offset)?;
188                Some(CudaSurface {
189                    device_ptr,
190                    _marker: core::marker::PhantomData,
191                    stats: self.stats,
192                })
193            }
194            Storage::Host(_) => None,
195        }
196        #[cfg(not(feature = "cuda-runtime"))]
197        {
198            let _ = self.stats;
199            None
200        }
201    }
202
203    /// Download a sequence of surfaces into a tightly concatenated output buffer.
204    ///
205    /// CUDA surfaces produced from one contiguous batch allocation are copied
206    /// with one device-to-host transfer. Other layouts fall back to downloading
207    /// each surface tightly in order.
208    pub fn download_batch_tight(surfaces: &[Self]) -> Result<Vec<u8>, Error> {
209        let required = batch_tight_required_len(surfaces)?;
210        if required == 0 {
211            return Ok(Vec::new());
212        }
213
214        #[cfg(feature = "cuda-runtime")]
215        if let Some((buffer, offset)) = contiguous_cuda_batch_range(surfaces) {
216            let mut out = try_vec_with_capacity(required, "j2k CUDA contiguous batch download")?;
217            buffer
218                .copy_range_to_host_uninit(offset, &mut out.spare_capacity_mut()[..required])
219                .map_err(cuda_error)?;
220            // SAFETY: the CUDA copy above initialized exactly `required`
221            // bytes in this Vec's spare capacity and returned success.
222            unsafe {
223                out.set_len(required);
224            }
225            return Ok(out);
226        }
227
228        let mut out = try_vec_filled(required, 0u8, "j2k CUDA batch download")?;
229        Self::download_batch_tight_into(surfaces, &mut out)?;
230        Ok(out)
231    }
232
233    /// Download a sequence of surfaces into a tightly concatenated output buffer.
234    ///
235    /// CUDA surfaces produced from one contiguous batch allocation are copied
236    /// with one device-to-host transfer. Other layouts fall back to downloading
237    /// each surface tightly in order.
238    pub fn download_batch_tight_into(surfaces: &[Self], out: &mut [u8]) -> Result<(), Error> {
239        let required = batch_tight_required_len(surfaces)?;
240        if out.len() < required {
241            return Err(BufferError::OutputTooSmall {
242                required,
243                have: out.len(),
244            }
245            .into());
246        }
247        if required == 0 {
248            return Ok(());
249        }
250
251        #[cfg(feature = "cuda-runtime")]
252        if let Some((buffer, offset)) = contiguous_cuda_batch_range(surfaces) {
253            return buffer
254                .copy_range_to_host(offset, &mut out[..required])
255                .map_err(cuda_error);
256        }
257
258        let mut cursor = 0usize;
259        for surface in surfaces {
260            let len = surface.byte_len();
261            surface.download_into(&mut out[cursor..cursor + len], surface.pitch_bytes)?;
262            cursor += len;
263        }
264        Ok(())
265    }
266}
267
268fn batch_tight_required_len(surfaces: &[Surface]) -> Result<usize, Error> {
269    surfaces
270        .iter()
271        .try_fold(0usize, |sum, surface| sum.checked_add(surface.byte_len()))
272        .ok_or(BufferError::SizeOverflow {
273            what: "tight batch surface output",
274        })
275        .map_err(Error::from)
276}
277
278#[cfg(feature = "cuda-runtime")]
279pub(crate) fn cuda_range_storage(
280    buffer: Arc<CudaDeviceBuffer>,
281    offset: usize,
282    len: usize,
283) -> Storage {
284    Storage::CudaRange {
285        buffer,
286        offset,
287        len,
288    }
289}
290
291#[cfg(feature = "cuda-runtime")]
292fn contiguous_cuda_batch_range(surfaces: &[Surface]) -> Option<(&CudaDeviceBuffer, usize)> {
293    let first = surfaces.first()?;
294    let Storage::CudaRange {
295        buffer,
296        offset,
297        len,
298    } = &first.storage
299    else {
300        return None;
301    };
302    let first_buffer = buffer;
303    let first_offset = *offset;
304    let mut expected_offset = first_offset.checked_add(*len)?;
305    for surface in &surfaces[1..] {
306        let Storage::CudaRange {
307            buffer,
308            offset,
309            len,
310        } = &surface.storage
311        else {
312            return None;
313        };
314        if !Arc::ptr_eq(first_buffer, buffer) || *offset != expected_offset {
315            return None;
316        }
317        expected_offset = expected_offset.checked_add(*len)?;
318    }
319    Some((first_buffer.as_ref(), first_offset))
320}
321
322#[cfg(any(feature = "cuda-runtime", test))]
323fn tight_cuda_download_len(
324    byte_len: usize,
325    pitch_bytes: usize,
326    stride: usize,
327    out_len: usize,
328) -> Option<usize> {
329    (stride == pitch_bytes && out_len >= byte_len).then_some(byte_len)
330}
331
332#[doc(hidden)]
333impl DeviceSurface for Surface {
334    fn backend_kind(&self) -> BackendKind {
335        self.metadata().backend
336    }
337
338    fn residency(&self) -> j2k_core::SurfaceResidency {
339        self.metadata().residency
340    }
341
342    fn dimensions(&self) -> (u32, u32) {
343        self.metadata().dimensions
344    }
345
346    fn pixel_format(&self) -> PixelFormat {
347        self.metadata().pixel_format
348    }
349
350    fn byte_len(&self) -> usize {
351        self.metadata().byte_len()
352    }
353
354    fn execution_stats(&self) -> ExecutionStats {
355        ExecutionStats {
356            kernel_dispatches: self.stats.total as u64,
357            ..ExecutionStats::default()
358        }
359    }
360
361    fn memory_range(&self) -> Option<DeviceMemoryRange> {
362        match &self.storage {
363            Storage::Host(_) => None,
364            #[cfg(feature = "cuda-runtime")]
365            Storage::Cuda(buffer) => Some(DeviceMemoryRange::new(
366                BackendKind::Cuda,
367                buffer.device_ptr(),
368                0,
369                self.byte_len(),
370            )),
371            #[cfg(feature = "cuda-runtime")]
372            Storage::CudaRange {
373                buffer,
374                offset,
375                len,
376            } => Some(DeviceMemoryRange::new(
377                BackendKind::Cuda,
378                buffer.device_ptr(),
379                *offset,
380                *len,
381            )),
382        }
383    }
384}
385
386#[cfg(test)]
387mod tests {
388    use super::{tight_cuda_download_len, CudaSurfaceStats, Storage, Surface, SurfaceResidency};
389    use j2k_core::{BackendKind, PixelFormat};
390
391    #[test]
392    fn tight_cuda_download_len_accepts_exact_tight_output() {
393        assert_eq!(tight_cuda_download_len(32, 8, 8, 32), Some(32));
394    }
395
396    #[test]
397    fn download_batch_tight_returns_tightly_concatenated_host_surfaces() {
398        let surfaces = [
399            Surface {
400                backend: BackendKind::Cpu,
401                residency: SurfaceResidency::Host,
402                dimensions: (2, 1),
403                fmt: PixelFormat::Gray8,
404                pitch_bytes: 2,
405                stats: CudaSurfaceStats::default(),
406                storage: Storage::Host(vec![1, 2]),
407            },
408            Surface {
409                backend: BackendKind::Cpu,
410                residency: SurfaceResidency::Host,
411                dimensions: (1, 1),
412                fmt: PixelFormat::Rgb8,
413                pitch_bytes: 3,
414                stats: CudaSurfaceStats::default(),
415                storage: Storage::Host(vec![3, 4, 5]),
416            },
417        ];
418
419        let tight = Surface::download_batch_tight(&surfaces).expect("batch download");
420
421        assert_eq!(tight, vec![1, 2, 3, 4, 5]);
422    }
423}
j2k_cuda/surface.rs

j2k_cuda/
surface.rs
