From 3b210e67f249576e2bd4cd81e6fa2ef98d00bf49 Mon Sep 17 00:00:00 2001
From: Jan Ferdinand Sauer <ferdinand@neptune.cash>
Date: Thu, 6 Jun 2024 14:40:35 +0200
Subject: [PATCH] refactor(bench): Streamline Fibonacci benchmark

- Increase accuracy of VM performance profile by accumulating times
  across multiple benchmark runs.
- Simplify code flow.
---
 triton-vm/benches/prove_fib.rs | 74 +++++++++++-----------------------
 1 file changed, 23 insertions(+), 51 deletions(-)

diff --git a/triton-vm/benches/prove_fib.rs b/triton-vm/benches/prove_fib.rs
index 645731e08..5dbd584f3 100644
--- a/triton-vm/benches/prove_fib.rs
+++ b/triton-vm/benches/prove_fib.rs
@@ -1,70 +1,42 @@
 use criterion::criterion_group;
 use criterion::criterion_main;
-use criterion::BenchmarkId;
 use criterion::Criterion;
 
-use triton_vm::aet::AlgebraicExecutionTrace;
 use triton_vm::example_programs::FIBONACCI_SEQUENCE;
 use triton_vm::prelude::*;
-use triton_vm::profiler::VMPerformanceProfile;
 
-const FIBONACCI_INDEX: BFieldElement = BFieldElement::new(100);
+const FIBONACCI_INDEX: u32 = 100;
 
-/// cargo criterion --bench prove_fib
-fn prove_fib(criterion: &mut Criterion) {
-    let (claim, aet) = trace_execution();
-    fib_benchmark_group(criterion, &claim, &aet);
-
-    let profile = prover_performance_profile(&claim, &aet);
-    eprintln!("{profile}");
-}
-
-fn fib_benchmark_group(criterion: &mut Criterion, claim: &Claim, aet: &AlgebraicExecutionTrace) {
-    let bench_id_name = format!("ProveFib{FIBONACCI_INDEX}");
-    let bench_group_name = format!("prove_fib_{FIBONACCI_INDEX}");
-    let bench_id = BenchmarkId::new(bench_id_name, 0);
-    let mut group = criterion.benchmark_group(bench_group_name);
-
-    let stark = Stark::default();
-    group.bench_function(bench_id, |bencher| bencher.iter(|| stark.prove(claim, aet)));
-    group.finish();
-}
-
-fn prover_performance_profile(
-    claim: &Claim,
-    aet: &AlgebraicExecutionTrace,
-) -> VMPerformanceProfile {
-    let profile_name = format!("Prove Fibonacci {FIBONACCI_INDEX}");
-    triton_vm::profiler::start(profile_name);
-    let stark = Stark::default();
-    let proof = stark.prove(claim, aet).unwrap();
-    let profile = triton_vm::profiler::finish();
-
-    let padded_height = proof.padded_height().unwrap();
-    let fri = stark.derive_fri(padded_height).unwrap();
-    profile
-        .with_cycle_count(aet.processor_trace.nrows())
-        .with_padded_height(padded_height)
-        .with_fri_domain_len(fri.domain.length)
+criterion_main!(benches);
+criterion_group! {
+    name = benches;
+    config = Criterion::default();
+    targets = prove_fib
 }
 
-fn trace_execution() -> (Claim, AlgebraicExecutionTrace) {
+/// cargo criterion --bench prove_fib
+fn prove_fib(c: &mut Criterion) {
     let program = FIBONACCI_SEQUENCE.clone();
-    let public_input = PublicInput::from([FIBONACCI_INDEX]);
+    let public_input = PublicInput::new(bfe_vec![FIBONACCI_INDEX]);
     let (aet, output) = program
-        .trace_execution(public_input.clone(), [].into())
+        .trace_execution(public_input.clone(), NonDeterminism::default())
         .unwrap();
 
     let claim = Claim::about_program(&program)
         .with_input(public_input.individual_tokens)
         .with_output(output);
-    (claim, aet)
-}
 
-criterion_group! {
-    name = benches;
-    config = Criterion::default();
-    targets = prove_fib
-}
+    let stark = Stark::default();
+    let bench_id = format!("Prove Fibonacci {FIBONACCI_INDEX}");
+    triton_vm::profiler::start(&bench_id);
+    c.bench_function(&bench_id, |b| b.iter(|| stark.prove(&claim, &aet).unwrap()));
+    let profile = triton_vm::profiler::finish();
 
-criterion_main!(benches);
+    let padded_height = aet.padded_height();
+    let fri = stark.derive_fri(padded_height).unwrap();
+    let profile = profile
+        .with_cycle_count(aet.processor_trace.nrows())
+        .with_padded_height(padded_height)
+        .with_fri_domain_len(fri.domain.length);
+    eprintln!("{profile}");
+}