From ef648a59b916dec7370f4d0a6f1d82328241e531 Mon Sep 17 00:00:00 2001 From: gabrielbosio Date: Mon, 13 Apr 2026 12:25:12 -0300 Subject: [PATCH 1/7] Cache LDE from Round 1 to skip recomputation in Rounds 2-4 --- crypto/stark/src/prover.rs | 164 ++++++++++++++-------- scripts/bench_timing_profile.sh | 237 ++++++++++++++++++++++++++++++++ 2 files changed, 341 insertions(+), 60 deletions(-) create mode 100755 scripts/bench_timing_profile.sh diff --git a/crypto/stark/src/prover.rs b/crypto/stark/src/prover.rs index 50af95525..93e4f8c5a 100644 --- a/crypto/stark/src/prover.rs +++ b/crypto/stark/src/prover.rs @@ -124,11 +124,13 @@ where precomputed_tree: Option>>, precomputed_root: Option, num_precomputed_cols: usize, + /// Cached main LDE columns from Phase A, used in Phase D to skip recomputation. + cached_main_lde: Vec>>, } -/// Metadata from Round 1 commitments — stores Merkle trees and roots, but not LDE evaluations. -/// LDE evaluations are recomputed from the trace in Rounds 2-4; Merkle trees are retained -/// to avoid expensive Keccak re-hashing. +/// Metadata from Round 1 commitments — stores Merkle trees, roots, and cached LDE evaluations. +/// LDE evaluations are cached from Phase A/C and consumed in Phase D (Rounds 2-4), +/// eliminating the expensive recomputation (iFFT + coset shift + FFT per column). pub struct Round1Metadata where Field: IsFFTField + IsSubFieldOf, @@ -155,6 +157,10 @@ where rap_challenges: Vec>, /// Bus interaction public inputs (initial and final aux column values). bus_public_inputs: Option>, + /// Cached main LDE columns from Phase A (consumed by Phase D). + cached_main_lde: Vec>>, + /// Cached aux LDE columns from Phase C (consumed by Phase D). + cached_aux_lde: Vec>>, } /// Pre-computed twiddle factors and coset weights for a given domain size. @@ -1515,33 +1521,27 @@ pub trait IsStarkProver< let mut twiddle_caches: Vec> = Vec::with_capacity(num_airs); let mut max_main_cols = 0usize; let mut max_aux_cols = 0usize; - let mut max_lde_size = 0usize; for (air, trace, _pub_inputs) in &*air_trace_pairs { let trace_length = trace.num_rows(); let domain = new_domain(*air, trace_length); - - let lde_size = domain.interpolation_domain_size * domain.blowup_factor; let twiddles = LdeTwiddles::new(&domain); max_main_cols = max_main_cols.max(trace.num_main_columns); max_aux_cols = max_aux_cols.max(air.num_auxiliary_rap_columns()); - max_lde_size = max_lde_size.max(lde_size); domains.push(domain); twiddle_caches.push(twiddles); } // Allocate K independent LDE column buffer pool sets for parallel table processing. + // Buffers start empty; each table allocates on demand during extract+expand. + // After commit, LDE columns are moved into per-table cache (not reused across chunks). let k = table_parallelism().min(num_airs).max(1); let mut pool_sets: Vec> = (0..k) .map(|_| PoolSet { - main: (0..max_main_cols) - .map(|_| Vec::with_capacity(max_lde_size)) - .collect(), - aux: (0..max_aux_cols) - .map(|_| Vec::with_capacity(max_lde_size)) - .collect(), + main: (0..max_main_cols).map(|_| Vec::new()).collect(), + aux: (0..max_aux_cols).map(|_| Vec::new()).collect(), }) .collect(); @@ -1575,7 +1575,7 @@ pub trait IsStarkProver< let domain = &domains[idx]; let twiddles = &twiddle_caches[idx]; - if air.is_preprocessed() { + let result = if air.is_preprocessed() { Self::commit_preprocessed_trace( *trace, domain, @@ -1586,13 +1586,23 @@ pub trait IsStarkProver< ) } else { Self::commit_main_trace(*trace, domain, twiddles, &mut pool.main) - } + }?; + + // Cache main LDE columns from pool (zero-copy move). + // Pool slots become empty; next chunk allocates on demand. + let num_main_cols = trace.num_main_columns; + let cached_main: Vec<_> = pool.main[..num_main_cols] + .iter_mut() + .map(std::mem::take) + .collect(); + + Ok((result, cached_main)) }) .collect(); // Sequential: append roots to shared transcript (Fiat-Shamir ordering) for result in chunk_results { - let (tree, root, pre_tree, pre_root, n_pre) = result?; + let ((tree, root, pre_tree, pre_root, n_pre), cached_main) = result?; if let Some(ref pre_r) = pre_root { transcript.append_bytes(pre_r); } @@ -1603,6 +1613,7 @@ pub trait IsStarkProver< precomputed_tree: pre_tree.map(Arc::new), precomputed_root: pre_root, num_precomputed_cols: n_pre, + cached_main_lde: cached_main, }); } } @@ -1677,6 +1688,7 @@ pub trait IsStarkProver< let mut aux_results: Vec<( Option>>, Option, + Vec>>, )> = Vec::with_capacity(num_airs); for chunk_start in (0..num_airs).step_by(k) { @@ -1715,41 +1727,50 @@ pub trait IsStarkProver< .ok_or(ProvingError::EmptyCommitment)?; #[cfg(feature = "instruments")] crate::instruments::accum_r1_aux(aux_lde_dur, t_sub.elapsed()); - Ok((Some(Arc::new(tree)), Some(root))) + + // Cache aux LDE columns from pool (zero-copy move). + let cached_aux: Vec<_> = pool.aux[..num_aux_cols] + .iter_mut() + .map(std::mem::take) + .collect(); + + Ok((Some(Arc::new(tree)), Some(root), cached_aux)) } else { - Ok((None, None)) + Ok((None, None, Vec::new())) } }) .collect(); // Sequential: append aux roots to forked transcripts for (j, result) in chunk_aux.into_iter().enumerate() { - let (aux_tree, aux_root) = result?; + let (aux_tree, aux_root, cached_aux) = result?; if let Some(ref root) = aux_root { table_transcripts[chunk_start + j].append_bytes(root); } - aux_results.push((aux_tree, aux_root)); + aux_results.push((aux_tree, aux_root, cached_aux)); } } // Build metadata sequentially from main_commits + aux_results + bus_inputs let mut metadatas: Vec> = Vec::with_capacity(num_airs); - for ((main_commit, (aux_tree, aux_root)), bus_public_inputs) in main_commits + for ((main_commit, (aux_tree, aux_root, cached_aux)), bus_public_inputs) in main_commits .into_iter() .zip(aux_results) .zip(bus_inputs_vec) { metadatas.push(Round1Metadata { - main_merkle_tree: Arc::clone(&main_commit.main_tree), + main_merkle_tree: main_commit.main_tree, main_merkle_root: main_commit.main_root, - precomputed_merkle_tree: main_commit.precomputed_tree.as_ref().map(Arc::clone), + precomputed_merkle_tree: main_commit.precomputed_tree, precomputed_merkle_root: main_commit.precomputed_root, num_precomputed_cols: main_commit.num_precomputed_cols, aux_merkle_tree: aux_tree, aux_merkle_root: aux_root, rap_challenges: lookup_challenges.clone(), bus_public_inputs, + cached_main_lde: main_commit.cached_main_lde, + cached_aux_lde: cached_aux, }); } @@ -1769,11 +1790,15 @@ pub trait IsStarkProver< ); } + // Pool sets are no longer needed — Phase D uses cached LDE from metadata. + drop(pool_sets); + // ===================================================================== // Rounds 2-4: Parallel per-table proving in chunks of K // ===================================================================== - // Each chunk of K tables is processed in parallel. Each worker gets its - // own pool set and transcript fork. Pool sets are reused across chunks. + // Each chunk of K tables is processed in parallel. Cached LDE columns + // from Phase A/C are consumed here (zero-copy move), eliminating the + // expensive reconstruct_round1 recomputation. #[cfg(feature = "instruments")] let phase_start = Instant::now(); @@ -1788,45 +1813,79 @@ pub trait IsStarkProver< let mut proofs = Vec::with_capacity(num_airs); for chunk_start in (0..num_airs).step_by(k) { let chunk_end = (chunk_start + k).min(num_airs); - let chunk_size = chunk_end - chunk_start; + let chunk_metadatas = &mut metadatas[chunk_start..chunk_end]; let chunk_transcripts = &mut table_transcripts[chunk_start..chunk_end]; #[cfg(feature = "parallel")] - let iter = pool_sets[..chunk_size] + let iter = chunk_metadatas .par_iter_mut() .zip(chunk_transcripts.par_iter_mut()) .enumerate(); #[cfg(not(feature = "parallel"))] - let iter = pool_sets[..chunk_size] + let iter = chunk_metadatas .iter_mut() .zip(chunk_transcripts.iter_mut()) .enumerate(); let chunk_results: Vec> = iter - .map(|(j, (pool, table_transcript))| { + .map(|(j, (metadata, table_transcript))| { let idx = chunk_start + j; let (air, trace, pub_inputs) = &air_trace_pairs[idx]; - let metadata = &metadatas[idx]; + let _ = trace; // used by instruments let domain = &domains[idx]; - let twiddles = &twiddle_caches[idx]; #[cfg(feature = "instruments")] let table_start = Instant::now(); - #[cfg(feature = "instruments")] - let lde_start = Instant::now(); - let round_1_result = Self::reconstruct_round1( - *air, - *trace, - domain, - metadata, - twiddles, - &mut pool.main, - &mut pool.aux, - )?; - #[cfg(feature = "instruments")] - let lde_dur = lde_start.elapsed(); + // Build Round1 from cached LDE (zero-copy move, no recomputation). + let cached_main = std::mem::take(&mut metadata.cached_main_lde); + let cached_aux = std::mem::take(&mut metadata.cached_aux_lde); + + let lde_trace = LDETraceTable::from_columns( + cached_main, + cached_aux, + air.step_size(), + domain.blowup_factor, + ); + + let main = Round1CommitmentData:: { + lde_trace_merkle_tree: Arc::clone(&metadata.main_merkle_tree), + lde_trace_merkle_root: metadata.main_merkle_root, + precomputed_merkle_tree: metadata + .precomputed_merkle_tree + .as_ref() + .map(Arc::clone), + precomputed_merkle_root: metadata.precomputed_merkle_root, + num_precomputed_cols: metadata.num_precomputed_cols, + }; + + let aux = if air.has_aux_trace() { + Some(Round1CommitmentData:: { + lde_trace_merkle_tree: Arc::clone( + metadata + .aux_merkle_tree + .as_ref() + .expect("aux tree must exist when has_trace_interaction"), + ), + lde_trace_merkle_root: metadata + .aux_merkle_root + .expect("aux root must exist when has_trace_interaction"), + precomputed_merkle_tree: None, + precomputed_merkle_root: None, + num_precomputed_cols: 0, + }) + } else { + None + }; + + let round_1_result = Round1 { + lde_trace, + main, + aux, + rap_challenges: metadata.rap_challenges.clone(), + bus_public_inputs: metadata.bus_public_inputs.clone(), + }; if let Some(ref bpi) = round_1_result.bus_public_inputs { table_transcript.append_field_element(&bpi.table_contribution); @@ -1840,15 +1899,9 @@ pub trait IsStarkProver< domain, )?; - // Collect per-table sub-op timing via TLS. - // Both the store (inside prove_rounds_2_to_4) and this take run on the - // same rayon worker thread, so sub-ops are valid in both sequential and - // parallel mode. #[cfg(feature = "instruments")] let table_timing = { - let mut sub_ops = - crate::instruments::take_round_sub_ops().unwrap_or_default(); - sub_ops.trace_lde += lde_dur; + let sub_ops = crate::instruments::take_round_sub_ops().unwrap_or_default(); ( air.name().to_string(), trace.num_rows(), @@ -1857,15 +1910,6 @@ pub trait IsStarkProver< ) }; - // Return column Vecs to pool (zero-copy move back) - let (main_cols, aux_cols) = round_1_result.lde_trace.into_columns(); - for (slot, col) in pool.main.iter_mut().zip(main_cols) { - *slot = col; - } - for (slot, col) in pool.aux.iter_mut().zip(aux_cols) { - *slot = col; - } - #[cfg(feature = "instruments")] return Ok((proof, table_timing)); #[cfg(not(feature = "instruments"))] diff --git a/scripts/bench_timing_profile.sh b/scripts/bench_timing_profile.sh new file mode 100755 index 000000000..671118662 --- /dev/null +++ b/scripts/bench_timing_profile.sh @@ -0,0 +1,237 @@ +#!/bin/bash +# Per-phase timing profile across program sizes. +# Shows how each proving phase and sub-operation scales with program size. +# +# Usage: bench_timing_profile.sh [--no-build] [--programs "1M 2M 4M 8M"] +# +# Requires: instruments feature. +# Timing is deterministic with instruments, so 1 run per size is enough. + +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +ROOT_DIR="$(cd "$SCRIPT_DIR/.." && pwd)" +TMP_DIR="/tmp/bench_timing_profile" +ELF_DIR="$ROOT_DIR/executor/program_artifacts/asm" + +GREEN='\033[0;32m' +BOLD='\033[1m' +NC='\033[0m' + +BUILD=true +PROGRAMS="500k 1M 2M 4M" + +while [[ $# -gt 0 ]]; do + case $1 in + --no-build) BUILD=false; shift ;; + --programs) PROGRAMS="$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: $1" >&2; exit 1 ;; + esac +} + +rm -rf "$TMP_DIR" && mkdir -p "$TMP_DIR" + +if $BUILD; then + echo -e "${GREEN}Building CLI with instruments...${NC}" + cargo build --release -p cli --features instruments \ + --manifest-path "$ROOT_DIR/Cargo.toml" 2>&1 | tail -1 +fi +CLI="$ROOT_DIR/target/release/cli" + +# Parse instruments stderr into key=value pairs. +# Uses occurrence counting to disambiguate expand_pool_to_lde and commit (Merkle). +parse_timing() { + awk ' + BEGIN { lde_n = 0; merkle_n = 0 } + 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 "execute=" v } + /^ Trace build/ { v = secs(); if (v) print "trace_build=" v } + /^ AIR construction/ { v = secs(); if (v) print "air=" v } + /^ Pre-pass/ { v = secs(); if (v) print "prepass=" v } + /^ Round 1 / { v = secs(); if (v) print "round1=" v } + /Main trace commits/ { v = secs(); if (v) print "main_commits=" v } + /Aux trace build/ { v = secs(); if (v) print "aux_build=" v } + /Aux trace commit/ { v = secs(); if (v) print "aux_commit=" v } + /Rounds 2/ { v = secs(); if (v) print "rounds24=" v } + /expand_pool_to_lde/ && !/R1 expand/ { + lde_n++; v = secs() + if (v && lde_n == 1) print "main_lde=" v + if (v && lde_n == 2) print "aux_lde=" v + } + /commit \(Merkle\)/ { + merkle_n++; v = secs() + if (v && merkle_n == 1) print "main_merkle=" v + if (v && merkle_n == 2) print "aux_merkle=" v + } + /R1 expand_pool_to_lde/ { v = secs(); if (v) print "r1_lde=" v } + /R2 evaluate/ { v = secs(); if (v) print "r2_evaluate=" v } + /R2 decompose/ { v = secs(); if (v) print "r2_decompose=" v } + /R2 commit_comp/ { v = secs(); if (v) print "r2_commit_comp=" v } + /R3 OOD/ { v = secs(); if (v) print "r3_ood=" v } + /R4 deep_comp/ { v = secs(); if (v) print "r4_deep_comp=" v } + /R4 interpolate/ { v = secs(); if (v) print "r4_interp=" v } + /R4 fri::commit/ { v = secs(); if (v) print "r4_fri_commit=" v } + /R4 queries/ { v = secs(); if (v) print "r4_queries=" v } + /^ Total FFT/ { v = secs(); if (v) print "total_fft=" v } + /^ Total Merkle/ { v = secs(); if (v) print "total_merkle=" v } + /^ TOTAL / { v = secs(); if (v) print "total=" v } + ' "$1" +} + +for size in $PROGRAMS; 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}...${NC}" + + STDERR="$TMP_DIR/${size}_stderr.txt" + "$CLI" prove "$ELF" -o "$TMP_DIR/proof.bin" 2>"$STDERR" >/dev/null + rm -f "$TMP_DIR/proof.bin" + + echo "steps=$steps" > "$TMP_DIR/${size}_data.txt" + parse_timing "$STDERR" >> "$TMP_DIR/${size}_data.txt" +done + +# --- Display --------------------------------------------------------------- + +get_val() { + grep "^${2}=" "$1" 2>/dev/null | cut -d= -f2 +} + +print_section() { + local title=$1; shift + + echo "" + echo -e " ${BOLD}${title}${NC}" + printf " %-30s" "" + for size in $PROGRAMS; do printf " %9s" "$size"; done + echo "" + printf " %-30s" "" + for size in $PROGRAMS; do printf " %9s" "─────────"; done + echo "" + + for spec in "$@"; do + local label="${spec%%:*}" + local key="${spec#*:}" + printf " %-30s" "$label" + for size in $PROGRAMS; do + local DATA="$TMP_DIR/${size}_data.txt" + local val + val=$(get_val "$DATA" "$key") + if [ -n "$val" ]; then + printf " %8.2fs" "$val" + else + printf " %9s" "-" + fi + done + echo "" + done +} + +echo "" +echo -e "${BOLD}=== TIMING PROFILE ACROSS SIZES ===${NC}" + +TABLE_K="${TABLE_PARALLELISM:-default (cores/3)}" +echo -e " TABLE_PARALLELISM=$TABLE_K" + +print_section "Phase (wall time)" \ + "TOTAL:total" \ + "Execute:execute" \ + "Trace build:trace_build" \ + "AIR construction:air" \ + "Pre-pass:prepass" \ + "Round 1:round1" \ + "Rounds 2-4:rounds24" + +print_section "Round 1 breakdown (CPU sum)" \ + "Main LDE:main_lde" \ + "Main Merkle:main_merkle" \ + "Aux build (wall):aux_build" \ + "Aux LDE:aux_lde" \ + "Aux Merkle:aux_merkle" + +print_section "Rounds 2-4 sub-ops (CPU sum)" \ + "R1 LDE (reconstruct):r1_lde" \ + "R2 decompose+extend:r2_decompose" \ + "R2 evaluate:r2_evaluate" \ + "R2 commit comp poly:r2_commit_comp" \ + "R3 OOD evaluation:r3_ood" \ + "R4 deep comp evals:r4_deep_comp" \ + "R4 interpolate+evaluate:r4_interp" \ + "R4 FRI commit:r4_fri_commit" \ + "R4 queries & openings:r4_queries" + +print_section "Cross-round totals (CPU sum)" \ + "Total FFT:total_fft" \ + "Total Merkle:total_merkle" + +# --- Growth rate ----------------------------------------------------------- + +echo "" +echo -e "${BOLD}=== GROWTH RATE (per 1M steps) ===${NC}" +echo "" + +for spec in \ + "TOTAL:total" \ + "Trace build:trace_build" \ + "Round 1:round1" \ + "Rounds 2-4:rounds24" \ + "Main LDE:main_lde" \ + "Aux LDE:aux_lde" \ + "R1 LDE (reconstruct):r1_lde" \ + "R2 decompose+extend:r2_decompose" \ + "R3 OOD evaluation:r3_ood" \ + "Total FFT:total_fft" \ + "Total Merkle:total_merkle" +do + label="${spec%%:*}" + key="${spec#*:}" + + PAIRS="" + for size in $PROGRAMS; do + DATA="$TMP_DIR/${size}_data.txt" + [ -f "$DATA" ] || continue + steps=$(get_val "$DATA" "steps") + val=$(get_val "$DATA" "$key") + [ -z "$val" ] && continue + steps_m=$(awk "BEGIN {printf \"%.2f\", $steps / 1000000}") + PAIRS="$PAIRS $steps_m $val" + done + + echo "$PAIRS" | awk -v label="$label" '{ + n = NF / 2 + if (n < 2) { printf " %-30s (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 + printf " %-30s %+.2fs/M steps (base: %.2fs, R\xc2\xb2=%.3f)\n", label, b, a, r2 + }' +done + +echo "" +echo "Raw data in $TMP_DIR/" From fe4f3409aa54ebdb614cd382c35d7bafae645de9 Mon Sep 17 00:00:00 2001 From: gabrielbosio Date: Mon, 13 Apr 2026 18:18:58 -0300 Subject: [PATCH 2/7] Remove timing profile script (moved to separate PR) --- scripts/bench_timing_profile.sh | 237 -------------------------------- 1 file changed, 237 deletions(-) delete mode 100755 scripts/bench_timing_profile.sh diff --git a/scripts/bench_timing_profile.sh b/scripts/bench_timing_profile.sh deleted file mode 100755 index 671118662..000000000 --- a/scripts/bench_timing_profile.sh +++ /dev/null @@ -1,237 +0,0 @@ -#!/bin/bash -# Per-phase timing profile across program sizes. -# Shows how each proving phase and sub-operation scales with program size. -# -# Usage: bench_timing_profile.sh [--no-build] [--programs "1M 2M 4M 8M"] -# -# Requires: instruments feature. -# Timing is deterministic with instruments, so 1 run per size is enough. - -set -euo pipefail - -SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" -ROOT_DIR="$(cd "$SCRIPT_DIR/.." && pwd)" -TMP_DIR="/tmp/bench_timing_profile" -ELF_DIR="$ROOT_DIR/executor/program_artifacts/asm" - -GREEN='\033[0;32m' -BOLD='\033[1m' -NC='\033[0m' - -BUILD=true -PROGRAMS="500k 1M 2M 4M" - -while [[ $# -gt 0 ]]; do - case $1 in - --no-build) BUILD=false; shift ;; - --programs) PROGRAMS="$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: $1" >&2; exit 1 ;; - esac -} - -rm -rf "$TMP_DIR" && mkdir -p "$TMP_DIR" - -if $BUILD; then - echo -e "${GREEN}Building CLI with instruments...${NC}" - cargo build --release -p cli --features instruments \ - --manifest-path "$ROOT_DIR/Cargo.toml" 2>&1 | tail -1 -fi -CLI="$ROOT_DIR/target/release/cli" - -# Parse instruments stderr into key=value pairs. -# Uses occurrence counting to disambiguate expand_pool_to_lde and commit (Merkle). -parse_timing() { - awk ' - BEGIN { lde_n = 0; merkle_n = 0 } - 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 "execute=" v } - /^ Trace build/ { v = secs(); if (v) print "trace_build=" v } - /^ AIR construction/ { v = secs(); if (v) print "air=" v } - /^ Pre-pass/ { v = secs(); if (v) print "prepass=" v } - /^ Round 1 / { v = secs(); if (v) print "round1=" v } - /Main trace commits/ { v = secs(); if (v) print "main_commits=" v } - /Aux trace build/ { v = secs(); if (v) print "aux_build=" v } - /Aux trace commit/ { v = secs(); if (v) print "aux_commit=" v } - /Rounds 2/ { v = secs(); if (v) print "rounds24=" v } - /expand_pool_to_lde/ && !/R1 expand/ { - lde_n++; v = secs() - if (v && lde_n == 1) print "main_lde=" v - if (v && lde_n == 2) print "aux_lde=" v - } - /commit \(Merkle\)/ { - merkle_n++; v = secs() - if (v && merkle_n == 1) print "main_merkle=" v - if (v && merkle_n == 2) print "aux_merkle=" v - } - /R1 expand_pool_to_lde/ { v = secs(); if (v) print "r1_lde=" v } - /R2 evaluate/ { v = secs(); if (v) print "r2_evaluate=" v } - /R2 decompose/ { v = secs(); if (v) print "r2_decompose=" v } - /R2 commit_comp/ { v = secs(); if (v) print "r2_commit_comp=" v } - /R3 OOD/ { v = secs(); if (v) print "r3_ood=" v } - /R4 deep_comp/ { v = secs(); if (v) print "r4_deep_comp=" v } - /R4 interpolate/ { v = secs(); if (v) print "r4_interp=" v } - /R4 fri::commit/ { v = secs(); if (v) print "r4_fri_commit=" v } - /R4 queries/ { v = secs(); if (v) print "r4_queries=" v } - /^ Total FFT/ { v = secs(); if (v) print "total_fft=" v } - /^ Total Merkle/ { v = secs(); if (v) print "total_merkle=" v } - /^ TOTAL / { v = secs(); if (v) print "total=" v } - ' "$1" -} - -for size in $PROGRAMS; 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}...${NC}" - - STDERR="$TMP_DIR/${size}_stderr.txt" - "$CLI" prove "$ELF" -o "$TMP_DIR/proof.bin" 2>"$STDERR" >/dev/null - rm -f "$TMP_DIR/proof.bin" - - echo "steps=$steps" > "$TMP_DIR/${size}_data.txt" - parse_timing "$STDERR" >> "$TMP_DIR/${size}_data.txt" -done - -# --- Display --------------------------------------------------------------- - -get_val() { - grep "^${2}=" "$1" 2>/dev/null | cut -d= -f2 -} - -print_section() { - local title=$1; shift - - echo "" - echo -e " ${BOLD}${title}${NC}" - printf " %-30s" "" - for size in $PROGRAMS; do printf " %9s" "$size"; done - echo "" - printf " %-30s" "" - for size in $PROGRAMS; do printf " %9s" "─────────"; done - echo "" - - for spec in "$@"; do - local label="${spec%%:*}" - local key="${spec#*:}" - printf " %-30s" "$label" - for size in $PROGRAMS; do - local DATA="$TMP_DIR/${size}_data.txt" - local val - val=$(get_val "$DATA" "$key") - if [ -n "$val" ]; then - printf " %8.2fs" "$val" - else - printf " %9s" "-" - fi - done - echo "" - done -} - -echo "" -echo -e "${BOLD}=== TIMING PROFILE ACROSS SIZES ===${NC}" - -TABLE_K="${TABLE_PARALLELISM:-default (cores/3)}" -echo -e " TABLE_PARALLELISM=$TABLE_K" - -print_section "Phase (wall time)" \ - "TOTAL:total" \ - "Execute:execute" \ - "Trace build:trace_build" \ - "AIR construction:air" \ - "Pre-pass:prepass" \ - "Round 1:round1" \ - "Rounds 2-4:rounds24" - -print_section "Round 1 breakdown (CPU sum)" \ - "Main LDE:main_lde" \ - "Main Merkle:main_merkle" \ - "Aux build (wall):aux_build" \ - "Aux LDE:aux_lde" \ - "Aux Merkle:aux_merkle" - -print_section "Rounds 2-4 sub-ops (CPU sum)" \ - "R1 LDE (reconstruct):r1_lde" \ - "R2 decompose+extend:r2_decompose" \ - "R2 evaluate:r2_evaluate" \ - "R2 commit comp poly:r2_commit_comp" \ - "R3 OOD evaluation:r3_ood" \ - "R4 deep comp evals:r4_deep_comp" \ - "R4 interpolate+evaluate:r4_interp" \ - "R4 FRI commit:r4_fri_commit" \ - "R4 queries & openings:r4_queries" - -print_section "Cross-round totals (CPU sum)" \ - "Total FFT:total_fft" \ - "Total Merkle:total_merkle" - -# --- Growth rate ----------------------------------------------------------- - -echo "" -echo -e "${BOLD}=== GROWTH RATE (per 1M steps) ===${NC}" -echo "" - -for spec in \ - "TOTAL:total" \ - "Trace build:trace_build" \ - "Round 1:round1" \ - "Rounds 2-4:rounds24" \ - "Main LDE:main_lde" \ - "Aux LDE:aux_lde" \ - "R1 LDE (reconstruct):r1_lde" \ - "R2 decompose+extend:r2_decompose" \ - "R3 OOD evaluation:r3_ood" \ - "Total FFT:total_fft" \ - "Total Merkle:total_merkle" -do - label="${spec%%:*}" - key="${spec#*:}" - - PAIRS="" - for size in $PROGRAMS; do - DATA="$TMP_DIR/${size}_data.txt" - [ -f "$DATA" ] || continue - steps=$(get_val "$DATA" "steps") - val=$(get_val "$DATA" "$key") - [ -z "$val" ] && continue - steps_m=$(awk "BEGIN {printf \"%.2f\", $steps / 1000000}") - PAIRS="$PAIRS $steps_m $val" - done - - echo "$PAIRS" | awk -v label="$label" '{ - n = NF / 2 - if (n < 2) { printf " %-30s (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 - printf " %-30s %+.2fs/M steps (base: %.2fs, R\xc2\xb2=%.3f)\n", label, b, a, r2 - }' -done - -echo "" -echo "Raw data in $TMP_DIR/" From 9135f118c20ae0bf6074691852dc1b60286c4ce3 Mon Sep 17 00:00:00 2001 From: gabrielbosio Date: Mon, 13 Apr 2026 18:22:17 -0300 Subject: [PATCH 3/7] Document memory trade-off, fix expect messages --- crypto/stark/src/prover.rs | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/crypto/stark/src/prover.rs b/crypto/stark/src/prover.rs index 93e4f8c5a..b4978b788 100644 --- a/crypto/stark/src/prover.rs +++ b/crypto/stark/src/prover.rs @@ -131,6 +131,10 @@ where /// Metadata from Round 1 commitments — stores Merkle trees, roots, and cached LDE evaluations. /// LDE evaluations are cached from Phase A/C and consumed in Phase D (Rounds 2-4), /// eliminating the expensive recomputation (iFFT + coset shift + FFT per column). +/// +/// Memory trade-off: all N tables' LDE columns are live simultaneously between Phase A/C +/// and Phase D (O(N × cols × lde_size)), but pool pre-allocation is removed, so peak heap +/// is lower overall (~70 GB vs ~220 GB at 8M cycles). pub struct Round1Metadata where Field: IsFFTField + IsSubFieldOf, @@ -598,11 +602,11 @@ pub trait IsStarkProver< metadata .aux_merkle_tree .as_ref() - .expect("aux tree must exist when has_trace_interaction"), + .expect("aux tree must exist when has_aux_trace"), ), lde_trace_merkle_root: metadata .aux_merkle_root - .expect("aux root must exist when has_trace_interaction"), + .expect("aux root must exist when has_aux_trace"), precomputed_merkle_tree: None, precomputed_merkle_root: None, num_precomputed_cols: 0, @@ -1866,11 +1870,11 @@ pub trait IsStarkProver< metadata .aux_merkle_tree .as_ref() - .expect("aux tree must exist when has_trace_interaction"), + .expect("aux tree must exist when has_aux_trace"), ), lde_trace_merkle_root: metadata .aux_merkle_root - .expect("aux root must exist when has_trace_interaction"), + .expect("aux root must exist when has_aux_trace"), precomputed_merkle_tree: None, precomputed_merkle_root: None, num_precomputed_cols: 0, From c1146e7d4dc8b0aaee2715ac0b033d72640f7d13 Mon Sep 17 00:00:00 2001 From: gabrielbosio Date: Mon, 13 Apr 2026 18:52:12 -0300 Subject: [PATCH 4/7] Fix stale trace_lde comment in prove_rounds_2_to_4 --- crypto/stark/src/prover.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/crypto/stark/src/prover.rs b/crypto/stark/src/prover.rs index b4978b788..b053ee0c6 100644 --- a/crypto/stark/src/prover.rs +++ b/crypto/stark/src/prover.rs @@ -2086,7 +2086,7 @@ pub trait IsStarkProver< let (r4_fft, r4_merkle, r4_deep_comp, r4_queries) = crate::instruments::take_r4_sub().unwrap_or((zero, zero, zero, zero)); crate::instruments::store_round_sub_ops(crate::instruments::TableSubOps { - trace_lde: std::time::Duration::ZERO, // added by caller from lde_dur + trace_lde: std::time::Duration::ZERO, // LDE timing in Phase A/C (accum_r1_main/aux) constraints: r2_constraints, comp_decompose: r2_fft, comp_commit: r2_merkle, From f8342ffe1a0f94b2aa4c9aefa81884acd326930b Mon Sep 17 00:00:00 2001 From: gabrielbosio Date: Tue, 14 Apr 2026 16:52:24 -0300 Subject: [PATCH 5/7] Remove dead trace_lde telemetry field --- crypto/stark/src/instruments.rs | 2 -- crypto/stark/src/prover.rs | 1 - prover/src/instruments.rs | 8 +------- 3 files changed, 1 insertion(+), 10 deletions(-) diff --git a/crypto/stark/src/instruments.rs b/crypto/stark/src/instruments.rs index 3c13ef2e1..dd566b1ae 100644 --- a/crypto/stark/src/instruments.rs +++ b/crypto/stark/src/instruments.rs @@ -5,8 +5,6 @@ use std::time::Duration; /// Sub-operation timing breakdown for a single table in Rounds 2-4. #[derive(Clone, Debug, Default)] pub struct TableSubOps { - /// reconstruct_round1 (expand_pool_to_lde) - pub trace_lde: Duration, /// evaluator.evaluate() pub constraints: Duration, /// decompose_and_extend_d2 diff --git a/crypto/stark/src/prover.rs b/crypto/stark/src/prover.rs index b053ee0c6..86e760dbb 100644 --- a/crypto/stark/src/prover.rs +++ b/crypto/stark/src/prover.rs @@ -2086,7 +2086,6 @@ pub trait IsStarkProver< let (r4_fft, r4_merkle, r4_deep_comp, r4_queries) = crate::instruments::take_r4_sub().unwrap_or((zero, zero, zero, zero)); crate::instruments::store_round_sub_ops(crate::instruments::TableSubOps { - trace_lde: std::time::Duration::ZERO, // LDE timing in Phase A/C (accum_r1_main/aux) constraints: r2_constraints, comp_decompose: r2_fft, comp_commit: r2_merkle, diff --git a/prover/src/instruments.rs b/prover/src/instruments.rs index e3db38b95..ac517e85b 100644 --- a/prover/src/instruments.rs +++ b/prover/src/instruments.rs @@ -96,7 +96,6 @@ pub fn print_report( entry.total_dur += *dur; entry.total_rows += rows; entry.count += 1; - entry.sub_ops.trace_lde += sub_ops.trace_lde; 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; @@ -134,7 +133,6 @@ pub fn print_report( } // Sub-operation totals across all tables - let mut total_trace_lde = Duration::ZERO; let mut total_constraints = Duration::ZERO; let mut total_comp_decompose = Duration::ZERO; let mut total_comp_commit = Duration::ZERO; @@ -144,7 +142,6 @@ pub fn print_report( let mut total_fri_commit = Duration::ZERO; let mut total_queries = Duration::ZERO; for (_, t) in &sorted { - total_trace_lde += t.sub_ops.trace_lde; total_constraints += t.sub_ops.constraints; total_comp_decompose += t.sub_ops.comp_decompose; total_comp_commit += t.sub_ops.comp_commit; @@ -155,8 +152,7 @@ pub fn print_report( total_queries += t.sub_ops.queries; } - let sub_ops_sum = total_trace_lde - + total_constraints + let sub_ops_sum = total_constraints + total_comp_decompose + total_comp_commit + total_ood @@ -166,7 +162,6 @@ pub fn print_report( + total_queries; if sub_ops_sum > Duration::ZERO { let mut sub_ops: Vec<(&str, Duration)> = vec![ - ("R1 expand_pool_to_lde", total_trace_lde), ("R2 evaluate", total_constraints), ("R2 decompose_and_extend_d2", total_comp_decompose), ("R2 commit_composition_poly", total_comp_commit), @@ -189,7 +184,6 @@ pub fn print_report( // Cross-round totals: all FFT work and all Merkle work let total_fft = mp.round1_sub.main_lde + mp.round1_sub.aux_lde - + total_trace_lde + total_comp_decompose + total_deep_extend; let total_merkle = mp.round1_sub.main_merkle From 6b8be5b82c6b8b9d2ac02971263061718b46ba8f Mon Sep 17 00:00:00 2001 From: Gabriel Bosio <38794644+gabrielbosio@users.noreply.github.com> Date: Tue, 14 Apr 2026 18:32:10 -0300 Subject: [PATCH 6/7] Remove LDE column pool, commit functions own their buffers (#498) --- crypto/math/src/fft/polynomial.rs | 2 +- crypto/stark/src/instruments.rs | 4 +- crypto/stark/src/prover.rs | 217 ++++++++++-------------------- crypto/stark/src/table.rs | 30 ++--- crypto/stark/src/trace.rs | 23 ++-- prover/src/instruments.rs | 4 +- 6 files changed, 98 insertions(+), 182 deletions(-) diff --git a/crypto/math/src/fft/polynomial.rs b/crypto/math/src/fft/polynomial.rs index 627bfa3ba..9157903fd 100644 --- a/crypto/math/src/fft/polynomial.rs +++ b/crypto/math/src/fft/polynomial.rs @@ -215,7 +215,7 @@ impl Polynomial> { /// /// Unlike [`coset_lde_full_into`], this skips the `clear + extend_from_slice` step /// since data is already in the buffer. Used for transpose elimination: columns are - /// extracted directly into pool buffers, then expanded in-place. + /// extracted directly into owned buffers, then expanded in-place. pub fn coset_lde_full_expand + Send + Sync>( buffer: &mut Vec>, blowup_factor: usize, diff --git a/crypto/stark/src/instruments.rs b/crypto/stark/src/instruments.rs index dd566b1ae..3c9f432d7 100644 --- a/crypto/stark/src/instruments.rs +++ b/crypto/stark/src/instruments.rs @@ -26,11 +26,11 @@ pub struct TableSubOps { /// Sub-operation breakdown for Round 1 aux commit pass. #[derive(Clone, Debug, Default)] pub struct Round1SubOps { - /// Main trace: expand_pool_to_lde (LDE/FFT) + /// Main trace: expand_columns_to_lde (LDE/FFT) pub main_lde: Duration, /// Main trace: commit_columns_bit_reversed (Merkle) pub main_merkle: Duration, - /// Aux trace: expand_pool_to_lde (LDE/FFT) + /// Aux trace: expand_columns_to_lde (LDE/FFT) pub aux_lde: Duration, /// Aux trace: commit_columns_bit_reversed (Merkle) pub aux_merkle: Duration, diff --git a/crypto/stark/src/prover.rs b/crypto/stark/src/prover.rs index 86e760dbb..4f2962c0b 100644 --- a/crypto/stark/src/prover.rs +++ b/crypto/stark/src/prover.rs @@ -133,8 +133,7 @@ where /// eliminating the expensive recomputation (iFFT + coset shift + FFT per column). /// /// Memory trade-off: all N tables' LDE columns are live simultaneously between Phase A/C -/// and Phase D (O(N × cols × lde_size)), but pool pre-allocation is removed, so peak heap -/// is lower overall (~70 GB vs ~220 GB at 8M cycles). +/// and Phase D (O(N × cols × lde_size)). pub struct Round1Metadata where Field: IsFFTField + IsSubFieldOf, @@ -175,7 +174,7 @@ where /// /// The `coset_weights` vector stores `[n_inv, n_inv*g, n_inv*g², ..., n_inv*g^{n-1}]` /// where `g` is the coset offset and `n_inv = 1/n`. These are used in the iFFT+coset-shift -/// step of `expand_pool_to_lde`. +/// step of `expand_columns_to_lde`. pub struct LdeTwiddles { inv: LayerTwiddles, fwd: LayerTwiddles, @@ -234,12 +233,6 @@ fn table_parallelism() -> usize { } } -/// A set of LDE column buffer pools for one concurrent table slot. -struct PoolSet { - main: Vec>>, - aux: Vec>>, -} - /// A container for the results of the second round of the STARK Prove protocol. pub struct Round2 where @@ -424,14 +417,12 @@ pub trait IsStarkProver< .expect("coset LDE computation") } - /// Expand pool buffers in-place from N column evaluations to N×blowup LDE evaluations. + /// Expand each column in-place from N evaluations to N×blowup LDE evaluations. /// - /// The pool buffers already contain column data extracted via `extract_columns_*_into`. - /// This performs iFFT + coset shift + FFT in-place, eliminating the T1 transpose copy. - /// Coset weights are pre-cached in `LdeTwiddles` to avoid recomputation across phases. - fn expand_pool_to_lde( - pool: &mut [Vec>], - num_cols: usize, + /// Performs iFFT + coset shift + FFT in place. Coset weights are pre-cached in + /// `LdeTwiddles` to avoid recomputation across phases. + fn expand_columns_to_lde( + columns: &mut [Vec>], domain: &Domain, twiddles: &LdeTwiddles, ) where @@ -439,14 +430,14 @@ pub trait IsStarkProver< E: IsSubFieldOf + IsField + Send + Sync, FieldElement: Send + Sync, { - if num_cols == 0 { + if columns.is_empty() { return; } #[cfg(feature = "parallel")] - let iter = pool[..num_cols].par_iter_mut(); + let iter = columns.par_iter_mut(); #[cfg(not(feature = "parallel"))] - let iter = pool[..num_cols].iter_mut(); + let iter = columns.iter_mut(); iter.for_each(|buf| { Polynomial::coset_lde_full_expand::( buf, @@ -459,14 +450,13 @@ pub trait IsStarkProver< }); } - /// Compute main LDE, commit, return tree and root. - /// Uses the provided pool buffers to avoid allocation; the pool retains capacity for reuse. + /// Compute main LDE, commit, and return the Merkle tree/root along with the + /// owned LDE columns (consumed later in Phase D). #[allow(clippy::type_complexity)] fn commit_main_trace( trace: &TraceTable, domain: &Domain, twiddles: &LdeTwiddles, - main_pool: &mut [Vec>], ) -> Result< ( BatchedMerkleTree, @@ -474,6 +464,7 @@ pub trait IsStarkProver< Option>, Option, usize, + Vec>>, ), ProvingError, > @@ -481,27 +472,25 @@ pub trait IsStarkProver< FieldElement: AsBytes, FieldElement: AsBytes, { - let num_cols = trace.num_main_columns; - trace.extract_columns_main_into(main_pool); + let lde_size = domain.interpolation_domain_size * domain.blowup_factor; + let mut columns = trace.extract_columns_main(lde_size); #[cfg(feature = "instruments")] let t_sub = Instant::now(); - Self::expand_pool_to_lde::(main_pool, num_cols, domain, twiddles); + Self::expand_columns_to_lde::(&mut columns, domain, twiddles); #[cfg(feature = "instruments")] let main_lde_dur = t_sub.elapsed(); #[cfg(feature = "instruments")] let t_sub = Instant::now(); - let (tree, root) = Self::commit_columns_bit_reversed(&main_pool[..num_cols]) - .ok_or(ProvingError::EmptyCommitment)?; + let (tree, root) = + Self::commit_columns_bit_reversed(&columns).ok_or(ProvingError::EmptyCommitment)?; #[cfg(feature = "instruments")] crate::instruments::accum_r1_main(main_lde_dur, t_sub.elapsed()); - // Pool buffers retain capacity; tree is returned to caller - Ok((tree, root, None, None, 0)) + Ok((tree, root, None, None, 0, columns)) } /// Commit preprocessed trace: precomputed and multiplicity columns get separate trees. - /// Uses pool buffers to avoid allocation. #[allow(clippy::type_complexity)] fn commit_preprocessed_trace( trace: &TraceTable, @@ -509,7 +498,6 @@ pub trait IsStarkProver< precomputed_commitment: Commitment, num_precomputed_cols: usize, twiddles: &LdeTwiddles, - main_pool: &mut [Vec>], ) -> Result< ( BatchedMerkleTree, @@ -517,6 +505,7 @@ pub trait IsStarkProver< Option>, Option, usize, + Vec>>, ), ProvingError, > @@ -525,21 +514,22 @@ pub trait IsStarkProver< FieldElement: AsBytes, { let num_cols = trace.num_main_columns; - trace.extract_columns_main_into(main_pool); + let lde_size = domain.interpolation_domain_size * domain.blowup_factor; + let mut columns = trace.extract_columns_main(lde_size); #[cfg(feature = "instruments")] let t_sub = Instant::now(); - Self::expand_pool_to_lde::(main_pool, num_cols, domain, twiddles); + Self::expand_columns_to_lde::(&mut columns, domain, twiddles); #[cfg(feature = "instruments")] let main_lde_dur = t_sub.elapsed(); #[cfg(feature = "instruments")] let t_sub = Instant::now(); let (precomputed_tree, precomputed_root) = - Self::commit_columns_bit_reversed(&main_pool[..num_precomputed_cols]) + Self::commit_columns_bit_reversed(&columns[..num_precomputed_cols]) .ok_or(ProvingError::EmptyCommitment)?; let (mult_tree, mult_root) = - Self::commit_columns_bit_reversed(&main_pool[num_precomputed_cols..num_cols]) + Self::commit_columns_bit_reversed(&columns[num_precomputed_cols..num_cols]) .ok_or(ProvingError::EmptyCommitment)?; #[cfg(feature = "instruments")] crate::instruments::accum_r1_main(main_lde_dur, t_sub.elapsed()); @@ -549,40 +539,36 @@ pub trait IsStarkProver< "Prover's precomputed commitment doesn't match hardcoded AIR commitment" ); - // Pool buffers retain capacity; trees are returned to caller Ok(( mult_tree, mult_root, Some(precomputed_tree), Some(precomputed_root), num_precomputed_cols, + columns, )) } - /// Reconstruct a full Round1 struct by recomputing LDE evaluations and using - /// the stored Merkle trees from metadata. Uses pool buffers to avoid allocation. + /// Recompute Round1 from the trace, reusing the Merkle trees stored in metadata. /// - /// The Merkle trees were already built during Phase A/C and are reused here, - /// eliminating redundant Keccak hashing. + /// Only used by `run_debug_checks` — Phase D consumes the cached LDE from + /// metadata directly and does not go through this path. + #[cfg(feature = "debug-checks")] fn reconstruct_round1( air: &dyn AIR, trace: &TraceTable, domain: &Domain, metadata: &Round1Metadata, twiddles: &LdeTwiddles, - main_pool: &mut [Vec>], - aux_pool: &mut [Vec>], ) -> Result, ProvingError> where FieldElement: AsBytes, FieldElement: AsBytes, { - // Recompute main LDE into pool buffers (extract columns directly, no T1 transpose) - let num_main_cols = trace.num_main_columns; - trace.extract_columns_main_into(main_pool); - Self::expand_pool_to_lde::(main_pool, num_main_cols, domain, twiddles); + let lde_size = domain.interpolation_domain_size * domain.blowup_factor; + let mut main_columns = trace.extract_columns_main(lde_size); + Self::expand_columns_to_lde::(&mut main_columns, domain, twiddles); - // Use stored Merkle trees from Phase A/C via Arc (pointer copy, no deep clone) let main = Round1CommitmentData:: { lde_trace_merkle_tree: Arc::clone(&metadata.main_merkle_tree), lde_trace_merkle_root: metadata.main_merkle_root, @@ -591,12 +577,9 @@ pub trait IsStarkProver< num_precomputed_cols: metadata.num_precomputed_cols, }; - // Recompute aux LDE into pool buffers, use stored aux Merkle tree - let (aux, num_aux_cols) = if air.has_aux_trace() { - let n_aux = trace.num_aux_columns; - trace.extract_columns_aux_into(aux_pool); - Self::expand_pool_to_lde::(aux_pool, n_aux, domain, twiddles); - // Safe: has_aux_trace() is true only when Phase C stored aux tree/root + let (aux, aux_columns) = if air.has_aux_trace() { + let mut aux_columns = trace.extract_columns_aux(lde_size); + Self::expand_columns_to_lde::(&mut aux_columns, domain, twiddles); let aux_commitment = Round1CommitmentData:: { lde_trace_merkle_tree: Arc::clone( metadata @@ -611,23 +594,17 @@ pub trait IsStarkProver< precomputed_merkle_root: None, num_precomputed_cols: 0, }; - (Some(aux_commitment), n_aux) + (Some(aux_commitment), aux_columns) } else { - (None, 0) + (None, Vec::new()) }; - // Take column Vecs from pool (zero-copy move) instead of cloning. - // After prove_rounds_2_to_4, columns are returned to the pool via into_columns. - let main_cols: Vec<_> = main_pool[..num_main_cols] - .iter_mut() - .map(std::mem::take) - .collect(); - let aux_cols: Vec<_> = aux_pool[..num_aux_cols] - .iter_mut() - .map(std::mem::take) - .collect(); - let lde_trace = - LDETraceTable::from_columns(main_cols, aux_cols, air.step_size(), domain.blowup_factor); + let lde_trace = LDETraceTable::from_columns( + main_columns, + aux_columns, + air.step_size(), + domain.blowup_factor, + ); Ok(Round1 { lde_trace, @@ -646,8 +623,6 @@ pub trait IsStarkProver< metadatas: &[Round1Metadata], domains: &[Domain], twiddle_caches: &[LdeTwiddles], - main_pool: &mut [Vec>], - aux_pool: &mut [Vec>], ) where FieldElement: AsBytes, FieldElement: AsBytes, @@ -660,10 +635,8 @@ pub trait IsStarkProver< .zip(metadatas.iter()) .zip(domains.iter().zip(twiddle_caches.iter())) { - let result = Self::reconstruct_round1( - *air, *trace, domain, metadata, twiddles, main_pool, aux_pool, - ) - .expect("reconstruct_round1 failed in debug-checks"); + let result = Self::reconstruct_round1(*air, *trace, domain, metadata, twiddles) + .expect("reconstruct_round1 failed in debug-checks"); temp_results.push(result); } @@ -1473,7 +1446,7 @@ pub trait IsStarkProver< openings } - // TODO: propagate errors instead of unwrap() in commit_columns, reconstruct_round1, and expand_pool_to_lde + // TODO: propagate errors instead of unwrap() in commit_columns, reconstruct_round1, and expand_columns_to_lde /// Generates STARK proofs for one or more AIRs with a shared transcript. /// /// # Multi-Table Proving with LogUp @@ -1515,7 +1488,7 @@ pub trait IsStarkProver< .any(|(air, _, _)| air.has_aux_trace()); // ===================================================================== - // Pre-pass: compute domains, twiddles, and max dimensions for pool allocation + // Pre-pass: compute domains and twiddles // ===================================================================== #[cfg(feature = "instruments")] @@ -1523,31 +1496,17 @@ pub trait IsStarkProver< let mut domains = Vec::with_capacity(num_airs); let mut twiddle_caches: Vec> = Vec::with_capacity(num_airs); - let mut max_main_cols = 0usize; - let mut max_aux_cols = 0usize; for (air, trace, _pub_inputs) in &*air_trace_pairs { let trace_length = trace.num_rows(); let domain = new_domain(*air, trace_length); let twiddles = LdeTwiddles::new(&domain); - max_main_cols = max_main_cols.max(trace.num_main_columns); - max_aux_cols = max_aux_cols.max(air.num_auxiliary_rap_columns()); - domains.push(domain); twiddle_caches.push(twiddles); } - // Allocate K independent LDE column buffer pool sets for parallel table processing. - // Buffers start empty; each table allocates on demand during extract+expand. - // After commit, LDE columns are moved into per-table cache (not reused across chunks). let k = table_parallelism().min(num_airs).max(1); - let mut pool_sets: Vec> = (0..k) - .map(|_| PoolSet { - main: (0..max_main_cols).map(|_| Vec::new()).collect(), - aux: (0..max_aux_cols).map(|_| Vec::new()).collect(), - }) - .collect(); #[cfg(feature = "instruments")] let prepass_elapsed = phase_start.elapsed(); @@ -1556,7 +1515,7 @@ pub trait IsStarkProver< // Round 1, Phase A: Commit all main traces (parallel in chunks of K) // ===================================================================== // All main trace commitments must be in the transcript before sampling - // LogUp challenges. Pool buffers are reused across chunks. + // LogUp challenges. #[cfg(feature = "instruments")] let phase_start = Instant::now(); @@ -1565,48 +1524,36 @@ pub trait IsStarkProver< for chunk_start in (0..num_airs).step_by(k) { let chunk_end = (chunk_start + k).min(num_airs); - let chunk_size = chunk_end - chunk_start; + let chunk_range = chunk_start..chunk_end; #[cfg(feature = "parallel")] - let iter = pool_sets[..chunk_size].par_iter_mut().enumerate(); + let iter = chunk_range.into_par_iter(); #[cfg(not(feature = "parallel"))] - let iter = pool_sets[..chunk_size].iter_mut().enumerate(); + let iter = chunk_range; let chunk_results: Vec> = iter - .map(|(j, pool)| { - let idx = chunk_start + j; + .map(|idx| { let (air, trace, _) = &air_trace_pairs[idx]; let domain = &domains[idx]; let twiddles = &twiddle_caches[idx]; - let result = if air.is_preprocessed() { + if air.is_preprocessed() { Self::commit_preprocessed_trace( *trace, domain, air.precomputed_commitment(), air.num_precomputed_columns(), twiddles, - &mut pool.main, ) } else { - Self::commit_main_trace(*trace, domain, twiddles, &mut pool.main) - }?; - - // Cache main LDE columns from pool (zero-copy move). - // Pool slots become empty; next chunk allocates on demand. - let num_main_cols = trace.num_main_columns; - let cached_main: Vec<_> = pool.main[..num_main_cols] - .iter_mut() - .map(std::mem::take) - .collect(); - - Ok((result, cached_main)) + Self::commit_main_trace(*trace, domain, twiddles) + } }) .collect(); // Sequential: append roots to shared transcript (Fiat-Shamir ordering) for result in chunk_results { - let ((tree, root, pre_tree, pre_root, n_pre), cached_main) = result?; + let (tree, root, pre_tree, pre_root, n_pre, cached_main) = result?; if let Some(ref pre_r) = pre_root { transcript.append_bytes(pre_r); } @@ -1646,7 +1593,7 @@ pub trait IsStarkProver< // // Split into two passes for parallelism: // Pass 1 (parallel): Build all auxiliary traces (fingerprint + batch inversion) - // Pass 2 (sequential): Fork transcript → extract → LDE → commit (shared pool) + // Pass 2 (parallel): Fork transcript → extract → LDE → commit // Pass 1: Build aux traces in parallel. // Each build_auxiliary_trace has internal parallelism (batch_inverse, par_chunks), @@ -1672,7 +1619,7 @@ pub trait IsStarkProver< let aux_build_elapsed = phase_start.elapsed(); // Pass 2: Parallel fork transcript → extract → LDE → commit in chunks of K. - // Each table gets its own transcript fork and pool set. + // Each table gets its own transcript fork. #[cfg(feature = "instruments")] let phase_start = Instant::now(); @@ -1697,28 +1644,26 @@ pub trait IsStarkProver< for chunk_start in (0..num_airs).step_by(k) { let chunk_end = (chunk_start + k).min(num_airs); - let chunk_size = chunk_end - chunk_start; + let chunk_range = chunk_start..chunk_end; #[cfg(feature = "parallel")] - let iter = pool_sets[..chunk_size].par_iter_mut().enumerate(); + let iter = chunk_range.into_par_iter(); #[cfg(not(feature = "parallel"))] - let iter = pool_sets[..chunk_size].iter_mut().enumerate(); + let iter = chunk_range; let chunk_aux: Vec> = iter - .map(|(j, pool)| { - let idx = chunk_start + j; + .map(|idx| { let (air, trace, _) = &air_trace_pairs[idx]; let domain = &domains[idx]; let twiddles = &twiddle_caches[idx]; if air.has_aux_trace() { - let num_aux_cols = trace.num_aux_columns; - trace.extract_columns_aux_into(&mut pool.aux); + let lde_size = domain.interpolation_domain_size * domain.blowup_factor; + let mut columns = trace.extract_columns_aux(lde_size); #[cfg(feature = "instruments")] let t_sub = Instant::now(); - Self::expand_pool_to_lde::( - &mut pool.aux, - num_aux_cols, + Self::expand_columns_to_lde::( + &mut columns, domain, twiddles, ); @@ -1726,19 +1671,12 @@ pub trait IsStarkProver< let aux_lde_dur = t_sub.elapsed(); #[cfg(feature = "instruments")] let t_sub = Instant::now(); - let (tree, root) = - Self::commit_columns_bit_reversed(&pool.aux[..num_aux_cols]) - .ok_or(ProvingError::EmptyCommitment)?; + let (tree, root) = Self::commit_columns_bit_reversed(&columns) + .ok_or(ProvingError::EmptyCommitment)?; #[cfg(feature = "instruments")] crate::instruments::accum_r1_aux(aux_lde_dur, t_sub.elapsed()); - // Cache aux LDE columns from pool (zero-copy move). - let cached_aux: Vec<_> = pool.aux[..num_aux_cols] - .iter_mut() - .map(std::mem::take) - .collect(); - - Ok((Some(Arc::new(tree)), Some(root), cached_aux)) + Ok((Some(Arc::new(tree)), Some(root), columns)) } else { Ok((None, None, Vec::new())) } @@ -1782,20 +1720,7 @@ pub trait IsStarkProver< let aux_commit_elapsed = phase_start.elapsed(); #[cfg(feature = "debug-checks")] - { - let debug_pool = &mut pool_sets[0]; - Self::run_debug_checks( - &air_trace_pairs, - &metadatas, - &domains, - &twiddle_caches, - &mut debug_pool.main, - &mut debug_pool.aux, - ); - } - - // Pool sets are no longer needed — Phase D uses cached LDE from metadata. - drop(pool_sets); + Self::run_debug_checks(&air_trace_pairs, &metadatas, &domains, &twiddle_caches); // ===================================================================== // Rounds 2-4: Parallel per-table proving in chunks of K diff --git a/crypto/stark/src/table.rs b/crypto/stark/src/table.rs index 4a946f2ad..352b9e6f7 100644 --- a/crypto/stark/src/table.rs +++ b/crypto/stark/src/table.rs @@ -83,29 +83,25 @@ impl Table { .collect() } - /// Extract columns directly into pre-allocated output buffers. + /// Extract columns as owned vectors, with each allocated at `capacity`. /// - /// Each `output[col_idx]` is cleared and filled with the column data. - /// When `output[col_idx].capacity() >= height`, no heap allocation occurs. - /// This eliminates the T1 transpose allocation that `columns()` performs. - pub fn extract_columns_into(&self, output: &mut [Vec>]) { - debug_assert!( - output.len() >= self.width, - "output has {} buffers but table has {} columns", - output.len(), - self.width - ); + /// `capacity` is a hint sized for downstream LDE expansion so the FFT grows + /// in place without a second allocation. Avoids the T1 transpose `columns()` + /// performs. + pub fn extract_columns(&self, capacity: usize) -> Vec>> { + let capacity = capacity.max(self.height); #[cfg(feature = "parallel")] - let iter = output[..self.width].par_iter_mut().enumerate(); + let iter = (0..self.width).into_par_iter(); #[cfg(not(feature = "parallel"))] - let iter = output[..self.width].iter_mut().enumerate(); - iter.for_each(|(col_idx, buf)| { - buf.clear(); - buf.reserve(self.height.saturating_sub(buf.capacity())); + let iter = 0..self.width; + iter.map(|col_idx| { + let mut buf = Vec::with_capacity(capacity); for row_idx in 0..self.height { buf.push(self.data[row_idx * self.width + col_idx].clone()); } - }); + buf + }) + .collect() } /// Given row and column indexes, returns the stored field element in that position of the table. diff --git a/crypto/stark/src/trace.rs b/crypto/stark/src/trace.rs index 0025c22ea..d172c80f2 100644 --- a/crypto/stark/src/trace.rs +++ b/crypto/stark/src/trace.rs @@ -167,20 +167,15 @@ where .unwrap() } - /// Extract main columns directly into pre-allocated output buffers. - /// - /// Eliminates the T1 transpose allocation that `columns_main()` performs. - /// When `output` buffers have sufficient capacity, no heap allocation occurs. - pub fn extract_columns_main_into(&self, output: &mut [Vec>]) { - self.main_table.extract_columns_into(output); - } - - /// Extract auxiliary columns directly into pre-allocated output buffers. - /// - /// Eliminates the T1 transpose allocation that `columns_aux()` performs. - /// When `output` buffers have sufficient capacity, no heap allocation occurs. - pub fn extract_columns_aux_into(&self, output: &mut [Vec>]) { - self.aux_table.extract_columns_into(output); + /// Extract main columns as owned vectors, each allocated at `capacity`. + /// Pass the LDE size so downstream FFT expansion is in-place. + pub fn extract_columns_main(&self, capacity: usize) -> Vec>> { + self.main_table.extract_columns(capacity) + } + + /// Extract auxiliary columns as owned vectors, each allocated at `capacity`. + pub fn extract_columns_aux(&self, capacity: usize) -> Vec>> { + self.aux_table.extract_columns(capacity) } } /// Column-major LDE trace table. diff --git a/prover/src/instruments.rs b/prover/src/instruments.rs index ac517e85b..e087cf55a 100644 --- a/prover/src/instruments.rs +++ b/prover/src/instruments.rs @@ -75,11 +75,11 @@ 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(" expand_pool_to_lde", mp.round1_sub.main_lde, 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(" 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(" expand_columns_to_lde", mp.round1_sub.aux_lde, total); row_sub(" commit (Merkle)", mp.round1_sub.aux_merkle, total); row_top("Rounds 2\u{2013}4", mp.rounds_2_4, total); From 97e0e9102eb0dae260949e489014f1844de477f9 Mon Sep 17 00:00:00 2001 From: Mauro Toscano <12560266+MauroToscano@users.noreply.github.com> Date: Wed, 15 Apr 2026 17:16:54 -0300 Subject: [PATCH 7/7] Simplification (#500) * Simplification * Rename metadata -> cache * Rework * CacheLDE -> LDE * Drop has_aux_trace param from build_round1 * Route reconstruct_round1 through build_round1 * Revert "Drop has_aux_trace param from build_round1" This reverts commit 164c7c0484532248a58b32b0838a9567482d1841. * Pass has_aux_trace to build_round1 from reconstruct_round1 --------- Co-authored-by: gabrielbosio --- crypto/stark/src/prover.rs | 247 ++++++++++++++++++------------------- 1 file changed, 120 insertions(+), 127 deletions(-) diff --git a/crypto/stark/src/prover.rs b/crypto/stark/src/prover.rs index 4f2962c0b..41d27fc66 100644 --- a/crypto/stark/src/prover.rs +++ b/crypto/stark/src/prover.rs @@ -81,7 +81,7 @@ where { /// The Merkle trees constructed to obtain the commitment of the entire trace table. /// For preprocessed tables, this contains only the multiplicity columns. - /// Wrapped in Arc to share with Round1Metadata without deep-cloning (~64MB per table). + /// Wrapped in Arc to share with Round1Commitments without deep-cloning (~64MB per table). pub(crate) lde_trace_merkle_tree: Arc>, /// The root of the Merkle tree in `lde_trace_merkle_tree`. pub(crate) lde_trace_merkle_root: Commitment, @@ -124,17 +124,11 @@ where precomputed_tree: Option>>, precomputed_root: Option, num_precomputed_cols: usize, - /// Cached main LDE columns from Phase A, used in Phase D to skip recomputation. - cached_main_lde: Vec>>, } -/// Metadata from Round 1 commitments — stores Merkle trees, roots, and cached LDE evaluations. -/// LDE evaluations are cached from Phase A/C and consumed in Phase D (Rounds 2-4), -/// eliminating the expensive recomputation (iFFT + coset shift + FFT per column). -/// -/// Memory trade-off: all N tables' LDE columns are live simultaneously between Phase A/C -/// and Phase D (O(N × cols × lde_size)). -pub struct Round1Metadata +/// Round 1 commitment artifacts — Merkle trees, roots, challenges, and bus inputs. +/// Borrowed (not consumed) when building `Round1` in Phase D. +pub struct Round1Commitments where Field: IsFFTField + IsSubFieldOf, FieldExtension: IsField, @@ -160,10 +154,68 @@ where rap_challenges: Vec>, /// Bus interaction public inputs (initial and final aux column values). bus_public_inputs: Option>, - /// Cached main LDE columns from Phase A (consumed by Phase D). - cached_main_lde: Vec>>, - /// Cached aux LDE columns from Phase C (consumed by Phase D). - cached_aux_lde: Vec>>, +} + +/// LDE columns for main (Phase A) and auxiliary (Phase C) traces, consumed by value in Phase D. +/// +/// Memory trade-off: all N tables' LDE columns are live simultaneously between Phase A/C +/// and Phase D (O(N × cols × lde_size)). +struct Lde { + main: Vec>>, + aux: Vec>>, +} + +impl Round1Commitments +where + Field: IsFFTField + IsSubFieldOf + Send + Sync, + FieldExtension: IsField + Send + Sync, + FieldElement: AsBytes, + FieldElement: AsBytes, +{ + /// Build a `Round1` by consuming a `Lde` and borrowing commitment data. + fn build_round1( + &self, + lde: Lde, + step_size: usize, + blowup_factor: usize, + has_aux_trace: bool, + ) -> Round1 { + let lde_trace = LDETraceTable::from_columns(lde.main, lde.aux, step_size, blowup_factor); + + let main = Round1CommitmentData:: { + lde_trace_merkle_tree: Arc::clone(&self.main_merkle_tree), + lde_trace_merkle_root: self.main_merkle_root, + precomputed_merkle_tree: self.precomputed_merkle_tree.as_ref().map(Arc::clone), + precomputed_merkle_root: self.precomputed_merkle_root, + num_precomputed_cols: self.num_precomputed_cols, + }; + + let aux = if has_aux_trace { + Some(Round1CommitmentData:: { + lde_trace_merkle_tree: Arc::clone( + self.aux_merkle_tree + .as_ref() + .expect("aux tree must exist when has_aux_trace"), + ), + lde_trace_merkle_root: self + .aux_merkle_root + .expect("aux root must exist when has_aux_trace"), + precomputed_merkle_tree: None, + precomputed_merkle_root: None, + num_precomputed_cols: 0, + }) + } else { + None + }; + + Round1 { + lde_trace, + main, + aux, + rap_challenges: self.rap_challenges.clone(), + bus_public_inputs: self.bus_public_inputs.clone(), + } + } } /// Pre-computed twiddle factors and coset weights for a given domain size. @@ -513,7 +565,6 @@ pub trait IsStarkProver< FieldElement: AsBytes, FieldElement: AsBytes, { - let num_cols = trace.num_main_columns; let lde_size = domain.interpolation_domain_size * domain.blowup_factor; let mut columns = trace.extract_columns_main(lde_size); #[cfg(feature = "instruments")] @@ -529,7 +580,7 @@ pub trait IsStarkProver< .ok_or(ProvingError::EmptyCommitment)?; let (mult_tree, mult_root) = - Self::commit_columns_bit_reversed(&columns[num_precomputed_cols..num_cols]) + Self::commit_columns_bit_reversed(&columns[num_precomputed_cols..]) .ok_or(ProvingError::EmptyCommitment)?; #[cfg(feature = "instruments")] crate::instruments::accum_r1_main(main_lde_dur, t_sub.elapsed()); @@ -549,16 +600,16 @@ pub trait IsStarkProver< )) } - /// Recompute Round1 from the trace, reusing the Merkle trees stored in metadata. + /// Recompute Round1 from the trace, reusing the Merkle trees stored in commitments. /// - /// Only used by `run_debug_checks` — Phase D consumes the cached LDE from - /// metadata directly and does not go through this path. + /// Only used by `run_debug_checks` — Phase D consumes the cached LDE + /// directly and does not go through this path. #[cfg(feature = "debug-checks")] fn reconstruct_round1( air: &dyn AIR, trace: &TraceTable, domain: &Domain, - metadata: &Round1Metadata, + commitment: &Round1Commitments, twiddles: &LdeTwiddles, ) -> Result, ProvingError> where @@ -566,53 +617,23 @@ pub trait IsStarkProver< FieldElement: AsBytes, { let lde_size = domain.interpolation_domain_size * domain.blowup_factor; - let mut main_columns = trace.extract_columns_main(lde_size); - Self::expand_columns_to_lde::(&mut main_columns, domain, twiddles); - - let main = Round1CommitmentData:: { - lde_trace_merkle_tree: Arc::clone(&metadata.main_merkle_tree), - lde_trace_merkle_root: metadata.main_merkle_root, - precomputed_merkle_tree: metadata.precomputed_merkle_tree.as_ref().map(Arc::clone), - precomputed_merkle_root: metadata.precomputed_merkle_root, - num_precomputed_cols: metadata.num_precomputed_cols, - }; + let mut main = trace.extract_columns_main(lde_size); + Self::expand_columns_to_lde::(&mut main, domain, twiddles); - let (aux, aux_columns) = if air.has_aux_trace() { - let mut aux_columns = trace.extract_columns_aux(lde_size); - Self::expand_columns_to_lde::(&mut aux_columns, domain, twiddles); - let aux_commitment = Round1CommitmentData:: { - lde_trace_merkle_tree: Arc::clone( - metadata - .aux_merkle_tree - .as_ref() - .expect("aux tree must exist when has_aux_trace"), - ), - lde_trace_merkle_root: metadata - .aux_merkle_root - .expect("aux root must exist when has_aux_trace"), - precomputed_merkle_tree: None, - precomputed_merkle_root: None, - num_precomputed_cols: 0, - }; - (Some(aux_commitment), aux_columns) + let aux = if air.has_aux_trace() { + let mut aux = trace.extract_columns_aux(lde_size); + Self::expand_columns_to_lde::(&mut aux, domain, twiddles); + aux } else { - (None, Vec::new()) + Vec::new() }; - let lde_trace = LDETraceTable::from_columns( - main_columns, - aux_columns, + Ok(commitment.build_round1( + Lde { main, aux }, air.step_size(), domain.blowup_factor, - ); - - Ok(Round1 { - lde_trace, - main, - aux, - rap_challenges: metadata.rap_challenges.clone(), - bus_public_inputs: metadata.bus_public_inputs.clone(), - }) + air.has_aux_trace(), + )) } /// Reconstruct Round1 for every table, print the bus balance report, and @@ -620,7 +641,7 @@ pub trait IsStarkProver< #[cfg(feature = "debug-checks")] fn run_debug_checks( air_trace_pairs: &[AirTracePair<'_, Field, FieldExtension, PI>], - metadatas: &[Round1Metadata], + commitments: &[Round1Commitments], domains: &[Domain], twiddle_caches: &[LdeTwiddles], ) where @@ -630,12 +651,12 @@ pub trait IsStarkProver< { let mut temp_results: Vec> = Vec::with_capacity(air_trace_pairs.len()); - for (((air, trace, _), metadata), (domain, twiddles)) in air_trace_pairs + for (((air, trace, _), commitment), (domain, twiddles)) in air_trace_pairs .iter() - .zip(metadatas.iter()) + .zip(commitments.iter()) .zip(domains.iter().zip(twiddle_caches.iter())) { - let result = Self::reconstruct_round1(*air, *trace, domain, metadata, twiddles) + let result = Self::reconstruct_round1(*air, *trace, domain, commitment, twiddles) .expect("reconstruct_round1 failed in debug-checks"); temp_results.push(result); } @@ -1521,6 +1542,7 @@ pub trait IsStarkProver< let phase_start = Instant::now(); let mut main_commits: Vec> = Vec::with_capacity(num_airs); + let mut main_ldes: Vec>>> = Vec::with_capacity(num_airs); for chunk_start in (0..num_airs).step_by(k) { let chunk_end = (chunk_start + k).min(num_airs); @@ -1564,8 +1586,8 @@ pub trait IsStarkProver< precomputed_tree: pre_tree.map(Arc::new), precomputed_root: pre_root, num_precomputed_cols: n_pre, - cached_main_lde: cached_main, }); + main_ldes.push(cached_main); } } @@ -1693,15 +1715,19 @@ pub trait IsStarkProver< } } - // Build metadata sequentially from main_commits + aux_results + bus_inputs - let mut metadatas: Vec> = + // Build commitments and cached LDEs as separate vecs: + // commitments are borrowed in Phase D, LDEs are consumed by value. + let mut commitments: Vec> = Vec::with_capacity(num_airs); - for ((main_commit, (aux_tree, aux_root, cached_aux)), bus_public_inputs) in main_commits - .into_iter() - .zip(aux_results) - .zip(bus_inputs_vec) + let mut cached_ldes: Vec> = Vec::with_capacity(num_airs); + for (((main_commit, main_lde), (aux_tree, aux_root, cached_aux)), bus_public_inputs) in + main_commits + .into_iter() + .zip(main_ldes) + .zip(aux_results) + .zip(bus_inputs_vec) { - metadatas.push(Round1Metadata { + commitments.push(Round1Commitments { main_merkle_tree: main_commit.main_tree, main_merkle_root: main_commit.main_root, precomputed_merkle_tree: main_commit.precomputed_tree, @@ -1711,8 +1737,10 @@ pub trait IsStarkProver< aux_merkle_root: aux_root, rap_challenges: lookup_challenges.clone(), bus_public_inputs, - cached_main_lde: main_commit.cached_main_lde, - cached_aux_lde: cached_aux, + }); + cached_ldes.push(Lde { + main: main_lde, + aux: cached_aux, }); } @@ -1720,7 +1748,7 @@ pub trait IsStarkProver< let aux_commit_elapsed = phase_start.elapsed(); #[cfg(feature = "debug-checks")] - Self::run_debug_checks(&air_trace_pairs, &metadatas, &domains, &twiddle_caches); + Self::run_debug_checks(&air_trace_pairs, &commitments, &domains, &twiddle_caches); // ===================================================================== // Rounds 2-4: Parallel per-table proving in chunks of K @@ -1740,25 +1768,31 @@ pub trait IsStarkProver< )> = Vec::with_capacity(num_airs); let mut proofs = Vec::with_capacity(num_airs); + let mut lde_drain = cached_ldes.into_iter(); for chunk_start in (0..num_airs).step_by(k) { let chunk_end = (chunk_start + k).min(num_airs); + let chunk_size = chunk_end - chunk_start; - let chunk_metadatas = &mut metadatas[chunk_start..chunk_end]; + let chunk_ldes: Vec> = + lde_drain.by_ref().take(chunk_size).collect(); + let chunk_commitments = &commitments[chunk_start..chunk_end]; let chunk_transcripts = &mut table_transcripts[chunk_start..chunk_end]; #[cfg(feature = "parallel")] - let iter = chunk_metadatas - .par_iter_mut() + let iter = chunk_ldes + .into_par_iter() + .zip(chunk_commitments.par_iter()) .zip(chunk_transcripts.par_iter_mut()) .enumerate(); #[cfg(not(feature = "parallel"))] - let iter = chunk_metadatas - .iter_mut() + let iter = chunk_ldes + .into_iter() + .zip(chunk_commitments.iter()) .zip(chunk_transcripts.iter_mut()) .enumerate(); let chunk_results: Vec> = iter - .map(|(j, (metadata, table_transcript))| { + .map(|(j, ((lde, commitment), table_transcript))| { let idx = chunk_start + j; let (air, trace, pub_inputs) = &air_trace_pairs[idx]; let _ = trace; // used by instruments @@ -1767,55 +1801,14 @@ pub trait IsStarkProver< #[cfg(feature = "instruments")] let table_start = Instant::now(); - // Build Round1 from cached LDE (zero-copy move, no recomputation). - let cached_main = std::mem::take(&mut metadata.cached_main_lde); - let cached_aux = std::mem::take(&mut metadata.cached_aux_lde); - - let lde_trace = LDETraceTable::from_columns( - cached_main, - cached_aux, + // Build Round1 from cached LDE (consumed by value, no recomputation). + let round_1_result = commitment.build_round1( + lde, air.step_size(), domain.blowup_factor, + air.has_aux_trace(), ); - let main = Round1CommitmentData:: { - lde_trace_merkle_tree: Arc::clone(&metadata.main_merkle_tree), - lde_trace_merkle_root: metadata.main_merkle_root, - precomputed_merkle_tree: metadata - .precomputed_merkle_tree - .as_ref() - .map(Arc::clone), - precomputed_merkle_root: metadata.precomputed_merkle_root, - num_precomputed_cols: metadata.num_precomputed_cols, - }; - - let aux = if air.has_aux_trace() { - Some(Round1CommitmentData:: { - lde_trace_merkle_tree: Arc::clone( - metadata - .aux_merkle_tree - .as_ref() - .expect("aux tree must exist when has_aux_trace"), - ), - lde_trace_merkle_root: metadata - .aux_merkle_root - .expect("aux root must exist when has_aux_trace"), - precomputed_merkle_tree: None, - precomputed_merkle_root: None, - num_precomputed_cols: 0, - }) - } else { - None - }; - - let round_1_result = Round1 { - lde_trace, - main, - aux, - rap_challenges: metadata.rap_challenges.clone(), - bus_public_inputs: metadata.bus_public_inputs.clone(), - }; - if let Some(ref bpi) = round_1_result.bus_public_inputs { table_transcript.append_field_element(&bpi.table_contribution); }