j2k_cuda_runtime/
memory.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3mod pinned_staging;
4mod pool;
5mod ranges;
6
7pub(crate) use self::pinned_staging::PinnedUploadStagingPool;
8pub use self::pinned_staging::{
9    CudaPinnedUploadOperationGuard, CudaPinnedUploadStagingCheckout,
10    CudaPinnedUploadStagingPoolDiagnostics, CudaPinnedUploadStagingPoolLimits,
11};
12#[cfg(test)]
13pub(crate) use self::pool::copy_pooled_bytes_to_vec_uninit;
14#[cfg(test)]
15pub(crate) use self::pool::pool_fit_buffer_index_by_len;
16pub(crate) use self::pool::{
17    copy_pooled_bytes_to_vec_uninit_with_budget, pooled_device_buffer, CudaBufferPoolReuseGuard,
18};
19pub use self::pool::{
20    CudaBufferPool, CudaBufferPoolDiagnostics, CudaBufferPoolLimits, CudaBufferPoolTakeTrace,
21    CudaPooledDeviceBuffer,
22};
23pub(crate) use self::ranges::CheckedDeviceBufferRanges;
24
25#[cfg(test)]
26use crate::context::validate_non_null_pinned_host_allocation;
27use crate::{
28    bytes::f32_slice_as_bytes, context::CudaContext, driver::CuDevicePtr, error::CudaError,
29};
30use std::ffi::c_void;
31
32impl CudaContext {
33    /// Upload host bytes into a CUDA device buffer.
34    pub fn upload(&self, bytes: &[u8]) -> Result<CudaDeviceBuffer, CudaError> {
35        let mut ptr = 0;
36        let buffer = if bytes.is_empty() {
37            self.inner.set_current()?;
38            CudaDeviceBuffer {
39                context: self.clone(),
40                ptr,
41                len: bytes.len(),
42            }
43        } else {
44            self.inner.with_current_stateful_operation(|| {
45                // SAFETY: CUDA writes a device pointer for the requested byte
46                // size while this context's lifecycle gate is held.
47                self.inner.driver.check("cuMemAlloc_v2", unsafe {
48                    (self.inner.driver.cu_mem_alloc)(&raw mut ptr, bytes.len())
49                })?;
50                crate::context::validate_device_allocation(ptr, bytes.len())
51            })?;
52
53            CudaDeviceBuffer {
54                context: self.clone(),
55                ptr,
56                len: bytes.len(),
57            }
58        };
59
60        if !bytes.is_empty() {
61            self.inner.with_current_resource_operation(|| {
62                // SAFETY: ptr is a valid device allocation of bytes.len(), the
63                // host pointer covers that length, and the lifecycle gate is held.
64                self.inner.driver.check("cuMemcpyHtoD_v2", unsafe {
65                    (self.inner.driver.cu_memcpy_htod)(
66                        ptr,
67                        bytes.as_ptr().cast::<c_void>(),
68                        bytes.len(),
69                    )
70                })
71            })?;
72        }
73
74        Ok(buffer)
75    }
76
77    /// Upload host `f32` samples into a CUDA device buffer.
78    pub fn upload_f32(&self, samples: &[f32]) -> Result<CudaDeviceBuffer, CudaError> {
79        self.upload(f32_slice_as_bytes(samples))
80    }
81
82    /// Allocate an uninitialized CUDA device buffer.
83    pub fn allocate(&self, len: usize) -> Result<CudaDeviceBuffer, CudaError> {
84        let mut ptr = 0;
85        if len != 0 {
86            self.inner.with_current_stateful_operation(|| {
87                // SAFETY: CUDA writes a device pointer for the requested byte
88                // size while this context's lifecycle gate is held.
89                self.inner.driver.check("cuMemAlloc_v2", unsafe {
90                    (self.inner.driver.cu_mem_alloc)(&raw mut ptr, len)
91                })?;
92                crate::context::validate_device_allocation(ptr, len)
93            })?;
94        } else {
95            self.inner.set_current()?;
96        }
97        Ok(CudaDeviceBuffer {
98            context: self.clone(),
99            ptr,
100            len,
101        })
102    }
103
104    /// Allocate page-locked host memory for host-to-device staging.
105    #[cfg(test)]
106    pub(crate) fn pinned_host_buffer(&self, len: usize) -> Result<CudaPinnedHostBuffer, CudaError> {
107        let mut ptr = std::ptr::null_mut();
108        if len != 0 {
109            self.inner.with_current_stateful_operation(|| {
110                // SAFETY: CUDA writes a page-locked host pointer for the requested
111                // byte length. The allocation is freed by CudaPinnedHostBuffer.
112                self.inner.driver.check("cuMemHostAlloc", unsafe {
113                    (self.inner.driver.cu_mem_host_alloc)(&raw mut ptr, len, 0)
114                })?;
115                validate_non_null_pinned_host_allocation(ptr.cast::<u8>(), len).map(|_| ())
116            })?;
117        } else {
118            self.inner.set_current()?;
119        }
120        Ok(CudaPinnedHostBuffer {
121            context: self.clone(),
122            ptr: ptr.cast::<u8>(),
123            len,
124        })
125    }
126
127    /// Create a reusable device-buffer pool for this context.
128    pub fn buffer_pool(&self) -> CudaBufferPool {
129        CudaBufferPool::new(self.clone())
130    }
131
132    /// Create a reusable best-fit device-buffer pool for workloads with many
133    /// same-sized intermediate buffers.
134    pub fn best_fit_buffer_pool(&self) -> CudaBufferPool {
135        CudaBufferPool::new_size_buckets(self.clone())
136    }
137}
138
139/// Page-locked host staging buffer.
140#[cfg(test)]
141#[derive(Debug)]
142pub(crate) struct CudaPinnedHostBuffer {
143    pub(crate) context: CudaContext,
144    pub(crate) ptr: *mut u8,
145    pub(crate) len: usize,
146}
147
148#[cfg(test)]
149impl CudaPinnedHostBuffer {
150    /// Immutable byte view of the pinned allocation.
151    pub(crate) fn as_slice(&self) -> &[u8] {
152        if self.len == 0 {
153            &[]
154        } else {
155            // SAFETY: ptr is a live pinned allocation of len bytes.
156            unsafe { std::slice::from_raw_parts(self.ptr.cast_const(), self.len) }
157        }
158    }
159
160    /// Mutable byte view of the pinned allocation.
161    pub(crate) fn as_mut_slice(&mut self) -> &mut [u8] {
162        if self.len == 0 {
163            &mut []
164        } else {
165            // SAFETY: ptr is uniquely borrowed through &mut self and covers len
166            // bytes allocated by CUDA.
167            unsafe { std::slice::from_raw_parts_mut(self.ptr, self.len) }
168        }
169    }
170}
171
172#[cfg(test)]
173impl Drop for CudaPinnedHostBuffer {
174    fn drop(&mut self) {
175        if !self.ptr.is_null() {
176            let free_result = self.context.inner.with_current_stateful_operation(|| {
177                // SAFETY: ptr was returned by cuMemHostAlloc for this process,
178                // and the context lifecycle gate is held during destruction.
179                self.context.inner.driver.check("cuMemFreeHost", unsafe {
180                    (self.context.inner.driver.cu_mem_free_host)(self.ptr.cast())
181                })
182            });
183            if free_result.is_err() {
184                std::mem::forget(self.context.clone());
185            }
186        }
187    }
188}
189
190// SAFETY: The pinned allocation is owned by this value and CUDA frees it on
191// drop. Mutable access still requires &mut self.
192#[cfg(test)]
193unsafe impl Send for CudaPinnedHostBuffer {}
194
195/// Owned CUDA device buffer.
196#[derive(Debug)]
197pub struct CudaDeviceBuffer {
198    pub(crate) context: CudaContext,
199    pub(crate) ptr: CuDevicePtr,
200    pub(crate) len: usize,
201}
202
203#[doc(hidden)]
204/// Typed immutable device buffer view.
205#[derive(Clone, Copy, Debug)]
206pub struct CudaDeviceBufferView<'a, T> {
207    pub(crate) ptr: CuDevicePtr,
208    pub(crate) len: usize,
209    pub(crate) _marker: std::marker::PhantomData<&'a T>,
210}
211
212impl<T> CudaDeviceBufferView<'_, T> {
213    /// Raw CUDA device pointer value for kernel argument binding.
214    pub fn device_ptr(&self) -> u64 {
215        self.ptr
216    }
217
218    /// Number of typed elements in this view.
219    pub fn len(&self) -> usize {
220        self.len
221    }
222
223    /// Whether this view has no elements.
224    pub fn is_empty(&self) -> bool {
225        self.len == 0
226    }
227}
228
229#[doc(hidden)]
230/// Typed mutable device buffer view.
231#[derive(Debug)]
232pub struct CudaDeviceBufferViewMut<'a, T> {
233    pub(crate) ptr: CuDevicePtr,
234    pub(crate) len: usize,
235    pub(crate) _marker: std::marker::PhantomData<&'a mut T>,
236}
237
238impl<T> CudaDeviceBufferViewMut<'_, T> {
239    /// Raw CUDA device pointer value for kernel argument binding.
240    pub fn device_ptr(&self) -> u64 {
241        self.ptr
242    }
243
244    /// Number of typed elements in this view.
245    pub fn len(&self) -> usize {
246        self.len
247    }
248
249    /// Whether this view has no elements.
250    pub fn is_empty(&self) -> bool {
251        self.len == 0
252    }
253}
254
255#[doc(hidden)]
256/// One byte range inside a contiguous CUDA batch output allocation.
257#[derive(Clone, Copy, Debug, Eq, PartialEq)]
258pub struct CudaDeviceBufferRange {
259    /// Byte offset from the start of the contiguous allocation.
260    pub offset: usize,
261    /// Byte length for this output item.
262    pub len: usize,
263}
264
265impl CudaDeviceBuffer {
266    pub(crate) fn is_owned_by(&self, context: &CudaContext) -> bool {
267        self.context.is_same_context(context)
268    }
269
270    /// CUDA context that owns this allocation.
271    pub fn context(&self) -> CudaContext {
272        self.context.clone()
273    }
274
275    /// Raw CUDA device pointer value.
276    pub fn device_ptr(&self) -> u64 {
277        self.ptr
278    }
279
280    /// Device allocation length in bytes.
281    pub fn byte_len(&self) -> usize {
282        self.len
283    }
284
285    #[doc(hidden)]
286    /// Borrow this allocation as a typed immutable device view.
287    pub fn typed_view<T>(&self) -> Result<CudaDeviceBufferView<'_, T>, CudaError> {
288        let element_size = std::mem::size_of::<T>();
289        if element_size == 0 || !self.len.is_multiple_of(element_size) {
290            return Err(CudaError::LengthNotElementAligned {
291                bytes: self.len,
292                element_size,
293            });
294        }
295        Ok(CudaDeviceBufferView {
296            ptr: self.ptr,
297            len: self.len / element_size,
298            _marker: std::marker::PhantomData,
299        })
300    }
301
302    #[doc(hidden)]
303    /// Borrow this allocation as a typed mutable device view.
304    pub fn typed_view_mut<T>(&mut self) -> Result<CudaDeviceBufferViewMut<'_, T>, CudaError> {
305        let element_size = std::mem::size_of::<T>();
306        if element_size == 0 || !self.len.is_multiple_of(element_size) {
307            return Err(CudaError::LengthNotElementAligned {
308                bytes: self.len,
309                element_size,
310            });
311        }
312        Ok(CudaDeviceBufferViewMut {
313            ptr: self.ptr,
314            len: self.len / element_size,
315            _marker: std::marker::PhantomData,
316        })
317    }
318
319    /// Copy device bytes into caller-owned host output.
320    pub fn copy_to_host(&self, out: &mut [u8]) -> Result<(), CudaError> {
321        if out.len() < self.len {
322            return Err(CudaError::OutputTooSmall {
323                required: self.len,
324                have: out.len(),
325            });
326        }
327        if self.len == 0 {
328            return Ok(());
329        }
330
331        self.context.inner.with_current_resource_operation(|| {
332            // SAFETY: ptr is a live device allocation of self.len bytes, out
333            // covers that range, and the context lifecycle gate is held.
334            self.context.inner.driver.check("cuMemcpyDtoH_v2", unsafe {
335                (self.context.inner.driver.cu_memcpy_dtoh)(
336                    out.as_mut_ptr().cast::<c_void>(),
337                    self.ptr,
338                    self.len,
339                )
340            })
341        })
342    }
343
344    /// Copy a byte range from this device buffer into caller-owned host output.
345    pub fn copy_range_to_host(&self, offset: usize, out: &mut [u8]) -> Result<(), CudaError> {
346        self.copy_byte_range_to_host_elements(offset, out)
347    }
348
349    /// Copy a byte range from this device buffer into uninitialized host output.
350    pub fn copy_range_to_host_uninit(
351        &self,
352        offset: usize,
353        out: &mut [std::mem::MaybeUninit<u8>],
354    ) -> Result<(), CudaError> {
355        self.copy_byte_range_to_host_elements(offset, out)
356    }
357
358    fn copy_byte_range_to_host_elements<T>(
359        &self,
360        offset: usize,
361        out: &mut [T],
362    ) -> Result<(), CudaError> {
363        let byte_len = out
364            .len()
365            .checked_mul(std::mem::size_of::<T>())
366            .ok_or(CudaError::LengthTooLarge { len: usize::MAX })?;
367        let end = offset
368            .checked_add(byte_len)
369            .ok_or(CudaError::LengthTooLarge { len: usize::MAX })?;
370        if end > self.len {
371            return Err(CudaError::OutputTooSmall {
372                required: end,
373                have: self.len,
374            });
375        }
376        if byte_len == 0 {
377            return Ok(());
378        }
379
380        let source = self
381            .ptr
382            .checked_add(
383                u64::try_from(offset).map_err(|_| CudaError::LengthTooLarge { len: offset })?,
384            )
385            .ok_or(CudaError::LengthTooLarge { len: usize::MAX })?;
386        self.context.inner.with_current_resource_operation(|| {
387            // SAFETY: `source` is inside this live device allocation, `out`
388            // covers exactly `byte_len` bytes, and the lifecycle gate is held.
389            self.context.inner.driver.check("cuMemcpyDtoH_v2", unsafe {
390                (self.context.inner.driver.cu_memcpy_dtoh)(
391                    out.as_mut_ptr().cast::<c_void>(),
392                    source,
393                    byte_len,
394                )
395            })
396        })
397    }
398}
399
400impl Drop for CudaDeviceBuffer {
401    fn drop(&mut self) {
402        if self.ptr != 0 {
403            let free_result = self.context.inner.with_current_stateful_operation(|| {
404                // SAFETY: ptr was allocated by this CUDA context. The context
405                // lifetime gate is held while the allocation is destroyed.
406                let status = unsafe { (self.context.inner.driver.cu_mem_free)(self.ptr) };
407                self.context.inner.driver.check("cuMemFree_v2", status)
408            });
409            if free_result.is_err() {
410                // Retain the context so neither this allocation nor any
411                // potentially in-flight work is torn down after completion
412                // became uncertain.
413                std::mem::forget(self.context.clone());
414            }
415        }
416    }
417}
418
419pub(crate) fn checked_image_words(
420    width: u32,
421    height: u32,
422    channels: usize,
423) -> Result<usize, CudaError> {
424    width
425        .try_into()
426        .ok()
427        .and_then(|width: usize| width.checked_mul(height as usize))
428        .and_then(|pixels| pixels.checked_mul(channels))
429        .ok_or(CudaError::ImageTooLarge {
430            width,
431            height,
432            channels,
433        })
434}
j2k_cuda_runtime/memory.rs

j2k_cuda_runtime/
memory.rs
