Add per-phase heap profiling to prover instruments

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

@@ -262,6 +262,12 @@ fn cmd_prove(elf_path: PathBuf, output_path: PathBuf, blowup: Option<u8>, time:
     #[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().unwrap_or(0)
+    });
+
     let start = Instant::now();
     let proof = match blowup {
         Some(b) => {

crypto/stark/src/instruments.rs

@@ -1,7 +1,30 @@
 use std::cell::RefCell;
+use std::sync::OnceLock;
 use std::sync::atomic::{AtomicU64, Ordering};
 use std::time::Duration;
 
+static HEAP_READER: OnceLock<fn() -> usize> = OnceLock::new();
+
+pub fn set_heap_reader(f: fn() -> usize) {
+    let _ = HEAP_READER.set(f);
+}
+
+pub fn heap_bytes() -> Option<usize> {
+    HEAP_READER.get().map(|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 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 {
@@ -49,6 +72,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

@@ -1496,6 +1496,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();
 
@@ -1547,6 +1549,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)
@@ -1609,6 +1615,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
@@ -1655,6 +1665,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 and pool set.
@@ -1755,6 +1769,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")]
         {
@@ -1897,6 +1915,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,8 @@ pub fn print_report(
     trace_build: Duration,
     air_construction: Duration,
     total: Duration,
+    heap_before: Option<usize>,
+    heap_profile: &stark::instruments::ProveHeapProfile,
 ) {
     let mp = stark::instruments::take();
 
@@ -69,18 +71,18 @@ 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_pool_to_lde", mp.round1_sub.main_lde, total);
-        row_sub("    commit (Merkle)", mp.round1_sub.main_merkle, total);
+        row_sub("    Main expand_pool_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_pool_to_lde", mp.round1_sub.aux_lde, total);
-        row_sub("    commit (Merkle)", mp.round1_sub.aux_merkle, total);
+        row_sub("    Aux expand_pool_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
@@ -214,4 +216,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_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

@@ -494,6 +494,8 @@ pub fn prove_with_options(
 ) -> 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")]
@@ -507,6 +509,8 @@ pub fn prove_with_options(
 
     #[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")]
@@ -518,6 +522,8 @@ pub fn prove_with_options(
 
     #[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")]
@@ -534,6 +540,8 @@ pub fn prove_with_options(
 
     #[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();
 
@@ -551,6 +559,12 @@ pub fn prove_with_options(
             trace_build_elapsed,
             air_elapsed,
             total_start.elapsed(),
+            heap_before,
+            &stark::instruments::ProveHeapProfile {
+                after_execute: heap_after_execute,
+                after_trace_build: heap_after_trace,
+                after_air: heap_after_air,
+            },
         );
     }