(slightly) Faster ML evaluations (#342)

* refactor: Refactor polynomial evaluation and proof verification in Spartan

- Refactor the `evaluate` function in `MultilinearPolynomial` struct for better clarity and performance.
- Adjust the `prove` function in `ppsnark.rs` to utilize `chis` method when evaluating a multilinear polynomial.
- Modify the `RelaxedR1CSSNARK` struct's prove and verify methods.
- Add more claims regarding the W and E polynomials to the verification process.
- Implement batch evaluation for verifying the W and E polynomials.

* fix: eval_table

src/spartan/batched.rs

@@ -434,9 +434,15 @@ impl<E: Engine, EE: EvaluationEngineTrait<E>> BatchedRelaxedR1CSSNARKTrait<E>
       )
     }
 
+    let chis_r_x = r_x
+      .par_iter()
+      .map(|r_x| EqPolynomial::evals_from_points(r_x))
+      .collect::<Vec<_>>();
+
     // Evaluate τ(rₓ) for each instance
-    let evals_tau = zip_with!(iter, (polys_tau, r_x), |poly_tau, r_x| poly_tau
-      .evaluate(r_x));
+    let evals_tau = zip_with!(iter, (polys_tau, chis_r_x), |poly_tau, er_x| {
+      MultilinearPolynomial::evaluate_with_chis(poly_tau.evaluations(), er_x)
+    });
 
     // Compute expected claim for all instances ∑ᵢ rⁱ⋅τ(rₓ)⋅(Azᵢ⋅Bzᵢ − uᵢ⋅Czᵢ − Eᵢ)
     let claim_outer_final_expected = zip_with!(
@@ -503,7 +509,7 @@ impl<E: Engine, EE: EvaluationEngineTrait<E>> BatchedRelaxedR1CSSNARKTrait<E>
 
     // compute evaluations of R1CS matrices M(r_x, r_y) = eq(r_y)ᵀ⋅M⋅eq(r_x)
     let multi_evaluate = |M_vec: &[&SparseMatrix<E::Scalar>],
-                          r_x: &[E::Scalar],
+                          chi_r_x: &[E::Scalar],
                           r_y: &[E::Scalar]|
      -> Vec<E::Scalar> {
       let evaluate_with_table =
@@ -520,21 +526,19 @@ impl<E: Engine, EE: EvaluationEngineTrait<E>> BatchedRelaxedR1CSSNARKTrait<E>
             .sum()
         };
 
-      let (T_x, T_y) = rayon::join(
-        || EqPolynomial::evals_from_points(r_x),
-        || EqPolynomial::evals_from_points(r_y),
-      );
+      let T_x = chi_r_x;
+      let T_y = EqPolynomial::evals_from_points(r_y);
 
       M_vec
         .par_iter()
-        .map(|&M_vec| evaluate_with_table(M_vec, &T_x, &T_y))
+        .map(|&M_vec| evaluate_with_table(M_vec, T_x, &T_y))
         .collect()
     };
 
     // Compute inner claim ∑ᵢ r³ⁱ⋅(Aᵢ(r_x, r_y) + r⋅Bᵢ(r_x, r_y) + r²⋅Cᵢ(r_x, r_y))⋅Zᵢ(r_y)
     let claim_inner_final_expected = zip_with!(
       iter,
-      (vk.S, r_x, r_y, evals_Z, inner_r_powers),
+      (vk.S, chis_r_x, r_y, evals_Z, inner_r_powers),
       |S, r_x, r_y, eval_Z, r_i| {
         let evals = multi_evaluate(&[&S.A, &S.B, &S.C], r_x, r_y);
         let eval = evals[0] + inner_r * evals[1] + inner_r_square * evals[2];

src/spartan/polys/multilinear.rs

@@ -105,25 +105,18 @@ impl<Scalar: PrimeField> MultilinearPolynomial<Scalar> {
   pub fn evaluate(&self, r: &[Scalar]) -> Scalar {
     // r must have a value for each variable
     assert_eq!(r.len(), self.get_num_vars());
-    let chis = EqPolynomial::evals_from_points(r);
-
-    zip_with!(
-      (chis.into_par_iter(), self.Z.par_iter()),
-      |chi_i, Z_i| chi_i * Z_i
-    )
-    .sum()
+    Self::evaluate_with(&self.Z, r)
   }
 
   /// Evaluates the polynomial with the given evaluations and point.
   pub fn evaluate_with(Z: &[Scalar], r: &[Scalar]) -> Scalar {
-    zip_with!(
-      (
-        EqPolynomial::evals_from_points(r).into_par_iter(),
-        Z.par_iter()
-      ),
-      |a, b| a * b
-    )
-    .sum()
+    let chis = EqPolynomial::evals_from_points(r);
+    Self::evaluate_with_chis(Z, &chis)
+  }
+
+  /// Evaluates the polynomial with the given evaluations and chi coefficients
+  pub fn evaluate_with_chis(Z: &[Scalar], chis: &[Scalar]) -> Scalar {
+    zip_with!(par_iter, (chis, Z), |a, b| *a * b).sum()
   }
 }
 

src/spartan/ppsnark.rs

@@ -553,10 +553,11 @@ impl<E: Engine, EE: EvaluationEngineTrait<E>> RelaxedR1CSSNARKTrait<E> for Relax
 
       (Az, Bz, Cz, W, E)
     };
+    let chis_taus = EqPolynomial::evals_from_points(&tau_coords);
     let (eval_Az_at_tau, eval_Bz_at_tau, eval_Cz_at_tau) = {
       let evals_at_tau = [&Az, &Bz, &Cz]
         .into_par_iter()
-        .map(|p| MultilinearPolynomial::evaluate_with(p, &tau_coords))
+        .map(|p| MultilinearPolynomial::evaluate_with_chis(p, &chis_taus))
         .collect::<Vec<E::Scalar>>();
       (evals_at_tau[0], evals_at_tau[1], evals_at_tau[2])
     };

src/spartan/snark.rs

@@ -181,8 +181,10 @@ impl<E: Engine, EE: EvaluationEngineTrait<E>> RelaxedR1CSSNARKTrait<E> for Relax
 
     // claims from the end of sum-check
     let (claim_Az, claim_Bz): (E::Scalar, E::Scalar) = (claims_outer[1], claims_outer[2]);
-    let claim_Cz = poly_Cz.evaluate(&r_x);
-    let eval_E = MultilinearPolynomial::evaluate_with(&W.E, &r_x);
+    let chis_r_x = EqPolynomial::evals_from_points(&r_x);
+
+    let claim_Cz = MultilinearPolynomial::evaluate_with_chis(poly_Cz.evaluations(), &chis_r_x);
+    let eval_E = MultilinearPolynomial::evaluate_with_chis(&W.E, &chis_r_x);
     transcript.absorb(
       b"claims_outer",
       &[claim_Az, claim_Bz, claim_Cz, eval_E].as_slice(),

src/supernova/circuit.rs

@@ -1,6 +1,6 @@
 //! Supernova implementation support arbitrary argumented circuits and running instances.
 //! There are two Verification Circuits for each argumented circuit: The primary and the secondary.
-//! Each of them is over a Pasta curve but
+//! Each of them is over a cycle curve but
 //! only the primary executes the next step of the computation.
 //! Each circuit takes as input 2 hashes.
 //! Each circuit folds the last invocation of the other into the respective running instance, specified by `augmented_circuit_index`
@@ -10,7 +10,7 @@
 //!    1. Ui[] are contained in X[0] hash pre-image.
 //!    2. R1CS Instance u is folded into Ui[augmented_circuit_index] correctly; just like Nova IVC.
 //!    3. (optional by F logic) F circuit might check `program_counter_{i}` invoked current F circuit is legal or not.
-//!    3. F circuit produce `program_counter_{i+1}` and sent to next round for optionally constraint the next F' argumented circuit.
+//!    3. F circuit produce `program_counter_{i+1}` and sent to next round to optionally constraint the next F' argumented circuit.
 use crate::{
   constants::{NIO_NOVA_FOLD, NUM_HASH_BITS},
   gadgets::{

src/supernova/utils.rs

@@ -37,6 +37,7 @@ pub fn get_from_vec_alloc_relaxed_r1cs<E: Engine, CS: ConstraintSystem<<E as Eng
   // Since `selector_vec` is correct, only one entry is 1.
   // If selector_vec[0] is 1, then all `conditionally_select` will return `first`.
   // Otherwise, the correct instance will be selected.
+  // TODO: reformulate when iterator_try_reduce stabilizes
   let selected = a
     .iter()
     .zip_eq(selector_vec.iter())