openvm_recursion_circuit/batch_constraint/eq_airs/eq_sharp_uni/
trace.rs

1use std::borrow::BorrowMut;
2
3use openvm_stark_backend::{
4    keygen::types::{MultiStarkVerifyingKey, MultiStarkVerifyingKey0},
5    poly_common::Squarable,
6};
7use openvm_stark_sdk::config::baby_bear_poseidon2::{BabyBearPoseidon2Config, EF, F};
8use p3_field::{BasedVectorSpace, Field, PrimeCharacteristicRing, TwoAdicField};
9use p3_matrix::dense::RowMajorMatrix;
10use p3_maybe_rayon::prelude::*;
11
12use crate::{
13    batch_constraint::eq_airs::eq_sharp_uni::air::{EqSharpUniCols, EqSharpUniReceiverCols},
14    system::Preflight,
15    tracegen::RowMajorChip,
16    utils::MultiProofVecVec,
17};
18
19#[repr(C)]
20#[derive(Debug, Clone, Copy)]
21struct EqSharpUniRecord {
22    xi: EF,
23    product: EF,
24
25    root: F,
26    root_pow: F,
27
28    // PERF: we can only store `xi_idx: u32` and later take `xi` from the transcript
29    xi_idx: u32,
30    root_half_order: u32,
31    iter_idx: u32,
32}
33
34#[derive(Debug, Clone)]
35pub struct EqSharpUniBlob {
36    records: MultiProofVecVec<EqSharpUniRecord>,
37    final_products: MultiProofVecVec<EF>,
38    rs: Vec<EF>,
39}
40
41impl EqSharpUniBlob {
42    fn new(l_skip: usize, num_proofs: usize) -> Self {
43        Self {
44            records: MultiProofVecVec::with_capacity(num_proofs * ((1 << l_skip) - 1)),
45            final_products: MultiProofVecVec::with_capacity(num_proofs << l_skip),
46            rs: Vec::with_capacity(num_proofs),
47        }
48    }
49}
50
51pub fn generate_eq_sharp_uni_blob(
52    vk: &MultiStarkVerifyingKey0<BabyBearPoseidon2Config>,
53    preflights: &[&Preflight],
54) -> EqSharpUniBlob {
55    let l_skip = vk.params.l_skip;
56    let mut blob = EqSharpUniBlob::new(l_skip, preflights.len());
57    let mut products = vec![EF::ONE; 1 << l_skip];
58    let roots = F::two_adic_generator(l_skip)
59        .inverse()
60        .exp_powers_of_2()
61        .take(l_skip)
62        .collect::<Vec<_>>();
63    for preflight in preflights.iter() {
64        products[0] = EF::ONE;
65        for i in 0..l_skip {
66            let xi_idx = l_skip - 1 - i;
67            let xi = preflight.gkr.xi[xi_idx].1;
68            let root = roots[l_skip - 1 - i];
69            let mut root_pow = F::ONE;
70            for (j, &product) in products.iter().take(1 << i).enumerate() {
71                blob.records.push(EqSharpUniRecord {
72                    xi,
73                    product,
74                    root,
75                    root_pow,
76                    xi_idx: xi_idx as u32,
77                    root_half_order: 1 << i,
78                    iter_idx: j as u32,
79                });
80                root_pow *= root;
81            }
82            for j in (0..(1 << i)).rev() {
83                let value = products[j];
84                let root_pow = blob.records.data()[blob.records.len() - (1 << i) + j].root_pow;
85                let second = xi * root_pow;
86                products[j + (1 << i)] = value * (EF::ONE - xi - second);
87                products[j] = value * (EF::ONE - xi + second);
88            }
89        }
90        blob.records.end_proof();
91        blob.final_products.extend_from_slice(&products);
92        blob.final_products.end_proof();
93        blob.rs.push(preflight.batch_constraint.sumcheck_rnd[0]);
94    }
95    blob
96}
97
98pub struct EqSharpUniTraceGenerator;
99
100impl RowMajorChip<F> for EqSharpUniTraceGenerator {
101    type Ctx<'a> = (
102        &'a MultiStarkVerifyingKey<BabyBearPoseidon2Config>,
103        &'a EqSharpUniBlob,
104        &'a [&'a Preflight],
105    );
106
107    #[tracing::instrument(level = "trace", skip_all)]
108    fn generate_trace(
109        &self,
110        ctx: &Self::Ctx<'_>,
111        required_height: Option<usize>,
112    ) -> Option<RowMajorMatrix<F>> {
113        let (vk, blob, preflights) = ctx;
114        let width = EqSharpUniCols::<F>::width();
115        let l_skip = vk.inner.params.l_skip;
116        let one_height = (1 << l_skip) - 1;
117        let total_height = one_height * preflights.len();
118
119        let padded_height = if let Some(height) = required_height {
120            if height < total_height {
121                return None;
122            }
123            height
124        } else {
125            total_height.next_power_of_two()
126        };
127        let mut trace = vec![F::ZERO; padded_height * width];
128
129        for pidx in 0..preflights.len() {
130            let records = &blob.records[pidx];
131            trace[(pidx * one_height * width)..((pidx + 1) * one_height * width)]
132                .par_chunks_exact_mut(width)
133                .zip(records.par_iter())
134                .enumerate()
135                .for_each(|(i, (chunk, record))| {
136                    let cols: &mut EqSharpUniCols<_> = chunk.borrow_mut();
137                    cols.is_valid = F::ONE;
138                    cols.is_first = F::from_bool(i == 0);
139                    cols.proof_idx = F::from_usize(pidx);
140                    cols.xi_idx = F::from_u32(record.xi_idx);
141                    cols.xi
142                        .copy_from_slice(record.xi.as_basis_coefficients_slice());
143                    cols.iter_idx = F::from_u32(record.iter_idx);
144                    cols.is_first_iter = F::from_bool(record.iter_idx == 0);
145                    cols.product_before
146                        .copy_from_slice(record.product.as_basis_coefficients_slice());
147                    cols.root = record.root;
148                    cols.root_pow = record.root_pow;
149                    cols.root_half_order = F::from_u32(record.root_half_order);
150                });
151        }
152
153        trace[total_height * width..]
154            .par_chunks_mut(width)
155            .enumerate()
156            .for_each(|(i, chunk)| {
157                let cols: &mut EqSharpUniCols<F> = chunk.borrow_mut();
158                cols.proof_idx = F::from_usize(preflights.len() + i);
159            });
160
161        Some(RowMajorMatrix::new(trace, width))
162    }
163}
164
165pub struct EqSharpUniReceiverTraceGenerator;
166
167impl RowMajorChip<F> for EqSharpUniReceiverTraceGenerator {
168    type Ctx<'a> = (
169        &'a MultiStarkVerifyingKey<BabyBearPoseidon2Config>,
170        &'a EqSharpUniBlob,
171        &'a [&'a Preflight],
172    );
173
174    #[tracing::instrument(level = "trace", skip_all)]
175    fn generate_trace(
176        &self,
177        ctx: &Self::Ctx<'_>,
178        required_height: Option<usize>,
179    ) -> Option<RowMajorMatrix<F>> {
180        let (vk, blob, preflights) = ctx;
181        let l_skip = vk.inner.params.l_skip;
182
183        let width = EqSharpUniReceiverCols::<F>::width();
184        let one_height = 1 << l_skip;
185        let total_height = one_height * preflights.len();
186        let padded_height = if let Some(height) = required_height {
187            if height < total_height {
188                return None;
189            }
190            height
191        } else {
192            total_height.next_power_of_two()
193        };
194        let mut trace = vec![F::ZERO; padded_height * width];
195
196        for pidx in 0..preflights.len() {
197            let products = &blob.final_products[pidx];
198            let r = blob.rs[pidx];
199            trace[(pidx * one_height * width)..((pidx + 1) * one_height * width)]
200                .par_chunks_exact_mut(width)
201                .zip(products.par_iter())
202                .enumerate()
203                .for_each(|(i, (chunk, product))| {
204                    let cols: &mut EqSharpUniReceiverCols<_> = chunk.borrow_mut();
205                    cols.is_valid = F::ONE;
206                    cols.is_first = F::from_bool(i == 0);
207                    cols.is_last = F::from_bool(i + 1 == one_height);
208                    cols.proof_idx = F::from_usize(pidx);
209                    cols.coeff
210                        .copy_from_slice(product.as_basis_coefficients_slice());
211                    cols.r.copy_from_slice(r.as_basis_coefficients_slice());
212                    cols.idx = F::from_usize(i);
213                });
214            let mut cur_sum = EF::ZERO;
215            trace[(pidx * one_height * width)..((pidx + 1) * one_height * width)]
216                .chunks_exact_mut(width)
217                .rev()
218                .for_each(|chunk| {
219                    let cols: &mut EqSharpUniReceiverCols<_> = chunk.borrow_mut();
220                    cur_sum = cur_sum * r + EF::from_basis_coefficients_slice(&cols.coeff).unwrap();
221                    cols.cur_sum
222                        .copy_from_slice(cur_sum.as_basis_coefficients_slice());
223                });
224        }
225
226        Some(RowMajorMatrix::new(trace, width))
227    }
228}
openvm_recursion_circuit/batch_constraint/eq_airs/eq_sharp_uni/trace.rs

openvm_recursion_circuit/batch_constraint/eq_airs/eq_sharp_uni/
trace.rs
