merge latest changes from main

Author: TomWambsgans

Cargo.lock

@@ -20,6 +20,7 @@ members = [
     "crates/backend/sumcheck",
     "crates/backend/system-info",
     "crates/backend/zk-alloc",
+    "crates/backend/parallel",
 ]
 
 [workspace.lints]
@@ -63,12 +64,12 @@ leansig_wrapper = { path = "crates/leansig_wrapper" }
 backend = { path = "crates/backend" }
 zk-alloc = { path = "crates/backend/zk-alloc" }
 system-info = { path = "crates/backend/system-info" }
+parallel = { path = "crates/backend/parallel" }
 
 # 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"
@@ -86,12 +87,12 @@ include_dir = "0.7"
 [features]
 prox-gaps-conjecture = ["rec_aggregation/prox-gaps-conjecture"]
 test-config = ["rec_aggregation/test-config"]
-standard-alloc = ["rec_aggregation/standard-alloc"]
 
 [dependencies]
 clap.workspace = true
 rec_aggregation.workspace = true
 zk-alloc.workspace = true
+parallel.workspace = true
 rand.workspace = true
 sub_protocols.workspace = true
 utils.workspace = true
@@ -107,3 +108,4 @@ serde.workspace = true
 
 [profile.release]
 lto = "thin"
+codegen-units = 1

Cargo.toml

@@ -9,9 +9,10 @@ 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" }
 koala-bear = { path = "koala-bear", package = "mt-koala-bear" }
 utils = { path = "utils", package = "mt-utils" }
+zk-alloc.workspace = true

crates/backend/Cargo.toml

@@ -47,14 +47,14 @@ where
         self.shift
     }
 
-    #[inline]
+    #[inline(always)]
     fn assert_zero(&mut self, x: IF) {
         let alpha_power = self.extra_data.alpha_powers()[self.constraint_index];
         self.accumulator += alpha_power * x;
         self.constraint_index += 1;
     }
 
-    #[inline]
+    #[inline(always)]
     fn assert_zero_ef(&mut self, x: EF) {
         let alpha_power = self.extra_data.alpha_powers()[self.constraint_index];
         self.accumulator += alpha_power * x;

crates/backend/air/src/constraint_folder/normal.rs

@@ -57,21 +57,21 @@ where
         self.shift
     }
 
-    #[inline]
+    #[inline(always)]
     fn assert_zero(&mut self, x: IF) {
         let alpha_power = self.extra_data.alpha_powers()[self.constraint_index];
         self.accumulator += EFPacking::<EF>::from(alpha_power) * x;
         self.constraint_index += 1;
     }
 
-    #[inline]
+    #[inline(always)]
     fn assert_zero_ef(&mut self, x: EFPacking<EF>) {
         let alpha_power = self.extra_data.alpha_powers()[self.constraint_index];
         self.accumulator += EFPacking::<EF>::from(alpha_power) * x;
         self.constraint_index += 1;
     }
 
-    #[inline]
+    #[inline(always)]
     fn assert_eq_low(&mut self, x: IF, y: IF) {
         let alpha_power = self.extra_data.alpha_powers()[self.constraint_index];
         let contrib = EFPacking::<EF>::from(alpha_power) * (x - y);
@@ -80,7 +80,7 @@ where
         self.constraint_index += 1;
     }
 
-    #[inline]
+    #[inline(always)]
     fn low_degree_block<F>(&mut self, state: &mut [IF], block: F)
     where
         F: FnOnce(&mut Self, &mut [IF]),

crates/backend/air/src/constraint_folder/packed.rs

@@ -56,20 +56,24 @@ pub trait AirBuilder: Sized {
     fn assert_zero(&mut self, x: Self::IF);
     fn assert_zero_ef(&mut self, x: Self::EF);
 
+    #[inline(always)]
     fn assert_eq(&mut self, x: Self::IF, y: Self::IF) {
         self.assert_zero(x - y);
     }
 
+    #[inline(always)]
     fn assert_bool(&mut self, x: Self::IF) {
         self.assert_zero(x.bool_check());
     }
 
+    #[inline(always)]
     fn assert_eq_low(&mut self, x: Self::IF, y: Self::IF) {
         self.assert_eq(x, y);
     }
 
     /// Execute `block` as a low-degree sub-region whose post-state is "cacheable"
     /// = linear in z without the low-degree constraints
+    #[inline(always)]
     fn low_degree_block<F>(&mut self, state: &mut [Self::IF], block: F)
     where
         F: FnOnce(&mut Self, &mut [Self::IF]),

crates/backend/air/src/lib.rs

@@ -92,9 +92,15 @@ fn alloc_node<F: Field>(node: SymbolicNode<F>) -> u32 {
     })
 }
 
-pub fn get_node<F: Field>(idx: u32) -> SymbolicNode<F> {
+/// # Safety
+/// `idx` must be an offset returned by `alloc_node::<F>` for the current (same `F`, uncleared) arena.
+pub unsafe fn get_node<F: Field>(idx: u32) -> SymbolicNode<F> {
     ARENA.with(|arena| {
         let bytes = arena.borrow();
+        assert!(
+            idx as usize + std::mem::size_of::<SymbolicNode<F>>() <= bytes.len(),
+            "arena index out of bounds"
+        );
         unsafe { std::ptr::read_unaligned(bytes.as_ptr().add(idx as usize) as *const SymbolicNode<F>) }
     })
 }

crates/backend/air/src/symbolic.rs

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

crates/backend/fiat-shamir/Cargo.toml

@@ -9,8 +9,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};
 
@@ -132,9 +131,16 @@ 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| {
+        // Work-stealing parallel search: each worker pulls batches from a shared counter and
+        // stops once any worker has found a witness (`found`).
+        let next_batch = AtomicU64::new(0);
+        let found = AtomicBool::new(false);
+        parallel::for_each_index(parallel::num_threads(), |_| {
+            while !found.load(Ordering::Relaxed) {
+                let batch = next_batch.fetch_add(1, Ordering::Relaxed);
+                if batch >= num_batches {
+                    break;
+                }
                 let base = batch * lanes as u64;
 
                 let packed_witnesses = Packed::<EF>::from_fn(|lane| {
@@ -159,14 +165,14 @@ 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);
+                        break;
                     }
                 }
-                false
-            })
-            .expect("failed to find witness");
+            }
+        });
 
-        let witness = witness_found.lock().unwrap().unwrap();
+        let witness = witness_found.lock().unwrap().expect("failed to find witness");
 
         self.challenger.observe_many(&[witness]);
         assert!(self.challenger.state[CAPACITY].as_canonical_u64() & ((1 << bits) - 1) == 0);

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

@@ -10,6 +10,6 @@ itertools.workspace = true
 num-bigint = "*"
 paste = "*"
 rand.workspace = true
-rayon.workspace = true
+parallel.workspace = true
 serde.workspace = true
 tracing.workspace = true