1use core::{cmp::min, iter::zip};
2use std::borrow::BorrowMut;
3
4use itertools::Itertools;
5use openvm_circuit_primitives::encoder::Encoder;
6use openvm_stark_backend::{
7 air_builders::symbolic::{symbolic_variable::Entry, SymbolicExpressionNode},
8 keygen::types::MultiStarkVerifyingKey,
9 poly_common::{eval_eq_uni_at_one, Squarable},
10};
11use openvm_stark_sdk::config::baby_bear_poseidon2::{BabyBearPoseidon2Config, D_EF, EF, F};
12use p3_field::{BasedVectorSpace, PrimeCharacteristicRing, PrimeField32, TwoAdicField};
13use p3_matrix::{dense::RowMajorMatrix, Matrix};
14use p3_maybe_rayon::prelude::*;
15use strum::EnumCount;
16
17use crate::{
18 batch_constraint::expr_eval::{
19 default_poseidon2_sub_chip, generate_dag_commit_info,
20 symbolic_expression::air::{
21 CachedSymbolicExpressionColumns, NodeKind, SingleMainSymbolicExpressionColumns,
22 ENCODER_MAX_DEGREE, NUM_FLAGS,
23 },
24 DagCommitCols, DagCommitInfo,
25 },
26 system::Preflight,
27 tracegen::RowMajorChip,
28 utils::{interaction_length, MultiVecWithBounds},
29};
30
31pub struct SymbolicExpressionTraceGenerator {
32 pub max_num_proofs: usize,
33 pub has_cached: bool,
34}
35
36pub(crate) struct SymbolicExpressionCtx<'a> {
37 pub vk: &'a MultiStarkVerifyingKey<BabyBearPoseidon2Config>,
38 pub preflights: &'a [&'a Preflight],
39 pub expr_evals: &'a MultiVecWithBounds<EF, 2>,
40 pub cached_trace_record: &'a Option<&'a CachedTraceRecord>,
41}
42
43impl RowMajorChip<F> for SymbolicExpressionTraceGenerator {
44 type Ctx<'a> = SymbolicExpressionCtx<'a>;
45
46 #[tracing::instrument(level = "trace", skip_all)]
47 fn generate_trace(
48 &self,
49 ctx: &Self::Ctx<'_>,
50 required_height: Option<usize>,
51 ) -> Option<RowMajorMatrix<F>> {
52 let child_vk = ctx.vk;
53 let preflights = ctx.preflights;
54 let max_num_proofs = self.max_num_proofs;
55 let has_cached = self.has_cached;
56 let expr_evals = ctx.expr_evals;
57 let trace_height = required_height;
58 let l_skip = child_vk.inner.params.l_skip;
59
60 let single_main_width = SingleMainSymbolicExpressionColumns::<F>::width();
61 let dag_commit_width = DagCommitCols::<F>::width();
62 let main_width =
63 single_main_width * max_num_proofs + if has_cached { 0 } else { dag_commit_width };
64
65 struct Record {
66 args: [F; 2 * D_EF],
67 sort_idx: usize,
68 n_abs: usize,
69 is_n_neg: usize,
70 }
71 let mut records = vec![];
72
73 for (proof_idx, preflight) in preflights.iter().enumerate() {
74 let rs = &preflight.batch_constraint.sumcheck_rnd;
75 let (&rs_0, rs_rest) = rs.split_first().unwrap();
76 let mut is_first_uni_by_log_height = vec![];
77 let mut is_last_uni_by_log_height = vec![];
78
79 for (log_height, &r_pow) in rs_0
80 .exp_powers_of_2()
81 .take(l_skip + 1)
82 .collect::<Vec<_>>()
83 .iter()
84 .rev()
85 .enumerate()
86 {
87 is_first_uni_by_log_height.push(eval_eq_uni_at_one(log_height, r_pow));
88 is_last_uni_by_log_height.push(eval_eq_uni_at_one(
89 log_height,
90 r_pow * F::two_adic_generator(log_height),
91 ));
92 }
93 let mut is_first_mle_by_n = vec![EF::ONE];
94 let mut is_last_mle_by_n = vec![EF::ONE];
95 for (i, &r) in rs_rest.iter().enumerate() {
96 is_first_mle_by_n.push(is_first_mle_by_n[i] * (EF::ONE - r));
97 is_last_mle_by_n.push(is_last_mle_by_n[i] * r);
98 }
99
100 for (air_idx, vk) in child_vk.inner.per_air.iter().enumerate() {
101 let constraints = &vk.symbolic_constraints.constraints;
102 let expr_evals = &expr_evals[[proof_idx, air_idx]];
103
104 if expr_evals.is_empty() {
106 let n = constraints.nodes.len()
107 + vk.symbolic_constraints
108 .interactions
109 .iter()
110 .map(interaction_length)
111 .sum::<usize>()
112 + vk.unused_variables.len();
113
114 records.resize_with(records.len() + n, || None);
115 continue;
116 }
117
118 let (sort_idx, trace_vdata) = preflight
119 .proof_shape
120 .sorted_trace_vdata
121 .iter()
122 .enumerate()
123 .find_map(|(sort_idx, (idx, vdata))| {
124 (*idx == air_idx).then_some((sort_idx, vdata))
125 })
126 .unwrap();
127
128 let log_height = trace_vdata.log_height;
129 let (n_abs, is_n_neg) = if log_height < l_skip {
130 (l_skip - log_height, 1)
131 } else {
132 (log_height - l_skip, 0)
133 };
134
135 for (node_idx, node) in constraints.nodes.iter().enumerate() {
136 let mut record = Record {
137 args: [F::ZERO; 2 * D_EF],
138 sort_idx,
139 n_abs,
140 is_n_neg,
141 };
142 match node {
143 SymbolicExpressionNode::Variable(var) => match var.entry {
144 Entry::Preprocessed { .. } => {
145 record.args[..D_EF].copy_from_slice(
146 expr_evals[node_idx].as_basis_coefficients_slice(),
147 );
148 }
149 Entry::Main { .. } => {
150 record.args[..D_EF].copy_from_slice(
151 expr_evals[node_idx].as_basis_coefficients_slice(),
152 );
153 }
154 Entry::Public => record.args[..D_EF].copy_from_slice(
155 expr_evals[node_idx].as_basis_coefficients_slice(),
156 ),
157 Entry::Challenge => unreachable!(),
158 },
159 SymbolicExpressionNode::IsFirstRow => {
160 record.args[..D_EF].copy_from_slice(
161 is_first_uni_by_log_height[min(log_height, l_skip)]
162 .as_basis_coefficients_slice(),
163 );
164 record.args[D_EF..2 * D_EF].copy_from_slice(
165 is_first_mle_by_n[log_height.saturating_sub(l_skip)]
166 .as_basis_coefficients_slice(),
167 );
168 }
169 SymbolicExpressionNode::IsLastRow
170 | SymbolicExpressionNode::IsTransition => {
171 record.args[..D_EF].copy_from_slice(
172 is_last_uni_by_log_height[min(log_height, l_skip)]
173 .as_basis_coefficients_slice(),
174 );
175 record.args[D_EF..2 * D_EF].copy_from_slice(
176 is_last_mle_by_n[log_height.saturating_sub(l_skip)]
177 .as_basis_coefficients_slice(),
178 );
179 }
180 SymbolicExpressionNode::Constant(_) => {}
181 SymbolicExpressionNode::Add {
182 left_idx,
183 right_idx,
184 degree_multiple: _,
185 } => {
186 record.args[..D_EF].copy_from_slice(
187 expr_evals[*left_idx].as_basis_coefficients_slice(),
188 );
189 record.args[D_EF..2 * D_EF].copy_from_slice(
190 expr_evals[*right_idx].as_basis_coefficients_slice(),
191 );
192 }
193 SymbolicExpressionNode::Sub {
194 left_idx,
195 right_idx,
196 degree_multiple: _,
197 } => {
198 record.args[..D_EF].copy_from_slice(
199 expr_evals[*left_idx].as_basis_coefficients_slice(),
200 );
201 record.args[D_EF..2 * D_EF].copy_from_slice(
202 expr_evals[*right_idx].as_basis_coefficients_slice(),
203 );
204 }
205 SymbolicExpressionNode::Neg {
206 idx,
207 degree_multiple: _,
208 } => {
209 record.args[..D_EF]
210 .copy_from_slice(expr_evals[*idx].as_basis_coefficients_slice());
211 }
212 SymbolicExpressionNode::Mul {
213 left_idx,
214 right_idx,
215 degree_multiple: _,
216 } => {
217 record.args[..D_EF].copy_from_slice(
218 expr_evals[*left_idx].as_basis_coefficients_slice(),
219 );
220 record.args[D_EF..2 * D_EF].copy_from_slice(
221 expr_evals[*right_idx].as_basis_coefficients_slice(),
222 );
223 }
224 };
225 records.push(Some(record));
226 }
227 for interaction in &vk.symbolic_constraints.interactions {
228 let mut args = [F::ZERO; 2 * D_EF];
229 args[..D_EF].copy_from_slice(
230 expr_evals[interaction.count].as_basis_coefficients_slice(),
231 );
232 records.push(Some(Record {
233 args,
234 sort_idx,
235 n_abs,
236 is_n_neg,
237 }));
238
239 for &node_idx in &interaction.message {
240 let mut args = [F::ZERO; 2 * D_EF];
241 args[..D_EF]
242 .copy_from_slice(expr_evals[node_idx].as_basis_coefficients_slice());
243 records.push(Some(Record {
244 args,
245 sort_idx,
246 n_abs,
247 is_n_neg,
248 }));
249 }
250
251 args.fill(F::ZERO);
252 args[0] = F::from_u16(interaction.bus_index + 1);
253 records.push(Some(Record {
254 args,
255 sort_idx,
256 n_abs,
257 is_n_neg,
258 }));
259 }
260 let mut node_idx = constraints.nodes.len();
261 for unused_var in &vk.unused_variables {
262 match unused_var.entry {
263 Entry::Preprocessed { .. } => {
264 let mut args = [F::ZERO; 2 * D_EF];
265 args[..D_EF].copy_from_slice(
266 expr_evals[node_idx].as_basis_coefficients_slice(),
267 );
268 records.push(Some(Record {
269 args,
270 sort_idx,
271 n_abs,
272 is_n_neg,
273 }));
274 }
275 Entry::Main { .. } => {
276 let mut args = [F::ZERO; 2 * D_EF];
277 args[..D_EF].copy_from_slice(
278 expr_evals[node_idx].as_basis_coefficients_slice(),
279 );
280 records.push(Some(Record {
281 args,
282 sort_idx,
283 n_abs,
284 is_n_neg,
285 }));
286 }
287 Entry::Public | Entry::Challenge => {
288 unreachable!()
289 }
290 }
291 node_idx += 1;
292 }
293 }
294 }
295
296 let num_valid_rows = records.len() / preflights.len();
299 let height = if let Some(height) = trace_height {
300 if height < num_valid_rows {
301 return None;
302 }
303 height
304 } else {
305 num_valid_rows.next_power_of_two()
306 };
307 let mut main_trace = F::zero_vec(main_width * height);
308
309 let (encoder, cached_records, poseidon2_rows) = if has_cached {
310 (None, None, None)
311 } else {
312 let encoder = Encoder::new(NodeKind::COUNT, ENCODER_MAX_DEGREE, true);
313 assert_eq!(encoder.width(), NUM_FLAGS);
314
315 let cached_trace_record = ctx.cached_trace_record.unwrap();
316 debug_assert_eq!(cached_trace_record.records.len(), num_valid_rows);
317
318 let poseidon2_subchip = default_poseidon2_sub_chip();
319 let poseidon2_trace = poseidon2_subchip.generate_trace(
320 cached_trace_record
321 .dag_commit_info
322 .as_ref()
323 .unwrap()
324 .poseidon2_inputs
325 .clone(),
326 );
327 let poseidon2_rows = poseidon2_trace
328 .rows()
329 .map(|row| row.collect_vec())
330 .collect_vec();
331
332 (
333 Some(encoder),
334 Some(&cached_trace_record.records),
335 Some(poseidon2_rows),
336 )
337 };
338
339 main_trace
340 .par_chunks_exact_mut(main_width)
341 .enumerate()
342 .for_each(|(row_idx, row)| {
343 let main_offset = if has_cached {
344 0
345 } else {
346 let poseidon2_row = &poseidon2_rows.as_ref().unwrap()[row_idx];
349 row[..poseidon2_row.len()].copy_from_slice(poseidon2_row);
350
351 if row_idx < num_valid_rows {
352 let record = &cached_records.as_ref().unwrap()[row_idx];
353 let encoder = encoder.as_ref().unwrap();
354 let cols: &mut DagCommitCols<_> = row[..dag_commit_width].borrow_mut();
355 for (i, x) in encoder
356 .get_flag_pt(record.kind as usize)
357 .into_iter()
358 .enumerate()
359 {
360 cols.flags[i] = F::from_u32(x);
361 }
362 cols.is_constraint = F::from_bool(record.is_constraint);
363 }
364
365 dag_commit_width
366 };
367
368 for proof_idx in 0..max_num_proofs {
369 if proof_idx >= preflights.len() {
370 continue;
371 }
372
373 let start = main_offset + proof_idx * single_main_width;
374 let end = start + single_main_width;
375 let cols: &mut SingleMainSymbolicExpressionColumns<_> =
376 row[start..end].borrow_mut();
377
378 if row_idx >= num_valid_rows {
379 cols.slot_state = F::ONE;
382 continue;
383 }
384
385 let record_idx = proof_idx * num_valid_rows + row_idx;
386 let Some(record) = records[record_idx].as_ref() else {
387 cols.slot_state = F::ONE;
389 continue;
390 };
391
392 cols.slot_state = F::TWO;
393 cols.args = record.args;
394 cols.sort_idx = F::from_usize(record.sort_idx);
395 cols.n_abs = F::from_usize(record.n_abs);
396 cols.is_n_neg = F::from_usize(record.is_n_neg);
397 }
398 });
399
400 Some(RowMajorMatrix::new(main_trace, main_width))
401 }
402}
403
404#[derive(Debug, Clone, Copy)]
405pub(crate) struct CachedRecord {
406 pub(crate) kind: NodeKind,
407 pub(crate) air_idx: usize,
408 pub(crate) node_idx: usize,
409 pub(crate) attrs: [usize; 3],
410 pub(crate) is_constraint: bool,
411 pub(crate) constraint_idx: usize,
412 pub(crate) fanout: usize,
413}
414
415#[derive(Debug, Clone)]
416pub struct CachedTraceRecord {
417 pub(crate) records: Vec<CachedRecord>,
418 pub dag_commit_info: Option<DagCommitInfo<F>>,
419}
420
421pub(crate) fn build_cached_trace_record(
422 child_vk: &MultiStarkVerifyingKey<BabyBearPoseidon2Config>,
423 has_cached: bool,
424) -> CachedTraceRecord {
425 let mut fanout_per_air = Vec::with_capacity(child_vk.inner.per_air.len());
426 for vk in &child_vk.inner.per_air {
427 let nodes = &vk.symbolic_constraints.constraints.nodes;
428 let mut fanout = vec![0usize; nodes.len()];
429
430 for node in nodes.iter() {
431 match node {
432 SymbolicExpressionNode::Add {
433 left_idx,
434 right_idx,
435 ..
436 }
437 | SymbolicExpressionNode::Sub {
438 left_idx,
439 right_idx,
440 ..
441 }
442 | SymbolicExpressionNode::Mul {
443 left_idx,
444 right_idx,
445 ..
446 } => {
447 fanout[*left_idx] += 1;
448 fanout[*right_idx] += 1;
449 }
450 SymbolicExpressionNode::Neg { idx, .. } => {
451 fanout[*idx] += 1;
452 }
453 _ => {}
454 }
455 }
456 for interaction in vk.symbolic_constraints.interactions.iter() {
457 fanout[interaction.count] += 1;
458 for &node_idx in &interaction.message {
459 fanout[node_idx] += 1;
460 }
461 }
462 fanout_per_air.push(fanout);
463 }
464
465 let mut records = vec![];
466 for (air_idx, (vk, fanout_per_node)) in
467 zip(child_vk.inner.per_air.iter(), fanout_per_air.into_iter()).enumerate()
468 {
469 let constraints = &vk.symbolic_constraints.constraints;
470 let constraint_idxs = &constraints.constraint_idx;
471
472 #[cfg(debug_assertions)]
473 {
474 for i in 1..constraint_idxs.len() {
475 debug_assert!(constraint_idxs[i - 1] < constraint_idxs[i]);
476 }
477 }
478
479 let mut j = 0;
480
481 for (node_idx, (node, &fanout)) in
482 zip(constraints.nodes.iter(), fanout_per_node.iter()).enumerate()
483 {
484 if j < constraint_idxs.len() && constraint_idxs[j] < node_idx {
485 j += 1;
486 }
487 let is_constraint = j < constraint_idxs.len() && constraint_idxs[j] == node_idx;
488
489 let mut record = CachedRecord {
490 kind: NodeKind::Constant,
491 air_idx,
492 node_idx,
493 attrs: [0; 3],
494 is_constraint,
495 constraint_idx: if !is_constraint { 0 } else { j },
496 fanout,
497 };
498
499 match node {
500 SymbolicExpressionNode::Variable(var) => {
501 record.attrs[0] = var.index;
502 match var.entry {
503 Entry::Preprocessed { offset } => {
504 record.kind = NodeKind::VarPreprocessed;
505 record.attrs[1] = 1;
506 record.attrs[2] = offset;
507 }
508 Entry::Main { part_index, offset } => {
509 record.kind = NodeKind::VarMain;
510 record.attrs[1] = vk.dag_main_part_index_to_commit_index(part_index);
511 record.attrs[2] = offset;
512 }
513 Entry::Public => {
514 record.kind = NodeKind::VarPublicValue;
515 }
516 Entry::Challenge => unreachable!(),
517 }
518 }
519 SymbolicExpressionNode::IsFirstRow => {
520 record.kind = NodeKind::SelIsFirst;
521 }
522 SymbolicExpressionNode::IsLastRow => {
523 record.kind = NodeKind::SelIsLast;
524 }
525 SymbolicExpressionNode::IsTransition => {
526 record.kind = NodeKind::SelIsTransition;
527 }
528 SymbolicExpressionNode::Constant(val) => {
529 record.kind = NodeKind::Constant;
530 record.attrs[0] = val.as_canonical_u32() as usize;
531 }
532 SymbolicExpressionNode::Add {
533 left_idx,
534 right_idx,
535 degree_multiple: _,
536 } => {
537 record.kind = NodeKind::Add;
538 record.attrs[0] = *left_idx;
539 record.attrs[1] = *right_idx;
540 }
541 SymbolicExpressionNode::Sub {
542 left_idx,
543 right_idx,
544 degree_multiple: _,
545 } => {
546 record.kind = NodeKind::Sub;
547 record.attrs[0] = *left_idx;
548 record.attrs[1] = *right_idx;
549 }
550 SymbolicExpressionNode::Neg {
551 idx,
552 degree_multiple: _,
553 } => {
554 record.kind = NodeKind::Neg;
555 record.attrs[0] = *idx;
556 }
557 SymbolicExpressionNode::Mul {
558 left_idx,
559 right_idx,
560 degree_multiple: _,
561 } => {
562 record.kind = NodeKind::Mul;
563 record.attrs[0] = *left_idx;
564 record.attrs[1] = *right_idx;
565 }
566 };
567 records.push(record);
568 }
569 for (interaction_idx, interaction) in
570 vk.symbolic_constraints.interactions.iter().enumerate()
571 {
572 records.push(CachedRecord {
573 kind: NodeKind::InteractionMult,
574 air_idx,
575 node_idx: interaction_idx,
576 attrs: [interaction.count, 0, 0],
577 is_constraint: false,
578 constraint_idx: 0,
579 fanout: 0,
580 });
581 for (idx_in_message, &node_idx) in interaction.message.iter().enumerate() {
582 records.push(CachedRecord {
583 kind: NodeKind::InteractionMsgComp,
584 air_idx,
585 node_idx: interaction_idx,
586 attrs: [node_idx, idx_in_message, 0],
587 is_constraint: false,
588 constraint_idx: 0,
589 fanout: 0,
590 });
591 }
592 records.push(CachedRecord {
593 kind: NodeKind::InteractionBusIndex,
594 air_idx,
595 node_idx: interaction_idx,
596 attrs: [interaction.bus_index as usize, 0, 0],
597 is_constraint: false,
598 constraint_idx: 0,
599 fanout: 0,
600 });
601 }
602 let mut node_idx = constraints.nodes.len();
603 for unused_var in &vk.unused_variables {
604 let record = match unused_var.entry {
605 Entry::Preprocessed { offset } => CachedRecord {
606 kind: NodeKind::VarPreprocessed,
607 air_idx,
608 node_idx,
609 attrs: [unused_var.index, 1, offset],
610 is_constraint: false,
611 constraint_idx: 0,
612 fanout: 0,
613 },
614 Entry::Main { part_index, offset } => {
615 let part = vk.dag_main_part_index_to_commit_index(part_index);
616 CachedRecord {
617 kind: NodeKind::VarMain,
618 air_idx,
619 node_idx,
620 attrs: [unused_var.index, part, offset],
621 is_constraint: false,
622 constraint_idx: 0,
623 fanout: 0,
624 }
625 }
626 Entry::Public | Entry::Challenge => {
627 unreachable!()
628 }
629 };
630 node_idx += 1;
631 records.push(record);
632 }
633 }
634
635 let dag_commit_info = (!has_cached).then(|| {
636 let encoder = Encoder::new(NodeKind::COUNT, ENCODER_MAX_DEGREE, true);
637 generate_dag_commit_info(&records, encoder)
638 });
639
640 CachedTraceRecord {
641 records,
642 dag_commit_info,
643 }
644}
645
646#[tracing::instrument(
648 name = "generate_cached_trace",
649 skip_all,
650 fields(air = "SymbolicExpressionAir")
651)]
652pub(crate) fn generate_symbolic_expr_cached_trace(
653 cached_trace_record: &CachedTraceRecord,
654) -> RowMajorMatrix<F> {
655 let encoder = Encoder::new(NodeKind::COUNT, ENCODER_MAX_DEGREE, true);
660 assert_eq!(encoder.width(), NUM_FLAGS);
661
662 let cached_width = CachedSymbolicExpressionColumns::<F>::width();
663 let records = &cached_trace_record.records;
664
665 let height = records.len().next_power_of_two();
666 let mut cached_trace = F::zero_vec(cached_width * height);
667 cached_trace
668 .par_chunks_exact_mut(cached_width)
669 .zip(records)
670 .for_each(|(row, record)| {
671 let cols: &mut CachedSymbolicExpressionColumns<_> = row.borrow_mut();
672
673 for (i, x) in encoder
674 .get_flag_pt(record.kind as usize)
675 .into_iter()
676 .enumerate()
677 {
678 cols.flags[i] = F::from_u32(x);
679 }
680 cols.air_idx = F::from_usize(record.air_idx);
681 cols.node_or_interaction_idx = F::from_usize(record.node_idx);
682 cols.attrs = record.attrs.map(F::from_usize);
683 cols.is_constraint = F::from_bool(record.is_constraint);
684 cols.constraint_idx = F::from_usize(record.constraint_idx);
685 cols.fanout = F::from_usize(record.fanout);
686 });
687
688 RowMajorMatrix::new(cached_trace, cached_width)
689}
690
691#[cfg(feature = "cuda")]
692pub(in crate::batch_constraint) mod cuda {
693
694 use openvm_cuda_backend::{base::DeviceMatrix, prelude::F, GpuBackend};
695 use openvm_cuda_common::{copy::MemCopyH2D, d_buffer::DeviceBuffer, stream::GpuDeviceCtx};
696 use openvm_stark_backend::prover::AirProvingContext;
697
698 use super::*;
699 use crate::{
700 batch_constraint::{cuda_abi::sym_expr_common_tracegen, cuda_utils::*},
701 cuda::{preflight::PreflightGpu, proof::ProofGpu, to_device_or_nullptr_on},
702 tracegen::ModuleChip,
703 };
704
705 pub struct SymbolicExpressionGpuCtx<'a> {
706 pub vk: &'a MultiStarkVerifyingKey<BabyBearPoseidon2Config>,
707 pub proofs: &'a [ProofGpu],
708 pub preflights: &'a [PreflightGpu],
709 pub expr_evals: &'a MultiVecWithBounds<openvm_cuda_backend::prelude::EF, 2>,
710 pub cached_trace_record: &'a Option<&'a CachedTraceRecord>,
711 pub device_ctx: &'a GpuDeviceCtx,
712 }
713
714 impl ModuleChip<GpuBackend> for SymbolicExpressionTraceGenerator {
715 type Ctx<'a> = SymbolicExpressionGpuCtx<'a>;
716
717 #[tracing::instrument(name = "generate_trace", level = "trace", skip_all)]
718 fn generate_proving_ctx(
719 &self,
720 ctx: &Self::Ctx<'_>,
721 required_height: Option<usize>,
722 ) -> Option<AirProvingContext<GpuBackend>> {
723 let child_vk = ctx.vk;
724 let proofs = ctx.proofs;
725 let preflights = ctx.preflights;
726 let max_num_proofs = self.max_num_proofs;
727 let has_cached = self.has_cached;
728 let expr_evals = ctx.expr_evals;
729 let device_ctx = ctx.device_ctx;
730
731 debug_assert_eq!(proofs.len(), preflights.len());
732
733 let num_airs = child_vk.inner.per_air.len();
734
735 let mut constraint_nodes = MultiVecWithBounds::<_, 1>::new();
736
737 let mut interactions = MultiVecWithBounds::<_, 1>::new();
738
739 let mut interaction_messages = Vec::new();
740
741 let mut unused_variables = MultiVecWithBounds::<_, 1>::new();
742
743 let mut record_bounds = Vec::with_capacity(num_airs + 1);
744 record_bounds.push(0);
745
746 let mut total_rows = 0;
747
748 for vk in &child_vk.inner.per_air {
749 let constraints = &vk.symbolic_constraints.constraints;
750 for node in &constraints.nodes {
751 constraint_nodes.push(flatten_constraint_node(vk, node));
752 }
753 constraint_nodes.close_level(0);
754
755 for interaction in &vk.symbolic_constraints.interactions {
756 let message_start = interaction_messages.len();
757 interaction_messages.extend(&interaction.message);
758 interactions.push(FlatInteraction {
759 count: interaction.count as u32,
760 message_start: message_start as u32,
761 message_len: interaction.message.len() as u32,
762 bus_index: u32::from(interaction.bus_index),
763 count_weight: interaction.count_weight,
764 });
765 }
766 interactions.close_level(0);
767
768 for unused in &vk.unused_variables {
769 unused_variables.push(flatten_unused_symbolic_variable(unused));
770 }
771 unused_variables.close_level(0);
772
773 let interaction_message_rows: usize = vk
774 .symbolic_constraints
775 .interactions
776 .iter()
777 .map(|interaction| interaction.message.len())
778 .sum();
779 let rows_for_air = constraints.nodes.len() + 2 * vk.symbolic_constraints.interactions.len() + interaction_message_rows + vk.unused_variables.len(); total_rows += rows_for_air;
784 record_bounds.push(total_rows as u32);
785 }
786
787 let mut air_ids_per_record = vec![0; total_rows];
788 for i in 0..(record_bounds.len() - 1) {
789 air_ids_per_record[(record_bounds[i] as usize)..(record_bounds[i + 1] as usize)]
790 .fill(i as u32);
791 }
792
793 let height = if let Some(height) = required_height {
794 if height < total_rows {
795 return None;
796 }
797 height
798 } else {
799 total_rows.max(1).next_power_of_two()
800 };
801 let commit_width = DagCommitCols::<F>::width();
802 let width = SingleMainSymbolicExpressionColumns::<F>::width() * max_num_proofs
803 + if has_cached { 0 } else { commit_width };
804 let trace = DeviceMatrix::with_capacity_on(height, width, device_ctx);
805
806 let d_log_heights = proofs
807 .iter()
808 .flat_map(|proof| {
809 proof
810 .cpu
811 .trace_vdata
812 .iter()
813 .map(|v| v.as_ref().map_or(0, |td| td.log_height))
814 })
815 .collect::<Vec<_>>()
816 .to_device_on(device_ctx)
817 .unwrap();
818
819 let mut sort_idx_by_air_idx = vec![0usize; num_airs * proofs.len()];
820 for (chunk, preflight) in sort_idx_by_air_idx
821 .chunks_exact_mut(num_airs)
822 .zip(preflights.iter())
823 {
824 for (sort_idx, (air_idx, _)) in preflight
825 .cpu
826 .proof_shape
827 .sorted_trace_vdata
828 .iter()
829 .enumerate()
830 {
831 chunk[*air_idx] = sort_idx;
832 }
833 }
834 let d_sort_idx_by_air_idx = sort_idx_by_air_idx.to_device_on(device_ctx).unwrap();
835
836 let d_expr_evals = expr_evals.data.to_device_on(device_ctx).unwrap();
837 let d_ee_bounds_0 = expr_evals.bounds[0].to_device_on(device_ctx).unwrap();
838 let d_ee_bounds_1 = expr_evals.bounds[1].to_device_on(device_ctx).unwrap();
839
840 let d_constraint_nodes = constraint_nodes.data.to_device_on(device_ctx).unwrap();
841 let d_constraint_nodes_bounds =
842 constraint_nodes.bounds[0].to_device_on(device_ctx).unwrap();
843 let d_interactions = to_device_or_nullptr_on(&interactions.data, device_ctx).unwrap();
844 let d_interactions_bounds = interactions.bounds[0].to_device_on(device_ctx).unwrap();
845 let d_interaction_messages =
846 to_device_or_nullptr_on(&interaction_messages, device_ctx).unwrap();
847 let d_unused_variables =
848 to_device_or_nullptr_on(&unused_variables.data, device_ctx).unwrap();
849 let d_unused_variables_bounds =
850 unused_variables.bounds[0].to_device_on(device_ctx).unwrap();
851 let d_record_bounds = record_bounds.to_device_on(device_ctx).unwrap();
852 let d_air_ids_per_record = air_ids_per_record.to_device_on(device_ctx).unwrap();
853
854 let mut sumcheck_data = Vec::new();
855 let mut sumcheck_bounds = Vec::with_capacity(preflights.len() + 1);
856 sumcheck_bounds.push(0);
857 for preflight in preflights {
858 sumcheck_data.extend_from_slice(&preflight.cpu.batch_constraint.sumcheck_rnd);
859 sumcheck_bounds.push(sumcheck_data.len());
860 }
861 let d_sumcheck_rnds = if sumcheck_data.is_empty() {
862 DeviceBuffer::new()
863 } else {
864 sumcheck_data.to_device_on(device_ctx).unwrap()
865 };
866 let d_sumcheck_bounds = sumcheck_bounds.to_device_on(device_ctx).unwrap();
867 let d_cached_records = ctx.cached_trace_record.map(|data| {
868 build_cached_gpu_records(data)
869 .unwrap()
870 .to_device_on(device_ctx)
871 .unwrap()
872 });
873
874 unsafe {
875 sym_expr_common_tracegen(
876 trace.buffer(),
877 height,
878 child_vk.inner.params.l_skip,
879 &d_log_heights,
880 &d_sort_idx_by_air_idx,
881 num_airs,
882 proofs.len(),
883 max_num_proofs,
884 &d_expr_evals,
885 &d_ee_bounds_0,
886 &d_ee_bounds_1,
887 &d_constraint_nodes,
888 &d_constraint_nodes_bounds,
889 &d_interactions,
890 &d_interactions_bounds,
891 &d_interaction_messages,
892 &d_unused_variables,
893 &d_unused_variables_bounds,
894 &d_record_bounds,
895 &d_air_ids_per_record,
896 total_rows,
897 &d_sumcheck_rnds,
898 &d_sumcheck_bounds,
899 d_cached_records.as_ref(),
900 device_ctx.stream.as_raw(),
901 )
902 .unwrap();
903 }
904
905 Some(AirProvingContext::simple_no_pis(trace))
906 }
907 }
908}