Review fixes for prove-and-retire (#994)

bin/cli/Cargo.toml

@@ -14,12 +14,19 @@ clap = { version = "4.3.10", features = ["derive"] }
 rkyv = { version = "0.8.10", default-features = false, features = ["alloc", "bytecheck", "aligned", "pointer_width_64"] }
 tempfile = "3"
 tikv-jemallocator = "0.6"
-tikv-jemalloc-ctl = { version = "0.6", features = ["stats"] }
+# No `stats` by default: that feature propagates to tikv-jemalloc-sys and builds
+# jemalloc with `--enable-stats`, i.e. counters on the malloc fast path, for every
+# binary. `keep_large_buffers_warm` needs only `raw`/`mallctl`; the heap tracker is
+# what needs the counters, and it is behind `jemalloc-stats`.
+tikv-jemalloc-ctl = { version = "0.6" }
 tikv-jemalloc-sys = "0.6"
 env_logger = "0.11"
+# Used by `keep_large_buffers_warm`, whose body is Linux-only — so a macOS build
+# will not catch its absence.
+log = "0.4"
 
 [features]
-jemalloc-stats = []
+jemalloc-stats = ["tikv-jemalloc-ctl/stats"]
 disk-spill = ["prover/disk-spill"]
 instruments = ["prover/instruments", "stark/instruments"]
 # GPU profiling build (Nsight): CUDA prover + instruments spans + NVTX ranges.

bin/cli/src/main.rs

@@ -18,8 +18,18 @@ static ALLOC: tikv_jemallocator::Jemalloc = tikv_jemallocator::Jemalloc;
 /// Disable the arena's decay and purge it on our own clock instead: hot buffers are
 /// reused across threads, cold ones still go back to the OS.
 ///
+/// The index is `opt.narenas`: jemalloc 5 reserves the slot immediately after the
+/// automatic arenas for the oversize arena (`arena_init_huge`, `huge_arena_ind =
+/// narenas_total_get()`), and `opt.oversize_threshold` defaults to the same 8 MiB.
+///
 /// Linux only: elsewhere jemalloc is built without background threads and the decay
 /// mallctl traps.
+///
+/// Called only from the paths that allocate trace-sized buffers, not from `main`:
+/// it disables an arena's decay for the life of the process and leaves a purge
+/// thread behind, which is not something `cli verify` or `--help` should pay, and
+/// it would otherwise change the allocator under every baseline measured with
+/// this binary.
 fn keep_large_buffers_warm() {
     #[cfg(target_os = "linux")]
     {
@@ -35,13 +45,28 @@ fn keep_large_buffers_warm() {
         // SAFETY: `opt.narenas` is `unsigned`, the decay knob is `ssize_t`, and `purge`
         // takes no value.
         unsafe {
-            let Ok(huge_arena) = raw::read::<u32>(b"opt.narenas\0") else {
-                return;
+            let huge_arena = match raw::read::<u32>(b"opt.narenas\0") {
+                Ok(n) => n,
+                Err(e) => {
+                    // Not fatal, but the run is then indistinguishable from one
+                    // where the knob worked — which is exactly what makes a
+                    // memory measurement unreadable. Say so.
+                    log::warn!(
+                        "keep_large_buffers_warm: cannot read opt.narenas ({e}); \
+                                the oversize arena keeps jemalloc's default decay"
+                    );
+                    return;
+                }
             };
             let decay = format!("arena.{huge_arena}.dirty_decay_ms\0");
-            if raw::write(decay.as_bytes(), -1i64).is_err() {
+            if let Err(e) = raw::write(decay.as_bytes(), -1i64) {
+                log::warn!(
+                    "keep_large_buffers_warm: cannot disable decay on arena \
+                            {huge_arena} ({e}); the oversize arena keeps jemalloc's default"
+                );
                 return;
             }
+            log::debug!("keep_large_buffers_warm: decay disabled on arena {huge_arena}");
             let purge = CString::new(format!("arena.{huge_arena}.purge")).unwrap();
             std::thread::spawn(move || {
                 loop {
@@ -356,7 +381,6 @@ enum Stage {
 }
 
 fn main() -> ExitCode {
-    keep_large_buffers_warm();
     env_logger::init();
     let cli = Cli::parse();
 
@@ -1258,7 +1282,7 @@ fn report_batched_size(proof: &prover::logup_phase::BatchedProof, tables: usize,
 
 /// Where the time went, summed per span label.
 ///
-/// The prover's own spans are per table and there are 227 of them, so the raw
+/// The prover's own spans are per table and there are hundreds of them, so the raw
 /// timeline is unreadable; what answers "where is the time" is the total per
 /// label. Sums exceed wall time, because tables run concurrently — the ratios
 /// between labels are the point, not the absolute figures.
@@ -1338,6 +1362,23 @@ fn cmd_trace_build(
     verify: bool,
     output: Option<PathBuf>,
 ) -> ExitCode {
+    // Only the per-table proof is serializable today, so `--output` with any
+    // other stage would walk the whole execution and then write nothing. Say so
+    // before spending the walk rather than exiting 0 in silence.
+    if output.is_some() && prove_and_retire && through != Stage::Logup {
+        eprintln!(
+            "--output writes the per-table proof, which only `--through logup` assembles; \
+             `--through {}` has no serializable proof yet (see docs/prove_and_retire_design.md \u{00A7}10).",
+            match through {
+                Stage::Walk => "walk",
+                Stage::Commit => "commit",
+                Stage::Challenge => "challenge",
+                Stage::Batched => "batched",
+                Stage::Logup => unreachable!(),
+            }
+        );
+        return ExitCode::FAILURE;
+    }
     let elf_data = match std::fs::read(&elf_path) {
         Ok(data) => data,
         Err(e) => {
@@ -1373,6 +1414,9 @@ fn cmd_trace_build(
         }
     };
     let outcome = if prove_and_retire {
+        // The walk allocates and drops one trace-sized buffer after another,
+        // which is the pattern this works around.
+        keep_large_buffers_warm();
         run_approach_1(
             &elf,
             &elf_data,

crypto/crypto/src/merkle_tree/merkle.rs

@@ -330,15 +330,20 @@ where
         leaves_len: usize,
         sibling_leaf: B::Node,
     ) -> Option<Proof<B::Node>> {
-        if leaves_len <= 1 || pos >= leaves_len {
+        if leaves_len <= 1 || pos >= leaves_len || self.is_root_only() {
             return None;
         }
         let mut merkle_path = Vec::with_capacity(leaves_len.trailing_zeros() as usize);
         merkle_path.push(sibling_leaf);
 
         let mut node = parent_index(pos + leaves_len - 1);
         while node != ROOT {
-            merkle_path.push(self.nodes.get(sibling_index(node))?.clone());
+            // `node_get`, not `self.nodes` directly: every other read in this
+            // file goes through it for the disk-spill mmap indirection. The two
+            // are mutually exclusive today — `drop_leaves` refuses an mmap-backed
+            // tree — but a direct read would silently yield `None` here if that
+            // ever stopped holding, and the prover's opening path unwraps this.
+            merkle_path.push(self.node_get(sibling_index(node))?.clone());
             node = parent_index(node);
         }
         self.create_proof(merkle_path)

crypto/stark/src/batched_verifier.rs

@@ -36,6 +36,10 @@ struct GroupReplay<F: IsFFTField, E: IsField> {
     iotas: Vec<usize>,
     acc: Vec<FieldElement<E>>,
     acc_sym: Vec<FieldElement<E>>,
+    /// The first table of the group, whose AIR supplied the domain. Kept from
+    /// the scan that already found it rather than looked up again: the second
+    /// scan is what forced an `expect` into verifier code.
+    member: usize,
 }
 
 /// Verify the rounds after round 1 of every table and each group's FRI.
@@ -90,12 +94,33 @@ where
     let num_queries = first_air.options().fri_number_of_queries;
     let grinding_factor = first_air.context().proof_options.grinding_factor;
 
-    // Every table's domain is its group's.
-    for (idx, (table, &g)) in tables.iter().zip(group_of).enumerate() {
+    // Every table's domain is its group's, and every block and opening has the
+    // shape its AIR declares.
+    //
+    // Both run before the fold seed below reads a single out-of-domain value.
+    // That seed indexes the blocks at the dimensions the *proof* advertises, so
+    // a proof whose advertised dimensions disagree with its data length has to
+    // be rejected here rather than panic there — the rule `ood_blocks_well_formed`
+    // documents, and the one the ordinary verifier keeps by running it in round 1.
+    for (idx, ((air, table), &g)) in airs.iter().zip(tables).zip(group_of).enumerate() {
         if table.trace_length == 0 || table.trace_length != groups[g].trace_rows {
             error!("batched: table {idx} does not live on its group's domain");
             return false;
         }
+        let view = StarkProofView::Owned(table);
+        // `trace_length` is non-zero by the check above, so the division is safe.
+        if table.composition_poly_parts_ood_evaluation.len()
+            != air.composition_poly_degree_bound(table.trace_length) / table.trace_length
+            || !V::<Field, FieldExtension, PI>::ood_blocks_well_formed(*air, view)
+            || !V::<Field, FieldExtension, PI>::trace_opening_widths_well_formed(
+                *air,
+                view,
+                num_queries,
+            )
+        {
+            error!("batched: table {idx}'s blocks or openings are malformed");
+            return false;
+        }
     }
 
     // The fold coefficients, from the seed, in the prover's order.
@@ -182,6 +207,7 @@ where
             iotas,
             acc: vec![FieldElement::zero(); num_queries],
             acc_sym: vec![FieldElement::zero(); num_queries],
+            member,
         });
     }
 
@@ -199,27 +225,17 @@ where
         if let Some(ref bpi) = table.bus_public_inputs {
             fork.append_field_element(&bpi.table_contribution);
         }
+        // Shapes were pinned to the AIR in the pre-pass above, before the fold
+        // seed read any of this table's blocks.
         let domain = new_verifier_domain(*air, table.trace_length);
-        if table.composition_poly_parts_ood_evaluation.len()
-            != air.composition_poly_degree_bound(table.trace_length) / table.trace_length
-            || !V::<Field, FieldExtension, PI>::ood_blocks_well_formed(*air, view)
-            || !V::<Field, FieldExtension, PI>::trace_opening_widths_well_formed(
-                *air,
-                view,
-                num_queries,
-            )
-        {
-            error!("batched: table {idx}'s blocks or openings are malformed");
-            return false;
-        }
         let layout = V::<Field, FieldExtension, PI>::ood_layout(*air);
         let RoundsChallenges {
             z,
             boundary_coeffs,
             transition_coeffs,
             trace_term_coeffs,
             gammas,
-        } = V::<Field, FieldExtension, PI>::replay_rounds_2_and_3(
+        } = V::<Field, FieldExtension, PI>::replay_rounds_2_to_4(
             *air,
             view,
             &public_inputs[idx],
@@ -291,10 +307,7 @@ where
 
     // Each group's FRI, from the folded first layer down to the final polynomial.
     for (g, (group, replay)) in groups.iter().zip(replays.iter()).enumerate() {
-        let member = group_of
-            .iter()
-            .position(|&h| h == g)
-            .expect("checked above");
+        let member = replay.member;
         let fri = group.fri;
         let synthetic = StarkProof::<Field, FieldExtension, PI> {
             trace_length: group.trace_rows,

crypto/stark/src/prover.rs

@@ -131,10 +131,12 @@ impl<F: IsField> TableCommit<F>
 where
     FieldElement<F>: AsBytes,
 {
-    /// Build a `TableCommit` for a plain (non-preprocessed) table.
     /// Roots without a tree, for a pass that will not open against it. The
     /// tree is the expensive half; a caller that already knows the roots and
     /// only needs them to travel should not pay for one.
+    ///
+    /// Serves preprocessed and plain tables alike — `precomputed` is `Some`
+    /// exactly when the AIR is preprocessed.
     fn known_roots(root: Commitment, precomputed: Option<Commitment>) -> Self {
         Self {
             tree: Arc::new(BatchedMerkleTree::from_root(root)),
@@ -146,6 +148,7 @@ where
         }
     }
 
+    /// Build a `TableCommit` for a plain (non-preprocessed) table.
     fn plain(#[allow(unused_mut)] mut tree: BatchedMerkleTree<F>, root: Commitment) -> Self {
         let leaves_dropped = Self::retire_leaves(&mut tree);
         Self {
@@ -646,26 +649,6 @@ fn host_cores() -> usize {
         .unwrap_or(4)
 }
 
-/// Number of tables `multi_prove` proves concurrently, out of `num_airs` of
-/// them.
-///
-/// Defaults: **every table** under `cuda`, `num_cores / 3` on CPU builds
-/// (benchmarked optimal on both M3 Pro and EPYC 9454P — every table there is
-/// pure host work, so `k` genuinely competes for cores). Both arms are
-/// overridden by the `TABLE_PARALLELISM` env var, and the result is clamped to
-/// `1..=num_airs`. Without the `parallel` feature this is 1 and the env var is
-/// ignored.
-///
-/// # Why the `cuda` arm has no core term
-///
-/// Measured over 881 runs on two RTX 5090 boxes (sweep record linked from
-/// PR #911): the work `k` divides is device- and workload-bound — invariant to
-/// host core count over an 8× range — so `available_parallelism()` is the
-/// wrong quantity to scale `k` by. `k` is not a thread count; it counts
-/// concurrent drivers whose per-table work all runs on the one global rayon
-/// pool. Worst case against the best measured `k`: `num_airs` +1.6 % (inside
-/// noise), the old `cores*2/3` +13.0 %. Bounding concurrency is memory
-/// admission's job (`VramGate`), not this count's.
 /// A table's Round 1 roots, in the order Fiat-Shamir absorbs them.
 ///
 /// A plain table contributes one root. A preprocessed one contributes two: its
@@ -708,10 +691,6 @@ pub struct TableDeep<FieldExtension: IsField> {
     pub trace_rows: usize,
     /// The DEEP composition codeword, `lde_size` long.
     pub deep: Vec<FieldElement<FieldExtension>>,
-    /// Where the table sits in the AIR order, carried so the batch can say
-    /// which group each table ended up in once the codewords are sorted by
-    /// domain rather than by position.
-    pub air_index: usize,
     /// What the batch's coefficient is drawn from: this table's public round-3
     /// data, in the order a transcript absorbs it.
     ///
@@ -826,6 +805,29 @@ const RETIRE_OVERRIDE_OFF: u8 = 1;
 const RETIRE_OVERRIDE_ON: u8 = 2;
 static RETIRE_LDE_OVERRIDE: AtomicU8 = AtomicU8::new(RETIRE_OVERRIDE_UNSET);
 
+/// Number of tables `multi_prove` proves concurrently, out of `num_airs` of
+/// them.
+///
+/// Defaults: **every table** under `cuda`, `num_cores / 3` on CPU builds
+/// (benchmarked optimal on both M3 Pro and EPYC 9454P — every table there is
+/// pure host work, so `k` genuinely competes for cores). Both arms are
+/// overridden by the `TABLE_PARALLELISM` env var, and the result is clamped to
+/// `1..=num_airs`. Without the `parallel` feature this is 1 and the env var is
+/// ignored.
+///
+/// Not to be confused with `prover::pass::table_parallelism`, which sizes the
+/// prove-and-retire walk's batches and reads `A1_TABLE_PARALLELISM`.
+///
+/// # Why the `cuda` arm has no core term
+///
+/// Measured over 881 runs on two RTX 5090 boxes (sweep record linked from
+/// PR #911): the work `k` divides is device- and workload-bound — invariant to
+/// host core count over an 8× range — so `available_parallelism()` is the
+/// wrong quantity to scale `k` by. `k` is not a thread count; it counts
+/// concurrent drivers whose per-table work all runs on the one global rayon
+/// pool. Worst case against the best measured `k`: `num_airs` +1.6 % (inside
+/// noise), the old `cores*2/3` +13.0 %. Bounding concurrency is memory
+/// admission's job (`VramGate`), not this count's.
 pub fn table_parallelism(num_airs: usize) -> usize {
     #[cfg(feature = "parallel")]
     {
@@ -2102,7 +2104,7 @@ pub trait IsStarkProver<
     ///
     /// The codeword is kept rather than the LDEs it came from. That is the
     /// whole reason this split is affordable: on the ethrex block the trace and
-    /// composition LDEs of all 227 tables are tens of gigabytes, while their
+    /// composition LDEs of every table are tens of gigabytes, while their
     /// DEEP codewords together are about 6.5 GB — one extension element per row
     /// instead of every column. Holding them is what saves walking the
     /// execution again just to recompute them once the coefficient is known.
@@ -2180,7 +2182,6 @@ pub trait IsStarkProver<
             lde_size: domain.interpolation_domain_size * domain.blowup_factor,
             trace_rows: domain.interpolation_domain_size,
             deep,
-            air_index: usize::MAX,
             bus_contribution: round_1_result
                 .bus_public_inputs
                 .as_ref()
@@ -2262,7 +2263,7 @@ pub trait IsStarkProver<
     /// The accumulator is the batch polynomial the spec describes. A member is
     /// added and dropped, so what is held is one codeword per distinct domain
     /// rather than one per table — which on the ethrex block is 13 instead of
-    /// 227, and about a gigabyte instead of eight and a half.
+    /// one per table, and about a gigabyte instead of eight and a half.
     fn accumulate(
         acc: &mut Vec<FieldElement<FieldExtension>>,
         coefficient: &FieldElement<FieldExtension>,