openvm_stark_backend/prover/
matrix.rs

1use std::borrow::Borrow;
2
3use getset::CopyGetters;
4use p3_field::Field;
5use p3_matrix::dense::{DenseMatrix, RowMajorMatrix};
6use p3_maybe_rayon::prelude::*;
7use serde::{Deserialize, Serialize};
8
9use crate::prover::MatrixDimensions;
10
11impl<T, V: Borrow<[T]>> MatrixDimensions for DenseMatrix<T, V> {
12    fn height(&self) -> usize {
13        self.values.borrow().len() / self.width
14    }
15    fn width(&self) -> usize {
16        self.width
17    }
18}
19
20/// Trait to consolidate different virtual matrices that may not be backed by an in-memory matrix
21/// with a standard column-major or row-major layout.
22///
23/// Note: for performance, users should still be aware of the underlying memory layout. This trait
24/// is just an organizational convenience.
25pub trait MatrixView<F>: MatrixDimensions {
26    fn get(&self, row_idx: usize, col_idx: usize) -> Option<&F> {
27        if col_idx >= self.width() || row_idx >= self.height() {
28            None
29        } else {
30            // SAFETY: bounds checked above
31            Some(unsafe { self.get_unchecked(row_idx, col_idx) })
32        }
33    }
34    /// Get a reference to an element without bounds checking.
35    ///
36    /// # Safety
37    ///
38    /// The caller must ensure that `col_idx` and `row_idx` are within the bounds of the matrix.
39    unsafe fn get_unchecked(&self, row_idx: usize, col_idx: usize) -> &F;
40}
41
42#[derive(Clone, Debug, Serialize, Deserialize)]
43pub struct ColMajorMatrix<F> {
44    pub values: Vec<F>,
45    width: usize,
46    height: usize,
47}
48
49#[derive(Clone, Copy, Debug)]
50pub struct ColMajorMatrixView<'a, F> {
51    pub values: &'a [F],
52    width: usize,
53    height: usize,
54}
55
56/// Vertically strided column-major matrix view.
57#[derive(Clone, Copy, Debug, CopyGetters)]
58pub struct StridedColMajorMatrixView<'a, F> {
59    #[getset(get_copy = "pub")]
60    values: &'a [F],
61    width: usize,
62    height: usize,
63    /// Row stride
64    #[getset(get_copy = "pub")]
65    stride: usize,
66}
67
68impl<F> ColMajorMatrix<F> {
69    pub fn new(values: Vec<F>, width: usize) -> Self {
70        assert_eq!(values.len() % width, 0);
71        let height = values.len() / width;
72        assert!(height.is_power_of_two());
73        Self {
74            values,
75            width,
76            height,
77        }
78    }
79
80    pub fn column(&self, col_idx: usize) -> &[F] {
81        let start = col_idx * self.height;
82        let end = start + self.height;
83        &self.values[start..end]
84    }
85
86    pub fn columns(&self) -> impl Iterator<Item = &[F]> {
87        self.values.chunks_exact(self.height)
88    }
89
90    pub fn as_view(&self) -> ColMajorMatrixView<'_, F> {
91        ColMajorMatrixView {
92            values: &self.values,
93            width: self.width,
94            height: self.height,
95        }
96    }
97
98    pub fn from_row_major(mat: &RowMajorMatrix<F>) -> Self
99    where
100        F: Field,
101    {
102        let mut values = F::zero_vec(mat.values.len());
103        let width = mat.width;
104        let height = mat.values.len() / mat.width;
105        values.par_iter_mut().enumerate().for_each(|(idx, value)| {
106            let r = idx % height;
107            let c = idx / height;
108            // SAFETY: index is in bounds for row-major matrix
109            *value = unsafe { *mat.values.get_unchecked(r * width + c) };
110        });
111        Self {
112            values,
113            width,
114            height,
115        }
116    }
117}
118
119impl<F: Sync> ColMajorMatrix<F> {
120    pub fn par_columns(&self) -> impl ParallelIterator<Item = &[F]> {
121        self.values.par_chunks_exact(self.height)
122    }
123}
124
125impl<F> MatrixDimensions for ColMajorMatrix<F> {
126    fn width(&self) -> usize {
127        self.width
128    }
129    fn height(&self) -> usize {
130        self.height
131    }
132}
133
134impl<F> MatrixView<F> for ColMajorMatrix<F> {
135    unsafe fn get_unchecked(&self, row_idx: usize, col_idx: usize) -> &F {
136        debug_assert!(col_idx < self.width);
137        debug_assert!(row_idx < self.height);
138        self.values.get_unchecked(col_idx * self.height + row_idx)
139    }
140}
141
142impl<'a, F> ColMajorMatrixView<'a, F> {
143    pub fn new(values: &'a [F], width: usize) -> Self {
144        assert_eq!(values.len() % width, 0);
145        let height = values.len() / width;
146        debug_assert!(height == 0 || height.is_power_of_two());
147        Self {
148            values,
149            width,
150            height,
151        }
152    }
153
154    pub fn column(&self, col_idx: usize) -> &[F] {
155        let start = col_idx * self.height;
156        let end = start + self.height;
157        &self.values[start..end]
158    }
159
160    pub fn columns(&self) -> impl Iterator<Item = &[F]> {
161        self.values.chunks_exact(self.height)
162    }
163}
164
165impl<F> MatrixDimensions for ColMajorMatrixView<'_, F> {
166    fn width(&self) -> usize {
167        self.width
168    }
169    fn height(&self) -> usize {
170        self.height
171    }
172}
173
174impl<F> MatrixView<F> for ColMajorMatrixView<'_, F> {
175    unsafe fn get_unchecked(&self, row_idx: usize, col_idx: usize) -> &F {
176        debug_assert!(col_idx < self.width);
177        debug_assert!(row_idx < self.height);
178        self.values
179            .get_unchecked(col_maj_idx(row_idx, col_idx, self.height))
180    }
181}
182
183impl<'a, F> StridedColMajorMatrixView<'a, F> {
184    pub fn new(values: &'a [F], width: usize, stride: usize) -> Self {
185        assert_eq!(values.len() % (width * stride), 0);
186        let height = values.len() / (width * stride);
187        debug_assert!(height == 0 || height.is_power_of_two());
188        Self {
189            values,
190            width,
191            height,
192            stride,
193        }
194    }
195
196    pub fn to_matrix(&self) -> ColMajorMatrix<F>
197    where
198        F: Field,
199    {
200        let values: Vec<_> = (0..self.width * self.height)
201            .into_par_iter()
202            .map(|i| {
203                let r = i % self.height;
204                let c = i / self.height;
205                unsafe { *self.get_unchecked(r, c) }
206            })
207            .collect();
208        ColMajorMatrix::new(values, self.width)
209    }
210
211    pub fn to_row_major_matrix(&self) -> RowMajorMatrix<F>
212    where
213        F: Field,
214    {
215        let values: Vec<_> = (0..self.width * self.height)
216            .into_par_iter()
217            .map(|i| {
218                let r = i / self.width;
219                let c = i % self.width;
220                unsafe { *self.get_unchecked(r, c) }
221            })
222            .collect();
223        RowMajorMatrix::new(values, self.width)
224    }
225}
226
227impl<F> MatrixDimensions for StridedColMajorMatrixView<'_, F> {
228    fn width(&self) -> usize {
229        self.width
230    }
231    fn height(&self) -> usize {
232        self.height
233    }
234}
235
236impl<F> MatrixView<F> for StridedColMajorMatrixView<'_, F> {
237    unsafe fn get_unchecked(&self, row_idx: usize, col_idx: usize) -> &F {
238        debug_assert!(col_idx < self.width);
239        debug_assert!(row_idx < self.height);
240        self.values.get_unchecked(col_maj_idx(
241            row_idx * self.stride,
242            col_idx,
243            self.height * self.stride,
244        ))
245    }
246}
247
248impl<'a, F> From<ColMajorMatrixView<'a, F>> for StridedColMajorMatrixView<'a, F> {
249    fn from(mat: ColMajorMatrixView<'a, F>) -> Self {
250        Self::new(mat.values, mat.width, 1)
251    }
252}
253
254#[inline(always)]
255pub(crate) fn col_maj_idx(row_idx: usize, col_idx: usize, height: usize) -> usize {
256    col_idx * height + row_idx
257}
openvm_stark_backend/prover/matrix.rs

openvm_stark_backend/prover/
matrix.rs
