From 0f2f29302e2ba1d196db8f2c7e47e010ef4ba9bf Mon Sep 17 00:00:00 2001 From: diegokingston Date: Tue, 4 Aug 2026 19:02:53 -0300 Subject: [PATCH] perf(prover): pull-cursor CPU scheduler instead of fixed chunks, plus tests MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two follow-ups to the CPU-aware scheduler. **Keep the pull cursor.** The CPU `run_admitted` batched `order` into `chunks(workers)` and ran each chunk to completion, so every table waited for the slowest member of its chunk — the behaviour the doc comment directly above `run_admitted` says `heaviest_first` plus the cursor were introduced to remove. Sorting heaviest-first hides this while the weight estimate is accurate: chunks then hold similar-weight work. It bites when the estimate mispredicts, which is plausible here since the weight is `estimate_table_vram_bytes` — a device footprint used as a proxy for host time. This keeps the original CUDA scheduler's shape (atomic cursor, heaviest-first, concurrency bounded by `workers`) but spawns the workers as Rayon tasks rather than OS threads. That preserves the fix this branch is for — nested per-table Rayon work stays inside the pool and work-steals — while a worker that finishes a small table pulls the next one immediately. `TABLE_PARALLELISM` keeps its meaning and the number of simultaneously live per-table working sets is unchanged. Blocking a pool worker is only acceptable because, unlike the CUDA path, there is no admission gate to wait on. **Add coverage.** `run_admitted` is now three cfg-gated implementations with no tests. These are cfg-agnostic, so they exercise whichever one the build selects: * results land in the slot named by their own index, not by position in `order` (those differ once the order is heaviest-first); * a batch does not serialize when the estimate mispredicts — verified to FAIL at 565ms against the chunked implementation and pass at ~130ms here; * `heaviest_first` is a descending permutation. Not addressed here, to keep the diff reviewable: the CPU staging block is a near-duplicate of the `debug-checks` one, and `VramGate` is still constructed and threaded as `_gate` on CPU rather than not built. --- crypto/stark/src/prover.rs | 127 ++++++++++++++++++++++++++++++++++--- 1 file changed, 119 insertions(+), 8 deletions(-) diff --git a/crypto/stark/src/prover.rs b/crypto/stark/src/prover.rs index efd86b342..2db84781b 100644 --- a/crypto/stark/src/prover.rs +++ b/crypto/stark/src/prover.rs @@ -19,8 +19,6 @@ use math::{ polynomial::Polynomial, }; -#[cfg(all(feature = "parallel", not(feature = "cuda")))] -use rayon::prelude::IntoParallelRefIterator; #[cfg(feature = "parallel")] use rayon::prelude::{IntoParallelIterator, ParallelIterator}; @@ -745,13 +743,40 @@ fn run_admitted( .map(|_| std::sync::Mutex::new(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].lock().unwrap() = Some(result); + // Same pull-cursor as the CUDA scheduler, but the workers are Rayon tasks + // rather than OS threads. That keeps both properties that matter here: + // + // * the nested Rayon work each table runs stays *inside* the pool, so it + // work-steals normally (spawning OS driver threads is what this file's + // CPU path was fixing); + // * a worker that finishes a small table immediately pulls the next one. + // Batching `order` into fixed chunks instead would make every table wait + // for the slowest of its chunk — precisely the behaviour `heaviest_first` + // and the cursor were introduced to remove. + // + // Concurrency stays bounded by `workers`, so `TABLE_PARALLELISM` keeps its + // meaning and the number of simultaneously live per-table working sets is + // unchanged. Blocking is safe to do here only because, unlike the CUDA + // path, there is no admission gate for a worker to block on. + let cursor = std::sync::atomic::AtomicUsize::new(0); + let workers = workers.max(1).min(order.len().max(1)); + rayon::scope(|scope| { + for _ in 0..workers { + let cursor = &cursor; + let results = &results; + let task = &task; + scope.spawn(move |_| { + loop { + let pos = cursor.fetch_add(1, std::sync::atomic::Ordering::Relaxed); + let Some(&idx) = order.get(pos) else { + return; + }; + let out = task(idx); + *results[idx].lock().unwrap() = Some(out); + } + }); } - } + }); results .into_iter() @@ -4076,3 +4101,89 @@ fn print_bus_balance_report( } } } + +#[cfg(test)] +mod scheduler_tests { + use super::{VramGate, heaviest_first, run_admitted}; + + /// Every index must be filled exactly once, and each result must land in the + /// slot named by its own index — not by its position in `order`. The + /// scheduler hands out work heaviest-first, so those two orders differ. + #[test] + fn run_admitted_fills_every_slot_by_index() { + let estimates: Vec = vec![7, 1, 9, 3, 5, 5, 0, 11]; + let order = heaviest_first(&estimates); + let gate = VramGate::new(u64::MAX); + + for workers in [1usize, 2, 3, 8, 64] { + let out = run_admitted(&order, &estimates, &gate, workers, |idx| idx * 10); + assert_eq!(out.len(), estimates.len(), "workers={workers}"); + for (idx, slot) in out.into_iter().enumerate() { + assert_eq!(slot, Some(idx * 10), "slot {idx}, workers={workers}"); + } + } + } + + /// A batch must not serialize when the *weight estimate mispredicts* actual + /// cost. Sorting heaviest-first makes fixed chunks look fine as long as the + /// estimate is accurate — each chunk then holds similar-weight work. The + /// failure mode is a chunk that contains one genuinely slow task the + /// estimate did not flag: everything else in that chunk finishes and waits. + /// + /// Here the estimates are uniform (so the order is the identity) while every + /// 4th task is slow. With `workers = 4` and fixed chunking that is one slow + /// task per chunk and the batch costs 4 x SLOW; a pull cursor runs the four + /// slow tasks concurrently and costs about 1 x SLOW. + #[test] + fn run_admitted_does_not_serialize_when_estimates_mispredict() { + use std::time::{Duration, Instant}; + const N: usize = 16; + const WORKERS: usize = 4; + const SLOW: Duration = Duration::from_millis(120); + const FAST: Duration = Duration::from_millis(5); + + let estimates: Vec = vec![1; N]; + // Explicit identity order: this test is about the scheduler, not about + // how `heaviest_first` happens to break ties. + let order: Vec = (0..N).collect(); + let gate = VramGate::new(u64::MAX); + + let start = Instant::now(); + let out = run_admitted(&order, &estimates, &gate, WORKERS, |idx| { + std::thread::sleep(if idx % WORKERS == WORKERS - 1 { + SLOW + } else { + FAST + }); + idx + }); + let elapsed = start.elapsed(); + + assert_eq!(out.len(), N); + for (idx, slot) in out.into_iter().enumerate() { + assert_eq!(slot, Some(idx)); + } + // Fixed chunking costs >= 4 * 120ms = 480ms. A cursor overlaps the slow + // tasks and lands near 120ms; 300ms leaves generous CI headroom while + // still failing the chunked behaviour. + assert!( + elapsed < Duration::from_millis(300), + "took {elapsed:?}; a fixed-chunk scheduler would need >= 480ms here" + ); + } + + #[test] + fn heaviest_first_orders_by_descending_estimate() { + let estimates: Vec = vec![3, 10, 1, 10, 4]; + let order = heaviest_first(&estimates); + assert_eq!(order.len(), estimates.len()); + let weights: Vec = order.iter().map(|&i| estimates[i]).collect(); + assert!( + weights.windows(2).all(|w| w[0] >= w[1]), + "not descending: {weights:?}" + ); + let mut seen = order.clone(); + seen.sort_unstable(); + assert_eq!(seen, (0..estimates.len()).collect::>()); + } +}