Merge pull request #728 from nagatoism/disable-rayon

  Add feature "multicore"  and you can disable rayon by disabling the  "multicore" feature.

README.md

@@ -14,6 +14,10 @@ minor version bump.
 `halo2` currently uses [rayon](https://github.com/rayon-rs/rayon) for parallel computation.
 The `RAYON_NUM_THREADS` environment variable can be used to set the number of threads.
 
+You can disable `rayon` by disabling the `"multicore"` feature.
+Warning! Halo2 will lose access to parallelism if you disable the `"multicore"` feature.
+This will significantly degrade performance.
+
 ## License
 
 Licensed under either of

halo2/Cargo.toml

@@ -19,7 +19,7 @@ all-features = true
 rustdoc-args = ["--cfg", "docsrs", "--html-in-header", "katex-header.html"]
 
 [dependencies]
-halo2_proofs = { version = "0.2", path = "../halo2_proofs" }
+halo2_proofs = { version = "0.2", path = "../halo2_proofs", default-features = false }
 
 [lib]
 bench = false

halo2_gadgets/Cargo.toml

@@ -26,7 +26,7 @@ arrayvec = "0.7.0"
 bitvec = "1"
 ff = "0.13"
 group = "0.13"
-halo2_proofs = { version = "0.2", path = "../halo2_proofs" }
+halo2_proofs = { version = "0.2", path = "../halo2_proofs", default-features = false }
 lazy_static = "1"
 pasta_curves = "0.5"
 proptest = { version = "1.0.0", optional = true }

halo2_proofs/Cargo.toml

@@ -45,13 +45,13 @@ harness = false
 
 [dependencies]
 backtrace = { version = "0.3", optional = true }
-rayon = "1.5.1"
 ff = "0.13"
 group = "0.13"
 pasta_curves = "0.5"
 rand_core = { version = "0.6", default-features = false }
 tracing = "0.1"
 blake2b_simd = "1"
+maybe-rayon = {version = "0.1.0", default-features = false}
 
 # Developer tooling dependencies
 plotters = { version = "0.3.0", default-features = false, optional = true }
@@ -68,7 +68,8 @@ rand_core = { version = "0.6", default-features = false, features = ["getrandom"
 getrandom = { version = "0.2", features = ["js"] }
 
 [features]
-default = ["batch"]
+default = ["batch", "multicore"]
+multicore = ["maybe-rayon/threads"]
 dev-graph = ["plotters", "tabbycat"]
 test-dev-graph = [
     "dev-graph",

halo2_proofs/README.md

@@ -15,6 +15,10 @@ minor version bump.
 computation. The `RAYON_NUM_THREADS` environment variable can be used to set the number of
 threads.
 
+You can disable `rayon` by disabling the `"multicore"` feature.
+Warning! Halo2 will lose access to parallelism if you disable the `"multicore"` feature.
+This will significantly degrade performance.
+
 ## License
 
 Licensed under either of

halo2_proofs/src/multicore.rs

@@ -1,9 +1,73 @@
 //! An interface for dealing with the kinds of parallel computations involved in
-//! `halo2`. It's currently just a (very!) thin wrapper around [`rayon`] but may
+//! `halo2`. It's currently just a (very!) thin wrapper around `rayon` but may
 //! be extended in the future to allow for various parallelism strategies.
 
-pub use rayon::{
-    current_num_threads,
-    iter::{IndexedParallelIterator, IntoParallelIterator, ParallelIterator},
+#[cfg(all(
+    feature = "multicore",
+    target_arch = "wasm32",
+    not(target_feature = "atomics")
+))]
+compile_error!(
+    "The multicore feature flag is not supported on wasm32 architectures without atomics"
+);
+
+pub use maybe_rayon::{
+    iter::{IntoParallelIterator, ParallelIterator},
     join, scope, Scope,
 };
+
+#[cfg(feature = "multicore")]
+pub use maybe_rayon::{current_num_threads, iter::IndexedParallelIterator};
+
+#[cfg(not(feature = "multicore"))]
+pub fn current_num_threads() -> usize {
+    1
+}
+
+#[cfg(not(feature = "multicore"))]
+pub trait IndexedParallelIterator: std::iter::Iterator {}
+
+pub trait TryFoldAndReduce<T, E> {
+    /// Implements `iter.try_fold().try_reduce()` for `rayon::iter::ParallelIterator`,
+    /// falling back on `Iterator::try_fold` when the `multicore` feature flag is
+    /// disabled.
+    /// The `try_fold_and_reduce` function can only be called by a iter with
+    /// `Result<T, E>` item type because the `fold_op` must meet the trait
+    /// bounds of both `try_fold` and `try_reduce` from rayon.   
+    fn try_fold_and_reduce(
+        self,
+        identity: impl Fn() -> T + Send + Sync,
+        fold_op: impl Fn(T, Result<T, E>) -> Result<T, E> + Send + Sync,
+    ) -> Result<T, E>;
+}
+
+#[cfg(feature = "multicore")]
+impl<T, E, I> TryFoldAndReduce<T, E> for I
+where
+    T: Send + Sync,
+    E: Send + Sync,
+    I: maybe_rayon::iter::ParallelIterator<Item = Result<T, E>>,
+{
+    fn try_fold_and_reduce(
+        self,
+        identity: impl Fn() -> T + Send + Sync,
+        fold_op: impl Fn(T, Result<T, E>) -> Result<T, E> + Send + Sync,
+    ) -> Result<T, E> {
+        self.try_fold(&identity, &fold_op)
+            .try_reduce(&identity, |a, b| fold_op(a, Ok(b)))
+    }
+}
+
+#[cfg(not(feature = "multicore"))]
+impl<T, E, I> TryFoldAndReduce<T, E> for I
+where
+    I: std::iter::Iterator<Item = Result<T, E>>,
+{
+    fn try_fold_and_reduce(
+        mut self,
+        identity: impl Fn() -> T + Send + Sync,
+        fold_op: impl Fn(T, Result<T, E>) -> Result<T, E> + Send + Sync,
+    ) -> Result<T, E> {
+        self.try_fold(identity(), fold_op)
+    }
+}

halo2_proofs/src/plonk/verifier/batch.rs

@@ -4,12 +4,15 @@ use rand_core::OsRng;
 
 use super::{verify_proof, VerificationStrategy};
 use crate::{
-    multicore::{IndexedParallelIterator, IntoParallelIterator, ParallelIterator},
+    multicore::{IntoParallelIterator, TryFoldAndReduce},
     plonk::{Error, VerifyingKey},
     poly::commitment::{Guard, Params, MSM},
     transcript::{Blake2bRead, EncodedChallenge},
 };
 
+#[cfg(feature = "multicore")]
+use crate::multicore::{IndexedParallelIterator, ParallelIterator};
+
 /// A proof verification strategy that returns the proof's MSM.
 ///
 /// `BatchVerifier` handles the accumulation of the MSMs for the batched proofs.
@@ -108,11 +111,10 @@ where
                     e
                 })
             })
-            .try_fold(
+            .try_fold_and_reduce(
                 || params.empty_msm(),
-                |msm, res| res.map(|proof_msm| accumulate_msm(msm, proof_msm)),
-            )
-            .try_reduce(|| params.empty_msm(), |a, b| Ok(accumulate_msm(a, b)));
+                |acc, res| res.map(|proof_msm| accumulate_msm(acc, proof_msm)),
+            );
 
         match final_msm {
             Ok(msm) => msm.eval(),