Skip to main content

singe_cusparse/
scalar.rs

1use std::{marker::PhantomData, ptr};
2
3use singe_cuda::types::DevicePtr;
4
5use crate::types::PointerMode;
6
7#[derive(Debug, Clone, Copy)]
8#[non_exhaustive]
9pub enum Scalar<'a, T> {
10    Host(&'a T),
11    Device(DeviceScalar<T>),
12}
13
14#[derive(Debug)]
15#[non_exhaustive]
16pub enum ScalarMut<'a, T> {
17    Host(&'a mut T),
18    Device(DeviceScalarMut<T>),
19}
20
21#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
22pub struct DeviceScalar<T> {
23    ptr: DevicePtr,
24    _t: PhantomData<*const T>,
25}
26
27#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
28pub struct DeviceScalarMut<T> {
29    ptr: DevicePtr,
30    _t: PhantomData<*mut T>,
31}
32
33impl<'a, T> Scalar<'a, T> {
34    pub const fn host(value: &'a T) -> Self {
35        Self::Host(value)
36    }
37
38    /// Creates a device-resident scalar from a raw device pointer.
39    ///
40    /// # Safety
41    ///
42    /// `ptr` must point to a valid device allocation containing a `T` for every
43    /// cuSPARSE operation using this scalar.
44    pub const unsafe fn device(ptr: DevicePtr) -> Self {
45        Self::Device(unsafe { DeviceScalar::new(ptr) })
46    }
47
48    pub(crate) const fn pointer_mode(self) -> PointerMode {
49        match self {
50            Self::Host(_) => PointerMode::Host,
51            Self::Device(_) => PointerMode::Device,
52        }
53    }
54
55    pub(crate) const fn ptr(self) -> *const T {
56        match self {
57            Self::Host(value) => ptr::from_ref(value),
58            Self::Device(value) => value.as_ptr(),
59        }
60    }
61}
62
63impl<'a, T> ScalarMut<'a, T> {
64    pub const fn host(value: &'a mut T) -> Self {
65        Self::Host(value)
66    }
67
68    /// Creates a mutable device-resident scalar from a raw device pointer.
69    ///
70    /// # Safety
71    ///
72    /// `ptr` must point to a valid writable device allocation containing a `T`
73    /// for every cuSPARSE operation using this scalar.
74    pub const unsafe fn device(ptr: DevicePtr) -> Self {
75        Self::Device(unsafe { DeviceScalarMut::new(ptr) })
76    }
77
78    pub(crate) const fn pointer_mode(&self) -> PointerMode {
79        match self {
80            Self::Host(_) => PointerMode::Host,
81            Self::Device(_) => PointerMode::Device,
82        }
83    }
84
85    pub(crate) fn as_mut_ptr(&mut self) -> *mut T {
86        match self {
87            Self::Host(value) => ptr::from_mut(*value),
88            Self::Device(value) => value.as_mut_ptr(),
89        }
90    }
91}
92
93impl<'a, T> From<&'a T> for Scalar<'a, T> {
94    fn from(value: &'a T) -> Self {
95        Self::host(value)
96    }
97}
98
99impl<T> From<DeviceScalar<T>> for Scalar<'_, T> {
100    fn from(value: DeviceScalar<T>) -> Self {
101        Self::Device(value)
102    }
103}
104
105impl<T> DeviceScalar<T> {
106    /// Creates a device scalar from a raw device pointer.
107    ///
108    /// # Safety
109    ///
110    /// `ptr` must point to a valid device allocation containing a `T` for every
111    /// cuSPARSE operation using this scalar.
112    pub const unsafe fn new(ptr: DevicePtr) -> Self {
113        Self {
114            ptr,
115            _t: PhantomData,
116        }
117    }
118
119    pub const fn as_device_ptr(self) -> DevicePtr {
120        self.ptr
121    }
122
123    pub const fn as_ptr(self) -> *const T {
124        self.ptr.cast()
125    }
126}
127
128impl<T> DeviceScalarMut<T> {
129    /// Creates a mutable device scalar from a raw device pointer.
130    ///
131    /// # Safety
132    ///
133    /// `ptr` must point to a valid writable device allocation containing a `T`
134    /// for every cuSPARSE operation using this scalar.
135    pub const unsafe fn new(ptr: DevicePtr) -> Self {
136        Self {
137            ptr,
138            _t: PhantomData,
139        }
140    }
141
142    pub const fn as_device_ptr(&self) -> DevicePtr {
143        self.ptr
144    }
145
146    pub const fn as_mut_ptr(&self) -> *mut T {
147        self.ptr.cast()
148    }
149}
150
151impl<'a, T> From<&'a mut T> for ScalarMut<'a, T> {
152    fn from(value: &'a mut T) -> Self {
153        Self::host(value)
154    }
155}
156
157impl<T> From<DeviceScalarMut<T>> for ScalarMut<'_, T> {
158    fn from(value: DeviceScalarMut<T>) -> Self {
159        Self::Device(value)
160    }
161}