bench: Tidy up PLONK benchmark

- Added keygen benchmark.
- Refactored to extract common prover and verifier logic.
- Benchmarks are now grouped.
- Prover (and keygen) benchmarks now only take 10 samples instead of 100
  (to make them feasible to collect in CI).

benches/plonk.rs

@@ -11,16 +11,13 @@ use halo2::transcript::{Blake2bRead, Blake2bWrite, Challenge255};
 
 use std::marker::PhantomData;
 
-use criterion::Criterion;
+use criterion::{BenchmarkId, Criterion};
 
-fn bench_with_k(name: &str, k: u32, c: &mut Criterion) {
+fn criterion_benchmark(c: &mut Criterion) {
     /// This represents an advice column at a certain row in the ConstraintSystem
     #[derive(Copy, Clone, Debug)]
     pub struct Variable(Column<Advice>, usize);
 
-    // Initialize the polynomial commitment parameters
-    let params: Params<EqAffine> = Params::new(k);
-
     #[derive(Clone)]
     struct PlonkConfig {
         a: Column<Advice>,
@@ -247,61 +244,74 @@ fn bench_with_k(name: &str, k: u32, c: &mut Criterion) {
         }
     }
 
-    let empty_circuit: MyCircuit<Fp> = MyCircuit { a: None, k };
+    fn keygen(k: u32) -> (Params<EqAffine>, ProvingKey<EqAffine>) {
+        let params: Params<EqAffine> = Params::new(k);
+        let empty_circuit: MyCircuit<Fp> = MyCircuit { a: None, k };
+        let vk = keygen_vk(&params, &empty_circuit).expect("keygen_vk should not fail");
+        let pk = keygen_pk(&params, vk, &empty_circuit).expect("keygen_pk should not fail");
+        (params, pk)
+    }
 
-    // Initialize the proving key
-    let vk = keygen_vk(&params, &empty_circuit).expect("keygen_vk should not fail");
-    let pk = keygen_pk(&params, vk, &empty_circuit).expect("keygen_pk should not fail");
+    fn prover(k: u32, params: &Params<EqAffine>, pk: &ProvingKey<EqAffine>) -> Vec<u8> {
+        let circuit: MyCircuit<Fp> = MyCircuit {
+            a: Some(Fp::rand()),
+            k,
+        };
 
-    let prover_name = name.to_string() + "-prover";
-    let verifier_name = name.to_string() + "-verifier";
+        let mut transcript = Blake2bWrite::<_, _, Challenge255<_>>::init(vec![]);
+        create_proof(params, pk, &[circuit], &[&[]], &mut transcript)
+            .expect("proof generation should not fail");
+        transcript.finalize()
+    }
 
-    c.bench_function(&prover_name, |b| {
-        b.iter(|| {
-            let circuit: MyCircuit<Fp> = MyCircuit {
-                a: Some(Fp::rand()),
-                k,
-            };
+    fn verifier(params: &Params<EqAffine>, vk: &VerifyingKey<EqAffine>, proof: &[u8]) {
+        let msm = params.empty_msm();
+        let mut transcript = Blake2bRead::<_, _, Challenge255<_>>::init(proof);
+        let guard = verify_proof(params, vk, msm, &[&[]], &mut transcript).unwrap();
+        let msm = guard.clone().use_challenges();
+        assert!(msm.eval());
+    }
 
-            // Create a proof
-            let mut transcript = Blake2bWrite::<_, _, Challenge255<_>>::init(vec![]);
-            create_proof(&params, &pk, &[circuit], &[&[]], &mut transcript)
-                .expect("proof generation should not fail")
+    let k_range = 8..=16;
+
+    let mut keygen_group = c.benchmark_group("plonk-keygen");
+    keygen_group.sample_size(10);
+    for k in k_range.clone() {
+        keygen_group.bench_with_input(BenchmarkId::from_parameter(k), &k, |b, &k| {
+            b.iter(|| keygen(k));
         });
-    });
+    }
+    keygen_group.finish();
 
-    let circuit: MyCircuit<Fp> = MyCircuit {
-        a: Some(Fp::rand()),
-        k,
-    };
+    let mut prover_group = c.benchmark_group("plonk-prover");
+    prover_group.sample_size(10);
+    for k in k_range.clone() {
+        let (params, pk) = keygen(k);
 
-    // Create a proof
-    let mut transcript = Blake2bWrite::<_, _, Challenge255<_>>::init(vec![]);
-    create_proof(&params, &pk, &[circuit], &[&[]], &mut transcript)
-        .expect("proof generation should not fail");
-    let proof = transcript.finalize();
+        prover_group.bench_with_input(
+            BenchmarkId::from_parameter(k),
+            &(k, &params, &pk),
+            |b, &(k, params, pk)| {
+                b.iter(|| prover(k, params, pk));
+            },
+        );
+    }
+    prover_group.finish();
 
-    c.bench_function(&verifier_name, |b| {
-        b.iter(|| {
-            let msm = params.empty_msm();
-            let mut transcript = Blake2bRead::<_, _, Challenge255<_>>::init(&proof[..]);
-            let guard = verify_proof(&params, pk.get_vk(), msm, &[&[]], &mut transcript).unwrap();
-            let msm = guard.clone().use_challenges();
-            assert!(msm.eval());
-        });
-    });
-}
+    let mut verifier_group = c.benchmark_group("plonk-verifier");
+    for k in k_range {
+        let (params, pk) = keygen(k);
+        let proof = prover(k, &params, &pk);
 
-fn criterion_benchmark(c: &mut Criterion) {
-    bench_with_k("plonk-k=8", 8, c);
-    bench_with_k("plonk-k=9", 9, c);
-    bench_with_k("plonk-k=10", 10, c);
-    bench_with_k("plonk-k=11", 11, c);
-    bench_with_k("plonk-k=12", 12, c);
-    bench_with_k("plonk-k=13", 13, c);
-    bench_with_k("plonk-k=14", 14, c);
-    bench_with_k("plonk-k=15", 15, c);
-    bench_with_k("plonk-k=16", 16, c);
+        verifier_group.bench_with_input(
+            BenchmarkId::from_parameter(k),
+            &(&params, pk.get_vk(), &proof[..]),
+            |b, &(params, vk, proof)| {
+                b.iter(|| verifier(params, vk, proof));
+            },
+        );
+    }
+    verifier_group.finish();
 }
 
 criterion_group!(benches, criterion_benchmark);