openvm_recursion_circuit/batch_constraint/expr_eval/constraints_folding/
trace.rs1use std::borrow::BorrowMut;
2
3use itertools::Itertools;
4use openvm_stark_backend::keygen::types::MultiStarkVerifyingKey0;
5use openvm_stark_sdk::config::baby_bear_poseidon2::{BabyBearPoseidon2Config, D_EF, EF, F};
6use p3_field::{BasedVectorSpace, PrimeCharacteristicRing};
7use p3_matrix::dense::RowMajorMatrix;
8use p3_maybe_rayon::prelude::*;
9
10use crate::{
11 batch_constraint::expr_eval::constraints_folding::air::ConstraintsFoldingCols,
12 system::Preflight,
13 tracegen::RowMajorChip,
14 utils::{MultiProofVecVec, MultiVecWithBounds},
15};
16
17#[derive(Copy, Clone)]
18#[repr(C)]
19pub(crate) struct ConstraintsFoldingRecord {
20 sort_idx: usize,
21 air_idx: usize,
22 constraint_idx: usize,
23 node_idx: usize,
24 is_first_in_air: bool,
25 value: EF,
26}
27
28pub(crate) struct ConstraintsFoldingBlob {
29 pub(crate) records: MultiProofVecVec<ConstraintsFoldingRecord>,
30 pub(crate) folded_claims: MultiProofVecVec<(isize, EF)>,
32}
33
34impl ConstraintsFoldingBlob {
35 pub fn new(
36 vk: &MultiStarkVerifyingKey0<BabyBearPoseidon2Config>,
37 expr_evals: &MultiVecWithBounds<EF, 2>,
38 preflights: &[&Preflight],
39 ) -> Self {
40 let constraints = vk
41 .per_air
42 .iter()
43 .map(|vk| vk.symbolic_constraints.constraints.constraint_idx.clone())
44 .collect_vec();
45
46 let mut records = MultiProofVecVec::new();
47 let mut folded = MultiProofVecVec::new();
48 for (pidx, preflight) in preflights.iter().enumerate() {
49 let lambda_tidx = preflight.batch_constraint.lambda_tidx;
50 let lambda = EF::from_basis_coefficients_slice(
51 &preflight.transcript.values()[lambda_tidx..lambda_tidx + D_EF],
52 )
53 .unwrap();
54
55 let vdata = &preflight.proof_shape.sorted_trace_vdata;
56 for (sort_idx, (air_idx, v)) in vdata.iter().enumerate() {
57 let constrs = &constraints[*air_idx];
58 records.push(ConstraintsFoldingRecord {
59 sort_idx,
61 air_idx: *air_idx,
62 constraint_idx: 0,
63 node_idx: 0,
64 is_first_in_air: true,
65 value: EF::ZERO,
66 });
67 let mut folded_claim = EF::ZERO;
68 let mut lambda_pow = EF::ONE;
69 for (constraint_idx, &constr) in constrs.iter().enumerate() {
70 let value = expr_evals[[pidx, *air_idx]][constr];
71 folded_claim += lambda_pow * value;
72 lambda_pow *= lambda;
73 records.push(ConstraintsFoldingRecord {
74 sort_idx,
75 air_idx: *air_idx,
76 constraint_idx: constraint_idx + 1,
77 node_idx: constr,
78 is_first_in_air: false,
79 value,
80 });
81 }
82 let n_lift = v.log_height.saturating_sub(vk.params.l_skip);
83 let n = v.log_height as isize - vk.params.l_skip as isize;
84 folded.push((
85 n,
86 folded_claim * preflight.batch_constraint.eq_ns_frontloaded[n_lift],
87 ));
88 }
89 records.end_proof();
90 folded.end_proof();
91 }
92 Self {
93 records,
94 folded_claims: folded,
95 }
96 }
97}
98
99pub struct ConstraintsFoldingTraceGenerator;
100
101impl RowMajorChip<F> for ConstraintsFoldingTraceGenerator {
102 type Ctx<'a> = (&'a ConstraintsFoldingBlob, &'a [&'a Preflight]);
103
104 #[tracing::instrument(level = "trace", skip_all)]
105 fn generate_trace(
106 &self,
107 ctx: &Self::Ctx<'_>,
108 required_height: Option<usize>,
109 ) -> Option<RowMajorMatrix<F>> {
110 let (blob, preflights) = ctx;
111 let width = ConstraintsFoldingCols::<F>::width();
112
113 let total_height = blob.records.len();
114 debug_assert!(total_height > 0);
115 let padding_height = if let Some(height) = required_height {
116 if height < total_height {
117 return None;
118 }
119 height
120 } else {
121 total_height.next_power_of_two()
122 };
123 let mut trace = vec![F::ZERO; padding_height * width];
124
125 let mut cur_height = 0;
126 for (pidx, preflight) in preflights.iter().enumerate() {
127 let lambda_tidx = preflight.batch_constraint.lambda_tidx;
128 let lambda_slice = &preflight.transcript.values()[lambda_tidx..lambda_tidx + D_EF];
129 let records = &blob.records[pidx];
130
131 trace[cur_height * width..(cur_height + records.len()) * width]
132 .par_chunks_exact_mut(width)
133 .zip(records.par_iter())
134 .for_each(|(chunk, record)| {
135 let cols: &mut ConstraintsFoldingCols<_> = chunk.borrow_mut();
136 let n_lift = preflight.proof_shape.sorted_trace_vdata[record.sort_idx]
137 .1
138 .log_height
139 .saturating_sub(preflight.proof_shape.l_skip);
140
141 cols.is_valid = F::ONE;
142 cols.proof_idx = F::from_usize(pidx);
143 cols.air_idx = F::from_usize(record.air_idx);
144 cols.sort_idx = F::from_usize(record.sort_idx);
145 cols.constraint_idx = F::from_usize(record.constraint_idx);
146 cols.n_lift = F::from_usize(n_lift);
147 cols.lambda_tidx = F::from_usize(lambda_tidx);
148 cols.lambda.copy_from_slice(lambda_slice);
149 cols.value
150 .copy_from_slice(record.value.as_basis_coefficients_slice());
151 cols.eq_n.copy_from_slice(
152 preflight.batch_constraint.eq_ns_frontloaded[n_lift]
153 .as_basis_coefficients_slice(),
154 );
155 cols.is_first_in_air = F::from_bool(record.is_first_in_air);
156 });
157
158 let mut cur_sum = EF::ZERO;
160 let lambda = EF::from_basis_coefficients_slice(lambda_slice).unwrap();
161 trace[cur_height * width..(cur_height + records.len()) * width]
162 .chunks_exact_mut(width)
163 .rev()
164 .for_each(|chunk| {
165 let cols: &mut ConstraintsFoldingCols<_> = chunk.borrow_mut();
166 cur_sum =
167 cur_sum * lambda + EF::from_basis_coefficients_slice(&cols.value).unwrap();
168 cols.cur_sum
169 .copy_from_slice(cur_sum.as_basis_coefficients_slice());
170 if cols.is_first_in_air == F::ONE {
171 cur_sum = EF::ZERO;
172 }
173 });
174
175 {
176 let cols: &mut ConstraintsFoldingCols<_> =
177 trace[cur_height * width..(cur_height + 1) * width].borrow_mut();
178 cols.is_first = F::ONE;
179 }
180 cur_height += records.len();
181 }
182 Some(RowMajorMatrix::new(trace, width))
183 }
184}
185
186#[cfg(feature = "cuda")]
187pub(in crate::batch_constraint) mod cuda {
188 use openvm_circuit_primitives::cuda_abi::UInt2;
189 use openvm_cuda_backend::{base::DeviceMatrix, GpuBackend};
190 use openvm_cuda_common::{copy::MemCopyH2D, d_buffer::DeviceBuffer, stream::GpuDeviceCtx};
191 use openvm_stark_backend::prover::AirProvingContext;
192
193 use super::*;
194 use crate::{
195 batch_constraint::cuda_abi::{
196 constraints_folding_tracegen, constraints_folding_tracegen_temp_bytes, AffineFpExt,
197 FpExtWithTidx,
198 },
199 cuda::{preflight::PreflightGpu, vk::VerifyingKeyGpu},
200 tracegen::ModuleChip,
201 };
202
203 pub struct ConstraintsFoldingBlobGpu {
204 pub values: Vec<Vec<Vec<EF>>>,
206 pub constraints_folding_per_proof: Vec<FpExtWithTidx>,
208 pub folded_claims: MultiProofVecVec<(isize, EF)>,
210 }
211
212 impl ConstraintsFoldingBlobGpu {
213 pub fn new(
214 vk: &VerifyingKeyGpu,
215 expr_evals: &MultiVecWithBounds<EF, 2>,
216 preflights: &[PreflightGpu],
217 _device_ctx: &GpuDeviceCtx,
218 ) -> Self {
219 let constraints = vk
220 .cpu
221 .inner
222 .per_air
223 .iter()
224 .map(|vk| vk.symbolic_constraints.constraints.constraint_idx.clone())
225 .collect_vec();
226
227 let mut values = Vec::with_capacity(preflights.len());
228 let mut constraints_folding_per_proof = Vec::with_capacity(preflights.len());
229 let mut folded_claims = MultiProofVecVec::new();
230
231 for (pidx, preflight) in preflights.iter().enumerate() {
232 let lambda_tidx = preflight.cpu.batch_constraint.lambda_tidx;
233 let lambda = EF::from_basis_coefficients_slice(
234 &preflight.cpu.transcript.values()[lambda_tidx..lambda_tidx + D_EF],
235 )
236 .unwrap();
237
238 let vdata = &preflight.cpu.proof_shape.sorted_trace_vdata;
239 let mut proof_values = Vec::with_capacity(vdata.len());
240
241 for (air_idx, v) in vdata.iter() {
242 let mut folded_claim = EF::ZERO;
243 let mut lambda_pow = EF::ONE;
244
245 let air_values = std::iter::once(EF::ZERO)
246 .chain(constraints[*air_idx].iter().map(|&constr| {
247 let value = expr_evals[[pidx, *air_idx]][constr];
248 folded_claim += lambda_pow * value;
249 lambda_pow *= lambda;
250 value
251 }))
252 .collect_vec();
253 proof_values.push(air_values);
254
255 let n_lift = v.log_height.saturating_sub(vk.system_params.l_skip);
256 let n = v.log_height as isize - vk.system_params.l_skip as isize;
257 folded_claims.push((
258 n,
259 folded_claim * preflight.cpu.batch_constraint.eq_ns_frontloaded[n_lift],
260 ));
261 }
262
263 values.push(proof_values);
264 constraints_folding_per_proof.push(FpExtWithTidx {
265 value: lambda,
266 tidx: lambda_tidx as u32,
267 });
268 folded_claims.end_proof();
269 }
270
271 Self {
272 values,
273 constraints_folding_per_proof,
274 folded_claims,
275 }
276 }
277 }
278
279 impl ModuleChip<GpuBackend> for ConstraintsFoldingTraceGenerator {
280 type Ctx<'a> = (
281 &'a VerifyingKeyGpu,
282 &'a [PreflightGpu],
283 &'a ConstraintsFoldingBlobGpu,
284 &'a GpuDeviceCtx,
285 );
286
287 #[tracing::instrument(level = "trace", skip_all)]
288 fn generate_proving_ctx(
289 &self,
290 ctx: &Self::Ctx<'_>,
291 required_height: Option<usize>,
292 ) -> Option<AirProvingContext<GpuBackend>> {
293 let (child_vk, preflights_gpu, blob, device_ctx) = ctx;
294
295 let mut num_valid_rows = 0u32;
296 let mut row_bounds = Vec::with_capacity(preflights_gpu.len());
297 let mut constraint_bounds = Vec::with_capacity(preflights_gpu.len());
298 let mut proof_and_sort_idxs = vec![];
299
300 let flat_values = blob
301 .values
302 .iter()
303 .enumerate()
304 .flat_map(|(proof_idx, proof_values)| {
305 let mut num_constraints_in_proof = 0;
306 let mut proof_constraint_bounds = Vec::with_capacity(proof_values.len());
307 for (sort_idx, air_values) in proof_values.iter().enumerate() {
308 let num_constraints = air_values.len();
309 num_constraints_in_proof += num_constraints as u32;
310 proof_constraint_bounds.push(num_constraints_in_proof);
311 proof_and_sort_idxs.extend(std::iter::repeat_n(
312 UInt2 {
313 x: proof_idx as u32,
314 y: sort_idx as u32,
315 },
316 num_constraints,
317 ));
318 }
319 num_valid_rows += num_constraints_in_proof;
320 row_bounds.push(num_valid_rows);
321 constraint_bounds
322 .push(proof_constraint_bounds.to_device_on(device_ctx).unwrap());
323 proof_values.iter().flatten().copied()
324 })
325 .collect_vec();
326 let eq_ns = preflights_gpu
327 .iter()
328 .map(|preflight| {
329 preflight
330 .cpu
331 .batch_constraint
332 .eq_ns_frontloaded
333 .to_device_on(device_ctx)
334 .unwrap()
335 })
336 .collect_vec();
337
338 let height = if let Some(height) = required_height {
339 if height < num_valid_rows as usize {
340 return None;
341 }
342 height
343 } else {
344 (num_valid_rows as usize).next_power_of_two()
345 };
346 let width = ConstraintsFoldingCols::<F>::width();
347 let d_trace = DeviceMatrix::<F>::with_capacity_on(height, width, device_ctx);
348
349 let d_proof_and_sort_idxs = proof_and_sort_idxs.to_device_on(device_ctx).unwrap();
350 let d_values = flat_values.to_device_on(device_ctx).unwrap();
351 let d_cur_sum_evals =
352 DeviceBuffer::<AffineFpExt>::with_capacity_on(d_values.len(), device_ctx);
353
354 let d_constraint_bounds = constraint_bounds.iter().map(|b| b.as_ptr()).collect_vec();
355 let d_sorted_trace_heights = preflights_gpu
356 .iter()
357 .map(|preflight| preflight.proof_shape.sorted_trace_heights.as_ptr())
358 .collect_vec();
359 let d_eq_ns = eq_ns.iter().map(|b| b.as_ptr()).collect_vec();
360
361 let d_per_proof = blob
362 .constraints_folding_per_proof
363 .to_device_on(device_ctx)
364 .unwrap();
365
366 unsafe {
367 let temp_bytes = constraints_folding_tracegen_temp_bytes(
368 &d_proof_and_sort_idxs,
369 &d_cur_sum_evals,
370 num_valid_rows,
371 device_ctx.stream.as_raw(),
372 )
373 .unwrap();
374 let d_temp_buffer = DeviceBuffer::<u8>::with_capacity_on(temp_bytes, device_ctx);
375 constraints_folding_tracegen(
376 d_trace.buffer(),
377 height,
378 width,
379 &d_proof_and_sort_idxs,
380 &d_cur_sum_evals,
381 &d_values,
382 &row_bounds,
383 d_constraint_bounds,
384 d_sorted_trace_heights,
385 d_eq_ns,
386 &d_per_proof,
387 preflights_gpu.len() as u32,
388 child_vk.per_air.len() as u32,
389 num_valid_rows,
390 child_vk.system_params.l_skip as u32,
391 &d_temp_buffer,
392 temp_bytes,
393 device_ctx.stream.as_raw(),
394 )
395 .unwrap();
396 }
397
398 Some(AirProvingContext::simple_no_pis(d_trace))
399 }
400 }
401}