feat: add initial bench_builder cmd to BaseTester

halo2-base/Cargo.toml

@@ -15,6 +15,7 @@ serde = { version = "1.0", features = ["derive"] }
 serde_json = "1.0"
 log = "0.4"
 getset = "0.1.2"
+ark-std = { version = "0.3.0", features = ["print-trace"], optional = true }
 
 # Use Axiom's custom halo2 monorepo for faster proving when feature = "halo2-axiom" is on
 halo2_proofs_axiom = { git = "https://github.com/axiom-crypto/halo2.git", package = "halo2_proofs", optional = true }
@@ -60,7 +61,7 @@ halo2-pse = ["halo2_proofs/circuit-params"]
 halo2-axiom = ["halo2_proofs_axiom"]
 display = []
 profile = ["halo2_proofs_axiom?/profile"]
-test-utils = ["dep:rand"]
+test-utils = ["dep:rand", "ark-std"]
 
 [[bench]]
 name = "mul"
@@ -69,3 +70,7 @@ harness = false
 [[bench]]
 name = "inner_product"
 harness = false
+
+[[example]]
+name = "inner_product"
+features = ["test-utils"]

halo2-base/examples/inner_product.rs

@@ -1,54 +1,39 @@
-use halo2_base::gates::builder::{GateThreadBuilder, RangeCircuitBuilder};
+#![cfg(feature = "test-utils")]
 use halo2_base::gates::flex_gate::{GateChip, GateInstructions};
-use halo2_base::halo2_proofs::{
-    arithmetic::Field,
-    dev::MockProver,
-    halo2curves::bn256::{Bn256, Fr},
-    plonk::*,
-    poly::kzg::commitment::ParamsKZG,
-};
-use halo2_base::utils::testing::{check_proof, gen_proof};
+use halo2_base::halo2_proofs::{arithmetic::Field, halo2curves::bn256::Fr};
+use halo2_base::safe_types::RangeInstructions;
+use halo2_base::utils::testing::base_test;
 use halo2_base::utils::ScalarField;
 use halo2_base::{Context, QuantumCell::Existing};
 use itertools::Itertools;
 use rand::rngs::OsRng;
 
 const K: u32 = 19;
 
-fn inner_prod_bench<F: ScalarField>(ctx: &mut Context<F>, a: Vec<F>, b: Vec<F>) {
+fn inner_prod_bench<F: ScalarField>(
+    ctx: &mut Context<F>,
+    gate: &GateChip<F>,
+    a: Vec<F>,
+    b: Vec<F>,
+) {
     assert_eq!(a.len(), b.len());
     let a = ctx.assign_witnesses(a);
     let b = ctx.assign_witnesses(b);
 
-    let chip = GateChip::default();
     for _ in 0..(1 << K) / 16 - 10 {
-        chip.inner_product(ctx, a.clone(), b.clone().into_iter().map(Existing));
+        gate.inner_product(ctx, a.clone(), b.clone().into_iter().map(Existing));
     }
 }
 
 fn main() {
-    let k = 10u32;
-    // create circuit for keygen
-    let mut builder = GateThreadBuilder::new(false);
-    inner_prod_bench(builder.main(0), vec![Fr::zero(); 5], vec![Fr::zero(); 5]);
-    let config_params = builder.config(k as usize, Some(20));
-    let circuit = RangeCircuitBuilder::mock(builder, config_params.clone());
-
-    // check the circuit is correct just in case
-    MockProver::run(k, &circuit, vec![]).unwrap().assert_satisfied();
-
-    let params = ParamsKZG::<Bn256>::setup(k, OsRng);
-    let vk = keygen_vk(&params, &circuit).expect("vk should not fail");
-    let pk = keygen_pk(&params, vk, &circuit).expect("pk should not fail");
-
-    let break_points = circuit.0.break_points.take();
-
-    let mut builder = GateThreadBuilder::new(true);
-    let a = (0..5).map(|_| Fr::random(OsRng)).collect_vec();
-    let b = (0..5).map(|_| Fr::random(OsRng)).collect_vec();
-    inner_prod_bench(builder.main(0), a, b);
-    let circuit = RangeCircuitBuilder::prover(builder, config_params, break_points);
-
-    let proof = gen_proof(&params, &pk, circuit);
-    check_proof(&params, pk.get_vk(), &proof, true);
+    base_test().k(12).bench_builder(
+        (vec![Fr::ZERO; 5], vec![Fr::ZERO; 5]),
+        (
+            (0..5).map(|_| Fr::random(OsRng)).collect_vec(),
+            (0..5).map(|_| Fr::random(OsRng)).collect_vec(),
+        ),
+        |builder, range, (a, b)| {
+            inner_prod_bench(builder.main(0), range.gate(), a, b);
+        },
+    );
 }

halo2-base/src/utils/testing.rs

@@ -1,7 +1,7 @@
 //! Utilities for testing
 use crate::{
     gates::{
-        builder::{GateThreadBuilder, RangeCircuitBuilder},
+        builder::{BaseConfigParams, GateThreadBuilder, RangeCircuitBuilder},
         GateChip,
     },
     halo2_proofs::{
@@ -20,8 +20,12 @@ use crate::{
     safe_types::RangeChip,
     Context,
 };
+use ark_std::{end_timer, perf_trace::TimerInfo, start_timer};
+use halo2_proofs_axiom::plonk::{keygen_pk, keygen_vk};
 use rand::{rngs::StdRng, SeedableRng};
 
+use super::fs::gen_srs;
+
 /// Helper function to generate a proof with real prover using SHPLONK KZG multi-open polynomical commitment scheme
 /// and Blake2b as the hash function for Fiat-Shamir.
 pub fn gen_proof_with_instances(
@@ -149,7 +153,6 @@ impl BaseTester {
     }
 
     /// Run a mock test by providing a closure that uses a `builder` and `RangeChip`.
-    /// - `expect_satisfied`: flag for whether you expect the test to pass or fail. Failure means a constraint system failure -- the tester does not catch system panics.
     pub fn run_builder<R>(
         &self,
         f: impl FnOnce(&mut GateThreadBuilder<Fr>, &RangeChip<Fr>) -> R,
@@ -179,4 +182,79 @@ impl BaseTester {
         }
         res
     }
+
+    /// Runs keygen, real prover, and verifier by providing a closure that uses a `builder` and `RangeChip`.
+    ///
+    /// Must provide `init_input` for use during key generation, which is preferably not equal to `logic_input`.
+    /// These are the inputs to the closure, not necessary public inputs to the circuit.
+    ///
+    /// Currently for testing, no public instances.
+    pub fn bench_builder<I: Clone>(
+        &self,
+        init_input: I,
+        logic_input: I,
+        f: impl Fn(&mut GateThreadBuilder<Fr>, &RangeChip<Fr>, I),
+    ) -> BenchStats {
+        let mut builder = GateThreadBuilder::keygen();
+        let range = RangeChip::default(self.lookup_bits.unwrap_or(0));
+        // run the function, mutating `builder`
+        f(&mut builder, &range, init_input);
+
+        // helper check: if your function didn't use lookups, turn lookup table "off"
+        let t_cells_lookup = builder
+            .threads
+            .iter()
+            .map(|t| t.iter().map(|ctx| ctx.cells_to_lookup.len()).sum::<usize>())
+            .sum::<usize>();
+        let lookup_bits = if t_cells_lookup == 0 { None } else { self.lookup_bits };
+
+        // configure the circuit shape, 9 blinding rows seems enough
+        let mut config_params = builder.config(self.k as usize, Some(9));
+        config_params.lookup_bits = lookup_bits;
+        dbg!(&config_params);
+        let circuit = RangeCircuitBuilder::keygen(builder, config_params.clone());
+
+        let params = gen_srs(config_params.k as u32);
+        let vk_time = start_timer!(|| "Generating vkey");
+        let vk = keygen_vk(&params, &circuit).unwrap();
+        end_timer!(vk_time);
+        let pk_time = start_timer!(|| "Generating pkey");
+        let pk = keygen_pk(&params, vk, &circuit).unwrap();
+        end_timer!(pk_time);
+
+        let break_points = circuit.0.break_points.borrow().clone();
+        drop(circuit);
+        // create real proof
+        let proof_time = start_timer!(|| "Proving time");
+        let mut builder = GateThreadBuilder::prover();
+        let range = RangeChip::default(self.lookup_bits.unwrap_or(0));
+        f(&mut builder, &range, logic_input);
+        let circuit = RangeCircuitBuilder::prover(builder, config_params.clone(), break_points);
+        let proof = gen_proof(&params, &pk, circuit);
+        end_timer!(proof_time);
+
+        let proof_size = proof.len();
+
+        let verify_time = start_timer!(|| "Verify time");
+        check_proof(&params, pk.get_vk(), &proof, self.expect_satisfied);
+        end_timer!(verify_time);
+
+        BenchStats { config_params, vk_time, pk_time, proof_time, proof_size, verify_time }
+    }
+}
+
+/// Bench stats
+pub struct BenchStats {
+    /// Config params
+    pub config_params: BaseConfigParams,
+    /// Vkey gen time
+    pub vk_time: TimerInfo,
+    /// Pkey gen time
+    pub pk_time: TimerInfo,
+    /// Proving time
+    pub proof_time: TimerInfo,
+    /// Proof size in bytes
+    pub proof_size: usize,
+    /// Verify time
+    pub verify_time: TimerInfo,
 }