openvm_cuda_backend/
base.rs

1use std::{fmt::Debug, marker::PhantomData, sync::Arc};
2
3use openvm_cuda_common::{
4    copy::MemCopyD2H, d_buffer::DeviceBuffer, error::MemCopyError, stream::GpuDeviceCtx,
5};
6use openvm_stark_backend::prover::MatrixDimensions;
7
8pub struct DeviceMatrix<T> {
9    buffer: Arc<DeviceBuffer<T>>,
10    height: usize,
11    width: usize,
12}
13
14unsafe impl<T> Send for DeviceMatrix<T> {}
15unsafe impl<T> Sync for DeviceMatrix<T> {}
16
17impl<T> Clone for DeviceMatrix<T> {
18    fn clone(&self) -> Self {
19        Self {
20            buffer: Arc::clone(&self.buffer),
21            height: self.height,
22            width: self.width,
23        }
24    }
25}
26
27impl<T> Drop for DeviceMatrix<T> {
28    fn drop(&mut self) {
29        tracing::debug!(
30            "Dropping DeviceMatrix of size {} with Arc strong count={}",
31            self.buffer.len(),
32            self.strong_count()
33        );
34    }
35}
36
37impl<T> DeviceMatrix<T> {
38    pub fn new(buffer: Arc<DeviceBuffer<T>>, height: usize, width: usize) -> Self {
39        assert_ne!(
40            height * width,
41            0,
42            "Zero dimensions h {} w {} are wrong",
43            height,
44            width
45        );
46        assert_eq!(
47            buffer.len(),
48            height * width,
49            "Buffer size must match dimensions"
50        );
51        Self {
52            buffer,
53            height,
54            width,
55        }
56    }
57
58    pub fn with_capacity_on(height: usize, width: usize, device_ctx: &GpuDeviceCtx) -> Self {
59        let buffer = DeviceBuffer::with_capacity_on(height * width, device_ctx);
60        Self::new(Arc::new(buffer), height, width)
61    }
62
63    pub fn dummy() -> Self {
64        Self {
65            buffer: Arc::new(DeviceBuffer::new()),
66            height: 0,
67            width: 0,
68        }
69    }
70
71    pub fn buffer(&self) -> &DeviceBuffer<T> {
72        &self.buffer
73    }
74
75    pub fn strong_count(&self) -> usize {
76        Arc::strong_count(&self.buffer)
77    }
78
79    pub fn as_view(&self) -> DeviceMatrixView<'_, T> {
80        // SAFETY: buffer is borrowed for lifetime 'a of the view
81        unsafe { DeviceMatrixView::from_raw_parts(self.buffer.as_ptr(), self.height, self.width) }
82    }
83}
84
85impl<T> MatrixDimensions for DeviceMatrix<T> {
86    #[inline]
87    fn height(&self) -> usize {
88        self.height
89    }
90
91    #[inline]
92    fn width(&self) -> usize {
93        self.width
94    }
95}
96
97impl<T> MemCopyD2H<T> for DeviceMatrix<T> {
98    fn to_host_on(&self, device_ctx: &GpuDeviceCtx) -> Result<Vec<T>, MemCopyError> {
99        self.buffer.to_host_on(device_ctx)
100    }
101}
102
103impl<T: Debug> Debug for DeviceMatrix<T> {
104    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
105        write!(
106            f,
107            "DeviceMatrix (height = {}, width = {}): {:?}",
108            self.height(),
109            self.width(),
110            self.buffer()
111        )
112    }
113}
114
115/// View of a device matrix. Dropping does not free memory.
116#[derive(Clone, Copy)]
117pub struct DeviceMatrixView<'a, T> {
118    ptr: *const T,
119    height: usize,
120    width: usize,
121    _ptr_lifetime: PhantomData<&'a T>,
122}
123
124unsafe impl<T> Send for DeviceMatrixView<'_, T> {}
125unsafe impl<T> Sync for DeviceMatrixView<'_, T> {}
126
127impl<T> DeviceMatrixView<'_, T> {
128    /// # Safety
129    /// - The pointer must be valid for the lifetime of the view.
130    /// - The pointer must have memory allocated for the following `height * width` elements of `T`.
131    pub unsafe fn from_raw_parts(ptr: *const T, height: usize, width: usize) -> Self {
132        Self {
133            ptr,
134            height,
135            width,
136            _ptr_lifetime: PhantomData,
137        }
138    }
139
140    pub fn as_ptr(&self) -> *const T {
141        self.ptr
142    }
143}
144
145impl<T> MatrixDimensions for DeviceMatrixView<'_, T> {
146    #[inline]
147    fn height(&self) -> usize {
148        self.height
149    }
150
151    #[inline]
152    fn width(&self) -> usize {
153        self.width
154    }
155}
156
157/// The following trait and types are borrowed from [halo2](https:://github.com/zcash/halo2).
158/// The basis over which a polynomial is described.
159pub trait Basis: Copy + Debug + Send + Sync {}
160
161/// The polynomial is defined as coefficients
162#[derive(Clone, Copy, Debug)]
163pub struct Coeff;
164impl Basis for Coeff {}
165
166/// The polynomial is defined as coefficients of Lagrange basis polynomials
167#[derive(Clone, Copy, Debug)]
168pub struct LagrangeCoeff;
169impl Basis for LagrangeCoeff {}
170
171/// The polynomial is defined as coefficients of Lagrange basis polynomials in
172/// an extended size domain which supports multiplication
173#[derive(Clone, Copy, Debug)]
174pub struct ExtendedLagrangeCoeff;
175impl Basis for ExtendedLagrangeCoeff {}
176
177pub struct DevicePoly<T, B> {
178    pub is_bit_reversed: bool,
179    pub coeff: DeviceBuffer<T>,
180    _marker: PhantomData<B>,
181}
182
183impl<T, B> DevicePoly<T, B> {
184    pub fn new(is_bit_reversed: bool, coeff: DeviceBuffer<T>) -> Self {
185        Self {
186            is_bit_reversed,
187            coeff,
188            _marker: PhantomData,
189        }
190    }
191
192    #[inline]
193    pub fn len(&self) -> usize {
194        self.coeff.len()
195    }
196
197    #[inline]
198    pub fn is_empty(&self) -> bool {
199        self.len() == 0
200    }
201}
202
203#[cfg(test)]
204mod tests {
205    use super::*;
206
207    #[test]
208    fn test_device_matrix() {
209        let device_ctx = GpuDeviceCtx::for_current_device().unwrap();
210        let buffer = Arc::new(DeviceBuffer::<i32>::with_capacity_on(12, &device_ctx));
211        let matrix = DeviceMatrix::<i32>::new(buffer, 3, 4);
212        assert_eq!(matrix.height(), 3);
213        assert_eq!(matrix.width(), 4);
214    }
215}
openvm_cuda_backend/base.rs

openvm_cuda_backend/
base.rs
