openvm_recursion_circuit/batch_constraint/eq_airs/eq_3b/
trace.rs1use std::borrow::BorrowMut;
2
3use openvm_stark_backend::keygen::types::{MultiStarkVerifyingKey, MultiStarkVerifyingKey0};
4use openvm_stark_sdk::config::baby_bear_poseidon2::{BabyBearPoseidon2Config, EF, F};
5use p3_field::{BasedVectorSpace, PrimeCharacteristicRing};
6use p3_matrix::dense::RowMajorMatrix;
7
8use crate::{
9 batch_constraint::eq_airs::eq_3b::air::Eq3bColumns, system::Preflight, tracegen::RowMajorChip,
10 utils::MultiProofVecVec,
11};
12
13#[repr(C)]
14#[derive(Debug, Copy, Clone)]
15pub(crate) struct StackedIdxRecord {
16 pub sort_idx: u32,
17 pub interaction_idx: u32,
18 pub stacked_idx: u32,
19 pub n_lift: u32,
20 pub is_last_in_air: bool,
21 pub no_interactions: bool,
22}
23
24#[repr(C)]
25#[derive(Debug, Copy, Clone)]
26pub(crate) struct RecordIdx {
27 pub record_idx: u32,
28 pub local_n: u32,
29}
30
31impl StackedIdxRecord {
32 pub fn from_usizes(
33 sort_idx: usize,
34 interaction_idx: usize,
35 stacked_idx: usize,
36 n_lift: usize,
37 is_last_in_air: bool,
38 no_interactions: bool,
39 ) -> Self {
40 Self {
41 sort_idx: sort_idx as u32,
42 interaction_idx: interaction_idx as u32,
43 stacked_idx: stacked_idx as u32,
44 n_lift: n_lift as u32,
45 is_last_in_air,
46 no_interactions,
47 }
48 }
49
50 pub fn eq_mle(&self, xi: &[EF], l_skip: usize, n_logup: usize) -> EF {
51 xi[l_skip + (self.n_lift as usize)..l_skip + n_logup]
52 .iter()
53 .enumerate()
54 .map(|(i, &x)| {
55 if self.stacked_idx & (1 << (l_skip + (self.n_lift as usize) + i)) > 0 {
56 x
57 } else {
58 EF::ONE - x
59 }
60 })
61 .fold(EF::ONE, |acc, x| acc * x)
62 }
63}
64
65pub struct Eq3bBlob {
66 pub(crate) all_stacked_ids: MultiProofVecVec<StackedIdxRecord>,
67 pub(crate) record_idxs: MultiProofVecVec<RecordIdx>,
68}
69
70impl Eq3bBlob {
71 fn new() -> Self {
72 Self {
73 all_stacked_ids: MultiProofVecVec::new(),
74 record_idxs: MultiProofVecVec::new(),
75 }
76 }
77}
78
79pub(crate) fn generate_eq_3b_blob(
80 vk: &MultiStarkVerifyingKey0<BabyBearPoseidon2Config>,
81 preflights: &[&Preflight],
82) -> Eq3bBlob {
83 let l_skip = vk.params.l_skip;
84 let mut blob = Eq3bBlob::new();
85 for preflight in preflights.iter() {
86 let mut row_idx = 0;
87 let mut record_idx = 0;
88 let n_logup = preflight.proof_shape.n_logup;
89 for (sort_idx, (air_idx, vdata)) in
90 preflight.proof_shape.sorted_trace_vdata.iter().enumerate()
91 {
92 let n_lift = vdata.log_height.saturating_sub(l_skip);
93 let num_interactions = vk.per_air[*air_idx].num_interactions();
94 if num_interactions == 0 {
95 blob.record_idxs.push(RecordIdx {
96 record_idx,
97 local_n: 0,
98 });
99 blob.all_stacked_ids.push(StackedIdxRecord::from_usizes(
100 sort_idx, 0, row_idx, n_lift, true, true,
101 ));
102 record_idx += 1;
103 } else {
104 for i in 0..num_interactions {
105 blob.record_idxs
106 .extend((0..=n_logup).map(|local_n| RecordIdx {
107 record_idx,
108 local_n: local_n as u32,
109 }));
110 blob.all_stacked_ids.push(StackedIdxRecord::from_usizes(
111 sort_idx,
112 i,
113 row_idx,
114 n_lift,
115 i + 1 == num_interactions,
116 false,
117 ));
118 row_idx += 1 << (l_skip + n_lift);
119 record_idx += 1;
120 }
121 }
122 }
123 debug_assert!(row_idx <= 1 << (l_skip + n_logup));
124 blob.all_stacked_ids.end_proof();
125 blob.record_idxs.end_proof();
126 }
127 blob
128}
129
130pub struct Eq3bTraceGenerator;
131
132impl RowMajorChip<F> for Eq3bTraceGenerator {
133 type Ctx<'a> = (
134 &'a MultiStarkVerifyingKey<BabyBearPoseidon2Config>,
135 &'a Eq3bBlob,
136 &'a [&'a Preflight],
137 );
138
139 #[tracing::instrument(level = "trace", skip_all)]
140 fn generate_trace(
141 &self,
142 ctx: &Self::Ctx<'_>,
143 required_height: Option<usize>,
144 ) -> Option<RowMajorMatrix<F>> {
145 let (vk, blob, preflights) = ctx;
146 let width = Eq3bColumns::<F>::width();
147 let l_skip = vk.inner.params.l_skip;
148
149 let total_valid = preflights
150 .iter()
151 .enumerate()
152 .map(|(i, p)| {
153 let num_dummy_records = blob.all_stacked_ids[i]
154 .iter()
155 .filter(|record| record.no_interactions)
156 .count();
157 (p.proof_shape.n_logup + 1) * (blob.all_stacked_ids[i].len() - num_dummy_records)
158 + num_dummy_records
159 })
160 .sum();
161 let padding_height = if let Some(height) = required_height {
162 if height < total_valid {
163 return None;
164 }
165 height
166 } else {
167 total_valid.next_power_of_two()
168 };
169
170 let mut trace = vec![F::ZERO; padding_height * width];
171
172 let mut cur_height = 0;
173 for (pidx, preflight) in preflights.iter().enumerate() {
174 let xi = &preflight.batch_constraint.xi;
175 let n_logup = preflight.proof_shape.n_logup;
176
177 let stacked_ids = &blob.all_stacked_ids[pidx];
178 let one_height = n_logup + 1;
179
180 for (j, record) in stacked_ids.iter().enumerate() {
181 let this_height = if record.no_interactions {
182 1
183 } else {
184 one_height
185 };
186 let shifted_idx = record.stacked_idx >> l_skip;
187 let mut cur_eq = EF::ONE;
188 trace[cur_height * width..(cur_height + this_height) * width]
189 .chunks_exact_mut(width)
190 .enumerate()
191 .for_each(|(n, chunk)| {
192 let cols: &mut Eq3bColumns<_> = chunk.borrow_mut();
193 cols.is_valid = F::ONE;
194 cols.is_first = F::from_bool(j == 0 && n == 0);
195 cols.proof_idx = F::from_usize(pidx);
196
197 cols.sort_idx = F::from_u32(record.sort_idx);
198 cols.interaction_idx = F::from_u32(record.interaction_idx);
199 cols.n_lift = F::from_u32(record.n_lift);
200 cols.n_logup = F::from_usize(n_logup);
201 cols.two_to_the_n_lift = F::from_usize(1 << record.n_lift);
202 cols.n = F::from_usize(n);
203 cols.n_at_least_n_lift = F::from_bool(n >= record.n_lift as usize);
204 cols.has_no_interactions = F::from_bool(record.no_interactions);
205 cols.hypercube_volume = F::from_usize(1 << n);
206 cols.is_first_in_air = F::from_bool(
207 record.interaction_idx == 0 && n == 0 || record.no_interactions,
208 );
209 cols.is_first_in_interaction =
210 F::from_bool(n == 0 || record.no_interactions);
211 cols.idx = F::from_u32(shifted_idx & ((1 << n) - 1));
212 cols.running_idx = F::from_u32(shifted_idx);
213 let nth_bit = (shifted_idx & (1 << n)) > 0;
214 cols.nth_bit = F::from_bool(nth_bit);
215 let xi = if (record.n_lift as usize..n_logup).contains(&n) {
216 xi[l_skip + n]
217 } else if nth_bit {
218 EF::ONE
219 } else {
220 EF::ZERO
221 };
222 cols.xi.copy_from_slice(xi.as_basis_coefficients_slice());
223 cols.eq
224 .copy_from_slice(cur_eq.as_basis_coefficients_slice());
225 cur_eq *= if nth_bit { xi } else { EF::ONE - xi };
226 });
227 cur_height += this_height;
228 }
229 }
230
231 Some(RowMajorMatrix::new(trace, width))
232 }
233}
234
235#[cfg(feature = "cuda")]
236pub(in crate::batch_constraint) mod cuda {
237 use openvm_cuda_backend::{base::DeviceMatrix, prelude::F, GpuBackend};
238 use openvm_cuda_common::{copy::MemCopyH2D, stream::GpuDeviceCtx};
239 use openvm_stark_backend::prover::AirProvingContext;
240
241 use super::*;
242 use crate::{
243 batch_constraint::cuda_abi::eq_3b_tracegen,
244 cuda::{preflight::PreflightGpu, to_device_or_nullptr_on},
245 tracegen::ModuleChip,
246 utils::MultiVecWithBounds,
247 };
248
249 impl ModuleChip<GpuBackend> for Eq3bTraceGenerator {
250 type Ctx<'a> = (
251 &'a MultiStarkVerifyingKey<BabyBearPoseidon2Config>,
252 &'a Eq3bBlob,
253 &'a [PreflightGpu],
254 &'a GpuDeviceCtx,
255 );
256
257 #[tracing::instrument(name = "generate_trace", level = "trace", skip_all)]
258 fn generate_proving_ctx(
259 &self,
260 ctx: &Self::Ctx<'_>,
261 required_height: Option<usize>,
262 ) -> Option<AirProvingContext<GpuBackend>> {
263 let (child_vk, blob, preflights, device_ctx) = ctx;
264 debug_assert_eq!(blob.all_stacked_ids.num_proofs(), preflights.len());
265
266 let num_proofs = preflights.len();
267 let width = Eq3bColumns::<F>::width();
268 let l_skip = child_vk.inner.params.l_skip;
269
270 let mut device_records =
271 MultiVecWithBounds::<_, 1>::with_capacity(blob.all_stacked_ids.len());
272 let mut record_idxs = MultiVecWithBounds::<_, 1>::with_capacity(blob.record_idxs.len());
273
274 let mut rows_per_proof_bounds = Vec::with_capacity(num_proofs + 1);
275 rows_per_proof_bounds.push(0);
276
277 let mut n_logups = Vec::with_capacity(num_proofs);
278 let mut all_xi = MultiVecWithBounds::<_, 1>::new();
279
280 let mut num_valid_rows = 0usize;
281
282 for (pidx, preflight) in preflights.iter().enumerate() {
283 let records = &blob.all_stacked_ids[pidx];
284
285 device_records.extend(records.iter().cloned());
286 device_records.close_level(0);
287 record_idxs.extend(blob.record_idxs[pidx].iter().cloned());
288 record_idxs.close_level(0);
289
290 let n_logup = preflight.cpu.proof_shape.n_logup;
291 n_logups.push(n_logup);
292
293 let one_height = n_logup + 1;
294 let rows_for_proof = records.iter().fold(0, |acc, record| {
295 acc + if record.no_interactions {
296 1
297 } else {
298 one_height
299 }
300 });
301 num_valid_rows += rows_for_proof;
302 rows_per_proof_bounds.push(num_valid_rows);
303
304 all_xi.extend(preflight.cpu.batch_constraint.xi.iter().cloned());
305 all_xi.close_level(0);
306 }
307
308 let height = if let Some(height) = required_height {
309 if height < num_valid_rows {
310 return None;
311 }
312 height
313 } else {
314 num_valid_rows.max(1).next_power_of_two()
315 };
316 let d_trace = DeviceMatrix::with_capacity_on(height, width, device_ctx);
317
318 let d_records = to_device_or_nullptr_on(&device_records.data, device_ctx).unwrap();
319 let d_record_bounds = device_records.bounds[0].to_device_on(device_ctx).unwrap();
320 let d_record_idxs = to_device_or_nullptr_on(&record_idxs.data, device_ctx).unwrap();
321 let d_record_idxs_bounds = record_idxs.bounds[0].to_device_on(device_ctx).unwrap();
322 let d_rows_per_proof_bounds = rows_per_proof_bounds.to_device_on(device_ctx).unwrap();
323 let d_n_logups = n_logups.to_device_on(device_ctx).unwrap();
324 let d_xis = to_device_or_nullptr_on(&all_xi.data, device_ctx).unwrap();
325 let d_xi_bounds = all_xi.bounds[0].to_device_on(device_ctx).unwrap();
326
327 unsafe {
328 eq_3b_tracegen(
329 d_trace.buffer(),
330 num_valid_rows,
331 height,
332 num_proofs,
333 l_skip,
334 &d_records,
335 &d_record_bounds,
336 &d_record_idxs,
337 &d_record_idxs_bounds,
338 &d_rows_per_proof_bounds,
339 &d_n_logups,
340 &d_xis,
341 &d_xi_bounds,
342 device_ctx.stream.as_raw(),
343 )
344 .unwrap();
345 }
346
347 Some(AirProvingContext::simple_no_pis(d_trace))
348 }
349 }
350}