feat: parallelizing VID polynomial evaluation

primitives/benches/advz.rs

@@ -24,7 +24,7 @@ fn advz<E, H>(c: &mut Criterion, pairing_name: &str)
 where
     E: Pairing,
     // TODO(Gus) clean up nasty trait bounds upstream
-    H: Digest + DynDigest + Default + Clone + Write,
+    H: Digest + DynDigest + Default + Clone + Write + Send + Sync,
     <<H as OutputSizeUser>::OutputSize as ArrayLength<u8>>::ArrayType: Copy,
 {
     // play with these items

primitives/src/merkle_tree/hasher.rs

@@ -98,7 +98,7 @@ pub type GenericHasherMerkleTree<H, E, I, Arity> =
 /// # use jf_primitives::merkle_tree::hasher::HasherDigest;
 /// fn generic_over_hasher<H>()
 /// where
-///     H: Digest + Write,
+///     H: Digest + Write + Send + Sync,
 ///     <<H as OutputSizeUser>::OutputSize as ArrayLength<u8>>::ArrayType: Copy,
 /// {
 ///     let my_data = [1, 2, 3, 4, 5, 6, 7, 8, 9];
@@ -114,21 +114,21 @@ pub type GenericHasherMerkleTree<H, E, I, Arity> =
 /// # use jf_primitives::merkle_tree::hasher::HasherDigest;
 /// fn generic_over_hasher<H>()
 /// where
-///     H: Digest + Write,
+///     H: Digest + Write + Send + Sync,
 /// {
 ///     let my_data = [1, 2, 3, 4, 5, 6, 7, 8, 9];
 ///     let mt = HasherMerkleTree::<H, usize>::from_elems(None, &my_data).unwrap();
 /// }
 /// ```
-pub trait HasherDigest: Digest<OutputSize = Self::Foo> + Write {
+pub trait HasherDigest: Digest<OutputSize = Self::Foo> + Write + Send + Sync {
     /// Associated type needed to express trait bounds.
     type Foo: ArrayLength<u8, ArrayType = Self::Bar>;
     /// Associated type needed to express trait bounds.
     type Bar: Copy;
 }
 impl<T> HasherDigest for T
 where
-    T: Digest + Write,
+    T: Digest + Write + Send + Sync,
     <T::OutputSize as ArrayLength<u8>>::ArrayType: Copy,
 {
     type Foo = T::OutputSize;

primitives/src/vid/advz.rs

@@ -42,8 +42,12 @@ use bytes_to_field::{bytes_to_field, field_to_bytes};
 use core::mem;
 use derivative::Derivative;
 use digest::crypto_common::Output;
-use itertools::Itertools;
-use jf_utils::canonical;
+use jf_utils::{
+    canonical,
+    par_utils::{parallelizable_chunks, parallelizable_slice_iter},
+};
+#[cfg(feature = "parallel")]
+use rayon::prelude::ParallelIterator;
 use serde::{Deserialize, Serialize};
 
 mod bytes_to_field;
@@ -555,32 +559,34 @@ where
         H: HasherDigest,
     {
         let code_word_size = self.num_storage_nodes * self.multiplicity;
-        let all_storage_node_evals = {
-            let mut all_storage_node_evals = vec![Vec::with_capacity(polys.len()); code_word_size];
+        let mut all_storage_node_evals = vec![Vec::with_capacity(polys.len()); code_word_size];
 
-            for poly in polys.iter() {
-                let poly_evals = UnivariateKzgPCS::<E>::multi_open_rou_evals(
+        let all_poly_evals = parallelizable_slice_iter(polys)
+            .map(|poly| {
+                UnivariateKzgPCS::<E>::multi_open_rou_evals(
                     poly,
                     code_word_size,
                     &self.multi_open_domain,
                 )
-                .map_err(vid)?;
+                .map_err(vid)
+            })
+            .collect::<Result<Vec<_>, VidError>>()?;
 
-                for (storage_node_evals, poly_eval) in
-                    all_storage_node_evals.iter_mut().zip(poly_evals)
-                {
-                    storage_node_evals.push(poly_eval);
-                }
+        for poly_evals in all_poly_evals {
+            for (storage_node_evals, poly_eval) in all_storage_node_evals
+                .iter_mut()
+                .zip(poly_evals.into_iter())
+            {
+                storage_node_evals.push(poly_eval);
             }
+        }
 
-            // sanity checks
-            assert_eq!(all_storage_node_evals.len(), code_word_size);
-            for storage_node_evals in all_storage_node_evals.iter() {
-                assert_eq!(storage_node_evals.len(), polys.len());
-            }
+        // sanity checks
+        assert_eq!(all_storage_node_evals.len(), code_word_size);
+        for storage_node_evals in all_storage_node_evals.iter() {
+            assert_eq!(storage_node_evals.len(), polys.len());
+        }
 
-            all_storage_node_evals
-        };
         Ok(all_storage_node_evals)
     }
 
@@ -616,10 +622,9 @@ where
         E: Pairing,
     {
         let chunk_size = self.payload_chunk_size * self.multiplicity;
-        bytes_to_field::<_, KzgEval<E>>(payload)
-            .chunks(chunk_size)
-            .into_iter()
-            .map(|evals_iter| self.polynomial(evals_iter))
+        let elem_bytes_len = bytes_to_field::elem_byte_capacity::<<E as Pairing>::ScalarField>();
+        parallelizable_chunks(payload, chunk_size * elem_bytes_len)
+            .map(|chunk| self.polynomial(bytes_to_field::<_, KzgEval<E>>(chunk)))
             .collect()
     }
 
@@ -803,7 +808,7 @@ mod tests {
         let srs = init_srs(payload_chunk_size, &mut rng);
         let advz =
             Advz::<Bls12_381, Sha256>::new(payload_chunk_size, num_storage_nodes, 1, srs).unwrap();
-        let payload_random = init_random_payload(1 << 20, &mut rng);
+        let payload_random = init_random_payload(1 << 25, &mut rng);
 
         let _ = advz.disperse(payload_random);
     }

utilities/src/par_utils.rs

@@ -27,3 +27,14 @@ pub fn parallelizable_slice_iter<T: Sync>(data: &[T]) -> rayon::slice::Iter<T> {
 pub fn parallelizable_slice_iter<T>(data: &[T]) -> ark_std::slice::Iter<T> {
     data.iter()
 }
+
+#[cfg(feature = "parallel")]
+pub fn parallelizable_chunks<T: Sync>(data: &[T], chunk_size: usize) -> rayon::slice::Chunks<T> {
+    use rayon::slice::ParallelSlice;
+    data.par_chunks(chunk_size)
+}
+
+#[cfg(not(feature = "parallel"))]
+pub fn parallelizable_chunks<T>(data: &[T], chunk_size: usize) -> ark_std::slice::Chunks<T> {
+    data.chunks(chunk_size)
+}