feat: using plonky3 maybe-rayon instead of rayon for optional parallelism

Makefile.toml

@@ -4,6 +4,7 @@ CARGO_MAKE_EXTEND_WORKSPACE_MAKEFILE = true
 [tasks.tests]
 args = [
   "test",
+  "--all-features",
   # Run everything but 'benches'.
   "--lib",
   "--bins",
@@ -16,6 +17,14 @@ env = { RUST_MIN_STACK = "33554432" }
 workspace = false
 
 [tasks.clippy]
-args = ["clippy", "--workspace", "--all-targets", "--", "-D", "warnings"]
+args = [
+  "clippy",
+  "--workspace",
+  "--all-targets",
+  "--all-features",
+  "--",
+  "-D",
+  "warnings",
+]
 command = "cargo"
 workspace = false

crates/curves/Cargo.toml

@@ -24,9 +24,9 @@ serde = { workspace = true, features = ["derive"] }
 typenum = "1.17.0"
 
 cfg-if.workspace = true
-ff_ext = { version = "0.1.0", path = "../ff_ext" }
-itertools = { workspace = true }
-multilinear_extensions = { version = "0.1.0", path = "../multilinear_extensions" }
+ff_ext.workspace = true
+itertools.workspace = true
+multilinear_extensions.workspace = true
 rug = { version = "1.26.1", optional = true }
 
 [dev-dependencies]

crates/mpcs/Cargo.toml

@@ -14,7 +14,7 @@ bincode = "1.3.3"
 clap.workspace = true
 ff_ext.workspace = true
 itertools.workspace = true
-whir.workspace = true
+whir = { workspace = true, optional = true }
 multilinear_extensions.workspace = true
 num-integer = "0.1"
 p3.workspace = true
@@ -32,18 +32,9 @@ witness.workspace = true
 criterion.workspace = true
 
 [features]
-benchmark = ["parallel"]
-default = ["parallel"] # Add "sanity-check" to debug
-nightly-features = [
-  "p3/nightly-features",
-  "ff_ext/nightly-features",
-  "multilinear_extensions/nightly-features",
-  "sumcheck/nightly-features",
-  "transcript/nightly-features",
-  "whir/nightly-features",
-  "witness/nightly-features",
-]
-parallel = ["dep:rayon"]
+nightly-features = ["ff_ext/nightly-features"]
+parallel = ["p3/parallel", "dep:rayon"]
+whir = ["dep:whir"]
 print-trace = ["whir/print-trace"]
 sanity-check = []
 

crates/mpcs/benches/interpolate.rs

@@ -7,9 +7,9 @@ use itertools::Itertools;
 use mpcs::util::arithmetic::interpolate_field_type_over_boolean_hypercube;
 
 use multilinear_extensions::mle::{FieldType, MultilinearExtension};
+use p3::maybe_rayon::prelude::*;
 use rand::{SeedableRng, rngs::OsRng};
 use rand_chacha::ChaCha8Rng;
-use rayon::iter::{IntoParallelRefIterator, ParallelIterator};
 
 type E = GoldilocksExt2;
 

crates/mpcs/src/basefold/commit_phase.rs

@@ -38,11 +38,7 @@ use multilinear_extensions::{
     mle::{IntoMLE, MultilinearExtension},
     virtual_poly::build_eq_x_r_vec,
 };
-use rayon::{
-    iter::{IntoParallelIterator, IntoParallelRefMutIterator},
-    prelude::{IndexedParallelIterator, ParallelIterator},
-    slice::ParallelSlice,
-};
+use p3::maybe_rayon::prelude::*;
 
 use super::structure::BasefoldCommitmentWithWitness;
 

crates/mpcs/src/lib.rs

@@ -258,6 +258,7 @@ pub enum Error {
     PolynomialSizesNotEqual,
     MerkleRootMismatch,
     PointEvalMismatch(String),
+    #[cfg(feature = "whir")]
     WhirError(whir_external::error::Error),
 }
 
@@ -266,9 +267,11 @@ pub use basefold::{
     Basefold, BasefoldCommitment, BasefoldCommitmentWithWitness, BasefoldDefault, BasefoldParams,
     BasefoldRSParams, BasefoldSpec, EncodingScheme, RSCode, RSCodeDefaultSpec,
 };
+#[cfg(feature = "whir")]
 extern crate whir as whir_external;
 mod whir;
 use multilinear_extensions::mle::ArcMultilinearExtension;
+#[cfg(feature = "whir")]
 pub use whir::{Whir, WhirDefault, WhirDefaultSpec};
 
 // TODO: Need to use some functions here in the integration benchmarks. But

crates/mpcs/src/util.rs

@@ -101,7 +101,7 @@ pub(crate) fn codeword_fold_with_challenge<E: ExtensionField>(
     lo + challenge * (hi - lo)
 }
 
-#[cfg(any(test, feature = "benchmark"))]
+#[cfg(test)]
 pub mod test {
     use ff_ext::FromUniformBytes;
     use rand::{

crates/multilinear_extensions/Cargo.toml

@@ -26,9 +26,8 @@ log = "0.4"
 serde_json.workspace = true
 
 [features]
-default = ["parallel"]
-nightly-features = ["p3/nightly-features", "ff_ext/nightly-features"]
-parallel = []
+nightly-features = ["ff_ext/nightly-features"]
+parallel = ["p3/parallel"]
 
 [[bench]]
 harness = false

crates/multilinear_extensions/src/expression.rs

@@ -9,8 +9,7 @@ use crate::{
 };
 use ff_ext::{ExtensionField, SmallField};
 use itertools::{Either, Itertools, chain, izip};
-use p3::field::FieldAlgebra;
-use rayon::iter::{IntoParallelIterator, IntoParallelRefIterator, ParallelIterator};
+use p3::{field::FieldAlgebra, maybe_rayon::prelude::*};
 use serde::de::DeserializeOwned;
 use std::{
     cmp::max,

crates/multilinear_extensions/src/mle.rs

@@ -9,15 +9,13 @@ use crate::{
 };
 use either::Either;
 use ff_ext::{ExtensionField, FromUniformBytes};
-use p3::field::{Field, FieldAlgebra};
-use rand::Rng;
-use rayon::{
-    iter::{
-        IndexedParallelIterator, IntoParallelIterator, IntoParallelRefIterator,
-        IntoParallelRefMutIterator, ParallelIterator,
-    },
-    slice::ParallelSliceMut,
+#[cfg(not(feature = "parallel"))]
+use itertools::Itertools;
+use p3::{
+    field::{Field, FieldAlgebra},
+    maybe_rayon::prelude::*,
 };
+use rand::Rng;
 use serde::{Deserialize, Serialize, de::DeserializeOwned};
 use std::fmt::Debug;
 
@@ -612,10 +610,8 @@ impl<'a, E: ExtensionField> MultilinearExtension<'a, E> {
                         Cow::Owned(MultilinearExtension::from_evaluations_ext_vec(
                             self.num_vars() - 1,
                             evaluations
-                                .par_iter()
-                                .chunks(2)
-                                .with_min_len(64)
-                                .map(|buf| *point * (*buf[1] - *buf[0]) + *buf[0])
+                                .par_chunks(2)
+                                .map(|buf| *point * (buf[1] - buf[0]) + buf[0])
                                 .collect(),
                         ))
                     });
@@ -645,10 +641,8 @@ impl<'a, E: ExtensionField> MultilinearExtension<'a, E> {
             match &mut self.evaluations {
                 FieldType::Base(slice) => {
                     let slice_ext = slice
-                        .par_iter()
-                        .chunks(2)
-                        .with_min_len(64)
-                        .map(|buf| *point * (*buf[1] - *buf[0]) + *buf[0])
+                        .par_chunks(2)
+                        .map(|buf| *point * (buf[1] - buf[0]) + buf[0])
                         .collect();
                     let _ = mem::replace(
                         &mut self.evaluations,
@@ -658,10 +652,8 @@ impl<'a, E: ExtensionField> MultilinearExtension<'a, E> {
                 FieldType::Ext(slice) => {
                     let slice_mut = slice.to_mut();
                     slice_mut
-                        .par_iter_mut()
-                        .chunks(2)
-                        .with_min_len(64)
-                        .for_each(|mut buf| *buf[0] = *buf[0] + (*buf[1] - *buf[0]) * *point);
+                        .par_chunks_mut(2)
+                        .for_each(|buf| buf[0] = buf[0] + (buf[1] - buf[0]) * *point);
 
                     // sequentially update buf[b1, b2,..bt] = buf[b1, b2,..bt, 0]
                     for index in 0..1 << (max_log2_size - 1) {

crates/multilinear_extensions/src/virtual_poly.rs

@@ -10,12 +10,8 @@ use crate::{
 use either::Either;
 use ff_ext::ExtensionField;
 use itertools::Itertools;
-use p3::field::Field;
+use p3::{field::Field, maybe_rayon::prelude::*};
 use rand::Rng;
-use rayon::{
-    iter::{IndexedParallelIterator, IntoParallelIterator, ParallelIterator},
-    slice::ParallelSliceMut,
-};
 use serde::{Deserialize, Serialize, de::DeserializeOwned};
 
 pub type MonomialTermsType<'a, E> =

crates/p3/Cargo.toml

@@ -27,7 +27,6 @@ p3-symmetric.workspace = true
 p3-util.workspace = true
 
 [features]
-default = ["parallel"]
 nightly-features = [
   "p3-baby-bear/nightly-features",
   "p3-dft/nightly-features",

crates/poseidon/Cargo.toml

@@ -18,4 +18,4 @@ serde.workspace = true
 rand.workspace = true
 
 [features]
-nightly-features = ["p3/nightly-features", "ff_ext/nightly-features"]
+nightly-features = ["ff_ext/nightly-features"]

crates/sumcheck/Cargo.toml

@@ -13,16 +13,14 @@ version.workspace = true
 either.workspace = true
 ff_ext.workspace = true
 itertools.workspace = true
-multilinear_extensions = { path = "../multilinear_extensions", features = [
-  "parallel",
-] }
+multilinear_extensions.workspace = true
 p3.workspace = true
 rayon.workspace = true
 serde.workspace = true
-sumcheck_macro = { path = "../sumcheck_macro" }
+sumcheck_macro.workspace = true
 thiserror.workspace = true
 tracing.workspace = true
-transcript = { path = "../transcript" }
+transcript.workspace = true
 
 [dev-dependencies]
 criterion.workspace = true
@@ -31,13 +29,9 @@ poseidon.workspace = true
 rand.workspace = true
 
 [features]
-nightly-features = [
-  "p3/nightly-features",
-  "ff_ext/nightly-features",
-  "multilinear_extensions/nightly-features",
-  "sumcheck_macro/nightly-features",
-  "transcript/nightly-features",
-]
+nightly-features = ["ff_ext/nightly-features"]
+parallel = ["p3/parallel", "multilinear_extensions/parallel"]
+
 
 [[bench]]
 harness = false

crates/sumcheck/src/prover.rs

@@ -10,10 +10,7 @@ use multilinear_extensions::{
     virtual_poly::{MonomialTerms, VirtualPolynomial},
     virtual_polys::{PolyMeta, VirtualPolynomials},
 };
-use rayon::{
-    iter::{IndexedParallelIterator, IntoParallelRefIterator, IntoParallelRefMutIterator},
-    prelude::{IntoParallelIterator, ParallelIterator},
-};
+use p3::maybe_rayon::prelude::*;
 use sumcheck_macro::sumcheck_code_gen;
 use transcript::{Challenge, Transcript};
 
@@ -41,12 +38,14 @@ impl<'a, E: ExtensionField> Phase1Workers<'a, E> {
     ) -> (Vec<IOPProverState<'a, E>>, Vec<IOPProverMessage<E>>) {
         let mut prover_msgs = Vec::with_capacity(num_variables);
         for _ in 0..num_variables {
-            let evaluations = self
-                .workers_states
-                .par_iter_mut()
-                .map(|state| state.run_round())
-                .reduce(|| AdditiveVec::new(max_degree + 1), |a, b| a + b);
-
+            let evaluations: AdditiveVec<E> = self.workers_states.par_iter_mut().par_fold_reduce(
+                || AdditiveVec::new(max_degree + 1),
+                |mut acc, state| {
+                    acc += state.run_round();
+                    acc
+                },
+                |a, b| a + b,
+            );
             transcript.append_field_element_exts(&evaluations.0);
 
             let challenge = transcript.sample_and_append_challenge(b"Internal round");
@@ -166,8 +165,8 @@ impl<'a, E: ExtensionField> IOPProverState<'a, E> {
                 max_thread_id,
                 num_variables,
                 poly_meta,
-                polys,
                 max_degree,
+                polys,
                 transcript,
             );
             exit_span!(span);
@@ -234,8 +233,8 @@ impl<'a, E: ExtensionField> IOPProverState<'a, E> {
         max_thread_id: usize,
         num_variables: usize,
         poly_meta: Vec<PolyMeta>,
-        mut polys: Vec<VirtualPolynomial<'a, E>>,
         max_degree: usize,
+        mut polys: Vec<VirtualPolynomial<'a, E>>,
         transcript: &mut impl Transcript<E>,
     ) -> (Vec<IOPProverState<'a, E>>, Vec<IOPProverMessage<E>>) {
         let log2_max_thread_id = ceil_log2(max_thread_id); // do not support SIZE not power of 2
@@ -599,8 +598,8 @@ impl<'a, E: ExtensionField> IOPProverState<'a, E> {
             .poly
             .products
             .par_iter()
-            .fold_with(
-                AdditiveVec::new(self.poly.aux_info.max_degree + 1),
+            .par_fold_reduce(
+                || AdditiveVec::new(self.poly.aux_info.max_degree + 1),
                 |mut uni_polys, MonomialTerms { terms }| {
                     for Term {
                         scalar,
@@ -642,9 +641,9 @@ impl<'a, E: ExtensionField> IOPProverState<'a, E> {
                     }
                     uni_polys
                 },
-            )
-            .reduce_with(|acc, item| acc + item)
-            .unwrap();
+                |acc, item| acc + item,
+            );
+
         exit_span!(span);
 
         exit_span!(start);

crates/sumcheck_macro/Cargo.toml

@@ -21,9 +21,9 @@ rand.workspace = true
 syn = { version = "2.0", features = ["full"] }
 
 [features]
-nightly-features = ["p3/nightly-features"]
+nightly-features = ["ff_ext/nightly-features"]
 
 [dev-dependencies]
 ff_ext.workspace = true
-multilinear_extensions = { path = "../multilinear_extensions" }
+multilinear_extensions.workspace = true
 sumcheck.workspace = true

crates/sumcheck_macro/src/lib.rs

@@ -237,7 +237,7 @@ pub fn sumcheck_code_gen(input: proc_macro::TokenStream) -> proc_macro::TokenStr
         };
 
         let iter = if parallalize {
-            quote! {.into_par_iter().step_by(2).rev().with_min_len(64)}
+            quote! {.into_par_iter().step_by(2).rev()}
         } else {
             quote! {.step_by(2).rev()}
         };

crates/transcript/Cargo.toml

@@ -16,8 +16,4 @@ p3.workspace = true
 poseidon.workspace = true
 
 [features]
-nightly-features = [
-  "p3/nightly-features",
-  "ff_ext/nightly-features",
-  "poseidon/nightly-features",
-]
+nightly-features = ["ff_ext/nightly-features"]