1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116

use std::alloc::Layout;

use crate::ptr::{non_null::NonNull, reef::Ref};

pub mod cpu;

#[cfg(feature = "cuda")]
pub mod cuda;

/// Represents a physical device that can host memory.
/// For example: [`cpu::Cpu`], [`cuda::Cuda`]
pub trait Device {
    #![allow(clippy::missing_safety_doc)]

    type Ptr<T: ?Sized>: DevicePtr<T>;
    const IS_CPU: bool = false;

    fn copy_from_host<T: Copy>(from: &[T], to: &mut Ref<[T], Self>);
    fn copy_to_host<T: Copy>(from: &Ref<[T], Self>, to: &mut [T]);
    fn copy<T: Copy>(from: &Ref<[T], Self>, to: &mut Ref<[T], Self>);
}

/// Defines the default allocator for a device.
/// For instance, The default allocator for [`cpu::Cpu`] is [`std::alloc::Global`]
pub trait DefaultDeviceAllocator: Device {
    type Alloc: DeviceAllocator<Device = Self> + Default;
}

/// Represents a type safe device-based pointer.
/// For the CPU, this will be just `*mut T`.
pub trait DevicePtr<T: ?Sized>: Copy {
    fn as_raw(self) -> *mut T;
    fn from_raw(ptr: *mut T) -> Self;

    /// # Safety
    /// Pointer must be valid and aligned
    unsafe fn write(self, val: T)
    where
        T: Sized;

    /// # Safety
    /// Offset should not overflow isize.
    /// Resulting pointer should not overflow usize.
    /// Resulting pointer must be in bounds of an allocated buffer.
    #[must_use]
    unsafe fn add(self, count: usize) -> Self
    where
        T: Sized,
    {
        Self::from_raw(self.as_raw().add(count))
    }

    /// # Safety
    /// Resulting pointer should not underflow usize.
    /// Resulting pointer must be in bounds of an allocated buffer.
    #[must_use]
    unsafe fn sub(self, count: usize) -> Self
    where
        T: Sized,
    {
        Self::from_raw(self.as_raw().sub(count))
    }

    /// # Safety
    /// Resulting pointer should not overflow usize.
    /// Resulting pointer must be in bounds of an allocated buffer.
    #[must_use]
    unsafe fn offset(self, count: isize) -> Self
    where
        T: Sized,
    {
        Self::from_raw(self.as_raw().offset(count))
    }
}

/// An Allocator in a specific device.
/// All [`std::alloc::Allocator`]s are [`DeviceAllocator<Device=cpu::Cpu>`]
pub trait DeviceAllocator {
    #![allow(clippy::missing_safety_doc)]

    /// Error returned when failing to allocate
    type AllocError: std::error::Error;

    /// Device that the allocations will live in.
    type Device: Device;

    fn allocate(&self, layout: Layout) -> Result<NonNull<[u8], Self::Device>, Self::AllocError>;

    fn allocate_zeroed(
        &self,
        layout: Layout,
    ) -> Result<NonNull<[u8], Self::Device>, Self::AllocError>;

    unsafe fn deallocate(&self, ptr: NonNull<u8, Self::Device>, layout: Layout);

    unsafe fn grow(
        &self,
        ptr: NonNull<u8, Self::Device>,
        old_layout: Layout,
        new_layout: Layout,
    ) -> Result<NonNull<[u8], Self::Device>, Self::AllocError>;

    unsafe fn grow_zeroed(
        &self,
        ptr: NonNull<u8, Self::Device>,
        old_layout: Layout,
        new_layout: Layout,
    ) -> Result<NonNull<[u8], Self::Device>, Self::AllocError>;

    unsafe fn shrink(
        &self,
        ptr: NonNull<u8, Self::Device>,
        old_layout: Layout,
        new_layout: Layout,
    ) -> Result<NonNull<[u8], Self::Device>, Self::AllocError>;
}