mpi with rayon

Cargo.toml

@@ -11,8 +11,9 @@ members = [
     "circuit",
     "config", # gkr_field_config + pcs_config + transcript_config
     "config/gkr_field_config", # definitions of all field types used in gkr and pcs
-    "config/mpi_config", # definitions of mpi communication toolkit
+    # "config/mpi_config", # definitions of mpi communication toolkit
     "config/config_macros", # proc macros used to declare a new config, this has to a separate crate due to rust compilation issues
+    "config/mpi_config", # definitions of mpi communication toolkit via rayon backend
     "gkr",
     "poly_commit",
     "sumcheck",
@@ -40,6 +41,7 @@ mpi = "0.8.0"
 rand = "0.8.5"
 rayon = "1.10"
 sha2 = "0.10.8"
+serial_test = "2.0"
 tiny-keccak = { version = "2.0.2", features = [ "sha3", "keccak" ] }
 tokio = { version = "1.38.0", features = ["full"] }
 tynm = { version = "0.1.6", default-features = false }

arith/polynomials/src/ref_mle.rs

@@ -9,6 +9,7 @@ pub trait MultilinearExtension<F: Field>: Index<usize, Output = F> {
 
     fn num_vars(&self) -> usize;
 
+    #[inline]
     fn hypercube_size(&self) -> usize {
         1 << self.num_vars()
     }
@@ -18,6 +19,11 @@ pub trait MultilinearExtension<F: Field>: Index<usize, Output = F> {
     fn hypercube_basis_ref(&self) -> &Vec<F>;
 
     fn interpolate_over_hypercube(&self) -> Vec<F>;
+
+    #[inline]
+    fn serialized_size(&self) -> usize {
+        self.hypercube_size() * F::SIZE
+    }
 }
 
 #[derive(Debug, Clone)]

config/mpi_config/Cargo.toml

@@ -4,8 +4,9 @@ version = "0.1.0"
 edition = "2021"
 
 [dependencies]
+rayon.workspace = true
+
 arith = { path = "../../arith" }
-mpi.workspace = true
 
 [dev-dependencies]
-mersenne31 = { path = "../../arith/mersenne31"}
+serial_test.workspace = true

config/mpi_config/src/atomic_vec.rs

@@ -0,0 +1,74 @@
+use std::sync::atomic::{AtomicUsize, Ordering};
+
+/// A lock-free append-only vector implementation
+/// credit: Claude
+#[derive(Debug)]
+pub struct AtomicVec<T> {
+    // The actual data storage
+    data: Vec<T>,
+    // Current length of valid data
+    len: AtomicUsize,
+}
+
+impl<T: Clone> AtomicVec<T> {
+    #[inline]
+    pub fn new(capacity: usize) -> Self {
+        let mut data = Vec::with_capacity(capacity);
+        // Pre-fill with default values to avoid reallocation
+        data.resize_with(capacity, || unsafe { std::mem::zeroed() });
+        Self {
+            data,
+            len: AtomicUsize::new(0),
+        }
+    }
+
+    /// Append data to the vector
+    /// Returns the start index where data was appended
+    pub fn append(&self, items: &[T]) -> Option<usize> {
+        let old_len = self.len.fetch_add(items.len(), Ordering::AcqRel);
+        if old_len + items.len() > self.data.capacity() {
+            // Restore the length if we would exceed capacity
+            self.len.fetch_sub(items.len(), Ordering::Release);
+            return None;
+        }
+
+        // Safe because:
+        // 1. We've pre-allocated the space
+        // 2. Each thread writes to its own section
+        // 3. The atomic len ensures no overlapping writes
+        unsafe {
+            let ptr = self.data.as_ptr().add(old_len) as *mut T;
+            for (i, item) in items.iter().enumerate() {
+                std::ptr::write(ptr.add(i), item.clone());
+            }
+        }
+
+        Some(old_len)
+    }
+
+    /// Read a slice of data
+    #[inline]
+    pub fn get_slice(&self, start: usize, end: usize) -> Option<&[T]> {
+        let current_len = self.len.load(Ordering::Acquire);
+        if start >= current_len || end > current_len || start > end {
+            return None;
+        }
+
+        // Safe because:
+        // 1. We've checked the bounds
+        // 2. No data is ever modified after being written
+        Some(unsafe { std::slice::from_raw_parts(self.data.as_ptr().add(start), end - start) })
+    }
+
+    /// Get current length
+    #[inline]
+    pub fn len(&self) -> usize {
+        self.len.load(Ordering::Acquire)
+    }
+
+    /// Check if empty
+    #[inline]
+    pub fn is_empty(&self) -> bool {
+        self.len() == 0
+    }
+}