From c19b24bc13b92718f06ee1aa33737fab833da7d3 Mon Sep 17 00:00:00 2001 From: gabrielbosio Date: Fri, 17 Apr 2026 16:37:05 -0300 Subject: [PATCH 1/7] Add per-phase heap profiling to prover instruments --- bin/cli/Cargo.toml | 2 +- bin/cli/src/main.rs | 6 +++++ crypto/stark/src/instruments.rs | 24 +++++++++++++++++++ crypto/stark/src/prover.rs | 19 +++++++++++++++ prover/src/instruments.rs | 41 +++++++++++++++++++++++++++++---- prover/src/lib.rs | 14 +++++++++++ 6 files changed, 100 insertions(+), 6 deletions(-) diff --git a/bin/cli/Cargo.toml b/bin/cli/Cargo.toml index 4bfcb7956..fdc8eab8c 100644 --- a/bin/cli/Cargo.toml +++ b/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"] diff --git a/bin/cli/src/main.rs b/bin/cli/src/main.rs index 725f0de5f..96568aa9b 100644 --- a/bin/cli/src/main.rs +++ b/bin/cli/src/main.rs @@ -262,6 +262,12 @@ fn cmd_prove(elf_path: PathBuf, output_path: PathBuf, blowup: Option, 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) => { diff --git a/crypto/stark/src/instruments.rs b/crypto/stark/src/instruments.rs index 3c9f432d7..3d8d84a35 100644 --- a/crypto/stark/src/instruments.rs +++ 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 usize> = OnceLock::new(); + +pub fn set_heap_reader(f: fn() -> usize) { + let _ = HEAP_READER.set(f); +} + +pub fn heap_bytes() -> Option { + HEAP_READER.get().map(|f| f()) +} + +pub type HeapSnapshot = (&'static str, usize); + +pub fn snap(label: &'static str) -> Option { + heap_bytes().map(|b| (label, b)) +} + +pub struct ProveHeapProfile { + pub after_execute: Option, + pub after_trace_build: Option, + pub after_air: Option, +} + /// Sub-operation timing breakdown for a single table in Rounds 2-4. #[derive(Clone, Debug, Default)] pub struct TableSubOps { @@ -47,6 +70,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, } /// Round 1 sub-timings: atomics so parallel rayon workers can accumulate safely. diff --git a/crypto/stark/src/prover.rs b/crypto/stark/src/prover.rs index 8e59807c1..b96f127a7 100644 --- a/crypto/stark/src/prover.rs +++ b/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 = 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, }); } diff --git a/prover/src/instruments.rs b/prover/src/instruments.rs index e087cf55a..d94a55bec 100644 --- a/prover/src/instruments.rs +++ b/prover/src/instruments.rs @@ -57,6 +57,8 @@ pub fn print_report( trace_build: Duration, air_construction: Duration, total: Duration, + heap_before: Option, + 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_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 +210,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| { + 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!(); + } } diff --git a/prover/src/lib.rs b/prover/src/lib.rs index 7cb947f6f..b5d0bfd17 100644 --- a/prover/src/lib.rs +++ b/prover/src/lib.rs @@ -494,6 +494,8 @@ pub fn prove_with_options( ) -> Result { #[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, + }, ); } From 68e45c9f7ba9cddcbe1b4c727610d8bb66871fb7 Mon Sep 17 00:00:00 2001 From: gabrielbosio Date: Fri, 17 Apr 2026 16:37:41 -0300 Subject: [PATCH 2/7] Add bench_prover_scaling.sh --- scripts/bench_prover_scaling.sh | 262 ++++++++++++++++++++++++++++++++ 1 file changed, 262 insertions(+) create mode 100755 scripts/bench_prover_scaling.sh diff --git a/scripts/bench_prover_scaling.sh b/scripts/bench_prover_scaling.sh new file mode 100755 index 000000000..04f045333 --- /dev/null +++ b/scripts/bench_prover_scaling.sh @@ -0,0 +1,262 @@ +#!/bin/bash +# Prover scaling benchmark across fib_iterative sizes. +# Shows how per-phase timing and heap grow with program size, plus linear regression. +# +# Usage: bench_prover_scaling.sh [--sizes "500k 1M 2M 4M"] [--runs N] +# +# Requires: instruments + jemalloc-stats features. + +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +ROOT_DIR="$(cd "$SCRIPT_DIR/.." && pwd)" +TMP_DIR="/tmp/bench_prover_scaling" +ELF_DIR="$ROOT_DIR/executor/program_artifacts/asm" + +GREEN='\033[0;32m' +BOLD='\033[1m' +NC='\033[0m' + +SIZES="500k 1M 2M 4M" +RUNS=1 + +while [[ $# -gt 0 ]]; do + case $1 in + --sizes) SIZES="$2"; shift 2 ;; + --runs) RUNS="$2"; shift 2 ;; + *) echo "Unknown arg: $1"; exit 1 ;; + esac +done + +suffix_to_steps() { + case $1 in + 160k) echo 160000 ;; 250k) echo 250000 ;; 372k) echo 372000 ;; + 500k) echo 500000 ;; 1M) echo 1000000 ;; 1200k) echo 1200000 ;; + 2M) echo 2000000 ;; 4M) echo 4000000 ;; 8M) echo 8000000 ;; + 16M) echo 16000000 ;; 32M) echo 32000000 ;; 64M) echo 64000000 ;; 128M) echo 128000000 ;; + *) echo "Unknown size: $1" >&2; exit 1 ;; + esac +} + +rm -rf "$TMP_DIR" && mkdir -p "$TMP_DIR" + +echo -e "${GREEN}Building CLI with jemalloc-stats + instruments...${NC}" +cargo build --release -p cli --features jemalloc-stats,instruments \ + --manifest-path "$ROOT_DIR/Cargo.toml" 2>&1 | tail -1 +CLI="$ROOT_DIR/target/release/cli" + +# Parse timing (seconds) and heap (MB) from one run's stdout + stderr. +# Emits key=value lines to stdout. +parse_run() { + local stdout=$1 stderr=$2 + awk ' + function secs( s) { + if (match($0, /[0-9]+\.?[0-9]*s/)) { + s = substr($0, RSTART, RLENGTH - 1) + return s + } + return "" + } + /^ Execute / { v = secs(); if (v) print "t_execute=" v } + /^ Trace build/ { v = secs(); if (v) print "t_trace_build=" v } + /^ AIR construction/ { v = secs(); if (v) print "t_air=" v } + /^ Pre-pass/ { v = secs(); if (v) print "t_prepass=" v } + /^ Round 1 / { v = secs(); if (v) print "t_round1=" v } + /Main trace commits/ { v = secs(); if (v) print "t_main_commits="v } + /Aux trace build/ { v = secs(); if (v) print "t_aux_build=" v } + /Aux trace commit/ { v = secs(); if (v) print "t_aux_commit=" v } + /Rounds 2/ { v = secs(); if (v) print "t_rounds24=" v } + /Main expand_pool_to_lde/{ v = secs(); if (v) print "t_main_lde=" v } + /Aux expand_pool_to_lde/ { v = secs(); if (v) print "t_aux_lde=" v } + /Main commit \(Merkle\)/ { v = secs(); if (v) print "t_main_merkle=" v } + /Aux commit \(Merkle\)/ { v = secs(); if (v) print "t_aux_merkle=" v } + /^ Total FFT/ { v = secs(); if (v) print "t_total_fft=" v } + /^ Total Merkle/ { v = secs(); if (v) print "t_total_merkle="v } + /^ TOTAL / { v = secs(); if (v) print "t_total=" v } + /After execute/ { print "h_execute=" $(NF-1) } + /After trace build/ { print "h_trace_build=" $(NF-1) } + /After AIR/ { print "h_air=" $(NF-1) } + /after pool alloc/ { print "h_pool_alloc=" $(NF-1) } + /after main commits/ { print "h_main_commits=" $(NF-1) } + /after aux build/ { print "h_aux_build=" $(NF-1) } + /after aux commit/ { print "h_aux_commit=" $(NF-1) } + ' "$stderr" + + grep -o 'Peak heap: [0-9]*' "$stdout" | awk '{print "peak=" $3}' + grep -o 'Proving time: [0-9.]*' "$stdout" | awk '{print "wall=" $3}' +} + +get_val() { + awk -F= -v k="$2" '$1==k {print $2; exit}' "$1" 2>/dev/null +} + +# Pick the run whose wall-clock is the median across RUNS; copy its data file. +# Writes final per-size data to ${size}_data.txt. +select_median_run() { + local size=$1 + if [ "$RUNS" -eq 1 ]; then + cp "$TMP_DIR/${size}_run_1.txt" "$TMP_DIR/${size}_data.txt" + return + fi + # Collect (wall, run_index), sort by wall, pick middle + local pairs="" + for i in $(seq 1 "$RUNS"); do + local w + w=$(get_val "$TMP_DIR/${size}_run_${i}.txt" wall) + pairs+=" $w $i" + done + local median_idx + median_idx=$(echo "$pairs" | tr ' ' '\n' | awk 'NF' | paste -d' ' - - | \ + sort -n | awk -v n="$RUNS" 'NR == int((n+1)/2) {print $2}') + cp "$TMP_DIR/${size}_run_${median_idx}.txt" "$TMP_DIR/${size}_data.txt" +} + +for size in $SIZES; do + ELF="$ELF_DIR/fib_iterative_${size}.elf" + [ -f "$ELF" ] || { echo "Missing: $ELF"; continue; } + steps=$(suffix_to_steps "$size") + echo -e "${GREEN}Running fib_iterative_${size} (${RUNS} runs)...${NC}" + + for i in $(seq 1 "$RUNS"); do + STDOUT="$TMP_DIR/${size}_run_${i}_stdout.txt" + STDERR="$TMP_DIR/${size}_run_${i}_stderr.txt" + "$CLI" prove "$ELF" -o "$TMP_DIR/proof.bin" --time >"$STDOUT" 2>"$STDERR" + rm -f "$TMP_DIR/proof.bin" + { + echo "steps=$steps" + parse_run "$STDOUT" "$STDERR" + } > "$TMP_DIR/${size}_run_${i}.txt" + done + select_median_run "$size" +done + +# --------------------------------------------------------------------------- +# Rendering +# --------------------------------------------------------------------------- + +print_row() { + local label=$1 key=$2 unit=$3 + printf " %-26s" "$label" + for size in $SIZES; do + local data="$TMP_DIR/${size}_data.txt" + if [ ! -f "$data" ]; then + printf " %10s" "-"; continue + fi + local v + v=$(get_val "$data" "$key") + if [ -z "$v" ]; then + printf " %10s" "-" + elif [ "$unit" = "s" ]; then + printf " %9.2fs" "$v" + else + printf " %10d" "$v" + fi + done + echo "" +} + +print_header() { + printf " %-26s" "" + for size in $SIZES; do printf " %10s" "$size"; done + echo "" + printf " %-26s" "" + for size in $SIZES; do printf " %10s" "──────────"; done + echo "" +} + +echo "" +echo -e "${BOLD}=== TIMING (seconds) ===${NC}" +print_header +print_row "Execute" t_execute s +print_row "Trace build" t_trace_build s +print_row "AIR construction" t_air s +print_row "Pre-pass" t_prepass s +print_row "Round 1" t_round1 s +print_row " Main trace commits" t_main_commits s +print_row " Main LDE" t_main_lde s +print_row " Main Merkle" t_main_merkle s +print_row " Aux trace build" t_aux_build s +print_row " Aux trace commit" t_aux_commit s +print_row " Aux LDE" t_aux_lde s +print_row " Aux Merkle" t_aux_merkle s +print_row "Rounds 2-4" t_rounds24 s +print_row "Total FFT (all rounds)" t_total_fft s +print_row "Total Merkle" t_total_merkle s +print_row "TOTAL" t_total s + +echo "" +echo -e "${BOLD}=== HEAP (MB absolute) ===${NC}" +print_header +print_row "After execute" h_execute mb +print_row "After trace build" h_trace_build mb +print_row "After AIR construction" h_air mb +print_row "After pool alloc" h_pool_alloc mb +print_row "After main commits" h_main_commits mb +print_row "After aux build" h_aux_build mb +print_row "After aux commit" h_aux_commit mb +print_row "Peak heap" peak mb + +# --------------------------------------------------------------------------- +# Linear regression per metric: y = a + b * (steps / 1M) +# --------------------------------------------------------------------------- + +regress() { + local label=$1 key=$2 unit=$3 + local pairs="" + for size in $SIZES; do + local data="$TMP_DIR/${size}_data.txt" + [ -f "$data" ] || continue + local steps v + steps=$(get_val "$data" steps) + v=$(get_val "$data" "$key") + [ -z "$v" ] && continue + [ -z "$steps" ] && continue + local steps_m + steps_m=$(awk -v s="$steps" 'BEGIN {printf "%.6f", s / 1000000}') + pairs+=" $steps_m $v" + done + echo "$pairs" | awk -v label="$label" -v unit="$unit" '{ + n = NF / 2 + if (n < 2) { printf " %-26s (insufficient data)\n", label; next } + for (i = 0; i < n; i++) { x[i] = $(2*i+1); y[i] = $(2*i+2) } + sx=0; sy=0; sxx=0; sxy=0 + for (i=0;i0) ? 1 - ss_res/ss_tot : 1 + if (unit == "s") + printf " %-26s %+7.2fs/M (base %6.2fs, R\xc2\xb2=%.3f)\n", label, b, a, r2 + else + printf " %-26s %+7.0f MB/M (base %6.0f MB, R\xc2\xb2=%.3f)\n", label, b, a, r2 + }' +} + +echo "" +echo -e "${BOLD}=== TIMING GROWTH (per 1M steps) ===${NC}" +regress "Execute" t_execute s +regress "Trace build" t_trace_build s +regress "AIR construction" t_air s +regress "Pre-pass" t_prepass s +regress "Round 1" t_round1 s +regress "Rounds 2-4" t_rounds24 s +regress "Total FFT" t_total_fft s +regress "Total Merkle" t_total_merkle s +regress "TOTAL" t_total s + +echo "" +echo -e "${BOLD}=== HEAP GROWTH (per 1M steps) ===${NC}" +regress "After execute" h_execute mb +regress "After trace build" h_trace_build mb +regress "After AIR construction" h_air mb +regress "After pool alloc" h_pool_alloc mb +regress "After main commits" h_main_commits mb +regress "After aux build" h_aux_build mb +regress "After aux commit" h_aux_commit mb +regress "Peak heap" peak mb + +echo "" +echo "Raw data: $TMP_DIR/" From 67e3e18b151218933789bcf59b6978f4a303c215 Mon Sep 17 00:00:00 2001 From: gabrielbosio Date: Fri, 17 Apr 2026 16:44:49 -0300 Subject: [PATCH 3/7] Match expand_columns_to_lde label in bench_prover_scaling.sh --- scripts/bench_prover_scaling.sh | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/scripts/bench_prover_scaling.sh b/scripts/bench_prover_scaling.sh index 04f045333..ac26f90f8 100755 --- a/scripts/bench_prover_scaling.sh +++ b/scripts/bench_prover_scaling.sh @@ -66,8 +66,8 @@ parse_run() { /Aux trace build/ { v = secs(); if (v) print "t_aux_build=" v } /Aux trace commit/ { v = secs(); if (v) print "t_aux_commit=" v } /Rounds 2/ { v = secs(); if (v) print "t_rounds24=" v } - /Main expand_pool_to_lde/{ v = secs(); if (v) print "t_main_lde=" v } - /Aux expand_pool_to_lde/ { v = secs(); if (v) print "t_aux_lde=" v } + /Main expand_columns_to_lde/{ v = secs(); if (v) print "t_main_lde=" v } + /Aux expand_columns_to_lde/ { v = secs(); if (v) print "t_aux_lde=" v } /Main commit \(Merkle\)/ { v = secs(); if (v) print "t_main_merkle=" v } /Aux commit \(Merkle\)/ { v = secs(); if (v) print "t_aux_merkle=" v } /^ Total FFT/ { v = secs(); if (v) print "t_total_fft=" v } From 6967e536270a21f7ffad7081a5bef09de340891e Mon Sep 17 00:00:00 2001 From: gabrielbosio Date: Fri, 17 Apr 2026 18:11:02 -0300 Subject: [PATCH 4/7] Capitalize multi_prove heap snapshot labels --- crypto/stark/src/prover.rs | 8 ++++---- scripts/bench_prover_scaling.sh | 8 ++++---- 2 files changed, 8 insertions(+), 8 deletions(-) diff --git a/crypto/stark/src/prover.rs b/crypto/stark/src/prover.rs index b96f127a7..41ccb8366 100644 --- a/crypto/stark/src/prover.rs +++ b/crypto/stark/src/prover.rs @@ -1540,7 +1540,7 @@ 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") { + if let Some(s) = crate::instruments::snap("After pool alloc") { heap_snaps.push(s); } @@ -1606,7 +1606,7 @@ 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") { + if let Some(s) = crate::instruments::snap("After main commits") { heap_snaps.push(s); } @@ -1656,7 +1656,7 @@ 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") { + if let Some(s) = crate::instruments::snap("After aux build") { heap_snaps.push(s); } @@ -1767,7 +1767,7 @@ 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") { + if let Some(s) = crate::instruments::snap("After aux commit") { heap_snaps.push(s); } diff --git a/scripts/bench_prover_scaling.sh b/scripts/bench_prover_scaling.sh index ac26f90f8..a5390e656 100755 --- a/scripts/bench_prover_scaling.sh +++ b/scripts/bench_prover_scaling.sh @@ -76,10 +76,10 @@ parse_run() { /After execute/ { print "h_execute=" $(NF-1) } /After trace build/ { print "h_trace_build=" $(NF-1) } /After AIR/ { print "h_air=" $(NF-1) } - /after pool alloc/ { print "h_pool_alloc=" $(NF-1) } - /after main commits/ { print "h_main_commits=" $(NF-1) } - /after aux build/ { print "h_aux_build=" $(NF-1) } - /after aux commit/ { print "h_aux_commit=" $(NF-1) } + /After pool alloc/ { print "h_pool_alloc=" $(NF-1) } + /After main commits/ { print "h_main_commits=" $(NF-1) } + /After aux build/ { print "h_aux_build=" $(NF-1) } + /After aux commit/ { print "h_aux_commit=" $(NF-1) } ' "$stderr" grep -o 'Peak heap: [0-9]*' "$stdout" | awk '{print "peak=" $3}' From 3e3495d8b4688d9a88b9c0a70b97b10aaa320ce9 Mon Sep 17 00:00:00 2001 From: gabrielbosio Date: Fri, 17 Apr 2026 18:40:59 -0300 Subject: [PATCH 5/7] Bundle heap_before into ProveHeapProfile --- crypto/stark/src/instruments.rs | 1 + prover/src/instruments.rs | 15 +++++++++++---- prover/src/lib.rs | 2 +- 3 files changed, 13 insertions(+), 5 deletions(-) diff --git a/crypto/stark/src/instruments.rs b/crypto/stark/src/instruments.rs index 3d8d84a35..a4a2b283a 100644 --- a/crypto/stark/src/instruments.rs +++ b/crypto/stark/src/instruments.rs @@ -20,6 +20,7 @@ pub fn snap(label: &'static str) -> Option { } pub struct ProveHeapProfile { + pub before: Option, pub after_execute: Option, pub after_trace_build: Option, pub after_air: Option, diff --git a/prover/src/instruments.rs b/prover/src/instruments.rs index d94a55bec..f15223e18 100644 --- a/prover/src/instruments.rs +++ b/prover/src/instruments.rs @@ -57,7 +57,6 @@ pub fn print_report( trace_build: Duration, air_construction: Duration, total: Duration, - heap_before: Option, heap_profile: &stark::instruments::ProveHeapProfile, ) { let mp = stark::instruments::take(); @@ -77,11 +76,19 @@ pub fn print_report( 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(" Main expand_columns_to_lde", mp.round1_sub.main_lde, 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(" Aux expand_columns_to_lde", mp.round1_sub.aux_lde, 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); @@ -212,7 +219,7 @@ pub fn print_report( eprintln!(); let mb = |b: usize| b / (1024 * 1024); - if let Some(before) = heap_before { + if let Some(before) = heap_profile.before { eprintln!("=== HEAP PROFILE (MB) ==="); eprintln!(" {:<36} {:>8} {:>8}", "Phase", "Heap", "Delta"); eprintln!(" {}", "─".repeat(56)); diff --git a/prover/src/lib.rs b/prover/src/lib.rs index b5d0bfd17..beced633f 100644 --- a/prover/src/lib.rs +++ b/prover/src/lib.rs @@ -559,8 +559,8 @@ pub fn prove_with_options( trace_build_elapsed, air_elapsed, total_start.elapsed(), - heap_before, &stark::instruments::ProveHeapProfile { + before: heap_before, after_execute: heap_after_execute, after_trace_build: heap_after_trace, after_air: heap_after_air, From 8828e65384a2541e5d9139c803e6aa7004c5aca4 Mon Sep 17 00:00:00 2001 From: gabrielbosio Date: Fri, 17 Apr 2026 18:44:08 -0300 Subject: [PATCH 6/7] Propagate None on heap reader failure --- bin/cli/src/main.rs | 2 +- crypto/stark/src/instruments.rs | 6 +++--- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/bin/cli/src/main.rs b/bin/cli/src/main.rs index 96568aa9b..3eec1be51 100644 --- a/bin/cli/src/main.rs +++ b/bin/cli/src/main.rs @@ -265,7 +265,7 @@ fn cmd_prove(elf_path: PathBuf, output_path: PathBuf, blowup: Option, time: #[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) + tikv_jemalloc_ctl::stats::allocated::read().ok() }); let start = Instant::now(); diff --git a/crypto/stark/src/instruments.rs b/crypto/stark/src/instruments.rs index a4a2b283a..16ff95082 100644 --- a/crypto/stark/src/instruments.rs +++ b/crypto/stark/src/instruments.rs @@ -3,14 +3,14 @@ use std::sync::OnceLock; use std::sync::atomic::{AtomicU64, Ordering}; use std::time::Duration; -static HEAP_READER: OnceLock usize> = OnceLock::new(); +static HEAP_READER: OnceLock Option> = OnceLock::new(); -pub fn set_heap_reader(f: fn() -> usize) { +pub fn set_heap_reader(f: fn() -> Option) { let _ = HEAP_READER.set(f); } pub fn heap_bytes() -> Option { - HEAP_READER.get().map(|f| f()) + HEAP_READER.get().and_then(|f| f()) } pub type HeapSnapshot = (&'static str, usize); From 5050fedf54dd1196f659dd29849334bbb58b222a Mon Sep 17 00:00:00 2001 From: gabrielbosio Date: Fri, 17 Apr 2026 19:13:03 -0300 Subject: [PATCH 7/7] Require time|heap mode argument for bench script --- scripts/bench_prover_scaling.sh | 69 ++++++++++++++++++++------------- 1 file changed, 43 insertions(+), 26 deletions(-) diff --git a/scripts/bench_prover_scaling.sh b/scripts/bench_prover_scaling.sh index a5390e656..c1196d76e 100755 --- a/scripts/bench_prover_scaling.sh +++ b/scripts/bench_prover_scaling.sh @@ -1,10 +1,12 @@ #!/bin/bash # Prover scaling benchmark across fib_iterative sizes. -# Shows how per-phase timing and heap grow with program size, plus linear regression. +# Shows how per-phase timing or heap grows with program size, plus linear regression. # -# Usage: bench_prover_scaling.sh [--sizes "500k 1M 2M 4M"] [--runs N] +# Usage: bench_prover_scaling.sh [--sizes "500k 1M 2M 4M"] [--runs N] # -# Requires: instruments + jemalloc-stats features. +# time mode: builds with `instruments` only (clean timings, no heap data). +# heap mode: builds with `instruments,jemalloc-stats` (timings + per-phase heap; +# jemalloc-stats adds a few percent of timing overhead). set -euo pipefail @@ -17,6 +19,17 @@ GREEN='\033[0;32m' BOLD='\033[1m' NC='\033[0m' +if [[ $# -lt 1 ]]; then + echo "Usage: $0 [--sizes \"...\"] [--runs N]" >&2 + exit 1 +fi +MODE=$1; shift +case $MODE in + time) FEATURES="instruments" ;; + heap) FEATURES="instruments,jemalloc-stats" ;; + *) echo "Unknown mode: $MODE (expected time|heap)" >&2; exit 1 ;; +esac + SIZES="500k 1M 2M 4M" RUNS=1 @@ -40,8 +53,8 @@ suffix_to_steps() { rm -rf "$TMP_DIR" && mkdir -p "$TMP_DIR" -echo -e "${GREEN}Building CLI with jemalloc-stats + instruments...${NC}" -cargo build --release -p cli --features jemalloc-stats,instruments \ +echo -e "${GREEN}Building CLI with features: ${FEATURES}...${NC}" +cargo build --release -p cli --features "$FEATURES" \ --manifest-path "$ROOT_DIR/Cargo.toml" 2>&1 | tail -1 CLI="$ROOT_DIR/target/release/cli" @@ -184,17 +197,19 @@ print_row "Total FFT (all rounds)" t_total_fft s print_row "Total Merkle" t_total_merkle s print_row "TOTAL" t_total s -echo "" -echo -e "${BOLD}=== HEAP (MB absolute) ===${NC}" -print_header -print_row "After execute" h_execute mb -print_row "After trace build" h_trace_build mb -print_row "After AIR construction" h_air mb -print_row "After pool alloc" h_pool_alloc mb -print_row "After main commits" h_main_commits mb -print_row "After aux build" h_aux_build mb -print_row "After aux commit" h_aux_commit mb -print_row "Peak heap" peak mb +if [[ "$MODE" == "heap" ]]; then + echo "" + echo -e "${BOLD}=== HEAP (MB absolute) ===${NC}" + print_header + print_row "After execute" h_execute mb + print_row "After trace build" h_trace_build mb + print_row "After AIR construction" h_air mb + print_row "After pool alloc" h_pool_alloc mb + print_row "After main commits" h_main_commits mb + print_row "After aux build" h_aux_build mb + print_row "After aux commit" h_aux_commit mb + print_row "Peak heap" peak mb +fi # --------------------------------------------------------------------------- # Linear regression per metric: y = a + b * (steps / 1M) @@ -247,16 +262,18 @@ regress "Total FFT" t_total_fft s regress "Total Merkle" t_total_merkle s regress "TOTAL" t_total s -echo "" -echo -e "${BOLD}=== HEAP GROWTH (per 1M steps) ===${NC}" -regress "After execute" h_execute mb -regress "After trace build" h_trace_build mb -regress "After AIR construction" h_air mb -regress "After pool alloc" h_pool_alloc mb -regress "After main commits" h_main_commits mb -regress "After aux build" h_aux_build mb -regress "After aux commit" h_aux_commit mb -regress "Peak heap" peak mb +if [[ "$MODE" == "heap" ]]; then + echo "" + echo -e "${BOLD}=== HEAP GROWTH (per 1M steps) ===${NC}" + regress "After execute" h_execute mb + regress "After trace build" h_trace_build mb + regress "After AIR construction" h_air mb + regress "After pool alloc" h_pool_alloc mb + regress "After main commits" h_main_commits mb + regress "After aux build" h_aux_build mb + regress "After aux commit" h_aux_commit mb + regress "Peak heap" peak mb +fi echo "" echo "Raw data: $TMP_DIR/"