#![cfg(feature = "cuda")]
use core::ffi::c_int;
use core::marker::PhantomData;
use core::mem::size_of;
use core::ops::{Deref, DerefMut};
use crate::ffi;
use crate::GpuError;
pub struct PinnedHostBuf<T: Copy + Default> {
ptr: *mut T,
len: usize,
_marker: PhantomData<T>,
}
impl<T: Copy + Default> PinnedHostBuf<T> {
pub fn new(len: usize) -> Result<Self, GpuError> {
if len == 0 {
return Ok(Self {
ptr: core::ptr::null_mut(),
len: 0,
_marker: PhantomData,
});
}
let size = (len as u64) * (size_of::<T>() as u64);
let mut raw: *mut u8 = core::ptr::null_mut();
#[allow(unsafe_code)]
let status: c_int = unsafe { ffi::dsfb_gpu_alloc_pinned_bytes(size, &mut raw) };
if status != 0 || raw.is_null() {
return Err(GpuError::KernelFailed(status));
}
let ptr = raw.cast::<T>();
#[allow(unsafe_code)]
unsafe {
for i in 0..len {
ptr.add(i).write(T::default());
}
}
Ok(Self {
ptr,
len,
_marker: PhantomData,
})
}
#[must_use]
pub fn len(&self) -> usize {
self.len
}
#[must_use]
pub fn is_empty(&self) -> bool {
self.len == 0
}
#[must_use]
pub fn as_ptr(&self) -> *const T {
self.ptr
}
#[must_use]
pub fn as_mut_ptr(&mut self) -> *mut T {
self.ptr
}
#[must_use]
pub fn as_slice(&self) -> &[T] {
if self.len == 0 {
&[]
} else {
#[allow(unsafe_code)]
unsafe {
core::slice::from_raw_parts(self.ptr, self.len)
}
}
}
#[must_use]
pub fn as_mut_slice(&mut self) -> &mut [T] {
if self.len == 0 {
&mut []
} else {
#[allow(unsafe_code)]
unsafe {
core::slice::from_raw_parts_mut(self.ptr, self.len)
}
}
}
}
impl<T: Copy + Default> Deref for PinnedHostBuf<T> {
type Target = [T];
fn deref(&self) -> &[T] {
self.as_slice()
}
}
impl<T: Copy + Default> DerefMut for PinnedHostBuf<T> {
fn deref_mut(&mut self) -> &mut [T] {
self.as_mut_slice()
}
}
impl<T: Copy + Default> Drop for PinnedHostBuf<T> {
fn drop(&mut self) {
if !self.ptr.is_null() {
#[allow(unsafe_code)]
unsafe {
let _ = ffi::dsfb_gpu_free_pinned_bytes(self.ptr.cast::<u8>());
}
self.ptr = core::ptr::null_mut();
self.len = 0;
}
}
}
#[cfg(test)]
#[allow(clippy::unwrap_used, clippy::expect_used)]
mod tests {
use super::*;
#[test]
fn pinned_host_buf_round_trip_u32() {
let mut buf: PinnedHostBuf<u32> = PinnedHostBuf::new(64).expect("allocate pinned");
assert_eq!(buf.len(), 64);
assert!(!buf.is_empty());
for (i, slot) in buf.as_mut_slice().iter_mut().enumerate() {
*slot = (i as u32) * 7 + 13;
}
for (i, value) in buf.as_slice().iter().enumerate() {
assert_eq!(*value, (i as u32) * 7 + 13);
}
}
#[test]
fn pinned_host_buf_zero_length_is_empty() {
let buf: PinnedHostBuf<u32> = PinnedHostBuf::new(0).expect("zero-length allocate");
assert_eq!(buf.len(), 0);
assert!(buf.is_empty());
assert!(buf.as_ptr().is_null());
assert_eq!(buf.as_slice(), &[] as &[u32]);
}
}