Feat/read pk buffer capacity (privacy-scaling-explorations#24)

* feat: change default `read_pk` buffer capacity to 1MB

* feat: add bench for read_pk

snark-verifier-sdk/Cargo.toml

@@ -69,3 +69,8 @@ harness = false
 name = "zkevm_plus_state"
 required-features = ["loader_halo2", "loader_evm", "zkevm", "halo2-pse"]
 harness = false
+
+[[bench]]
+name = "read_pk"
+required-features = ["loader_halo2"]
+harness = false

snark-verifier-sdk/benches/read_pk.rs

@@ -0,0 +1,216 @@
+use ark_std::{end_timer, start_timer};
+use criterion::Criterion;
+use criterion::{criterion_group, criterion_main};
+use halo2_base::gates::builder::BASE_CONFIG_PARAMS;
+use halo2_base::halo2_proofs;
+use halo2_base::utils::fs::gen_srs;
+use halo2_proofs::halo2curves as halo2_curves;
+use halo2_proofs::{halo2curves::bn256::Bn256, poly::kzg::commitment::ParamsKZG};
+use pprof::criterion::{Output, PProfProfiler};
+use rand::rngs::OsRng;
+
+use snark_verifier_sdk::halo2::aggregation::AggregationConfigParams;
+use snark_verifier_sdk::{
+    gen_pk,
+    halo2::{aggregation::AggregationCircuit, gen_snark_shplonk},
+    Snark,
+};
+use snark_verifier_sdk::{read_pk_with_capacity, SHPLONK};
+use std::path::Path;
+
+mod application {
+    use super::halo2_curves::bn256::Fr;
+    use super::halo2_proofs::{
+        circuit::{Layouter, SimpleFloorPlanner, Value},
+        plonk::{Advice, Circuit, Column, ConstraintSystem, Error, Fixed, Instance},
+        poly::Rotation,
+    };
+    use rand::RngCore;
+    use snark_verifier_sdk::CircuitExt;
+
+    #[derive(Clone, Copy)]
+    pub struct StandardPlonkConfig {
+        a: Column<Advice>,
+        b: Column<Advice>,
+        c: Column<Advice>,
+        q_a: Column<Fixed>,
+        q_b: Column<Fixed>,
+        q_c: Column<Fixed>,
+        q_ab: Column<Fixed>,
+        constant: Column<Fixed>,
+        #[allow(dead_code)]
+        instance: Column<Instance>,
+    }
+
+    impl StandardPlonkConfig {
+        fn configure(meta: &mut ConstraintSystem<Fr>) -> Self {
+            let [a, b, c] = [(); 3].map(|_| meta.advice_column());
+            let [q_a, q_b, q_c, q_ab, constant] = [(); 5].map(|_| meta.fixed_column());
+            let instance = meta.instance_column();
+
+            [a, b, c].map(|column| meta.enable_equality(column));
+
+            meta.create_gate(
+                "q_a·a + q_b·b + q_c·c + q_ab·a·b + constant + instance = 0",
+                |meta| {
+                    let [a, b, c] =
+                        [a, b, c].map(|column| meta.query_advice(column, Rotation::cur()));
+                    let [q_a, q_b, q_c, q_ab, constant] = [q_a, q_b, q_c, q_ab, constant]
+                        .map(|column| meta.query_fixed(column, Rotation::cur()));
+                    let instance = meta.query_instance(instance, Rotation::cur());
+                    Some(
+                        q_a * a.clone()
+                            + q_b * b.clone()
+                            + q_c * c
+                            + q_ab * a * b
+                            + constant
+                            + instance,
+                    )
+                },
+            );
+
+            StandardPlonkConfig { a, b, c, q_a, q_b, q_c, q_ab, constant, instance }
+        }
+    }
+
+    #[derive(Clone, Default)]
+    pub struct StandardPlonk(Fr);
+
+    impl StandardPlonk {
+        pub fn rand<R: RngCore>(mut rng: R) -> Self {
+            Self(Fr::from(rng.next_u32() as u64))
+        }
+    }
+
+    impl CircuitExt<Fr> for StandardPlonk {
+        fn num_instance(&self) -> Vec<usize> {
+            vec![1]
+        }
+
+        fn instances(&self) -> Vec<Vec<Fr>> {
+            vec![vec![self.0]]
+        }
+    }
+
+    impl Circuit<Fr> for StandardPlonk {
+        type Config = StandardPlonkConfig;
+        type FloorPlanner = SimpleFloorPlanner;
+
+        fn without_witnesses(&self) -> Self {
+            Self::default()
+        }
+
+        fn configure(meta: &mut ConstraintSystem<Fr>) -> Self::Config {
+            meta.set_minimum_degree(4);
+            StandardPlonkConfig::configure(meta)
+        }
+
+        fn synthesize(
+            &self,
+            config: Self::Config,
+            mut layouter: impl Layouter<Fr>,
+        ) -> Result<(), Error> {
+            layouter.assign_region(
+                || "",
+                |mut region| {
+                    #[cfg(feature = "halo2-pse")]
+                    {
+                        region.assign_advice(|| "", config.a, 0, || Value::known(self.0))?;
+                        region.assign_fixed(|| "", config.q_a, 0, || Value::known(-Fr::one()))?;
+                        region.assign_advice(
+                            || "",
+                            config.a,
+                            1,
+                            || Value::known(-Fr::from(5u64)),
+                        )?;
+                        for (idx, column) in (1..).zip([
+                            config.q_a,
+                            config.q_b,
+                            config.q_c,
+                            config.q_ab,
+                            config.constant,
+                        ]) {
+                            region.assign_fixed(
+                                || "",
+                                column,
+                                1,
+                                || Value::known(Fr::from(idx as u64)),
+                            )?;
+                        }
+                        let a =
+                            region.assign_advice(|| "", config.a, 2, || Value::known(Fr::one()))?;
+                        a.copy_advice(|| "", &mut region, config.b, 3)?;
+                        a.copy_advice(|| "", &mut region, config.c, 4)?;
+                    }
+                    #[cfg(feature = "halo2-axiom")]
+                    {
+                        region.assign_advice(config.a, 0, Value::known(self.0));
+                        region.assign_fixed(config.q_a, 0, -Fr::one());
+                        region.assign_advice(config.a, 1, Value::known(-Fr::from(5u64)));
+                        for (idx, column) in (1..).zip([
+                            config.q_a,
+                            config.q_b,
+                            config.q_c,
+                            config.q_ab,
+                            config.constant,
+                        ]) {
+                            region.assign_fixed(column, 1, Fr::from(idx as u64));
+                        }
+
+                        let a = region.assign_advice(config.a, 2, Value::known(Fr::one()));
+                        a.copy_advice(&mut region, config.b, 3);
+                        a.copy_advice(&mut region, config.c, 4);
+                    }
+
+                    Ok(())
+                },
+            )
+        }
+    }
+}
+
+fn gen_application_snark(params: &ParamsKZG<Bn256>) -> Snark {
+    let circuit = application::StandardPlonk::rand(OsRng);
+
+    let pk = gen_pk(params, &circuit, Some(Path::new("examples/app.pk")));
+    gen_snark_shplonk(params, &pk, circuit, Some(Path::new("examples/app.snark")))
+}
+
+fn bench(c: &mut Criterion) {
+    let path = "./configs/example_evm_accumulator.json";
+    let params_app = gen_srs(8);
+
+    let snarks = [(); 3].map(|_| gen_application_snark(&params_app));
+    let agg_config = AggregationConfigParams::from_path(path);
+    BASE_CONFIG_PARAMS.with(|params| *params.borrow_mut() = agg_config.into());
+    let params = gen_srs(agg_config.degree);
+
+    let agg_circuit = AggregationCircuit::keygen::<SHPLONK>(&params, snarks);
+
+    let start0 = start_timer!(|| "gen vk & pk");
+    gen_pk(&params, &agg_circuit, Some(Path::new("examples/agg.pk")));
+    end_timer!(start0);
+
+    let mut group = c.benchmark_group("read-pk");
+    group.sample_size(10);
+    group.bench_with_input("1mb", &(1024 * 1024), |b, &c| {
+        b.iter(|| read_pk_with_capacity::<AggregationCircuit>(c, "examples/agg.pk"))
+    });
+    group.bench_with_input("10mb", &(10 * 1024 * 1024), |b, &c| {
+        b.iter(|| read_pk_with_capacity::<AggregationCircuit>(c, "examples/agg.pk"))
+    });
+    group.bench_with_input("100mb", &(100 * 1024 * 1024), |b, &c| {
+        b.iter(|| read_pk_with_capacity::<AggregationCircuit>(c, "examples/agg.pk"))
+    });
+    group.bench_with_input("1gb", &(1024 * 1024 * 1024), |b, &c| {
+        b.iter(|| read_pk_with_capacity::<AggregationCircuit>(c, "examples/agg.pk"))
+    });
+    group.finish();
+}
+
+criterion_group! {
+    name = benches;
+    config = Criterion::default().with_profiler(PProfProfiler::new(10, Output::Flamegraph(None)));
+    targets = bench
+}
+criterion_main!(benches);

snark-verifier-sdk/src/lib.rs

@@ -33,6 +33,8 @@ pub mod halo2;
 pub const LIMBS: usize = 3;
 pub const BITS: usize = 88;
 
+const BUFFER_SIZE: usize = 1024 * 1024; // 1MB
+
 /// AS stands for accumulation scheme.
 /// AS can be either `Kzg<Bn256, Gwc19>` (the original PLONK KZG multi-open) or `Kzg<Bn256, Bdfg21>` (SHPLONK)
 pub type PlonkVerifier<AS> = verifier::plonk::PlonkVerifier<AS, LimbsEncoding<LIMBS, BITS>>;
@@ -77,12 +79,19 @@ pub trait CircuitExt<F: Field>: Circuit<F> {
 }
 
 pub fn read_pk<C: Circuit<Fr>>(path: &Path) -> io::Result<ProvingKey<G1Affine>> {
-    let f = File::open(path)?;
+    read_pk_with_capacity::<C>(BUFFER_SIZE, path)
+}
+
+pub fn read_pk_with_capacity<C: Circuit<Fr>>(
+    capacity: usize,
+    path: impl AsRef<Path>,
+) -> io::Result<ProvingKey<G1Affine>> {
+    let f = File::open(path.as_ref())?;
     #[cfg(feature = "display")]
-    let read_time = start_timer!(|| format!("Reading pkey from {path:?}"));
+    let read_time = start_timer!(|| format!("Reading pkey from {:?}", path.as_ref()));
 
     // BufReader is indeed MUCH faster than Read
-    let mut bufreader = BufReader::new(f);
+    let mut bufreader = BufReader::with_capacity(capacity, f);
     // But it's even faster to load the whole file into memory first and then process,
     // HOWEVER this requires twice as much memory to initialize
     // let initial_buffer_size = f.metadata().map(|m| m.len() as usize + 1).unwrap_or(0);
@@ -121,7 +130,7 @@ pub fn gen_pk<C: Circuit<Fr>>(
         let write_time = start_timer!(|| format!("Writing pkey to {path:?}"));
 
         path.parent().and_then(|dir| fs::create_dir_all(dir).ok()).unwrap();
-        let mut f = BufWriter::new(File::create(path).unwrap());
+        let mut f = BufWriter::with_capacity(BUFFER_SIZE, File::create(path).unwrap());
         pk.write(&mut f, SerdeFormat::RawBytesUnchecked).unwrap();
 
         #[cfg(feature = "display")]