Add prover benchmark to measure scaling (#507)

* Add per-phase heap profiling to prover instruments

* Add bench_prover_scaling.sh

* Match expand_columns_to_lde label in bench_prover_scaling.sh

* Capitalize multi_prove heap snapshot labels

* Bundle heap_before into ProveHeapProfile

* Propagate None on heap reader failure

* Require time|heap mode argument for bench script

Author: gabrielbosio

bin/cli/Cargo.toml

@@ -14,4 +14,4 @@ tikv-jemalloc-ctl = { version = "0.6", features = ["stats"], optional = true }
 
 [features]
 jemalloc-stats = ["dep:tikv-jemalloc-ctl"]
-instruments = ["prover/instruments"]
+instruments = ["prover/instruments", "stark/instruments"]

bin/cli/src/main.rs

@@ -346,6 +346,12 @@ fn cmd_prove(
     #[cfg(feature = "jemalloc-stats")]
     let tracker = heap_tracker::HeapTracker::start();
 
+    #[cfg(all(feature = "jemalloc-stats", feature = "instruments"))]
+    stark::instruments::set_heap_reader(|| {
+        tikv_jemalloc_ctl::epoch::advance().ok();
+        tikv_jemalloc_ctl::stats::allocated::read().ok()
+    });
+
     let start = Instant::now();
     let proof = match blowup {
         Some(b) => {

crypto/stark/src/instruments.rs

@@ -1,7 +1,31 @@
 use std::cell::RefCell;
+use std::sync::OnceLock;
 use std::sync::atomic::{AtomicU64, Ordering};
 use std::time::Duration;
 
+static HEAP_READER: OnceLock<fn() -> Option<usize>> = OnceLock::new();
+
+pub fn set_heap_reader(f: fn() -> Option<usize>) {
+    let _ = HEAP_READER.set(f);
+}
+
+pub fn heap_bytes() -> Option<usize> {
+    HEAP_READER.get().and_then(|f| f())
+}
+
+pub type HeapSnapshot = (&'static str, usize);
+
+pub fn snap(label: &'static str) -> Option<HeapSnapshot> {
+    heap_bytes().map(|b| (label, b))
+}
+
+pub struct ProveHeapProfile {
+    pub before: Option<usize>,
+    pub after_execute: Option<usize>,
+    pub after_trace_build: Option<usize>,
+    pub after_air: Option<usize>,
+}
+
 /// Sub-operation timing breakdown for a single table in Rounds 2-4.
 #[derive(Clone, Debug, Default)]
 pub struct TableSubOps {
@@ -47,6 +71,7 @@ pub struct MultiProveTiming {
     pub round1_sub: Round1SubOps,
     /// (name, rows, duration, sub_ops) per table for rounds 2-4.
     pub table_timings: Vec<(String, usize, Duration, TableSubOps)>,
+    pub heap_snapshots: Vec<HeapSnapshot>,
 }
 
 /// Round 1 sub-timings: atomics so parallel rayon workers can accumulate safely.

crypto/stark/src/prover.rs

@@ -1506,6 +1506,8 @@ pub trait IsStarkProver<
 
         #[cfg(feature = "instruments")]
         crate::instruments::reset_all();
+        #[cfg(feature = "instruments")]
+        let mut heap_snaps: Vec<crate::instruments::HeapSnapshot> = Vec::new();
 
         let num_airs = air_trace_pairs.len();
 
@@ -1537,6 +1539,10 @@ pub trait IsStarkProver<
 
         #[cfg(feature = "instruments")]
         let prepass_elapsed = phase_start.elapsed();
+        #[cfg(feature = "instruments")]
+        if let Some(s) = crate::instruments::snap("After pool alloc") {
+            heap_snaps.push(s);
+        }
 
         // =====================================================================
         // Round 1, Phase A: Commit all main traces (parallel in chunks of K)
@@ -1599,6 +1605,10 @@ pub trait IsStarkProver<
 
         #[cfg(feature = "instruments")]
         let main_commits_elapsed = phase_start.elapsed();
+        #[cfg(feature = "instruments")]
+        if let Some(s) = crate::instruments::snap("After main commits") {
+            heap_snaps.push(s);
+        }
 
         // =====================================================================
         // Round 1, Phase B: Sample shared LogUp challenges
@@ -1645,6 +1655,10 @@ pub trait IsStarkProver<
 
         #[cfg(feature = "instruments")]
         let aux_build_elapsed = phase_start.elapsed();
+        #[cfg(feature = "instruments")]
+        if let Some(s) = crate::instruments::snap("After aux build") {
+            heap_snaps.push(s);
+        }
 
         // Pass 2: Parallel fork transcript → extract → LDE → commit in chunks of K.
         // Each table gets its own transcript fork.
@@ -1752,6 +1766,10 @@ pub trait IsStarkProver<
 
         #[cfg(feature = "instruments")]
         let aux_commit_elapsed = phase_start.elapsed();
+        #[cfg(feature = "instruments")]
+        if let Some(s) = crate::instruments::snap("After aux commit") {
+            heap_snaps.push(s);
+        }
 
         #[cfg(feature = "debug-checks")]
         Self::run_debug_checks(&air_trace_pairs, &commitments, &domains, &twiddle_caches);
@@ -1869,6 +1887,7 @@ pub trait IsStarkProver<
                 rounds_2_4: phase_start.elapsed(),
                 round1_sub: crate::instruments::take_r1_sub(),
                 table_timings,
+                heap_snapshots: heap_snaps,
             });
         }
 

prover/src/instruments.rs

@@ -57,6 +57,7 @@ pub fn print_report(
     trace_build: Duration,
     air_construction: Duration,
     total: Duration,
+    heap_profile: &stark::instruments::ProveHeapProfile,
 ) {
     let mp = stark::instruments::take();
 
@@ -69,18 +70,26 @@ pub fn print_report(
     row_top("Trace build", trace_build, total);
     row_top("AIR construction", air_construction, total);
 
-    if let Some(mp) = mp {
+    if let Some(ref mp) = mp {
         let round1 = mp.main_commits + mp.aux_build + mp.aux_commit;
 
         row_top("Pre-pass (domains/twiddles)", mp.prepass, total);
         row_top("Round 1", round1, total);
         row_sub("  Main trace commits", mp.main_commits, total);
-        row_sub("    expand_columns_to_lde", mp.round1_sub.main_lde, total);
-        row_sub("    commit (Merkle)", mp.round1_sub.main_merkle, total);
+        row_sub(
+            "    Main expand_columns_to_lde",
+            mp.round1_sub.main_lde,
+            total,
+        );
+        row_sub("    Main commit (Merkle)", mp.round1_sub.main_merkle, total);
         row_sub("  Aux trace build (parallel)", mp.aux_build, total);
         row_sub("  Aux trace commit", mp.aux_commit, total);
-        row_sub("    expand_columns_to_lde", mp.round1_sub.aux_lde, total);
-        row_sub("    commit (Merkle)", mp.round1_sub.aux_merkle, total);
+        row_sub(
+            "    Aux expand_columns_to_lde",
+            mp.round1_sub.aux_lde,
+            total,
+        );
+        row_sub("    Aux commit (Merkle)", mp.round1_sub.aux_merkle, total);
         row_top("Rounds 2\u{2013}4", mp.rounds_2_4, total);
 
         // Merge split tables: MEMW[0..4] → MEMW x5
@@ -208,4 +217,33 @@ pub fn print_report(
     eprintln!("  {}", "─".repeat(58));
     eprintln!("  {:<36} {:>7.2}s", "TOTAL", total.as_secs_f64());
     eprintln!();
+
+    let mb = |b: usize| b / (1024 * 1024);
+    if let Some(before) = heap_profile.before {
+        eprintln!("=== HEAP PROFILE (MB) ===");
+        eprintln!("  {:<36} {:>8} {:>8}", "Phase", "Heap", "Delta");
+        eprintln!("  {}", "─".repeat(56));
+        let mut prev = before;
+        let mut print_row = |label: &str, val: Option<usize>| {
+            if let Some(v) = val {
+                let cur = mb(v);
+                let delta = mb(v) as isize - mb(prev) as isize;
+                eprintln!("  {:<36} {:>7} {:>+8}", label, cur, delta);
+                prev = v;
+            }
+        };
+        print_row("After execute", heap_profile.after_execute);
+        print_row("After trace build", heap_profile.after_trace_build);
+        print_row("After AIR construction", heap_profile.after_air);
+        if let Some(ref mp_data) = mp {
+            for (label, bytes) in &mp_data.heap_snapshots {
+                let cur = mb(*bytes);
+                let delta = cur as isize - mb(prev) as isize;
+                eprintln!("  {:<36} {:>7} {:>+8}", label, cur, delta);
+                prev = *bytes;
+            }
+        }
+        eprintln!("  {}", "─".repeat(56));
+        eprintln!();
+    }
 }

prover/src/lib.rs

@@ -511,6 +511,8 @@ pub fn prove_with_options_and_inputs(
 ) -> Result<VmProof, Error> {
     #[cfg(feature = "instruments")]
     let total_start = std::time::Instant::now();
+    #[cfg(feature = "instruments")]
+    let heap_before = stark::instruments::heap_bytes();
 
     // Phase 1: Execute (ELF load + run)
     #[cfg(feature = "instruments")]
@@ -525,6 +527,8 @@ pub fn prove_with_options_and_inputs(
 
     #[cfg(feature = "instruments")]
     let execute_elapsed = phase_start.elapsed();
+    #[cfg(feature = "instruments")]
+    let heap_after_execute = stark::instruments::heap_bytes();
 
     // Phase 2: Trace build
     #[cfg(feature = "instruments")]
@@ -536,6 +540,8 @@ pub fn prove_with_options_and_inputs(
 
     #[cfg(feature = "instruments")]
     let trace_build_elapsed = phase_start.elapsed();
+    #[cfg(feature = "instruments")]
+    let heap_after_trace = stark::instruments::heap_bytes();
 
     // Phase 3: AIR construction
     #[cfg(feature = "instruments")]
@@ -552,6 +558,8 @@ pub fn prove_with_options_and_inputs(
 
     #[cfg(feature = "instruments")]
     let air_elapsed = phase_start.elapsed();
+    #[cfg(feature = "instruments")]
+    let heap_after_air = stark::instruments::heap_bytes();
 
     let runtime_page_ranges = traces.runtime_page_ranges();
 
@@ -569,6 +577,12 @@ pub fn prove_with_options_and_inputs(
             trace_build_elapsed,
             air_elapsed,
             total_start.elapsed(),
+            &stark::instruments::ProveHeapProfile {
+                before: heap_before,
+                after_execute: heap_after_execute,
+                after_trace_build: heap_after_trace,
+                after_air: heap_after_air,
+            },
         );
     }