tensorgraph_sys/ptr/
reef.rs

1use std::{
2    alloc::Allocator,
3    marker::PhantomData,
4    mem::MaybeUninit,
5    ops::{Deref, DerefMut, Index, IndexMut},
6};
7
8use crate::{
9    device::{cpu::Cpu, DefaultDeviceAllocator, Device, DevicePtr},
10    vec::Vec,
11};
12
13/// A reference type for devices. Should have the same representation as `&T` but
14/// is not safely transmutable. Device references can not be read directly safely since
15/// the host is not guaranteed to be on that device.
16pub struct Ref<T: ?Sized, D: Device> {
17    _device: PhantomData<D>,
18    inner: T,
19}
20impl<T: ?Sized, D: Device> Ref<T, D> {
21    /// # Safety
22    /// Refs should always point to properly allocated and initialised memory.
23    /// If any of the slice's values are not initialised or not in an allocated region,
24    /// it may result in undefined behaviour
25    pub unsafe fn from_ptr<'a>(ptr: D::Ptr<T>) -> &'a Self {
26        &*(ptr.as_raw() as *const Self)
27    }
28
29    /// # Safety
30    /// Refs should always point to properly allocated and initialised memory.
31    /// If any of the slice's values are not initialised or not in an allocated region,
32    /// it may result in undefined behaviour
33    ///
34    /// Also, you must make sure that only one mut reference is active to this pointer
35    pub unsafe fn from_ptr_mut<'a>(ptr: D::Ptr<T>) -> &'a mut Self {
36        &mut *(ptr.as_raw() as *mut Self)
37    }
38}
39
40impl<T, D: Device> Ref<[T], D> {
41    pub fn len(&self) -> usize {
42        self.inner.len()
43    }
44
45    pub fn is_empty(&self) -> bool {
46        self.inner.is_empty()
47    }
48
49    pub fn as_ptr(&self) -> D::Ptr<T> {
50        D::Ptr::from_raw(&self.inner as *const [T] as *mut T)
51    }
52
53    pub fn as_slice_ptr(&self) -> D::Ptr<[T]> {
54        D::Ptr::from_raw(&self.inner as *const [T] as *mut [T])
55    }
56}
57impl<'a, T, D: Device> Ref<[MaybeUninit<T>], D> {
58    /// # Safety
59    /// Contents must be initialised
60    pub unsafe fn assume_init(&self) -> &Ref<[T], D> {
61        &*(MaybeUninit::slice_assume_init_ref(&self.inner) as *const [T] as *const _)
62    }
63    /// # Safety
64    /// Contents must be initialised
65    pub unsafe fn assume_init_mut(&mut self) -> &mut Ref<[T], D> {
66        &mut *(MaybeUninit::slice_assume_init_mut(&mut self.inner) as *mut [T] as *mut _)
67    }
68}
69
70impl<T: Copy, D: Device> Ref<[T], D> {
71    pub fn copy_from_slice(&mut self, from: &Self) {
72        D::copy(from, self);
73    }
74
75    pub fn copy_from_host(&mut self, from: &[T]) {
76        D::copy_from_host(from, self);
77    }
78
79    pub fn copy_to_host(&self, to: &mut [T]) {
80        D::copy_to_host(self, to);
81    }
82}
83
84impl<T: Copy, D: Device> Ref<[MaybeUninit<T>], D> {
85    pub fn init_from_slice(&mut self, from: &Ref<[T], D>) {
86        unsafe {
87            let ptr = self as *mut _ as *mut Ref<[T], D>;
88            (*ptr).copy_from_slice(from);
89        }
90    }
91
92    pub fn init_from_host(&mut self, from: &[T]) {
93        unsafe {
94            let ptr = self as *mut _ as *mut Ref<[T], D>;
95            (*ptr).copy_from_host(from);
96        }
97    }
98}
99
100impl<T: Copy, D: DefaultDeviceAllocator> ToOwned for Ref<[T], D> {
101    type Owned = Vec<T, D::Alloc>;
102
103    fn to_owned(&self) -> Self::Owned {
104        unsafe {
105            let mut v = Vec::with_capacity_in(self.len(), D::Alloc::default());
106            let buf = &mut v.space_capacity_mut()[..self.len()];
107            buf.init_from_slice(self);
108            v.set_len(self.len());
109            v
110        }
111    }
112}
113
114impl<T> Deref for Ref<[T], Cpu> {
115    type Target = [T];
116
117    fn deref(&self) -> &Self::Target {
118        &self.inner
119    }
120}
121
122impl<T> DerefMut for Ref<[T], Cpu> {
123    fn deref_mut(&mut self) -> &mut Self::Target {
124        &mut self.inner
125    }
126}
127
128impl<'a, T, D: Device> AsRef<Ref<[T], D>> for &'a Ref<[T], D> {
129    fn as_ref(&self) -> &Ref<[T], D> {
130        self
131    }
132}
133
134impl<'a, T, D: Device> AsRef<Ref<[T], D>> for &'a mut Ref<[T], D> {
135    fn as_ref(&self) -> &Ref<[T], D> {
136        self
137    }
138}
139
140impl<'a, T, D: Device> AsMut<Ref<[T], D>> for &'a mut Ref<[T], D> {
141    fn as_mut(&mut self) -> &mut Ref<[T], D> {
142        self
143    }
144}
145
146impl<T, D: Device, S> Index<S> for Ref<[T], D>
147where
148    [T]: Index<S, Output = [T]>,
149{
150    type Output = Self;
151
152    fn index(&self, index: S) -> &Self::Output {
153        unsafe { &*(&self.inner[index] as *const [T] as *const Self) }
154    }
155}
156
157impl<T, D: Device, S> IndexMut<S> for Ref<[T], D>
158where
159    [T]: IndexMut<S, Output = [T]>,
160{
161    fn index_mut(&mut self, index: S) -> &mut Self::Output {
162        unsafe { &mut *(&mut self.inner[index] as *mut [T] as *mut Self) }
163    }
164}
165
166impl<T, const N: usize> AsRef<Ref<[T], Cpu>> for [T; N] {
167    fn as_ref(&self) -> &Ref<[T], Cpu> {
168        unsafe { &*(self.as_slice() as *const [T] as *const Ref<[T], Cpu>) }
169    }
170}
171
172impl<T, const N: usize> AsMut<Ref<[T], Cpu>> for [T; N] {
173    fn as_mut(&mut self) -> &mut Ref<[T], Cpu> {
174        unsafe { &mut *(self.as_mut_slice() as *mut [T] as *mut Ref<[T], Cpu>) }
175    }
176}
177
178impl<T, A: Allocator> AsRef<Ref<[T], Cpu>> for std::vec::Vec<T, A> {
179    fn as_ref(&self) -> &Ref<[T], Cpu> {
180        unsafe { &*(self.as_slice() as *const [T] as *const Ref<[T], Cpu>) }
181    }
182}
183
184impl<T, A: Allocator> AsMut<Ref<[T], Cpu>> for std::vec::Vec<T, A> {
185    fn as_mut(&mut self) -> &mut Ref<[T], Cpu> {
186        unsafe { &mut *(self.as_mut_slice() as *mut [T] as *mut Ref<[T], Cpu>) }
187    }
188}
189
190impl<T> AsRef<Ref<[T], Cpu>> for [T] {
191    fn as_ref(&self) -> &Ref<[T], Cpu> {
192        unsafe { &*(self as *const [T] as *const _) }
193    }
194}
195
196impl<T> AsMut<Ref<[T], Cpu>> for [T] {
197    fn as_mut(&mut self) -> &mut Ref<[T], Cpu> {
198        unsafe { &mut *(self as *mut [T] as *mut _) }
199    }
200}
tensorgraph_sys/ptr/reef.rs

tensorgraph_sys/ptr/
reef.rs
