diff --git a/Makefile b/Makefile index 3e4a88ecb..be92871a4 100644 --- a/Makefile +++ b/Makefile @@ -666,7 +666,7 @@ check: clippy: cargo clippy --workspace --all-targets -- -D warnings -A clippy::op_ref cargo clippy --workspace --all-targets --no-default-features --features lambda-vm-prover/debug-checks -- -D warnings -A clippy::op_ref - cargo clippy --workspace --all-targets --features lambda-vm-prover/disk-spill -- -D warnings -A clippy::op_ref + cargo clippy --workspace --all-targets --features lambda-vm-prover/disk-spill,lambda-vm-prover/instruments -- -D warnings -A clippy::op_ref fmt: cargo fmt --all @@ -676,7 +676,9 @@ lint: cargo fmt --check --all cargo clippy --workspace --all-targets -- -D warnings -A clippy::op_ref cargo clippy --workspace --all-targets --no-default-features --features lambda-vm-prover/debug-checks -- -D warnings -A clippy::op_ref - cargo clippy --workspace --all-targets --features lambda-vm-prover/disk-spill -- -D warnings -A clippy::op_ref + # `instruments` rides on this pass rather than adding a fifth one: it gates the per-table + # timing plumbing, which no other pass compiles, so breakage there used to reach main. + cargo clippy --workspace --all-targets --features lambda-vm-prover/disk-spill,lambda-vm-prover/instruments -- -D warnings -A clippy::op_ref # The cuda feature gates whole modules + cuda-only integration tests. build.rs emits empty # cubin stubs when nvcc is absent, so this checks on a GPU-less host (CI lint runner, dev laptop) # too — no GPU required. Catches cuda-gated breakage that the non-cuda passes above miss. diff --git a/crypto/stark/src/instruments.rs b/crypto/stark/src/instruments.rs index 0f68059f4..6ef06a387 100644 --- a/crypto/stark/src/instruments.rs +++ b/crypto/stark/src/instruments.rs @@ -362,8 +362,15 @@ pub fn take_r1_sub() -> Round1SubOps { /// Note: thread local stores (R2_SUB, R4_SUB, ROUND_SUB_OPS) are only cleared /// for the calling thread. Rayon worker threads are not reset, so stale data is /// possible if a previous run panicked without consuming stored values. -/// In practice this is safe because store/take pairs always execute within the -/// same rayon task closure. +/// +/// Each store/take pair must also stay adjacent on one thread, with no rayon +/// region in between. The CPU table scheduler runs a whole table's task as a +/// rayon item, so a worker that blocks inside one table's nested parallelism +/// can run a sibling table's task on top of that frame, consume this thread's +/// slot, and leave the outer table reporting zeros. Measured with a probe: no +/// loss while the tables in flight stay under the pool's thread count (the +/// default `k` does), 7-10 % of them lost once they exceed it — e.g. with +/// `TABLE_PARALLELISM` raised past the core count. pub fn reset_all() { R1_MAIN_LDE_US.store(0, Ordering::Relaxed); R1_MAIN_MERKLE_US.store(0, Ordering::Relaxed); diff --git a/crypto/stark/src/prover.rs b/crypto/stark/src/prover.rs index 5078ce290..6aee481b9 100644 --- a/crypto/stark/src/prover.rs +++ b/crypto/stark/src/prover.rs @@ -19,6 +19,8 @@ use math::{ polynomial::Polynomial, }; +#[cfg(all(feature = "parallel", not(feature = "cuda")))] +use rayon::prelude::IntoParallelRefIterator; #[cfg(feature = "parallel")] use rayon::prelude::{IntoParallelIterator, ParallelIterator}; @@ -700,17 +702,25 @@ fn estimate_table_vram_bytes(main_cols: usize, aux_cols: usize, lde_size: usize) /// Only OS driver threads block here (see `run_admitted`) — never rayon /// workers, whose pool the admitted tables use internally and which a /// blocked worker would starve. +/// +/// Device-only mechanism: CPU builds have no device allocation to admit +/// against, so their `run_admitted` arms take the gate and ignore it. That is +/// what the `allow(dead_code)` attributes here are for — `multi_prove` still +/// constructs the gate unconditionally so the three arms share one signature. +#[cfg_attr(not(feature = "cuda"), allow(dead_code))] struct VramGate { used: std::sync::Mutex, freed: std::sync::Condvar, budget: u64, } +#[cfg_attr(not(feature = "cuda"), allow(dead_code))] struct VramPermit<'a> { gate: &'a VramGate, bytes: u64, } +#[cfg_attr(not(feature = "cuda"), allow(dead_code))] impl VramGate { fn new(budget: u64) -> Self { Self { @@ -741,11 +751,32 @@ impl Drop for VramPermit<'_> { } } +/// Debug-only contract check for `run_admitted`: the result slots are addressed +/// by `order`'s *values*, so every value must index `estimates` and appear at +/// most once. Callers pass either a full permutation or one group of one, and a +/// grouping bug would otherwise surface as a confusing `take()` panic inside a +/// task rather than here. +fn debug_check_order(order: &[usize], estimates: &[u64]) { + if cfg!(debug_assertions) { + let mut seen = vec![false; estimates.len()]; + for &idx in order { + assert!( + idx < estimates.len(), + "run_admitted: order names table {idx}, estimates has {}", + estimates.len() + ); + assert!(!seen[idx], "run_admitted: order names table {idx} twice"); + seen[idx] = true; + } + } +} + /// Run `task` once per table index on `workers` OS driver threads, admitting /// each index through `gate` with its estimated bytes. `order` fixes the /// start order (heaviest table first, so the long pole starts early and small /// tables fill around it — the fixed chunks this replaces made every table /// wait for the slowest of its chunk). Returns one slot per original index. +#[cfg(feature = "cuda")] fn run_admitted( order: &[usize], estimates: &[u64], @@ -753,6 +784,7 @@ fn run_admitted( workers: usize, task: impl Fn(usize) -> T + Sync, ) -> Vec> { + debug_check_order(order, estimates); let results: Vec>> = estimates .iter() .map(|_| std::sync::Mutex::new(None)) @@ -781,6 +813,54 @@ fn run_admitted( .collect() } +/// CPU version of the table scheduler. There is no device admission to wait +/// on here, and the work underneath each table already uses Rayon internally. +/// Keep the outer scheduling inside the same Rayon pool instead of creating +/// driver OS threads: otherwise every table task launches nested Rayon work +/// from outside the pool and loses the work-stealing behavior of the original +/// phase scheduler. +#[cfg(all(not(feature = "cuda"), feature = "parallel"))] +fn run_admitted( + order: &[usize], + estimates: &[u64], + _gate: &VramGate, + workers: usize, + task: impl Fn(usize) -> T + Sync, +) -> Vec> { + debug_check_order(order, estimates); + // No interior mutability here, unlike the `cuda` arm: `collect()` joins + // each chunk, so every slot is written from this thread. + let mut results: Vec> = (0..estimates.len()).map(|_| None).collect(); + + for chunk in order.chunks(workers.max(1)) { + let chunk_results: Vec<(usize, T)> = + chunk.par_iter().map(|&idx| (idx, task(idx))).collect(); + for (idx, result) in chunk_results { + results[idx] = Some(result); + } + } + + results +} + +/// Sequential fallback when the prover is built without its default +/// `parallel` feature. +#[cfg(all(not(feature = "cuda"), not(feature = "parallel")))] +fn run_admitted( + order: &[usize], + estimates: &[u64], + _gate: &VramGate, + _workers: usize, + task: impl Fn(usize) -> T + Sync, +) -> Vec> { + debug_check_order(order, estimates); + let mut results: Vec> = (0..estimates.len()).map(|_| None).collect(); + for &idx in order { + results[idx] = Some(task(idx)); + } + results +} + /// Table indices sorted heaviest-first by estimate. fn heaviest_first(estimates: &[u64]) -> Vec { let mut order: Vec = (0..estimates.len()).collect(); @@ -4483,6 +4563,15 @@ pub trait IsStarkProver< &boundary_coefficients, )?; + // Round 2's sub-op timings travel in a thread-local, so take them here + // instead of after round 4: this thread re-enters rayon in rounds 3 and + // 4, and a worker that blocks there can run another table's whole + // scheduler task on top of this frame and consume the slot. Keeping + // store and take adjacent is what pins the pair to one thread — see + // `instruments::reset_all`. + #[cfg(feature = "instruments")] + let r2_sub = crate::instruments::take_r2_sub(); + // >>>> Send commitments: [H₁], [H₂] transcript.append_bytes(&round_2_result.composition_poly_root); @@ -4579,8 +4668,7 @@ pub trait IsStarkProver< #[cfg(feature = "instruments")] { let zero = Duration::ZERO; - let (r2_constraints, r2_fft, r2_merkle) = - crate::instruments::take_r2_sub().unwrap_or((zero, zero, zero)); + let (r2_constraints, r2_fft, r2_merkle) = r2_sub.unwrap_or((zero, zero, zero)); 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 { @@ -4740,3 +4828,73 @@ fn print_bus_balance_report( } } } + +#[cfg(test)] +mod scheduler_tests { + use super::*; + use std::sync::atomic::{AtomicUsize, Ordering}; + + /// Runs `run_admitted` over `order` with `n` tables and returns how many + /// times each table's task ran, plus the slots it wrote. + fn calls_and_results( + order: &[usize], + n: usize, + workers: usize, + ) -> (Vec, Vec>) { + let estimates = vec![1u64; n]; + let gate = VramGate::new(u64::MAX); + let calls: Vec = (0..n).map(|_| AtomicUsize::new(0)).collect(); + let out = run_admitted(order, &estimates, &gate, workers, |idx| { + calls[idx].fetch_add(1, Ordering::Relaxed); + idx + }); + ( + calls.iter().map(|c| c.load(Ordering::Relaxed)).collect(), + out, + ) + } + + /// The scheduler's contract, independent of how many tables run at once: + /// each table's task runs exactly once and its result lands on its own + /// index. Nothing else pins this, and a chunking or sizing bug would only + /// show up on hosts with a particular core count. + #[test] + fn every_table_runs_once_and_lands_on_its_own_slot() { + for n in [1usize, 2, 5, 21] { + // Heaviest-first, which is the order `multi_prove` passes. + let estimates: Vec = (0..n).map(|i| (n - i) as u64).collect(); + let order = heaviest_first(&estimates); + for workers in [1usize, 2, 3, n, n + 5] { + let (calls, out) = calls_and_results(&order, n, workers); + assert!( + calls.iter().all(|&c| c == 1), + "n={n} workers={workers} calls={calls:?}" + ); + for (idx, slot) in out.iter().enumerate() { + assert_eq!(*slot, Some(idx), "n={n} workers={workers}"); + } + } + } + } + + /// What the `Staged(depth)` arm passes once the barrier runs over groups: + /// a subset of the permutation. The untouched tables must stay empty. + #[test] + fn a_group_fills_only_its_own_slots() { + let n = 9; + let group = [7usize, 1, 4]; + let (calls, out) = calls_and_results(&group, n, 2); + for idx in 0..n { + let ran = group.contains(&idx); + assert_eq!(calls[idx], ran as usize, "table {idx}"); + assert_eq!(out[idx], ran.then_some(idx), "table {idx}"); + } + } + + #[test] + fn an_empty_order_runs_nothing() { + let (calls, out) = calls_and_results(&[], 4, 3); + assert!(calls.iter().all(|&c| c == 0)); + assert!(out.iter().all(|s| s.is_none())); + } +}