1use core::iter::zip;
2use std::sync::Arc;
3
4use itertools::Itertools;
5use openvm_cpu_backend::CpuBackend;
6use openvm_stark_backend::{
7 air_builders::symbolic::{symbolic_variable::Entry, SymbolicExpressionNode},
8 keygen::types::MultiStarkVerifyingKey,
9 poly_common::{eval_eq_sharp_uni, eval_eq_uni, eval_eq_uni_at_one},
10 proof::{column_openings_by_rot, BatchConstraintProof, Proof},
11 prover::{
12 AirProvingContext, ColMajorMatrix, CommittedTraceData, DeviceDataTransporter,
13 TraceCommitter,
14 },
15 AirRef, FiatShamirTranscript, StarkEngine, StarkProtocolConfig, TranscriptHistory,
16};
17use openvm_stark_sdk::config::baby_bear_poseidon2::{BabyBearPoseidon2Config, EF, F};
18use p3_baby_bear::BabyBear;
19use p3_field::{Field, PrimeCharacteristicRing, TwoAdicField};
20use p3_matrix::dense::RowMajorMatrix;
21use p3_maybe_rayon::prelude::{IntoParallelRefIterator, ParallelIterator};
22use strum::{EnumCount, EnumDiscriminants};
23
24use crate::{
25 batch_constraint::{
26 bus::{
27 BatchConstraintConductorBus, ConstraintsFoldingBus, Eq3bBus, EqNOuterBus,
28 EqNegInternalBus, EqSharpUniBus, EqZeroNBus, ExpressionClaimBus,
29 InteractionsFoldingBus, SumcheckClaimBus, SymbolicExpressionBus,
30 UnivariateSumcheckInputBus,
31 },
32 eq_airs::{
33 generate_eq_sharp_uni_blob, Eq3bAir, Eq3bBlob, EqNegAir, EqNegTraceGenerator, EqNsAir,
34 EqSharpUniAir, EqSharpUniBlob, EqSharpUniReceiverAir, EqUniAir,
35 },
36 expr_eval::{
37 CachedTraceRecord, ConstraintsFoldingAir, ConstraintsFoldingBlob, DagCommitSubAir,
38 InteractionsFoldingAir, InteractionsFoldingBlob, SymbolicExpressionAir,
39 },
40 expression_claim::{
41 generate_expression_claim_blob, ExpressionClaimAir, ExpressionClaimBlob,
42 ExpressionClaimTraceGenerator,
43 },
44 fractions_folder::{FractionsFolderAir, FractionsFolderTraceGenerator},
45 sumcheck::{
46 multilinear::MultilinearSumcheckTraceGenerator,
47 univariate::UnivariateSumcheckTraceGenerator, MultilinearSumcheckAir,
48 UnivariateSumcheckAir,
49 },
50 },
51 bus::{
52 AirPresenceBus, AirShapeBus, BatchConstraintModuleBus, ColumnClaimsBus,
53 ConstraintSumcheckRandomnessBus, ConstraintsFoldingInputBus, Eq3bShapeBus,
54 EqNegBaseRandBus, EqNegResultBus, EqNsNLogupMaxBus, ExpressionClaimNMaxBus,
55 FractionFolderInputBus, HyperdimBus, InteractionsFoldingInputBus, NLiftBus,
56 PublicValuesBus, SelHypercubeBus, SelUniBus, StackingModuleBus, TranscriptBus,
57 XiRandomnessBus,
58 },
59 primitives::{bus::PowerCheckerBus, pow::PowerCheckerCpuTraceGenerator},
60 system::{
61 AirModule, BatchConstraintPreflight, BusIndexManager, BusInventory, GlobalCtxCpu,
62 Preflight, TraceGenModule, POW_CHECKER_HEIGHT,
63 },
64 tracegen::{ModuleChip, RowMajorChip, StandardTracegenCtx},
65 utils::MultiVecWithBounds,
66};
67
68pub mod bus;
69pub mod eq_airs;
70pub mod expr_eval;
71pub mod expression_claim;
72pub mod fractions_folder;
73pub mod sumcheck;
74
75#[cfg(feature = "cuda")]
76mod cuda_abi;
77#[cfg(feature = "cuda")]
78mod cuda_utils;
79
80pub(crate) const LOCAL_SYMBOLIC_EXPRESSION_AIR_IDX: usize = 0;
82
83pub fn commit_child_vk<E>(
84 engine: &E,
85 child_vk: &MultiStarkVerifyingKey<BabyBearPoseidon2Config>,
86 has_cached: bool,
87) -> CommittedTraceData<E::PB>
88where
89 E: StarkEngine,
90 E::SC: StarkProtocolConfig<F = F>,
91{
92 let cached_trace_record = expr_eval::build_cached_trace_record(child_vk, has_cached);
93 let cached_trace = expr_eval::generate_symbolic_expr_cached_trace(&cached_trace_record);
94 let cached_trace = ColMajorMatrix::from_row_major(&cached_trace);
95 let cached_trace = engine.device().transport_matrix_to_device(&cached_trace);
96 let (commitment, data) = engine.device().commit(&[&cached_trace]).unwrap();
97 CommittedTraceData {
98 commitment,
99 data: Arc::new(data),
100 trace: cached_trace,
101 }
102}
103
104pub struct BatchConstraintModule {
105 transcript_bus: TranscriptBus,
106 constraint_sumcheck_randomness_bus: ConstraintSumcheckRandomnessBus,
107 xi_randomness_bus: XiRandomnessBus,
108 gkr_claim_bus: BatchConstraintModuleBus,
109 constraints_folding_input_bus: ConstraintsFoldingInputBus,
110 interactions_folding_input_bus: InteractionsFoldingInputBus,
111 fraction_folder_input_bus: FractionFolderInputBus,
112 univariate_sumcheck_input_bus: UnivariateSumcheckInputBus,
113 stacking_module_bus: StackingModuleBus,
114 column_opening_bus: ColumnClaimsBus,
115 air_shape_bus: AirShapeBus,
116 air_presence_bus: AirPresenceBus,
117 hyperdim_bus: HyperdimBus,
118 public_values_bus: PublicValuesBus,
119 sel_uni_bus: SelUniBus,
120 eq_n_outer_bus: EqNOuterBus,
121
122 batch_constraint_conductor_bus: BatchConstraintConductorBus,
123 sumcheck_bus: SumcheckClaimBus,
124 expression_claim_n_max_bus: ExpressionClaimNMaxBus,
125 n_lift_bus: NLiftBus,
126 eq_n_logup_n_max_bus: EqNsNLogupMaxBus,
127 eq_3b_shape_bus: Eq3bShapeBus,
128
129 zero_n_bus: EqZeroNBus,
130 eq_sharp_uni_bus: EqSharpUniBus,
131 eq_3b_bus: Eq3bBus,
132 sel_hypercube_bus: SelHypercubeBus,
133 eq_neg_result_bus: EqNegResultBus,
134 eq_neg_base_rand_bus: EqNegBaseRandBus,
135 eq_neg_internal_bus: EqNegInternalBus,
136
137 symbolic_expression_bus: SymbolicExpressionBus,
138 expression_claim_bus: ExpressionClaimBus,
139 interactions_folding_bus: InteractionsFoldingBus,
140 constraints_folding_bus: ConstraintsFoldingBus,
141 power_checker_bus: PowerCheckerBus,
142
143 l_skip: usize,
144 max_constraint_degree: usize,
145
146 max_num_proofs: usize,
147 pub(crate) has_cached: bool,
148}
149
150impl BatchConstraintModule {
151 pub fn new(
152 child_vk: &MultiStarkVerifyingKey<BabyBearPoseidon2Config>,
153 b: &mut BusIndexManager,
154 bus_inventory: BusInventory,
155 max_num_proofs: usize,
156 has_cached: bool,
157 ) -> Self {
158 let l_skip = child_vk.inner.params.l_skip;
159 let max_constraint_degree = child_vk.max_constraint_degree();
160 BatchConstraintModule {
161 transcript_bus: bus_inventory.transcript_bus,
162 constraint_sumcheck_randomness_bus: bus_inventory.constraint_randomness_bus,
163 xi_randomness_bus: bus_inventory.xi_randomness_bus,
164 gkr_claim_bus: bus_inventory.bc_module_bus,
165 constraints_folding_input_bus: bus_inventory.constraints_folding_input_bus,
166 interactions_folding_input_bus: bus_inventory.interactions_folding_input_bus,
167 fraction_folder_input_bus: bus_inventory.fraction_folder_input_bus,
168 stacking_module_bus: bus_inventory.stacking_module_bus,
169 column_opening_bus: bus_inventory.column_claims_bus,
170 air_shape_bus: bus_inventory.air_shape_bus,
171 air_presence_bus: bus_inventory.air_presence_bus,
172 hyperdim_bus: bus_inventory.hyperdim_bus,
173 public_values_bus: bus_inventory.public_values_bus,
174 sel_uni_bus: bus_inventory.sel_uni_bus,
175 eq_neg_base_rand_bus: bus_inventory.eq_neg_base_rand_bus,
176 eq_neg_result_bus: bus_inventory.eq_neg_result_bus,
177 expression_claim_n_max_bus: bus_inventory.expression_claim_n_max_bus,
178 n_lift_bus: bus_inventory.n_lift_bus,
179 eq_n_logup_n_max_bus: bus_inventory.eq_n_logup_n_max_bus,
180 eq_3b_shape_bus: bus_inventory.eq_3b_shape_bus,
181 batch_constraint_conductor_bus: BatchConstraintConductorBus::new(b.new_bus_idx()),
182 univariate_sumcheck_input_bus: UnivariateSumcheckInputBus::new(b.new_bus_idx()),
183 sumcheck_bus: SumcheckClaimBus::new(b.new_bus_idx()),
184
185 zero_n_bus: EqZeroNBus::new(b.new_bus_idx()),
186 eq_sharp_uni_bus: EqSharpUniBus::new(b.new_bus_idx()),
187 eq_3b_bus: Eq3bBus::new(b.new_bus_idx()),
188 eq_neg_internal_bus: EqNegInternalBus::new(b.new_bus_idx()),
189 sel_hypercube_bus: SelHypercubeBus::new(b.new_bus_idx()),
190 eq_n_outer_bus: EqNOuterBus::new(b.new_bus_idx()),
191 symbolic_expression_bus: SymbolicExpressionBus::new(b.new_bus_idx()),
193 expression_claim_bus: ExpressionClaimBus::new(b.new_bus_idx()),
194 interactions_folding_bus: InteractionsFoldingBus::new(b.new_bus_idx()),
195 constraints_folding_bus: ConstraintsFoldingBus::new(b.new_bus_idx()),
196 power_checker_bus: bus_inventory.power_checker_bus,
197 l_skip,
198 max_constraint_degree,
199 max_num_proofs,
200 has_cached,
201 }
202 }
203
204 #[tracing::instrument(level = "trace", skip_all)]
205 pub fn run_preflight<TS>(
206 &self,
207 child_vk: &MultiStarkVerifyingKey<BabyBearPoseidon2Config>,
208 proof: &Proof<BabyBearPoseidon2Config>,
209 preflight: &mut Preflight,
210 ts: &mut TS,
211 ) where
212 TS: FiatShamirTranscript<BabyBearPoseidon2Config> + TranscriptHistory,
213 {
214 let BatchConstraintProof {
215 numerator_term_per_air,
216 denominator_term_per_air,
217 univariate_round_coeffs,
218 sumcheck_round_polys,
219 column_openings,
220 } = &proof.batch_constraint_proof;
221
222 let mut sumcheck_rnd = vec![];
223
224 let mut xi = preflight.gkr.xi.iter().map(|(_, x)| *x).collect_vec();
225 let l_skip = preflight.proof_shape.l_skip;
226 let n_global = preflight.proof_shape.n_global();
227 for _ in xi.len()..(l_skip + n_global) {
228 xi.push(ts.sample_ext());
229 }
230
231 let lambda_tidx = ts.len();
233 let _lambda = ts.sample_ext();
234
235 for (sum_claim_p, sum_claim_q) in zip(numerator_term_per_air, denominator_term_per_air) {
236 ts.observe_ext(*sum_claim_p);
237 ts.observe_ext(*sum_claim_q);
238 }
239 let _mu = ts.sample_ext();
240
241 let tidx_before_univariate = ts.len();
242
243 for coef in univariate_round_coeffs {
245 ts.observe_ext(*coef);
246 }
247 let r0 = ts.sample_ext();
248 sumcheck_rnd.push(r0);
249
250 let tidx_before_multilinear = ts.len();
251
252 for polys in sumcheck_round_polys {
253 for eval in polys {
254 ts.observe_ext(*eval);
255 }
256 let ri = ts.sample_ext();
257 sumcheck_rnd.push(ri);
258 }
259
260 let tidx_before_column_openings = ts.len();
261
262 for (sort_idx, (air_id, _)) in preflight.proof_shape.sorted_trace_vdata.iter().enumerate() {
264 let need_rot = child_vk.inner.per_air[*air_id].params.need_rot;
265 for (col_opening, rot_opening) in
266 column_openings_by_rot(&column_openings[sort_idx][0], need_rot)
267 {
268 ts.observe_ext(col_opening);
269 ts.observe_ext(rot_opening);
270 }
271 }
272
273 for (sort_idx, (air_id, _)) in preflight.proof_shape.sorted_trace_vdata.iter().enumerate() {
274 let need_rot = child_vk.inner.per_air[*air_id].params.need_rot;
275 for part in column_openings[sort_idx].iter().skip(1) {
276 for (col_opening, rot_opening) in column_openings_by_rot(part, need_rot) {
277 ts.observe_ext(col_opening);
278 ts.observe_ext(rot_opening);
279 }
280 }
281 }
282
283 let omega_skip_pows = F::two_adic_generator(l_skip)
284 .powers()
285 .take(1 << l_skip)
286 .collect_vec();
287
288 let mut eq_ns = Vec::with_capacity(preflight.proof_shape.n_max + 1);
289 let mut eq_sharp_ns = Vec::with_capacity(preflight.proof_shape.n_max + 1);
290 let mut eq = eval_eq_uni(l_skip, xi[0], sumcheck_rnd[0]);
291 let mut eq_sharp = eval_eq_sharp_uni(&omega_skip_pows, &xi[..l_skip], sumcheck_rnd[0]);
292 eq_ns.push(eq);
293 eq_sharp_ns.push(eq_sharp);
294 for i in 0..preflight.proof_shape.n_max {
295 let mult = EF::ONE - xi[l_skip + i] - sumcheck_rnd[1 + i]
296 + (xi[l_skip + i] * sumcheck_rnd[1 + i]).double();
297 eq *= mult;
298 eq_sharp *= mult;
299 eq_ns.push(eq);
300 eq_sharp_ns.push(eq_sharp);
301 }
302
303 let mut r_rev_prod = sumcheck_rnd[preflight.proof_shape.n_max];
304 let mut eq_ns_frontloaded = Vec::with_capacity(preflight.proof_shape.n_max + 1);
305 let mut eq_sharp_ns_frontloaded = Vec::with_capacity(preflight.proof_shape.n_max + 1);
306 for i in (0..preflight.proof_shape.n_max).rev() {
309 eq_ns_frontloaded.push(eq_ns[i] * r_rev_prod);
310 eq_sharp_ns_frontloaded.push(eq_sharp_ns[i] * r_rev_prod);
311 r_rev_prod *= sumcheck_rnd[i];
312 }
313 eq_ns_frontloaded.reverse();
314 eq_sharp_ns_frontloaded.reverse();
315 eq_ns_frontloaded.push(eq_ns[preflight.proof_shape.n_max]);
316 eq_sharp_ns_frontloaded.push(eq_sharp_ns[preflight.proof_shape.n_max]);
317
318 preflight.batch_constraint = BatchConstraintPreflight {
319 lambda_tidx,
320 tidx_before_univariate,
321 tidx_before_multilinear,
322 tidx_before_column_openings,
323 post_tidx: ts.len(),
324 xi,
325 sumcheck_rnd,
326 eq_ns,
327 eq_sharp_ns,
328 eq_ns_frontloaded,
329 eq_sharp_ns_frontloaded,
330 }
331 }
332}
333
334impl AirModule for BatchConstraintModule {
335 fn num_airs(&self) -> usize {
336 BatchConstraintModuleChipDiscriminants::COUNT
337 }
338
339 fn airs<SC: StarkProtocolConfig<F = BabyBear>>(&self) -> Vec<AirRef<SC>> {
340 let l_skip = self.l_skip;
341
342 let symbolic_expression_air = SymbolicExpressionAir::<BabyBear> {
343 expr_bus: self.symbolic_expression_bus,
344 air_shape_bus: self.air_shape_bus,
345 air_presence_bus: self.air_presence_bus,
346 column_claims_bus: self.column_opening_bus,
347 interactions_folding_bus: self.interactions_folding_bus,
348 constraints_folding_bus: self.constraints_folding_bus,
349 hyperdim_bus: self.hyperdim_bus,
350 public_values_bus: self.public_values_bus,
351 sel_hypercube_bus: self.sel_hypercube_bus,
352 sel_uni_bus: self.sel_uni_bus,
353 cnt_proofs: self.max_num_proofs,
354 dag_commit_subair: (!self.has_cached).then_some(Arc::new(DagCommitSubAir::new())),
355 };
356 let fraction_folder_air = FractionsFolderAir {
357 transcript_bus: self.transcript_bus,
358 univariate_sumcheck_input_bus: self.univariate_sumcheck_input_bus,
359 fraction_folder_input_bus: self.fraction_folder_input_bus,
360 sumcheck_bus: self.sumcheck_bus,
361 mu_bus: self.batch_constraint_conductor_bus,
362 gkr_claim_bus: self.gkr_claim_bus,
363 };
364 let sumcheck_uni_air = UnivariateSumcheckAir {
365 l_skip,
366 univariate_deg: (self.max_constraint_degree + 1) * ((1 << l_skip) - 1),
367 univariate_sumcheck_input_bus: self.univariate_sumcheck_input_bus,
368 stacking_module_bus: self.stacking_module_bus,
369 claim_bus: self.sumcheck_bus,
370 transcript_bus: self.transcript_bus,
371 randomness_bus: self.constraint_sumcheck_randomness_bus,
372 batch_constraint_conductor_bus: self.batch_constraint_conductor_bus,
373 };
374 let sumcheck_lin_air = MultilinearSumcheckAir {
375 max_constraint_degree: self.max_constraint_degree,
376 claim_bus: self.sumcheck_bus,
377 transcript_bus: self.transcript_bus,
378 randomness_bus: self.constraint_sumcheck_randomness_bus,
379 batch_constraint_conductor_bus: self.batch_constraint_conductor_bus,
380 stacking_module_bus: self.stacking_module_bus,
381 };
382 let eq_ns_air = EqNsAir {
383 zero_n_bus: self.zero_n_bus,
384 xi_bus: self.xi_randomness_bus,
385 r_xi_bus: self.batch_constraint_conductor_bus,
386 sel_hypercube_bus: self.sel_hypercube_bus,
387 eq_n_outer_bus: self.eq_n_outer_bus,
388 eq_n_logup_n_max_bus: self.eq_n_logup_n_max_bus,
389 l_skip,
390 };
391 let eq_3b_air = Eq3bAir {
392 eq_3b_bus: self.eq_3b_bus,
393 eq_3b_shape_bus: self.eq_3b_shape_bus,
394 batch_constraint_conductor_bus: self.batch_constraint_conductor_bus,
395 l_skip,
396 };
397 let eq_sharp_uni_air = EqSharpUniAir {
398 xi_bus: self.xi_randomness_bus,
399 eq_bus: self.eq_sharp_uni_bus,
400 batch_constraint_conductor_bus: self.batch_constraint_conductor_bus,
401 l_skip,
402 canonical_inverse_generator: F::two_adic_generator(l_skip).inverse(),
403 };
404 let eq_sharp_uni_receiver_air = EqSharpUniReceiverAir {
405 r_bus: self.batch_constraint_conductor_bus,
406 eq_bus: self.eq_sharp_uni_bus,
407 zero_n_bus: self.zero_n_bus,
408 l_skip,
409 };
410 let eq_uni_air = EqUniAir {
411 r_xi_bus: self.batch_constraint_conductor_bus,
412 zero_n_bus: self.zero_n_bus,
413 l_skip,
414 };
415 let eq_neg_air = EqNegAir {
416 result_bus: self.eq_neg_result_bus,
417 base_rand_bus: self.eq_neg_base_rand_bus,
418 internal_bus: self.eq_neg_internal_bus,
419 sel_uni_bus: self.sel_uni_bus,
420 l_skip: self.l_skip,
421 };
422 let expression_claim_air = ExpressionClaimAir {
423 expression_claim_n_max_bus: self.expression_claim_n_max_bus,
424 expr_claim_bus: self.expression_claim_bus,
425 mu_bus: self.batch_constraint_conductor_bus,
426 sumcheck_claim_bus: self.sumcheck_bus,
427 eq_n_outer_bus: self.eq_n_outer_bus,
428 pow_checker_bus: self.power_checker_bus,
429 hyperdim_bus: self.hyperdim_bus,
430 };
431 let interactions_folding_air = InteractionsFoldingAir {
432 transcript_bus: self.transcript_bus,
433 air_shape_bus: self.air_shape_bus,
434 interaction_bus: self.interactions_folding_bus,
435 interactions_folding_input_bus: self.interactions_folding_input_bus,
436 expression_claim_bus: self.expression_claim_bus,
437 eq_3b_bus: self.eq_3b_bus,
438 };
439 let constraints_folding_air = ConstraintsFoldingAir {
440 transcript_bus: self.transcript_bus,
441 constraint_bus: self.constraints_folding_bus,
442 expression_claim_bus: self.expression_claim_bus,
443 eq_n_outer_bus: self.eq_n_outer_bus,
444 n_lift_bus: self.n_lift_bus,
445 air_shape_bus: self.air_shape_bus,
446 constraints_folding_input_bus: self.constraints_folding_input_bus,
447 };
448 vec![
450 Arc::new(symbolic_expression_air) as AirRef<_>,
451 Arc::new(fraction_folder_air) as AirRef<_>,
452 Arc::new(sumcheck_uni_air) as AirRef<_>,
453 Arc::new(sumcheck_lin_air) as AirRef<_>,
454 Arc::new(eq_ns_air) as AirRef<_>,
455 Arc::new(eq_3b_air) as AirRef<_>,
456 Arc::new(eq_sharp_uni_air) as AirRef<_>,
457 Arc::new(eq_sharp_uni_receiver_air) as AirRef<_>,
458 Arc::new(eq_uni_air) as AirRef<_>,
459 Arc::new(expression_claim_air) as AirRef<_>,
460 Arc::new(interactions_folding_air) as AirRef<_>,
461 Arc::new(constraints_folding_air) as AirRef<_>,
462 Arc::new(eq_neg_air) as AirRef<_>,
463 ]
464 }
465}
466
467pub(crate) struct BatchConstraintBlob {
468 pub expr_evals: MultiVecWithBounds<EF, 2>,
471 pub selector_counts: MultiVecWithBounds<SelectorCount, 1>,
473
474 pub eq_3b_blob: Eq3bBlob,
475 pub eq_sharp_uni_blob: EqSharpUniBlob,
476}
477
478#[derive(Clone, Copy, Debug, Default)]
479pub struct SelectorCount {
480 pub first: usize,
481 pub last: usize,
482 pub transition: usize,
483}
484
485impl BatchConstraintBlob {
486 pub fn new(
487 child_vk: &MultiStarkVerifyingKey<BabyBearPoseidon2Config>,
488 proofs: &[&Proof<BabyBearPoseidon2Config>],
489 preflights: &[&Preflight],
490 ) -> Self {
491 let child_vk = &child_vk.inner;
492 let params = &child_vk.params;
493
494 let mut expr_evals_per_proof = MultiVecWithBounds::new();
495 let mut eq_r_one_counts_per_proof = MultiVecWithBounds::new();
496 for (proof, preflight) in zip(proofs, preflights) {
497 let rs = &preflight.batch_constraint.sumcheck_rnd;
498
499 let (&rs_0, rs_rest) = rs.split_first().unwrap();
500 let mut is_first_row_by_log_height = vec![];
501 let mut is_last_row_by_log_height = vec![];
502 let n_max = preflight.proof_shape.n_max;
503 let mut selector_counts =
504 vec![SelectorCount::default(); child_vk.params.l_skip + n_max + 1];
505
506 let omega = F::two_adic_generator(params.l_skip);
507 for log_height in 0..=params.l_skip {
508 is_first_row_by_log_height.push(eval_eq_uni_at_one(
509 log_height,
510 rs_0.exp_power_of_2(params.l_skip - log_height),
511 ));
512 is_last_row_by_log_height.push(eval_eq_uni_at_one(
513 log_height,
514 (rs_0 * omega).exp_power_of_2(params.l_skip - log_height),
515 ));
516 }
517 for (i, &r) in rs_rest.iter().enumerate() {
518 is_first_row_by_log_height
519 .push(is_first_row_by_log_height[params.l_skip + i] * (EF::ONE - r));
520 is_last_row_by_log_height.push(is_last_row_by_log_height[params.l_skip + i] * r);
521 }
522
523 let mut expr_evals_per_air = vec![];
524 for (air_idx, vk) in child_vk.per_air.iter().enumerate() {
525 if proof.trace_vdata[air_idx].is_none() {
526 expr_evals_per_air.push(vec![]);
527 continue;
528 }
529
530 let need_rot = child_vk.per_air[air_idx].params.need_rot;
531 let openings_per_col = if need_rot { 2 } else { 1 };
532 let openings = &proof.batch_constraint_proof.column_openings;
533 let (sorted_idx, vdata) = preflight
534 .proof_shape
535 .sorted_trace_vdata
536 .iter()
537 .enumerate()
538 .find_map(|(sorted_idx, (idx, vdata))| {
539 if air_idx == *idx {
540 Some((sorted_idx, vdata))
541 } else {
542 None
543 }
544 })
545 .unwrap();
546
547 let constraints = &vk.symbolic_constraints.constraints;
548 let mut expr_evals =
549 vec![EF::ZERO; constraints.nodes.len() + vk.unused_variables.len()];
550 let log_height = proof.trace_vdata[air_idx].as_ref().unwrap().log_height;
551
552 for (node_idx, node) in constraints.nodes.iter().enumerate() {
553 match node {
554 SymbolicExpressionNode::Variable(var) => match var.entry {
555 Entry::Preprocessed { offset } => {
556 debug_assert!(offset < openings_per_col);
557 expr_evals[node_idx] =
558 openings[sorted_idx][1][var.index * openings_per_col + offset];
559 }
560 Entry::Main { part_index, offset } => {
561 let part = vk.dag_main_part_index_to_commit_index(part_index);
562 debug_assert!(offset < openings_per_col);
563 expr_evals[node_idx] = openings[sorted_idx][part]
564 [var.index * openings_per_col + offset];
565 }
566 Entry::Public => {
567 expr_evals[node_idx] =
568 EF::from(proof.public_values[air_idx][var.index]);
569 }
570 Entry::Challenge => unreachable!(),
571 },
572 SymbolicExpressionNode::IsFirstRow => {
573 expr_evals[node_idx] = is_first_row_by_log_height[vdata.log_height];
574 selector_counts[log_height].first += 1;
575 }
576 SymbolicExpressionNode::IsLastRow => {
577 expr_evals[node_idx] = is_last_row_by_log_height[vdata.log_height];
578 selector_counts[log_height].last += 1;
579 }
580 SymbolicExpressionNode::IsTransition => {
581 expr_evals[node_idx] =
582 EF::ONE - is_last_row_by_log_height[vdata.log_height];
583 selector_counts[log_height].transition += 1;
584 }
585 SymbolicExpressionNode::Constant(val) => {
586 expr_evals[node_idx] = EF::from(*val);
587 }
588 SymbolicExpressionNode::Add {
589 left_idx,
590 right_idx,
591 degree_multiple: _,
592 } => {
593 debug_assert!(*left_idx < node_idx);
594 debug_assert!(*right_idx < node_idx);
595 expr_evals[node_idx] = expr_evals[*left_idx] + expr_evals[*right_idx];
596 }
597 SymbolicExpressionNode::Sub {
598 left_idx,
599 right_idx,
600 degree_multiple: _,
601 } => {
602 debug_assert!(*left_idx < node_idx);
603 debug_assert!(*right_idx < node_idx);
604 expr_evals[node_idx] = expr_evals[*left_idx] - expr_evals[*right_idx];
605 }
606 SymbolicExpressionNode::Neg {
607 idx,
608 degree_multiple: _,
609 } => {
610 debug_assert!(*idx < node_idx);
611 expr_evals[node_idx] = -expr_evals[*idx];
612 }
613 SymbolicExpressionNode::Mul {
614 left_idx,
615 right_idx,
616 degree_multiple: _,
617 } => {
618 debug_assert!(*left_idx < node_idx);
619 debug_assert!(*right_idx < node_idx);
620 expr_evals[node_idx] = expr_evals[*left_idx] * expr_evals[*right_idx];
621 }
622 };
623 }
624 let mut node_idx = constraints.nodes.len();
625 for unused_var in &vk.unused_variables {
626 match unused_var.entry {
627 Entry::Preprocessed { offset } => {
628 debug_assert!(offset < openings_per_col);
629 expr_evals[node_idx] = openings[sorted_idx][1]
630 [unused_var.index * openings_per_col + offset];
631 }
632 Entry::Main { part_index, offset } => {
633 let part = vk.dag_main_part_index_to_commit_index(part_index);
634 debug_assert!(offset < openings_per_col);
635 expr_evals[node_idx] = openings[sorted_idx][part]
636 [unused_var.index * openings_per_col + offset];
637 }
638 Entry::Public | Entry::Challenge => {
639 unreachable!()
640 }
641 }
642 node_idx += 1;
643 }
644 expr_evals_per_air.push(expr_evals);
645 }
646 for v in expr_evals_per_air {
647 expr_evals_per_proof.extend(v);
648 expr_evals_per_proof.close_level(1);
649 }
650 expr_evals_per_proof.close_level(0);
651 eq_r_one_counts_per_proof.extend(selector_counts);
652 eq_r_one_counts_per_proof.close_level(0);
653 }
654 let eq_3b_blob = eq_airs::generate_eq_3b_blob(child_vk, preflights);
655 let eq_sharp_uni_blob = generate_eq_sharp_uni_blob(child_vk, preflights);
656 Self {
657 expr_evals: expr_evals_per_proof,
658 selector_counts: eq_r_one_counts_per_proof,
659 eq_3b_blob,
660 eq_sharp_uni_blob,
661 }
662 }
663}
664
665pub(crate) struct BatchConstraintBlobCpu {
666 pub common_blob: BatchConstraintBlob,
667 pub cf_blob: Option<ConstraintsFoldingBlob>,
668 pub if_blob: Option<InteractionsFoldingBlob>,
669 pub expr_claim_blob: ExpressionClaimBlob,
670}
671
672impl BatchConstraintBlobCpu {
673 #[tracing::instrument(name = "generate_blob", skip_all)]
674 pub fn new(
675 child_vk: &MultiStarkVerifyingKey<BabyBearPoseidon2Config>,
676 proofs: &[Proof<BabyBearPoseidon2Config>],
677 preflights: &[Preflight],
678 ) -> Self {
679 let proofs = proofs.iter().collect_vec();
680 let preflights = preflights.iter().collect_vec();
681 let common_blob = BatchConstraintBlob::new(child_vk, &proofs, &preflights);
682 let cf_blob =
683 ConstraintsFoldingBlob::new(&child_vk.inner, &common_blob.expr_evals, &preflights);
684 let if_blob = InteractionsFoldingBlob::new(
685 &child_vk.inner,
686 &common_blob.expr_evals,
687 &common_blob.eq_3b_blob,
688 &preflights,
689 );
690 let expr_claim_blob =
691 generate_expression_claim_blob(&cf_blob.folded_claims, &if_blob.folded_claims);
692 Self {
693 common_blob,
694 cf_blob: Some(cf_blob),
695 if_blob: Some(if_blob),
696 expr_claim_blob,
697 }
698 }
699}
700
701impl<SC: StarkProtocolConfig<F = F>> TraceGenModule<GlobalCtxCpu, CpuBackend<SC>>
702 for BatchConstraintModule
703{
704 type ModuleSpecificCtx<'a> = (
705 Option<&'a CachedTraceRecord>,
706 Arc<PowerCheckerCpuTraceGenerator<2, POW_CHECKER_HEIGHT>>,
707 );
708
709 #[tracing::instrument(skip_all)]
713 fn generate_proving_ctxs(
714 &self,
715 child_vk: &MultiStarkVerifyingKey<BabyBearPoseidon2Config>,
716 proofs: &[Proof<BabyBearPoseidon2Config>],
717 preflights: &[Preflight],
718 ctx: &Self::ModuleSpecificCtx<'_>,
719 required_heights: Option<&[usize]>,
720 ) -> Option<Vec<AirProvingContext<CpuBackend<SC>>>> {
721 let blob = BatchConstraintBlobCpu::new(child_vk, proofs, preflights);
722 let pow_checker = ctx.1.clone();
723 let ctx = (
724 StandardTracegenCtx {
725 vk: child_vk,
726 proofs: &proofs.iter().collect_vec(),
727 preflights: &preflights.iter().collect_vec(),
728 },
729 blob,
730 ctx.0,
731 );
732
733 let chips = [
734 BatchConstraintModuleChip::SymbolicExpression {
735 max_num_proofs: self.max_num_proofs,
736 has_cached: self.has_cached,
737 },
738 BatchConstraintModuleChip::FractionsFolder,
739 BatchConstraintModuleChip::SumcheckUni,
740 BatchConstraintModuleChip::SumcheckLin,
741 BatchConstraintModuleChip::EqNs,
742 BatchConstraintModuleChip::Eq3b,
743 BatchConstraintModuleChip::EqSharpUni,
744 BatchConstraintModuleChip::EqSharpUniReceiver,
745 BatchConstraintModuleChip::EqUni,
746 BatchConstraintModuleChip::ExpressionClaim { pow_checker },
747 BatchConstraintModuleChip::InteractionsFolding,
748 BatchConstraintModuleChip::ConstraintsFolding,
749 BatchConstraintModuleChip::EqNeg,
750 ];
751
752 let span = tracing::Span::current();
753 chips
754 .par_iter()
755 .map(|chip| {
756 let _guard = span.enter();
757 chip.generate_proving_ctx(
758 &ctx,
759 required_heights.map(|heights| heights[chip.index()]),
760 )
761 })
762 .collect::<Vec<_>>()
763 .into_iter()
764 .collect()
765 }
766}
767
768#[derive(strum_macros::Display, EnumDiscriminants)]
770#[strum_discriminants(derive(strum_macros::EnumCount))]
771#[strum_discriminants(repr(usize))]
772enum BatchConstraintModuleChip {
773 SymbolicExpression {
774 max_num_proofs: usize,
775 has_cached: bool,
776 },
777 FractionsFolder,
778 SumcheckUni,
779 SumcheckLin,
780 EqNs,
781 Eq3b,
782 EqSharpUni,
783 EqSharpUniReceiver,
784 EqUni,
785 ExpressionClaim {
786 pow_checker: Arc<PowerCheckerCpuTraceGenerator<2, POW_CHECKER_HEIGHT>>,
787 },
788 InteractionsFolding,
789 ConstraintsFolding,
790 EqNeg,
791}
792
793impl BatchConstraintModuleChip {
794 fn index(&self) -> usize {
795 BatchConstraintModuleChipDiscriminants::from(self) as usize
796 }
797}
798
799impl RowMajorChip<F> for BatchConstraintModuleChip {
800 type Ctx<'a> = (
801 StandardTracegenCtx<'a>,
802 BatchConstraintBlobCpu,
803 Option<&'a CachedTraceRecord>,
804 );
805
806 #[tracing::instrument(
807 name = "wrapper.generate_trace",
808 level = "trace",
809 skip_all,
810 fields(air = %self)
811 )]
812 fn generate_trace(
813 &self,
814 ctx: &Self::Ctx<'_>,
815 required_height: Option<usize>,
816 ) -> Option<RowMajorMatrix<F>> {
817 use BatchConstraintModuleChip::*;
818 let child_vk = ctx.0.vk;
819 let proofs = ctx.0.proofs;
820 let preflights = ctx.0.preflights;
821 let blob = &ctx.1;
822 let cached_trace_record = ctx.2;
823 match self {
824 FractionsFolder => {
825 FractionsFolderTraceGenerator.generate_trace(&ctx.0, required_height)
826 }
827 SumcheckUni => UnivariateSumcheckTraceGenerator.generate_trace(&ctx.0, required_height),
828 SumcheckLin => {
829 MultilinearSumcheckTraceGenerator.generate_trace(&ctx.0, required_height)
830 }
831 EqNs => eq_airs::EqNsTraceGenerator.generate_trace(
832 &(child_vk, preflights, &blob.common_blob.selector_counts),
833 required_height,
834 ),
835 Eq3b => eq_airs::Eq3bTraceGenerator.generate_trace(
836 &(child_vk, &blob.common_blob.eq_3b_blob, preflights),
837 required_height,
838 ),
839 EqSharpUni => eq_airs::EqSharpUniTraceGenerator.generate_trace(
840 &(child_vk, &blob.common_blob.eq_sharp_uni_blob, preflights),
841 required_height,
842 ),
843 EqSharpUniReceiver => eq_airs::EqSharpUniReceiverTraceGenerator.generate_trace(
844 &(child_vk, &blob.common_blob.eq_sharp_uni_blob, preflights),
845 required_height,
846 ),
847 EqUni => eq_airs::EqUniTraceGenerator.generate_trace(&ctx.0, required_height),
848 SymbolicExpression {
849 max_num_proofs,
850 has_cached,
851 } => expr_eval::SymbolicExpressionTraceGenerator {
852 max_num_proofs: *max_num_proofs,
853 has_cached: *has_cached,
854 }
855 .generate_trace(
856 &expr_eval::SymbolicExpressionCtx {
857 vk: child_vk,
858 preflights,
859 expr_evals: &blob.common_blob.expr_evals,
860 cached_trace_record: &cached_trace_record,
861 },
862 required_height,
863 ),
864 ExpressionClaim { pow_checker } => ExpressionClaimTraceGenerator.generate_trace(
865 &expression_claim::ExpressionClaimCtx {
866 blob: &blob.expr_claim_blob,
867 proofs,
868 preflights,
869 pow_checker: pow_checker.as_ref(),
870 },
871 required_height,
872 ),
873 InteractionsFolding => expr_eval::InteractionsFoldingTraceGenerator
874 .generate_trace(&(child_vk, blob, preflights), required_height),
875 ConstraintsFolding => expr_eval::ConstraintsFoldingTraceGenerator.generate_trace(
876 &(blob.cf_blob.as_ref().unwrap(), preflights),
877 required_height,
878 ),
879 EqNeg => EqNegTraceGenerator.generate_trace(
880 &(child_vk, preflights, &blob.common_blob.selector_counts),
881 required_height,
882 ),
883 }
884 }
885}
886
887#[cfg(feature = "cuda")]
888pub mod cuda_tracegen {
889 use openvm_cuda_backend::{data_transporter::transport_matrix_h2d_row, GpuBackend};
890 use openvm_cuda_common::stream::GpuDeviceCtx;
891
892 use super::*;
893 use crate::{
894 batch_constraint::expr_eval::{
895 constraints_folding::cuda::ConstraintsFoldingBlobGpu,
896 interactions_folding::cuda::InteractionsFoldingBlobGpu,
897 },
898 cuda::{preflight::PreflightGpu, proof::ProofGpu, vk::VerifyingKeyGpu, GlobalCtxGpu},
899 tracegen::cuda::StandardTracegenGpuCtx,
900 };
901
902 impl ModuleChip<GpuBackend> for BatchConstraintModuleChip {
903 type Ctx<'a> = (
904 StandardTracegenGpuCtx<'a>,
905 &'a BatchConstraintBlobGpu,
906 Option<&'a CachedTraceRecord>,
907 );
908
909 fn generate_proving_ctx(
910 &self,
911 ctx: &Self::Ctx<'_>,
912 required_height: Option<usize>,
913 ) -> Option<AirProvingContext<GpuBackend>> {
914 use BatchConstraintModuleChip::*;
915 let child_vk = ctx.0.vk;
916 let proofs = ctx.0.proofs;
917 let preflights = ctx.0.preflights;
918 let blob = ctx.1;
919 let cached_trace_record = ctx.2;
920 match self {
921 Eq3b => eq_airs::Eq3bTraceGenerator.generate_proving_ctx(
922 &(
923 &child_vk.cpu,
924 &blob.common_blob.eq_3b_blob,
925 preflights,
926 ctx.0.device_ctx,
927 ),
928 required_height,
929 ),
930 SymbolicExpression {
931 max_num_proofs,
932 has_cached,
933 } => expr_eval::SymbolicExpressionTraceGenerator {
934 max_num_proofs: *max_num_proofs,
935 has_cached: *has_cached,
936 }
937 .generate_proving_ctx(
938 &expr_eval::symbolic_expression::cuda::SymbolicExpressionGpuCtx {
939 vk: &child_vk.cpu,
940 proofs,
941 preflights,
942 expr_evals: &blob.common_blob.expr_evals,
943 cached_trace_record: &cached_trace_record,
944 device_ctx: ctx.0.device_ctx,
945 },
946 required_height,
947 ),
948 InteractionsFolding => expr_eval::InteractionsFoldingTraceGenerator
949 .generate_proving_ctx(
950 &(child_vk, preflights, &blob.if_blob, ctx.0.device_ctx),
951 required_height,
952 ),
953 ConstraintsFolding => expr_eval::ConstraintsFoldingTraceGenerator
954 .generate_proving_ctx(
955 &(child_vk, preflights, &blob.cf_blob, ctx.0.device_ctx),
956 required_height,
957 ),
958 _ => unreachable!(),
959 }
960 }
961 }
962
963 pub(in crate::batch_constraint) struct BatchConstraintBlobGpu {
964 pub common_blob: BatchConstraintBlob,
965 pub cf_blob: ConstraintsFoldingBlobGpu,
966 pub if_blob: InteractionsFoldingBlobGpu,
967 pub expr_claim_blob: ExpressionClaimBlob,
968 }
969
970 impl BatchConstraintBlobGpu {
971 #[tracing::instrument(name = "generate_blob", skip_all)]
972 pub fn new(
973 child_vk: &VerifyingKeyGpu,
974 proofs: &[ProofGpu],
975 preflights: &[PreflightGpu],
976 device_ctx: &GpuDeviceCtx,
977 ) -> Self {
978 let cpu_proofs = proofs.iter().map(|p| &p.cpu).collect_vec();
979 let cpu_preflights = preflights.iter().map(|p| &p.cpu).collect_vec();
980 let common_blob = BatchConstraintBlob::new(&child_vk.cpu, &cpu_proofs, &cpu_preflights);
981 let cf_blob = ConstraintsFoldingBlobGpu::new(
982 child_vk,
983 &common_blob.expr_evals,
984 preflights,
985 device_ctx,
986 );
987 let if_blob = InteractionsFoldingBlobGpu::new(
988 child_vk,
989 &common_blob.expr_evals,
990 &common_blob.eq_3b_blob,
991 preflights,
992 device_ctx,
993 );
994 let expr_claim_blob =
995 generate_expression_claim_blob(&cf_blob.folded_claims, &if_blob.folded_claims);
996 Self {
997 common_blob,
998 cf_blob,
999 if_blob,
1000 expr_claim_blob,
1001 }
1002 }
1003 }
1004
1005 impl TraceGenModule<GlobalCtxGpu, GpuBackend> for BatchConstraintModule {
1006 type ModuleSpecificCtx<'a> = (
1007 Option<&'a CachedTraceRecord>,
1008 Arc<PowerCheckerCpuTraceGenerator<2, POW_CHECKER_HEIGHT>>,
1009 &'a openvm_cuda_common::stream::GpuDeviceCtx,
1010 );
1011
1012 #[tracing::instrument(skip_all)]
1013 fn generate_proving_ctxs(
1014 &self,
1015 child_vk: &VerifyingKeyGpu,
1016 proofs: &[ProofGpu],
1017 preflights: &[PreflightGpu],
1018 module_ctx: &Self::ModuleSpecificCtx<'_>,
1019 required_heights: Option<&[usize]>,
1020 ) -> Option<Vec<AirProvingContext<GpuBackend>>> {
1021 let cached_trace_record = module_ctx.0;
1022 let pow_checker = module_ctx.1.clone();
1023 let device_ctx = module_ctx.2;
1024 let blob = BatchConstraintBlobGpu::new(child_vk, proofs, preflights, device_ctx);
1025 let ctx = (
1026 StandardTracegenGpuCtx {
1027 vk: child_vk,
1028 proofs,
1029 preflights,
1030 device_ctx,
1031 },
1032 &blob,
1033 cached_trace_record,
1034 );
1035
1036 let gpu_chips = [
1038 BatchConstraintModuleChip::SymbolicExpression {
1039 max_num_proofs: self.max_num_proofs,
1040 has_cached: self.has_cached,
1041 },
1042 BatchConstraintModuleChip::Eq3b,
1043 BatchConstraintModuleChip::ConstraintsFolding,
1044 BatchConstraintModuleChip::InteractionsFolding,
1045 ];
1046 let cpu_chips = [
1048 BatchConstraintModuleChip::FractionsFolder,
1049 BatchConstraintModuleChip::SumcheckUni,
1050 BatchConstraintModuleChip::SumcheckLin,
1051 BatchConstraintModuleChip::EqNs,
1052 BatchConstraintModuleChip::EqSharpUni,
1053 BatchConstraintModuleChip::EqSharpUniReceiver,
1054 BatchConstraintModuleChip::EqUni,
1055 BatchConstraintModuleChip::ExpressionClaim { pow_checker },
1056 BatchConstraintModuleChip::EqNeg,
1057 ];
1058 let span = tracing::Span::current();
1059 let indexed_gpu_traces = gpu_chips
1062 .iter()
1063 .map(|chip| {
1064 let _guard = span.enter();
1066 (
1067 chip.index(),
1068 chip.generate_proving_ctx(
1069 &ctx,
1070 required_heights.map(|heights| heights[chip.index()]),
1071 ),
1072 )
1073 })
1074 .collect_vec();
1075
1076 let blob = BatchConstraintBlobCpu {
1077 common_blob: blob.common_blob,
1078 cf_blob: None,
1079 if_blob: None,
1080 expr_claim_blob: blob.expr_claim_blob,
1081 };
1082 let cpu_proofs = proofs.iter().map(|p| &p.cpu).collect_vec();
1083 let cpu_preflights = preflights.iter().map(|p| &p.cpu).collect_vec();
1084 let cpu_ctx = (
1085 StandardTracegenCtx {
1086 vk: &child_vk.cpu,
1087 proofs: &cpu_proofs,
1088 preflights: &cpu_preflights,
1089 },
1090 blob,
1091 cached_trace_record,
1092 );
1093
1094 let indexed_cpu_rm_traces = cpu_chips
1096 .par_iter()
1097 .map(|chip| {
1098 let _guard = span.enter();
1099 (
1100 chip.index(),
1101 chip.generate_trace(
1102 &cpu_ctx,
1103 required_heights.map(|heights| heights[chip.index()]),
1104 ),
1105 )
1106 })
1107 .collect::<Vec<_>>();
1108
1109 let indexed_cpu_gpu_traces = indexed_cpu_rm_traces
1111 .into_iter()
1112 .map(|(idx, trace)| {
1113 (
1114 idx,
1115 trace.map(|m| {
1116 AirProvingContext::simple_no_pis(
1117 transport_matrix_h2d_row(&m, device_ctx).unwrap(),
1118 )
1119 }),
1120 )
1121 })
1122 .collect::<Vec<_>>();
1123
1124 indexed_gpu_traces
1125 .into_iter()
1126 .chain(indexed_cpu_gpu_traces)
1127 .sorted_by(|a, b| a.0.cmp(&b.0))
1128 .map(|(_index, ctx)| ctx)
1129 .collect()
1130 }
1131 }
1132}