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 162a201ef..731e71490 100644 --- a/bin/cli/src/main.rs +++ b/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) => { diff --git a/crypto/stark/src/instruments.rs b/crypto/stark/src/instruments.rs index 3c9f432d7..16ff95082 100644 --- a/crypto/stark/src/instruments.rs +++ 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 Option> = OnceLock::new(); + +pub fn set_heap_reader(f: fn() -> Option) { + let _ = HEAP_READER.set(f); +} + +pub fn heap_bytes() -> Option { + HEAP_READER.get().and_then(|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 before: Option, + 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 +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, } /// 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..41ccb8366 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..f15223e18 100644 --- a/prover/src/instruments.rs +++ b/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| { + 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 1c7bc05c9..3cbbcd668 100644 --- a/prover/src/lib.rs +++ b/prover/src/lib.rs @@ -511,6 +511,8 @@ pub fn prove_with_options_and_inputs( ) -> 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")] @@ -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, + }, ); } diff --git a/scripts/bench_prover_scaling.sh b/scripts/bench_prover_scaling.sh new file mode 100755 index 000000000..c1196d76e --- /dev/null +++ b/scripts/bench_prover_scaling.sh @@ -0,0 +1,279 @@ +#!/bin/bash +# Prover scaling benchmark across fib_iterative sizes. +# 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] +# +# 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 + +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' + +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 + +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 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" + +# 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_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 } + /^ 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 + +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) +# --------------------------------------------------------------------------- + +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 + +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/"