benchmark fix: use iter_custom to measure routine time exactly (#776)

### context
the issue was identified when running sumcheck benchmark with setup but
empty routine. The benchmark results shows there still ~60ms even
routine is empty due to time was also measure something hidden out from
desired routine

### solution
by criterion
[docs](https://github.com/bheisler/criterion.rs/blob/0.5.0/src/bencher.rs#L24-L33).
`iter_with_setup` will measure the time when setup object implementing
`Drop`, which might be the root cause. The most accurate way I found is
switch to use `iter_custom` so we can control elapsed time accurately.

ceno_zkvm/benches/fibonacci.rs

@@ -1,8 +1,4 @@
-use std::{
-    fs,
-    path::PathBuf,
-    time::{Duration, Instant},
-};
+use std::{fs, path::PathBuf, time::Duration};
 
 use ceno_emul::{Platform, Program};
 use ceno_zkvm::{
@@ -82,27 +78,28 @@ fn fibonacci_prove(c: &mut Criterion) {
                 format!("fibonacci_max_steps_{}", max_steps),
             ),
             |b| {
-                b.iter_with_setup(
-                    || {
-                        run_e2e_with_checkpoint::<E, Pcs>(
+                b.iter_custom(|iters| {
+                    let mut time = Duration::new(0, 0);
+                    for _ in 0..iters {
+                        let (_, run_e2e_proof) = run_e2e_with_checkpoint::<E, Pcs>(
                             program.clone(),
                             platform.clone(),
                             vec![],
                             max_steps,
                             Checkpoint::PrepE2EProving,
-                        )
-                    },
-                    |(_, run_e2e_proof)| {
-                        let timer = Instant::now();
-
+                        );
+                        let instant = std::time::Instant::now();
                         run_e2e_proof();
+                        let elapsed = instant.elapsed();
                         println!(
                             "Fibonacci::create_proof, max_steps = {}, time = {}",
                             max_steps,
-                            timer.elapsed().as_secs_f64()
+                            elapsed.as_secs_f64()
                         );
-                    },
-                );
+                        time += elapsed;
+                    }
+                    time
+                });
             },
         );
 

ceno_zkvm/benches/fibonacci_witness.rs

@@ -49,20 +49,23 @@ fn fibonacci_witness(c: &mut Criterion) {
             format!("fib_wit_max_steps_{}", max_steps),
         ),
         |b| {
-            b.iter_with_setup(
-                || {
-                    run_e2e_with_checkpoint::<E, Pcs>(
+            b.iter_custom(|iters| {
+                let mut time = Duration::new(0, 0);
+                for _ in 0..iters {
+                    let (_, generate_witness) = run_e2e_with_checkpoint::<E, Pcs>(
                         program.clone(),
                         platform.clone(),
                         vec![],
                         max_steps,
                         Checkpoint::PrepWitnessGen,
-                    )
-                },
-                |(_, generate_witness)| {
+                    );
+                    let instant = std::time::Instant::now();
                     generate_witness();
-                },
-            );
+                    let elapsed = instant.elapsed();
+                    time += elapsed;
+                }
+                time
+            });
         },
     );
 

ceno_zkvm/benches/riscv_add.rs

@@ -1,4 +1,4 @@
-use std::time::{Duration, Instant};
+use std::time::Duration;
 
 use ark_std::test_rng;
 use ceno_zkvm::{
@@ -70,22 +70,22 @@ fn bench_add(c: &mut Criterion) {
         group.bench_function(
             BenchmarkId::new("prove_add", format!("prove_add_log2_{}", instance_num_vars)),
             |b| {
-                b.iter_with_setup(
-                    || {
+                b.iter_custom(|iters| {
+                    let mut time = Duration::new(0, 0);
+                    for _ in 0..iters {
                         // generate mock witness
                         let mut rng = test_rng();
                         let num_instances = 1 << instance_num_vars;
-                        (0..num_witin as usize)
+                        let wits_in = (0..num_witin as usize)
                             .map(|_| {
                                 (0..num_instances)
                                     .map(|_| Goldilocks::random(&mut rng))
                                     .collect::<Vec<Goldilocks>>()
                                     .into_mle()
                             })
-                            .collect_vec()
-                    },
-                    |wits_in| {
-                        let timer = Instant::now();
+                            .collect_vec();
+
+                        let instant = std::time::Instant::now();
                         let num_instances = 1 << instance_num_vars;
                         let mut transcript = BasicTranscript::new(b"riscv");
                         let commit =
@@ -109,13 +109,16 @@ fn bench_add(c: &mut Criterion) {
                                 &challenges,
                             )
                             .expect("create_proof failed");
+                        let elapsed = instant.elapsed();
                         println!(
                             "AddInstruction::create_proof, instance_num_vars = {}, time = {}",
                             instance_num_vars,
-                            timer.elapsed().as_secs_f64()
+                            elapsed.as_secs_f64()
                         );
-                    },
-                );
+                        time += elapsed;
+                    }
+                    time
+                });
             },
         );
 

sumcheck/benches/devirgo_sumcheck.rs

@@ -1,7 +1,7 @@
 #![allow(clippy::manual_memcpy)]
 #![allow(clippy::needless_range_loop)]
 
-use std::array;
+use std::{array, time::Duration};
 
 use ark_std::test_rng;
 use criterion::*;
@@ -112,31 +112,23 @@ fn sumcheck_fn(c: &mut Criterion) {
         group.bench_function(
             BenchmarkId::new("prove_sumcheck", format!("sumcheck_nv_{}", nv)),
             |b| {
-                b.iter_with_setup(
-                    || {
-                        let prover_transcript = Transcript::<E>::new(b"test");
-                        let (asserted_sum, virtual_poly, virtual_poly_splitted) =
-                            { prepare_input(nv) };
-                        (
-                            prover_transcript,
-                            asserted_sum,
-                            virtual_poly,
-                            virtual_poly_splitted,
-                        )
-                    },
-                    |(
-                        mut prover_transcript,
-                        _asserted_sum,
-                        virtual_poly,
-                        _virtual_poly_splitted,
-                    )| {
+                b.iter_custom(|iters| {
+                    let mut time = Duration::new(0, 0);
+                    for _ in 0..iters {
+                        let mut prover_transcript = Transcript::<E>::new(b"test");
+                        let (_, virtual_poly, _) = { prepare_input(nv) };
+
+                        let instant = std::time::Instant::now();
                         #[allow(deprecated)]
                         let (_sumcheck_proof_v1, _) = IOPProverState::<E>::prove_parallel(
                             virtual_poly.clone(),
                             &mut prover_transcript,
                         );
-                    },
-                );
+                        let elapsed = instant.elapsed();
+                        time += elapsed;
+                    }
+                    time
+                });
             },
         );
 
@@ -157,31 +149,23 @@ fn devirgo_sumcheck_fn(c: &mut Criterion) {
         group.bench_function(
             BenchmarkId::new("prove_sumcheck", format!("devirgo_nv_{}", nv)),
             |b| {
-                b.iter_with_setup(
-                    || {
-                        let prover_transcript = Transcript::<E>::new(b"test");
-                        let (asserted_sum, virtual_poly, virtual_poly_splitted) =
-                            { prepare_input(nv) };
-                        (
-                            prover_transcript,
-                            asserted_sum,
-                            virtual_poly,
-                            virtual_poly_splitted,
-                        )
-                    },
-                    |(
-                        mut prover_transcript,
-                        _asserted_sum,
-                        _virtual_poly,
-                        virtual_poly_splitted,
-                    )| {
+                b.iter_custom(|iters| {
+                    let mut time = Duration::new(0, 0);
+                    for _ in 0..iters {
+                        let mut prover_transcript = Transcript::<E>::new(b"test");
+                        let (_, _, virtual_poly_splitted) = { prepare_input(nv) };
+
+                        let instant = std::time::Instant::now();
                         let (_sumcheck_proof_v2, _) = IOPProverState::<E>::prove_batch_polys(
                             threads,
                             virtual_poly_splitted,
                             &mut prover_transcript,
                         );
-                    },
-                );
+                        let elapsed = instant.elapsed();
+                        time += elapsed;
+                    }
+                    time
+                });
             },
         );