diff --git a/src/core/air/accumulation.rs b/src/core/air/accumulation.rs index a10398596..3e7ae9162 100644 --- a/src/core/air/accumulation.rs +++ b/src/core/air/accumulation.rs @@ -1,90 +1,69 @@ //! Accumulators for a random linear combination of circle polynomials. -//! Given N polynomials, sort them by size: u_0(P), ... u_{N-1}(P). -//! Given a random alpha, the combined polynomial is defined as -//! f(p) = sum_i alpha^{N-1-i} u_i (P). +//! Given N polynomials, u_0(P), ... u_{N-1}(P), and a random alpha, the combined polynomial is +//! defined as +//! f(p) = sum_i alpha^{N-1-i} u_i(P). + use itertools::Itertools; use crate::core::backend::cpu::CPUCircleEvaluation; use crate::core::backend::{Backend, CPUBackend}; use crate::core::fields::qm31::SecureField; use crate::core::fields::secure_column::SecureColumn; -use crate::core::fields::FieldExpOps; use crate::core::poly::circle::{CanonicCoset, CirclePoly, SecureCirclePoly}; use crate::core::poly::BitReversedOrder; +use crate::core::utils::generate_secure_powers; -/// Accumulates evaluations of u_i(P0) at a single point. +/// Accumulates N evaluations of u_i(P0) at a single point. /// Computes f(P0), the combined polynomial at that point. +/// For n accumulated evaluations, the i'th evaluation is multiplied by alpha^(N-1-i). pub struct PointEvaluationAccumulator { random_coeff: SecureField, - /// Accumulated evaluations for each log_size. - /// Each `sub_accumulation` holds `sum_{i=0}^{n-1} evaluation_i * alpha^(n-1-i)`, - /// where `n` is the number of accumulated evaluations for this log_size. - sub_accumulations: Vec, - /// Number of accumulated evaluations for each log_size. - n_accumulated: Vec, + accumulation: SecureField, } + impl PointEvaluationAccumulator { /// Creates a new accumulator. /// `random_coeff` should be a secure random field element, drawn from the channel. - /// `max_log_size` is the maximum log_size of the accumulated evaluations. - pub fn new(random_coeff: SecureField, max_log_size: u32) -> Self { - // TODO(spapini): Consider making all log_sizes usize. - let max_log_size = max_log_size as usize; + pub fn new(random_coeff: SecureField) -> Self { Self { random_coeff, - sub_accumulations: vec![SecureField::default(); max_log_size + 1], - n_accumulated: vec![0; max_log_size + 1], + accumulation: SecureField::default(), } } - /// Accumulates u_i(P0), a polynomial evaluation at a P0. - pub fn accumulate(&mut self, log_size: u32, evaluation: SecureField) { - assert!(log_size > 0 && log_size < self.sub_accumulations.len() as u32); - let sub_accumulation = &mut self.sub_accumulations[log_size as usize]; - *sub_accumulation = *sub_accumulation * self.random_coeff + evaluation; - - self.n_accumulated[log_size as usize] += 1; + /// Accumulates u_i(P0), a polynomial evaluation at a P0 in reverse order. + pub fn accumulate(&mut self, evaluation: SecureField) { + self.accumulation = self.accumulation * self.random_coeff + evaluation; } - /// Computes f(P0), the evaluation of the combined polynomial at P0. pub fn finalize(self) -> SecureField { - // Each `sub_accumulation` holds a linear combination of a consecutive slice of - // u_0(P0), ... u_{N-1}(P0): - // alpha^n_k u_i(P0) + alpha^{n_k-1} u_{i+1}(P0) + ... + alpha^0 u_{i+n_k-1}(P0). - // To combine all these slices, multiply an accumulator by alpha^k, and add the next slice. - self.sub_accumulations - .iter() - .zip(self.n_accumulated.iter()) - .fold(SecureField::default(), |total, (sub_accumulation, n_i)| { - total * self.random_coeff.pow(*n_i as u128) + *sub_accumulation - }) + self.accumulation } } +// TODO(ShaharS), rename terminology to constraints instead of columns. /// Accumulates evaluations of u_i(P), each at an evaluation domain of the size of that polynomial. /// Computes the coefficients of f(P). pub struct DomainEvaluationAccumulator { - pub random_coeff: SecureField, + random_coeff_powers: Vec, /// Accumulated evaluations for each log_size. - /// Each `sub_accumulation` holds `sum_{i=0}^{n-1} evaluation_i * alpha^(n-1-i)`, - /// where `n` is the number of accumulated evaluations for this log_size. + /// Each `sub_accumulation` holds the sum over all columns i of that log_size, of + /// `evaluation_i * alpha^(N - 1 - i)` + /// where `N` is the total number of evaluations. sub_accumulations: Vec>, - /// Number of accumulated evaluations for each log_size. - n_cols_per_size: Vec, } impl DomainEvaluationAccumulator { /// Creates a new accumulator. /// `random_coeff` should be a secure random field element, drawn from the channel. /// `max_log_size` is the maximum log_size of the accumulated evaluations. - pub fn new(random_coeff: SecureField, max_log_size: u32) -> Self { + pub fn new(random_coeff: SecureField, max_log_size: u32, total_columns: usize) -> Self { let max_log_size = max_log_size as usize; Self { - random_coeff, + random_coeff_powers: generate_secure_powers(random_coeff, total_columns), sub_accumulations: (0..(max_log_size + 1)) .map(|n| SecureColumn::zeros(1 << n)) .collect(), - n_cols_per_size: vec![0; max_log_size + 1], } } @@ -97,18 +76,19 @@ impl DomainEvaluationAccumulator { &mut self, n_cols_per_size: [(u32, usize); N], ) -> [ColumnAccumulator<'_, B>; N] { - n_cols_per_size.iter().for_each(|(log_size, n_col)| { - assert!(*log_size > 0 && *log_size < self.sub_accumulations.len() as u32); - self.n_cols_per_size[*log_size as usize] += n_col; - }); self.sub_accumulations .get_many_mut(n_cols_per_size.map(|(log_size, _)| log_size as usize)) .unwrap_or_else(|e| panic!("invalid log_sizes: {}", e)) .into_iter() .zip(n_cols_per_size) - .map(|(col, (_, count))| ColumnAccumulator { - random_coeff_pow: self.random_coeff.pow(count as u128), - col, + .map(|(col, (_, n_cols))| { + let random_coeffs = self + .random_coeff_powers + .split_off(self.random_coeff_powers.len() - n_cols); + ColumnAccumulator { + random_coeff_powers: random_coeffs, + col, + } }) .collect_vec() .try_into() @@ -124,17 +104,16 @@ impl DomainEvaluationAccumulator { impl DomainEvaluationAccumulator { /// Computes f(P) as coefficients. pub fn finalize(self) -> SecureCirclePoly { + assert_eq!( + self.random_coeff_powers.len(), + 0, + "not all random coefficients were used" + ); let mut res_coeffs = SecureColumn::::zeros(1 << self.log_size()); let res_log_size = self.log_size(); let res_size = 1 << res_log_size; - for ((log_size, values), n_cols) in self - .sub_accumulations - .into_iter() - .enumerate() - .zip(self.n_cols_per_size.iter()) - .skip(1) - { + for (log_size, values) in self.sub_accumulations.into_iter().enumerate().skip(1) { let coeffs = SecureColumn:: { columns: values.columns.map(|c| { CPUCircleEvaluation::<_, BitReversedOrder>::new( @@ -147,9 +126,8 @@ impl DomainEvaluationAccumulator { }), }; // Add column coefficients into result coefficients, element-wise, in-place. - let multiplier = self.random_coeff.pow(*n_cols as u128); for i in 0..res_size { - let res_coeff = res_coeffs.at(i) * multiplier + coeffs.at(i); + let res_coeff = res_coeffs.at(i) + coeffs.at(i); res_coeffs.set(i, res_coeff); } } @@ -158,14 +136,14 @@ impl DomainEvaluationAccumulator { } } -/// An domain accumulator for polynomials of a single size. +/// A domain accumulator for polynomials of a single size. pub struct ColumnAccumulator<'a, B: Backend> { - random_coeff_pow: SecureField, + pub random_coeff_powers: Vec, col: &'a mut SecureColumn, } impl<'a> ColumnAccumulator<'a, CPUBackend> { pub fn accumulate(&mut self, index: usize, evaluation: SecureField) { - let val = self.col.at(index) * self.random_coeff_pow + evaluation; + let val = self.col.at(index) + evaluation; self.col.set(index, val); } } @@ -202,18 +180,15 @@ mod tests { let alpha = qm31!(2, 3, 4, 5); // Use accumulator. - let mut accumulator = PointEvaluationAccumulator::new(alpha, MAX_LOG_SIZE); - for (log_size, evaluation) in log_sizes.iter().zip(evaluations.iter()) { - accumulator.accumulate(*log_size, (*evaluation).into()); + let mut accumulator = PointEvaluationAccumulator::new(alpha); + for (_, evaluation) in log_sizes.iter().zip(evaluations.iter()) { + accumulator.accumulate((*evaluation).into()); } let accumulator_res = accumulator.finalize(); // Use direct computation. let mut res = SecureField::default(); - // Sort evaluations by log_size. - let mut pairs = log_sizes.into_iter().zip(evaluations).collect::>(); - pairs.sort_by_key(|(log_size, _)| *log_size); - for (_, evaluation) in pairs.iter() { + for evaluation in evaluations.iter() { res = res * alpha + *evaluation; } @@ -227,9 +202,10 @@ mod tests { const LOG_SIZE_MIN: u32 = 4; const LOG_SIZE_BOUND: u32 = 10; const MASK: u32 = P; - let log_sizes = (0..100) + let mut log_sizes = (0..100) .map(|_| rng.gen_range(LOG_SIZE_MIN..LOG_SIZE_BOUND)) .collect::>(); + log_sizes.sort(); // Generate random evaluations. let evaluations = log_sizes @@ -243,7 +219,11 @@ mod tests { let alpha = qm31!(2, 3, 4, 5); // Use accumulator. - let mut accumulator = DomainEvaluationAccumulator::::new(alpha, LOG_SIZE_BOUND); + let mut accumulator = DomainEvaluationAccumulator::::new( + alpha, + LOG_SIZE_BOUND, + evaluations.len(), + ); let n_cols_per_size: [(u32, usize); (LOG_SIZE_BOUND - LOG_SIZE_MIN) as usize] = array::from_fn(|i| { let current_log_size = LOG_SIZE_MIN + i as u32; @@ -255,17 +235,27 @@ mod tests { (current_log_size, n_cols) }); let mut cols = accumulator.columns(n_cols_per_size); - for log_size in n_cols_per_size.iter().map(|(log_size, _)| *log_size) { + let mut eval_chunk_offset = 0; + for (log_size, n_cols) in n_cols_per_size.iter() { for index in 0..(1 << log_size) { let mut val = SecureField::zero(); - for (col_log_size, evaluation) in log_sizes.iter().zip(evaluations.iter()) { - if log_size != *col_log_size { + for (eval_index, (col_log_size, evaluation)) in + log_sizes.iter().zip(evaluations.iter()).enumerate() + { + if *log_size != *col_log_size { continue; } - val = val * alpha + evaluation[index]; + + // The random coefficient powers chunk is in regular order. + let random_coeff_chunk = + &cols[(log_size - LOG_SIZE_MIN) as usize].random_coeff_powers; + val += random_coeff_chunk + [random_coeff_chunk.len() - 1 - (eval_index - eval_chunk_offset)] + * evaluation[index]; } cols[(log_size - LOG_SIZE_MIN) as usize].accumulate(index, val); } + eval_chunk_offset += n_cols; } let accumulator_poly = accumulator.finalize(); @@ -273,13 +263,9 @@ mod tests { let point = CirclePoint::::get_point(98989892); let accumulator_res = accumulator_poly.eval_at_point(point); - // Sort evaluations by log_size. - let mut pairs = log_sizes.into_iter().zip(evaluations).collect::>(); - pairs.sort_by_key(|(log_size, _)| *log_size); - // Use direct computation. let mut res = SecureField::default(); - for (log_size, values) in pairs.into_iter() { + for (log_size, values) in log_sizes.into_iter().zip(evaluations) { res = res * alpha + CPUCircleEvaluation::new(CanonicCoset::new(log_size).circle_domain(), values) .interpolate() diff --git a/src/core/air/air_ext.rs b/src/core/air/air_ext.rs index 6fe86ac35..144c8fc41 100644 --- a/src/core/air/air_ext.rs +++ b/src/core/air/air_ext.rs @@ -32,8 +32,12 @@ pub trait AirExt: Air { random_coeff: SecureField, component_traces: &[ComponentTrace<'_, CPUBackend>], ) -> SecureCirclePoly { - let mut accumulator = - DomainEvaluationAccumulator::new(random_coeff, self.composition_log_degree_bound()); + let total_constraints: usize = self.components().iter().map(|c| c.n_constraints()).sum(); + let mut accumulator = DomainEvaluationAccumulator::new( + random_coeff, + self.composition_log_degree_bound(), + total_constraints, + ); zip(self.components(), component_traces).for_each(|(component, trace)| { component.evaluate_constraint_quotients_on_domain(trace, &mut accumulator) }); @@ -80,8 +84,7 @@ pub trait AirExt: Air { mask_values: &ComponentVec>, random_coeff: SecureField, ) -> SecureField { - let mut evaluation_accumulator = - PointEvaluationAccumulator::new(random_coeff, self.composition_log_degree_bound()); + let mut evaluation_accumulator = PointEvaluationAccumulator::new(random_coeff); zip(self.components(), &mask_values.0).for_each(|(component, mask)| { component.evaluate_constraint_quotients_at_point( point, diff --git a/src/core/air/mod.rs b/src/core/air/mod.rs index 7e99d4fc5..a052c3c07 100644 --- a/src/core/air/mod.rs +++ b/src/core/air/mod.rs @@ -58,6 +58,8 @@ impl Deref for Mask { /// A component is a set of trace columns of various sizes along with a set of /// constraints on them. pub trait Component { + fn n_constraints(&self) -> usize; + fn max_constraint_log_degree_bound(&self) -> u32; /// Returns the degree bounds of each trace column. diff --git a/src/core/prover/mod.rs b/src/core/prover/mod.rs index 41cc419e2..bd30b8c6c 100644 --- a/src/core/prover/mod.rs +++ b/src/core/prover/mod.rs @@ -253,6 +253,10 @@ mod tests { } impl Component for TestComponent { + fn n_constraints(&self) -> usize { + 0 + } + fn max_constraint_log_degree_bound(&self) -> u32 { self.max_constraint_log_degree_bound } @@ -279,7 +283,7 @@ mod tests { _mask: &crate::core::ColumnVec>, evaluation_accumulator: &mut PointEvaluationAccumulator, ) { - evaluation_accumulator.accumulate(1, qm31!(0, 0, 0, 1)) + evaluation_accumulator.accumulate(qm31!(0, 0, 0, 1)) } } diff --git a/src/examples/fibonacci/component.rs b/src/examples/fibonacci/component.rs index 885319fe1..a8e88db04 100644 --- a/src/examples/fibonacci/component.rs +++ b/src/examples/fibonacci/component.rs @@ -69,6 +69,10 @@ impl FibonacciComponent { } impl Component for FibonacciComponent { + fn n_constraints(&self) -> usize { + 2 + } + fn max_constraint_log_degree_bound(&self) -> u32 { // Step constraint is of degree 2. self.log_size + 1 @@ -83,7 +87,6 @@ impl Component for FibonacciComponent { trace: &ComponentTrace<'_, CPUBackend>, evaluation_accumulator: &mut DomainEvaluationAccumulator, ) { - let random_coeff = evaluation_accumulator.random_coeff; let poly = &trace.columns[0]; let trace_domain = CanonicCoset::new(self.log_size); let trace_eval_domain = CanonicCoset::new(self.log_size + 1).circle_domain(); @@ -104,9 +107,10 @@ impl Component for FibonacciComponent { let mul = trace_domain.step_size().div(point_coset.step_size); for (i, point) in point_coset.iter().enumerate() { let mask = [eval[i], eval[i as isize + mul], eval[i as isize + 2 * mul]]; - let res = self.step_constraint_eval_quotient_by_mask(point, &mask); - let res = res * random_coeff - + self.boundary_constraint_eval_quotient_by_mask(point, &[mask[0]]); + let mut res = self.boundary_constraint_eval_quotient_by_mask(point, &[mask[0]]) + * accum.random_coeff_powers[0]; + res += self.step_constraint_eval_quotient_by_mask(point, &mask) + * accum.random_coeff_powers[1]; accum.accumulate(bit_reverse_index(i + off, constraint_log_degree_bound), res); } } @@ -122,13 +126,10 @@ impl Component for FibonacciComponent { mask: &ColumnVec>, evaluation_accumulator: &mut PointEvaluationAccumulator, ) { - let constraints_log_degree_bound = self.log_size + 1; evaluation_accumulator.accumulate( - constraints_log_degree_bound, self.step_constraint_eval_quotient_by_mask(point, &mask[0][..].try_into().unwrap()), ); evaluation_accumulator.accumulate( - constraints_log_degree_bound, self.boundary_constraint_eval_quotient_by_mask( point, &mask[0][..1].try_into().unwrap(), diff --git a/src/examples/fibonacci/mod.rs b/src/examples/fibonacci/mod.rs index d179cc2d8..9603aa629 100644 --- a/src/examples/fibonacci/mod.rs +++ b/src/examples/fibonacci/mod.rs @@ -150,8 +150,7 @@ mod tests { .air .component .mask_points_and_values(point, &component_traces[0]); - let mut evaluation_accumulator = - PointEvaluationAccumulator::new(random_coeff, fib.air.composition_log_degree_bound()); + let mut evaluation_accumulator = PointEvaluationAccumulator::new(random_coeff); fib.air.component.evaluate_constraint_quotients_at_point( point, &mask_values, diff --git a/src/examples/wide_fibonacci/constraint_eval.rs b/src/examples/wide_fibonacci/constraint_eval.rs index 81eda35be..261b6437d 100644 --- a/src/examples/wide_fibonacci/constraint_eval.rs +++ b/src/examples/wide_fibonacci/constraint_eval.rs @@ -14,6 +14,10 @@ use crate::core::utils::bit_reverse_index; use crate::core::ColumnVec; impl Component for WideFibComponent { + fn n_constraints(&self) -> usize { + 62 + } + fn max_constraint_log_degree_bound(&self) -> u32 { self.log_size + 1 } @@ -52,9 +56,10 @@ impl Component for WideFibComponent { BaseField::batch_inverse(&denoms, &mut denom_inverses); let mut numerators = vec![SecureField::zero(); 1 << (self.max_constraint_log_degree_bound())]; - let random_coeff = evaluation_accumulator.random_coeff; - let [mut accum] = - evaluation_accumulator.columns([(self.max_constraint_log_degree_bound(), 64)]); + let [mut accum] = evaluation_accumulator + .columns([(self.max_constraint_log_degree_bound(), self.n_constraints())]); + // TODO (ShaharS) Change to get the correct power of random coeff inside the loop. + let random_coeff = accum.random_coeff_powers[1]; for (i, point_index) in eval_domain.iter_indices().enumerate() { numerators[i] = numerators[i] * random_coeff + (trace_evals[2].get_at(point_index) diff --git a/src/examples/wide_fibonacci/mod.rs b/src/examples/wide_fibonacci/mod.rs index 3680946f2..9923c499a 100644 --- a/src/examples/wide_fibonacci/mod.rs +++ b/src/examples/wide_fibonacci/mod.rs @@ -45,6 +45,7 @@ mod tests { let mut acc = DomainEvaluationAccumulator::new( QM31::from_u32_unchecked(1, 2, 3, 4), wide_fib.log_size + 1, + wide_fib.n_constraints(), ); let inputs = (0..1 << wide_fib.log_size) .map(|i| Input {