Constraint accumulator precompute random coeff powers (#554)

<!-- Reviewable:start -->
This change is [<img src="https://reviewable.io/review_button.svg" height="34" align="absmiddle" alt="Reviewable"/>](https://reviewable.io/reviews/starkware-libs/stwo/554)
<!-- Reviewable:end -->

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<SecureField>,
-    /// Number of accumulated evaluations for each log_size.
-    n_accumulated: Vec<usize>,
+    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<B: Backend> {
-    pub random_coeff: SecureField,
+    random_coeff_powers: Vec<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.
+    /// 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<SecureColumn<B>>,
-    /// Number of accumulated evaluations for each log_size.
-    n_cols_per_size: Vec<usize>,
 }
 
 impl<B: Backend> DomainEvaluationAccumulator<B> {
     /// 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<B: Backend> DomainEvaluationAccumulator<B> {
         &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<B: Backend> DomainEvaluationAccumulator<B> {
 impl DomainEvaluationAccumulator<CPUBackend> {
     /// 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::<CPUBackend>::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::<CPUBackend> {
                 columns: values.columns.map(|c| {
                     CPUCircleEvaluation::<_, BitReversedOrder>::new(
@@ -147,9 +126,8 @@ impl DomainEvaluationAccumulator<CPUBackend> {
                 }),
             };
             // 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<CPUBackend> {
     }
 }
 
-/// 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<SecureField>,
     col: &'a mut SecureColumn<B>,
 }
 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::<Vec<_>>();
-        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::<Vec<_>>();
+        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::<CPUBackend>::new(alpha, LOG_SIZE_BOUND);
+        let mut accumulator = DomainEvaluationAccumulator::<CPUBackend>::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,31 +235,37 @@ 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();
 
         // Pick an arbitrary sample point.
         let point = CirclePoint::<SecureField>::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::<Vec<_>>();
-        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()

src/core/air/air_ext.rs

@@ -32,8 +32,12 @@ pub trait AirExt: Air<CPUBackend> {
         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<CPUBackend> {
         mask_values: &ComponentVec<Vec<SecureField>>,
         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,

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<B: Backend> {
+    fn n_constraints(&self) -> usize;
+
     fn max_constraint_log_degree_bound(&self) -> u32;
 
     /// Returns the degree bounds of each trace column.

src/core/prover/mod.rs

@@ -253,6 +253,10 @@ mod tests {
     }
 
     impl Component<CPUBackend> 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<Vec<SecureField>>,
             evaluation_accumulator: &mut PointEvaluationAccumulator,
         ) {
-            evaluation_accumulator.accumulate(1, qm31!(0, 0, 0, 1))
+            evaluation_accumulator.accumulate(qm31!(0, 0, 0, 1))
         }
     }
 

src/examples/fibonacci/component.rs

@@ -69,6 +69,10 @@ impl FibonacciComponent {
 }
 
 impl Component<CPUBackend> 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<CPUBackend> for FibonacciComponent {
         trace: &ComponentTrace<'_, CPUBackend>,
         evaluation_accumulator: &mut DomainEvaluationAccumulator<CPUBackend>,
     ) {
-        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<CPUBackend> 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<CPUBackend> for FibonacciComponent {
         mask: &ColumnVec<Vec<SecureField>>,
         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(),

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,