Optimize prover boundary evaluations in construction of composition p…

…olynomial (#22)

* Compute boundary evaluations and batch inverse of zerofier outside of for loop in round 2

* Clippy and format

* evaluate boundary polys before collecting the polynomials

* add space

---------

Co-authored-by: Sergio Chouhy <sergio.chouhy@gmail.com>

src/air/constraints/evaluator.rs

@@ -5,7 +5,11 @@ use lambdaworks_math::{
 };
 
 use super::{boundary::BoundaryConstraints, evaluation_table::ConstraintEvaluationTable};
-use crate::air::{frame::Frame, trace::TraceTable, traits::AIR};
+use crate::{
+    air::{frame::Frame, trace::TraceTable, traits::AIR},
+    prover::evaluate_polynomial_on_lde_domain,
+    Domain,
+};
 use std::iter::zip;
 
 pub struct ConstraintEvaluator<'poly, F: IsFFTField, A: AIR> {
@@ -36,7 +40,7 @@ impl<'poly, F: IsFFTField, A: AIR + AIR<Field = F>> ConstraintEvaluator<'poly, F
     pub fn evaluate(
         &self,
         lde_trace: &TraceTable<F>,
-        lde_domain: &[FieldElement<F>],
+        domain: &Domain<F>,
         alpha_and_beta_transition_coefficients: &[(FieldElement<F>, FieldElement<F>)],
         alpha_and_beta_boundary_coefficients: &[(FieldElement<F>, FieldElement<F>)],
         rap_challenges: &A::RAPChallenges,
@@ -47,7 +51,7 @@ impl<'poly, F: IsFFTField, A: AIR + AIR<Field = F>> ConstraintEvaluator<'poly, F
         // The + 1 is for the boundary constraints column
         let mut evaluation_table = ConstraintEvaluationTable::new(
             self.air.context().num_transition_constraints() + 1,
-            lde_domain,
+            &domain.lde_roots_of_unity_coset,
         );
         let n_trace_colums = self.trace_polys.len();
         let boundary_constraints = &self.boundary_constraints;
@@ -56,38 +60,69 @@ impl<'poly, F: IsFFTField, A: AIR + AIR<Field = F>> ConstraintEvaluator<'poly, F
             boundary_constraints.generate_roots_of_unity(&self.primitive_root, n_trace_colums);
         let values = boundary_constraints.values(n_trace_colums);
 
-        let boundary_polys: Vec<Polynomial<FieldElement<F>>> = zip(domains, values)
+        #[cfg(debug_assertions)]
+        let mut boundary_polys = Vec::new();
+
+        let boundary_polys_evaluations: Vec<Vec<FieldElement<F>>> = zip(domains, values)
             .zip(self.trace_polys)
             .map(|((xs, ys), trace_poly)| {
-                trace_poly
+                let boundary_poly = trace_poly
                     - &Polynomial::interpolate(&xs, &ys)
-                        .expect("xs and ys have equal length and xs are unique")
+                        .expect("xs and ys have equal length and xs are unique");
+
+                #[cfg(debug_assertions)]
+                boundary_polys.push(boundary_poly.clone());
+
+                evaluate_polynomial_on_lde_domain(
+                    &boundary_poly,
+                    domain.blowup_factor,
+                    domain.interpolation_domain_size,
+                    &domain.coset_offset,
+                )
+                .unwrap()
             })
             .collect();
 
-        let boundary_zerofiers: Vec<Polynomial<FieldElement<F>>> = (0..n_trace_colums)
+        #[cfg(debug_assertions)]
+        let mut boundary_zerofiers = Vec::new();
+
+        let boundary_zerofiers_inverse_evaluations: Vec<Vec<FieldElement<F>>> = (0..n_trace_colums)
             .map(|col| {
-                self.boundary_constraints
-                    .compute_zerofier(&self.primitive_root, col)
+                let zerofier = self
+                    .boundary_constraints
+                    .compute_zerofier(&self.primitive_root, col);
+
+                #[cfg(debug_assertions)]
+                boundary_zerofiers.push(zerofier.clone());
+
+                let mut evals = evaluate_polynomial_on_lde_domain(
+                    &zerofier,
+                    domain.blowup_factor,
+                    domain.interpolation_domain_size,
+                    &domain.coset_offset,
+                )
+                .unwrap();
+                FieldElement::inplace_batch_inverse(&mut evals);
+                evals
             })
             .collect();
 
-        let blowup_factor = self.air.blowup_factor();
-
         #[cfg(debug_assertions)]
         for (poly, z) in boundary_polys.iter().zip(boundary_zerofiers.iter()) {
             let (_, b) = poly.clone().long_division_with_remainder(z);
             assert_eq!(b, Polynomial::zero());
         }
 
+        let blowup_factor = self.air.blowup_factor();
+
         #[cfg(debug_assertions)]
         let mut transition_evaluations = Vec::new();
 
         let divisors = self.air.transition_divisors();
         let boundary_term_degree_adjustment =
             self.air.composition_poly_degree_bound() - self.air.context().trace_length;
         // Iterate over trace and domain and compute transitions
-        for (i, d) in lde_domain.iter().enumerate() {
+        for (i, d) in domain.lde_roots_of_unity_coset.iter().enumerate() {
             let frame = Frame::read_from_trace(
                 lde_trace,
                 i,
@@ -109,16 +144,22 @@ impl<'poly, F: IsFFTField, A: AIR + AIR<Field = F>> ConstraintEvaluator<'poly, F
             );
 
             let d_adjustment_power = d.pow(boundary_term_degree_adjustment);
-            let boundary_evaluation = zip(&boundary_polys, &boundary_zerofiers)
-                .enumerate()
-                .map(|(index, (boundary_poly, boundary_zerofier))| {
+            let boundary_evaluation = zip(
+                &boundary_polys_evaluations,
+                &boundary_zerofiers_inverse_evaluations,
+            )
+            .enumerate()
+            .map(
+                |(index, (boundary_poly_evaluation, zerofier_inverse_evaluation))| {
                     let (boundary_alpha, boundary_beta) =
                         alpha_and_beta_boundary_coefficients[index].clone();
 
-                    (boundary_poly.evaluate(d) / boundary_zerofier.evaluate(d))
+                    &boundary_poly_evaluation[i]
+                        * &zerofier_inverse_evaluation[i]
                         * (&boundary_alpha * &d_adjustment_power + &boundary_beta)
-                })
-                .fold(FieldElement::<F>::zero(), |acc, eval| acc + eval);
+                },
+            )
+            .fold(FieldElement::<F>::zero(), |acc, eval| acc + eval);
 
             evaluations.push(boundary_evaluation);
 

src/prover.rs

@@ -92,7 +92,7 @@ where
     (trees, roots)
 }
 
-fn evaluate_polynomial_on_lde_domain<F>(
+pub fn evaluate_polynomial_on_lde_domain<F>(
     p: &Polynomial<FieldElement<F>>,
     blowup_factor: usize,
     domain_size: usize,
@@ -225,7 +225,7 @@ where
 
     let constraint_evaluations = evaluator.evaluate(
         &round_1_result.lde_trace,
-        &domain.lde_roots_of_unity_coset,
+        domain,
         transition_coeffs,
         boundary_coeffs,
         &round_1_result.rap_challenges,