From 2d3bb883f7fc227d609ec4323d53c202ce2ce0fe Mon Sep 17 00:00:00 2001 From: MauroFab Date: Mon, 20 Jul 2026 14:34:16 -0300 Subject: [PATCH] fix(prover): print per-table instruments report on the continuation path The `instruments` feature's per-table-kind timing report only printed on the monolithic `prove()` path. The continuation path (`prove_epoch`/`prove_global`) calls `Prover::multi_prove` directly; `multi_prove` still `store()`s the `MultiProveTiming` for each call, but nothing ever printed it, so `prove --continuations` with `--features instruments` produced no per-table report at all. Factor the mp-derived per-table block out of `instruments::print_report` into a shared `print_multi_prove()` helper, and add `print_epoch_report()` which drains the stored timing via `stark::instruments::take()` and formats it with that same helper. `prove_epoch`/`prove_global` now call it (feature-gated) right after each `multi_prove`, labeled "EPOCH N" / "GLOBAL". The monolithic path's output is unchanged (same statements, same order). All new code and both call sites are under `#[cfg(feature = "instruments")]`, so the feature-off build compiles identically with zero added overhead. --- prover/src/continuation.rs | 17 +- prover/src/instruments.rs | 368 ++++++++++++++++++++----------------- 2 files changed, 214 insertions(+), 171 deletions(-) diff --git a/prover/src/continuation.rs b/prover/src/continuation.rs index 169cd7278..db7fa1334 100644 --- a/prover/src/continuation.rs +++ b/prover/src/continuation.rs @@ -698,6 +698,12 @@ fn prove_epoch( ) .map_err(|e| Error::Prover(format!("{e:?}")))?; + // Emit this epoch's per-table timing report (same formatter as the monolithic + // path). `start.label` is 1-based (index + 1), so subtract 1 for the 0-based + // epoch number. Gated: zero overhead when the feature is off. + #[cfg(feature = "instruments")] + crate::instruments::print_epoch_report(&format!("EPOCH {}", start.label - 1)); + let l2g_root = proof .proofs .last() @@ -878,7 +884,7 @@ fn prove_global( pairs.push((air as AirRef, trace, &())); } - Prover::multi_prove( + let proof = Prover::multi_prove( pairs, &mut global_transcript( elf_bytes, @@ -890,7 +896,14 @@ fn prove_global( #[cfg(feature = "disk-spill")] stark::storage_mode::StorageMode::Ram, ) - .map_err(|e| Error::Prover(format!("{e:?}"))) + .map_err(|e| Error::Prover(format!("{e:?}")))?; + + // Cross-epoch global-memory proof: emit its per-table timing report too, + // using the shared monolithic formatter. Gated: zero overhead when off. + #[cfg(feature = "instruments")] + crate::instruments::print_epoch_report("GLOBAL"); + + Ok(proof) } #[allow(clippy::too_many_arguments)] diff --git a/prover/src/instruments.rs b/prover/src/instruments.rs index 0ea28273b..7f03f0a1d 100644 --- a/prover/src/instruments.rs +++ b/prover/src/instruments.rs @@ -71,175 +71,7 @@ pub fn print_report( row_top("AIR construction", air_construction, total); 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( - " Main LDE (fused GPU: LDE+Keccak+Merkle / CPU: LDE only)", - mp.round1_sub.main_lde, - total, - ); - row_sub( - " Main commit (Merkle, CPU only)", - mp.round1_sub.main_merkle, - total, - ); - row_sub(" Aux trace build (parallel)", mp.aux_build, total); - row_sub( - " LogUp fingerprint (CPU)", - mp.round1_sub.aux_fingerprint, - total, - ); - row_sub( - " LogUp batch invert (CPU)", - mp.round1_sub.aux_invert, - total, - ); - row_sub( - " LogUp term combine (CPU)", - mp.round1_sub.aux_term, - total, - ); - row_sub( - " LogUp accumulate scan (CPU)", - mp.round1_sub.aux_accumulate, - total, - ); - row_sub(" Aux trace commit", mp.aux_commit, total); - row_sub( - " Aux LDE (fused GPU: LDE+Keccak+Merkle / CPU: LDE only)", - mp.round1_sub.aux_lde, - total, - ); - row_sub( - " Aux commit (Merkle, CPU only)", - 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 - let mut merged: BTreeMap = BTreeMap::new(); - for (name, rows, dur, sub_ops) in &mp.table_timings { - let base = base_name(name).to_string(); - let entry = merged.entry(base).or_insert(MergedTable { - total_dur: Duration::ZERO, - total_rows: 0, - count: 0, - sub_ops: stark::instruments::TableSubOps::default(), - }); - entry.total_dur += *dur; - entry.total_rows += rows; - entry.count += 1; - entry.sub_ops.constraints += sub_ops.constraints; - entry.sub_ops.comp_decompose += sub_ops.comp_decompose; - entry.sub_ops.comp_commit += sub_ops.comp_commit; - entry.sub_ops.ood += sub_ops.ood; - entry.sub_ops.deep_comp += sub_ops.deep_comp; - entry.sub_ops.deep_extend += sub_ops.deep_extend; - entry.sub_ops.fri_commit += sub_ops.fri_commit; - entry.sub_ops.queries += sub_ops.queries; - } - - let mut sorted: Vec<_> = merged.into_iter().collect(); - sorted.sort_by(|a, b| b.1.total_dur.cmp(&a.1.total_dur)); - - let threshold = total.as_secs_f64() * 0.02; - let mut others_dur = Duration::ZERO; - let mut others_count = 0usize; - - for (name, t) in &sorted { - if t.total_dur.as_secs_f64() >= threshold { - let display_name = if t.count > 1 { - format!("{name} x{}", t.count) - } else { - name.clone() - }; - let label = format!(" {:<18} {:>6}", display_name, fmt_rows(t.total_rows),); - row_sub(&label, t.total_dur, total); - } else { - others_dur += t.total_dur; - others_count += 1; - } - } - if others_count > 0 { - let label = format!(" ({others_count} others)"); - row_sub(&label, others_dur, total); - } - - // Sub-operation totals across all tables - let mut total_constraints = Duration::ZERO; - let mut total_comp_decompose = Duration::ZERO; - let mut total_comp_commit = Duration::ZERO; - let mut total_ood = Duration::ZERO; - let mut total_deep_comp = Duration::ZERO; - let mut total_deep_extend = Duration::ZERO; - let mut total_fri_commit = Duration::ZERO; - let mut total_queries = Duration::ZERO; - for (_, t) in &sorted { - total_constraints += t.sub_ops.constraints; - total_comp_decompose += t.sub_ops.comp_decompose; - total_comp_commit += t.sub_ops.comp_commit; - total_ood += t.sub_ops.ood; - total_deep_comp += t.sub_ops.deep_comp; - total_deep_extend += t.sub_ops.deep_extend; - total_fri_commit += t.sub_ops.fri_commit; - total_queries += t.sub_ops.queries; - } - - let sub_ops_sum = total_constraints - + total_comp_decompose - + total_comp_commit - + total_ood - + total_deep_comp - + total_deep_extend - + total_fri_commit - + total_queries; - if sub_ops_sum > Duration::ZERO { - let mut sub_ops: Vec<(&str, Duration)> = vec![ - ("R2 evaluate", total_constraints), - ("R2 decompose_and_extend_d2", total_comp_decompose), - ("R2 commit_bit_reversed (comp-poly)", total_comp_commit), - ("R3 OOD evaluation", total_ood), - ("R4 deep_composition_poly_evals", total_deep_comp), - ("R4 interpolate+evaluate_fft", total_deep_extend), - ("R4 fri::commit_phase", total_fri_commit), - ("R4 queries & openings", total_queries), - ]; - sub_ops.sort_by(|a, b| b.1.cmp(&a.1)); - eprintln!( - " {}", - " \u{2500}\u{2500} sub-operation totals (all tables) \u{2500}\u{2500}", - ); - for (label, dur) in &sub_ops { - row_sub(&format!(" {label}"), *dur, total); - } - } - - // Cross-round totals: all FFT work and all Merkle work - let total_fft = mp.round1_sub.main_lde - + mp.round1_sub.aux_lde - + total_comp_decompose - + total_deep_extend; - let total_merkle = mp.round1_sub.main_merkle - + mp.round1_sub.aux_merkle - + total_comp_commit - + total_fri_commit; - eprintln!(); - eprintln!( - " {:<36} {:>7.2}s {:>5.1}%", - "Total FFT", - total_fft.as_secs_f64(), - pct(total_fft, total) - ); - eprintln!( - " {:<36} {:>7.2}s {:>5.1}%", - "Total Merkle", - total_merkle.as_secs_f64(), - pct(total_merkle, total) - ); + print_multi_prove(mp, total); } eprintln!(" {}", "─".repeat(58)); @@ -275,3 +107,201 @@ pub fn print_report( eprintln!(); } } + +/// Print the mp-derived per-table timing block (Round 1 sub-ops, Rounds 2-4, +/// merged per-table-kind timings, sub-operation totals, and cross-round FFT / +/// Merkle totals) to stderr, as a percentage of `total`. +/// +/// Shared by the monolithic [`print_report`] and the per-epoch continuation +/// [`print_epoch_report`], so both paths format identically from a single source. +fn print_multi_prove(mp: &stark::instruments::MultiProveTiming, total: Duration) { + 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( + " Main LDE (fused GPU: LDE+Keccak+Merkle / CPU: LDE only)", + mp.round1_sub.main_lde, + total, + ); + row_sub( + " Main commit (Merkle, CPU only)", + mp.round1_sub.main_merkle, + total, + ); + row_sub(" Aux trace build (parallel)", mp.aux_build, total); + row_sub( + " LogUp fingerprint (CPU)", + mp.round1_sub.aux_fingerprint, + total, + ); + row_sub( + " LogUp batch invert (CPU)", + mp.round1_sub.aux_invert, + total, + ); + row_sub( + " LogUp term combine (CPU)", + mp.round1_sub.aux_term, + total, + ); + row_sub( + " LogUp accumulate scan (CPU)", + mp.round1_sub.aux_accumulate, + total, + ); + row_sub(" Aux trace commit", mp.aux_commit, total); + row_sub( + " Aux LDE (fused GPU: LDE+Keccak+Merkle / CPU: LDE only)", + mp.round1_sub.aux_lde, + total, + ); + row_sub( + " Aux commit (Merkle, CPU only)", + 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 + let mut merged: BTreeMap = BTreeMap::new(); + for (name, rows, dur, sub_ops) in &mp.table_timings { + let base = base_name(name).to_string(); + let entry = merged.entry(base).or_insert(MergedTable { + total_dur: Duration::ZERO, + total_rows: 0, + count: 0, + sub_ops: stark::instruments::TableSubOps::default(), + }); + entry.total_dur += *dur; + entry.total_rows += rows; + entry.count += 1; + entry.sub_ops.constraints += sub_ops.constraints; + entry.sub_ops.comp_decompose += sub_ops.comp_decompose; + entry.sub_ops.comp_commit += sub_ops.comp_commit; + entry.sub_ops.ood += sub_ops.ood; + entry.sub_ops.deep_comp += sub_ops.deep_comp; + entry.sub_ops.deep_extend += sub_ops.deep_extend; + entry.sub_ops.fri_commit += sub_ops.fri_commit; + entry.sub_ops.queries += sub_ops.queries; + } + + let mut sorted: Vec<_> = merged.into_iter().collect(); + sorted.sort_by(|a, b| b.1.total_dur.cmp(&a.1.total_dur)); + + let threshold = total.as_secs_f64() * 0.02; + let mut others_dur = Duration::ZERO; + let mut others_count = 0usize; + + for (name, t) in &sorted { + if t.total_dur.as_secs_f64() >= threshold { + let display_name = if t.count > 1 { + format!("{name} x{}", t.count) + } else { + name.clone() + }; + let label = format!(" {:<18} {:>6}", display_name, fmt_rows(t.total_rows),); + row_sub(&label, t.total_dur, total); + } else { + others_dur += t.total_dur; + others_count += 1; + } + } + if others_count > 0 { + let label = format!(" ({others_count} others)"); + row_sub(&label, others_dur, total); + } + + // Sub-operation totals across all tables + let mut total_constraints = Duration::ZERO; + let mut total_comp_decompose = Duration::ZERO; + let mut total_comp_commit = Duration::ZERO; + let mut total_ood = Duration::ZERO; + let mut total_deep_comp = Duration::ZERO; + let mut total_deep_extend = Duration::ZERO; + let mut total_fri_commit = Duration::ZERO; + let mut total_queries = Duration::ZERO; + for (_, t) in &sorted { + total_constraints += t.sub_ops.constraints; + total_comp_decompose += t.sub_ops.comp_decompose; + total_comp_commit += t.sub_ops.comp_commit; + total_ood += t.sub_ops.ood; + total_deep_comp += t.sub_ops.deep_comp; + total_deep_extend += t.sub_ops.deep_extend; + total_fri_commit += t.sub_ops.fri_commit; + total_queries += t.sub_ops.queries; + } + + let sub_ops_sum = total_constraints + + total_comp_decompose + + total_comp_commit + + total_ood + + total_deep_comp + + total_deep_extend + + total_fri_commit + + total_queries; + if sub_ops_sum > Duration::ZERO { + let mut sub_ops: Vec<(&str, Duration)> = vec![ + ("R2 evaluate", total_constraints), + ("R2 decompose_and_extend_d2", total_comp_decompose), + ("R2 commit_bit_reversed (comp-poly)", total_comp_commit), + ("R3 OOD evaluation", total_ood), + ("R4 deep_composition_poly_evals", total_deep_comp), + ("R4 interpolate+evaluate_fft", total_deep_extend), + ("R4 fri::commit_phase", total_fri_commit), + ("R4 queries & openings", total_queries), + ]; + sub_ops.sort_by(|a, b| b.1.cmp(&a.1)); + eprintln!(" \u{2500}\u{2500} sub-operation totals (all tables) \u{2500}\u{2500}"); + for (label, dur) in &sub_ops { + row_sub(&format!(" {label}"), *dur, total); + } + } + + // Cross-round totals: all FFT work and all Merkle work + let total_fft = + mp.round1_sub.main_lde + mp.round1_sub.aux_lde + total_comp_decompose + total_deep_extend; + let total_merkle = + mp.round1_sub.main_merkle + mp.round1_sub.aux_merkle + total_comp_commit + total_fri_commit; + eprintln!(); + eprintln!( + " {:<36} {:>7.2}s {:>5.1}%", + "Total FFT", + total_fft.as_secs_f64(), + pct(total_fft, total) + ); + eprintln!( + " {:<36} {:>7.2}s {:>5.1}%", + "Total Merkle", + total_merkle.as_secs_f64(), + pct(total_merkle, total) + ); +} + +/// Print the per-table timing report for a single continuation-epoch prove. +/// +/// Feature-gated; called right after each `multi_prove` on the continuation +/// path (`prove_epoch`/`prove_global`). Drains the timing `store`d by that +/// `multi_prove` via [`stark::instruments::take`] and reuses the shared +/// [`print_multi_prove`] formatter, so per-epoch output matches the monolithic +/// per-table report (minus the execute/trace/AIR phases, which are not measured +/// per epoch). `total` is derived from the mp phase sum, so percentages are +/// relative to this epoch's proving time. `label` names the section, e.g. +/// "EPOCH 0" / "GLOBAL". +pub fn print_epoch_report(label: &str) { + let Some(mp) = stark::instruments::take() else { + return; + }; + let round1 = mp.main_commits + mp.aux_build + mp.aux_commit; + let total = mp.prepass + round1 + mp.rounds_2_4; + + eprintln!(); + eprintln!("=== {label} ==="); + eprintln!(" {:<36} {:>8} {:>5}", "Phase", "Wall", "%",); + eprintln!(" {}", "─".repeat(58)); + print_multi_prove(&mp, total); + eprintln!(" {}", "─".repeat(58)); + eprintln!(" {:<36} {:>7.2}s", "TOTAL", total.as_secs_f64()); + eprintln!(); +}