Merge pull request #51 from 0xfuturistic/feat/parallelize_to_grand_pr…

…oducts feat(parallelize): add parallelism to to_grand_products
a16z · Sep 25, 2023 · 6911db4 · 6911db4
2 parents e52d160 + ddcce1c
commit 6911db4
Show file tree

Hide file tree

Showing 4 changed files with 47 additions and 21 deletions.
diff --git a/src/lasso/memory_checking.rs b/src/lasso/memory_checking.rs
@@ -18,21 +18,22 @@ use ark_ff::{Field, PrimeField};
 use ark_serialize::*;
 use ark_std::{One, Zero};
 use merlin::Transcript;
+use std::marker::Sync;
 
 #[derive(Debug, CanonicalSerialize, CanonicalDeserialize)]
 pub struct MemoryCheckingProof<
   G: CurveGroup,
   const C: usize,
   const M: usize,
-  S: SubtableStrategy<G::ScalarField, C, M>,
+  S: SubtableStrategy<G::ScalarField, C, M> + Sync,
 > where
   [(); S::NUM_MEMORIES]: Sized,
 {
   proof_prod_layer: ProductLayerProof<G::ScalarField, { S::NUM_MEMORIES }>,
   proof_hash_layer: HashLayerProof<G, C, M, S>,
 }
 
-impl<G: CurveGroup, const C: usize, const M: usize, S: SubtableStrategy<G::ScalarField, C, M>>
+impl<G: CurveGroup, const C: usize, const M: usize, S: SubtableStrategy<G::ScalarField, C, M> + Sync>
   MemoryCheckingProof<G, C, M, S>
 where
   [(); S::NUM_SUBTABLES]: Sized,
@@ -666,7 +667,7 @@ impl<F: PrimeField, const NUM_MEMORIES: usize> ProductLayerProof<F, NUM_MEMORIES
   /// - `transcript`: The proof transcript, used for Fiat-Shamir.
   #[tracing::instrument(skip_all, name = "ProductLayer.prove")]
   pub fn prove<G>(
-    grand_products: &mut [GrandProducts<F>; NUM_MEMORIES],
+    grand_products: &mut Vec<GrandProducts<F>>,
     transcript: &mut Transcript,
   ) -> (Self, Vec<F>, Vec<F>)
   where

diff --git a/src/lasso/surge.rs b/src/lasso/surge.rs
@@ -20,6 +20,7 @@ use ark_serialize::*;
 
 use ark_std::log2;
 use merlin::Transcript;
+use std::marker::Sync;
 
 pub struct SparsePolyCommitmentGens<G> {
   pub gens_combined_l_variate: PolyCommitmentGens<G>,
@@ -93,7 +94,7 @@ pub struct SparsePolynomialEvaluationProof<
   G: CurveGroup,
   const C: usize,
   const M: usize,
-  S: SubtableStrategy<G::ScalarField, C, M>,
+  S: SubtableStrategy<G::ScalarField, C, M> + Sync,
 > where
   [(); S::NUM_MEMORIES]: Sized,
 {
@@ -102,7 +103,7 @@ pub struct SparsePolynomialEvaluationProof<
   memory_check: MemoryCheckingProof<G, C, M, S>,
 }
 
-impl<G: CurveGroup, const C: usize, const M: usize, S: SubtableStrategy<G::ScalarField, C, M>>
+impl<G: CurveGroup, const C: usize, const M: usize, S: SubtableStrategy<G::ScalarField, C, M> + Sync>
   SparsePolynomialEvaluationProof<G, C, M, S>
 where
   [(); S::NUM_SUBTABLES]: Sized,

diff --git a/src/subprotocols/sumcheck.rs b/src/subprotocols/sumcheck.rs
@@ -174,7 +174,7 @@ impl<F: PrimeField> SumcheckInstanceProof<F> {
       let iterator = (0..mle_half).into_par_iter();
 
       #[cfg(not(feature = "multicore"))]
-      let iterator = (0..mle_half).iter();
+      let iterator = 0..mle_half;
 
       let accum: Vec<Vec<F>> = iterator
         .map(|poly_term_i| {

diff --git a/src/subtables/mod.rs b/src/subtables/mod.rs
@@ -1,4 +1,4 @@
-use std::marker::PhantomData;
+use std::marker::{PhantomData, Sync};
 
 use ark_ec::CurveGroup;
 use ark_ff::PrimeField;
@@ -107,7 +107,7 @@ where
 /// Stores the non-sparse evaluations of T[k] for each of the 'c'-dimensions as DensePolynomials, enables combination and commitment.
 impl<F: PrimeField, const C: usize, const M: usize, S> Subtables<F, C, M, S>
 where
-  S: SubtableStrategy<F, C, M>,
+  S: SubtableStrategy<F, C, M> + Sync,
   [(); S::NUM_SUBTABLES]: Sized,
   [(); S::NUM_MEMORIES]: Sized,
 {
@@ -135,19 +135,43 @@ where
     &self,
     dense: &DensifiedRepresentation<F, C>,
     r_mem_check: &(F, F),
-  ) -> [GrandProducts<F>; S::NUM_MEMORIES] {
-    std::array::from_fn(|i| {
-      let subtable = &self.subtable_entries[S::memory_to_subtable_index(i)];
-      let j = S::memory_to_dimension_index(i);
-      GrandProducts::new(
-        subtable,
-        &dense.dim[j],
-        &dense.dim_usize[j],
-        &dense.read[j],
-        &dense.r#final[j],
-        r_mem_check,
-      )
-    })
+  ) -> Vec<GrandProducts<F>> {
+    #[cfg(feature = "multicore")]
+    {
+      (0..S::NUM_MEMORIES)
+        .into_par_iter()
+        .map(|i| {
+          let subtable = &self.subtable_entries[S::memory_to_subtable_index(i)];
+          let j = S::memory_to_dimension_index(i);
+          GrandProducts::new(
+            subtable,
+            &dense.dim[j],
+            &dense.dim_usize[j],
+            &dense.read[j],
+            &dense.r#final[j],
+            r_mem_check,
+          )
+        })
+        .collect::<Vec<_>>()
+    }
+
+    #[cfg(not(feature = "multicore"))]
+    {
+      (0..S::NUM_MEMORIES)
+        .map(|i| {
+          let subtable = &self.subtable_entries[S::memory_to_subtable_index(i)];
+          let j = S::memory_to_dimension_index(i);
+          GrandProducts::new(
+            subtable,
+            &dense.dim[j],
+            &dense.dim_usize[j],
+            &dense.read[j],
+            &dense.r#final[j],
+            r_mem_check,
+          )
+        })
+        .collect::<Vec<_>>()
+    }
   }
 
   #[tracing::instrument(skip_all, name = "Subtables.commit")]