prism-q 0.35.0

//! Typed device allocation wrapper with host-to-device and device-to-host copies.

use cudarc::driver::{CudaSlice, DeviceRepr, ValidAsZeroBits};

use crate::error::Result;

use super::device::GpuDevice;
use super::driver_err;

/// RAII wrapper over a typed device allocation.
pub struct GpuBuffer<T: DeviceRepr> {
    slice: CudaSlice<T>,
}

impl<T: DeviceRepr + ValidAsZeroBits> GpuBuffer<T> {
    /// Allocate `len` elements of `T` on the given device, zero-initialised.
    pub fn alloc_zeros(device: &GpuDevice, len: usize) -> Result<Self> {
        let stream = device.stream()?;
        let slice = stream
            .alloc_zeros::<T>(len)
            .map_err(|e| driver_err("alloc_zeros", e))?;
        Ok(Self { slice })
    }
}

impl<T: DeviceRepr> GpuBuffer<T> {
    /// Allocate a device buffer initialized from `host` in one transfer,
    /// skipping the zero fill of [`alloc_zeros`](Self::alloc_zeros).
    pub fn from_host(device: &GpuDevice, host: &[T]) -> Result<Self> {
        let stream = device.stream()?;
        let slice = stream
            .clone_htod(host)
            .map_err(|e| driver_err("from_host", e))?;
        Ok(Self { slice })
    }

    /// Copy `host.len()` elements from host to device (in-place).
    pub fn copy_from_host(&mut self, device: &GpuDevice, host: &[T]) -> Result<()> {
        let stream = device.stream()?;
        stream
            .memcpy_htod(host, &mut self.slice)
            .map_err(|e| driver_err("copy_from_host", e))
    }

    /// Copy all of `src` into the start of this buffer on the device.
    pub(crate) fn copy_from_device(
        &mut self,
        device: &GpuDevice,
        src: &GpuBuffer<T>,
    ) -> Result<()> {
        let stream = device.stream()?;
        stream
            .memcpy_dtod(&src.slice, &mut self.slice)
            .map_err(|e| driver_err("copy_from_device", e))
    }

    /// Copy `host.len()` elements from device into the provided host buffer.
    pub fn copy_to_host(&self, device: &GpuDevice, host: &mut [T]) -> Result<()> {
        let stream = device.stream()?;
        stream
            .memcpy_dtoh(&self.slice, host)
            .map_err(|e| driver_err("copy_to_host", e))
    }
}

impl<T: DeviceRepr> GpuBuffer<T> {
    pub fn len(&self) -> usize {
        self.slice.len()
    }

    pub(crate) fn raw(&self) -> &CudaSlice<T> {
        &self.slice
    }

    pub(crate) fn raw_mut(&mut self) -> &mut CudaSlice<T> {
        &mut self.slice
    }
}

impl<T: DeviceRepr> std::fmt::Debug for GpuBuffer<T> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("GpuBuffer")
            .field("len", &self.slice.len())
            .finish()
    }
}