Use rayon for parallelization

Author: ThomasdenH

Cargo.lock

@@ -20,6 +20,11 @@ rand_core = "0.9"
 bitvec = "1"
 deepfold = { path = "deepfold/deepfold" }
 util = { path = "deepfold/util" }
+rayon = { version = "1.12", optional = true }
+
+[features]
+default = [] # Add "rayon" here if you want it to be enabled by default
+rayon = ["dep:rayon"]
 
 [dev-dependencies]
 criterion = "0.3"

Cargo.toml

@@ -1,3 +1,6 @@
+#[cfg(feature = "rayon")]
+use rayon::prelude::*;
+
 mod keys;
 pub mod polynomials;
 
@@ -672,8 +675,13 @@ pub fn verify_pok<
 }
 
 pub fn interpolate_polynomial<F: Field>(evaluations: &[F]) -> MultilinearPolynomial<F> {
-    let mut coefficients: Vec<F> = vec![F::ZERO; evaluations.len()];
-    for (index, evaluation) in evaluations.iter().enumerate() {
+    #[cfg(feature = "rayon")]
+    let iter_1 = (0..evaluations.len()).into_par_iter();
+    #[cfg(not(feature = "rayon"))]
+    let iter_1 = 0..evaluations.len();
+
+    let coefficients = iter_1
+        .map(|coef_index|
         // Add eq(i,j) * evaluations[j] to the coefficients.
         //
         // We can interpret index as a sequence of bits, specifying the
@@ -687,21 +695,17 @@ pub fn interpolate_polynomial<F: Field>(evaluations: &[F]) -> MultilinearPolynom
         // - Evaluation point bit is 0, and coef bit is 1, -1.
         // - Evaluation point bit is 1, and coef bit is 0, 0.
         // - Evaluation point bit is 1, and coef bit is 1, 1.
-        for (coef_index, coefficient) in coefficients.iter_mut().enumerate() {
-            // If there is any bit where evaluation is 1, and coef is 0, set to 0.
-            if (!coef_index) & index != 0 {
-                continue;
-            }
-
-            let mut value = *evaluation;
-            // For every bit where evaluation is 0, and coef is 1, invert.
+        // If there is any bit where evaluation is 1, and coef is 0, set to 0.
+        evaluations.iter().copied().enumerate()
+            .filter(|(index, _)| (!coef_index) & index == 0)
+            .map(|(index, evaluation)|
             if ((!index) & coef_index).count_ones() % 2 == 1 {
-                value = -value;
+                -evaluation
+            } else {
+                evaluation
             }
-
-            *coefficient += value;
-        }
-    }
+        ).sum())
+        .collect();
     MultilinearPolynomial::new(coefficients)
 }
 

src/lib.rs

@@ -1,4 +1,6 @@
 //! These are the polynomials that are custom for our PIOP.
+#[cfg(feature = "rayon")]
+use rayon::prelude::*;
 use std::{marker::PhantomData, ops::MulAssign};
 
 use crate::{
@@ -15,7 +17,7 @@ mod q;
 
 pub use mask::Mask;
 pub use q::Q;
-use spongefish::{ProofResult, VerifierState, codecs::arkworks_algebra::{FieldToUnitDeserialize}};
+use spongefish::{codecs::arkworks_algebra::FieldToUnitDeserialize, ProofResult, VerifierState};
 
 #[derive(Clone)]
 pub struct P<
@@ -43,7 +45,7 @@ pub struct P<
     j_invariant_checkers: [PK::Checker; 2],
 }
 
-pub trait Checker<F: Field> {
+pub trait Checker<F: Field>: Sync {
     fn eval_at_x(&self, x: F) -> F;
     fn fix_variable(&mut self, x: F);
     fn value(&self) -> F;
@@ -66,8 +68,7 @@ impl<F: Field> Checker<F> for ACChecker<F> {
     fn eval_at_x(&self, _x: F) -> F {
         F::ZERO
     }
-    fn fix_variable(&mut self, _x: F) {
-    }
+    fn fix_variable(&mut self, _x: F) {}
     fn value(&self) -> F {
         F::ZERO
     }
@@ -227,8 +228,11 @@ impl<
         if self.variable_count() == VARIABLE_COUNT {
             let x = f;
             let i_hypercube_variables_mask = (1 << LOG_2_PATH_LENGTH) - 1;
-            (0..(1 << LOG_2_PATH_LENGTH))
-                .map(HypercubePoint)
+            #[cfg(feature = "rayon")]
+            let iter = (0..(1 << LOG_2_PATH_LENGTH)).into_par_iter();
+            #[cfg(not(feature = "rayon"))]
+            let iter = 0..(1 << LOG_2_PATH_LENGTH);
+            iter.map(HypercubePoint)
                 .map(|j| {
                     // We can compute i, since we already know that Q will be zero otherwise
                     let i_hypercube_remainder =
@@ -261,8 +265,11 @@ impl<
             let i_hypercube_variables_mask = (1 << remaining_i_vars) - 1;
             // Evaluate i
             // We can sum over j, and determine the qualifying i from this
-            (0..(1 << LOG_2_PATH_LENGTH))
-                .map(HypercubePoint)
+            #[cfg(feature = "rayon")]
+            let iter = (0..(1 << LOG_2_PATH_LENGTH)).into_par_iter();
+            #[cfg(not(feature = "rayon"))]
+            let iter = 0..(1 << LOG_2_PATH_LENGTH);
+            iter.map(HypercubePoint)
                 .map(|j| {
                     // Now compute j, the first couple of variables are free,
                     // the others are determined by i.
@@ -295,8 +302,11 @@ impl<
             debug_assert!(self.b_0_j.variable_count() > 0);
             debug_assert!(self.b_1_j.variable_count() > 0);
             let variables_j = self.b_1_j.variable_count() - 1;
-            (0usize..(1 << variables_j))
-                .map(HypercubePoint)
+            #[cfg(feature = "rayon")]
+            let iter = (0..(1 << variables_j)).into_par_iter();
+            #[cfg(not(feature = "rayon"))]
+            let iter = 0..(1 << variables_j);
+            iter.map(HypercubePoint)
                 .map(|j_hypercube_remainder| {
                     let b_0_i = self.b_0_i.get_as_const();
                     let b_1_i = self.b_1_i.get_as_const();
@@ -396,7 +406,8 @@ impl<
             self.e_0
                 // x has been absorbed in e_0
                 * self.q.eval(q_eval_point)
-                * (self.a_a_j * b_0_j + self.a_a_i * b_0_i) * (b_0_j - b_0_i) - (b_1_i + b_1_j)
+                * (self.a_a_j * b_0_j + self.a_a_i * b_0_i) * (b_0_j - b_0_i)
+                - (b_1_i + b_1_j)
                 + self.mask.eval(point)
                 + self
                     .j_invariant_checkers

src/polynomials/piop_polynomials.rs

@@ -174,8 +174,9 @@ impl<F: Field> JInvariantChecker<F> {
         let mut val = d5 * (F::from(4) * c - d1) - F::ONE;
         val = val * h + c * d4 - F::ONE;
         val = val * h + j * d2 * (F::from(4) * c - d1)
-            + F::from(-256)
-                * (d3 * (F::from(9) * c - d1) + F::from(27) * d2 * (d1 - c));
+            - F::from(256)
+                * (-d1 * d3 + F::from(3 * 3) * c * d3 - F::from(9 * 3) * d2 * d1
+                    + F::from(27) * d2 * c);
         val = val * h + d3 - d1.pow([2]);
         val = val * h + d2 - c.pow([2]);
         val = val * h + d1 - a.pow([2]);