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