From 99c3a1173754a88d2dcb2a22e729763834f0a611 Mon Sep 17 00:00:00 2001 From: Joaquin Carletti Date: Wed, 29 Jul 2026 14:53:51 -0300 Subject: [PATCH 01/12] new opt --- crypto/math-cuda/kernels/fri.cu | 17 ++ crypto/math-cuda/kernels/inverse.cu | 35 +++ crypto/math-cuda/kernels/keccak.cu | 37 +++- crypto/math-cuda/kernels/ntt.cu | 57 +++++ crypto/math-cuda/src/device.rs | 49 +++++ crypto/math-cuda/src/fri.rs | 157 ++++++++------ crypto/math-cuda/src/inverse.rs | 69 ++++-- crypto/math-cuda/src/lde.rs | 120 ++++++++--- crypto/math-cuda/src/merkle.rs | 70 ++++++ crypto/stark/src/fri/fri_commitment.rs | 7 + crypto/stark/src/fri/mod.rs | 2 +- crypto/stark/src/gpu_lde.rs | 194 ++++++++++++++--- crypto/stark/src/logup_gpu.rs | 17 +- crypto/stark/src/lookup.rs | 23 +- crypto/stark/src/prover.rs | 287 +++++++++++++++++++------ 15 files changed, 919 insertions(+), 222 deletions(-) diff --git a/crypto/math-cuda/kernels/fri.cu b/crypto/math-cuda/kernels/fri.cu index 63d72cef1..bcc8f9e40 100644 --- a/crypto/math-cuda/kernels/fri.cu +++ b/crypto/math-cuda/kernels/fri.cu @@ -59,3 +59,20 @@ extern "C" __global__ void fri_update_twiddles( uint64_t old = tw_in[2 * j]; tw_out[j] = goldilocks::mul(old, old); } + +// Gather interleaved ext3 elements at arbitrary positions: one thread per +// query copies evals[positions[i]] (3 u64) into out[i]. Serves the FRI query +// phase's symmetric-eval reads off the resident layer buffers. +extern "C" __global__ void gather_ext3_at( + const uint64_t *evals, + const uint32_t *positions, + uint64_t q, + uint64_t *out +) { + uint64_t i = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x; + if (i >= q) return; + uint64_t p = positions[i]; + out[i * 3] = evals[p * 3]; + out[i * 3 + 1] = evals[p * 3 + 1]; + out[i * 3 + 2] = evals[p * 3 + 2]; +} diff --git a/crypto/math-cuda/kernels/inverse.cu b/crypto/math-cuda/kernels/inverse.cu index 4ee228d8c..65d54afc0 100644 --- a/crypto/math-cuda/kernels/inverse.cu +++ b/crypto/math-cuda/kernels/inverse.cu @@ -309,3 +309,38 @@ extern "C" __global__ void batch_inverse_combine_ext3( out_base[1] = res.b; out_base[2] = res.c; } + +// --------------------------------------------------------------------------- +// 7. invert_total_ext3 +// +// One-thread Fermat inversion of the scan total: out = src[n-1]^(p^3 - 2). +// Replaces the host round-trip (D2H + host Fermat + H2D + stream sync) so the +// whole batch inverse stays stream-ordered. The 192-bit exponent arrives as +// three little-endian u64 limbs. +// --------------------------------------------------------------------------- +extern "C" __global__ void invert_total_ext3( + const uint64_t *src, // 3 * n u64 (reads element n-1) + uint64_t n, + uint64_t e0, // exponent limbs, little-endian + uint64_t e1, + uint64_t e2, + uint64_t *out // 3 u64 +) { + if (blockIdx.x != 0 || threadIdx.x != 0) return; + const uint64_t *base = src + (n - 1) * 3; + ext3::Fe3 a = {base[0], base[1], base[2]}; + ext3::Fe3 r = ext3::one(); + uint64_t limbs[3] = {e0, e1, e2}; + for (int li = 2; li >= 0; --li) { + uint64_t bits = limbs[li]; + for (int b = 63; b >= 0; --b) { + r = ext3::mul(r, r); + if ((bits >> b) & 1) { + r = ext3::mul(r, a); + } + } + } + out[0] = r.a; + out[1] = r.b; + out[2] = r.c; +} diff --git a/crypto/math-cuda/kernels/keccak.cu b/crypto/math-cuda/kernels/keccak.cu index 7b62789f9..b026ff2b6 100644 --- a/crypto/math-cuda/kernels/keccak.cu +++ b/crypto/math-cuda/kernels/keccak.cu @@ -366,13 +366,11 @@ extern "C" __global__ void keccak_fri_leaves_ext3( // concatenation of two 32-byte siblings, identical to // `FieldElementVectorBackend::hash_new_parent` on host. // --------------------------------------------------------------------------- -extern "C" __global__ void keccak_merkle_level( +__device__ __forceinline__ void hash_merkle_parent( uint8_t *nodes, uint64_t parent_begin, // node index (counted in 32-byte nodes) - uint64_t n_pairs) { - uint64_t tid = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x; - if (tid >= n_pairs) return; - + uint64_t n_pairs, + uint64_t tid) { uint64_t st[25]; #pragma unroll for (int i = 0; i < 25; ++i) st[i] = 0; @@ -393,6 +391,35 @@ extern "C" __global__ void keccak_merkle_level( finalize_keccak256(st, rate_pos, nodes + (parent_begin + tid) * 32); } +extern "C" __global__ void keccak_merkle_level( + uint8_t *nodes, + uint64_t parent_begin, // node index (counted in 32-byte nodes) + uint64_t n_pairs) { + uint64_t tid = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x; + if (tid >= n_pairs) return; + hash_merkle_parent(nodes, parent_begin, n_pairs, tid); +} + +// Build every remaining level (from `level_begin` up to the root) in ONE +// single-block launch: each level's pairs are grid-strided over the block, +// with a __syncthreads() barrier between levels. Replaces log2 launches of +// `keccak_merkle_level` for the small top levels of the tree, whose per-level +// work is dwarfed by launch overhead. +extern "C" __global__ void keccak_merkle_tail( + uint8_t *nodes, + uint64_t level_begin) { + uint64_t lb = level_begin; + while (lb != 0) { + uint64_t nb = lb / 2; + uint64_t n_pairs = lb - nb; + for (uint64_t tid = threadIdx.x; tid < n_pairs; tid += blockDim.x) { + hash_merkle_parent(nodes, nb, n_pairs, tid); + } + __syncthreads(); + lb = nb; + } +} + // Gather Merkle authentication paths for a batch of leaf positions, reading the // resident tree `nodes` (32-byte nodes; layout: inner nodes [0..leaves_len-1], // root at 0, leaves at [leaves_len-1..]). One thread per query walks leaf->root, diff --git a/crypto/math-cuda/kernels/ntt.cu b/crypto/math-cuda/kernels/ntt.cu index 13c1af688..35a4f7b20 100644 --- a/crypto/math-cuda/kernels/ntt.cu +++ b/crypto/math-cuda/kernels/ntt.cu @@ -411,3 +411,60 @@ extern "C" __global__ void matrix_transpose_strided( __syncthreads(); } } + +// First-8-levels fused DIT on row-major data: one block stages 256 consecutive +// rows x blockDim.x columns in shmem and runs levels 0..min(8,log_n) with +// __syncthreads between levels (row-major analog of ntt_dit_8_levels_batched +// with base_step == 0, whose twiddle math this reuses verbatim). Grid: +// x = column tiles, y = n/256 row blocks. Requires n >= 256. Shmem tile is +// padded (pitch = T+1) to break bank conflicts on the butterfly accesses. +extern "C" __global__ void ntt_dit_8_levels_row_major(uint64_t *data, + const uint64_t *tw, + uint64_t n, + uint64_t log_n, + uint64_t m) +{ + extern __shared__ uint64_t tile[]; + uint32_t T = blockDim.x; + uint32_t pitch = T + 1; + uint64_t col = (uint64_t)blockIdx.x * T + threadIdx.x; + bool live = col < m; + uint64_t row_base = (uint64_t)blockIdx.y * 256; + + for (uint32_t r = threadIdx.y; r < 256; r += blockDim.y) { + if (live) tile[r * pitch + threadIdx.x] = data[(row_base + r) * m + col]; + } + __syncthreads(); + + uint32_t n_loc_steps = (uint32_t)min((uint64_t)8, log_n); + uint32_t remaining_high_bits = (uint32_t)(log_n - 1); + uint32_t high_mask = (1u << remaining_high_bits) - 1u; + + for (uint32_t loc_step = 0; loc_step < n_loc_steps; ++loc_step) { + for (uint32_t i = threadIdx.y; i < 128; i += blockDim.y) { + uint32_t half = 1u << loc_step; + uint32_t grp = i >> loc_step; + uint32_t grp_pos = i & (half - 1); + uint32_t idx1 = (grp << (loc_step + 1)) + grp_pos; + uint32_t idx2 = idx1 + half; + + uint32_t gs = loc_step; + uint32_t ggp = ((uint32_t)blockIdx.y << 7) + i; + ggp = (ggp & high_mask) + (ggp >> remaining_high_bits); + ggp = ggp & ((1u << gs) - 1u); + uint64_t factor = tw[(uint64_t)ggp * (n >> (gs + 1))]; + + if (live) { + uint64_t u = tile[idx1 * pitch + threadIdx.x]; + uint64_t v = mul(tile[idx2 * pitch + threadIdx.x], factor); + tile[idx1 * pitch + threadIdx.x] = add(u, v); + tile[idx2 * pitch + threadIdx.x] = sub(u, v); + } + } + __syncthreads(); + } + + for (uint32_t r = threadIdx.y; r < 256; r += blockDim.y) { + if (live) data[(row_base + r) * m + col] = tile[r * pitch + threadIdx.x]; + } +} diff --git a/crypto/math-cuda/src/device.rs b/crypto/math-cuda/src/device.rs index 8fd7f13de..292026401 100644 --- a/crypto/math-cuda/src/device.rs +++ b/crypto/math-cuda/src/device.rs @@ -186,6 +186,7 @@ pub struct Backend { // row-major NTT kernels pub bit_reverse_row_major: CudaFunction, pub ntt_dit_level_row_major: CudaFunction, + pub ntt_dit_8_levels_row_major: CudaFunction, pub pointwise_mul_row_major: CudaFunction, pub matrix_transpose_strided: CudaFunction, @@ -198,6 +199,7 @@ pub struct Backend { pub keccak_comp_poly_leaves_ext3: CudaFunction, pub keccak_fri_leaves_ext3: CudaFunction, pub keccak_merkle_level: CudaFunction, + pub keccak_merkle_tail: CudaFunction, pub merkle_gather_paths: CudaFunction, // barycentric.cubin @@ -214,6 +216,7 @@ pub struct Backend { // fri.cubin pub fri_fold_ext3: CudaFunction, + pub gather_ext3_at: CudaFunction, pub fri_update_twiddles: CudaFunction, // inverse.cubin @@ -223,6 +226,7 @@ pub struct Backend { pub block_inclusive_scan_rev_ext3: CudaFunction, pub apply_block_offsets_rev_ext3: CudaFunction, pub batch_inverse_combine_ext3: CudaFunction, + pub invert_total_ext3: CudaFunction, pub logup_fingerprint_ext3: CudaFunction, pub logup_term_ext3: CudaFunction, pub logup_row_sum_ext3: CudaFunction, @@ -412,6 +416,7 @@ impl Backend { scalar_mul_batched: ntt.load_function("scalar_mul_batched")?, bit_reverse_row_major: ntt.load_function("bit_reverse_row_major")?, ntt_dit_level_row_major: ntt.load_function("ntt_dit_level_row_major")?, + ntt_dit_8_levels_row_major: ntt.load_function("ntt_dit_8_levels_row_major")?, pointwise_mul_row_major: ntt.load_function("pointwise_mul_row_major")?, matrix_transpose_strided: ntt.load_function("matrix_transpose_strided")?, keccak256_leaves_base_row_major_row_pair: keccak @@ -425,6 +430,7 @@ impl Backend { keccak_comp_poly_leaves_ext3: keccak.load_function("keccak_comp_poly_leaves_ext3")?, keccak_fri_leaves_ext3: keccak.load_function("keccak_fri_leaves_ext3")?, keccak_merkle_level: keccak.load_function("keccak_merkle_level")?, + keccak_merkle_tail: keccak.load_function("keccak_merkle_tail")?, merkle_gather_paths: keccak.load_function("merkle_gather_paths")?, barycentric_base_batched: bary.load_function("barycentric_base_batched")?, barycentric_ext3_batched: bary.load_function("barycentric_ext3_batched")?, @@ -437,6 +443,7 @@ impl Backend { deep_composition_ext3_row: deep.load_function("deep_composition_ext3_row")?, bit_reverse_ext3_kernel: deep.load_function("bit_reverse_ext3_interleaved")?, fri_fold_ext3: fri.load_function("fri_fold_ext3")?, + gather_ext3_at: fri.load_function("gather_ext3_at")?, fri_update_twiddles: fri.load_function("fri_update_twiddles")?, compute_denoms_ext3: inverse.load_function("compute_denoms_ext3")?, block_inclusive_scan_fwd_ext3: inverse @@ -446,6 +453,7 @@ impl Backend { .load_function("block_inclusive_scan_rev_ext3")?, apply_block_offsets_rev_ext3: inverse.load_function("apply_block_offsets_rev_ext3")?, batch_inverse_combine_ext3: inverse.load_function("batch_inverse_combine_ext3")?, + invert_total_ext3: inverse.load_function("invert_total_ext3")?, logup_fingerprint_ext3: logup.load_function("logup_fingerprint_ext3")?, logup_term_ext3: logup.load_function("logup_term_ext3")?, logup_row_sum_ext3: logup.load_function("logup_row_sum_ext3")?, @@ -631,6 +639,47 @@ impl Drop for PendingD2H<'_> { /// stream-ordered on its own stream, so a `src` allocated on `stream` may be /// dropped after this call — the free queues behind the copy. Do NOT pass a /// `src` owned by a *different* stream and drop it before waiting. +/// Host→device copy staged through the pinned slot: one host memcpy into +/// pinned memory + one async DMA, instead of the driver's internal pageable +/// staging (small chunks; 2-3x slower for multi-hundred-MB traces and it +/// convoys under multi-thread load). Blocks until the DMA lands, so the slot +/// and `src_host` are both reusable on return. +pub fn htod_via( + stream: &Arc, + slot: &Mutex, + ctx: &CudaContext, + src_host: &[T], + dst: &mut cudarc::driver::CudaViewMut<'_, T>, +) -> Result<()> { + use cudarc::driver::DevicePtrMut; + let n_bytes = std::mem::size_of_val(src_host); + let u64_len = n_bytes.div_ceil(8); + let mut staging = slot.lock().unwrap(); + staging.ensure_capacity(u64_len, ctx)?; + ctx.bind_to_thread()?; + // SAFETY: the pinned allocation is stable while the lock is held and at + // least `n_bytes` long (`ensure_capacity`). The DMA reads it after the + // host memcpy (program order); `device_ptr_mut` orders the device write + // on `stream`. + unsafe { + std::ptr::copy_nonoverlapping( + src_host.as_ptr() as *const u8, + staging.ptr as *mut u8, + n_bytes, + ); + let (dst_ptr, _record) = dst.device_ptr_mut(stream); + cudarc::driver::sys::cuMemcpyHtoDAsync_v2( + dst_ptr, + staging.ptr as *const core::ffi::c_void, + n_bytes, + stream.cu_stream(), + ) + .result()?; + } + staging.record_event(stream)?; + staging.sync_event() +} + pub fn async_dtoh_via<'a, T: cudarc::driver::DeviceRepr>( stream: &Arc, slot: &'a Mutex, diff --git a/crypto/math-cuda/src/fri.rs b/crypto/math-cuda/src/fri.rs index 8a477e1ee..8da80472d 100644 --- a/crypto/math-cuda/src/fri.rs +++ b/crypto/math-cuda/src/fri.rs @@ -40,23 +40,19 @@ fn check_fault_injection() -> Result<()> { Ok(()) } -/// Device-side state across FRI commit iterations. Owns two ext3 eval -/// buffers (flip-flopped as layer input / output) and the inv_twiddles -/// buffer. Freed when dropped. +/// Device-side state across FRI commit iterations. Owns the current fold +/// input (the previous layer's evals, Arc-shared with that layer's retained +/// `gpu_evals`) and the inv_twiddles buffer. Freed when dropped. pub struct FriCommitState { pub stream: Arc, - // Ping-pong evaluation buffers. Both sized `3 * n0` u64 at init. Each - // successive fold uses half the space. Cheap to pre-allocate vs. per- - // layer alloc. - evals_a: CudaSlice, - evals_b: CudaSlice, + /// Current fold input. Each fold allocates a fresh output buffer that is + /// both returned to the caller (kept resident for the query phase) and + /// becomes the next fold's input. + current: Arc>, /// Base-field inv_twiddles; `n0 / 2` u64 at init, halved each layer. inv_tw: CudaSlice, - /// Number of ext3 elements in the buffer currently acting as fold input - /// (`evals_a` or `evals_b`, selected by `a_is_input`). + /// Number of ext3 elements in `current`. pub current_n: usize, - /// Which buffer holds the current layer's input. Toggles each fold. - a_is_input: bool, } impl FriCommitState { @@ -71,20 +67,16 @@ impl FriCommitState { let be = backend()?; let stream = be.next_stream(); - // SAFETY: every byte of evals_a is overwritten by the H2D below. - // evals_b is written by the first fold before it is read. - let mut evals_a = unsafe { stream.alloc::(3 * n0) }?; - let evals_b = unsafe { stream.alloc::(3 * n0) }?; - stream.memcpy_htod(evals_host, &mut evals_a)?; + // SAFETY: every byte of evals is overwritten by the H2D below. + let mut evals = unsafe { stream.alloc::(3 * n0) }?; + stream.memcpy_htod(evals_host, &mut evals)?; let inv_tw = stream.clone_htod(inv_tw_host)?; Ok(Self { stream, - evals_a, - evals_b, + current: Arc::new(evals), inv_tw, current_n: n0, - a_is_input: true, }) } @@ -96,31 +88,32 @@ impl FriCommitState { assert_eq!(buf.len(), 3 * n); assert_eq!(inv_tw_host.len(), n / 2); - // SAFETY: evals_b is written by the first fold before it is read. - let evals_b = unsafe { stream.alloc::(3 * n) }?; let inv_tw = stream.clone_htod(inv_tw_host)?; Ok(Self { stream, - evals_a: buf, - evals_b, + current: Arc::new(buf), inv_tw, current_n: n, - a_is_input: true, }) } - /// Fold the current layer using `zeta`, run the row-pair Keccak leaves - /// + pair-hash Merkle tree kernels on the result, and D2H: - /// - the new root (32 bytes) - /// - the new layer's evals (3 * (current_n / 2) u64s) - /// - the new layer's Merkle tree nodes (standard layout, byte-packed) + /// Fold the current layer using `zeta`, run the row-pair Keccak leaves and + /// pair-hash Merkle tree kernels on the result, and return the layer's + /// evals — device-resident Arc, plus a host copy only when `want_host` — + /// with its resident Merkle tree (root D2H'd, 32 bytes). /// /// Also advances the internal twiddle factors for the next layer. + #[allow(clippy::type_complexity)] pub fn fold_and_commit_layer( &mut self, zeta_raw: [u64; 3], - ) -> Result<(Vec, crate::lde::GpuMerkleTree)> { + want_host: bool, + ) -> Result<( + Option>, + Arc>, + crate::lde::GpuMerkleTree, + )> { #[cfg(feature = "test-faults")] check_fault_injection()?; let be = backend()?; @@ -147,15 +140,11 @@ impl FriCommitState { }; let n_out_u64 = n_out as u64; - // Split the eval buffers into (input, output) based on a_is_input. - // Disjoint-field borrow is fine since evals_a and evals_b are - // separate fields. - let (input_evals, output_evals): (&CudaSlice, &mut CudaSlice) = if self.a_is_input - { - (&self.evals_a, &mut self.evals_b) - } else { - (&self.evals_b, &mut self.evals_a) - }; + // Fresh output buffer per layer: it is retained by the caller for the + // query phase and becomes the next fold's input. + // SAFETY: the fold kernel writes all 3 * n_out slots before any read. + let mut out = unsafe { self.stream.alloc::(3 * n_out) }?; + let input_evals: &CudaSlice = self.current.as_ref(); unsafe { self.stream .launch_builder(&be.fri_fold_ext3) @@ -163,7 +152,7 @@ impl FriCommitState { .arg(&n_out_u64) .arg(&self.inv_tw) .arg(&zeta_dev) - .arg(output_evals) + .arg(&mut out) .launch(cfg)?; } @@ -182,17 +171,10 @@ impl FriCommitState { block_dim: (128, 1, 1), shared_mem_bytes: 0, }; - // Leaves read from the layer's OUTPUT eval buffer (the buffer - // we just wrote to above). - let output_evals: &CudaSlice = if self.a_is_input { - &self.evals_b - } else { - &self.evals_a - }; unsafe { self.stream .launch_builder(&be.keccak_fri_leaves_ext3) - .arg(output_evals) + .arg(&out) .arg(&num_leaves_u64) .arg(&mut leaves_view) .launch(kcfg)?; @@ -225,39 +207,40 @@ impl FriCommitState { self.inv_tw = tw_out; } - // Sync and D2H. - self.stream.synchronize()?; - - // Layer evals: 3 * n_out u64 from the output buffer, staged through - // the per-worker pinned slab (async DMA) instead of a blocking - // pageable copy. The wait is deferred past the root copy below. + // Layer evals to host only when a host copy is wanted (fallback + // consumers), staged through the per-worker pinned slab (async DMA); + // the wait is deferred past the root copy below. let n_evals = 3 * n_out; - let pending = { - let output_evals: &CudaSlice = if self.a_is_input { - &self.evals_b - } else { - &self.evals_a - }; - crate::device::async_dtoh_via( + let pending = if want_host { + Some(crate::device::async_dtoh_via( &self.stream, be.pinned_staging(), &be.ctx, - output_evals, + &out, n_evals, - )? + )?) + } else { + None }; // Keep the layer tree resident on device; copy only the 32-byte root so // R4 query openings gather paths on device instead of copying the tree. - // This pageable copy drains the stream (including the evals DMA above), + // This pageable copy drains the stream (including any evals DMA above), // so the pending wait after it is instant — one block covers both. let mut root = [0u8; 32]; self.stream .memcpy_dtoh(&nodes_dev.slice(0..32), &mut root)?; - let mut layer_evals = vec![0u64; n_evals]; - pending.wait_into_u64(&mut layer_evals)?; + let layer_evals = match pending { + Some(p) => { + let mut v = vec![0u64; n_evals]; + p.wait_into_u64(&mut v)?; + Some(v) + } + None => None, + }; - self.a_is_input = !self.a_is_input; + let out = Arc::new(out); + self.current = Arc::clone(&out); self.current_n = n_out; let tree = crate::lde::GpuMerkleTree { @@ -265,6 +248,42 @@ impl FriCommitState { leaves_len: num_leaves, root, }; - Ok((layer_evals, tree)) + Ok((layer_evals, out, tree)) + } +} + +/// Gather interleaved ext3 elements at `positions` from a resident evals +/// buffer — a small D2H of only the queried values (the FRI query phase's +/// `evaluation[index ^ 1]` reads). +pub fn gather_ext3_at( + evals: &CudaSlice, + positions: &[u32], + stream: &Arc, +) -> Result> { + let q = positions.len(); + if q == 0 { + return Ok(Vec::new()); + } + let be = backend()?; + let pos_dev = stream.clone_htod(positions)?; + // SAFETY: the gather kernel writes all 3 * q slots. + let mut out_dev = unsafe { stream.alloc::(3 * q) }?; + let cfg = LaunchConfig { + grid_dim: ((q as u32).div_ceil(128), 1, 1), + block_dim: (128, 1, 1), + shared_mem_bytes: 0, + }; + let q_u64 = q as u64; + unsafe { + stream + .launch_builder(&be.gather_ext3_at) + .arg(evals) + .arg(&pos_dev) + .arg(&q_u64) + .arg(&mut out_dev) + .launch(cfg)?; } + let out = stream.clone_dtoh(&out_dev)?; + stream.synchronize()?; + Ok(out) } diff --git a/crypto/math-cuda/src/inverse.rs b/crypto/math-cuda/src/inverse.rs index a59c3950c..4dd49556c 100644 --- a/crypto/math-cuda/src/inverse.rs +++ b/crypto/math-cuda/src/inverse.rs @@ -78,12 +78,56 @@ pub fn batch_inverse_ext3(a: &[u64]) -> Result> { Ok(out) } +/// `p^3 - 2` as little-endian u64 limbs: the Fermat exponent for inversion in +/// the Goldilocks cubic extension (`|F_{p^3}^*| = p^3 - 1`). +const EXT3_FERMAT_EXP: [u64; 3] = ext3_fermat_exponent(); + +const fn ext3_fermat_exponent() -> [u64; 3] { + const P: u128 = 0xFFFF_FFFF_0000_0001; + let p2 = P * P; + let m0 = ((p2 as u64) as u128) * P; + let m1 = (p2 >> 64) * P + (m0 >> 64); + let l0 = m0 as u64; + // p^3 mod 2^64 ends in ...0001, so subtracting 2 never borrows. + assert!(l0 >= 2); + [l0 - 2, m1 as u64, (m1 >> 64) as u64] +} + +/// One-thread Fermat inversion of `src[n-1]` into `out[0..3]`, stream-ordered. +fn launch_invert_total( + stream: &Arc, + be: &crate::device::Backend, + src: &CudaSlice, + n: usize, + out: &mut CudaSlice, +) -> Result<()> { + let cfg = LaunchConfig { + grid_dim: (1, 1, 1), + block_dim: (1, 1, 1), + shared_mem_bytes: 0, + }; + let n_u64 = n as u64; + let [e0, e1, e2] = EXT3_FERMAT_EXP; + unsafe { + stream + .launch_builder(&be.invert_total_ext3) + .arg(src) + .arg(&n_u64) + .arg(&e0) + .arg(&e1) + .arg(&e2) + .arg(&mut *out) + .launch(cfg)?; + } + Ok(()) +} + /// Device-input batch inverse. Allocates and returns a fresh `CudaSlice` /// of length `3 * n` holding the inverses. Requires `n >= 1`. /// -/// The caller's `stream` is used for every launch and synchronised at the -/// end (so the returned slice's data is committed before this function -/// returns). +/// Stream-ordered end to end: every launch (including the total's Fermat +/// inversion) goes on the caller's `stream`, so downstream same-stream +/// consumers need no synchronize. pub fn batch_inverse_ext3_dev( input: &CudaSlice, n: usize, @@ -101,13 +145,11 @@ pub fn batch_inverse_ext3_dev( )); } if n == 1 { - // Single element: D2H, host invert, H2D. Avoids running the - // scan + combine machinery for a degenerate case. - let host_view: Vec = stream.clone_dtoh(&input.slice(0..3))?; - stream.synchronize()?; - let inv = invert_ext3_host([host_view[0], host_view[1], host_view[2]])?; + // Single element: one-thread Fermat kernel, skipping the scan + + // combine machinery (and any host round-trip). + let be = backend()?; let mut out = unsafe { stream.alloc::(3) }?; - stream.memcpy_htod(&inv, &mut out)?; + launch_invert_total(stream, be, input, 1, &mut out)?; return Ok(out); } @@ -122,12 +164,11 @@ pub fn batch_inverse_ext3_dev( scan_into_fwd(stream, be, input, &mut prefix, n)?; scan_into_rev(stream, be, input, &mut suffix, n)?; - // total = prefix[n-1] = suffix[0]. Invert on host (one Fermat per batch). - let last_host: Vec = stream.clone_dtoh(&prefix.slice((n - 1) * 3..n * 3))?; - stream.synchronize()?; - let inv_total = invert_ext3_host([last_host[0], last_host[1], last_host[2]])?; + // total = prefix[n-1] = suffix[0]. One-thread Fermat inversion on device, + // keeping the whole batch inverse stream-ordered (the host round-trip here + // blocked the calling thread once per batch). let mut inv_total_dev = unsafe { stream.alloc::(3) }?; - stream.memcpy_htod(&inv_total, &mut inv_total_dev)?; + launch_invert_total(stream, be, &prefix, n, &mut inv_total_dev)?; // Combine: out[i] = prefix[i-1] * inv_total * suffix[i+1]. // SAFETY: the combine kernel writes every slot before any read. diff --git a/crypto/math-cuda/src/lde.rs b/crypto/math-cuda/src/lde.rs index 5f13161aa..ff05e7312 100644 --- a/crypto/math-cuda/src/lde.rs +++ b/crypto/math-cuda/src/lde.rs @@ -262,9 +262,32 @@ fn run_row_major_ntt_body( log_n: u64, m: u64, ) -> Result<()> { + // Levels 0..8 fused in shmem (one DRAM pass instead of eight); the + // remaining high-stride levels keep one kernel per level. + let mut first_level = 0u64; + if n >= 256 { + let t: u32 = 8.min(m as u32).max(1); + let cfg = LaunchConfig { + grid_dim: ((m as u32).div_ceil(t), (n / 256) as u32, 1), + block_dim: (t, 128, 1), + shared_mem_bytes: 256 * (t + 1) * 8, + }; + unsafe { + stream + .launch_builder(&be.ntt_dit_8_levels_row_major) + .arg(&mut *buf) + .arg(tw) + .arg(&n) + .arg(&log_n) + .arg(&m) + .launch(cfg)?; + } + first_level = 8.min(log_n); + } + let col_tile: u32 = 32.min(m as u32); let row_tile: u32 = (256 / col_tile).max(1); - for level in 0..log_n { + for level in first_level..log_n { let cfg = LaunchConfig { grid_dim: ( (m as u32).div_ceil(col_tile), @@ -437,8 +460,15 @@ fn expand_row_major_on_stream( // Fill a zeroed lde_size*total_cols buffer; only the first n*total_cols rows // carry data, the remainder are already zero (zero-padding for LDE). Host // input uploads (H2D); device input copies in place (D2D, no PCIe upload). + // Big host traces go through the pinned staging slot: the driver's + // internal pageable staging is 2-3x slower and convoys across threads. + const PINNED_H2D_MIN_U64: usize = 1 << 20; let mut buf = stream.alloc_zeros::(lde_size * total_cols)?; match input { + InnerInput::Host(h) if h.len() >= PINNED_H2D_MIN_U64 => { + let mut dst = buf.slice_mut(0..n * total_cols); + crate::device::htod_via(stream, be.pinned_staging(), &be.ctx, h, &mut dst)?; + } InnerInput::Host(h) => stream.memcpy_htod(h, &mut buf.slice_mut(0..n * total_cols))?, InnerInput::Dev(d) => stream.memcpy_dtod(d, &mut buf.slice_mut(0..n * total_cols))?, } @@ -694,7 +724,7 @@ pub fn coset_lde_row_major_split_trees( weights: &[u64], split_col: usize, build_precomputed: bool, -) -> Result<(Option>, Vec, GpuLdeBase, Vec)> { +) -> Result<(Option>, GpuLdeBase, Vec)> { assert!(split_col > 0 && split_col < m, "split inside the row"); assert!(n.is_power_of_two(), "n must be a power of two"); assert_eq!(weights.len(), n, "weights length must match n"); @@ -727,7 +757,7 @@ pub fn coset_lde_row_major_split_trees( )?; // One subset tree per column range, built sequentially on the stream. - let build_subset_tree = |col_start: u64, col_end: u64| -> Result> { + let build_subset_tree_dev = |col_start: u64, col_end: u64| -> Result> { let mut nodes_dev = unsafe { stream.alloc::(nodes_bytes) }?; { let mut leaves_view = @@ -745,17 +775,32 @@ pub fn coset_lde_row_major_split_trees( )?; } crate::merkle::build_inner_tree_levels(stream.as_ref(), be, &mut nodes_dev, num_leaves)?; - let mut nodes_host = vec![0u8; nodes_bytes]; - stream.memcpy_dtoh(&nodes_dev, &mut nodes_host)?; - Ok(nodes_host) + Ok(nodes_dev) }; + // Precomputed subset tree: full nodes to host (feeds the process-wide + // host tree cache keyed by root; built once per prove on cache miss). let precomputed_nodes = if build_precomputed { - Some(build_subset_tree(0, split_col as u64)?) + let nodes_dev = build_subset_tree_dev(0, split_col as u64)?; + let mut nodes_host = vec![0u8; nodes_bytes]; + stream.memcpy_dtoh(&nodes_dev, &mut nodes_host)?; + Some(nodes_host) } else { None }; - let mult_nodes = build_subset_tree(split_col as u64, cols_u64)?; + // Multiplicity subset tree: resident (per-epoch; the ~2x-leaves node + // download and host rebuild it used to pay are dropped — R4 openings + // gather paths on device). + let mult_tree = { + let nodes_dev = build_subset_tree_dev(split_col as u64, cols_u64)?; + let mut root = [0u8; 32]; + stream.memcpy_dtoh(&nodes_dev.slice(0..32), &mut root)?; + GpuMerkleTree { + nodes: Arc::new(nodes_dev), + leaves_len: num_leaves, + root, + } + }; // D2H the row-major LDE (preprocessed tables always keep the host copy — // they are excluded from the device-only gate). @@ -778,12 +823,12 @@ pub fn coset_lde_row_major_split_trees( buf: Arc::new(col_major_dev), m, lde_size, - tree: None, + tree: Some(mult_tree), ready: Some(Arc::new(ready)), trace_dev: trace_col_major.map(Arc::new), trace_rows: n, }; - Ok((precomputed_nodes, mult_nodes, handle, lde_out)) + Ok((precomputed_nodes, handle, lde_out)) } /// Row-major ext3 LDE + Keccak + Merkle, all on-device. @@ -1963,10 +2008,11 @@ pub fn coset_lde_batch_ext3_into( /// Batched ext3 coset LDE over columns ALREADY resident on device in slab /// layout (`3m` slabs of `lde_size` u64, first `n` of each filled, rest /// zero-padded), e.g. from the on-device degree-2 decomposition. Runs the -/// same butterfly pipeline as [`coset_lde_batch_ext3_into`], drains the -/// evaluations to `outputs` (interleaved ext3, `3*lde_size` u64 each), and -/// keeps the device buffer as a [`GpuLdeExt3`] handle (synchronized by the -/// drain, so `ready: None`). +/// same butterfly pipeline as [`coset_lde_batch_ext3_into`] and keeps the +/// device buffer as a [`GpuLdeExt3`] handle. With `outputs = Some(..)` the +/// evaluations are also drained to host (interleaved ext3, `3*lde_size` u64 +/// each; the drain synchronizes, so `ready: None`). With `None` nothing +/// leaves the device and the handle carries a `ready` event instead. pub fn coset_lde_batch_ext3_slabs_keep( stream: &Arc, mut buf: CudaSlice, @@ -1974,7 +2020,7 @@ pub fn coset_lde_batch_ext3_slabs_keep( n: usize, blowup_factor: usize, weights: &[u64], - outputs: &mut [&mut [u64]], + outputs: Option<&mut [&mut [u64]]>, ) -> Result { assert!(m > 0 && n.is_power_of_two(), "slab LDE shape"); assert_eq!(weights.len(), n, "weights length must match n"); @@ -1985,9 +2031,11 @@ pub fn coset_lde_batch_ext3_slabs_keep( let lde_size = n * blowup_factor; let mb = 3 * m; assert_eq!(buf.len(), mb * lde_size, "slab buffer shape"); - assert_eq!(outputs.len(), m, "outputs must match column count"); - for o in outputs.iter() { - assert_eq!(o.len(), 3 * lde_size, "each output must be 3*lde_size u64s"); + if let Some(outputs) = outputs.as_ref() { + assert_eq!(outputs.len(), m, "outputs must match column count"); + for o in outputs.iter() { + assert_eq!(o.len(), 3 * lde_size, "each output must be 3*lde_size u64s"); + } } assert_u32_domain(lde_size, "coset_lde_batch_ext3_slabs_keep lde_size"); let log_n = n.trailing_zeros() as u64; @@ -2049,22 +2097,38 @@ pub fn coset_lde_batch_ext3_slabs_keep( mb_u32, )?; - let pending = - crate::device::async_dtoh_via(stream, be.pinned_staging(), &be.ctx, &buf, mb * lde_size)?; - pending.wait_and_read(|bytes| { - // SAFETY: the pinned slab is u64-aligned by construction and the copy - // deposited exactly `mb * lde_size` u64s. - let pinned = - unsafe { std::slice::from_raw_parts(bytes.as_ptr() as *const u64, mb * lde_size) }; - unpack_pinned_slabs_to_ext3(pinned, outputs, lde_size); - })?; + let ready = match outputs { + Some(outputs) => { + let pending = crate::device::async_dtoh_via( + stream, + be.pinned_staging(), + &be.ctx, + &buf, + mb * lde_size, + )?; + pending.wait_and_read(|bytes| { + // SAFETY: the pinned slab is u64-aligned by construction and the + // copy deposited exactly `mb * lde_size` u64s. + let pinned = unsafe { + std::slice::from_raw_parts(bytes.as_ptr() as *const u64, mb * lde_size) + }; + unpack_pinned_slabs_to_ext3(pinned, outputs, lde_size); + })?; + None + } + None => { + let ready = be.take_event()?; + ready.event().record(stream)?; + Some(Arc::new(ready)) + } + }; Ok(GpuLdeExt3 { buf: Arc::new(buf), m, lde_size, tree: None, - ready: None, + ready, }) } diff --git a/crypto/math-cuda/src/merkle.rs b/crypto/math-cuda/src/merkle.rs index bfa756b13..1cb1c5b6f 100644 --- a/crypto/math-cuda/src/merkle.rs +++ b/crypto/math-cuda/src/merkle.rs @@ -158,10 +158,30 @@ pub(crate) fn build_inner_tree_levels( nodes_dev: &mut CudaSlice, leaves_len: usize, ) -> Result<()> { + // Once a level fits this many pairs, one single-block launch + // (`keccak_merkle_tail`) builds all remaining levels with barriers + // between them: the top ~11 levels of a big tree are each smaller than + // the per-launch overhead they used to pay. + const TAIL_MAX_PAIRS: u64 = 2048; let mut level_begin: u64 = (leaves_len - 1) as u64; while level_begin != 0 { let new_begin = level_begin / 2; let n_pairs = level_begin - new_begin; + if n_pairs <= TAIL_MAX_PAIRS { + let cfg = LaunchConfig { + grid_dim: (1, 1, 1), + block_dim: (KECCAK_BLOCK_DIM, 1, 1), + shared_mem_bytes: 0, + }; + unsafe { + stream + .launch_builder(&be.keccak_merkle_tail) + .arg(&mut *nodes_dev) + .arg(&level_begin) + .launch(cfg)?; + } + return Ok(()); + } let cfg = keccak_launch_cfg(n_pairs); unsafe { stream @@ -456,6 +476,56 @@ fn build_comp_poly_tree_nodes_dev( Ok((nodes_dev, num_leaves, stream)) } +/// Build the composition Merkle tree straight from a device-resident slab +/// buffer (`3*m` slabs of `lde_size` u64s, component `k` of part `c` at +/// `(c*3 + k) * lde_size` — the [`crate::lde::GpuLdeExt3`] layout). No host +/// staging and no H2D: the leaves kernel reads `buf` in place on `stream`. +pub fn build_comp_poly_tree_from_slabs_dev( + stream: &Arc, + buf: &CudaSlice, + m: usize, + lde_size: usize, +) -> Result { + assert!(m > 0); + assert!(lde_size.is_power_of_two() && lde_size >= 2); + assert_eq!(buf.len(), 3 * m * lde_size, "slab buffer shape"); + let num_leaves = lde_size / 2; + let tight_total_nodes = 2 * num_leaves - 1; + let be = backend()?; + + let mut nodes_dev = unsafe { stream.alloc::(tight_total_nodes * 32) }?; + let leaves_offset_bytes = (num_leaves - 1) * 32; + { + let mut leaves_view = + nodes_dev.slice_mut(leaves_offset_bytes..leaves_offset_bytes + num_leaves * 32); + let col_stride_u64 = lde_size as u64; + let num_parts_u64 = m as u64; + let num_rows_u64 = lde_size as u64; + let log_num_rows = lde_size.trailing_zeros() as u64; + let cfg = keccak_launch_cfg(num_leaves as u64); + unsafe { + stream + .launch_builder(&be.keccak_comp_poly_leaves_ext3) + .arg(buf) + .arg(&col_stride_u64) + .arg(&num_parts_u64) + .arg(&num_rows_u64) + .arg(&log_num_rows) + .arg(&mut leaves_view) + .launch(cfg)?; + } + } + build_inner_tree_levels(stream.as_ref(), be, &mut nodes_dev, num_leaves)?; + let mut root = [0u8; 32]; + stream.memcpy_dtoh(&nodes_dev.slice(0..32), &mut root)?; + stream.synchronize()?; + Ok(crate::lde::GpuMerkleTree { + nodes: Arc::new(nodes_dev), + leaves_len: num_leaves, + root, + }) +} + /// Build the comp poly Merkle tree on device and keep the nodes resident /// (returned as a [`crate::lde::GpuMerkleTree`] with its root), so R4 /// composition openings gather paths on device instead of copying the whole diff --git a/crypto/stark/src/fri/fri_commitment.rs b/crypto/stark/src/fri/fri_commitment.rs index 58c9eed77..1c199441d 100644 --- a/crypto/stark/src/fri/fri_commitment.rs +++ b/crypto/stark/src/fri/fri_commitment.rs @@ -18,6 +18,11 @@ where /// `merkle_tree` is a root only placeholder. `None` on the CPU path. #[cfg(feature = "cuda")] pub gpu_tree: Option, + /// The layer's evaluations kept resident on device (interleaved ext3, + /// `3 * len` u64). When `evaluation` is empty (device-only), the query + /// phase gathers `evaluation[index ^ 1]` from this buffer instead. + #[cfg(feature = "cuda")] + pub gpu_evals: Option>>, } impl FriLayer @@ -32,6 +37,8 @@ where merkle_tree, #[cfg(feature = "cuda")] gpu_tree: None, + #[cfg(feature = "cuda")] + gpu_evals: None, } } } diff --git a/crypto/stark/src/fri/mod.rs b/crypto/stark/src/fri/mod.rs index c3c16d123..1f53b51cf 100644 --- a/crypto/stark/src/fri/mod.rs +++ b/crypto/stark/src/fri/mod.rs @@ -150,7 +150,7 @@ where (final_poly_coeffs, fri_layer_list) } -pub fn query_phase( +pub fn query_phase( fri_layers: &[FriLayer>], iotas: &[usize], ) -> Vec> diff --git a/crypto/stark/src/gpu_lde.rs b/crypto/stark/src/gpu_lde.rs index 2167fcb94..8c87a8a3d 100644 --- a/crypto/stark/src/gpu_lde.rs +++ b/crypto/stark/src/gpu_lde.rs @@ -572,13 +572,16 @@ where /// Fully device-resident degree-2 decomposition + half extension: takes the /// resident composition evals `H`, decomposes into H0/H1 on device, LDE-extends -/// both, drains the evaluations to host (R3/openings still read them) and -/// keeps the de-interleaved parts buffer as a `GpuLdeExt3` for R4 DEEP. +/// both and keeps the de-interleaved parts buffer as a `GpuLdeExt3` (commit +/// tree, R3 OOD, R4 DEEP and openings all read the handle). With `want_host` +/// the evaluations are also drained to host for the fallback consumers; +/// without it (device-only) the returned part Vecs are empty placeholders. /// `None` → the caller downloads `H` and runs the host decompose path. pub(crate) fn try_decompose_extend_d2_dev( h: &math_cuda::constraint_interp::GpuCompH, inv_2x: &std::sync::Arc>>, weights: &[FieldElement], + want_host: bool, ) -> Option<(Vec>>, math_cuda::lde::GpuLdeExt3)> where F: IsField + 'static, @@ -615,11 +618,26 @@ where GPU_EXTEND_HALVES_CALLS.fetch_add(1, Ordering::Relaxed); GPU_LDE_CALLS.fetch_add(6, Ordering::Relaxed); - let mut lde_h0 = vec![FieldElement::::zero(); lde_size]; - let mut lde_h1 = vec![FieldElement::::zero(); lde_size]; // SAFETY: F == Goldilocks (repr u64); ext3 outputs are [u64; 3] per element. let weights_u64: &[u64] = unsafe { from_raw_parts(weights.as_ptr() as *const u64, weights.len()) }; + + if !want_host { + let handle = math_cuda::lde::coset_lde_batch_ext3_slabs_keep( + &stream, + slabs, + 2, + n, + 2, + weights_u64, + None, + ) + .ok()?; + return Some((vec![Vec::new(), Vec::new()], handle)); + } + + let mut lde_h0 = vec![FieldElement::::zero(); lde_size]; + let mut lde_h1 = vec![FieldElement::::zero(); lde_size]; let ext3_len = lde_size .checked_mul(3) .expect("ext3 output length overflow"); @@ -634,7 +652,7 @@ where n, 2, weights_u64, - &mut outputs, + Some(&mut outputs), ) .ok()?; @@ -800,7 +818,7 @@ where GPU_LEAF_HASH_CALLS.fetch_add(1 + build_precomputed as u64, Ordering::Relaxed); GPU_MERKLE_TREE_CALLS.fetch_add(1 + build_precomputed as u64, Ordering::Relaxed); - let (pre_nodes, mult_nodes, handle, lde_u64) = math_cuda::lde::coset_lde_row_major_split_trees( + let (pre_nodes, handle, lde_u64) = math_cuda::lde::coset_lde_row_major_split_trees( raw, n, m, @@ -815,7 +833,15 @@ where Some(nodes) => Some(tree_from_node_bytes::(nodes)?), None => None, }; - let mult_tree = tree_from_node_bytes::(mult_nodes)?; + // Mult tree resident in the handle: the host tree is root only and R4 + // openings gather authentication paths on device. + let mult_tree = MerkleTree::::from_root( + handle + .tree + .as_ref() + .expect("split path always builds the mult tree") + .root, + ); // Transmute Vec → Vec> (zero-copy, E == GoldilocksField). let lde_out: Vec> = unsafe { @@ -1127,6 +1153,38 @@ where Some((host, dev_tree)) } +/// Device-resident variant of [`try_build_comp_poly_tree_gpu`]: hashes the +/// composition tree straight from the resident R2 parts handle, skipping the +/// host pack + H2D re-upload of data that is already on device. +pub(crate) fn try_build_comp_poly_tree_gpu_from_dev( + handle: &math_cuda::lde::GpuLdeExt3, +) -> Option<(MerkleTree, math_cuda::lde::GpuMerkleTree)> +where + E: IsField + 'static, + B: IsMerkleTreeBackend, +{ + if TypeId::of::() != TypeId::of::() { + return None; + } + if handle.m == 0 || !handle.lde_size.is_power_of_two() || handle.lde_size < gpu_lde_threshold() + { + return None; + } + let be = math_cuda::device::backend().ok()?; + let stream = be.next_stream(); + handle.wait_ready_on(&stream).ok()?; + let dev_tree = math_cuda::merkle::build_comp_poly_tree_from_slabs_dev( + &stream, + handle.buf.as_ref(), + handle.m, + handle.lde_size, + ) + .ok()?; + GPU_COMP_POLY_TREE_CALLS.fetch_add(1, Ordering::Relaxed); + let host = MerkleTree::::from_root(dev_tree.root); + Some((host, dev_tree)) +} + /// R3 GPU dispatch: batched strided barycentric OOD evaluation over the main /// (base-field) LDE columns kept on device from R1. Operates on the /// device-resident LDE in place; only the coset points and inv_denoms are @@ -1230,6 +1288,37 @@ pub(crate) fn try_barycentric_ext3_on_handle( inv_denoms_host: &[FieldElement], r3_ctx: Option<(&R3DevContext, usize)>, ) -> Option>> +where + F: IsField + IsSubFieldOf + 'static, + E: IsField + 'static, +{ + try_barycentric_ext3_on_ext3_handle( + lde_trace.gpu_aux()?, + row_stride, + coset_points, + coset_offset_pow_n, + n_inv, + g_n_inv, + z_pow_n, + inv_denoms_host, + r3_ctx, + ) +} + +/// Same dispatch over an arbitrary resident ext3 handle (aux LDE or the R2 +/// composition parts). One column of OOD sums per handle column. +#[allow(clippy::too_many_arguments)] +pub(crate) fn try_barycentric_ext3_on_ext3_handle( + aux: &math_cuda::lde::GpuLdeExt3, + row_stride: usize, + coset_points: &[FieldElement], + coset_offset_pow_n: &FieldElement, + n_inv: &FieldElement, + g_n_inv: &FieldElement, + z_pow_n: &FieldElement, + inv_denoms_host: &[FieldElement], + r3_ctx: Option<(&R3DevContext, usize)>, +) -> Option>> where F: IsField + IsSubFieldOf + 'static, E: IsField + 'static, @@ -1240,7 +1329,6 @@ where if TypeId::of::() != TypeId::of::() { return None; } - let aux = lde_trace.gpu_aux()?; let num_cols = aux.m; if num_cols == 0 { return Some(Vec::new()); @@ -2137,6 +2225,7 @@ where Ok(s) => s, Err(_) => return None, }; + // Host-evals entry: the caller works with host copies, keep draining them. fri_commit_gpu_drive( state, transcript, @@ -2144,6 +2233,7 @@ where n0, blowup_log, final_poly_log_degree, + true, ) } @@ -2157,6 +2247,7 @@ pub(crate) fn try_fri_commit_gpu_from_dev( blowup_log: u32, final_poly_log_degree: u32, inv_twiddles: &[FieldElement], + want_host: bool, ) -> Option<( Vec>, Vec>>, @@ -2200,6 +2291,7 @@ where n0, blowup_log, final_poly_log_degree, + want_host, ) } @@ -2215,6 +2307,7 @@ fn fri_commit_gpu_drive( n0: usize, blowup_log: u32, final_poly_log_degree: u32, + want_host: bool, ) -> Option<( Vec>, Vec>>, @@ -2266,42 +2359,51 @@ where let zeta_ptr = &zeta as *const FieldElement as *const u64; let zeta_raw: [u64; 3] = unsafe { [*zeta_ptr, *zeta_ptr.add(1), *zeta_ptr.add(2)] }; - let (layer_evals_u64, dev_tree) = match state.fold_and_commit_layer(zeta_raw) { - Ok(v) => v, - Err(_) => { - *transcript = transcript_snapshot.clone(); - return None; - } - }; + let (layer_evals_u64, evals_dev, dev_tree) = + match state.fold_and_commit_layer(zeta_raw, want_host) { + Ok(v) => v, + Err(_) => { + *transcript = transcript_snapshot.clone(); + return None; + } + }; - // Build the FriLayer: ext3 evals and a root only host tree. The layer - // tree stays resident on device in `gpu_tree`; query openings gather - // paths from it via `gather_proofs_dev`. - let evaluation = u64_to_ext3_vec::(&layer_evals_u64); + // Build the FriLayer: a root only host tree, the tree and evals kept + // resident on device (`gpu_tree` / `gpu_evals`), and host evals only + // when a host copy was drained (fallback consumers). + let evaluation = layer_evals_u64 + .map(|v| u64_to_ext3_vec::(&v)) + .unwrap_or_default(); let root = dev_tree.root; let merkle_tree = MerkleTree::>::from_root(root); - let mut layer = FriLayer::new(&evaluation, merkle_tree); - layer.gpu_tree = Some(dev_tree); - fri_layer_list.push(layer); + fri_layer_list.push(FriLayer { + evaluation, + merkle_tree, + gpu_tree: Some(dev_tree), + gpu_evals: Some(evals_dev), + }); // >>>> Send commitment: [p_k] transcript.append_bytes(&root); } // Final (uncommitted) fold to the terminal codeword. n_out == terminal_len - // >= 2, so reuse fold_and_commit_layer and keep only its evaluations; the + // >= 2, so reuse fold_and_commit_layer and keep only its evaluations (the + // coefficient extraction below is host-side, so always drain them); the // Merkle root/nodes are discarded (the terminal layer is sent as coeffs). let zeta_final: FieldElement = transcript.sample_field_element(); let zeta_ptr = &zeta_final as *const FieldElement as *const u64; let zeta_raw: [u64; 3] = unsafe { [*zeta_ptr, *zeta_ptr.add(1), *zeta_ptr.add(2)] }; - let (terminal_evals_u64, _tree) = match state.fold_and_commit_layer(zeta_raw) { + let (terminal_evals_u64, _evals_dev, _tree) = match state.fold_and_commit_layer(zeta_raw, true) + { Ok(v) => v, Err(_) => { *transcript = transcript_snapshot; return None; } }; + let terminal_evals_u64 = terminal_evals_u64.expect("terminal fold drains to host"); debug_assert_eq!(terminal_evals_u64.len(), layout.terminal_len * 3); let terminal_codeword = u64_to_ext3_vec::(&terminal_evals_u64); @@ -2335,7 +2437,7 @@ pub(crate) fn try_fri_query_phase_gpu( iotas: &[usize], ) -> Option>> where - E: IsField, + E: IsField + 'static, FieldElement: AsBytes + Sync + Send, { if fri_layers.is_empty() { @@ -2376,6 +2478,30 @@ where ); } + // Symmetric evals per layer: read the host Vec when it was drained, + // otherwise a batched device gather off the resident layer evals + // (device-only, where no host copy exists). + let per_layer_syms: Vec>>> = fri_layers + .iter() + .enumerate() + .map(|(l, layer)| { + if !layer.evaluation.is_empty() { + return None; + } + let evals_dev = layer + .gpu_evals + .as_ref() + .expect("device-only FRI layer without resident evals"); + let positions: Vec = iotas.iter().map(|&iota| ((iota >> l) ^ 1) as u32).collect(); + let raw = math_cuda::fri::gather_ext3_at(evals_dev, &positions, &stream) + .expect("device FRI sym-eval gather failed; no host fallback"); + Some( + crate::constraint_ir::gpu_interp::ext3_u64_to_field::(&raw) + .expect("resident FRI evals are Goldilocks ext3"), + ) + }) + .collect(); + // Reassemble per-query decommitments, matching the host walk's order. let decommits = iotas .iter() @@ -2385,7 +2511,11 @@ where let mut layers_auth_paths = Vec::with_capacity(num_layers); let mut index = iota; for (l, layer) in fri_layers.iter().enumerate() { - layers_evaluations_sym.push(layer.evaluation[index ^ 1].clone()); + let sym = match &per_layer_syms[l] { + Some(v) => v[q].clone(), + None => layer.evaluation[index ^ 1].clone(), + }; + layers_evaluations_sym.push(sym); layers_auth_paths.push(per_layer_proofs[l][q].clone()); index >>= 1; } @@ -2458,25 +2588,31 @@ mod split_tree_tests { assert_eq!(mult_tree.root, cpu_mult_root, "multiplicity root"); // Openings must be byte-identical at scattered positions (pins the - // full node buffers, not just the roots). + // full node buffers, not just the roots). The mult tree is resident + // (host tree root only), so its paths come from the device gather — + // the exact production opening path. let num_leaves = n * blowup / 2; + let dev_tree = handle.tree.as_ref().expect("resident mult subset tree"); + let stream = math_cuda::device::backend().unwrap().next_stream(); for pos in [0usize, 1, 511, 12_345, num_leaves - 1] { assert_eq!( pre_tree.get_proof_by_pos(pos).unwrap().merkle_path, cpu_pre.get_proof_by_pos(pos).unwrap().merkle_path, "precomputed path at {pos}" ); + let dev_proofs = + gather_proofs_dev(dev_tree, &[pos], &stream).expect("device mult-tree path gather"); assert_eq!( - mult_tree.get_proof_by_pos(pos).unwrap().merkle_path, + dev_proofs[0].merkle_path, cpu_mult.get_proof_by_pos(pos).unwrap().merkle_path, "multiplicity path at {pos}" ); } + assert_eq!(mult_tree.root, dev_tree.root, "root-only host tree root"); // The handle must carry the column-major LDE for downstream rounds: // spot-check a few cells against the row-major host LDE. assert_eq!(handle.m, m); assert_eq!(handle.lde_size, n * blowup); - assert!(handle.tree.is_none(), "no device tree on the split path"); } } diff --git a/crypto/stark/src/logup_gpu.rs b/crypto/stark/src/logup_gpu.rs index 3fd49134d..9aed7c026 100644 --- a/crypto/stark/src/logup_gpu.rs +++ b/crypto/stark/src/logup_gpu.rs @@ -415,9 +415,10 @@ where /// straight to the aux LDE, no host round-trip) + the table contribution `L`. /// Returns `None` to fall back (non Goldilocks, below threshold, no GPU, GPU /// error). This is the residency path that avoids the term-column download. -pub fn try_build_aux_resident_gpu( +pub fn try_build_aux_resident_gpu<'a, F, E>( interactions: &[BusInteraction], - main_cols: &[Vec>], + num_cols: usize, + main_cols: impl FnOnce() -> &'a [Vec>], main_dev: Option<(&math_cuda::CudaSlice, usize)>, trace_len: usize, challenges: &[FieldElement], @@ -431,7 +432,7 @@ where { return None; } - if trace_len < GPU_LOGUP_MIN_ROWS || main_cols.is_empty() || interactions.is_empty() { + if trace_len < GPU_LOGUP_MIN_ROWS || num_cols == 0 || interactions.is_empty() { return None; } if std::env::var_os("LAMBDA_VM_NO_GPU_LOGUP").is_some() { @@ -442,18 +443,18 @@ where return None; } - let num_cols = main_cols.len(); desc.assert_columns_in_bounds(num_cols); // Reuse the resident main trace from the R1 main LDE (column-major - // `[col*trace_len + row]`, same column order as `main_cols`) when it matches - // this table exactly; otherwise flatten + upload the host columns. The - // resident buffer skips the ~3 GB main re-upload. + // `[col*trace_len + row]`, same column order as the host columns) when it + // matches this table exactly; otherwise materialize + flatten + upload the + // host columns. The resident buffer skips both the host transpose and the + // ~3 GB main re-upload. let resident_main = main_dev.filter(|&(buf, rows)| rows == trace_len && buf.len() == num_cols * trace_len); let mut main_flat = Vec::new(); if resident_main.is_none() { main_flat = vec![0u64; num_cols * trace_len]; - for (c, col) in main_cols.iter().enumerate() { + for (c, col) in main_cols().iter().enumerate() { for (r, e) in col.iter().enumerate() { main_flat[c * trace_len + r] = unsafe { *(e.value() as *const _ as *const u64) }; } diff --git a/crypto/stark/src/lookup.rs b/crypto/stark/src/lookup.rs index d376ebd1f..698e89ba9 100644 --- a/crypto/stark/src/lookup.rs +++ b/crypto/stark/src/lookup.rs @@ -1155,8 +1155,10 @@ where return None; } - // Clone main columns once (shared across all interactions) - let main_segment_cols = trace.columns_main(); + // Host main columns, materialized lazily: the resident GPU aux path + // reads the device main in place and must not pay this transpose. + let main_cols_cell: std::cell::OnceCell>>> = + std::cell::OnceCell::new(); let trace_len = trace.num_rows(); let _table_name = self.name.as_deref().unwrap_or("UNKNOWN"); @@ -1188,7 +1190,12 @@ where if trace.resident_aux_ok() && let Some(ra) = crate::logup_gpu::try_build_aux_resident_gpu::( interactions, - &main_segment_cols, + trace.num_main_columns, + || { + main_cols_cell + .get_or_init(|| trace.columns_main()) + .as_slice() + }, resident_main.as_ref().map(|r| (r.buf.as_ref(), r.rows)), trace_len, challenges, @@ -1201,12 +1208,14 @@ where return Some(BusPublicInputs { table_contribution }); } + let main_segment_cols = main_cols_cell.get_or_init(|| trace.columns_main()); + // GPU aux build (Goldilocks + ext3 + above threshold) computes all term // columns on device, byte identical, and falls back to the CPU build. #[cfg(feature = "cuda")] let gpu_term_cols = crate::logup_gpu::try_build_term_columns_gpu::( interactions, - &main_segment_cols, + main_segment_cols, trace_len, challenges, ); @@ -1220,7 +1229,7 @@ where let build_pair = |i: usize| { compute_logup_term_column( &[&interactions[i * 2], &interactions[i * 2 + 1]], - &main_segment_cols, + main_segment_cols, trace_len, challenges, _table_name, @@ -1248,7 +1257,7 @@ where &interactions[num_interactions - 2], &interactions[num_interactions - 1], ], - &main_segment_cols, + main_segment_cols, trace_len, challenges, _table_name, @@ -1256,7 +1265,7 @@ where } else { compute_logup_term_column( &[&interactions[num_interactions - 1]], - &main_segment_cols, + main_segment_cols, trace_len, challenges, _table_name, diff --git a/crypto/stark/src/prover.rs b/crypto/stark/src/prover.rs index 9a369b042..543ebf989 100644 --- a/crypto/stark/src/prover.rs +++ b/crypto/stark/src/prover.rs @@ -1481,10 +1481,12 @@ pub trait IsStarkProver< let mut gpu_composition_parts: Option = None; // Fully device-resident d=2 path: H stays on device through decompose + - // half extension, the parts handle feeds R4 DEEP, and only the final - // evaluations are drained to host (for the commit tree and openings). - // Any miss falls through to the host path below (downloading H when - // the evaluation itself already ran on device). + // half extension, and the parts handle feeds the commit tree, R3 OOD, + // R4 DEEP and the openings. The evaluations are drained to host only + // while a host trace copy exists (fallback consumers); under + // device-only nothing leaves the device and the placeholders below + // stay empty. Any miss falls through to the host path (downloading H + // when the evaluation itself already ran on device). #[cfg(feature = "cuda")] let mut precomputed_parts: Option>>> = None; #[cfg(feature = "cuda")] @@ -1502,6 +1504,7 @@ pub trait IsStarkProver< &h_dev, twiddles.inv_2x(domain), &twiddles.composition(domain).weights, + !round_1_result.lde_trace.host_trace_empty(), ) { Some((parts, handle)) => { gpu_composition_parts = Some(handle); @@ -1612,18 +1615,28 @@ pub trait IsStarkProver< #[cfg(feature = "instruments")] let t_sub = Instant::now(); - // GPU fast path for the comp-poly Merkle commit: row-pair Keccak - // leaves + device-side inner tree, both wrapping the host eval Vecs. - // GPU path keeps the composition tree resident on device (no whole tree - // copy) and returns a root only host tree. The device tree is threaded - // to R4 in `Round2.gpu_composition_tree`. + // GPU fast path for the comp-poly Merkle commit: hash straight from + // the resident parts handle when R2 kept one (no host pack + H2D + // re-upload); otherwise wrap the host eval Vecs. Either way the tree + // stays resident on device (no whole-tree copy), a root-only host tree + // is returned, and the device tree is threaded to R4 in + // `Round2.gpu_composition_tree`. #[cfg(feature = "cuda")] let (composition_poly_merkle_tree, composition_poly_root, gpu_composition_tree) = - match crate::gpu_lde::try_build_comp_poly_tree_gpu::< - FieldExtension, - BatchedMerkleTreeBackend, - >(&lde_composition_poly_parts_evaluations) - { + match gpu_composition_parts + .as_ref() + .and_then(|h| { + crate::gpu_lde::try_build_comp_poly_tree_gpu_from_dev::< + FieldExtension, + BatchedMerkleTreeBackend, + >(h) + }) + .or_else(|| { + crate::gpu_lde::try_build_comp_poly_tree_gpu::< + FieldExtension, + BatchedMerkleTreeBackend, + >(&lde_composition_poly_parts_evaluations) + }) { Some((host_tree, dev_tree)) => { let root = host_tree.root; (host_tree, root, Some(dev_tree)) @@ -1688,27 +1701,86 @@ pub trait IsStarkProver< // === Composition poly parts: barycentric evaluation at z^num_parts === let comp_z_pow_n = z_power.pow(domain_size); - let comp_inv_denoms = math::polynomial::barycentric_inv_denoms(&z_power, &dc.points); - let composition_poly_parts_ood_evaluation: Vec<_> = round_2_result - .lde_composition_poly_evaluations - .iter() - .map(|lde_evals| { - // Extract trace-size evaluations (stride = blowup_factor) - let evals: Vec> = (0..domain_size) - .map(|i| lde_evals[i * blowup_factor].clone()) - .collect(); - math::polynomial::interpolate_coset_eval_ext_with_g_n_inv( - &comp_z_pow_n, - &dc.offset_pow_n, - &dc.size_inv, - &dc.offset_pow_n_inv, - &dc.points, - &evals, - &comp_inv_denoms, - ) - }) - .collect(); + // GPU fast path: strided barycentric straight over the resident R2 + // parts handle (device inv_denoms for the single point z^P), skipping + // the host stride-extract and the sequential CPU fold per part. + #[cfg(feature = "cuda")] + let gpu_parts_ood: Option>> = + round_1_result + .lde_trace + .gpu_composition_parts() + .and_then(|parts_dev| { + let dispatch = |inv_host: &[FieldElement], + ctx: Option<(&crate::gpu_lde::R3DevContext, usize)>| { + crate::gpu_lde::try_barycentric_ext3_on_ext3_handle::( + parts_dev, + blowup_factor, + &dc.points, + &dc.offset_pow_n, + &dc.size_inv, + &dc.offset_pow_n_inv, + &comp_z_pow_n, + inv_host, + ctx, + ) + }; + match crate::gpu_lde::try_prep_r3_dev_context::( + &dc.points, + std::slice::from_ref(&z_power), + round_1_result.lde_trace.bound_stream(), + ) { + Some(ctx) => dispatch(&[], Some((&ctx, 0))), + // Below the dev-context threshold (single eval point): + // host inv_denoms + the same strided kernel, mirroring the + // trace OOD's mixed arm. + None => { + let inv = + math::polynomial::barycentric_inv_denoms(&z_power, &dc.points); + dispatch(&inv, None) + } + } + }); + #[cfg(not(feature = "cuda"))] + let gpu_parts_ood: Option>> = None; + + let composition_poly_parts_ood_evaluation: Vec<_> = match gpu_parts_ood { + Some(v) => v, + None => { + // The host part evals are empty under device-only (the R2 + // drain is skipped); reaching this arm there is a mis-gate. + #[cfg(feature = "cuda")] + assert!( + round_2_result + .lde_composition_poly_evaluations + .first() + .is_none_or(|p| !p.is_empty()), + "R3 parts OOD fell back to the host part evals, but they are \ + device-only (empty)" + ); + let comp_inv_denoms = + math::polynomial::barycentric_inv_denoms(&z_power, &dc.points); + round_2_result + .lde_composition_poly_evaluations + .iter() + .map(|lde_evals| { + // Extract trace-size evaluations (stride = blowup_factor) + let evals: Vec> = (0..domain_size) + .map(|i| lde_evals[i * blowup_factor].clone()) + .collect(); + math::polynomial::interpolate_coset_eval_ext_with_g_n_inv( + &comp_z_pow_n, + &dc.offset_pow_n, + &dc.size_inv, + &dc.offset_pow_n_inv, + &dc.points, + &evals, + &comp_inv_denoms, + ) + }) + .collect() + } + }; // === Trace polynomials: barycentric evaluation via LDE === let trace_ood_evaluations = crate::trace::get_trace_evaluations_from_lde( @@ -1812,6 +1884,7 @@ pub trait IsStarkProver< domain.blowup_factor.trailing_zeros(), air.options().fri_final_poly_log_degree as u32, domain.fri_inv_twiddles(), + !round_1_result.lde_trace.host_trace_empty(), ) }); #[cfg(not(feature = "cuda"))] @@ -2496,23 +2569,21 @@ pub trait IsStarkProver< // must succeed: there is no host tree to fall back to, so a gather error // is a hard abort. When the tree is not device resident the value is // `None` and the openings below walk the full host tree. + // For preprocessed tables the resident tree is the multiplicity subset + // tree (the host `main_commit.tree` is root only); values still come + // from the host LDE range gather below. #[cfg(feature = "cuda")] - let main_dev_proofs: Option>> = if is_preprocessed { - None - } else { - lde_trace - .gpu_main() - .and_then(|h| h.tree.as_ref()) - .map(|tree| { - let stream = lde_trace - .bound_stream() - .expect("bound stream for device-resident main-tree opening"); - // Row-pair leaves: one proof per query at position `challenge`. - crate::gpu_lde::gather_proofs_dev(tree, indexes_to_open, &stream).expect( - "device main-tree gather failed; resident tree has no host fallback", - ) - }) - }; + let main_dev_proofs: Option>> = lde_trace + .gpu_main() + .and_then(|h| h.tree.as_ref()) + .map(|tree| { + let stream = lde_trace + .bound_stream() + .expect("bound stream for device-resident main-tree opening"); + // Row-pair leaves: one proof per query at position `challenge`. + crate::gpu_lde::gather_proofs_dev(tree, indexes_to_open, &stream) + .expect("device main-tree gather failed; resident tree has no host fallback") + }); // Same for the aux trace tree, when it is device resident. #[cfg(feature = "cuda")] @@ -2554,8 +2625,10 @@ pub trait IsStarkProver< // *_dev_values.is_some()` on the Goldilocks path) and we never gather // rows for a tree that is not device resident. #[cfg(feature = "cuda")] - let main_dev_values: Option>> = - main_dev_proofs.as_ref().and_then(|_| { + let main_dev_values: Option>> = (!is_preprocessed) + .then_some(()) + .and(main_dev_proofs.as_ref()) + .and_then(|_| { lde_trace.gpu_main().and_then(|h| { Self::gather_query_rows_device( lde_trace, @@ -2595,12 +2668,65 @@ pub trait IsStarkProver< }) }); + // Composition part values off the resident R2 parts handle (one ext3 + // "column" per part), same row-pair gather as main/aux above. + #[cfg(feature = "cuda")] + let comp_num_parts = lde_trace + .gpu_composition_parts() + .map(|h| h.m) + .unwrap_or_else(|| round_2_result.lde_composition_poly_evaluations.len()); + #[cfg(feature = "cuda")] + let comp_dev_values: Option>> = + comp_dev_proofs.as_ref().and_then(|_| { + lde_trace.gpu_composition_parts().and_then(|h| { + Self::gather_query_rows_device( + lde_trace, + "composition", + |stream| { + math_cuda::barycentric::gather_rows_ext3_on_device( + h, + &query_rows, + stream, + ) + }, + |raw| { + crate::constraint_ir::gpu_interp::ext3_u64_to_field::( + raw, + ) + }, + ) + }) + }); + for (qi, index) in indexes_to_open.iter().enumerate() { #[cfg(not(feature = "cuda"))] let _ = qi; // For preprocessed tables, open the main split (multiplicities only); // for normal tables, open all main columns. let main_trace_opening = if is_preprocessed { + // Multiplicity subset: device proof (resident subset tree) + + // host range gather for the values. + #[cfg(feature = "cuda")] + { + match &main_dev_proofs { + Some(proofs) => Self::open_polys_with_proofs( + domain, + proofs[qi].clone(), + *index, + |row| { + lde_trace.gather_main_row_range( + row, + num_precomputed_cols, + total_cols, + ) + }, + ), + None => Self::open_polys_with(domain, &main_commit.tree, *index, |row| { + lde_trace.gather_main_row_range(row, num_precomputed_cols, total_cols) + }), + } + } + #[cfg(not(feature = "cuda"))] Self::open_polys_with(domain, &main_commit.tree, *index, |row| { lde_trace.gather_main_row_range(row, num_precomputed_cols, total_cols) }) @@ -2638,18 +2764,57 @@ pub trait IsStarkProver< let composition_openings = { #[cfg(feature = "cuda")] { - if let Some(proofs) = &comp_dev_proofs { - Self::open_composition_poly_with_proof( - proofs[qi].clone(), - &round_2_result.lde_composition_poly_evaluations, - *index, - ) - } else { - Self::open_composition_poly( + match (&comp_dev_proofs, &comp_dev_values) { + (Some(proofs), Some(vals)) => { + let (even, odd) = Self::device_row_pair(vals, qi, comp_num_parts); + // Cross-check against the host part evals while + // they are still resident (absent under full + // residency, where the gather is the only source). + if round_2_result + .lde_composition_poly_evaluations + .first() + .is_some_and(|p| !p.is_empty()) + { + let expected = Self::open_composition_poly_with_proof( + proofs[qi].clone(), + &round_2_result.lde_composition_poly_evaluations, + *index, + ); + assert_eq!( + even, expected.evaluations, + "device composition-row gather mismatch (even), query {qi}" + ); + assert_eq!( + odd, expected.evaluations_sym, + "device composition-row gather mismatch (odd), query {qi}" + ); + } + PolynomialOpenings { + proof: proofs[qi].clone(), + evaluations: even, + evaluations_sym: odd, + } + } + (Some(proofs), None) => { + assert!( + round_2_result + .lde_composition_poly_evaluations + .first() + .is_none_or(|p| !p.is_empty()), + "R4 composition opening fell back to the host part evals, \ + but they are device-only (empty)" + ); + Self::open_composition_poly_with_proof( + proofs[qi].clone(), + &round_2_result.lde_composition_poly_evaluations, + *index, + ) + } + _ => Self::open_composition_poly( &round_2_result.composition_poly_merkle_tree, &round_2_result.lde_composition_poly_evaluations, *index, - ) + ), } } #[cfg(not(feature = "cuda"))] From ac3dee9e36f7821c74648f61161682804583fe9c Mon Sep 17 00:00:00 2001 From: Joaquin Carletti Date: Wed, 29 Jul 2026 15:43:09 -0300 Subject: [PATCH 02/12] fix(gpu): harden round-2 residency paths after review Grid-stride the fused row-major NTT past gridDim.y (lde >= 2^24 silently fell back to CPU), assert the device-only contract in the R2 composition commit and preprocessed opening fallbacks, validate htod_via bounds, retain FRI device evals only under device-only, and move the inverse fault-injection hook so every batch-inverse entry is covered. --- crypto/crypto/src/merkle_tree/merkle.rs | 13 ++++++ crypto/math-cuda/kernels/ntt.cu | 61 ++++++++++++++----------- crypto/math-cuda/src/device.rs | 33 ++++++++----- crypto/math-cuda/src/inverse.rs | 9 +++- crypto/math-cuda/src/lde.rs | 19 ++++---- crypto/stark/src/gpu_lde.rs | 26 ++++++----- crypto/stark/src/lookup.rs | 2 +- crypto/stark/src/prover.rs | 28 ++++++++++-- 8 files changed, 125 insertions(+), 66 deletions(-) diff --git a/crypto/crypto/src/merkle_tree/merkle.rs b/crypto/crypto/src/merkle_tree/merkle.rs index d53f06f10..e9ecd0bb5 100644 --- a/crypto/crypto/src/merkle_tree/merkle.rs +++ b/crypto/crypto/src/merkle_tree/merkle.rs @@ -172,6 +172,19 @@ where /// but no nodes. Used when paths are gathered from a device resident copy /// (GPU) instead of this host tree, so the host nodes are never built. /// [`get_proof_by_pos`](Self::get_proof_by_pos) must NOT be called on it. + /// True when this tree carries only its root (the nodes live elsewhere, + /// e.g. device-resident): openings must not walk this tree. + pub fn is_root_only(&self) -> bool { + #[cfg(feature = "disk-spill")] + { + self.nodes.is_empty() && self.mmap_backing.is_none() + } + #[cfg(not(feature = "disk-spill"))] + { + self.nodes.is_empty() + } + } + pub fn from_root(root: B::Node) -> Self { MerkleTree { root, diff --git a/crypto/math-cuda/kernels/ntt.cu b/crypto/math-cuda/kernels/ntt.cu index 35a4f7b20..1e6c83f5c 100644 --- a/crypto/math-cuda/kernels/ntt.cu +++ b/crypto/math-cuda/kernels/ntt.cu @@ -429,42 +429,49 @@ extern "C" __global__ void ntt_dit_8_levels_row_major(uint64_t *data, uint32_t pitch = T + 1; uint64_t col = (uint64_t)blockIdx.x * T + threadIdx.x; bool live = col < m; - uint64_t row_base = (uint64_t)blockIdx.y * 256; - - for (uint32_t r = threadIdx.y; r < 256; r += blockDim.y) { - if (live) tile[r * pitch + threadIdx.x] = data[(row_base + r) * m + col]; - } - __syncthreads(); uint32_t n_loc_steps = (uint32_t)min((uint64_t)8, log_n); uint32_t remaining_high_bits = (uint32_t)(log_n - 1); uint32_t high_mask = (1u << remaining_high_bits) - 1u; - for (uint32_t loc_step = 0; loc_step < n_loc_steps; ++loc_step) { - for (uint32_t i = threadIdx.y; i < 128; i += blockDim.y) { - uint32_t half = 1u << loc_step; - uint32_t grp = i >> loc_step; - uint32_t grp_pos = i & (half - 1); - uint32_t idx1 = (grp << (loc_step + 1)) + grp_pos; - uint32_t idx2 = idx1 + half; + // Grid-stride over 256-row blocks: gridDim.y caps at 65535, so lde sizes + // >= 2^24 need more than one row block per y-slot. The trip count is + // uniform across the block, keeping every __syncthreads converged. + for (uint64_t rb = blockIdx.y; rb < (n >> 8); rb += gridDim.y) { + uint64_t row_base = rb * 256; - uint32_t gs = loc_step; - uint32_t ggp = ((uint32_t)blockIdx.y << 7) + i; - ggp = (ggp & high_mask) + (ggp >> remaining_high_bits); - ggp = ggp & ((1u << gs) - 1u); - uint64_t factor = tw[(uint64_t)ggp * (n >> (gs + 1))]; + for (uint32_t r = threadIdx.y; r < 256; r += blockDim.y) { + if (live) tile[r * pitch + threadIdx.x] = data[(row_base + r) * m + col]; + } + __syncthreads(); - if (live) { - uint64_t u = tile[idx1 * pitch + threadIdx.x]; - uint64_t v = mul(tile[idx2 * pitch + threadIdx.x], factor); - tile[idx1 * pitch + threadIdx.x] = add(u, v); - tile[idx2 * pitch + threadIdx.x] = sub(u, v); + for (uint32_t loc_step = 0; loc_step < n_loc_steps; ++loc_step) { + for (uint32_t i = threadIdx.y; i < 128; i += blockDim.y) { + uint32_t half = 1u << loc_step; + uint32_t grp = i >> loc_step; + uint32_t grp_pos = i & (half - 1); + uint32_t idx1 = (grp << (loc_step + 1)) + grp_pos; + uint32_t idx2 = idx1 + half; + + uint32_t gs = loc_step; + uint32_t ggp = ((uint32_t)rb << 7) + i; + ggp = (ggp & high_mask) + (ggp >> remaining_high_bits); + ggp = ggp & ((1u << gs) - 1u); + uint64_t factor = tw[(uint64_t)ggp * (n >> (gs + 1))]; + + if (live) { + uint64_t u = tile[idx1 * pitch + threadIdx.x]; + uint64_t v = mul(tile[idx2 * pitch + threadIdx.x], factor); + tile[idx1 * pitch + threadIdx.x] = add(u, v); + tile[idx2 * pitch + threadIdx.x] = sub(u, v); + } } + __syncthreads(); } - __syncthreads(); - } - for (uint32_t r = threadIdx.y; r < 256; r += blockDim.y) { - if (live) data[(row_base + r) * m + col] = tile[r * pitch + threadIdx.x]; + for (uint32_t r = threadIdx.y; r < 256; r += blockDim.y) { + if (live) data[(row_base + r) * m + col] = tile[r * pitch + threadIdx.x]; + } + __syncthreads(); } } diff --git a/crypto/math-cuda/src/device.rs b/crypto/math-cuda/src/device.rs index 292026401..c3dd4a2a9 100644 --- a/crypto/math-cuda/src/device.rs +++ b/crypto/math-cuda/src/device.rs @@ -611,7 +611,9 @@ pub fn backend() -> Result<&'static Backend> { /// /// Holding this value keeps the staging slot's mutex locked, which is what /// makes the whole scheme safe: no other caller (and no capacity growth) can -/// touch the slab while the DMA is in flight. +/// touch the slab while the DMA is in flight. Corollary: never call +/// `htod_via`/`async_dtoh_via` on the same slot from the thread holding a +/// live `PendingD2H` — the non-reentrant slot mutex self-deadlocks. pub struct PendingD2H<'a> { staging: std::sync::MutexGuard<'a, PinnedStaging>, n_bytes: usize, @@ -628,17 +630,6 @@ impl Drop for PendingD2H<'_> { } } -/// Enqueue an async D2H of `n_elems` of `src` into the pinned slab of `slot`, -/// without synchronizing the stream. Unlike `stream.memcpy_dtoh` into a plain -/// (pageable) slice — which the driver services synchronously — this returns -/// as soon as the copy is queued; the returned [`PendingD2H`] is awaited at -/// the point the host actually needs the bytes. -/// -/// SAFETY contract (upheld by construction for our callers): `src` must stay -/// alive until the copy completes. Dropping a `CudaSlice` frees it -/// stream-ordered on its own stream, so a `src` allocated on `stream` may be -/// dropped after this call — the free queues behind the copy. Do NOT pass a -/// `src` owned by a *different* stream and drop it before waiting. /// Host→device copy staged through the pinned slot: one host memcpy into /// pinned memory + one async DMA, instead of the driver's internal pageable /// staging (small chunks; 2-3x slower for multi-hundred-MB traces and it @@ -652,7 +643,14 @@ pub fn htod_via( dst: &mut cudarc::driver::CudaViewMut<'_, T>, ) -> Result<()> { use cudarc::driver::DevicePtrMut; + assert!( + dst.len() >= src_host.len(), + "htod_via: destination shorter than source" + ); let n_bytes = std::mem::size_of_val(src_host); + if n_bytes == 0 { + return Ok(()); + } let u64_len = n_bytes.div_ceil(8); let mut staging = slot.lock().unwrap(); staging.ensure_capacity(u64_len, ctx)?; @@ -680,6 +678,17 @@ pub fn htod_via( staging.sync_event() } +/// Enqueue an async D2H of `n_elems` of `src` into the pinned slab of `slot`, +/// without synchronizing the stream. Unlike `stream.memcpy_dtoh` into a plain +/// (pageable) slice — which the driver services synchronously — this returns +/// as soon as the copy is queued; the returned [`PendingD2H`] is awaited at +/// the point the host actually needs the bytes. +/// +/// SAFETY contract (upheld by construction for our callers): `src` must stay +/// alive until the copy completes. Dropping a `CudaSlice` frees it +/// stream-ordered on its own stream, so a `src` allocated on `stream` may be +/// dropped after this call — the free queues behind the copy. Do NOT pass a +/// `src` owned by a *different* stream and drop it before waiting. pub fn async_dtoh_via<'a, T: cudarc::driver::DeviceRepr>( stream: &Arc, slot: &'a Mutex, diff --git a/crypto/math-cuda/src/inverse.rs b/crypto/math-cuda/src/inverse.rs index 4dd49556c..9931e3a8c 100644 --- a/crypto/math-cuda/src/inverse.rs +++ b/crypto/math-cuda/src/inverse.rs @@ -94,6 +94,11 @@ const fn ext3_fermat_exponent() -> [u64; 3] { } /// One-thread Fermat inversion of `src[n-1]` into `out[0..3]`, stream-ordered. +/// +/// Unlike the host Fermat this used to call, a zero total maps silently to +/// zero instead of `Err`. Unreachable with honest inputs (LogUp/barycentric +/// denominators are nonzero w.h.p.); callers must not rely on a zero-total +/// error. fn launch_invert_total( stream: &Arc, be: &crate::device::Backend, @@ -133,6 +138,8 @@ pub fn batch_inverse_ext3_dev( n: usize, stream: &Arc, ) -> Result> { + #[cfg(feature = "test-faults")] + check_inverse_fault_injection()?; assert!(n >= 1, "batch_inverse_ext3_dev requires n >= 1"); // Runtime guard (not debug_assert): a u32 grid_dim is truncated past // u32::MAX / BLOCK_SIZE, which would silently launch too few blocks @@ -220,8 +227,6 @@ pub fn compute_and_invert_denoms_ext3_dev( sign: DenomSign, stream: &Arc, ) -> Result> { - #[cfg(feature = "test-faults")] - check_inverse_fault_injection()?; assert_eq!(z_scalars_host.len(), k_scalars * 3); assert!(n >= 1 && k_scalars >= 1); diff --git a/crypto/math-cuda/src/lde.rs b/crypto/math-cuda/src/lde.rs index ff05e7312..3d8bfa207 100644 --- a/crypto/math-cuda/src/lde.rs +++ b/crypto/math-cuda/src/lde.rs @@ -268,7 +268,7 @@ fn run_row_major_ntt_body( if n >= 256 { let t: u32 = 8.min(m as u32).max(1); let cfg = LaunchConfig { - grid_dim: ((m as u32).div_ceil(t), (n / 256) as u32, 1), + grid_dim: ((m as u32).div_ceil(t), ((n / 256) as u32).min(65535), 1), block_dim: (t, 128, 1), shared_mem_bytes: 256 * (t + 1) * 8, }; @@ -704,17 +704,16 @@ pub fn coset_lde_row_major_with_merkle_tree_keep( /// `[split_col, m)` commit to separate trees over the same row-major LDE, /// mirroring the CPU `commit_rows_bit_reversed_subset` pair. /// -/// Both trees' complete node buffers are downloaded to host -/// (`(2*num_leaves - 1) * 32` bytes each, inner nodes first, root at offset 0, +/// The precomputed tree's complete node buffer is downloaded to host +/// (`(2*num_leaves - 1) * 32` bytes, inner nodes first, root at offset 0, /// leaves at the tail — the exact `MerkleTree::from_precomputed_nodes` -/// layout), because preprocessed-table openings walk host trees. The -/// precomputed tree is only built when `build_precomputed` is true (the -/// caller skips it on a process-cache hit). +/// layout) because it feeds the process-wide host tree cache; it is only +/// built when `build_precomputed` is true (the caller skips it on a cache +/// hit). The multiplicity tree stays resident in `handle.tree` — openings +/// gather its paths on device. /// -/// Returns `(precomputed_nodes, mult_nodes, handle, row_major_lde)`. The -/// handle carries the column-major LDE + trace snapshot for downstream GPU -/// rounds but NO device tree (`tree: None`) — openings for preprocessed -/// tables never gather from device. +/// Returns `(precomputed_nodes, handle, row_major_lde)`. The handle also +/// carries the column-major LDE + trace snapshot for downstream GPU rounds. #[allow(clippy::type_complexity)] pub fn coset_lde_row_major_split_trees( row_major: &[u64], diff --git a/crypto/stark/src/gpu_lde.rs b/crypto/stark/src/gpu_lde.rs index 8c87a8a3d..98830fcc7 100644 --- a/crypto/stark/src/gpu_lde.rs +++ b/crypto/stark/src/gpu_lde.rs @@ -767,10 +767,11 @@ where /// one row-major GPU LDE of ALL columns plus TWO subset Merkle trees — the /// precomputed columns `[0, split_col)` and the multiplicity columns /// `[split_col, m)` — matching the CPU `commit_rows_bit_reversed_subset` -/// pair bit for bit. Trees come back as full HOST trees (openings for -/// preprocessed tables walk host trees); the handle keeps the column-major -/// LDE + trace snapshot device-resident for the downstream GPU rounds, with -/// no device tree. +/// pair bit for bit. The precomputed tree comes back as a full HOST tree +/// (it feeds the process-wide cache); the multiplicity tree stays resident +/// in the handle (root-only host tree, R4 openings gather paths on device). +/// The handle also keeps the column-major LDE + trace snapshot for the +/// downstream GPU rounds. /// /// `build_precomputed=false` skips the precomputed tree (process-cache hit); /// the first element is then `None`. @@ -1405,12 +1406,12 @@ pub fn gpu_fri_calls() -> u64 { /// Batch-invert dispatch counter (one per /// [`try_compute_and_invert_inv_denoms_dev`] call that actually built a -/// device handle). Fires at most twice per prove per table: once for R3 -/// OOD's `num_eval_points * trace_size` denominators and once for R4 -/// DEEP's `(1 + num_eval_points) * lde_size` denominators. R4 has two -/// chances at it (device-only DEEP, then the host DEEP arm), and both are -/// counted here, so a single failed dispatch does not necessarily lower the -/// total; R3's fallback is CPU-only, so a failure there does. +/// device handle). Fires up to three times per prove per table: R3 trace +/// OOD's `num_eval_points * trace_size` denominators, R3 parts OOD's single +/// point, and R4 DEEP's `(1 + num_eval_points) * lde_size` denominators. R4 +/// has two chances at it (device-only DEEP, then the host DEEP arm), and both +/// are counted here, so a single failed dispatch does not necessarily lower +/// the total; R3's fallbacks are CPU-only, so a failure there does. pub(crate) static GPU_BATCH_INVERT_CALLS: AtomicU64 = AtomicU64::new(0); pub fn gpu_batch_invert_calls() -> u64 { GPU_BATCH_INVERT_CALLS.load(Ordering::Relaxed) @@ -2376,11 +2377,14 @@ where .unwrap_or_default(); let root = dev_tree.root; let merkle_tree = MerkleTree::>::from_root(root); + // Retain the device evals only when no host copy exists (device-only): + // with a host copy the query phase reads it, and the retained buffer + // would be ~24 bytes/LDE-row of dead VRAM per table. fri_layer_list.push(FriLayer { evaluation, merkle_tree, gpu_tree: Some(dev_tree), - gpu_evals: Some(evals_dev), + gpu_evals: (!want_host).then_some(evals_dev), }); // >>>> Send commitment: [p_k] diff --git a/crypto/stark/src/lookup.rs b/crypto/stark/src/lookup.rs index 698e89ba9..ceda5417a 100644 --- a/crypto/stark/src/lookup.rs +++ b/crypto/stark/src/lookup.rs @@ -1286,7 +1286,7 @@ where let (per_bus_sums, per_bus_sender_sums, per_bus_receiver_sums) = compute_debug_bus_sums_batched( &self.auxiliary_trace_build_data.interactions, - &main_segment_cols, + main_segment_cols, trace_len, challenges, _table_name, diff --git a/crypto/stark/src/prover.rs b/crypto/stark/src/prover.rs index 543ebf989..8612bb20a 100644 --- a/crypto/stark/src/prover.rs +++ b/crypto/stark/src/prover.rs @@ -1642,6 +1642,14 @@ pub trait IsStarkProver< (host_tree, root, Some(dev_tree)) } None => { + // The host part evals are empty under device-only (the R2 + // drain is skipped); abort with the device-only contract's + // message instead of a misleading EmptyCommitment. + assert!( + !round_1_result.lde_trace.host_trace_empty(), + "R2 composition commit fell back to the host part evals, \ + but they are device-only (empty)" + ); let (tree, root) = crate::commitment::commit_bit_reversed( &lde_composition_poly_parts_evaluations, crate::commitment::ROWS_PER_LEAF, @@ -2721,9 +2729,23 @@ pub trait IsStarkProver< ) }, ), - None => Self::open_polys_with(domain, &main_commit.tree, *index, |row| { - lde_trace.gather_main_row_range(row, num_precomputed_cols, total_cols) - }), + None => { + // A root-only host tree means the nodes are + // device-resident: this arm would emit an empty + // path for query position 0 instead of failing. + assert!( + !main_commit.tree.is_root_only(), + "preprocessed opening fell back to the host tree, \ + but it is root-only (nodes device-resident)" + ); + Self::open_polys_with(domain, &main_commit.tree, *index, |row| { + lde_trace.gather_main_row_range( + row, + num_precomputed_cols, + total_cols, + ) + }) + } } } #[cfg(not(feature = "cuda"))] From 8a658846c143ea67f6c780f10ea27988d830fdf2 Mon Sep 17 00:00:00 2001 From: Joaquin Carletti Date: Fri, 31 Jul 2026 16:44:27 -0300 Subject: [PATCH 03/12] =?UTF-8?q?fix(gpu):=20address=20round-2=20review=20?= =?UTF-8?q?=E2=80=94=20guarded=20zero-inverse,=20retargeted=20fault=20hook?= =?UTF-8?q?,=20merkle=20root-only?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- crypto/crypto/src/merkle_tree/merkle.rs | 15 ++++++++---- crypto/math-cuda/src/inverse.rs | 32 +++++++++++++++++++++---- 2 files changed, 39 insertions(+), 8 deletions(-) diff --git a/crypto/crypto/src/merkle_tree/merkle.rs b/crypto/crypto/src/merkle_tree/merkle.rs index e9ecd0bb5..447654907 100644 --- a/crypto/crypto/src/merkle_tree/merkle.rs +++ b/crypto/crypto/src/merkle_tree/merkle.rs @@ -168,10 +168,6 @@ where }) } - /// Create a root only Merkle tree placeholder: stores the commitment root - /// but no nodes. Used when paths are gathered from a device resident copy - /// (GPU) instead of this host tree, so the host nodes are never built. - /// [`get_proof_by_pos`](Self::get_proof_by_pos) must NOT be called on it. /// True when this tree carries only its root (the nodes live elsewhere, /// e.g. device-resident): openings must not walk this tree. pub fn is_root_only(&self) -> bool { @@ -185,6 +181,10 @@ where } } + /// Create a root only Merkle tree placeholder: stores the commitment root + /// but no nodes. Used when paths are gathered from a device resident copy + /// (GPU) instead of this host tree, so the host nodes are never built. + /// [`get_proof_by_pos`](Self::get_proof_by_pos) must NOT be called on it. pub fn from_root(root: B::Node) -> Self { MerkleTree { root, @@ -266,7 +266,14 @@ where /// Returns a Merkle proof for the element/s at position pos /// For example, give me an inclusion proof for the 3rd element in the /// Merkle tree + /// + /// Returns `None` on a root-only tree ([`from_root`](Self::from_root)): + /// its nodes live elsewhere (e.g. device-resident), so a host path would + /// be a silently-empty bogus proof rather than an inclusion witness. pub fn get_proof_by_pos(&self, pos: usize) -> Option> { + if self.is_root_only() { + return None; + } let pos = pos + self.node_count() / 2; let Ok(merkle_path) = self.build_merkle_path(pos) else { return None; diff --git a/crypto/math-cuda/src/inverse.rs b/crypto/math-cuda/src/inverse.rs index 9931e3a8c..833bf2906 100644 --- a/crypto/math-cuda/src/inverse.rs +++ b/crypto/math-cuda/src/inverse.rs @@ -97,8 +97,11 @@ const fn ext3_fermat_exponent() -> [u64; 3] { /// /// Unlike the host Fermat this used to call, a zero total maps silently to /// zero instead of `Err`. Unreachable with honest inputs (LogUp/barycentric -/// denominators are nonzero w.h.p.); callers must not rely on a zero-total -/// error. +/// denominators are nonzero w.h.p. under random Fiat-Shamir challenges); +/// callers must not rely on a zero-total error. Debug builds add a D2H+sync +/// invertibility guard (see below) that panics on a zero total so a +/// construction/kernel bug fails loudly in tests; release elides it to keep +/// the batch inverse fully stream-ordered (no per-batch host round-trip). fn launch_invert_total( stream: &Arc, be: &crate::device::Backend, @@ -124,6 +127,23 @@ fn launch_invert_total( .arg(&mut *out) .launch(cfg)?; } + // Debug-only invertibility guard. The Fermat kernel maps a zero total + // (some denominator was zero) silently to zero, so the batch would ship + // all-zero "inverses" instead of erroring. A valid inverse is never zero, + // so `out == 0` unambiguously flags a zero total. Gated off release: the + // D2H+sync would reintroduce the per-batch host block this path exists to + // avoid, and a zero total is unreachable with honest inputs — a hit here + // is a construction or kernel bug, which tests/CI are the place to catch. + #[cfg(debug_assertions)] + { + let mut host = [0u64; 3]; + stream.memcpy_dtoh(&out.slice(0..3), &mut host)?; + stream.synchronize()?; + assert_ne!( + host, [0u64; 3], + "batch inverse: zero total has no inverse (a denominator was zero)" + ); + } Ok(()) } @@ -138,8 +158,6 @@ pub fn batch_inverse_ext3_dev( n: usize, stream: &Arc, ) -> Result> { - #[cfg(feature = "test-faults")] - check_inverse_fault_injection()?; assert!(n >= 1, "batch_inverse_ext3_dev requires n >= 1"); // Runtime guard (not debug_assert): a u32 grid_dim is truncated past // u32::MAX / BLOCK_SIZE, which would silently launch too few blocks @@ -227,6 +245,12 @@ pub fn compute_and_invert_denoms_ext3_dev( sign: DenomSign, stream: &Arc, ) -> Result> { + // Fault-injection hook lives here (not in the shared `batch_inverse_ext3_dev`) + // so `schedule_inverse_fault(N)` targets exactly the Nth R3/R4 denominator + // inversion the fallback test exercises — not the LogUp aux inverses that + // also route through `batch_inverse_ext3_dev` earlier in the prove. + #[cfg(feature = "test-faults")] + check_inverse_fault_injection()?; assert_eq!(z_scalars_host.len(), k_scalars * 3); assert!(n >= 1 && k_scalars >= 1); From 05fb8dd25b97e3097a5ac269c47fffdf0f92baa3 Mon Sep 17 00:00:00 2001 From: Joaquin Carletti Date: Fri, 31 Jul 2026 17:17:00 -0300 Subject: [PATCH 04/12] perf(gpu): stage htod_via in fixed 64MB chunks to bound pinned footprint --- crypto/math-cuda/src/device.rs | 89 +++++++++++++++++++++++----------- 1 file changed, 62 insertions(+), 27 deletions(-) diff --git a/crypto/math-cuda/src/device.rs b/crypto/math-cuda/src/device.rs index c3dd4a2a9..b0620c715 100644 --- a/crypto/math-cuda/src/device.rs +++ b/crypto/math-cuda/src/device.rs @@ -630,11 +630,26 @@ impl Drop for PendingD2H<'_> { } } -/// Host→device copy staged through the pinned slot: one host memcpy into -/// pinned memory + one async DMA, instead of the driver's internal pageable -/// staging (small chunks; 2-3x slower for multi-hundred-MB traces and it -/// convoys under multi-thread load). Blocks until the DMA lands, so the slot -/// and `src_host` are both reusable on return. +/// Chunk size for [`htod_via`]'s staged upload — the upper bound a single H2D +/// puts on a staging slot's page-locked footprint. 64 MB is large enough to +/// amortize the per-chunk DMA launch + event sync, small enough to keep the +/// pinned slab independent of trace size. +const HTOD_CHUNK_BYTES: usize = 64 << 20; // 64 MB + +/// Host→device copy staged through the pinned slot, in fixed-size chunks: each +/// chunk is one host memcpy into pinned memory + one async DMA, instead of the +/// driver's internal pageable staging (small chunks; 2-3x slower for +/// multi-hundred-MB traces and it convoys under multi-thread load). Blocks +/// until the last DMA lands, so the slot and `src_host` are both reusable on +/// return. +/// +/// Chunking caps the slot's page-locked footprint at [`HTOD_CHUNK_BYTES`] +/// regardless of trace size. This matters on the device-only path +/// (`retain_host_lde = false`): there is no [`async_dtoh_via`] drain to size +/// the slot, so `htod_via` is its only writer — an uncapped copy would grow +/// the per-worker slab to a whole trace and, being grow-only, never shrink it. +/// The host-retaining path is unaffected: its later `async_dtoh_via` grows the +/// same slot to the full LDE anyway, and we simply reuse the first chunk of it. pub fn htod_via( stream: &Arc, slot: &Mutex, @@ -651,31 +666,51 @@ pub fn htod_via( if n_bytes == 0 { return Ok(()); } - let u64_len = n_bytes.div_ceil(8); + let elem_size = std::mem::size_of::(); + // Chunk in whole elements so a `T` never straddles a chunk boundary. + let chunk_elems = (HTOD_CHUNK_BYTES / elem_size.max(1)).max(1); + let mut staging = slot.lock().unwrap(); - staging.ensure_capacity(u64_len, ctx)?; + // Only ask for a chunk's worth of pinned memory (or the whole copy when + // smaller). If another path (`async_dtoh_via` on the host-retaining flow) + // already grew this slot larger, it stays larger — grow-only — and we just + // use the first chunk of it. + let want_u64 = (chunk_elems * elem_size).div_ceil(8).min(n_bytes.div_ceil(8)); + staging.ensure_capacity(want_u64, ctx)?; ctx.bind_to_thread()?; - // SAFETY: the pinned allocation is stable while the lock is held and at - // least `n_bytes` long (`ensure_capacity`). The DMA reads it after the - // host memcpy (program order); `device_ptr_mut` orders the device write - // on `stream`. - unsafe { - std::ptr::copy_nonoverlapping( - src_host.as_ptr() as *const u8, - staging.ptr as *mut u8, - n_bytes, - ); - let (dst_ptr, _record) = dst.device_ptr_mut(stream); - cudarc::driver::sys::cuMemcpyHtoDAsync_v2( - dst_ptr, - staging.ptr as *const core::ffi::c_void, - n_bytes, - stream.cu_stream(), - ) - .result()?; + + // SAFETY: `device_ptr_mut` yields the destination base pointer and orders + // the device writes on `stream`; `dst.len() >= src_host.len()` (asserted), + // so every chunk's byte range stays within `dst`. + let (dst_base, _record) = dst.device_ptr_mut(stream); + let src = src_host.as_ptr() as *const u8; + let n_elems = src_host.len(); + let mut elem_off = 0usize; + while elem_off < n_elems { + let this_elems = (n_elems - elem_off).min(chunk_elems); + let this_bytes = this_elems * elem_size; + let byte_off = elem_off * elem_size; + // SAFETY: the pinned slab holds at least `chunk_elems * elem_size` + // bytes (or the whole copy when smaller). The previous chunk's DMA is + // synced below before this memcpy overwrites the slab, so the slab is + // never read (by an in-flight DMA) and written at the same time. + unsafe { + std::ptr::copy_nonoverlapping(src.add(byte_off), staging.ptr as *mut u8, this_bytes); + cudarc::driver::sys::cuMemcpyHtoDAsync_v2( + dst_base + byte_off as u64, + staging.ptr as *const core::ffi::c_void, + this_bytes, + stream.cu_stream(), + ) + .result()?; + } + // Single-buffered: wait for this chunk's DMA before the next memcpy + // reuses the slab. + staging.record_event(stream)?; + staging.sync_event()?; + elem_off += this_elems; } - staging.record_event(stream)?; - staging.sync_event() + Ok(()) } /// Enqueue an async D2H of `n_elems` of `src` into the pinned slab of `slot`, From 51536c2ae8a82b8ca7a7dd81f75ed5e4fd486073 Mon Sep 17 00:00:00 2001 From: Joaquin Carletti Date: Fri, 31 Jul 2026 17:24:52 -0300 Subject: [PATCH 05/12] style: rustfmt htod_via chunk-size expression --- crypto/math-cuda/src/device.rs | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/crypto/math-cuda/src/device.rs b/crypto/math-cuda/src/device.rs index b0620c715..b7b837393 100644 --- a/crypto/math-cuda/src/device.rs +++ b/crypto/math-cuda/src/device.rs @@ -675,7 +675,9 @@ pub fn htod_via( // smaller). If another path (`async_dtoh_via` on the host-retaining flow) // already grew this slot larger, it stays larger — grow-only — and we just // use the first chunk of it. - let want_u64 = (chunk_elems * elem_size).div_ceil(8).min(n_bytes.div_ceil(8)); + let want_u64 = (chunk_elems * elem_size) + .div_ceil(8) + .min(n_bytes.div_ceil(8)); staging.ensure_capacity(want_u64, ctx)?; ctx.bind_to_thread()?; From db04e557e57d7dded4ad0cd96292d06ecdcd3642 Mon Sep 17 00:00:00 2001 From: Joaquin Carletti Date: Fri, 31 Jul 2026 18:01:25 -0300 Subject: [PATCH 06/12] fix(gpu): gate R2 comp-tree host fallback on the parts, not host_trace_empty --- crypto/stark/src/prover.rs | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/crypto/stark/src/prover.rs b/crypto/stark/src/prover.rs index 8612bb20a..d8cf4bf87 100644 --- a/crypto/stark/src/prover.rs +++ b/crypto/stark/src/prover.rs @@ -1644,9 +1644,16 @@ pub trait IsStarkProver< None => { // The host part evals are empty under device-only (the R2 // drain is skipped); abort with the device-only contract's - // message instead of a misleading EmptyCommitment. + // message instead of a misleading EmptyCommitment. Gate on + // the parts the CPU fallback actually consumes, not on + // `host_trace_empty()`: the trace can stay device-resident + // while these parts were downloaded to the host anyway (the + // GPU decompose fell back to `decompose_and_extend_d2`), in + // which case this fallback is valid and must not panic. assert!( - !round_1_result.lde_trace.host_trace_empty(), + lde_composition_poly_parts_evaluations + .first() + .is_none_or(|p| !p.is_empty()), "R2 composition commit fell back to the host part evals, \ but they are device-only (empty)" ); From e75bcbed85110ad0347618075c94631f7287dadd Mon Sep 17 00:00:00 2001 From: Joaquin Carletti Date: Mon, 3 Aug 2026 14:50:48 -0300 Subject: [PATCH 07/12] =?UTF-8?q?chore(gpu):=20review=20follow-ups=20?= =?UTF-8?q?=E2=80=94=20gather=20bounds,=20release=20canaries,=20live=20zer?= =?UTF-8?q?o-total=20guard?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - gather_ext3_at asserts positions against the evals buffer host-side (same guard as gather_merkle_paths_dev). - The device-gather cross-checks keep query 0 as a release canary instead of paying every query; debug still checks all of them. - The batch-inverse zero-total guard also compiles under test-faults, so the GPU fallback suite (which runs --release) actually exercises it. - New htod_via round-trip test covering the 64 MB chunk loop and its partial tail. --- crypto/math-cuda/src/fri.rs | 8 +++++ crypto/math-cuda/src/inverse.rs | 15 +++++----- crypto/math-cuda/tests/htod_via.rs | 48 ++++++++++++++++++++++++++++++ crypto/stark/src/prover.rs | 17 +++++++---- 4 files changed, 75 insertions(+), 13 deletions(-) create mode 100644 crypto/math-cuda/tests/htod_via.rs diff --git a/crypto/math-cuda/src/fri.rs b/crypto/math-cuda/src/fri.rs index 8da80472d..a3c29dedc 100644 --- a/crypto/math-cuda/src/fri.rs +++ b/crypto/math-cuda/src/fri.rs @@ -264,6 +264,14 @@ pub fn gather_ext3_at( if q == 0 { return Ok(Vec::new()); } + // Guard the kernel's device reads: a position past the evals buffer would + // be a silent out-of-bounds read. Positions are valid by construction; + // this catches a caller bug host-side before it becomes device garbage + // (matching `gather_merkle_paths_dev`). + assert!( + positions.iter().all(|&p| (p as usize) < evals.len() / 3), + "gather_ext3_at: position >= evals length" + ); let be = backend()?; let pos_dev = stream.clone_htod(positions)?; // SAFETY: the gather kernel writes all 3 * q slots. diff --git a/crypto/math-cuda/src/inverse.rs b/crypto/math-cuda/src/inverse.rs index 833bf2906..1087e2ae4 100644 --- a/crypto/math-cuda/src/inverse.rs +++ b/crypto/math-cuda/src/inverse.rs @@ -127,14 +127,15 @@ fn launch_invert_total( .arg(&mut *out) .launch(cfg)?; } - // Debug-only invertibility guard. The Fermat kernel maps a zero total - // (some denominator was zero) silently to zero, so the batch would ship + // Invertibility guard. The Fermat kernel maps a zero total (some + // denominator was zero) silently to zero, so the batch would ship // all-zero "inverses" instead of erroring. A valid inverse is never zero, - // so `out == 0` unambiguously flags a zero total. Gated off release: the - // D2H+sync would reintroduce the per-batch host block this path exists to - // avoid, and a zero total is unreachable with honest inputs — a hit here - // is a construction or kernel bug, which tests/CI are the place to catch. - #[cfg(debug_assertions)] + // so `out == 0` unambiguously flags a zero total. Gated off plain release + // (the D2H+sync would reintroduce the per-batch host block this path + // exists to avoid); `test-faults` keeps it live in the GPU fallback + // suite, which runs --release — a hit is a construction or kernel bug, + // and that suite is where CI can actually catch it. + #[cfg(any(debug_assertions, feature = "test-faults"))] { let mut host = [0u64; 3]; stream.memcpy_dtoh(&out.slice(0..3), &mut host)?; diff --git a/crypto/math-cuda/tests/htod_via.rs b/crypto/math-cuda/tests/htod_via.rs new file mode 100644 index 000000000..6db1eb227 --- /dev/null +++ b/crypto/math-cuda/tests/htod_via.rs @@ -0,0 +1,48 @@ +//! Round-trip coverage for `htod_via`'s chunk loop: uploads larger than the +//! 64 MB pinned-staging chunk must arrive intact across every chunk boundary +//! (a stale slab or a bad byte offset would corrupt exactly one chunk). + +use math_cuda::device::{backend, htod_via}; +use rand::{Rng, SeedableRng}; +use rand_chacha::ChaCha8Rng; + +fn roundtrip(n_u64: usize, seed: u64) { + let mut rng = ChaCha8Rng::seed_from_u64(seed); + let src: Vec = (0..n_u64).map(|_| rng.r#gen::()).collect(); + + let be = backend().expect("cuda backend"); + let stream = be.next_stream(); + let mut dst = stream.alloc_zeros::(n_u64).expect("device alloc"); + htod_via( + &stream, + be.pinned_staging(), + &be.ctx, + &src, + &mut dst.slice_mut(0..n_u64), + ) + .expect("htod_via"); + + let back = stream.clone_dtoh(&dst).expect("dtoh"); + stream.synchronize().expect("sync"); + assert_eq!(src.len(), back.len()); + // Compare in chunks so a failure names the offset instead of dumping 100M+ values. + for (i, (a, b)) in src.iter().zip(back.iter()).enumerate() { + assert_eq!( + a, b, + "htod_via round-trip mismatch at u64 offset {i} (n={n_u64})" + ); + } +} + +#[test] +fn htod_via_single_chunk_roundtrip() { + // Below the 64 MB chunk: single iteration of the loop. + roundtrip(1 << 20, 42); +} + +#[test] +fn htod_via_multi_chunk_roundtrip() { + // 3 full chunks + a partial tail: exercises slab reuse across iterations + // and the final short chunk. 64 MB chunk = 2^23 u64s. + roundtrip((3 << 23) + 12345, 43); +} diff --git a/crypto/stark/src/prover.rs b/crypto/stark/src/prover.rs index d8cf4bf87..2e5a59016 100644 --- a/crypto/stark/src/prover.rs +++ b/crypto/stark/src/prover.rs @@ -2521,8 +2521,10 @@ pub trait IsStarkProver< // Cross-check the device gather against the host LDE. Skipped under // device-only (host trace empty): the gather was proven bit-identical // while the host copy was resident, and there is nothing to check - // against. - if !lde_trace.host_trace_empty() { + // against. Release keeps query 0 as a canary (the GPU test suites run + // --release, and gather failure modes — stride/offset/layout — are + // systematic, so one query catches them); debug checks every query. + if (cfg!(debug_assertions) || qi == 0) && !lde_trace.host_trace_empty() { let domain_size = domain.lde_roots_of_unity_coset.len() as u64; let r_even = reverse_index(challenge * 2, domain_size); let r_odd = reverse_index(challenge * 2 + 1, domain_size); @@ -2799,10 +2801,13 @@ pub trait IsStarkProver< // Cross-check against the host part evals while // they are still resident (absent under full // residency, where the gather is the only source). - if round_2_result - .lde_composition_poly_evaluations - .first() - .is_some_and(|p| !p.is_empty()) + // Query 0 stays a release canary, same rationale + // as `open_trace_polys_device`. + if (cfg!(debug_assertions) || qi == 0) + && round_2_result + .lde_composition_poly_evaluations + .first() + .is_some_and(|p| !p.is_empty()) { let expected = Self::open_composition_poly_with_proof( proofs[qi].clone(), From 4789c87be45c1a76d7d667d48476330729fbbbbc Mon Sep 17 00:00:00 2001 From: Mauro Toscano <12560266+MauroToscano@users.noreply.github.com> Date: Mon, 3 Aug 2026 16:12:27 -0300 Subject: [PATCH 08/12] fix(gpu): drain htod_via on error; narrow the merkle-tail threshold (#892) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix(gpu): drain htod_via on error, guard the R2 host-evaluator fallback Review follow-ups for the round-2 residency work, rebased onto e75bcbed — only the items that commit did not already cover. htod_via error path. Once a chunk's DMA is in flight, `record_event` / `sync_event` returning `Err` drops the staging `MutexGuard` with the device still reading the pinned slab, so the next locker's `ensure_capacity` can `cuMemFreeHost` it mid-copy. `async_dtoh_via` already guards this exact hazard and the file ships a `DrainOnErr` helper for it; `htod_via` was the one site not using it. R2 host-evaluator fallback. If the device decompose and the `H` download both fail under device-only, control reaches the host evaluator, which reads the intentionally-empty trace and panics with a bare out-of-bounds. Assert the device-only contract instead, matching the other fallback arms. Coverage. `batch_inverse_ext3_dev`'s `n == 1` branch is never exercised — `batch_inverse_n1` goes through the host-only short circuit in `batch_inverse_ext3`, as its own comment says. Add a direct device test. Docs. The preprocessed split-tree comment still claimed both trees come back as full host trees (the multiplicity tree is root-only + device resident), and `FriCommitState`'s doc claimed its input is always Arc-shared with a retained `gpu_evals` (only true on the device-only path). * perf(gpu): set the merkle-tail threshold to the block width TAIL_MAX_PAIRS = 2048 overshoots. The tail grid-strides a single 128-thread block on one SM, so a level of k pairs is k/128 SEQUENTIAL keccak-f1600s where the per-level launches it replaces spread them over k/128 parallel blocks. At 2048 the first four levels alone are 16+8+4+2 = 30 serial permutations against 4 parallel waves — order +100 us per large tree to save 4 launches worth order 10 us, and it sits on the critical path because the caller's 32-byte root memcpy_dtoh host-blocks on everything queued before it. At the block width the entry level is exactly one permutation per thread, so the tail still collapses the top levels into one launch but adds no serialization at all. --- crypto/math-cuda/src/device.rs | 23 ++++++++++++++++++++--- crypto/math-cuda/src/fri.rs | 6 ++++-- crypto/math-cuda/src/merkle.rs | 17 ++++++++++++++--- crypto/math-cuda/tests/batch_inverse.rs | 25 +++++++++++++++++++++++++ crypto/stark/src/prover.rs | 24 ++++++++++++++++++++---- 5 files changed, 83 insertions(+), 12 deletions(-) diff --git a/crypto/math-cuda/src/device.rs b/crypto/math-cuda/src/device.rs index b7b837393..8bc140f21 100644 --- a/crypto/math-cuda/src/device.rs +++ b/crypto/math-cuda/src/device.rs @@ -685,6 +685,15 @@ pub fn htod_via( // the device writes on `stream`; `dst.len() >= src_host.len()` (asserted), // so every chunk's byte range stays within `dst`. let (dst_base, _record) = dst.device_ptr_mut(stream); + // Declared after the slot's MutexGuard so it drops FIRST: once a chunk's + // DMA is in flight, any `?`-return below must drain the stream before the + // guard releases the slot, or the next locker's `ensure_capacity` could + // `cuMemFreeHost` the slab while the device is still reading it. Same + // hazard `async_dtoh_via` guards against on its record-event failure. + let mut drain = DrainOnErr { + stream, + armed: false, + }; let src = src_host.as_ptr() as *const u8; let n_elems = src_host.len(); let mut elem_off = 0usize; @@ -698,18 +707,26 @@ pub fn htod_via( // never read (by an in-flight DMA) and written at the same time. unsafe { std::ptr::copy_nonoverlapping(src.add(byte_off), staging.ptr as *mut u8, this_bytes); - cudarc::driver::sys::cuMemcpyHtoDAsync_v2( + let r = cudarc::driver::sys::cuMemcpyHtoDAsync_v2( dst_base + byte_off as u64, staging.ptr as *const core::ffi::c_void, this_bytes, stream.cu_stream(), ) - .result()?; + .result(); + // Armed even on failure: the driver may have enqueued the copy + // before reporting the error. + drain.armed = true; + r?; } // Single-buffered: wait for this chunk's DMA before the next memcpy - // reuses the slab. + // reuses the slab. Both calls can fail with the DMA still in flight, + // which is what `drain` covers. staging.record_event(stream)?; staging.sync_event()?; + // This chunk has landed; nothing is reading the slab until the next + // iteration re-arms. + drain.armed = false; elem_off += this_elems; } Ok(()) diff --git a/crypto/math-cuda/src/fri.rs b/crypto/math-cuda/src/fri.rs index a3c29dedc..533ff6e32 100644 --- a/crypto/math-cuda/src/fri.rs +++ b/crypto/math-cuda/src/fri.rs @@ -41,8 +41,10 @@ fn check_fault_injection() -> Result<()> { } /// Device-side state across FRI commit iterations. Owns the current fold -/// input (the previous layer's evals, Arc-shared with that layer's retained -/// `gpu_evals`) and the inv_twiddles buffer. Freed when dropped. +/// input (the previous layer's evals) and the inv_twiddles buffer. The input +/// is an `Arc` because the caller may also retain it as that layer's +/// `gpu_evals` — it does so only on the device-only path, where no host copy +/// of the evals exists. Freed when the last holder drops. pub struct FriCommitState { pub stream: Arc, /// Current fold input. Each fold allocates a fresh output buffer that is diff --git a/crypto/math-cuda/src/merkle.rs b/crypto/math-cuda/src/merkle.rs index 1cb1c5b6f..c499df702 100644 --- a/crypto/math-cuda/src/merkle.rs +++ b/crypto/math-cuda/src/merkle.rs @@ -160,9 +160,20 @@ pub(crate) fn build_inner_tree_levels( ) -> Result<()> { // Once a level fits this many pairs, one single-block launch // (`keccak_merkle_tail`) builds all remaining levels with barriers - // between them: the top ~11 levels of a big tree are each smaller than - // the per-launch overhead they used to pay. - const TAIL_MAX_PAIRS: u64 = 2048; + // between them: the top levels of a big tree are each smaller than the + // per-launch overhead they used to pay. + // + // Set to the block width, so the entry level is exactly one permutation + // per thread and the tail adds NO serialization over the per-level + // launches it replaces. Going wider is not free: the tail grid-strides a + // single 128-thread block on one SM, so a level of `k` pairs costs + // `k / 128` *sequential* keccak-f1600s where separate launches would have + // spread them over `k / 128` parallel blocks. At 2048 the first four + // levels alone are 16+8+4+2 = 30 serial permutations against 4 parallel + // waves — order +100 us per large tree, to save 4 launches worth order + // 10 us. It stays on the critical path because the caller's 32-byte root + // `memcpy_dtoh` host-blocks on everything queued before it. + const TAIL_MAX_PAIRS: u64 = KECCAK_BLOCK_DIM as u64; let mut level_begin: u64 = (leaves_len - 1) as u64; while level_begin != 0 { let new_begin = level_begin / 2; diff --git a/crypto/math-cuda/tests/batch_inverse.rs b/crypto/math-cuda/tests/batch_inverse.rs index bc52f9fcb..087a0b082 100644 --- a/crypto/math-cuda/tests/batch_inverse.rs +++ b/crypto/math-cuda/tests/batch_inverse.rs @@ -72,6 +72,31 @@ fn batch_inverse_n1() { run(1, 1); } +/// `batch_inverse_ext3_dev`'s own `n == 1` branch, which the host entry point +/// above never reaches: `batch_inverse_ext3` short-circuits n==1 to +/// `invert_ext3_host`, so only a direct device call exercises the single +/// `invert_total_ext3` launch that serves this case. +#[test] +fn batch_inverse_dev_n1() { + let mut rng = ChaCha8Rng::seed_from_u64(7); + let x = rand_fp3_nonzero(&mut rng); + let expected = x.inv().expect("nonzero is invertible"); + + let be = math_cuda::device::backend().expect("cuda backend"); + let stream = be.next_stream(); + let input = stream.clone_htod(&ext3_to_u64s(&[x])).unwrap(); + + let out_dev = math_cuda::inverse::batch_inverse_ext3_dev(&input, 1, &stream).unwrap(); + let got = stream.clone_dtoh(&out_dev).unwrap(); + stream.synchronize().unwrap(); + + assert_eq!( + canon3(&got), + canon3(&ext3_to_u64s(&[expected])), + "device n==1 inverse" + ); +} + #[test] fn batch_inverse_single_block() { // All single-block sizes (no recursion). diff --git a/crypto/stark/src/prover.rs b/crypto/stark/src/prover.rs index 2e5a59016..42142f770 100644 --- a/crypto/stark/src/prover.rs +++ b/crypto/stark/src/prover.rs @@ -1027,10 +1027,12 @@ pub trait IsStarkProver< // Fused GPU split path for preprocessed tables (cuda only): one // row-major LDE of ALL columns plus two subset Merkle trees // (precomputed / multiplicity) built on device — leaves and levels are - // bit-identical to `commit_rows_bit_reversed_subset`, and the trees - // come back as full host trees so the preprocessed opening path and - // the process-wide precomputed-tree cache work unchanged. The handle - // keeps the LDE device-resident for the downstream GPU rounds. + // bit-identical to `commit_rows_bit_reversed_subset`. The precomputed + // tree comes back as a full host tree, so the process-wide + // precomputed-tree cache works unchanged; the multiplicity tree stays + // device-resident behind a root-only host tree and its opening paths + // are gathered on device. The handle keeps the LDE device-resident for + // the downstream GPU rounds. #[cfg(feature = "cuda")] if let Some((expected_precomputed_root, num_precomputed)) = precomputed { let (trace_slice, num_cols) = trace.main_data_row_major(); @@ -1528,6 +1530,20 @@ pub trait IsStarkProver< #[cfg(feature = "instruments")] let t_sub = Instant::now(); + // Every arm below runs the HOST evaluator, which reads `get_main` / + // `get_aux`. Under device-only those buffers are intentionally empty, + // so landing here means the device decompose AND the `H` download both + // failed. Abort with the device-only contract's message rather than a + // bare index-out-of-bounds from somewhere inside the evaluator. + #[cfg(feature = "cuda")] + if precomputed_parts.is_none() { + assert!( + !round_1_result.lde_trace.host_trace_empty(), + "R2 composition fell back to the host evaluator, but the trace \ + is device-only (empty)" + ); + } + let lde_composition_poly_parts_evaluations = if let Some(parts) = precomputed_parts { parts } else if number_of_parts == 2 { From 9d2140b1ee36580303c831567c151411a863c812 Mon Sep 17 00:00:00 2001 From: Joaquin Carletti Date: Wed, 29 Jul 2026 16:43:43 -0300 Subject: [PATCH 09/12] perf(prover): replace table chunks with a VRAM-admitted per-table scheduler MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Fiat-Shamir only requires the main roots absorbed in index order before the shared challenges; past that fork every table's chain is independent. Phase A now runs all main commits under a byte-budget admission gate (no chunk barriers), and aux build, aux commit and rounds 2-4 run fused as one task per table, heaviest first — while a big table works through a host-bound stretch, the other tables' GPU stages fill the device. GPU builds default TABLE_PARALLELISM to 2/3 of the cores (swept flat at 10 on a 16-core RTX 5090). ethrex 10tx continuations on RTX 5090: 10.64s -> 8.54s (-19.7%, 8 ABBA pairs). --- crypto/stark/src/prover.rs | 704 ++++++++++++++++++++----------------- 1 file changed, 384 insertions(+), 320 deletions(-) diff --git a/crypto/stark/src/prover.rs b/crypto/stark/src/prover.rs index 42142f770..7d979bcbc 100644 --- a/crypto/stark/src/prover.rs +++ b/crypto/stark/src/prover.rs @@ -20,10 +20,7 @@ use math::{ }; #[cfg(feature = "parallel")] -use rayon::prelude::{ - IndexedParallelIterator, IntoParallelIterator, IntoParallelRefIterator, - IntoParallelRefMutIterator, ParallelIterator, -}; +use rayon::prelude::{IntoParallelIterator, ParallelIterator}; #[cfg(feature = "debug-checks")] use crate::debug::validate_trace; @@ -581,7 +578,18 @@ pub fn table_parallelism() -> usize { let cores = std::thread::available_parallelism() .map(|n| n.get()) .unwrap_or(4); - (cores / 3).max(1) + // GPU builds: with the admission scheduler most in-flight + // tables sit in GPU waits, so more of them pay (swept flat at + // ~2/3 of the cores on a 16-core/RTX 5090 box). CPU builds + // stay at cores/3 — every table is pure host work there. + #[cfg(feature = "cuda")] + { + (cores * 2 / 3).max(1) + } + #[cfg(not(feature = "cuda"))] + { + (cores / 3).max(1) + } }) } #[cfg(not(feature = "parallel"))] @@ -613,28 +621,100 @@ fn estimate_table_vram_bytes(main_cols: usize, aux_cols: usize, lde_size: usize) /// or VRAM not binding) chunks fall back to fixed size `k`, identical to the /// old `step_by(k)`, so scheduling and the proof are unchanged. Returns /// `(start, end)` half open ranges covering `0..estimates.len()` in order. -fn plan_table_chunks(estimates: &[u64], k: usize, budget: u64) -> Vec<(usize, usize)> { - let n = estimates.len(); - let k = k.max(1); - let budget = budget as u128; - let mut chunks = Vec::new(); - let mut start = 0; - while start < n { - let mut end = start; - let mut acc: u128 = 0; - while end < n { - let next = estimates[end] as u128; - // Always admit at least one table per chunk (oversized → solo). - if end > start && (end - start >= k || acc + next > budget) { - break; +/// Byte-budget admission gate for concurrently proven tables. `acquire` +/// blocks until the requested bytes fit under the budget, releasing on +/// permit drop. An oversized request is admitted alone (when nothing else +/// holds bytes), so tables larger than the whole budget still prove. +/// +/// 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. +struct VramGate { + used: std::sync::Mutex, + freed: std::sync::Condvar, + budget: u64, +} + +struct VramPermit<'a> { + gate: &'a VramGate, + bytes: u64, +} + +impl VramGate { + fn new(budget: u64) -> Self { + Self { + used: std::sync::Mutex::new(0), + freed: std::sync::Condvar::new(), + budget, + } + } + + fn acquire(&self, bytes: u64) -> VramPermit<'_> { + let mut used = self.used.lock().unwrap(); + loop { + if *used == 0 || used.saturating_add(bytes) <= self.budget { + *used = used.saturating_add(bytes); + return VramPermit { gate: self, bytes }; } - acc += next; - end += 1; + used = self.freed.wait(used).unwrap(); } - chunks.push((start, end)); - start = end; } - chunks +} + +impl Drop for VramPermit<'_> { + fn drop(&mut self) { + let mut used = self.gate.used.lock().unwrap(); + *used = used.saturating_sub(self.bytes); + drop(used); + self.gate.freed.notify_all(); + } +} + +/// 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. +fn run_admitted( + order: &[usize], + estimates: &[u64], + gate: &VramGate, + workers: usize, + task: impl Fn(usize) -> T + Sync, +) -> Vec> { + let results: Vec>> = estimates + .iter() + .map(|_| std::sync::Mutex::new(None)) + .collect(); + let cursor = std::sync::atomic::AtomicUsize::new(0); + std::thread::scope(|scope| { + for _ in 0..workers.max(1).min(order.len().max(1)) { + scope.spawn(|| { + loop { + let pos = cursor.fetch_add(1, std::sync::atomic::Ordering::Relaxed); + if pos >= order.len() { + return; + } + let idx = order[pos]; + let permit = gate.acquire(estimates[idx]); + let out = task(idx); + *results[idx].lock().unwrap() = Some(out); + drop(permit); + } + }); + } + }); + results + .into_iter() + .map(|m| m.into_inner().unwrap()) + .collect() +} + +/// Table indices sorted heaviest-first by estimate. +fn heaviest_first(estimates: &[u64]) -> Vec { + let mut order: Vec = (0..estimates.len()).collect(); + order.sort_by_key(|&i| std::cmp::Reverse(estimates[i])); + order } /// A container for the results of the second round of the STARK Prove protocol. @@ -1321,7 +1401,7 @@ pub trait IsStarkProver< /// validate each trace. Called once after Phase C commits. #[cfg(feature = "debug-checks")] fn run_debug_checks( - air_trace_pairs: &[AirTracePair<'_, Field, FieldExtension, PI>], + pair_cells: &[std::sync::Mutex>], commitments: &[Round1Commitments], domains: &[Arc>], twiddle_caches: &[Arc>], @@ -1331,13 +1411,15 @@ pub trait IsStarkProver< PI: Send + Sync + Clone, { let mut temp_results: Vec> = - Vec::with_capacity(air_trace_pairs.len()); - for (((air, trace, _), commitment), (domain, twiddles)) in air_trace_pairs + Vec::with_capacity(pair_cells.len()); + for ((cell, commitment), (domain, twiddles)) in pair_cells .iter() .zip(commitments.iter()) .zip(domains.iter().zip(twiddle_caches.iter())) { - let result = Self::reconstruct_round1(*air, *trace, domain, commitment, twiddles) + let pair = cell.lock().unwrap(); + let (air, trace, _) = &*pair; + let result = Self::reconstruct_round1(*air, trace, domain, commitment, twiddles) .expect("reconstruct_round1 failed in debug-checks"); temp_results.push(result); } @@ -1348,15 +1430,17 @@ pub trait IsStarkProver< .collect(); print_bus_balance_report(&all_bus_public_inputs); - for (((air, trace, pub_inputs), round_1_result), domain) in air_trace_pairs + for ((cell, round_1_result), domain) in pair_cells .iter() .zip(temp_results.iter()) .zip(domains.iter()) { + let pair = cell.lock().unwrap(); + let (air, trace, pub_inputs) = &*pair; validate_trace( *air, *pub_inputs, - *trace, + trace, domain, &round_1_result.rap_challenges, round_1_result.bus_public_inputs.as_ref(), @@ -2933,7 +3017,7 @@ pub trait IsStarkProver< /// /// The transcript must be safely initialized before passing it to this method. fn multi_prove( - mut air_trace_pairs: Vec>, + #[allow(unused_mut)] mut air_trace_pairs: Vec>, transcript: &mut (impl IsStarkTranscript + Clone + Send), #[cfg(feature = "disk-spill")] storage_mode: StorageMode, ) -> Result, ProvingError> @@ -3000,20 +3084,18 @@ pub trait IsStarkProver< // don't re-add pre-sizing without a shared-slab design that bounds the // number of allocations. + let vram_gate = VramGate::new(vram_budget); + // R1 main commit: only the main LDE and its Merkle scratch are resident, // so the aux columns add nothing to this phase's working set. - let main_chunks = { - let estimates: Vec = air_trace_pairs - .iter() - .enumerate() - .map(|(idx, (_, trace, _))| { - let lde_size = - domains[idx].interpolation_domain_size * domains[idx].blowup_factor; - estimate_table_vram_bytes(trace.num_main_columns, 0, lde_size) - }) - .collect(); - plan_table_chunks(&estimates, k, vram_budget) - }; + let main_estimates: Vec = air_trace_pairs + .iter() + .enumerate() + .map(|(idx, (_, trace, _))| { + let lde_size = domains[idx].interpolation_domain_size * domains[idx].blowup_factor; + estimate_table_vram_bytes(trace.num_main_columns, 0, lde_size) + }) + .collect(); // Spill main traces to mmap before Round 1 LDE. #[cfg(feature = "disk-spill")] @@ -3055,51 +3137,55 @@ pub trait IsStarkProver< let mut main_gpu_handles: Vec> = Vec::with_capacity(num_airs); - for &(chunk_start, chunk_end) in &main_chunks { - let chunk_range = chunk_start..chunk_end; - - let chunk_results: Vec> = - crate::par::par_map_collect(chunk_range, |idx| { - let (air, trace, _) = &air_trace_pairs[idx]; - let domain = &domains[idx]; - let twiddles = &twiddle_caches[idx]; - - let precomputed = air - .is_preprocessed() - .then(|| (air.precomputed_commitment(), air.num_precomputed_columns())); + // All main commits with continuous VRAM admission (no chunk barriers); + // the transcript only needs the roots absorbed in index order, done + // sequentially below once every commit completed — the one ordering + // Fiat-Shamir requires before sampling the shared challenges. + let main_results = run_admitted( + &heaviest_first(&main_estimates), + &main_estimates, + &vram_gate, + k, + |idx| { + let (air, trace, _) = &air_trace_pairs[idx]; + let domain = &domains[idx]; + let twiddles = &twiddle_caches[idx]; + + let precomputed = air + .is_preprocessed() + .then(|| (air.precomputed_commitment(), air.num_precomputed_columns())); + + // Stage-3 device-only gate: when it holds, `commit_main_trace` + // keeps the R1 LDE device-resident and skips the host D2H. + #[cfg(feature = "cuda")] + let device_only = Self::device_only_for(*air, domain); - // Stage-3 device-only gate: when it holds, `commit_main_trace` - // keeps the R1 LDE device-resident and skips the host D2H. + Self::commit_main_trace( + *trace, + domain, + twiddles, + precomputed, #[cfg(feature = "cuda")] - let device_only = Self::device_only_for(*air, domain); - - Self::commit_main_trace( - *trace, - domain, - twiddles, - precomputed, - #[cfg(feature = "cuda")] - device_only, - #[cfg(feature = "disk-spill")] - storage_mode, - ) - }); - - // Sequential: append roots to shared transcript (Fiat-Shamir ordering) - for result in chunk_results { - #[cfg(feature = "cuda")] - let (commit, cached_main, gpu_main) = result?; - #[cfg(not(feature = "cuda"))] - let (commit, cached_main) = result?; - if let Some(ref pre_root) = commit.precomputed_root { - transcript.append_bytes(pre_root); - } - transcript.append_bytes(&commit.root); - main_commits.push(commit); - main_ldes.push(cached_main); - #[cfg(feature = "cuda")] - main_gpu_handles.push(gpu_main); + device_only, + #[cfg(feature = "disk-spill")] + storage_mode, + ) + }, + ); + for result in main_results { + let result = result.expect("run_admitted fills every slot"); + #[cfg(feature = "cuda")] + let (commit, cached_main, gpu_main) = result?; + #[cfg(not(feature = "cuda"))] + let (commit, cached_main) = result?; + if let Some(ref pre_root) = commit.precomputed_root { + transcript.append_bytes(pre_root); } + transcript.append_bytes(&commit.root); + main_commits.push(commit); + main_ldes.push(cached_main); + #[cfg(feature = "cuda")] + main_gpu_handles.push(gpu_main); } #[cfg(feature = "instruments")] @@ -3134,13 +3220,9 @@ pub trait IsStarkProver< // Pass 1 (parallel): Build all auxiliary traces (fingerprint + batch inversion) // 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), - // but outer parallelism over 12 tables also helps on high-core-count machines. - #[cfg(feature = "instruments")] - let phase_start = Instant::now(); - #[cfg(feature = "instruments")] - let __sp = crate::instruments::span("r1_aux_build"); + // Aux build, aux commit and rounds 2-4 run FUSED per table below (one + // driver chains all three for its table, so tables never wait on a + // phase barrier); only this sequential prep runs here. // Disk-spill needs the aux columns in the host trace to spill them, so // disable the GPU-resident aux build (it would keep them device-only). @@ -3165,67 +3247,13 @@ pub trait IsStarkProver< } } - #[cfg(feature = "parallel")] - let aux_iter = air_trace_pairs.par_iter_mut(); - #[cfg(not(feature = "parallel"))] - let aux_iter = air_trace_pairs.iter_mut(); - let bus_inputs_vec: Vec>> = aux_iter - .map(|(air, trace, _)| { - if air.has_aux_trace() { - air.build_auxiliary_trace(*trace, &lookup_challenges) - } else { - None - } - }) - .collect(); - - // The trace-domain snapshots retained by the R1 main LDE (both Arcs: - // trace.main_trace_dev and GpuLdeBase.trace_dev) have exactly one - // consumer — the aux build above. Drop them now so the main-trace-sized - // device buffers are reclaimed before the aux-commit + DEEP/FRI VRAM - // peak instead of living to the end of the proof. - #[cfg(feature = "cuda")] - { - for (_, trace, _) in air_trace_pairs.iter_mut() { - trace.clear_main_trace_dev(); - } - for handle in main_gpu_handles.iter_mut().flatten() { - handle.trace_dev = None; - handle.trace_rows = 0; - } - } - - // Spill all aux trace tables to mmap before any Round 1 aux LDE work. - #[cfg(feature = "disk-spill")] - if storage_mode == StorageMode::Disk { - crate::par::par_try_for_each_mut(&mut air_trace_pairs, |(air, trace, _)| { - if air.has_aux_trace() { - trace - .spill_aux_to_disk() - .map_err(|e| ProvingError::DiskSpill(format!("aux trace: {e}")))?; - } - Ok::<(), ProvingError>(()) - })?; - } - - #[cfg(feature = "instruments")] - drop(__sp); - #[cfg(feature = "instruments")] - let aux_build_elapsed = phase_start.elapsed(); - #[cfg(feature = "instruments")] - if let Some(s) = crate::instruments::snap("After aux build") { - heap_snaps.push(s); - } - - // Pass 2: Parallel fork transcript → extract → LDE → commit in chunks of K. - // Each table gets its own transcript fork. + // The per-table aux build (inside the fused chain) reports through the + // per-table spans; the phase-level buckets are folded into rounds_2_4. #[cfg(feature = "instruments")] - let phase_start = Instant::now(); - #[cfg(feature = "instruments")] - let __sp = crate::instruments::span("r1_aux_commit"); + let aux_build_elapsed = Duration::ZERO; // Pre-fork all transcripts (cheap, sequential — must match verifier ordering) - let mut table_transcripts: Vec<_> = (0..num_airs) + let table_transcripts: Vec<_> = (0..num_airs) .map(|idx| { let mut t = transcript.clone(); if num_airs > 1 { @@ -3247,40 +3275,100 @@ pub trait IsStarkProver< ); #[cfg(not(feature = "cuda"))] type AuxResult = (Option>, (Vec>, usize)); - #[allow(clippy::type_complexity)] - let mut aux_results: Vec> = Vec::with_capacity(num_airs); - // R1 aux commit and rounds 2 to 4 share the peak working set: the main // and aux LDEs are co-resident, plus the composition and Merkle - // transients (in the scratch factor). `num_aux_columns` is populated by - // the aux build above, so this estimate is accurate for both phases. - let peak_chunks = { - let estimates: Vec = air_trace_pairs - .iter() - .enumerate() - .map(|(idx, (_, trace, _))| { - let lde_size = - domains[idx].interpolation_domain_size * domains[idx].blowup_factor; - estimate_table_vram_bytes( - trace.num_main_columns, - trace.num_aux_columns, - lde_size, - ) - }) + // transients (in the scratch factor). The aux width comes from the AIR + // layout (the aux build itself runs inside the admitted chain below). + let peak_estimates: Vec = air_trace_pairs + .iter() + .enumerate() + .map(|(idx, (air, trace, _))| { + let lde_size = domains[idx].interpolation_domain_size * domains[idx].blowup_factor; + let (_, aux_cols) = air.trace_layout(); + estimate_table_vram_bytes(trace.num_main_columns, aux_cols, lde_size) + }) + .collect(); + + // Per-table slots for the fused chain: each driver takes or locks only + // its own index, so every mutex is uncontended by construction. + let pair_cells: Vec>> = + air_trace_pairs + .into_iter() + .map(std::sync::Mutex::new) .collect(); - plan_table_chunks(&estimates, k, vram_budget) - }; + let main_commit_cells: Vec>>> = main_commits + .into_iter() + .map(|c| std::sync::Mutex::new(Some(c))) + .collect(); + #[allow(clippy::type_complexity)] + let main_lde_cells: Vec< + std::sync::Mutex>, usize)>>, + > = main_ldes + .into_iter() + .map(|l| std::sync::Mutex::new(Some(l))) + .collect(); + #[cfg(feature = "cuda")] + let gpu_main_cells: Vec>> = + main_gpu_handles + .into_iter() + .map(std::sync::Mutex::new) + .collect(); + let transcript_cells: Vec<_> = table_transcripts + .into_iter() + .map(std::sync::Mutex::new) + .collect(); + #[cfg(feature = "instruments")] + #[allow(clippy::type_complexity)] + let table_timings_mx: std::sync::Mutex< + Vec<(String, usize, Duration, crate::instruments::TableSubOps)>, + > = std::sync::Mutex::new(Vec::new()); - for &(chunk_start, chunk_end) in &peak_chunks { - let chunk_range = chunk_start..chunk_end; + // Fused chain, stage 1: aux build → aux commit → aux root into the + // table's transcript fork → Round1 assembly. + #[allow(clippy::type_complexity)] + let aux_stage = |idx: usize| -> Result< + ( + Round1Commitments, + Lde, + ), + ProvingError, + > { + let mut pair = pair_cells[idx].lock().unwrap(); + let (air, trace, _) = &mut *pair; + let domain = &domains[idx]; + let twiddles = &twiddle_caches[idx]; - #[allow(clippy::type_complexity)] - let chunk_aux: Vec, ProvingError>> = - crate::par::par_map_collect(chunk_range, |idx| { - let (air, trace, _) = &air_trace_pairs[idx]; - let domain = &domains[idx]; - let twiddles = &twiddle_caches[idx]; + #[cfg(feature = "instruments")] + let __sp = crate::instruments::span("r1_aux_build"); + let bus_public_inputs = if air.has_aux_trace() { + air.build_auxiliary_trace(*trace, &lookup_challenges) + } else { + None + }; + // The trace-domain snapshot retained by the R1 main LDE has exactly + // one consumer — the aux build above. Reclaim it before this + // table's aux-commit + DEEP/FRI VRAM peak. + #[cfg(feature = "cuda")] + { + trace.clear_main_trace_dev(); + if let Some(handle) = gpu_main_cells[idx].lock().unwrap().as_mut() { + handle.trace_dev = None; + handle.trace_rows = 0; + } + } + #[cfg(feature = "disk-spill")] + if storage_mode == StorageMode::Disk && air.has_aux_trace() { + trace + .spill_aux_to_disk() + .map_err(|e| ProvingError::DiskSpill(format!("aux trace: {e}")))?; + } + #[cfg(feature = "instruments")] + drop(__sp); + #[cfg(feature = "instruments")] + let __sp = crate::instruments::span("r1_aux_commit"); + let aux_full: AuxResult = + (|| -> Result, ProvingError> { if air.has_aux_trace() { let lde_size = domain.interpolation_domain_size * domain.blowup_factor; @@ -3418,179 +3506,155 @@ pub trait IsStarkProver< #[cfg(not(feature = "cuda"))] Ok((None, (Vec::new(), 0))) } - }); - - // Sequential: append aux roots to forked transcripts. - for (j, result) in chunk_aux.into_iter().enumerate() { - let aux_full = result?; - // Tuple shape is cfg-gated; `.0` is the optional TableCommit - // in both variants. - if let Some(ref c) = aux_full.0 { - table_transcripts[chunk_start + j].append_bytes(&c.root); - } - aux_results.push(aux_full); + })()?; + // Tuple shape is cfg-gated; `.0` is the optional TableCommit in + // both variants. Aux roots go to the table's OWN fork, so no + // cross-table ordering is needed here. + if let Some(ref c) = aux_full.0 { + transcript_cells[idx].lock().unwrap().append_bytes(&c.root); } - } - - // 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); - let mut cached_ldes: Vec> = Vec::with_capacity(num_airs); - // Under cuda, fold main_gpu_handles into the zip chain so each handle - // stays paired with its table by construction. - #[cfg(feature = "cuda")] - let main_iter = main_commits - .into_iter() - .zip(main_ldes) - .zip(main_gpu_handles); - #[cfg(not(feature = "cuda"))] - let main_iter = main_commits.into_iter().zip(main_ldes); + #[cfg(feature = "instruments")] + drop(__sp); - for ((main_pack, aux_full), bus_public_inputs) in - main_iter.zip(aux_results).zip(bus_inputs_vec) - { - #[cfg(feature = "cuda")] - let ((main_commit, main_lde), gpu_main) = main_pack; - #[cfg(not(feature = "cuda"))] - let (main_commit, main_lde) = main_pack; #[cfg(feature = "cuda")] let (aux_commit, cached_aux, gpu_aux) = aux_full; #[cfg(not(feature = "cuda"))] let (aux_commit, cached_aux) = aux_full; - commitments.push(Round1Commitments { + let main_commit = main_commit_cells[idx] + .lock() + .unwrap() + .take() + .expect("main commit consumed once per table"); + let main_lde = main_lde_cells[idx] + .lock() + .unwrap() + .take() + .expect("main lde consumed once per table"); + #[cfg(feature = "cuda")] + let gpu_main = gpu_main_cells[idx].lock().unwrap().take(); + let commitment = Round1Commitments { main: main_commit, aux: aux_commit, rap_challenges: lookup_challenges.clone(), bus_public_inputs, - }); + }; #[cfg(feature = "cuda")] - cached_ldes.push(Lde { + let lde = Lde { main: main_lde, aux: cached_aux, gpu_main, gpu_aux, - }); + }; #[cfg(not(feature = "cuda"))] - cached_ldes.push(Lde { + let lde = Lde { main: main_lde, aux: cached_aux, - }); - } + }; + Ok((commitment, lde)) + }; - #[cfg(feature = "instruments")] - drop(__sp); - #[cfg(feature = "instruments")] - let aux_commit_elapsed = phase_start.elapsed(); - #[cfg(feature = "instruments")] - if let Some(s) = crate::instruments::snap("After aux commit") { - heap_snaps.push(s); - } + // Fused chain, stage 2: Round1 from the cached LDE (consumed by value, + // no recomputation) → rounds 2-4 against the table's transcript fork. + let rounds_stage = |idx: usize, + commitment: Round1Commitments, + lde: Lde| + -> Result, ProvingError> { + let pair = pair_cells[idx].lock().unwrap(); + let (air, trace, pub_inputs) = &*pair; + let _ = trace; // used by instruments + let domain = &domains[idx]; - #[cfg(feature = "debug-checks")] - Self::run_debug_checks(&air_trace_pairs, &commitments, &domains, &twiddle_caches); + #[cfg(feature = "instruments")] + let __sp = crate::instruments::span("rounds_2to4"); + #[cfg(feature = "instruments")] + let table_start = Instant::now(); - // ===================================================================== - // Rounds 2-4: Parallel per-table proving in chunks of K - // ===================================================================== - // 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. + let mut round_1_result = + commitment.build_round1(lde, air.step_size(), domain.blowup_factor); + + let mut tguard = transcript_cells[idx].lock().unwrap(); + if let Some(ref bpi) = round_1_result.bus_public_inputs { + tguard.append_field_element(&bpi.table_contribution); + } + + let proof = Self::prove_rounds_2_to_4( + *air, + *pub_inputs, + &mut round_1_result, + &mut *tguard, + domain, + &twiddle_caches[idx], + )?; + + #[cfg(feature = "instruments")] + { + let sub_ops = crate::instruments::take_round_sub_ops().unwrap_or_default(); + table_timings_mx.lock().unwrap().push(( + air.name().to_string(), + trace.num_rows(), + table_start.elapsed(), + sub_ops, + )); + } + Ok(proof) + }; #[cfg(feature = "instruments")] let phase_start = Instant::now(); #[cfg(feature = "instruments")] - let __sp = crate::instruments::span("rounds_2to4"); - #[cfg(feature = "instruments")] - let mut table_timings: Vec<( - String, - usize, - Duration, - crate::instruments::TableSubOps, - )> = Vec::with_capacity(num_airs); + let aux_commit_elapsed = Duration::ZERO; + + let peak_order = heaviest_first(&peak_estimates); + + // One fused task per table: while a heavy table works through a + // host-bound stretch, the others' GPU stages fill the device. The + // shared transcript is untouched past this point (each fork is + // per-table), so any order is sound; proofs are drained in index order. + #[cfg(not(feature = "debug-checks"))] + let table_results = run_admitted(&peak_order, &peak_estimates, &vram_gate, k, |idx| { + let (commitment, lde) = aux_stage(idx)?; + rounds_stage(idx, commitment, lde) + }); - let mut proofs = Vec::with_capacity(num_airs); - let mut lde_drain = cached_ldes.into_iter(); - for &(chunk_start, chunk_end) in &peak_chunks { - let chunk_size = chunk_end - chunk_start; - - 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_ldes - .into_par_iter() - .zip(chunk_commitments.par_iter()) - .zip(chunk_transcripts.par_iter_mut()) - .enumerate(); - #[cfg(not(feature = "parallel"))] - let iter = chunk_ldes + // debug-checks needs every table's commitments and traces between the + // aux and rounds stages (cross-table bus balance), so it splits the + // fused chain into two admitted passes around the check. + #[cfg(feature = "debug-checks")] + let table_results = { + let aux_outs = run_admitted(&peak_order, &peak_estimates, &vram_gate, k, aux_stage); + let mut commitments = Vec::with_capacity(num_airs); + let mut ldes = Vec::with_capacity(num_airs); + for out in aux_outs { + let (c, l) = out.expect("run_admitted fills every slot")?; + commitments.push(c); + ldes.push(l); + } + Self::run_debug_checks(&pair_cells, &commitments, &domains, &twiddle_caches); + #[allow(clippy::type_complexity)] + let staged: Vec< + std::sync::Mutex< + Option<( + Round1Commitments, + Lde, + )>, + >, + > = commitments .into_iter() - .zip(chunk_commitments.iter()) - .zip(chunk_transcripts.iter_mut()) - .enumerate(); - - let chunk_results: Vec> = iter - .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 - let domain = &domains[idx]; - - #[cfg(feature = "instruments")] - let table_start = Instant::now(); - - // Build Round1 from cached LDE (consumed by value, no recomputation). - let mut round_1_result = - commitment.build_round1(lde, air.step_size(), domain.blowup_factor); - - if let Some(ref bpi) = round_1_result.bus_public_inputs { - table_transcript.append_field_element(&bpi.table_contribution); - } - - let proof = Self::prove_rounds_2_to_4( - *air, - *pub_inputs, - &mut round_1_result, - table_transcript, - domain, - &twiddle_caches[idx], - )?; - - #[cfg(feature = "instruments")] - let table_timing = { - let sub_ops = crate::instruments::take_round_sub_ops().unwrap_or_default(); - ( - air.name().to_string(), - trace.num_rows(), - table_start.elapsed(), - sub_ops, - ) - }; - - #[cfg(feature = "instruments")] - return Ok((proof, table_timing)); - #[cfg(not(feature = "instruments"))] - Ok(proof) - }) + .zip(ldes) + .map(|p| std::sync::Mutex::new(Some(p))) .collect(); + run_admitted(&peak_order, &peak_estimates, &vram_gate, k, |idx| { + let (c, l) = staged[idx].lock().unwrap().take().unwrap(); + rounds_stage(idx, c, l) + }) + }; - for result in chunk_results { - #[cfg(feature = "instruments")] - { - let (proof, timing) = result?; - proofs.push(proof); - table_timings.push(timing); - } - #[cfg(not(feature = "instruments"))] - proofs.push(result?); - } + let mut proofs = Vec::with_capacity(num_airs); + for result in table_results { + proofs.push(result.expect("run_admitted fills every slot")?); } - #[cfg(feature = "instruments")] - drop(__sp); + let table_timings = table_timings_mx.into_inner().unwrap(); #[cfg(feature = "instruments")] { // Store timing data for the top-level report in prove_with_options. From 8f91e7c70096b073a22b476c1fc7b5332e35e201 Mon Sep 17 00:00:00 2001 From: Mauro Toscano <12560266+MauroToscano@users.noreply.github.com> Date: Mon, 3 Aug 2026 17:42:07 -0300 Subject: [PATCH 10/12] fix(prover): repair the instruments span tree and timing report under the per-table scheduler (#893) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix(instruments): nest per-table spans under their real parent The per-table scheduler moved `r1_aux_build`, `r1_aux_commit` and `rounds_2to4` inside closures that run on `std::thread::scope` worker threads. `SPAN_DEPTH` is thread-local and a fresh OS thread starts at 0, so all three were stamped `depth = 0` and recorded as root siblings of `prove_total` instead of children of `proving`. That happens even at k = 1. Downstream, `scripts/profiling/phase_table.py` reconstructs the tree with `del stack[d:]`, so a depth-0 span empties the ancestor stack and `prove_total` stops being an ancestor of anything — the "% of total" column documented in `scripts/profiling/README.md` becomes meaningless. `run_admitted` now reads the spawning thread's depth and seeds each driver with it via new `instruments::current_depth` / `enter_depth`. Both call sites are `#[cfg(feature = "instruments")]`, so non-instrumented builds are byte-identical. Also correct the module contract doc: per-table spans genuinely do overlap now — that is inherent to running one driver per in-flight table, not a bug to code around. Only the top-level phase spans remain a strict latency breakdown. * fix(prover): report aux build/commit where they actually accrue `aux_build_elapsed` / `aux_commit_elapsed` were hardcoded to `Duration::ZERO`, but `prover/src/instruments.rs` still computed `round1 = main_commits + aux_build + aux_commit` and still printed the "Aux trace build" / "Aux trace commit" rows. Since `accum_r1_aux` keeps firing, the report showed nonzero LogUp and Aux-LDE/Merkle children under zero parents, and all the aux time silently landed in "Rounds 2-4". Time both stages inside the fused chain and sum them across drivers (`instruments::accum_aux_phases` / `take_aux_phases`), then restructure the report to match what the scheduler actually does: - "Round 1 (main trace commits)" is now exactly the main commits — the last phase-wide barrier, since the main roots must all be in the transcript before the shared LogUp challenges are sampled. - Aux build, aux commit and rounds 2-4 sit under one wall-clock parent, "Rounds 2-4 (aux build+commit fused)", with the aux rows marked as summed across concurrent drivers — they may exceed that wall, the same convention the existing accum_* sub-rows already use. No zero parents over nonzero children remain. Verified on fib_iterative_1M: Round 1 1.68s, Rounds 2-4 6.89s wall, aux build 2.91s and aux commit 3.38s summed over 5 drivers. * fix(bench): drop the heap guards whose snapshots no longer exist The scheduler removed `instruments::snap("After aux build")` and `snap("After aux commit")`. `bench_prover_scaling.sh` still parsed them, printed them and ran heap-growth regressions on them, so two regression guards were comparing nothing and dropping out without complaint. Re-adding a snapshot inside a per-table task would be meaningless — with k tables in flight there is no single moment at which aux build or aux commit has finished — so remove the two rows and their `regress` calls, with a NOTE recording why and pointing at the guards that still cover the fused region ("After main commits" and "Peak heap"). Also repoint the timing regexes at the labels the report actually prints. `Main expand_columns_to_lde` / `Aux expand_columns_to_lde` and `Main commit (Merkle)` / `Aux commit (Merkle)` had not matched since the labels gained their GPU/CPU suffixes, so t_main_lde, t_aux_lde, t_main_merkle and t_aux_merkle silently printed blank. All four populate again — checked by running the script's own awk over a real report. * docs(prover): refresh the comments the per-table scheduler invalidated Nothing functional. All of these described structure the scheduler removed: - `VramGate`'s rustdoc opened with the deleted `plan_table_chunks`'s doc comment ("Plan contiguous table chunks... Returns (start, end) half open ranges"), left behind and contiguous with `VramGate`'s own. - `Lde`'s doc claimed all N tables' LDE columns are live simultaneously. Only the main LDEs still are — the Round 1 main commit is a phase-wide barrier. Each aux LDE is produced and consumed inside one fused task, so at most k coexist. That is a memory improvement the PR made and did not claim; state the real, asymmetric bound. - A "Split into two passes for parallelism: Pass 1 ... Pass 2 ..." block sat two lines above the new comment saying the opposite. - `table_parallelism`'s doc still gave only `num_cores / 3`. Document both arms, that `TABLE_PARALLELISM` overrides both, and that without the `parallel` feature it is hardcoded to 1. - `run_debug_checks` said "called once after Phase C commits"; it now runs between two `run_admitted` passes. Document that, and the "each driver locks only its own index" contract its new `&[Mutex]` parameter relies on. - `auto_storage::peak_bytes` described phase D and a "worst possible chunk assignment". With `heaviest_first` the top-k is the set actually admitted first, not a worst case. Also document that `table_parallelism()` is not only the prover's k: `decide` feeds it into the RAM-vs-Disk choice, so the cuda arm's `cores * 2 / 3` doubles that transient term and makes `Disk` likelier. The direction is safe (it over-estimates) but was undocumented. - Remaining "Phase A/B/D" references, plus the "chunks of K" banner and the "Phase D's zip chain" handle comments. * test(prover): cover VramGate, run_admitted and heaviest_first These three had zero direct tests, and PR CI never exercises them concurrently: `ubuntu-latest` has 2-4 vCPU so `cores / 3` floors to k = 1, and `VramGate` is inert on non-cuda builds because `vram_budget = u64::MAX` makes `acquire`'s admit condition always true, so the condvar is never waited on. They are free functions over `&[u64]` with no field, AIR or GPU dependency, so a plain `#[cfg(test)] mod` pins them without a device: - `heaviest_first` returns a permutation of `0..n`, descending by estimate, with ties broken by index (stable sort — so the admission order does not vary run to run). - `run_admitted` fills every slot exactly once, including `order.len() == 0`, `workers > order.len()`, `workers == 1` and `workers == 0`. - `VramGate` admits an over-budget request alone rather than deadlocking, never lets concurrent admissions push `used` past the budget, and wakes waiters on permit drop. - A `u64::MAX` budget never blocks, including when the byte sum saturates. - `run_admitted` seeds its drivers' span depth, which guards the regression fixed earlier in this branch. Reads the depth directly rather than the global span timeline, which other tests in this binary also write to. Deterministic and fast: no sleeps as synchronization: channel rendezvous for ordering, and `recv_timeout` only as a failure deadline so a regression fails instead of hanging. --- crypto/stark/src/instruments.rs | 74 +++++++- crypto/stark/src/prover.rs | 309 +++++++++++++++++++++++++++----- prover/src/auto_storage.rs | 21 ++- prover/src/instruments.rs | 27 +-- scripts/bench_prover_scaling.sh | 34 ++-- 5 files changed, 389 insertions(+), 76 deletions(-) diff --git a/crypto/stark/src/instruments.rs b/crypto/stark/src/instruments.rs index 21866c465..790c06cb0 100644 --- a/crypto/stark/src/instruments.rs +++ b/crypto/stark/src/instruments.rs @@ -6,12 +6,20 @@ use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH}; // Wall clock span timeline: the trustworthy per step latency breakdown. // -// Spans open and close on the main thread at phase boundaries. They do not -// overlap and sum to their parent, so the tree is a true latency breakdown -// (unlike the accum_* thread local sub timers below, which sum per worker CPU -// time across rayon threads and can exceed 100%). A parallel region is one span -// around the blocking call; its internal split is reported separately as CPU -// time, never mixed into the wall tree. +// Top level phase spans open and close on the main thread at phase boundaries. +// They do not overlap and sum to their parent, so that part of the tree is a +// true latency breakdown (unlike the accum_* thread local sub timers below, +// which sum per worker CPU time across rayon threads and can exceed 100%). A +// parallel region is one span around the blocking call; its internal split is +// reported separately as CPU time, never mixed into the wall tree. +// +// Per table spans are the exception. `multi_prove` runs one driver thread per +// in flight table (`r1_aux_build`, `r1_aux_commit`, `rounds_2to4`), so up to +// `table_parallelism()` of them are open at once: they DO overlap in wall time +// and their sum exceeds their parent. They still nest correctly, because each +// driver seeds its span depth from the spawning thread (`current_depth` / +// `enter_depth`) instead of starting a fresh thread at depth 0 — but read them +// as per table wall time, not as a share of the enclosing phase. // // let _s = instruments::span("trace_build"); // RAII, stops on drop // @@ -36,6 +44,32 @@ thread_local! { static SPAN_DEPTH: std::cell::Cell = const { std::cell::Cell::new(0) }; } +/// Depth the next span opened on this thread would be stamped with. +/// +/// Read this on the thread that spawns workers and hand the value to +/// [`enter_depth`] inside each worker: a freshly spawned OS thread starts at +/// depth 0, so without it every span opened off the main thread is recorded as +/// a root sibling of `prove_total` and the tree stops being a breakdown. +pub fn current_depth() -> u16 { + SPAN_DEPTH.with(|d| d.get()) +} + +/// Restores the span depth this thread had before [`enter_depth`]. +#[must_use] +pub struct DepthGuard(u16); + +/// Seed this thread's span depth from a parent thread, restoring the previous +/// value when the returned guard drops. +pub fn enter_depth(depth: u16) -> DepthGuard { + DepthGuard(SPAN_DEPTH.with(|d| d.replace(depth))) +} + +impl Drop for DepthGuard { + fn drop(&mut self) { + SPAN_DEPTH.with(|d| d.set(self.0)); + } +} + #[must_use] pub struct SpanGuard { label: &'static str, @@ -263,8 +297,15 @@ pub struct Round1SubOps { pub struct MultiProveTiming { pub prepass: Duration, pub main_commits: Duration, + /// Aux build wall time summed over the concurrent per-table drivers. It is + /// no longer a phase of its own — the fused chain runs it inside + /// `rounds_2_4`, so it overlaps itself and is a subset of that wall time, + /// not an addend. pub aux_build: Duration, + /// Aux commit, same accounting as `aux_build`. pub aux_commit: Duration, + /// Wall clock of the fused per-table region: aux build + aux commit + + /// rounds 2-4, all of it concurrent across `table_parallelism()` drivers. pub rounds_2_4: Duration, /// Sub-op breakdown for Round 1 (main + aux LDE vs Merkle). pub round1_sub: Round1SubOps, @@ -278,6 +319,11 @@ static R1_MAIN_LDE_US: AtomicU64 = AtomicU64::new(0); static R1_MAIN_MERKLE_US: AtomicU64 = AtomicU64::new(0); static R1_AUX_LDE_US: AtomicU64 = AtomicU64::new(0); static R1_AUX_MERKLE_US: AtomicU64 = AtomicU64::new(0); +// Aux build / aux commit wall time per table, summed across the concurrent +// per-table drivers of the fused chain (so, like the sub-timers, this is a +// CPU-style total that can exceed the fused region's wall clock). +static AUX_BUILD_US: AtomicU64 = AtomicU64::new(0); +static AUX_COMMIT_US: AtomicU64 = AtomicU64::new(0); // Aux build (LogUp) sub-phases, CPU time accumulated across tables/chunks. static AUX_FINGERPRINT_US: AtomicU64 = AtomicU64::new(0); static AUX_INVERT_US: AtomicU64 = AtomicU64::new(0); @@ -326,6 +372,20 @@ pub fn accum_aux_accumulate(d: Duration) { AUX_ACCUM_US.fetch_add(d.as_micros() as u64, Ordering::Relaxed); } +/// One table's aux build and aux commit wall time, from its fused-chain driver. +pub fn accum_aux_phases(build: Duration, commit: Duration) { + AUX_BUILD_US.fetch_add(build.as_micros() as u64, Ordering::Relaxed); + AUX_COMMIT_US.fetch_add(commit.as_micros() as u64, Ordering::Relaxed); +} + +/// Drain the summed per-table aux build / aux commit times. +pub fn take_aux_phases() -> (Duration, Duration) { + ( + Duration::from_micros(AUX_BUILD_US.swap(0, Ordering::Relaxed)), + Duration::from_micros(AUX_COMMIT_US.swap(0, Ordering::Relaxed)), + ) +} + pub fn take_r1_sub() -> Round1SubOps { Round1SubOps { main_lde: Duration::from_micros(R1_MAIN_LDE_US.swap(0, Ordering::Relaxed)), @@ -352,6 +412,8 @@ pub fn reset_all() { R1_MAIN_MERKLE_US.store(0, Ordering::Relaxed); R1_AUX_LDE_US.store(0, Ordering::Relaxed); R1_AUX_MERKLE_US.store(0, Ordering::Relaxed); + AUX_BUILD_US.store(0, Ordering::Relaxed); + AUX_COMMIT_US.store(0, Ordering::Relaxed); AUX_FINGERPRINT_US.store(0, Ordering::Relaxed); AUX_INVERT_US.store(0, Ordering::Relaxed); AUX_TERM_US.store(0, Ordering::Relaxed); diff --git a/crypto/stark/src/prover.rs b/crypto/stark/src/prover.rs index 7d979bcbc..4199d0a80 100644 --- a/crypto/stark/src/prover.rs +++ b/crypto/stark/src/prover.rs @@ -246,7 +246,7 @@ type MainCommitTuple = ( type MainCommitTuple = (TableCommit, (Vec>, usize)); /// Round 1 commitment artifacts — Merkle trees, roots, challenges, and bus inputs. -/// Borrowed (not consumed) when building `Round1` in Phase D. +/// Borrowed (not consumed) when building `Round1`. pub(crate) struct Round1Commitments where Field: IsFFTField + IsSubFieldOf, @@ -260,10 +260,18 @@ where bus_public_inputs: Option>, } -/// LDE columns for main (Phase A) and auxiliary (Phase C) traces, consumed by value in Phase D. +/// Main and auxiliary LDE columns, consumed by value when the table's `Round1` +/// is assembled. /// -/// Memory trade-off: all N tables' LDE columns are live simultaneously between Phase A/C -/// and Phase D (O(N × cols × lde_size)). +/// Memory trade-off, asymmetric since the per-table scheduler fused aux build, +/// aux commit and rounds 2-4 into one task: +/// - main: produced by the Round 1 main commit, which is a phase-wide barrier, +/// so all N tables' main LDEs are live at once (O(N × main_cols × lde_size)). +/// - aux: produced and consumed inside the same fused task, so at most +/// `table_parallelism()` of them coexist (O(k × aux_cols × lde_size)). +/// +/// Under `debug-checks` the fused task is split around the cross-table bus +/// balance check, so there the aux LDEs are all-N-live like the main ones. struct Lde { /// Row-major main LDE buffer + its column count. main: (Vec>, usize), @@ -566,8 +574,17 @@ where } /// Number of tables to process concurrently in `multi_prove`. -/// Default: num_cores / 3 (benchmarked optimal on both M3 Pro and EPYC 9454P). -/// Override with `TABLE_PARALLELISM` env var. +/// +/// Defaults: `num_cores / 3` on CPU builds (benchmarked optimal on both M3 Pro +/// and EPYC 9454P — every table there is pure host work), `num_cores * 2 / 3` +/// under `cuda`, where most in-flight tables sit in GPU waits so more of them +/// pay (swept flat at ~2/3 of the cores on a 16-core/RTX 5090 box). Both arms +/// are overridden by the `TABLE_PARALLELISM` env var. Without the `parallel` +/// feature this is hardcoded to 1 and the env var is ignored. +/// +/// Not only the prover's `k`: `auto_storage::decide` feeds this into the +/// RAM-vs-Disk storage estimate, so the `cuda` arm also doubles that transient +/// term (see `peak_bytes`). pub fn table_parallelism() -> usize { #[cfg(feature = "parallel")] { @@ -615,12 +632,6 @@ fn estimate_table_vram_bytes(main_cols: usize, aux_cols: usize, lde_size: usize) lde_term.saturating_add(tree_term) } -/// Plan contiguous table chunks for parallel proving. A chunk grows until it -/// hits `k` tables or its summed VRAM estimate would exceed `budget`; a single -/// table larger than `budget` runs solo. With `budget == u64::MAX` (non-cuda, -/// or VRAM not binding) chunks fall back to fixed size `k`, identical to the -/// old `step_by(k)`, so scheduling and the proof are unchanged. Returns -/// `(start, end)` half open ranges covering `0..estimates.len()` in order. /// Byte-budget admission gate for concurrently proven tables. `acquire` /// blocks until the requested bytes fit under the budget, releasing on /// permit drop. An oversized request is admitted alone (when nothing else @@ -687,9 +698,17 @@ fn run_admitted( .map(|_| std::sync::Mutex::new(None)) .collect(); let cursor = std::sync::atomic::AtomicUsize::new(0); + // Spans opened inside `task` run on these worker threads, and a fresh OS + // thread's span depth starts at 0 — which would record every per-table + // span as a root instead of a child of the phase that spawned it. Carry + // the spawning thread's depth across the scope boundary. + #[cfg(feature = "instruments")] + let parent_depth = crate::instruments::current_depth(); std::thread::scope(|scope| { for _ in 0..workers.max(1).min(order.len().max(1)) { scope.spawn(|| { + #[cfg(feature = "instruments")] + let _depth = crate::instruments::enter_depth(parent_depth); loop { let pos = cursor.fetch_add(1, std::sync::atomic::Ordering::Relaxed); if pos >= order.len() { @@ -1035,9 +1054,9 @@ pub trait IsStarkProver< } /// Compute the main-trace LDE and commit. Returns a `TableCommit` along - /// with the owned LDE columns (consumed later in Phase D) and (under - /// cuda) the optional device LDE buffer kept alive for downstream rounds - /// when the R1 fused GPU pipeline ran. + /// with the owned LDE columns (consumed later by the table's fused task) + /// and (under cuda) the optional device LDE buffer kept alive for + /// downstream rounds when the R1 fused GPU pipeline ran. /// /// `precomputed`: if present, the leading `num_cols` columns are committed /// as a separate Merkle tree (the precomputed split for preprocessed @@ -1321,8 +1340,8 @@ pub trait IsStarkProver< /// Recompute Round1 from the trace, reusing the Merkle trees stored in commitments. /// - /// Only used by `run_debug_checks` — Phase D consumes the cached LDE - /// directly and does not go through this path. + /// Only used by `run_debug_checks` — the production path consumes the + /// cached LDE directly and does not go through here. #[cfg(feature = "debug-checks")] fn reconstruct_round1( air: &dyn AIR, @@ -1398,7 +1417,15 @@ pub trait IsStarkProver< } /// Reconstruct Round1 for every table, print the bus balance report, and - /// validate each trace. Called once after Phase C commits. + /// validate each trace. + /// + /// Cross-table (bus balance) checks need every table's commitments at once, + /// so under `debug-checks` the fused per-table chain is split into two + /// admitted passes and this runs once, on the main thread, between them. + /// + /// `pair_cells` is the same per-table slot vector the drivers use. Each + /// driver only ever locks its own index, so locking them here — after the + /// aux pass has joined and before the rounds pass starts — is uncontended. #[cfg(feature = "debug-checks")] fn run_debug_checks( pair_cells: &[std::sync::Mutex>], @@ -3068,7 +3095,7 @@ pub trait IsStarkProver< // of the tables proved concurrently so large blocks don't exhaust VRAM. // It is an extra ceiling on top of `k` (it never raises concurrency). On // non-cuda builds, or when the budget can't be queried, it is `u64::MAX` - // and chunking falls back to fixed size `k`. + // and the gate is inert — concurrency is then bounded by `k` alone. #[cfg(feature = "cuda")] let vram_budget = math_cuda::device::backend() .map(|b| b.vram_budget_bytes()) @@ -3118,7 +3145,7 @@ pub trait IsStarkProver< } // ===================================================================== - // Round 1, Phase A: Commit all main traces (parallel in chunks of K) + // Round 1: Commit all main traces (VRAM-admitted, up to K concurrent) // ===================================================================== // All main trace commitments must be in the transcript before sampling // LogUp challenges. @@ -3131,8 +3158,9 @@ pub trait IsStarkProver< let mut main_commits: Vec> = Vec::with_capacity(num_airs); let mut main_ldes: Vec<(Vec>, usize)> = Vec::with_capacity(num_airs); // Optional device-side LDE handle per table, populated only when the - // R1 fused GPU pipeline produced one. Threaded through Phase D's zip - // chain so each handle stays paired with its table by construction. + // R1 fused GPU pipeline produced one. Indexed by table, and moved into + // the per-table `gpu_main_cells` slots below so each handle stays + // paired with its table across the fused chain. #[cfg(feature = "cuda")] let mut main_gpu_handles: Vec> = Vec::with_capacity(num_airs); @@ -3198,7 +3226,7 @@ pub trait IsStarkProver< } // ===================================================================== - // Round 1, Phase B: Sample shared LogUp challenges + // Round 1: Sample shared LogUp challenges // ===================================================================== let lookup_challenges: Vec> = if needs_lookup_challenges { @@ -3210,16 +3238,13 @@ pub trait IsStarkProver< }; // ===================================================================== - // Phase C + Rounds 2-4: Forked per table + // Aux build + aux commit + Rounds 2-4: fused per table // ===================================================================== // Each table gets an independent transcript fork (cloned from the shared - // state after Phase B, domain-separated by table index). This matches - // the verifier's forking and makes per-table proving independent. + // state after the LogUp challenges, domain-separated by table index). + // This matches the verifier's forking and makes per-table proving + // independent. // - // Split into two passes for parallelism: - // Pass 1 (parallel): Build all auxiliary traces (fingerprint + batch inversion) - // Pass 2 (parallel): Fork transcript → extract → LDE → commit - // Aux build, aux commit and rounds 2-4 run FUSED per table below (one // driver chains all three for its table, so tables never wait on a // phase barrier); only this sequential prep runs here. @@ -3247,11 +3272,6 @@ pub trait IsStarkProver< } } - // The per-table aux build (inside the fused chain) reports through the - // per-table spans; the phase-level buckets are folded into rounds_2_4. - #[cfg(feature = "instruments")] - let aux_build_elapsed = Duration::ZERO; - // Pre-fork all transcripts (cheap, sequential — must match verifier ordering) let table_transcripts: Vec<_> = (0..num_airs) .map(|idx| { @@ -3263,10 +3283,10 @@ pub trait IsStarkProver< }) .collect(); - // Parallel aux commit in chunks of K. The closure returns a cfg-gated - // AuxResult. Under cuda it carries the optional ext3 GPU LDE handle as - // a third element, so Phase D's zip chain keeps it paired with its - // table without a separate handle vector. + // The aux stage of the fused chain returns a cfg-gated AuxResult. Under + // cuda it carries the optional ext3 GPU LDE handle as a third element, + // so the handle stays inside its own table's task and never needs a + // separate handle vector. #[cfg(feature = "cuda")] type AuxResult = ( Option>, @@ -3340,6 +3360,8 @@ pub trait IsStarkProver< #[cfg(feature = "instruments")] let __sp = crate::instruments::span("r1_aux_build"); + #[cfg(feature = "instruments")] + let t_aux_build = Instant::now(); let bus_public_inputs = if air.has_aux_trace() { air.build_auxiliary_trace(*trace, &lookup_challenges) } else { @@ -3363,16 +3385,20 @@ pub trait IsStarkProver< .map_err(|e| ProvingError::DiskSpill(format!("aux trace: {e}")))?; } #[cfg(feature = "instruments")] + let aux_build_dur = t_aux_build.elapsed(); + #[cfg(feature = "instruments")] drop(__sp); #[cfg(feature = "instruments")] let __sp = crate::instruments::span("r1_aux_commit"); + #[cfg(feature = "instruments")] + let t_aux_commit = Instant::now(); let aux_full: AuxResult = (|| -> Result, ProvingError> { if air.has_aux_trace() { let lde_size = domain.interpolation_domain_size * domain.blowup_factor; - // Same gate as the main commit (Phase A): skip the aux + // Same gate as the Round 1 main commit: skip the aux // host D2H when device-only, so both buffers are left // empty together for this table. #[cfg(feature = "cuda")] @@ -3514,6 +3540,8 @@ pub trait IsStarkProver< transcript_cells[idx].lock().unwrap().append_bytes(&c.root); } #[cfg(feature = "instruments")] + crate::instruments::accum_aux_phases(aux_build_dur, t_aux_commit.elapsed()); + #[cfg(feature = "instruments")] drop(__sp); #[cfg(feature = "cuda")] @@ -3601,8 +3629,6 @@ pub trait IsStarkProver< #[cfg(feature = "instruments")] let phase_start = Instant::now(); - #[cfg(feature = "instruments")] - let aux_commit_elapsed = Duration::ZERO; let peak_order = heaviest_first(&peak_estimates); @@ -3657,6 +3683,10 @@ pub trait IsStarkProver< let table_timings = table_timings_mx.into_inner().unwrap(); #[cfg(feature = "instruments")] { + // Aux build/commit are no longer phases of their own: each table's + // driver runs them inside the fused region, so these are sums over + // concurrent drivers and are a subset of `rounds_2_4`, not addends. + let (aux_build_elapsed, aux_commit_elapsed) = crate::instruments::take_aux_phases(); // Store timing data for the top-level report in prove_with_options. // Uses a thread-local to avoid changing multi_prove's return type. crate::instruments::store(crate::instruments::MultiProveTiming { @@ -3982,3 +4012,198 @@ fn print_bus_balance_report( } } } + +/// Scheduling primitives only. These are free functions over `&[u64]` with no +/// AIR, field or device dependency, so they are testable without a GPU — which +/// matters because CI never exercises them concurrently: `ubuntu-latest` has +/// 2-4 vCPU, so `table_parallelism()` floors to 1, and on non-cuda builds the +/// budget is `u64::MAX`, which makes `VramGate` inert. +#[cfg(test)] +mod scheduler_tests { + use super::{VramGate, heaviest_first, run_admitted}; + use std::sync::atomic::{AtomicUsize, Ordering}; + use std::time::Duration; + + /// Bound on how long a correct implementation may take to wake a waiter. + /// Only a failure deadline — never used to sequence the test. + const WAKE_DEADLINE: Duration = Duration::from_secs(10); + + #[test] + fn heaviest_first_is_a_descending_permutation() { + let empty: [u64; 0] = []; + assert!(heaviest_first(&empty).is_empty()); + assert_eq!(heaviest_first(&[3, 1, 2]), vec![0, 2, 1]); + + let estimates = [7u64, 0, 7, 3, 100, 1]; + let order = heaviest_first(&estimates); + let mut seen = order.clone(); + seen.sort_unstable(); + assert_eq!( + seen, + (0..estimates.len()).collect::>(), + "must be a permutation of 0..n" + ); + for w in order.windows(2) { + assert!( + estimates[w[0]] >= estimates[w[1]], + "must be descending by estimate, got {order:?}" + ); + } + } + + #[test] + fn heaviest_first_breaks_ties_by_index() { + // `sort_by_key` is stable, so equal estimates keep ascending index + // order: the admission order is a pure function of the estimates, and + // does not vary run to run. + assert_eq!(heaviest_first(&[5, 5, 5]), vec![0, 1, 2]); + assert_eq!(heaviest_first(&[1, 5, 5, 1]), vec![1, 2, 0, 3]); + } + + #[test] + fn run_admitted_fills_every_slot_exactly_once() { + // Includes the degenerate shapes: no work, one worker, more workers + // than tables, and `workers == 0` (clamped to 1 inside). + for (n, workers) in [(0, 4), (1, 1), (1, 8), (5, 0), (5, 1), (5, 8), (9, 4)] { + let estimates: Vec = (0..n).map(|i| i as u64 + 1).collect(); + let runs: Vec = (0..n).map(|_| AtomicUsize::new(0)).collect(); + let gate = VramGate::new(u64::MAX); + let out = run_admitted( + &heaviest_first(&estimates), + &estimates, + &gate, + workers, + |idx| { + runs[idx].fetch_add(1, Ordering::SeqCst); + idx * 10 + }, + ); + assert_eq!(out.len(), n, "one slot per table (n={n})"); + for i in 0..n { + assert_eq!( + runs[i].load(Ordering::SeqCst), + 1, + "table {i} ran exactly once (n={n}, workers={workers})" + ); + assert_eq!( + out[i], + Some(i * 10), + "slot {i} holds its own result (n={n}, workers={workers})" + ); + } + assert_eq!(*gate.used.lock().unwrap(), 0, "all permits released"); + } + } + + #[test] + fn concurrent_admissions_stay_under_budget() { + const BUDGET: u64 = 100; + const EACH: u64 = 40; // 2 fit, 3 do not + let estimates = vec![EACH; 16]; + let gate = VramGate::new(BUDGET); + let in_flight = AtomicUsize::new(0); + let max_in_flight = AtomicUsize::new(0); + run_admitted(&heaviest_first(&estimates), &estimates, &gate, 8, |_| { + let now = in_flight.fetch_add(1, Ordering::SeqCst) + 1; + max_in_flight.fetch_max(now, Ordering::SeqCst); + // Read under a held permit: the gate's own counter must never + // exceed the budget while anything is admitted. + assert!( + *gate.used.lock().unwrap() <= BUDGET, + "admitted bytes exceeded the budget" + ); + in_flight.fetch_sub(1, Ordering::SeqCst); + }); + assert!( + max_in_flight.load(Ordering::SeqCst) <= (BUDGET / EACH) as usize, + "more tables were in flight than the budget allows" + ); + assert_eq!(*gate.used.lock().unwrap(), 0); + } + + #[test] + fn oversized_request_is_admitted_alone_and_wakes_waiters() { + // A table bigger than the whole budget must still prove: it is admitted + // when nothing else holds bytes, rather than deadlocking forever. + let gate = VramGate::new(10); + let big = gate.acquire(100); + assert_eq!( + *gate.used.lock().unwrap(), + 100, + "oversized request admitted alone" + ); + + let (ready_tx, ready_rx) = std::sync::mpsc::channel(); + let (done_tx, done_rx) = std::sync::mpsc::channel(); + std::thread::scope(|s| { + s.spawn(|| { + ready_tx.send(()).unwrap(); + let permit = gate.acquire(5); + done_tx.send(()).unwrap(); + drop(permit); + }); + ready_rx.recv().unwrap(); + drop(big); + assert!( + done_rx.recv_timeout(WAKE_DEADLINE).is_ok(), + "dropping a permit must wake a waiter" + ); + }); + assert_eq!( + *gate.used.lock().unwrap(), + 0, + "permits release their bytes on drop" + ); + } + + /// Guards the instruments span tree: a driver thread starts at span depth + /// 0, so without the seeding in `run_admitted` every per-table span is + /// recorded as a root sibling of `prove_total` and the "% of total" column + /// stops meaning anything. Reads the depth directly instead of the global + /// span timeline, which other tests in this binary also write to. + #[cfg(feature = "instruments")] + #[test] + fn run_admitted_seeds_worker_span_depth() { + const PARENT_DEPTH: u16 = 3; + let outer = crate::instruments::enter_depth(PARENT_DEPTH); + let n = 6; + let estimates = vec![1u64; n]; + let gate = VramGate::new(u64::MAX); + let seen = run_admitted(&heaviest_first(&estimates), &estimates, &gate, 4, |_| { + crate::instruments::current_depth() + }); + for (i, depth) in seen.iter().enumerate() { + assert_eq!( + *depth, + Some(PARENT_DEPTH), + "table {i}'s driver must inherit the spawning thread's span depth" + ); + } + drop(outer); + assert_eq!( + crate::instruments::current_depth(), + 0, + "the calling thread's depth is restored" + ); + } + + #[test] + fn max_budget_gate_never_blocks() { + // The non-cuda / unqueryable-VRAM configuration. Two acquires that + // saturate `u64` must both be admitted. + let gate = VramGate::new(u64::MAX); + let held = gate.acquire(u64::MAX / 2); + let (tx, rx) = std::sync::mpsc::channel(); + std::thread::scope(|s| { + s.spawn(|| { + let _permit = gate.acquire(u64::MAX / 2 + 1000); + tx.send(()).unwrap(); + }); + assert!( + rx.recv_timeout(WAKE_DEADLINE).is_ok(), + "a u64::MAX budget must never block an acquire" + ); + }); + drop(held); + } +} diff --git a/prover/src/auto_storage.rs b/prover/src/auto_storage.rs index 49707cb4c..6b5ed8a5d 100644 --- a/prover/src/auto_storage.rs +++ b/prover/src/auto_storage.rs @@ -48,7 +48,8 @@ pub const SAFETY_FRACTION_DEN: u64 = 10; /// `(rows, main_cols, aux_cols, num_main_merkle_trees)` for a single table. type TableSpec = (u64, u64, u64, u64); -/// Bytes alive for the duration of phase D (LDE columns + main/aux Merkle). +/// Bytes counted as alive for the whole proof (LDE columns + main/aux Merkle). +/// Deliberately an over-estimate for the aux half — see `peak_bytes`. fn persistent_per_table(spec: TableSpec, blowup: u64) -> u64 { let (rows, main_cols, aux_cols, main_trees) = spec; let main_lde = rows @@ -228,19 +229,31 @@ pub fn decide(lengths: &TableLengths, blowup_factor: u8) -> StorageMode { } /// Peak RAM estimate in bytes for a proof whose trace shape matches `lengths`. +/// +/// `table_parallelism` is the prover's `k` (`stark::prover::table_parallelism`), +/// and it is not only a prover knob: `decide` feeds it in here, so the `cuda` +/// arm's `cores * 2 / 3` doubles the transient term below versus the CPU arm's +/// `cores / 3` and makes `Disk` more likely. That direction is safe (it +/// over-estimates), but it means a change to `k` changes the storage decision. pub fn peak_bytes(lengths: &TableLengths, blowup_factor: u8, table_parallelism: usize) -> u64 { let blowup = blowup_factor as u64; let k = table_parallelism.max(1); let specs = table_specs(lengths); - // Persistent: every table's LDE + main/aux Merkle is alive across phase D. + // Persistent: every table's main LDE + Merkle really is alive at once (the + // Round 1 main commit is a phase-wide barrier). The aux LDE no longer is — + // it is produced and consumed inside one table's fused task, so at most k + // coexist — but it is still counted for every table here, which keeps this + // an over-estimate rather than making the bound unsound. let persistent_total: u64 = specs .iter() .map(|s| persistent_per_table(*s, blowup)) .fold(0u64, u64::saturating_add); - // Transient: only k tables run round 2-4 in parallel. Conservative bound is - // the top-k tables by transient bytes (worst possible chunk assignment). + // Transient: only k tables run the fused aux+rounds task at a time. The + // top-k tables by transient bytes bound it; with the scheduler's + // heaviest-first admission that top-k is also the set actually admitted + // first, so this is the realistic peak, not a worst case. let mut transient_per: Vec = specs .iter() .map(|s| transient_per_table(*s, blowup)) diff --git a/prover/src/instruments.rs b/prover/src/instruments.rs index 0ea28273b..4663f092c 100644 --- a/prover/src/instruments.rs +++ b/prover/src/instruments.rs @@ -71,11 +71,13 @@ pub fn print_report( row_top("AIR construction", air_construction, total); if let Some(ref mp) = mp { - let round1 = mp.main_commits + mp.aux_build + mp.aux_commit; - + // Round 1's main commits are the last phase-level barrier: every main + // root must be in the transcript before the shared LogUp challenges are + // sampled. Aux build, aux commit and rounds 2-4 are fused per table, so + // they are reported under one wall-clock parent below, with their + // components as concurrent (summed-over-drivers) sub-rows. 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_top("Round 1 (main trace commits)", mp.main_commits, total); row_sub( " Main LDE (fused GPU: LDE+Keccak+Merkle / CPU: LDE only)", mp.round1_sub.main_lde, @@ -86,7 +88,15 @@ pub fn print_report( mp.round1_sub.main_merkle, total, ); - row_sub(" Aux trace build (parallel)", mp.aux_build, total); + row_top( + "Rounds 2\u{2013}4 (aux build+commit fused)", + mp.rounds_2_4, + total, + ); + eprintln!( + " \u{2500}\u{2500} below: summed across concurrent drivers \u{2500}\u{2500}" + ); + row_sub(" Aux trace build", mp.aux_build, total); row_sub( " LogUp fingerprint (CPU)", mp.round1_sub.aux_fingerprint, @@ -118,7 +128,7 @@ pub fn print_report( mp.round1_sub.aux_merkle, total, ); - row_top("Rounds 2\u{2013}4", mp.rounds_2_4, total); + eprintln!(" \u{2500}\u{2500} per table (R2\u{2013}4) \u{2500}\u{2500}"); // Merge split tables: MEMW[0..4] → MEMW x5 let mut merged: BTreeMap = BTreeMap::new(); @@ -209,10 +219,7 @@ pub fn print_report( ("R4 queries & openings", total_queries), ]; sub_ops.sort_by(|a, b| b.1.cmp(&a.1)); - eprintln!( - " {}", - " \u{2500}\u{2500} sub-operation totals (all tables) \u{2500}\u{2500}", - ); + eprintln!(" \u{2500}\u{2500} sub-operation totals (all tables) \u{2500}\u{2500}"); for (label, dur) in &sub_ops { row_sub(&format!(" {label}"), *dur, total); } diff --git a/scripts/bench_prover_scaling.sh b/scripts/bench_prover_scaling.sh index c1196d76e..520d13492 100755 --- a/scripts/bench_prover_scaling.sh +++ b/scripts/bench_prover_scaling.sh @@ -74,15 +74,17 @@ parse_run() { /^ Trace build/ { v = secs(); if (v) print "t_trace_build=" v } /^ AIR construction/ { v = secs(); if (v) print "t_air=" v } /^ Pre-pass/ { v = secs(); if (v) print "t_prepass=" v } + # "Round 1 (main trace commits)" is the whole of round 1 now; aux build and + # aux commit are fused into the per-table region and reported under + # "Rounds 2-4" as sums over the concurrent drivers, so they can exceed it. /^ Round 1 / { v = secs(); if (v) print "t_round1=" v } - /Main trace commits/ { v = secs(); if (v) print "t_main_commits="v } /Aux trace build/ { v = secs(); if (v) print "t_aux_build=" v } /Aux trace commit/ { v = secs(); if (v) print "t_aux_commit=" v } /Rounds 2/ { v = secs(); if (v) print "t_rounds24=" v } - /Main expand_columns_to_lde/{ v = secs(); if (v) print "t_main_lde=" v } - /Aux expand_columns_to_lde/ { v = secs(); if (v) print "t_aux_lde=" v } - /Main commit \(Merkle\)/ { v = secs(); if (v) print "t_main_merkle=" v } - /Aux commit \(Merkle\)/ { v = secs(); if (v) print "t_aux_merkle=" v } + /Main LDE/ { v = secs(); if (v) print "t_main_lde=" v } + /Aux LDE/ { v = secs(); if (v) print "t_aux_lde=" v } + /Main commit \(Merkle/ { v = secs(); if (v) print "t_main_merkle=" v } + /Aux commit \(Merkle/ { v = secs(); if (v) print "t_aux_merkle=" v } /^ Total FFT/ { v = secs(); if (v) print "t_total_fft=" v } /^ Total Merkle/ { v = secs(); if (v) print "t_total_merkle="v } /^ TOTAL / { v = secs(); if (v) print "t_total=" v } @@ -90,9 +92,13 @@ parse_run() { /After trace build/ { print "h_trace_build=" $(NF-1) } /After AIR/ { print "h_air=" $(NF-1) } /After pool alloc/ { print "h_pool_alloc=" $(NF-1) } + # NOTE: the "After aux build" / "After aux commit" heap snapshots were + # removed when aux build/commit were fused into the per-table scheduler -- + # with k tables in flight there is no single moment at which either has + # finished, so the snapshot had no meaning. "After main commits" is the + # last phase-wide barrier, and "Peak heap" still guards the region below + # it, so the heap-growth regressions keep coverage of the fused region. /After main commits/ { print "h_main_commits=" $(NF-1) } - /After aux build/ { print "h_aux_build=" $(NF-1) } - /After aux commit/ { print "h_aux_commit=" $(NF-1) } ' "$stderr" grep -o 'Peak heap: [0-9]*' "$stdout" | awk '{print "peak=" $3}' @@ -184,15 +190,14 @@ print_row "Execute" t_execute s print_row "Trace build" t_trace_build s print_row "AIR construction" t_air s print_row "Pre-pass" t_prepass s -print_row "Round 1" t_round1 s -print_row " Main trace commits" t_main_commits s +print_row "Round 1 (main commits)" t_round1 s print_row " Main LDE" t_main_lde s print_row " Main Merkle" t_main_merkle s +print_row "Rounds 2-4 (fused)" t_rounds24 s print_row " Aux trace build" t_aux_build s print_row " Aux trace commit" t_aux_commit s print_row " Aux LDE" t_aux_lde s print_row " Aux Merkle" t_aux_merkle s -print_row "Rounds 2-4" t_rounds24 s print_row "Total FFT (all rounds)" t_total_fft s print_row "Total Merkle" t_total_merkle s print_row "TOTAL" t_total s @@ -206,8 +211,8 @@ if [[ "$MODE" == "heap" ]]; then print_row "After AIR construction" h_air mb print_row "After pool alloc" h_pool_alloc mb print_row "After main commits" h_main_commits mb - print_row "After aux build" h_aux_build mb - print_row "After aux commit" h_aux_commit mb + # "After aux build" / "After aux commit" intentionally absent: see the NOTE + # in parse_run. Peak heap covers the fused region they used to bracket. print_row "Peak heap" peak mb fi @@ -270,8 +275,9 @@ if [[ "$MODE" == "heap" ]]; then regress "After AIR construction" h_air mb regress "After pool alloc" h_pool_alloc mb regress "After main commits" h_main_commits mb - regress "After aux build" h_aux_build mb - regress "After aux commit" h_aux_commit mb + # The "After aux build" / "After aux commit" heap-growth guards were dropped + # with their snapshots (see the NOTE in parse_run). "Peak heap" is the + # remaining regression guard over the fused per-table region. regress "Peak heap" peak mb fi From 5d452853f5dbe4f3fa3ab23cc8e3d857ed048782 Mon Sep 17 00:00:00 2001 From: Joaquin Carletti Date: Mon, 3 Aug 2026 18:02:17 -0300 Subject: [PATCH 11/12] docs(gpu): record why scheduler drivers share pinned-staging slot 0 Per-driver slots were measured: repeated pinned allocation costs more than the shared mutex, whose transfers cross-table overlap already hides. --- crypto/math-cuda/src/device.rs | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/crypto/math-cuda/src/device.rs b/crypto/math-cuda/src/device.rs index 8bc140f21..a7c129cc8 100644 --- a/crypto/math-cuda/src/device.rs +++ b/crypto/math-cuda/src/device.rs @@ -509,6 +509,12 @@ impl Backend { /// Map `rayon::current_thread_index()` to a slot index, with a defensive /// clamp in case the rayon pool grew past the Vec we sized at init. + /// + /// The per-table scheduler's driver threads are not rayon workers: they + /// all resolve to slot 0 and deliberately share one slab. Spreading them + /// over per-driver slots costs more in repeated pinned allocation than + /// the shared mutex does — the staged transfers are already hidden by + /// cross-table overlap. fn worker_slot(&self, len: usize) -> usize { let idx = rayon::current_thread_index().unwrap_or(0); // Should be unreachable with rayon's fixed default pool, but if a From cd1cfe46cc1626c6d14bf0c433175e0157dd95df Mon Sep 17 00:00:00 2001 From: Mauro Toscano <12560266+MauroToscano@users.noreply.github.com> Date: Mon, 3 Aug 2026 18:37:26 -0300 Subject: [PATCH 12/12] fix(instruments): restore the rounds 2-4 phase wall, make the prover timing report honest (#895) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix(instruments): restore the rounds 2-4 phase span instead of plumbing depth Supersedes the approach in #893. Adversarial review showed the depth field was never the defect. `phase_table.py:121` takes its denominator from `max(s["wall_ns"] for _, s in pathed)` — the longest span, not the root of the ancestor stack — so depth-0 records never broke the "% of total" column, and `scripts/profiling/README.md:77` was accurate all along. `prover/src/continuation.rs` has also recorded spans from worker threads since long before this branch (:1146, :1205, :1299, :1328, :1415), with the comment at :1051-1053 saying so. Seeding worker depth was therefore work that bought nothing, and it would have left overlapping siblings looking like a clean tree — a subtler lie. Removed (`instruments::current_depth` / `enter_depth` / `DepthGuard` and the seeding in `run_admitted`). The real defect is label collision under summing. `phase_table.py:129` does `e["wall_ns"] += s["wall_ns"]`, so spans sharing a label are summed. On origin/main `rounds_2to4` was ONE span around the chunk loop (prover.rs:3503) and measured the phase; this branch made it one span per table, so the row became the sum of N concurrent tables — up to k times the real wall, able to exceed 100% — and no span measured the phase at all. `r1_aux_build` and `r1_aux_commit` were phase spans on main too (:3143, :3225). So: reopen `rounds_2to4` on the calling thread around the whole fused region, and rename the per-table spans `*_table` so a per-instance label can never be summed into a phase row. This also repairs `LAMBDA_VM_NSYS_CAPTURE_SPAN=rounds_2to4` (README.md:115), which with the label on the per-table span had N driver threads calling cuProfilerStart/Stop, the first to finish ending the capture. The report follows, and is compile-coupled to the same change. #893 added per-driver aux timers to fill the zeroed `aux_build` / `aux_commit` buckets; the fused stages have no wall-clock phase of their own any more, so reporting one invites exactly the misreading the label summing caused. Both timers and both `MultiProveTiming` fields are gone. The report now shows only the two phases that remain — "Round 1 (main trace commits)" and "Rounds 2-4 (aux build+commit fused in)" — with the aux CPU-time rows grouped under the fused phase behind headers stating they are summed over tables. That still fixes what #893 set out to fix: no row prints a fabricated 0.00s over live children, and "Round 1" no longer duplicates its own child. Verified on fib_iterative_1M: phase spans sum to their parent (r1_prepass 0.148 + r1_main_commit 2.493 + rounds_2to4 8.943 = 11.584 vs proving 11.585). * revert: trim the bench script back to the minimum #893 also repointed four timing regexes in `scripts/bench_prover_scaling.sh` that had gone stale earlier and independently of this branch. That is unrelated churn in a script with no Makefile target and no workflow referencing it, so it is reverted. What stays removed: the two dead heap rows and their `regress` calls (their `snap()` sources no longer exist and cannot be recreated with k tables in flight) and the two aux timing rows, which follow the report. The NOTE explaining why is kept. Nothing here was failing silently, contrary to the original review note: `regress` prints "(insufficient data)" for a missing key and `print_row` prints "-". * test: drop the scheduler unit tests * ci: force k > 1 on one prover shard so the scheduler runs concurrently Replaces the `VramGate` / `run_admitted` / `heaviest_first` unit tests added in #893 (removed in the previous commit). Every assertion they made was guaranteed by construction, already covered end to end, or unreachable from the call sites: `heaviest_first` is `(0..n).collect()` plus `sort_by_key`; a slot mixup in `run_admitted` is schedule independent, so it trips one of the three `.expect("run_admitted fills every slot")` sites or fails `multi_verify` on every PR today; `order.len() == 0` and `workers > order.len()` cannot happen, since `k` is `.max(1)`'d and `order` is always a full permutation. The one property with teeth — an over-budget table admitted alone — HANGS rather than fails if it regresses, which on an 8-10 minute shard burns to the job timeout unless wrapped in a watchdog. That was ~50 lines of permanent maintenance against approximately zero risk. The actual PR-time gap is that the scheduler never runs concurrently. `table_parallelism()` defaults to `(cores / 3).max(1)` and every job in this workflow is `runs-on: ubuntu-latest` with no larger-runner label, so PR CI proves with exactly one driver thread; `VramGate` is additionally inert on non-cuda builds, where `vram_budget` is `u64::MAX` and `acquire`'s condition always holds. `TABLE_PARALLELISM: 6` on shard 1 only is the smallest change that puts several real table closures in flight at once. The other three shards keep default-k coverage — the expression yields an empty string there, which fails to parse and falls back to the default. `prover/Cargo.toml:8` is `default = ["parallel"]`, so the env arm is the live one. Not a substitute for GPU coverage: `gpu-tests.yml` on merge_group rents a >=16-core RTX 5090, taking the cuda arm (`cores * 2 / 3`) with a finite VRAM budget, so both the concurrent and blocking paths already run before merge. This closes the PR-time gap only. * fix(bench): drop the row killed by the Round 1 relabel "Round 1 (main trace commits)" is lowercase, so `/Main trace commits/` stopped matching, and the row would have printed "-". It is also now redundant: with the aux stages fused out of round 1, `t_main_commits` and `t_round1` are the same number by construction. --- .github/workflows/pr_main.yaml | 9 ++ crypto/stark/src/instruments.rs | 94 +++--------- crypto/stark/src/prover.rs | 253 +++----------------------------- prover/src/instruments.rs | 22 ++- scripts/bench_prover_scaling.sh | 36 ++--- 5 files changed, 76 insertions(+), 338 deletions(-) diff --git a/.github/workflows/pr_main.yaml b/.github/workflows/pr_main.yaml index 08879ec3f..1ff124048 100644 --- a/.github/workflows/pr_main.yaml +++ b/.github/workflows/pr_main.yaml @@ -406,6 +406,15 @@ jobs: name: prover-tests - name: Run prover tests (shard ${{ matrix.partition }}/4) + # Shard 1 only: force k > 1 so the per-table admission scheduler really + # runs several table closures concurrently. ubuntu-latest has 2-4 vCPU + # and table_parallelism() defaults to (cores / 3).max(1), so every + # other shard proves with a single driver thread and never exercises + # the concurrent path or VramGate's blocking path on a PR. The other + # three shards keep the default-k coverage. An empty value on those + # fails to parse and falls back to the default, so this is inert there. + env: + TABLE_PARALLELISM: ${{ matrix.partition == 1 && '6' || '' }} run: | cargo nextest run \ --archive-file prover-tests.tar.zst \ diff --git a/crypto/stark/src/instruments.rs b/crypto/stark/src/instruments.rs index 790c06cb0..796aaf46f 100644 --- a/crypto/stark/src/instruments.rs +++ b/crypto/stark/src/instruments.rs @@ -4,22 +4,26 @@ use std::sync::OnceLock; use std::sync::atomic::{AtomicU64, Ordering}; use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH}; -// Wall clock span timeline: the trustworthy per step latency breakdown. +// Wall clock span timeline: the per step latency breakdown. // -// Top level phase spans open and close on the main thread at phase boundaries. -// They do not overlap and sum to their parent, so that part of the tree is a -// true latency breakdown (unlike the accum_* thread local sub timers below, -// which sum per worker CPU time across rayon threads and can exceed 100%). A -// parallel region is one span around the blocking call; its internal split is -// reported separately as CPU time, never mixed into the wall tree. +// Phase spans open and close on the thread that drives the phase, at phase +// boundaries. Those are a true latency breakdown: they do not overlap and they +// sum to their parent, unlike the accum_* thread local sub timers below, which +// sum per worker CPU time across rayon threads and can exceed 100%. A parallel +// region is one span around the blocking call; its internal split is reported +// separately as CPU time, never mixed into the wall tree. // -// Per table spans are the exception. `multi_prove` runs one driver thread per -// in flight table (`r1_aux_build`, `r1_aux_commit`, `rounds_2to4`), so up to -// `table_parallelism()` of them are open at once: they DO overlap in wall time -// and their sum exceeds their parent. They still nest correctly, because each -// driver seeds its span depth from the spawning thread (`current_depth` / -// `enter_depth`) instead of starting a fresh thread at depth 0 — but read them -// as per table wall time, not as a share of the enclosing phase. +// Two properties of the recorded data are easy to misread: +// +// - Spans are ALSO opened on worker threads, not only on the main thread — +// the per table drivers in `multi_prove` (`*_table` labels) and the +// per stage workers in `continuation.rs`. `SPAN_DEPTH` is thread local and +// a fresh thread starts at 0, so those records carry depth 0 and their +// siblings overlap in wall time. Read them as per instance wall time. +// - `scripts/profiling/phase_table.py` SUMS spans that share a label, so a +// label used once per table reports the sum over all tables, which can +// exceed the enclosing phase's wall clock by up to `table_parallelism()`. +// Give a per instance span its own label; never reuse a phase label for it. // // let _s = instruments::span("trace_build"); // RAII, stops on drop // @@ -44,32 +48,6 @@ thread_local! { static SPAN_DEPTH: std::cell::Cell = const { std::cell::Cell::new(0) }; } -/// Depth the next span opened on this thread would be stamped with. -/// -/// Read this on the thread that spawns workers and hand the value to -/// [`enter_depth`] inside each worker: a freshly spawned OS thread starts at -/// depth 0, so without it every span opened off the main thread is recorded as -/// a root sibling of `prove_total` and the tree stops being a breakdown. -pub fn current_depth() -> u16 { - SPAN_DEPTH.with(|d| d.get()) -} - -/// Restores the span depth this thread had before [`enter_depth`]. -#[must_use] -pub struct DepthGuard(u16); - -/// Seed this thread's span depth from a parent thread, restoring the previous -/// value when the returned guard drops. -pub fn enter_depth(depth: u16) -> DepthGuard { - DepthGuard(SPAN_DEPTH.with(|d| d.replace(depth))) -} - -impl Drop for DepthGuard { - fn drop(&mut self) { - SPAN_DEPTH.with(|d| d.set(self.0)); - } -} - #[must_use] pub struct SpanGuard { label: &'static str, @@ -296,16 +274,13 @@ pub struct Round1SubOps { /// Timing data collected inside `multi_prove`. pub struct MultiProveTiming { pub prepass: Duration, + /// Round 1 main trace commits. The last phase-wide barrier — every main + /// root must be absorbed before the shared LogUp challenges are sampled. pub main_commits: Duration, - /// Aux build wall time summed over the concurrent per-table drivers. It is - /// no longer a phase of its own — the fused chain runs it inside - /// `rounds_2_4`, so it overlaps itself and is a subset of that wall time, - /// not an addend. - pub aux_build: Duration, - /// Aux commit, same accounting as `aux_build`. - pub aux_commit: Duration, - /// Wall clock of the fused per-table region: aux build + aux commit + - /// rounds 2-4, all of it concurrent across `table_parallelism()` drivers. + /// Wall clock of the fused per-table region: aux build, aux commit and + /// rounds 2-4, which run as one task per table across `table_parallelism()` + /// drivers. There is no phase-level wall for the aux stages on their own + /// any more; their CPU time shows up in `round1_sub`. pub rounds_2_4: Duration, /// Sub-op breakdown for Round 1 (main + aux LDE vs Merkle). pub round1_sub: Round1SubOps, @@ -319,11 +294,6 @@ static R1_MAIN_LDE_US: AtomicU64 = AtomicU64::new(0); static R1_MAIN_MERKLE_US: AtomicU64 = AtomicU64::new(0); static R1_AUX_LDE_US: AtomicU64 = AtomicU64::new(0); static R1_AUX_MERKLE_US: AtomicU64 = AtomicU64::new(0); -// Aux build / aux commit wall time per table, summed across the concurrent -// per-table drivers of the fused chain (so, like the sub-timers, this is a -// CPU-style total that can exceed the fused region's wall clock). -static AUX_BUILD_US: AtomicU64 = AtomicU64::new(0); -static AUX_COMMIT_US: AtomicU64 = AtomicU64::new(0); // Aux build (LogUp) sub-phases, CPU time accumulated across tables/chunks. static AUX_FINGERPRINT_US: AtomicU64 = AtomicU64::new(0); static AUX_INVERT_US: AtomicU64 = AtomicU64::new(0); @@ -372,20 +342,6 @@ pub fn accum_aux_accumulate(d: Duration) { AUX_ACCUM_US.fetch_add(d.as_micros() as u64, Ordering::Relaxed); } -/// One table's aux build and aux commit wall time, from its fused-chain driver. -pub fn accum_aux_phases(build: Duration, commit: Duration) { - AUX_BUILD_US.fetch_add(build.as_micros() as u64, Ordering::Relaxed); - AUX_COMMIT_US.fetch_add(commit.as_micros() as u64, Ordering::Relaxed); -} - -/// Drain the summed per-table aux build / aux commit times. -pub fn take_aux_phases() -> (Duration, Duration) { - ( - Duration::from_micros(AUX_BUILD_US.swap(0, Ordering::Relaxed)), - Duration::from_micros(AUX_COMMIT_US.swap(0, Ordering::Relaxed)), - ) -} - pub fn take_r1_sub() -> Round1SubOps { Round1SubOps { main_lde: Duration::from_micros(R1_MAIN_LDE_US.swap(0, Ordering::Relaxed)), @@ -412,8 +368,6 @@ pub fn reset_all() { R1_MAIN_MERKLE_US.store(0, Ordering::Relaxed); R1_AUX_LDE_US.store(0, Ordering::Relaxed); R1_AUX_MERKLE_US.store(0, Ordering::Relaxed); - AUX_BUILD_US.store(0, Ordering::Relaxed); - AUX_COMMIT_US.store(0, Ordering::Relaxed); AUX_FINGERPRINT_US.store(0, Ordering::Relaxed); AUX_INVERT_US.store(0, Ordering::Relaxed); AUX_TERM_US.store(0, Ordering::Relaxed); diff --git a/crypto/stark/src/prover.rs b/crypto/stark/src/prover.rs index 4199d0a80..4047458bc 100644 --- a/crypto/stark/src/prover.rs +++ b/crypto/stark/src/prover.rs @@ -698,17 +698,9 @@ fn run_admitted( .map(|_| std::sync::Mutex::new(None)) .collect(); let cursor = std::sync::atomic::AtomicUsize::new(0); - // Spans opened inside `task` run on these worker threads, and a fresh OS - // thread's span depth starts at 0 — which would record every per-table - // span as a root instead of a child of the phase that spawned it. Carry - // the spawning thread's depth across the scope boundary. - #[cfg(feature = "instruments")] - let parent_depth = crate::instruments::current_depth(); std::thread::scope(|scope| { for _ in 0..workers.max(1).min(order.len().max(1)) { scope.spawn(|| { - #[cfg(feature = "instruments")] - let _depth = crate::instruments::enter_depth(parent_depth); loop { let pos = cursor.fetch_add(1, std::sync::atomic::Ordering::Relaxed); if pos >= order.len() { @@ -1417,15 +1409,9 @@ pub trait IsStarkProver< } /// Reconstruct Round1 for every table, print the bus balance report, and - /// validate each trace. - /// - /// Cross-table (bus balance) checks need every table's commitments at once, - /// so under `debug-checks` the fused per-table chain is split into two - /// admitted passes and this runs once, on the main thread, between them. - /// - /// `pair_cells` is the same per-table slot vector the drivers use. Each - /// driver only ever locks its own index, so locking them here — after the - /// aux pass has joined and before the rounds pass starts — is uncontended. + /// validate each trace. Called once after every table's aux commit, which + /// under `debug-checks` means between the fused chain's two admitted + /// passes — cross-table bus balance needs all the commitments at once. #[cfg(feature = "debug-checks")] fn run_debug_checks( pair_cells: &[std::sync::Mutex>], @@ -3158,9 +3144,10 @@ pub trait IsStarkProver< let mut main_commits: Vec> = Vec::with_capacity(num_airs); let mut main_ldes: Vec<(Vec>, usize)> = Vec::with_capacity(num_airs); // Optional device-side LDE handle per table, populated only when the - // R1 fused GPU pipeline produced one. Indexed by table, and moved into - // the per-table `gpu_main_cells` slots below so each handle stays - // paired with its table across the fused chain. + // R1 fused GPU pipeline produced one. Pairing is by index: this vector + // is moved into the per-table `gpu_main_cells` mutex slots below, and + // each driver only ever touches `gpu_main_cells[idx]` for its own + // table. (It used to ride a zip chain through the old phase D.) #[cfg(feature = "cuda")] let mut main_gpu_handles: Vec> = Vec::with_capacity(num_airs); @@ -3359,9 +3346,7 @@ pub trait IsStarkProver< let twiddles = &twiddle_caches[idx]; #[cfg(feature = "instruments")] - let __sp = crate::instruments::span("r1_aux_build"); - #[cfg(feature = "instruments")] - let t_aux_build = Instant::now(); + let __sp = crate::instruments::span("r1_aux_build_table"); let bus_public_inputs = if air.has_aux_trace() { air.build_auxiliary_trace(*trace, &lookup_challenges) } else { @@ -3385,14 +3370,10 @@ pub trait IsStarkProver< .map_err(|e| ProvingError::DiskSpill(format!("aux trace: {e}")))?; } #[cfg(feature = "instruments")] - let aux_build_dur = t_aux_build.elapsed(); - #[cfg(feature = "instruments")] drop(__sp); #[cfg(feature = "instruments")] - let __sp = crate::instruments::span("r1_aux_commit"); - #[cfg(feature = "instruments")] - let t_aux_commit = Instant::now(); + let __sp = crate::instruments::span("r1_aux_commit_table"); let aux_full: AuxResult = (|| -> Result, ProvingError> { if air.has_aux_trace() { @@ -3540,8 +3521,6 @@ pub trait IsStarkProver< transcript_cells[idx].lock().unwrap().append_bytes(&c.root); } #[cfg(feature = "instruments")] - crate::instruments::accum_aux_phases(aux_build_dur, t_aux_commit.elapsed()); - #[cfg(feature = "instruments")] drop(__sp); #[cfg(feature = "cuda")] @@ -3593,7 +3572,7 @@ pub trait IsStarkProver< let domain = &domains[idx]; #[cfg(feature = "instruments")] - let __sp = crate::instruments::span("rounds_2to4"); + let __sp = crate::instruments::span("rounds_2to4_table"); #[cfg(feature = "instruments")] let table_start = Instant::now(); @@ -3629,6 +3608,15 @@ pub trait IsStarkProver< #[cfg(feature = "instruments")] let phase_start = Instant::now(); + // Phase-level span for the whole fused region, opened here on the + // calling thread. The per-table spans inside it (`*_table`) are one + // instance per table and `phase_table.py` sums same-label spans, so + // they cannot stand in for the phase wall: their sum runs up to `k` + // times over it. This is also the span `LAMBDA_VM_NSYS_CAPTURE_SPAN` + // brackets, which needs exactly one instance to start/stop the + // profiler around. + #[cfg(feature = "instruments")] + let __sp = crate::instruments::span("rounds_2to4"); let peak_order = heaviest_first(&peak_estimates); @@ -3680,20 +3668,16 @@ pub trait IsStarkProver< proofs.push(result.expect("run_admitted fills every slot")?); } #[cfg(feature = "instruments")] + drop(__sp); + #[cfg(feature = "instruments")] let table_timings = table_timings_mx.into_inner().unwrap(); #[cfg(feature = "instruments")] { - // Aux build/commit are no longer phases of their own: each table's - // driver runs them inside the fused region, so these are sums over - // concurrent drivers and are a subset of `rounds_2_4`, not addends. - let (aux_build_elapsed, aux_commit_elapsed) = crate::instruments::take_aux_phases(); // Store timing data for the top-level report in prove_with_options. // Uses a thread-local to avoid changing multi_prove's return type. crate::instruments::store(crate::instruments::MultiProveTiming { prepass: prepass_elapsed, main_commits: main_commits_elapsed, - aux_build: aux_build_elapsed, - aux_commit: aux_commit_elapsed, rounds_2_4: phase_start.elapsed(), round1_sub: crate::instruments::take_r1_sub(), table_timings, @@ -4012,198 +3996,3 @@ fn print_bus_balance_report( } } } - -/// Scheduling primitives only. These are free functions over `&[u64]` with no -/// AIR, field or device dependency, so they are testable without a GPU — which -/// matters because CI never exercises them concurrently: `ubuntu-latest` has -/// 2-4 vCPU, so `table_parallelism()` floors to 1, and on non-cuda builds the -/// budget is `u64::MAX`, which makes `VramGate` inert. -#[cfg(test)] -mod scheduler_tests { - use super::{VramGate, heaviest_first, run_admitted}; - use std::sync::atomic::{AtomicUsize, Ordering}; - use std::time::Duration; - - /// Bound on how long a correct implementation may take to wake a waiter. - /// Only a failure deadline — never used to sequence the test. - const WAKE_DEADLINE: Duration = Duration::from_secs(10); - - #[test] - fn heaviest_first_is_a_descending_permutation() { - let empty: [u64; 0] = []; - assert!(heaviest_first(&empty).is_empty()); - assert_eq!(heaviest_first(&[3, 1, 2]), vec![0, 2, 1]); - - let estimates = [7u64, 0, 7, 3, 100, 1]; - let order = heaviest_first(&estimates); - let mut seen = order.clone(); - seen.sort_unstable(); - assert_eq!( - seen, - (0..estimates.len()).collect::>(), - "must be a permutation of 0..n" - ); - for w in order.windows(2) { - assert!( - estimates[w[0]] >= estimates[w[1]], - "must be descending by estimate, got {order:?}" - ); - } - } - - #[test] - fn heaviest_first_breaks_ties_by_index() { - // `sort_by_key` is stable, so equal estimates keep ascending index - // order: the admission order is a pure function of the estimates, and - // does not vary run to run. - assert_eq!(heaviest_first(&[5, 5, 5]), vec![0, 1, 2]); - assert_eq!(heaviest_first(&[1, 5, 5, 1]), vec![1, 2, 0, 3]); - } - - #[test] - fn run_admitted_fills_every_slot_exactly_once() { - // Includes the degenerate shapes: no work, one worker, more workers - // than tables, and `workers == 0` (clamped to 1 inside). - for (n, workers) in [(0, 4), (1, 1), (1, 8), (5, 0), (5, 1), (5, 8), (9, 4)] { - let estimates: Vec = (0..n).map(|i| i as u64 + 1).collect(); - let runs: Vec = (0..n).map(|_| AtomicUsize::new(0)).collect(); - let gate = VramGate::new(u64::MAX); - let out = run_admitted( - &heaviest_first(&estimates), - &estimates, - &gate, - workers, - |idx| { - runs[idx].fetch_add(1, Ordering::SeqCst); - idx * 10 - }, - ); - assert_eq!(out.len(), n, "one slot per table (n={n})"); - for i in 0..n { - assert_eq!( - runs[i].load(Ordering::SeqCst), - 1, - "table {i} ran exactly once (n={n}, workers={workers})" - ); - assert_eq!( - out[i], - Some(i * 10), - "slot {i} holds its own result (n={n}, workers={workers})" - ); - } - assert_eq!(*gate.used.lock().unwrap(), 0, "all permits released"); - } - } - - #[test] - fn concurrent_admissions_stay_under_budget() { - const BUDGET: u64 = 100; - const EACH: u64 = 40; // 2 fit, 3 do not - let estimates = vec![EACH; 16]; - let gate = VramGate::new(BUDGET); - let in_flight = AtomicUsize::new(0); - let max_in_flight = AtomicUsize::new(0); - run_admitted(&heaviest_first(&estimates), &estimates, &gate, 8, |_| { - let now = in_flight.fetch_add(1, Ordering::SeqCst) + 1; - max_in_flight.fetch_max(now, Ordering::SeqCst); - // Read under a held permit: the gate's own counter must never - // exceed the budget while anything is admitted. - assert!( - *gate.used.lock().unwrap() <= BUDGET, - "admitted bytes exceeded the budget" - ); - in_flight.fetch_sub(1, Ordering::SeqCst); - }); - assert!( - max_in_flight.load(Ordering::SeqCst) <= (BUDGET / EACH) as usize, - "more tables were in flight than the budget allows" - ); - assert_eq!(*gate.used.lock().unwrap(), 0); - } - - #[test] - fn oversized_request_is_admitted_alone_and_wakes_waiters() { - // A table bigger than the whole budget must still prove: it is admitted - // when nothing else holds bytes, rather than deadlocking forever. - let gate = VramGate::new(10); - let big = gate.acquire(100); - assert_eq!( - *gate.used.lock().unwrap(), - 100, - "oversized request admitted alone" - ); - - let (ready_tx, ready_rx) = std::sync::mpsc::channel(); - let (done_tx, done_rx) = std::sync::mpsc::channel(); - std::thread::scope(|s| { - s.spawn(|| { - ready_tx.send(()).unwrap(); - let permit = gate.acquire(5); - done_tx.send(()).unwrap(); - drop(permit); - }); - ready_rx.recv().unwrap(); - drop(big); - assert!( - done_rx.recv_timeout(WAKE_DEADLINE).is_ok(), - "dropping a permit must wake a waiter" - ); - }); - assert_eq!( - *gate.used.lock().unwrap(), - 0, - "permits release their bytes on drop" - ); - } - - /// Guards the instruments span tree: a driver thread starts at span depth - /// 0, so without the seeding in `run_admitted` every per-table span is - /// recorded as a root sibling of `prove_total` and the "% of total" column - /// stops meaning anything. Reads the depth directly instead of the global - /// span timeline, which other tests in this binary also write to. - #[cfg(feature = "instruments")] - #[test] - fn run_admitted_seeds_worker_span_depth() { - const PARENT_DEPTH: u16 = 3; - let outer = crate::instruments::enter_depth(PARENT_DEPTH); - let n = 6; - let estimates = vec![1u64; n]; - let gate = VramGate::new(u64::MAX); - let seen = run_admitted(&heaviest_first(&estimates), &estimates, &gate, 4, |_| { - crate::instruments::current_depth() - }); - for (i, depth) in seen.iter().enumerate() { - assert_eq!( - *depth, - Some(PARENT_DEPTH), - "table {i}'s driver must inherit the spawning thread's span depth" - ); - } - drop(outer); - assert_eq!( - crate::instruments::current_depth(), - 0, - "the calling thread's depth is restored" - ); - } - - #[test] - fn max_budget_gate_never_blocks() { - // The non-cuda / unqueryable-VRAM configuration. Two acquires that - // saturate `u64` must both be admitted. - let gate = VramGate::new(u64::MAX); - let held = gate.acquire(u64::MAX / 2); - let (tx, rx) = std::sync::mpsc::channel(); - std::thread::scope(|s| { - s.spawn(|| { - let _permit = gate.acquire(u64::MAX / 2 + 1000); - tx.send(()).unwrap(); - }); - assert!( - rx.recv_timeout(WAKE_DEADLINE).is_ok(), - "a u64::MAX budget must never block an acquire" - ); - }); - drop(held); - } -} diff --git a/prover/src/instruments.rs b/prover/src/instruments.rs index 4663f092c..f15a8a824 100644 --- a/prover/src/instruments.rs +++ b/prover/src/instruments.rs @@ -71,11 +71,12 @@ pub fn print_report( row_top("AIR construction", air_construction, total); if let Some(ref mp) = mp { - // Round 1's main commits are the last phase-level barrier: every main - // root must be in the transcript before the shared LogUp challenges are - // sampled. Aux build, aux commit and rounds 2-4 are fused per table, so - // they are reported under one wall-clock parent below, with their - // components as concurrent (summed-over-drivers) sub-rows. + // Only two wall-clock phases are left. Round 1's main commits are the + // last phase-wide barrier (every main root must be absorbed before the + // shared LogUp challenges are sampled); everything after it — aux + // build, aux commit, rounds 2-4 — runs as one fused task per table, so + // those three have no wall-clock phase of their own to report. Their + // CPU time is listed under the fused phase instead. row_top("Pre-pass (domains/twiddles)", mp.prepass, total); row_top("Round 1 (main trace commits)", mp.main_commits, total); row_sub( @@ -89,14 +90,11 @@ pub fn print_report( total, ); row_top( - "Rounds 2\u{2013}4 (aux build+commit fused)", + "Rounds 2\u{2013}4 (aux build+commit fused in)", mp.rounds_2_4, total, ); - eprintln!( - " \u{2500}\u{2500} below: summed across concurrent drivers \u{2500}\u{2500}" - ); - row_sub(" Aux trace build", mp.aux_build, total); + eprintln!(" \u{2500}\u{2500} aux build (CPU, summed over tables) \u{2500}\u{2500}"); row_sub( " LogUp fingerprint (CPU)", mp.round1_sub.aux_fingerprint, @@ -117,7 +115,7 @@ pub fn print_report( mp.round1_sub.aux_accumulate, total, ); - row_sub(" Aux trace commit", mp.aux_commit, total); + eprintln!(" \u{2500}\u{2500} aux commit (CPU, summed over tables) \u{2500}\u{2500}"); row_sub( " Aux LDE (fused GPU: LDE+Keccak+Merkle / CPU: LDE only)", mp.round1_sub.aux_lde, @@ -128,7 +126,7 @@ pub fn print_report( mp.round1_sub.aux_merkle, total, ); - eprintln!(" \u{2500}\u{2500} per table (R2\u{2013}4) \u{2500}\u{2500}"); + eprintln!(" \u{2500}\u{2500} per table (R2\u{2013}4 wall) \u{2500}\u{2500}"); // Merge split tables: MEMW[0..4] → MEMW x5 let mut merged: BTreeMap = BTreeMap::new(); diff --git a/scripts/bench_prover_scaling.sh b/scripts/bench_prover_scaling.sh index 520d13492..88824729c 100755 --- a/scripts/bench_prover_scaling.sh +++ b/scripts/bench_prover_scaling.sh @@ -74,17 +74,12 @@ parse_run() { /^ Trace build/ { v = secs(); if (v) print "t_trace_build=" v } /^ AIR construction/ { v = secs(); if (v) print "t_air=" v } /^ Pre-pass/ { v = secs(); if (v) print "t_prepass=" v } - # "Round 1 (main trace commits)" is the whole of round 1 now; aux build and - # aux commit are fused into the per-table region and reported under - # "Rounds 2-4" as sums over the concurrent drivers, so they can exceed it. /^ Round 1 / { v = secs(); if (v) print "t_round1=" v } - /Aux trace build/ { v = secs(); if (v) print "t_aux_build=" v } - /Aux trace commit/ { v = secs(); if (v) print "t_aux_commit=" v } /Rounds 2/ { v = secs(); if (v) print "t_rounds24=" v } - /Main LDE/ { v = secs(); if (v) print "t_main_lde=" v } - /Aux LDE/ { v = secs(); if (v) print "t_aux_lde=" v } - /Main commit \(Merkle/ { v = secs(); if (v) print "t_main_merkle=" v } - /Aux commit \(Merkle/ { v = secs(); if (v) print "t_aux_merkle=" v } + /Main expand_columns_to_lde/{ v = secs(); if (v) print "t_main_lde=" v } + /Aux expand_columns_to_lde/ { v = secs(); if (v) print "t_aux_lde=" v } + /Main commit \(Merkle\)/ { v = secs(); if (v) print "t_main_merkle=" v } + /Aux commit \(Merkle\)/ { v = secs(); if (v) print "t_aux_merkle=" v } /^ Total FFT/ { v = secs(); if (v) print "t_total_fft=" v } /^ Total Merkle/ { v = secs(); if (v) print "t_total_merkle="v } /^ TOTAL / { v = secs(); if (v) print "t_total=" v } @@ -92,13 +87,13 @@ parse_run() { /After trace build/ { print "h_trace_build=" $(NF-1) } /After AIR/ { print "h_air=" $(NF-1) } /After pool alloc/ { print "h_pool_alloc=" $(NF-1) } - # NOTE: the "After aux build" / "After aux commit" heap snapshots were - # removed when aux build/commit were fused into the per-table scheduler -- - # with k tables in flight there is no single moment at which either has - # finished, so the snapshot had no meaning. "After main commits" is the - # last phase-wide barrier, and "Peak heap" still guards the region below - # it, so the heap-growth regressions keep coverage of the fused region. /After main commits/ { print "h_main_commits=" $(NF-1) } + # No "After aux build"/"After aux commit" rows: aux build and aux commit are + # fused into the per-table scheduler, so with k tables in flight there is no + # single moment at which either has finished, and the prover no longer takes + # those snapshots. "Aux trace build"/"Aux trace commit" timing rows are gone + # for the same reason. "After main commits" and "Peak heap" still bracket + # the fused region. ' "$stderr" grep -o 'Peak heap: [0-9]*' "$stdout" | awk '{print "peak=" $3}' @@ -190,14 +185,12 @@ print_row "Execute" t_execute s print_row "Trace build" t_trace_build s print_row "AIR construction" t_air s print_row "Pre-pass" t_prepass s -print_row "Round 1 (main commits)" t_round1 s +print_row "Round 1" t_round1 s print_row " Main LDE" t_main_lde s print_row " Main Merkle" t_main_merkle s -print_row "Rounds 2-4 (fused)" t_rounds24 s -print_row " Aux trace build" t_aux_build s -print_row " Aux trace commit" t_aux_commit s print_row " Aux LDE" t_aux_lde s print_row " Aux Merkle" t_aux_merkle s +print_row "Rounds 2-4" t_rounds24 s print_row "Total FFT (all rounds)" t_total_fft s print_row "Total Merkle" t_total_merkle s print_row "TOTAL" t_total s @@ -211,8 +204,6 @@ if [[ "$MODE" == "heap" ]]; then print_row "After AIR construction" h_air mb print_row "After pool alloc" h_pool_alloc mb print_row "After main commits" h_main_commits mb - # "After aux build" / "After aux commit" intentionally absent: see the NOTE - # in parse_run. Peak heap covers the fused region they used to bracket. print_row "Peak heap" peak mb fi @@ -275,9 +266,6 @@ if [[ "$MODE" == "heap" ]]; then regress "After AIR construction" h_air mb regress "After pool alloc" h_pool_alloc mb regress "After main commits" h_main_commits mb - # The "After aux build" / "After aux commit" heap-growth guards were dropped - # with their snapshots (see the NOTE in parse_run). "Peak heap" is the - # remaining regression guard over the fused per-table region. regress "Peak heap" peak mb fi