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

crates/mpcs/Cargo.toml

@@ -18,6 +18,7 @@ whir.workspace = true
 multilinear_extensions.workspace = true
 num-integer = "0.1"
 p3.workspace = true
+p3-maybe-rayon.workspace = true
 rand.workspace = true
 rand_chacha.workspace = true
 rayon = { workspace = true, optional = true }
@@ -32,8 +33,6 @@ 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",
@@ -43,7 +42,7 @@ nightly-features = [
   "whir/nightly-features",
   "witness/nightly-features",
 ]
-parallel = ["dep:rayon"]
+parallel = ["p3-maybe-rayon/parallel", "dep:rayon"]
 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/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

@@ -14,6 +14,7 @@ either.workspace = true
 ff_ext.workspace = true
 itertools.workspace = true
 p3.workspace = true
+p3-maybe-rayon.workspace = true
 rand.workspace = true
 rayon.workspace = true
 serde.workspace = true
@@ -26,9 +27,8 @@ log = "0.4"
 serde_json.workspace = true
 
 [features]
-default = ["parallel"]
 nightly-features = ["p3/nightly-features", "ff_ext/nightly-features"]
-parallel = []
+parallel = ["p3-maybe-rayon/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,11 @@ 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,
+use p3::{
+    field::{Field, FieldAlgebra},
+    maybe_rayon::prelude::*,
 };
+use rand::Rng;
 use serde::{Deserialize, Serialize, de::DeserializeOwned};
 use std::fmt::Debug;
 

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/whir/Cargo.toml

@@ -18,6 +18,7 @@ clap.workspace = true
 ff_ext = { path = "../ff_ext" }
 multilinear_extensions = { path = "../multilinear_extensions" }
 p3.workspace = true
+p3-maybe-rayon.workspace = true
 rand = "0.8"
 rand_chacha = "0.3"
 rayon = { workspace = true, optional = true }
@@ -42,6 +43,6 @@ nightly-features = [
   "transcript/nightly-features",
   "witness/nightly-features",
 ]
-parallel = ["dep:rayon"]
+parallel = ["p3-maybe-rayon/parallel", "dep:rayon"]
 print-trace = ["tracing/log"]
 rayon = ["dep:rayon"]

crates/whir/src/crypto/mod.rs

@@ -5,8 +5,8 @@ use p3::{
         Dimensions,
         dense::{DenseMatrix, RowMajorMatrix},
     },
+    maybe_rayon::prelude::*,
 };
-use rayon::iter::{IndexedParallelIterator, IntoParallelRefIterator, ParallelIterator};
 use serde::{Deserialize, Serialize, de::DeserializeOwned};
 use transcript::Transcript;
 

crates/whir/src/ntt/mod.rs

@@ -13,9 +13,8 @@ use p3::{
     dft::{Radix2DitParallel, TwoAdicSubgroupDft},
     field::TwoAdicField,
     matrix::Matrix,
+    maybe_rayon::prelude::*,
 };
-#[cfg(feature = "parallel")]
-use rayon::prelude::*;
 use tracing::instrument;
 use witness::{InstancePaddingStrategy, RowMajorMatrix};
 

crates/whir/src/ntt/ntt_impl.rs

@@ -14,16 +14,15 @@ use p3::{
         Matrix,
         dense::{DenseMatrix, RowMajorMatrix},
     },
+    maybe_rayon::prelude::*,
 };
 use std::{
     any::{Any, TypeId},
+    cmp::max,
     collections::HashMap,
     sync::{Arc, LazyLock, Mutex, RwLock, RwLockReadGuard},
 };
 
-#[cfg(feature = "parallel")]
-use {rayon::prelude::*, std::cmp::max};
-
 /// Global cache for NTT engines, indexed by field.
 // TODO: Skip `LazyLock` when `HashMap::with_hasher` becomes const.
 // see https://github.com/rust-lang/rust/issues/102575

crates/whir/src/ntt/transpose.rs

@@ -3,10 +3,10 @@ use crate::ntt::matrix_skip::MatrixMutSkip;
 use super::{super::utils::is_power_of_two, MatrixMut, utils::workload_size};
 use std::mem::swap;
 
-use p3::matrix::{Matrix, dense::RowMajorMatrix};
-use rayon::iter::{IndexedParallelIterator, IntoParallelRefMutIterator, ParallelIterator};
-#[cfg(feature = "parallel")]
-use rayon::join;
+use p3::{
+    matrix::{Matrix, dense::RowMajorMatrix},
+    maybe_rayon::prelude::*,
+};
 use sumcheck::macros::{entered_span, exit_span};
 
 // NOTE: The assumption that rows and cols are a power of two are actually only relevant for the square matrix case.

crates/whir/src/ntt/wavelet.rs

@@ -2,8 +2,7 @@ use super::{transpose, utils::workload_size};
 use p3::field::Field;
 use std::cmp::max;
 
-#[cfg(feature = "parallel")]
-use rayon::prelude::*;
+use p3::maybe_rayon::prelude::*;
 
 /// Fast Wavelet Transform.
 ///

crates/whir/src/sumcheck/prover_batched.rs

@@ -1,9 +1,7 @@
 use super::proof::SumcheckPolynomial;
 use crate::sumcheck::prover_single::SumcheckSingle;
 use ff_ext::ExtensionField;
-use p3::{field::Field, util::log2_strict_usize};
-#[cfg(feature = "parallel")]
-use rayon::{join, prelude::*};
+use p3::{field::Field, maybe_rayon::prelude::*, util::log2_strict_usize};
 
 pub struct SumcheckBatched<F: ExtensionField> {
     // The evaluation on each p and eq

crates/whir/src/sumcheck/prover_single.rs

@@ -1,9 +1,7 @@
 use super::proof::SumcheckPolynomial;
 
 use ff_ext::ExtensionField;
-use p3::util::log2_strict_usize;
-#[cfg(feature = "parallel")]
-use rayon::{join, prelude::*};
+use p3::{maybe_rayon::prelude::*, util::log2_strict_usize};
 
 pub struct SumcheckSingle<E: ExtensionField> {
     // The evaluation of p

crates/whir/src/utils.rs

@@ -4,10 +4,7 @@ use multilinear_extensions::mle::FieldType;
 use p3::{
     field::Field,
     matrix::{Matrix, dense::RowMajorMatrix},
-};
-use rayon::{
-    iter::ParallelIterator,
-    slice::{ParallelSlice, ParallelSliceMut},
+    maybe_rayon::prelude::*,
 };
 use std::collections::BTreeSet;
 use sumcheck::macros::{entered_span, exit_span};

crates/whir/src/whir/batch/committer.rs

@@ -13,14 +13,12 @@ use derive_more::Debug;
 use ff_ext::ExtensionField;
 use p3::{
     matrix::{Matrix, dense::RowMajorMatrix},
+    maybe_rayon::prelude::*,
     util::log2_strict_usize,
 };
 use sumcheck::macros::{entered_span, exit_span};
 use transcript::{BasicTranscript, Transcript};
 
-#[cfg(feature = "parallel")]
-use rayon::prelude::*;
-
 #[derive(Debug)]
 pub struct Witnesses<E: ExtensionField> {
     pub(crate) polys: Vec<Vec<E::BaseField>>,

crates/whir/src/whir/batch/prover.rs

@@ -21,13 +21,11 @@ use crate::{
 use ff_ext::{ExtensionField, PoseidonField};
 use itertools::zip_eq;
 use multilinear_extensions::mle::{FieldType, MultilinearExtension};
-use p3::{commit::Mmcs, matrix::dense::RowMajorMatrix};
+use p3::{commit::Mmcs, matrix::dense::RowMajorMatrix, maybe_rayon::prelude::*};
 use sumcheck::macros::{entered_span, exit_span};
 use transcript::Transcript;
 
 use crate::whir::fs_utils::get_challenge_stir_queries;
-#[cfg(feature = "parallel")]
-use rayon::prelude::*;
 
 struct RoundStateBatch<'a, E: ExtensionField> {
     round_state: RoundState<'a, E>,

crates/whir/src/whir/batch/verifier.rs

@@ -17,8 +17,7 @@ use crate::{
 use ff_ext::{ExtensionField, PoseidonField};
 use itertools::zip_eq;
 use multilinear_extensions::{mle::MultilinearExtension, virtual_poly::eq_eval};
-use p3::{commit::Mmcs, util::log2_strict_usize};
-use rayon::iter::{IntoParallelRefIterator, ParallelIterator};
+use p3::{commit::Mmcs, maybe_rayon::prelude::*, util::log2_strict_usize};
 use sumcheck::macros::{entered_span, exit_span};
 use transcript::Transcript;
 

crates/whir/src/whir/fold.rs

@@ -6,9 +6,8 @@ use crate::{
 use p3::{
     field::{Field, TwoAdicField},
     matrix::Matrix,
+    maybe_rayon::prelude::*,
 };
-#[cfg(feature = "parallel")]
-use rayon::prelude::*;
 
 /// Given the evaluation of f on the coset specified by coset_offset * <coset_gen>
 /// Compute the fold on that point

crates/whir/src/whir/prover.rs

@@ -14,8 +14,7 @@ use crate::{
 };
 use ff_ext::ExtensionField;
 use multilinear_extensions::mle::{FieldType, MultilinearExtension};
-use p3::matrix::dense::RowMajorMatrix;
-use rayon::iter::{IntoParallelRefIterator, ParallelIterator};
+use p3::{matrix::dense::RowMajorMatrix, maybe_rayon::prelude::*};
 use sumcheck::macros::{entered_span, exit_span};
 use transcript::Transcript;
 

crates/whir/src/whir/verifier.rs

@@ -7,8 +7,8 @@ use multilinear_extensions::{mle::MultilinearExtension, virtual_poly::eq_eval};
 use p3::{
     commit::Mmcs,
     field::{Field, FieldAlgebra},
+    maybe_rayon::prelude::*,
 };
-use rayon::iter::{IndexedParallelIterator, IntoParallelRefIterator, ParallelIterator};
 use serde::{Deserialize, Serialize, de::DeserializeOwned};
 use std::iter;
 use sumcheck::macros::{entered_span, exit_span};

crates/witness/src/lib.rs

@@ -2,12 +2,9 @@ use multilinear_extensions::mle::{IntoMLE, MultilinearExtension};
 use p3::{
     field::{Field, FieldAlgebra},
     matrix::Matrix,
+    maybe_rayon::prelude::*,
 };
 use rand::{Rng, distributions::Standard, prelude::Distribution};
-use rayon::{
-    iter::{IndexedParallelIterator, IntoParallelIterator, ParallelExtend, ParallelIterator},
-    slice::ParallelSliceMut,
-};
 use std::{
     ops::{Deref, DerefMut, Index},
     slice::{Chunks, ChunksMut},