Custom thread pool (#239)

removes rayon + increases perf

Author: TomWambsgans

Cargo.lock

@@ -19,6 +19,7 @@ members = [
     "crates/backend/fiat-shamir",
     "crates/backend/sumcheck",
     "crates/backend/system-info",
+    "crates/backend/parallel",
     "crates/backend/zk-alloc",
 ]
 
@@ -60,14 +61,14 @@ lean_compiler = { path = "crates/lean_compiler" }
 lean_prover = { path = "crates/lean_prover" }
 rec_aggregation = { path = "crates/rec_aggregation" }
 backend = { path = "crates/backend" }
-zk-alloc = { path = "crates/backend/zk-alloc" }
 system-info = { path = "crates/backend/system-info" }
+parallel = { path = "crates/backend/parallel" }
+zk-alloc = { path = "crates/backend/zk-alloc" }
 
 # External
 sha3 = "0.11.0"
 clap = { version = "4.5.59", features = ["derive"] }
 rand = "0.10.0"
-rayon = "1.11.0"
 pest = "2.7"
 pest_derive = "2.7"
 itertools = "0.14.0"
@@ -82,12 +83,14 @@ include_dir = "0.7"
 
 [features]
 prox-gaps-conjecture = ["rec_aggregation/prox-gaps-conjecture"]
+# Build with the plain system allocator instead of zk-alloc (for comparison/debugging).
 standard-alloc = ["rec_aggregation/standard-alloc"]
 
 [dependencies]
 clap.workspace = true
 rec_aggregation.workspace = true
 zk-alloc.workspace = true
+libc = "0.2"
 rand.workspace = true
 sub_protocols.workspace = true
 lean_vm.workspace = true

Cargo.toml

@@ -36,8 +36,8 @@ cargo run --release -- xmss --n-signatures 1550 --log-inv-rate 1
 
 | WHIR rate | Proven Regime         | Proximity Gaps Conjecture |
 | --------- | --------------------- | ------------------------- |
-| 1/2       | 1319 XMSS/s - 338 KiB | 1345 XMSS/s - 176 KiB     |
-| 1/4       | 961 XMSS/s - 228 KiB  | 969 XMSS/s - 126 KiB      |
+| 1/2       | 1453 XMSS/s - 344 KiB | 1500 XMSS/s - 178 KiB     |
+| 1/4       | 1058 XMSS/s - 229 KiB | 1065 XMSS/s - 127 KiB     |
 
 
 (Proving throughput - proof size)
@@ -54,15 +54,14 @@ cargo run --release -- recursion --n 2 --log-inv-rate 2
 
 | n   | WHIR rate | Proven Regime               | Proximity Gaps Conjecture   |
 | --- | --------- | --------------------------- | --------------------------- |
-| 1   | 1/2       | 0.39s = 1 x 0.39s - 278 KiB | 0.24s = 1 x 0.24s - 147 KiB |
-| 1   | 1/4       | 0.32s = 1 x 0.32s - 188 KiB | 0.27s = 1 x 0.27s - 100 KiB |
-| 2   | 1/2       | 0.7s = 2 x 0.35s - 293 KiB  | 0.43s = 2 x 0.21s - 157 KiB |
-| 2   | 1/4       | 0.56s = 2 x 0.28s - 194 KiB | 0.43s = 2 x 0.22s - 102 KiB |
-| 3   | 1/2       | 0.85s = 3 x 0.28s - 312 KiB | 0.63s = 3 x 0.21s - 150 KiB |
-| 3   | 1/4       | 0.94s = 3 x 0.31s - 203 KiB | 0.73s = 3 x 0.24s - 108 KiB |
-| 4   | 1/2       | 1.27s = 4 x 0.32s - 308 KiB | 0.78s = 4 x 0.2s - 166 KiB  |
-| 4   | 1/4       | 1.02s = 4 x 0.26s - 206 KiB | 0.79s = 4 x 0.2s - 108 KiB  |
-
+| 1   | 1/2       | 0.22s = 1 x 0.22s - 285 KiB | 0.15s = 1 x 0.15s - 143 KiB |
+| 1   | 1/4       | 0.24s = 1 x 0.24s - 189 KiB | 0.18s = 1 x 0.18s - 98 KiB  |
+| 2   | 1/2       | 0.52s = 2 x 0.26s - 282 KiB | 0.33s = 2 x 0.16s - 159 KiB |
+| 2   | 1/4       | 0.45s = 2 x 0.23s - 198 KiB | 0.33s = 2 x 0.16s - 102 KiB |
+| 3   | 1/2       | 0.7s = 3 x 0.23s - 317 KiB  | 0.47s = 3 x 0.16s - 151 KiB |
+| 3   | 1/4       | 0.67s = 3 x 0.22s - 191 KiB | 0.43s = 3 x 0.14s - 111 KiB |
+| 4   | 1/2       | 1.02s = 4 x 0.26s - 309 KiB | 0.61s = 4 x 0.15s - 168 KiB |
+| 4   | 1/4       | 0.85s = 4 x 0.21s - 208 KiB | 0.62s = 4 x 0.15s - 109 KiB |
 
 
 (time for n->1 recursive aggregation - proof size)

README.md

@@ -9,7 +9,7 @@ poly = { path = "poly", package = "mt-poly" }
 sumcheck = { path = "sumcheck", package = "mt-sumcheck" }
 field = { path = "field", package = "mt-field" }
 air = { path = "air", package = "mt-air" }
-rayon.workspace = true
+parallel.workspace = true
 whir = { path = "../whir", package = "mt-whir" }
 tracing.workspace = true
 fiat-shamir = { path = "fiat-shamir", package = "mt-fiat-shamir" }

crates/backend/Cargo.toml

@@ -10,4 +10,4 @@ symetric = { path = "../symetric", package = "mt-symetric" }
 utils = { path = "../utils", package = "utils" }
 tracing.workspace = true
 serde.workspace = true
-rayon.workspace = true
+parallel.workspace = true

crates/backend/fiat-shamir/Cargo.toml

@@ -6,8 +6,7 @@ use field::PrimeCharacteristicRing;
 use field::integers::QuotientMap;
 use field::{ExtensionField, PrimeField64};
 use koala_bear::symmetric::Permutation;
-use rayon::prelude::*;
-use std::sync::atomic::{AtomicU64, Ordering};
+use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
 use std::time::Duration;
 use std::{fmt::Debug, sync::Mutex, time::Instant};
 use symetric::CAPACITY;
@@ -132,9 +131,18 @@ where
         let witness_found = Mutex::<Option<PF<EF>>>::new(None);
         // each batch tests lanes witnesses simultaneously
         let num_batches = PF::<EF>::ORDER_U64.div_ceil(lanes as u64);
-        (0..num_batches)
-            .into_par_iter()
-            .find_any(|&batch| {
+
+        let next_batch = AtomicU64::new(0);
+        let found = AtomicBool::new(false);
+        parallel::for_each_index(parallel::num_threads(), |_| {
+            loop {
+                if found.load(Ordering::Relaxed) {
+                    return;
+                }
+                let batch = next_batch.fetch_add(1, Ordering::Relaxed);
+                if batch >= num_batches {
+                    return;
+                }
                 let base = batch * lanes as u64;
 
                 let packed_witnesses = Packed::<EF>::from_fn(|lane| {
@@ -159,12 +167,13 @@ where
                     let rand_usize = sample.as_canonical_u64() as usize;
                     if (rand_usize & ((1 << bits) - 1)) == 0 {
                         *witness_found.lock().unwrap() = Some(*witness);
-                        return true;
+                        found.store(true, Ordering::Relaxed);
+                        return;
                     }
                 }
-                false
-            })
-            .expect("failed to find witness");
+            }
+        });
+        assert!(found.load(Ordering::Relaxed), "failed to find witness");
 
         let witness = witness_found.lock().unwrap().unwrap();
 

crates/backend/fiat-shamir/src/prover.rs

@@ -9,7 +9,7 @@ utils = { path = "../utils", package = "utils" }
 itertools.workspace = true
 num-bigint = "*"
 paste = "*"
+parallel.workspace = true
 rand.workspace = true
-rayon.workspace = true
 serde.workspace = true
 tracing.workspace = true

crates/backend/field/Cargo.toml

@@ -9,7 +9,6 @@ use core::ops::{Add, AddAssign, Div, DivAssign, Mul, MulAssign, Neg, Sub, SubAss
 use core::{array, slice};
 
 use num_bigint::BigUint;
-use rayon::{current_num_threads, prelude::*};
 use serde::Serialize;
 use serde::de::DeserializeOwned;
 use utils::{flatten_to_base, iter_array_chunks_padded};
@@ -1020,24 +1019,21 @@ impl<F: Field> BoundedPowers<F> {
         let mut points_packed = F::Packing::zero_vec(num_packed);
 
         // Split computation evenly among threads
-        let num_threads = current_num_threads().max(1);
+        let num_threads = parallel::num_threads().max(1);
         let chunk_size = num_packed.div_ceil(num_threads);
 
         // Precompute base for each chunk.
         let base = self.iter.base;
         let chunk_base = base.exp_u64((chunk_size * width) as u64);
         let shift = self.iter.current;
 
-        points_packed
-            .par_chunks_mut(chunk_size)
-            .enumerate()
-            .for_each(|(chunk_idx, chunk_slice)| {
-                // First power in this chunk
-                let chunk_start = shift * chunk_base.exp_u64(chunk_idx as u64);
+        parallel::par_chunks_mut(&mut points_packed, chunk_size, |chunk_idx, chunk_slice| {
+            // First power in this chunk
+            let chunk_start = shift * chunk_base.exp_u64(chunk_idx as u64);
 
-                // Fill the chunk with packed powers.
-                F::Packing::packed_shifted_powers(base, chunk_start).fill(chunk_slice);
-            });
+            // Fill the chunk with packed powers.
+            F::Packing::packed_shifted_powers(base, chunk_start).fill(chunk_slice);
+        });
 
         // return the number of requested points, discarding the unused packed powers
         // SAFETY: size_of::<F::Packing> always divides size_of::<F::Packing>.

crates/backend/field/src/field.rs

@@ -8,7 +8,6 @@ field = { path = "../field", package = "mt-field" }
 utils = { path = "../utils", package = "utils" }
 
 rand.workspace = true
-rayon.workspace = true
 serde.workspace = true
 itertools.workspace = true
 tracing.workspace = true

crates/backend/koala-bear/Cargo.toml

@@ -0,0 +1,11 @@
+[package]
+name = "parallel"
+version.workspace = true
+edition.workspace = true
+description = "Minimal fixed-size thread pool for static data-parallel kernels"
+
+[dependencies]
+system-info.workspace = true
+
+[lints]
+workspace = true