p3_miden_dev_utils/
matrix.rs

1//! Matrix generation utilities for tests and benchmarks.
2
3use alloc::vec::Vec;
4
5use p3_dft::{Radix2DFTSmallBatch, TwoAdicSubgroupDft};
6use p3_field::{BasedVectorSpace, Field, PrimeCharacteristicRing, TwoAdicField};
7use p3_matrix::{Matrix, bitrev::BitReversibleMatrix, dense::RowMajorMatrix};
8use rand::{
9    SeedableRng,
10    distr::{Distribution, StandardUniform},
11    rngs::SmallRng,
12};
13
14use crate::fixtures::BENCH_SEED;
15
16// =============================================================================
17// Benchmark matrix generation
18// =============================================================================
19
20/// Generate benchmark matrices from relative specs.
21///
22/// Creates matrices with heights relative to `max_height = 1 << log_max_height`.
23/// Each spec `(offset, width)` creates a matrix with:
24/// - height = `max_height >> offset`
25/// - width = `width`
26///
27/// Matrices in each group are sorted by ascending height.
28pub fn generate_matrices_from_specs<F: Field>(
29    specs: &[&[(usize, usize)]],
30    log_max_height: u8,
31) -> Vec<Vec<RowMajorMatrix<F>>>
32where
33    StandardUniform: Distribution<F>,
34{
35    let rng = &mut SmallRng::seed_from_u64(BENCH_SEED);
36    let max_height = 1 << log_max_height as usize;
37
38    specs
39        .iter()
40        .map(|group_specs| {
41            let mut matrices: Vec<RowMajorMatrix<F>> = group_specs
42                .iter()
43                .map(|&(offset, width)| {
44                    let height = max_height >> offset;
45                    RowMajorMatrix::rand(rng, height, width)
46                })
47                .collect();
48            // Sort by ascending height (required by LMCS)
49            matrices.sort_by_key(|m| m.height());
50            matrices
51        })
52        .collect()
53}
54
55/// Generate a single flat matrix for FRI fold benchmarks.
56pub fn generate_flat_matrix<F: Field>(log_height: u8, width: usize) -> RowMajorMatrix<F>
57where
58    StandardUniform: Distribution<F>,
59{
60    let rng = &mut SmallRng::seed_from_u64(BENCH_SEED);
61    RowMajorMatrix::rand(rng, 1 << log_height as usize, width)
62}
63
64/// Calculate total elements across all matrices.
65pub fn total_elements<F: Clone + Send + Sync>(matrix_groups: &[Vec<RowMajorMatrix<F>>]) -> u64 {
66    matrix_groups
67        .iter()
68        .flat_map(|g| g.iter())
69        .map(|m| {
70            let dims = m.dimensions();
71            (dims.height * dims.width) as u64
72        })
73        .sum()
74}
75
76/// Calculate total elements for a flat matrix list.
77pub fn total_elements_flat<F: Clone + Send + Sync>(matrices: &[RowMajorMatrix<F>]) -> u64 {
78    matrices
79        .iter()
80        .map(|m| {
81            let dims = m.dimensions();
82            (dims.height * dims.width) as u64
83        })
84        .sum()
85}
86
87// =============================================================================
88// Test matrix utilities
89// =============================================================================
90
91/// Generate a matrix of LDE evaluations for random low-degree polynomials.
92///
93/// Each column is a polynomial of degree `poly_degree`, evaluated on the coset gK
94/// in bit-reversed order, where g = `shift` and K is a subgroup of order `lde_size`.
95///
96/// The coset evaluation is computed by scaling coefficients: for f(X) = Σ c_j X^j,
97/// the coset evaluations f(gX) = Σ (c_j g^j) X^j are obtained by DFT of scaled coefficients.
98pub fn random_lde_matrix<F, V>(
99    rng: &mut SmallRng,
100    log_poly_degree: u8,
101    log_blowup: u8,
102    num_columns: usize,
103    shift: F,
104) -> RowMajorMatrix<V>
105where
106    F: TwoAdicField,
107    V: BasedVectorSpace<F> + Clone + Send + Sync + Default,
108    StandardUniform: Distribution<V>,
109{
110    let poly_degree = 1 << log_poly_degree as usize;
111    let dft = Radix2DFTSmallBatch::<F>::default();
112
113    let evals = RowMajorMatrix::rand(rng, poly_degree, num_columns);
114    let lde = dft.coset_lde_algebra_batch(evals, log_blowup as usize, shift);
115    lde.bit_reverse_rows().to_row_major_matrix()
116}
117
118/// Concatenate matrices horizontally, padding each to a multiple of `R`.
119///
120/// All matrices are lifted to the maximum height first.
121pub fn concatenate_matrices<F: Field + PrimeCharacteristicRing, const R: usize>(
122    matrices: &[RowMajorMatrix<F>],
123) -> RowMajorMatrix<F> {
124    let max_height = matrices.last().unwrap().height();
125    let width: usize = matrices.iter().map(|m| aligned_len(m.width(), R)).sum();
126
127    let concatenated_data: Vec<_> = (0..max_height)
128        .flat_map(|idx| {
129            matrices.iter().flat_map(move |m| {
130                let mut row = m.row_slice(idx).unwrap().to_vec();
131                let padded_width = aligned_len(row.len(), R);
132                row.resize(padded_width, F::ZERO);
133                row
134            })
135        })
136        .collect();
137    RowMajorMatrix::new(concatenated_data, width)
138}
139
140#[inline]
141const fn aligned_len(len: usize, alignment: usize) -> usize {
142    if alignment <= 1 {
143        len
144    } else {
145        len.next_multiple_of(alignment)
146    }
147}
p3_miden_dev_utils/matrix.rs

p3_miden_dev_utils/
matrix.rs
