Optimize SRS serialization: parallel uncompressed format, file caching, buffered I/O

gkr_engine/src/poly_commit/definition.rs

@@ -1,10 +1,14 @@
 use polynomials::MultilinearExtension;
 use rand::RngCore;
 use serdes::ExpSerde;
+use std::io::{BufReader, BufWriter};
 use std::{fmt::Debug, str::FromStr};
 
 use crate::{ExpErrors, ExpanderSingleVarChallenge, FieldEngine, MPIEngine, Transcript};
 
+/// Buffer capacity for SRS file I/O (64 MB).
+const SRS_IO_BUFFER_CAPACITY: usize = 64 * 1024 * 1024;
+
 pub trait StructuredReferenceString {
     type PKey: Clone + Debug + ExpSerde + Send + Sync + 'static;
     type VKey: Clone + Debug + ExpSerde + Send + Sync + 'static;
@@ -63,18 +67,20 @@ pub trait ExpanderPCS<F: FieldEngine> {
         match path {
             Some(path) => {
                 match std::fs::File::open(path) {
-                    Ok(mut file) => {
+                    Ok(file) => {
                         // file exists; deserialize SRS from file
-                        Self::SRS::deserialize_from(&mut file).unwrap_or_else(|_| {
+                        let mut reader = BufReader::with_capacity(SRS_IO_BUFFER_CAPACITY, file);
+                        Self::SRS::deserialize_from(&mut reader).unwrap_or_else(|_| {
                             panic!("Failed to deserialize SRS for {} PCS", Self::NAME)
                         })
                     }
                     Err(_e) => {
                         // file does not exist; generate SRS and store to file
                         let srs = Self::gen_srs(params, mpi_engine, rng);
-                        let mut file =
+                        let file =
                             std::fs::File::create(path).expect("Failed to create SRS file");
-                        srs.serialize_into(&mut file)
+                        let mut writer = BufWriter::with_capacity(SRS_IO_BUFFER_CAPACITY, file);
+                        srs.serialize_into(&mut writer)
                             .expect("Failed to serialize SRS to file");
                         srs
                     }

poly_commit/src/kzg/uni_kzg/structs_kzg.rs

@@ -1,7 +1,11 @@
 use derivative::Derivative;
 use gkr_engine::StructuredReferenceString;
+use halo2curves::group::prime::PrimeCurveAffine;
+use halo2curves::group::UncompressedEncoding;
 use halo2curves::{pairing::Engine, CurveAffine};
+use rayon::prelude::*;
 use serdes::{ExpSerde, SerdeResult};
+use std::io::{Read, Write};
 
 #[derive(Clone, Copy, Debug, PartialEq, Eq, Derivative)]
 #[derivative(Default(bound = ""))]
@@ -34,7 +38,7 @@ where
 
 /// Structured reference string for univariate KZG polynomial commitment scheme.
 /// The univariate polynomial here is of coefficient form.
-#[derive(Clone, Debug, PartialEq, Eq, Derivative, ExpSerde)]
+#[derive(Clone, Debug, PartialEq, Eq, Derivative)]
 #[derivative(Default(bound = ""))]
 pub struct CoefFormUniKZGSRS<E: Engine>
 where
@@ -48,6 +52,79 @@ where
     pub tau_g2: E::G2Affine,
 }
 
+/// Custom ExpSerde implementation with parallel deserialization for fast SRS loading.
+/// Uses UNCOMPRESSED format to avoid expensive square root computation during load.
+impl<E: Engine> ExpSerde for CoefFormUniKZGSRS<E>
+where
+    E::G1Affine: ExpSerde
+        + CurveAffine<ScalarExt = E::Fr, CurveExt = E::G1>
+        + UncompressedEncoding
+        + Send
+        + Sync,
+    E::G2Affine: CurveAffine<ScalarExt = E::Fr, CurveExt = E::G2> + ExpSerde,
+{
+    fn serialize_into<W: Write>(&self, mut writer: W) -> SerdeResult<()> {
+        // Serialize length
+        self.powers_of_tau.len().serialize_into(&mut writer)?;
+
+        // Get uncompressed point size
+        let point_size =
+            std::mem::size_of::<<E::G1Affine as UncompressedEncoding>::Uncompressed>();
+        let total_size = self.powers_of_tau.len() * point_size;
+
+        // Pre-allocate buffer and convert all points to uncompressed format in parallel
+        let mut buffer = vec![0u8; total_size];
+        buffer
+            .par_chunks_mut(point_size)
+            .zip(self.powers_of_tau.par_iter())
+            .for_each(|(chunk, point)| {
+                let uncompressed = point.to_uncompressed();
+                chunk.copy_from_slice(uncompressed.as_ref());
+            });
+
+        // Write all at once
+        writer.write_all(&buffer)?;
+
+        // Serialize tau_g2
+        self.tau_g2.serialize_into(&mut writer)?;
+        Ok(())
+    }
+
+    fn deserialize_from<R: Read>(mut reader: R) -> SerdeResult<Self> {
+        // Read length
+        let len = usize::deserialize_from(&mut reader)?;
+
+        // Uncompressed G1 point size (e.g. BN256: 64 bytes)
+        let point_size =
+            std::mem::size_of::<<E::G1Affine as UncompressedEncoding>::Uncompressed>();
+        let total_bytes = len * point_size;
+
+        let mut buffer = vec![0u8; total_bytes];
+        reader.read_exact(&mut buffer)?;
+
+        // Parse points in parallel - using from_uncompressed_unchecked (no square root needed)
+        let powers_of_tau: Vec<E::G1Affine> = buffer
+            .par_chunks(point_size)
+            .map(|chunk| {
+                let mut uncompressed =
+                    <E::G1Affine as UncompressedEncoding>::Uncompressed::default();
+                uncompressed.as_mut().copy_from_slice(chunk);
+                E::G1Affine::from_uncompressed_unchecked(&uncompressed)
+                    .into_option()
+                    .ok_or(serdes::SerdeError::DeserializeError)
+            })
+            .collect::<SerdeResult<Vec<_>>>()?;
+
+        // Read tau_g2
+        let tau_g2 = E::G2Affine::deserialize_from(&mut reader)?;
+
+        Ok(Self {
+            powers_of_tau,
+            tau_g2,
+        })
+    }
+}
+
 impl<E: Engine> StructuredReferenceString for CoefFormUniKZGSRS<E>
 where
     <E as Engine>::G1Affine: ExpSerde + CurveAffine<ScalarExt = E::Fr, CurveExt = E::G1>,

poly_commit/src/utils.rs

@@ -5,6 +5,7 @@ use gkr_engine::{
 };
 use polynomials::{MultiLinearPoly, MultilinearExtension, MutableMultilinearExtension};
 
+/// Initialize PCS for testing without SRS caching (always regenerates SRS)
 #[allow(clippy::type_complexity)]
 pub fn expander_pcs_init_testing_only<FieldConfig: FieldEngine, PCS: ExpanderPCS<FieldConfig>>(
     n_input_vars: usize,
@@ -14,19 +15,44 @@ pub fn expander_pcs_init_testing_only<FieldConfig: FieldEngine, PCS: ExpanderPCS
     <PCS::SRS as StructuredReferenceString>::PKey,
     <PCS::SRS as StructuredReferenceString>::VKey,
     PCS::ScratchPad,
+) {
+    let srs_path = format!("/tmp/kzg_srs_{}.bin", n_input_vars);
+    expander_pcs_init_with_srs_path::<FieldConfig, PCS>(n_input_vars, mpi_config, Some(&srs_path))
+}
+
+/// Initialize PCS with optional SRS file caching
+///
+/// # Arguments
+/// * `n_input_vars` - Number of input variables (determines SRS size as 2^n)
+/// * `mpi_config` - MPI configuration
+/// * `srs_path` - Optional path to SRS cache file:
+///   - `Some(path)`: If file exists, load SRS from it; otherwise generate and save to it
+///   - `None`: Always regenerate SRS (no caching)
+#[allow(clippy::type_complexity)]
+pub fn expander_pcs_init_with_srs_path<FieldConfig: FieldEngine, PCS: ExpanderPCS<FieldConfig>>(
+    n_input_vars: usize,
+    mpi_config: &impl MPIEngine,
+    srs_path: Option<&str>,
+) -> (
+    PCS::Params,
+    <PCS::SRS as StructuredReferenceString>::PKey,
+    <PCS::SRS as StructuredReferenceString>::VKey,
+    PCS::ScratchPad,
 ) {
     let mut rng = test_rng();
 
     let pcs_params =
         <PCS as ExpanderPCS<FieldConfig>>::gen_params(n_input_vars, mpi_config.world_size());
+
     let pcs_setup = <PCS as ExpanderPCS<FieldConfig>>::gen_or_load_srs_for_testing(
         &pcs_params,
         mpi_config,
         &mut rng,
-        None,
+        srs_path,
     );
 
     let (pcs_proving_key, pcs_verification_key) = pcs_setup.into_keys();
+
     let pcs_scratch = <PCS as ExpanderPCS<FieldConfig>>::init_scratch_pad(&pcs_params, mpi_config);
 
     (

serdes/src/serdes.rs

@@ -68,8 +68,8 @@ impl<V: ExpSerde> ExpSerde for Vec<V> {
     }
 
     fn deserialize_from<R: Read>(mut reader: R) -> SerdeResult<Self> {
-        let mut v = Self::default();
         let len = usize::deserialize_from(&mut reader)?;
+        let mut v = Vec::with_capacity(len);
         for _ in 0..len {
             v.push(V::deserialize_from(&mut reader)?);
         }