From 1329fe0c3a9460501e3e12dd1eb00ebe672cfba2 Mon Sep 17 00:00:00 2001 From: jotabulacios Date: Tue, 21 Jul 2026 15:08:49 -0300 Subject: [PATCH 01/11] GPU composition residency + byte counters --- crypto/math-cuda/src/constraint_interp.rs | 1 + crypto/math-cuda/src/fri.rs | 1 + crypto/math-cuda/src/lde.rs | 46 +++++++- crypto/math-cuda/src/lib.rs | 1 + crypto/math-cuda/src/merkle.rs | 80 +++++++++++++- crypto/math-cuda/src/stagebytes.rs | 127 ++++++++++++++++++++++ crypto/math-cuda/tests/comp_poly_tree.rs | 33 ++++++ crypto/stark/src/gpu_lde.rs | 84 ++++++++++++-- crypto/stark/src/lib.rs | 3 + crypto/stark/src/prover.rs | 54 +++++++-- prover/benches/bench_continuation.rs | 33 +++++- 11 files changed, 434 insertions(+), 29 deletions(-) create mode 100644 crypto/math-cuda/src/stagebytes.rs diff --git a/crypto/math-cuda/src/constraint_interp.rs b/crypto/math-cuda/src/constraint_interp.rs index ba11b4f64..9b9c3874f 100644 --- a/crypto/math-cuda/src/constraint_interp.rs +++ b/crypto/math-cuda/src/constraint_interp.rs @@ -274,5 +274,6 @@ pub fn eval_composition_on_device( } let out = stream.clone_dtoh(&d_h)?; stream.synchronize()?; + crate::stagebytes::add_comp_dh_d2h(out.len() * 8); Ok(out) } diff --git a/crypto/math-cuda/src/fri.rs b/crypto/math-cuda/src/fri.rs index fb854d0a4..6bb22ff58 100644 --- a/crypto/math-cuda/src/fri.rs +++ b/crypto/math-cuda/src/fri.rs @@ -214,6 +214,7 @@ impl FriCommitState { let view = self.evals_a.slice(0..3 * n_out); self.stream.clone_dtoh(&view)? }; + crate::stagebytes::add_fri_layer_d2h(layer_evals.len() * 8); // 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. diff --git a/crypto/math-cuda/src/lde.rs b/crypto/math-cuda/src/lde.rs index 06c1a02cd..884c32366 100644 --- a/crypto/math-cuda/src/lde.rs +++ b/crypto/math-cuda/src/lde.rs @@ -1519,15 +1519,43 @@ pub fn coset_lde_batch_ext3_into( weights: &[u64], outputs: &mut [&mut [u64]], ) -> Result<()> { + coset_lde_batch_ext3_into_inner(columns, n, blowup_factor, weights, outputs, false)?; + Ok(()) +} + +/// Like [`coset_lde_batch_ext3_into`] but also retains the device buffer and +/// returns it as a [`GpuLdeExt3`] handle (slab layout `(c*3+k)*lde_size`), so +/// the composition Merkle commit and DEEP can reuse it without re-uploading it +/// from host. The host `outputs` are still filled (compat with the existing R4 +/// composition openings, which still read the host evals). +pub fn coset_lde_batch_ext3_into_keep( + columns: &[&[u64]], + n: usize, + blowup_factor: usize, + weights: &[u64], + outputs: &mut [&mut [u64]], +) -> Result { + let opt = coset_lde_batch_ext3_into_inner(columns, n, blowup_factor, weights, outputs, true)?; + Ok(opt.expect("keep_device_buf=true must return Some")) +} + +fn coset_lde_batch_ext3_into_inner( + columns: &[&[u64]], + n: usize, + blowup_factor: usize, + weights: &[u64], + outputs: &mut [&mut [u64]], + keep_device_buf: bool, +) -> Result> { if columns.is_empty() { - return Ok(()); + return Ok(None); } let m = columns.len(); assert_eq!(outputs.len(), m, "outputs must match columns count"); // Empty domain must short-circuit before the power-of-two assert // (is_power_of_two returns false for 0). if n == 0 { - return Ok(()); + return Ok(None); } assert!(n.is_power_of_two(), "n must be a power of two"); assert_eq!(weights.len(), n, "weights length must match n"); @@ -1631,7 +1659,19 @@ pub fn coset_lde_batch_ext3_into( // ext3-per-element layout. unpack_pinned_slabs_to_ext3(pinned, outputs, lde_size); drop(staging); - Ok(()) + if keep_device_buf { + // Same slab layout the comp-poly Merkle kernel and DEEP expect: + // `buf[(c*3 + k) * lde_size + r]`. + Ok(Some(GpuLdeExt3 { + buf: std::sync::Arc::new(buf), + m, + lde_size, + tree: None, + })) + } else { + drop(buf); + Ok(None) + } } /// Run the DIT butterfly body of a bit-reversed-input NTT over `m` batched diff --git a/crypto/math-cuda/src/lib.rs b/crypto/math-cuda/src/lib.rs index 493480991..98bba6df8 100644 --- a/crypto/math-cuda/src/lib.rs +++ b/crypto/math-cuda/src/lib.rs @@ -15,6 +15,7 @@ pub mod lde; pub mod logup; pub mod merkle; pub mod ntt; +pub mod stagebytes; // Re-exported for downstream crates so they can refer to CUDA primitive // types without depending on cudarc directly. diff --git a/crypto/math-cuda/src/merkle.rs b/crypto/math-cuda/src/merkle.rs index fb1125ea4..8699dae8a 100644 --- a/crypto/math-cuda/src/merkle.rs +++ b/crypto/math-cuda/src/merkle.rs @@ -400,7 +400,6 @@ fn build_comp_poly_tree_nodes_dev( let lde_size = ext3_elems; assert!(lde_size.is_power_of_two() && lde_size >= 2); let num_leaves = lde_size / 2; - let tight_total_nodes = 2 * num_leaves - 1; let be = backend()?; let stream = be.next_stream(); @@ -420,9 +419,28 @@ fn build_comp_poly_tree_nodes_dev( let mut buf = stream.alloc_zeros::(mb * lde_size)?; stream.memcpy_htod(&pinned[..mb * lde_size], &mut buf)?; stream.synchronize()?; + crate::stagebytes::add_comp_merkle_h2d(mb * lde_size * 8); drop(staging); - // Leaves into tail of a tight node buffer. + let nodes_dev = comp_poly_leaves_and_tree(&buf, m, lde_size, &stream)?; + Ok((nodes_dev, num_leaves, stream)) +} + +/// Row-pair Keccak leaves + inner tree levels over a device-resident comp-poly +/// LDE buffer in slab layout `buf[(part*3 + k) * lde_size + r]` (`m` parts). +/// Shared by the host-evals path (which packs + H2Ds into `buf` first) and the +/// device-buf path (which reuses an already-resident `GpuLdeExt3` buffer, so no +/// re-upload). Returns the tight node buffer (`(2*num_leaves - 1) * 32` bytes, +/// root at 0, leaves at the tail). +fn comp_poly_leaves_and_tree( + buf: &CudaSlice, + m: usize, + lde_size: usize, + stream: &Arc, +) -> Result> { + let be = backend()?; + let num_leaves = lde_size / 2; + let tight_total_nodes = 2 * num_leaves - 1; let mut nodes_dev = unsafe { stream.alloc::(tight_total_nodes * 32) }?; let leaves_offset_bytes = (num_leaves - 1) * 32; { @@ -436,7 +454,7 @@ fn build_comp_poly_tree_nodes_dev( unsafe { stream .launch_builder(&be.keccak_comp_poly_leaves_ext3) - .arg(&buf) + .arg(buf) .arg(&col_stride_u64) .arg(&num_parts_u64) .arg(&num_rows_u64) @@ -445,9 +463,37 @@ fn build_comp_poly_tree_nodes_dev( .launch(cfg)?; } } - build_inner_tree_levels(stream.as_ref(), be, &mut nodes_dev, num_leaves)?; - Ok((nodes_dev, num_leaves, stream)) + Ok(nodes_dev) +} + +/// Build the comp-poly Merkle tree directly from an already-resident device LDE +/// buffer (slab layout `buf[(part*3 + k) * lde_size + r]`, `m` parts), skipping +/// the host pack + H2D that [`build_comp_poly_tree_from_evals_ext3_keep`] does. +/// Used when the extended composition LDE is already on device (the `_keep` +/// extend handle). Returns the tree kept resident, root copied to host. +pub fn build_comp_poly_tree_from_dev_buf( + buf: &CudaSlice, + m: usize, + lde_size: usize, +) -> Result { + assert!(lde_size.is_power_of_two() && lde_size >= 2); + assert!( + buf.len() >= 3 * m * lde_size, + "device buf must hold at least 3*m*lde_size u64s" + ); + let be = backend()?; + let stream = be.next_stream(); + let num_leaves = lde_size / 2; + let nodes_dev = comp_poly_leaves_and_tree(buf, m, lde_size, &stream)?; + 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 @@ -468,6 +514,30 @@ pub fn build_comp_poly_tree_from_evals_ext3_keep( }) } +/// Parity harness for [`build_comp_poly_tree_from_dev_buf`]: pack `num_parts` +/// interleaved ext3 eval columns into slab layout, upload to a device buffer, +/// and build the tree via the device-buf path. Lets the parity suite exercise +/// the resident-handle Merkle commit (Step B) directly, for any `m`, without a +/// prior LDE. `parts_interleaved[i]` is `3*lde_size` u64s. +#[doc(hidden)] +pub fn build_comp_poly_tree_from_host_parts_via_dev_buf( + parts_interleaved: &[&[u64]], +) -> Result { + assert!(!parts_interleaved.is_empty()); + let m = parts_interleaved.len(); + let lde_size = parts_interleaved[0].len() / 3; + for p in parts_interleaved.iter() { + assert_eq!(p.len(), 3 * lde_size); + } + let be = backend()?; + let stream = be.next_stream(); + let mut host = vec![0u64; 3 * m * lde_size]; + pack_ext3_to_pinned_slabs(parts_interleaved, &mut host, lde_size); + let buf = stream.clone_htod(&host)?; + stream.synchronize()?; + build_comp_poly_tree_from_dev_buf(&buf, m, lde_size) +} + /// Test-only parity harness: build a FRI layer Merkle tree on device from an /// interleaved ext3 eval vector and return the full host node buffer so tests /// can compare it byte for byte against the CPU. Production folds and commits diff --git a/crypto/math-cuda/src/stagebytes.rs b/crypto/math-cuda/src/stagebytes.rs new file mode 100644 index 000000000..2b438f726 --- /dev/null +++ b/crypto/math-cuda/src/stagebytes.rs @@ -0,0 +1,127 @@ +//! Optional per-stage host↔device byte counters for profiling. +//! +//! Enabled by setting the `GPU_STAGE_BYTES` env var to any non-empty value +//! other than `0`. When disabled, every `add_*` call is a single relaxed +//! atomic load plus a branch, so the counters can be left compiled into the +//! `cuda` build without measurable cost. +//! +//! The counters attribute the round-2..4 host↔device round-trips the GPU +//! profiling found to dominate transfer, so composition-resident vs +//! FRI-resident can be prioritized with byte counts rather than inference: +//! * `comp_dh_d2h` — D2H of the composition evaluations `d_h` per table +//! * `comp_h01_h2d` — H2D of decomposed H₀/H₁ as input to the LDE extend +//! * `comp_h01_lde_d2h`— D2H of the extended H₀/H₁ LDE result back to host +//! * `comp_merkle_h2d` — H2D re-upload of those parts for the Merkle commit +//! * `comp_deep_h2d` — H2D re-upload of composition parts for DEEP fallback +//! * `fri_layer_d2h` — D2H of every FRI layer's evaluations +//! * `query_gather` — D2H of Merkle paths during query openings + +use std::sync::OnceLock; +use std::sync::atomic::{AtomicU64, Ordering}; + +static COMP_DH_D2H: AtomicU64 = AtomicU64::new(0); +static COMP_H01_H2D: AtomicU64 = AtomicU64::new(0); +static COMP_H01_LDE_D2H: AtomicU64 = AtomicU64::new(0); +static COMP_MERKLE_H2D: AtomicU64 = AtomicU64::new(0); +static COMP_DEEP_H2D: AtomicU64 = AtomicU64::new(0); +static FRI_LAYER_D2H: AtomicU64 = AtomicU64::new(0); +static QUERY_GATHER: AtomicU64 = AtomicU64::new(0); + +/// Cached once per process; the env var is read a single time. +pub fn enabled() -> bool { + static EN: OnceLock = OnceLock::new(); + *EN.get_or_init(|| { + std::env::var("GPU_STAGE_BYTES") + .map(|v| !v.is_empty() && v != "0") + .unwrap_or(false) + }) +} + +#[inline] +fn add(counter: &AtomicU64, bytes: usize) { + if enabled() { + counter.fetch_add(bytes as u64, Ordering::Relaxed); + } +} + +pub fn add_comp_dh_d2h(bytes: usize) { + add(&COMP_DH_D2H, bytes); +} +pub fn add_comp_h01_h2d(bytes: usize) { + add(&COMP_H01_H2D, bytes); +} +pub fn add_comp_h01_lde_d2h(bytes: usize) { + add(&COMP_H01_LDE_D2H, bytes); +} +pub fn add_comp_merkle_h2d(bytes: usize) { + add(&COMP_MERKLE_H2D, bytes); +} +pub fn add_comp_deep_h2d(bytes: usize) { + add(&COMP_DEEP_H2D, bytes); +} +pub fn add_fri_layer_d2h(bytes: usize) { + add(&FRI_LAYER_D2H, bytes); +} +pub fn add_query_gather(bytes: usize) { + add(&QUERY_GATHER, bytes); +} + +/// Zero every counter (call before a measured prove). +pub fn reset() { + for c in [ + &COMP_DH_D2H, + &COMP_H01_H2D, + &COMP_H01_LDE_D2H, + &COMP_MERKLE_H2D, + &COMP_DEEP_H2D, + &FRI_LAYER_D2H, + &QUERY_GATHER, + ] { + c.store(0, Ordering::Relaxed); + } +} + +/// Formatted summary, or `None` if counting is disabled. Composition rows are +/// grouped so the composition-vs-FRI split is read directly. +pub fn report() -> Option { + if !enabled() { + return None; + } + let mb = |c: &AtomicU64| c.load(Ordering::Relaxed) as f64 / 1e6; + let comp = mb(&COMP_DH_D2H) + + mb(&COMP_H01_H2D) + + mb(&COMP_H01_LDE_D2H) + + mb(&COMP_MERKLE_H2D) + + mb(&COMP_DEEP_H2D); + let fri = mb(&FRI_LAYER_D2H); + let q = mb(&QUERY_GATHER); + let mut s = String::from("GPU stage bytes (host<->device, MB):\n"); + s.push_str(&format!( + " composition d_h D2H {:>10.1}\n", + mb(&COMP_DH_D2H) + )); + s.push_str(&format!( + " composition H0/H1 in H2D {:>10.1}\n", + mb(&COMP_H01_H2D) + )); + s.push_str(&format!( + " composition H0/H1 LDE D2H {:>10.1}\n", + mb(&COMP_H01_LDE_D2H) + )); + s.push_str(&format!( + " composition Merkle H2D {:>10.1}\n", + mb(&COMP_MERKLE_H2D) + )); + s.push_str(&format!( + " composition DEEP H2D {:>10.1}\n", + mb(&COMP_DEEP_H2D) + )); + s.push_str(&format!(" composition SUBTOTAL {:>10.1}\n", comp)); + s.push_str(&format!(" FRI layer evals D2H {:>10.1}\n", fri)); + s.push_str(&format!(" query gather D2H {:>10.1}\n", q)); + s.push_str(&format!( + " TOTAL (counted) {:>10.1}\n", + comp + fri + q + )); + Some(s) +} diff --git a/crypto/math-cuda/tests/comp_poly_tree.rs b/crypto/math-cuda/tests/comp_poly_tree.rs index 51b826dd1..9959fbab5 100644 --- a/crypto/math-cuda/tests/comp_poly_tree.rs +++ b/crypto/math-cuda/tests/comp_poly_tree.rs @@ -233,3 +233,36 @@ fn comp_poly_tree_medium() { fn comp_poly_tree_large() { run_parity(14, 2, 4, 9999); } + +/// Direct parity for `build_comp_poly_tree_from_dev_buf` (the Step B +/// resident-handle Merkle commit): random ext3 eval columns → device buf → +/// tree root, compared against the CPU row-pair tree. Covers m = 1, 2, 3 (the +/// d==2 path uses m=2; d>2 would use m=3) across several sizes. +fn run_dev_buf_parity(lde_log: u32, num_parts: usize, seed: u64) { + let lde_size = 1usize << lde_log; + let mut rng = ChaCha8Rng::seed_from_u64(seed); + let parts: Vec> = (0..num_parts) + .map(|_| (0..lde_size).map(|_| rand_ext3(&mut rng)).collect()) + .collect(); + let cpu_nodes = cpu_tree_nodes(&parts); + + let parts_u64: Vec> = parts.iter().map(|c| ext3_to_u64s(c)).collect(); + let slices: Vec<&[u64]> = parts_u64.iter().map(|v| v.as_slice()).collect(); + let tree = math_cuda::merkle::build_comp_poly_tree_from_host_parts_via_dev_buf(&slices) + .expect("dev-buf comp-poly tree"); + + assert_eq!(tree.leaves_len, lde_size / 2); + assert_eq!( + tree.root, cpu_nodes[0], + "dev-buf root mismatch parts={num_parts} lde_log={lde_log}" + ); +} + +#[test] +fn comp_poly_tree_dev_buf_parity() { + for &parts in &[1usize, 2, 3] { + for lde_log in 2u32..=7 { + run_dev_buf_parity(lde_log, parts, 7000 + parts as u64 * 131 + lde_log as u64); + } + } +} diff --git a/crypto/stark/src/gpu_lde.rs b/crypto/stark/src/gpu_lde.rs index f962dc272..a2a226b7d 100644 --- a/crypto/stark/src/gpu_lde.rs +++ b/crypto/stark/src/gpu_lde.rs @@ -490,7 +490,11 @@ pub(crate) fn try_extend_two_halves_gpu( h0: &[FieldElement], h1: &[FieldElement], domain: &Domain, -) -> Option<(Vec>, Vec>)> +) -> Option<( + Vec>, + Vec>, + math_cuda::lde::GpuLdeExt3, +)> where F: IsFFTField + IsField + 'static, E: IsField + 'static, @@ -524,6 +528,7 @@ where }; let h0_raw = to_u64(h0); let h1_raw = to_u64(h1); + math_cuda::stagebytes::add_comp_h01_h2d((h0_raw.len() + h1_raw.len()) * 8); // weights[k] = g^(-k) / N as a u64. let inv_n = FieldElement::::from(n as u64).inv().expect("N nonzero"); @@ -545,7 +550,7 @@ where // Two ext3 columns (h0 + h1), each composed of 3 base-field components. const NUM_COLS: usize = 2; GPU_LDE_CALLS.fetch_add((NUM_COLS * 3) as u64, Ordering::Relaxed); - { + let handle = { let inputs: [&[u64]; 2] = [&h0_raw, &h1_raw]; // View each output Vec> as &mut [u64] of length 3*lde_size. let out0_ptr = lde_h0.as_mut_ptr() as *mut u64; @@ -558,14 +563,25 @@ where let out0_slice = unsafe { from_raw_parts_mut(out0_ptr, ext3_len) }; let out1_slice = unsafe { from_raw_parts_mut(out1_ptr, ext3_len) }; let mut outputs: [&mut [u64]; 2] = [out0_slice, out1_slice]; - if math_cuda::lde::coset_lde_batch_ext3_into(&inputs, n, blowup, &weights_u64, &mut outputs) - .is_err() - { - return None; - } - } + // `_keep` fills the host `outputs` (still consumed by R4 openings) AND + // retains the device buffer as a handle, so the Merkle commit and DEEP + // reuse it instead of re-uploading the parts. + let h = match math_cuda::lde::coset_lde_batch_ext3_into_keep( + &inputs, + n, + blowup, + &weights_u64, + &mut outputs, + ) { + Ok(h) => h, + Err(_) => return None, + }; + // The host D2H the `_into`/`_keep` still performs for compatibility. + math_cuda::stagebytes::add_comp_h01_lde_d2h((ext3_len * NUM_COLS) * 8); + h + }; - Some((lde_h0, lde_h1)) + Some((lde_h0, lde_h1, handle)) } pub(crate) static GPU_LEAF_HASH_CALLS: AtomicU64 = AtomicU64::new(0); @@ -941,6 +957,54 @@ where Some((host, dev_tree)) } +/// A-B-B-A toggles (profiling only): force the pre-residency behavior at runtime +/// so baseline / Step A / Step A+B can be measured interleaved in one binary. +/// `GPU_NO_RESIDENT_DEEP=1` makes DEEP re-upload the parts (disables Step A); +/// `GPU_NO_RESIDENT_MERKLE=1` makes the comp-poly commit re-upload from host +/// (disables Step B). Both default off (residency on). +pub fn deep_resident_disabled() -> bool { + static V: std::sync::OnceLock = std::sync::OnceLock::new(); + *V.get_or_init(|| { + std::env::var("GPU_NO_RESIDENT_DEEP") + .map(|v| v == "1") + .unwrap_or(false) + }) +} +pub fn merkle_resident_disabled() -> bool { + static V: std::sync::OnceLock = std::sync::OnceLock::new(); + *V.get_or_init(|| { + std::env::var("GPU_NO_RESIDENT_MERKLE") + .map(|v| v == "1") + .unwrap_or(false) + }) +} + +/// R2 GPU dispatch: build the composition-polynomial Merkle tree directly from +/// the device-resident extended-LDE handle (the `_keep` extend buffer), skipping +/// the host re-upload that [`try_build_comp_poly_tree_gpu`] performs. Same slab +/// layout and kernels, so the root is identical. Returns `None` on a math-cuda +/// `Err` (caller falls back to the host-evals path). +pub(crate) fn try_build_comp_poly_tree_from_handle( + handle: &math_cuda::lde::GpuLdeExt3, +) -> Option<(MerkleTree, math_cuda::lde::GpuMerkleTree)> +where + B: IsMerkleTreeBackend, +{ + if handle.lde_size < gpu_lde_threshold() { + return None; + } + let dev_tree = math_cuda::merkle::build_comp_poly_tree_from_dev_buf( + &handle.buf, + handle.m, + handle.lde_size, + ) + .ok()?; + debug_assert_eq!(dev_tree.leaves_len, handle.lde_size / 2); + 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 @@ -1509,6 +1573,7 @@ where packed[(p * 3 + 2) * lde_size + r] = chunk[2]; } } + math_cuda::stagebytes::add_comp_deep_h2d(packed.len() * 8); parts_host_packed = packed; // Host inv_denoms required when going through this path; we // validated the slice length above. @@ -1665,6 +1730,7 @@ pub(crate) fn gather_proofs_dev( stream, ) .ok()?; + math_cuda::stagebytes::add_query_gather(bytes.len()); let depth = tree.leaves_len.trailing_zeros() as usize; debug_assert_eq!(bytes.len(), positions.len() * depth * 32); let mut proofs = Vec::with_capacity(positions.len()); diff --git a/crypto/stark/src/lib.rs b/crypto/stark/src/lib.rs index 6f8e7c82e..118746dbf 100644 --- a/crypto/stark/src/lib.rs +++ b/crypto/stark/src/lib.rs @@ -17,6 +17,9 @@ pub mod frame; pub mod fri; #[cfg(feature = "cuda")] pub mod gpu_lde; +/// Re-export of the GPU per-stage byte counters (profiling only). +#[cfg(feature = "cuda")] +pub use math_cuda::stagebytes; pub mod grinding; #[cfg(feature = "instruments")] pub mod instruments; diff --git a/crypto/stark/src/prover.rs b/crypto/stark/src/prover.rs index 8c44c42a9..2ee9f0cca 100644 --- a/crypto/stark/src/prover.rs +++ b/crypto/stark/src/prover.rs @@ -1144,6 +1144,7 @@ pub trait IsStarkProver< constraint_evaluations: &[FieldElement], domain: &Domain, twiddles: &LdeTwiddles, + #[cfg(feature = "cuda")] gpu_parts_out: &mut Option, ) -> Vec>> where FieldElement: AsBytes + Sync + Send, @@ -1180,9 +1181,12 @@ pub trait IsStarkProver< // GPU fast path: batch both halves into one ext3 LDE call. Requires // `cuda` feature and a qualifying size. Falls through to CPU when not. #[cfg(feature = "cuda")] - if let Some((lde_h0, lde_h1)) = + if let Some((lde_h0, lde_h1, handle)) = crate::gpu_lde::try_extend_two_halves_gpu(&h0_evals, &h1_evals, domain) { + // Hand the resident extended-LDE buffer to the session so the + // composition Merkle commit and DEEP reuse it (mirrors the d>2 path). + *gpu_parts_out = Some(handle); return vec![lde_h0, lde_h1]; } @@ -1268,7 +1272,19 @@ pub trait IsStarkProver< // H₀(x²) = (H(x) + H(-x)) / 2 // H₁(x²) = (H(x) - H(-x)) / (2x) // On the LDE coset {g·ω^i}, we have -g·ω^i = g·ω^{i+N} since ω^N = -1. - Self::decompose_and_extend_d2(&constraint_evaluations, domain, twiddles) + #[cfg(feature = "cuda")] + { + Self::decompose_and_extend_d2( + &constraint_evaluations, + domain, + twiddles, + &mut gpu_composition_parts, + ) + } + #[cfg(not(feature = "cuda"))] + { + Self::decompose_and_extend_d2(&constraint_evaluations, domain, twiddles) + } } else if number_of_parts == 1 { // Degree bound equals trace length: constraint evals are the LDE directly. vec![constraint_evaluations] @@ -1330,12 +1346,25 @@ pub trait IsStarkProver< // copy) and returns a root only host tree. 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) - { + let (composition_poly_merkle_tree, composition_poly_root, gpu_composition_tree) = { + // Prefer the resident extended-LDE handle (no host re-upload); fall + // back to committing from the host eval Vecs when there is no handle + // (e.g. the CPU-extend path or the d>2 branch), then to full CPU. + // GPU_NO_RESIDENT_MERKLE=1 forces the host-evals re-upload (the + // pre-Step-B path) for interleaved A-B-B-A measurement. + let use_handle = + gpu_composition_parts.is_some() && !crate::gpu_lde::merkle_resident_disabled(); + let gpu_tree = if use_handle { + crate::gpu_lde::try_build_comp_poly_tree_from_handle::< + BatchedMerkleTreeBackend, + >(gpu_composition_parts.as_ref().unwrap()) + } else { + crate::gpu_lde::try_build_comp_poly_tree_gpu::< + FieldExtension, + BatchedMerkleTreeBackend, + >(&lde_composition_poly_parts_evaluations) + }; + match gpu_tree { Some((host_tree, dev_tree)) => { let root = host_tree.root; (host_tree, root, Some(dev_tree)) @@ -1348,7 +1377,8 @@ pub trait IsStarkProver< .ok_or(ProvingError::EmptyCommitment)?; (tree, root, None) } - }; + } + }; #[cfg(not(feature = "cuda"))] let (composition_poly_merkle_tree, composition_poly_root) = crate::commitment::commit_bit_reversed( @@ -1366,7 +1396,11 @@ pub trait IsStarkProver< // R2 to R4). The host evaluations stay in `Round2` for R4 openings. #[cfg(feature = "cuda")] if let Some(handle) = gpu_composition_parts { - round_1_result.lde_trace.set_gpu_composition_parts(handle); + // GPU_NO_RESIDENT_DEEP=1 withholds the handle so DEEP re-uploads + // (the pre-Step-A baseline), for interleaved A-B-B-A measurement. + if !crate::gpu_lde::deep_resident_disabled() { + round_1_result.lde_trace.set_gpu_composition_parts(handle); + } } Ok(Round2 { diff --git a/prover/benches/bench_continuation.rs b/prover/benches/bench_continuation.rs index c8638346f..ef8fc4b19 100644 --- a/prover/benches/bench_continuation.rs +++ b/prover/benches/bench_continuation.rs @@ -26,7 +26,7 @@ use std::time::Instant; fn main() { let args: Vec = std::env::args().collect(); if args.len() < 3 { - eprintln!("usage: bench_continuation [epoch_size]"); + eprintln!("usage: bench_continuation [epoch_size]"); std::process::exit(2); } let mode = args[1].as_str(); @@ -107,9 +107,15 @@ fn main() { } } "main" => { + #[cfg(feature = "cuda")] + stark::stagebytes::reset(); lambda_vm_prover::prove_with_inputs(&elf, &private_inputs) .expect("monolithic prove failed"); println!("main prove ok ({} bytes ELF)", elf.len()); + #[cfg(feature = "cuda")] + if let Some(r) = stark::stagebytes::report() { + print!("{r}"); + } } "cont" => { let epoch_size_log2: u32 = args @@ -132,8 +138,31 @@ fn main() { 1usize << epoch_size_log2 ); } + "warm" => { + // Run N proves in one process (arg 3, default 2). The first prove + // pays one-time costs (PTX/cubin load, pinned-buffer growth, + // stream-pool warm-up); later proves isolate steady-state cost. Under + // nsys the proves are separated by each prove's CPU-only + // execution+trace-gen gap, so the timeline can be split per prove + // (e.g. to check whether pinned-buffer allocation converges to ~0). + let n_proves: usize = args + .get(3) + .map(|s| s.parse().expect("bad prove count")) + .unwrap_or(2); + for i in 0..n_proves { + #[cfg(feature = "cuda")] + stark::stagebytes::reset(); + let t = Instant::now(); + lambda_vm_prover::prove_with_inputs(&elf, &private_inputs).expect("prove failed"); + println!("prove {} ok ({:.2}s)", i + 1, t.elapsed().as_secs_f64()); + #[cfg(feature = "cuda")] + if let Some(r) = stark::stagebytes::report() { + print!("{r}"); + } + } + } other => { - eprintln!("unknown mode {other:?}; use count|footprint|main|cont"); + eprintln!("unknown mode {other:?}; use count|footprint|main|warm|cont"); std::process::exit(2); } } From 4dd46399d3755a809fbd7889050f38f5edec8eee Mon Sep 17 00:00:00 2001 From: jotabulacios Date: Tue, 21 Jul 2026 15:30:16 -0300 Subject: [PATCH 02/11] Add monolithic prove+verify bench mode --- prover/benches/bench_continuation.rs | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/prover/benches/bench_continuation.rs b/prover/benches/bench_continuation.rs index ef8fc4b19..9214d8f22 100644 --- a/prover/benches/bench_continuation.rs +++ b/prover/benches/bench_continuation.rs @@ -117,6 +117,22 @@ fn main() { print!("{r}"); } } + "mainverify" => { + // Monolithic prove + verify (the wall workload), so the + // composition-resident change is checked end-to-end on the exact + // path measured, not only via continuation. + #[cfg(feature = "cuda")] + stark::stagebytes::reset(); + let proof = lambda_vm_prover::prove_with_inputs(&elf, &private_inputs) + .expect("monolithic prove failed"); + let ok = lambda_vm_prover::verify(&proof, &elf).expect("verify errored"); + assert!(ok, "monolithic proof did NOT verify"); + println!("main prove+verify ok"); + #[cfg(feature = "cuda")] + if let Some(r) = stark::stagebytes::report() { + print!("{r}"); + } + } "cont" => { let epoch_size_log2: u32 = args .get(3) From 2c643c996357f601140e08c03e03a9a8f6f1a88e Mon Sep 17 00:00:00 2001 From: jotabulacios Date: Tue, 21 Jul 2026 16:22:37 -0300 Subject: [PATCH 03/11] Read composition OOD/openings off GPU handle --- crypto/math-cuda/src/deep.rs | 11 +++ crypto/math-cuda/src/lde.rs | 21 +++++ crypto/math-cuda/src/merkle.rs | 8 +- crypto/math-cuda/tests/barycentric_strided.rs | 1 + crypto/math-cuda/tests/deep.rs | 1 + crypto/math-cuda/tests/gather_rows.rs | 1 + crypto/stark/src/gpu_lde.rs | 56 ++++++++++++ crypto/stark/src/prover.rs | 85 ++++++++++++++++++- crypto/stark/tests/gpu_constraint_interp.rs | 2 + 9 files changed, 183 insertions(+), 3 deletions(-) diff --git a/crypto/math-cuda/src/deep.rs b/crypto/math-cuda/src/deep.rs index 581fbc404..79df6b5d9 100644 --- a/crypto/math-cuda/src/deep.rs +++ b/crypto/math-cuda/src/deep.rs @@ -88,6 +88,11 @@ pub fn deep_composition_ext3_with_dev_parts( ) -> Result> { let be = backend()?; let stream = be.next_stream(); + // Order this stream against the producer's fill of the composition LDE + // (no-op while the producer still `synchronize()`s; real barrier at C4). + if let Some(ev) = h_parts_dev.ready.as_deref() { + stream.wait(ev)?; + } deep_composition_ext3_impl( &stream, main_lde, @@ -174,6 +179,12 @@ pub fn deep_composition_ext3_with_dev_parts_and_inv_denoms( let be = backend()?; + // Order the caller's stream against the producer's fill of the composition + // LDE (no-op while the producer still `synchronize()`s; real barrier at C4). + if let Some(ev) = h_parts_dev.ready.as_deref() { + stream.wait(ev)?; + } + // H2D only the small scalars on the caller's stream. let h_ood_dev = stream.clone_htod(h_ood)?; let trace_ood_dev = stream.clone_htod(trace_ood)?; diff --git a/crypto/math-cuda/src/lde.rs b/crypto/math-cuda/src/lde.rs index 884c32366..4bf2d5bdd 100644 --- a/crypto/math-cuda/src/lde.rs +++ b/crypto/math-cuda/src/lde.rs @@ -625,6 +625,9 @@ pub fn coset_lde_ext3_row_major_with_merkle_tree_keep( m, lde_size: n * blowup_factor, tree: Some(tree), + // This path already `synchronize()`s the fill before returning, so no + // cross-stream event is needed. + ready: None, }; Ok((handle, lde_out)) } @@ -656,6 +659,9 @@ pub fn coset_lde_ext3_row_major_with_merkle_tree_keep_dev( m, lde_size: n * blowup_factor, tree: Some(tree), + // This path already `synchronize()`s the fill before returning, so no + // cross-stream event is needed. + ready: None, }; Ok((handle, lde_out)) } @@ -692,6 +698,13 @@ pub struct GpuLdeExt3 { /// Optionally the aux or composition Merkle tree kept resident on device /// (the keep path), so R4 openings gather paths on device. None otherwise. pub tree: Option, + /// "LDE ready" event recorded on the producer stream after the NTT that + /// fills `buf`. A consumer on a different stream must `stream.wait(ready)` + /// before reading `buf`. `Some` only when the producer left a resident buffer + /// without a host-side `synchronize()` barrier (the composition path headed + /// for device-only openings); `None` when a `synchronize()` already ordered + /// the fill (all current paths), so waiting is unnecessary. + pub ready: Option>, } /// Merkle tree kept resident on device after a commit, so query openings gather @@ -1459,6 +1472,8 @@ fn evaluate_poly_coset_batch_ext3_into_inner( m, lde_size, tree: None, + // d>2 path: still `synchronize()`s the fill above; no event needed. + ready: None, })) } else { drop(buf); @@ -1662,11 +1677,17 @@ fn coset_lde_batch_ext3_into_inner( if keep_device_buf { // Same slab layout the comp-poly Merkle kernel and DEEP expect: // `buf[(c*3 + k) * lde_size + r]`. + // Record an "LDE ready" event on the producer stream so a consumer on + // another stream can order its reads of `buf` against the fill. Today a + // `synchronize()` above already orders it (belt-and-suspenders); it + // becomes the real barrier once Step C4 drops that host D2H+sync. + let ready = Some(std::sync::Arc::new(stream.record_event(None)?)); Ok(Some(GpuLdeExt3 { buf: std::sync::Arc::new(buf), m, lde_size, tree: None, + ready, })) } else { drop(buf); diff --git a/crypto/math-cuda/src/merkle.rs b/crypto/math-cuda/src/merkle.rs index 8699dae8a..08a604b6e 100644 --- a/crypto/math-cuda/src/merkle.rs +++ b/crypto/math-cuda/src/merkle.rs @@ -476,6 +476,7 @@ pub fn build_comp_poly_tree_from_dev_buf( buf: &CudaSlice, m: usize, lde_size: usize, + ready: Option<&cudarc::driver::CudaEvent>, ) -> Result { assert!(lde_size.is_power_of_two() && lde_size >= 2); assert!( @@ -484,6 +485,11 @@ pub fn build_comp_poly_tree_from_dev_buf( ); let be = backend()?; let stream = be.next_stream(); + // Order this (different) stream against the producer's fill of `buf`. A no-op + // while the producer still `synchronize()`s; the real barrier at Step C4. + if let Some(ev) = ready { + stream.wait(ev)?; + } let num_leaves = lde_size / 2; let nodes_dev = comp_poly_leaves_and_tree(buf, m, lde_size, &stream)?; let mut root = [0u8; 32]; @@ -535,7 +541,7 @@ pub fn build_comp_poly_tree_from_host_parts_via_dev_buf( pack_ext3_to_pinned_slabs(parts_interleaved, &mut host, lde_size); let buf = stream.clone_htod(&host)?; stream.synchronize()?; - build_comp_poly_tree_from_dev_buf(&buf, m, lde_size) + build_comp_poly_tree_from_dev_buf(&buf, m, lde_size, None) } /// Test-only parity harness: build a FRI layer Merkle tree on device from an diff --git a/crypto/math-cuda/tests/barycentric_strided.rs b/crypto/math-cuda/tests/barycentric_strided.rs index d96f7128b..38b4a9b9e 100644 --- a/crypto/math-cuda/tests/barycentric_strided.rs +++ b/crypto/math-cuda/tests/barycentric_strided.rs @@ -109,6 +109,7 @@ fn run_ext3(log_trace: u32, blowup: usize, num_cols: usize, seed: u64) { m: num_cols, lde_size, tree: None, + ready: None, }; // Pre-strided buffer for non-strided reference. diff --git a/crypto/math-cuda/tests/deep.rs b/crypto/math-cuda/tests/deep.rs index b7c027914..8a42fdc42 100644 --- a/crypto/math-cuda/tests/deep.rs +++ b/crypto/math-cuda/tests/deep.rs @@ -187,6 +187,7 @@ fn run_parity( m: num_aux, lde_size, tree: None, + ready: None, }) } else { drop(aux_dev); diff --git a/crypto/math-cuda/tests/gather_rows.rs b/crypto/math-cuda/tests/gather_rows.rs index 8b7f1b4a3..9aa03a15a 100644 --- a/crypto/math-cuda/tests/gather_rows.rs +++ b/crypto/math-cuda/tests/gather_rows.rs @@ -61,6 +61,7 @@ fn run_ext3(lde_size: usize, num_cols: usize, seed: u64) { m: num_cols, lde_size, tree: None, + ready: None, }; let rows: Vec = (0..9).map(|_| rng.gen_range(0..lde_size) as u32).collect(); diff --git a/crypto/stark/src/gpu_lde.rs b/crypto/stark/src/gpu_lde.rs index a2a226b7d..91a6e14dc 100644 --- a/crypto/stark/src/gpu_lde.rs +++ b/crypto/stark/src/gpu_lde.rs @@ -997,6 +997,7 @@ where &handle.buf, handle.m, handle.lde_size, + handle.ready.as_deref(), ) .ok()?; debug_assert_eq!(dev_tree.leaves_len, handle.lde_size / 2); @@ -1173,6 +1174,61 @@ where Some(apply_ext3_scalar::(&sums_raw, scalar, num_cols)) } +/// R3 OOD for the composition parts on the resident composition LDE handle. +/// Same barycentric core as [`try_barycentric_ext3_on_handle`] but reads an +/// explicit `GpuLdeExt3` (the composition parts) instead of the aux trace, so +/// R3 evaluates each part at `z^num_parts` on device. Returns one ext3 value per +/// part, or `None` to fall through to the CPU path. (Step C1: the caller keeps +/// the CPU result authoritative and cross-checks this against it.) +#[allow(clippy::too_many_arguments)] +pub(crate) fn try_barycentric_ext3_on_comp_handle( + handle: &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], +) -> Option>> +where + F: IsField + IsSubFieldOf + 'static, + E: IsField + 'static, +{ + if TypeId::of::() != TypeId::of::() { + return None; + } + if TypeId::of::() != TypeId::of::() { + return None; + } + let num_cols = handle.m; + if num_cols == 0 { + return Some(Vec::new()); + } + let n = coset_points.len(); + if !n.is_power_of_two() || n < gpu_bary_threshold() { + return None; + } + if handle.lde_size != n.checked_mul(row_stride)? { + return None; + } + if inv_denoms_host.len() != n { + return None; + } + let points_raw: &[u64] = unsafe { from_raw_parts(coset_points.as_ptr() as *const u64, n) }; + let inv_denoms_len = n.checked_mul(3).expect("inv_denoms u64 len overflow"); + let inv_denoms_raw: &[u64] = + unsafe { from_raw_parts(inv_denoms_host.as_ptr() as *const u64, inv_denoms_len) }; + let sums_raw = + match math_cuda::barycentric::barycentric_ext3_on_device(handle, row_stride, points_raw, inv_denoms_raw, n) { + Ok(v) => v, + Err(_) => return None, + }; + GPU_BARY_CALLS.fetch_add(1, Ordering::Relaxed); + let scalar = ood_ext3_scalar::(coset_offset_pow_n, n_inv, g_n_inv, z_pow_n); + Some(apply_ext3_scalar::(&sums_raw, scalar, num_cols)) +} + // ============================================================================ // R2 keep-handle variant, R4 DEEP composition, FRI commit dispatches // ============================================================================ diff --git a/crypto/stark/src/prover.rs b/crypto/stark/src/prover.rs index 2ee9f0cca..0ff39c1d6 100644 --- a/crypto/stark/src/prover.rs +++ b/crypto/stark/src/prover.rs @@ -1456,6 +1456,39 @@ pub trait IsStarkProver< }) .collect(); + // Step C1: compute the same composition-parts OOD on the resident device + // handle and cross-check it against the CPU result above (host stays + // authoritative this stage). Proves the device path before C4 makes it + // the sole source and the host composition LDE is dropped. + #[cfg(feature = "cuda")] + if let Some(handle) = round_1_result.lde_trace.gpu_composition_parts() { + if let Some(gpu_ood) = + crate::gpu_lde::try_barycentric_ext3_on_comp_handle::( + handle, + blowup_factor, + &dc.points, + &dc.offset_pow_n, + &dc.size_inv, + &dc.offset_pow_n_inv, + &comp_z_pow_n, + &comp_inv_denoms, + ) + { + assert_eq!( + gpu_ood.len(), + composition_poly_parts_ood_evaluation.len(), + "C1 composition OOD: GPU/CPU part count mismatch" + ); + for (j, (g, c)) in gpu_ood + .iter() + .zip(composition_poly_parts_ood_evaluation.iter()) + .enumerate() + { + assert_eq!(g, c, "C1 composition OOD: GPU != CPU at part {j}"); + } + } + } + // === Trace polynomials: barycentric evaluation via LDE === let trace_ood_evaluations = crate::trace::get_trace_evaluations_from_lde( &round_1_result.lde_trace, @@ -2228,6 +2261,34 @@ pub trait IsStarkProver< }) }); + // Step C2: gather the composition parts' query row-pairs off the resident + // handle (same `query_rows`, same lde_size as the trace). Cross-checked + // against the host `open_composition_poly_with_proof` values per query in + // the loop below (host authoritative this stage); becomes the sole source + // at C4 when the host composition LDE is dropped. + #[cfg(feature = "cuda")] + let comp_dev_values: Option>> = + comp_dev_proofs.as_ref().and_then(|_| { + round_1_result.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; @@ -2272,11 +2333,31 @@ pub trait IsStarkProver< #[cfg(feature = "cuda")] { if let Some(proofs) = &comp_dev_proofs { - Self::open_composition_poly_with_proof( + let host = Self::open_composition_poly_with_proof( proofs[qi].clone(), &round_2_result.lde_composition_poly_evaluations, *index, - ) + ); + // C2 cross-check: the values gathered off the resident + // handle must equal the host values (host authoritative). + if let Some(dev_vals) = comp_dev_values.as_ref() { + let num_parts_comp = + round_2_result.lde_composition_poly_evaluations.len(); + let (even, odd) = Self::device_row_pair::( + dev_vals, + qi, + num_parts_comp, + ); + assert_eq!( + even, host.evaluations, + "C2 composition open: GPU != CPU (even) query {qi}" + ); + assert_eq!( + odd, host.evaluations_sym, + "C2 composition open: GPU != CPU (sym) query {qi}" + ); + } + host } else { Self::open_composition_poly( &round_2_result.composition_poly_merkle_tree, diff --git a/crypto/stark/tests/gpu_constraint_interp.rs b/crypto/stark/tests/gpu_constraint_interp.rs index de7211dcd..7b0349a10 100644 --- a/crypto/stark/tests/gpu_constraint_interp.rs +++ b/crypto/stark/tests/gpu_constraint_interp.rs @@ -206,6 +206,7 @@ fn check_program(prog: &ConstraintProgram, label: &str, seed: u64) { m: aux_cols, lde_size, tree: None, + ready: None, }; // GPU: launch the interpreter over every row. @@ -421,6 +422,7 @@ fn check_composition(prog: &ConstraintProgram, label: &str, seed: u64) m: aux_cols, lde_size, tree: None, + ready: None, }; let inputs = CompositionInputs { From 1ee5a97d18bb95686d59f8936fd1d7dff2050418 Mon Sep 17 00:00:00 2001 From: jotabulacios Date: Tue, 21 Jul 2026 16:36:12 -0300 Subject: [PATCH 04/11] Add composition LDE-ready event waits + fmt --- crypto/math-cuda/src/barycentric.rs | 10 ++++++++++ crypto/stark/src/gpu_lde.rs | 15 ++++++++++----- 2 files changed, 20 insertions(+), 5 deletions(-) diff --git a/crypto/math-cuda/src/barycentric.rs b/crypto/math-cuda/src/barycentric.rs index 8fdea14f1..5cdbdfa30 100644 --- a/crypto/math-cuda/src/barycentric.rs +++ b/crypto/math-cuda/src/barycentric.rs @@ -257,6 +257,11 @@ pub fn barycentric_ext3_on_device( let be = backend()?; let stream = be.next_stream(); + // Order this stream against the producer's fill of the LDE (no-op unless the + // handle carries a ready event, i.e. the composition device-only path). + if let Some(ev) = aux_handle.ready.as_deref() { + stream.wait(ev)?; + } let points_dev = stream.clone_htod(coset_points)?; let inv_dev = stream.clone_htod(inv_denoms_ext3)?; @@ -390,6 +395,11 @@ pub fn gather_rows_ext3_on_device( return Ok(Vec::new()); } let be = backend()?; + // Order the caller's stream against the producer's fill (no-op unless the + // handle carries a ready event, i.e. the composition device-only path). + if let Some(ev) = aux.ready.as_deref() { + stream.wait(ev)?; + } let rows_dev = stream.clone_htod(rows)?; let mut out = stream.alloc_zeros::(rows.len() * num_cols * 3)?; let col_stride = aux.lde_size as u64; diff --git a/crypto/stark/src/gpu_lde.rs b/crypto/stark/src/gpu_lde.rs index 91a6e14dc..32c788dc7 100644 --- a/crypto/stark/src/gpu_lde.rs +++ b/crypto/stark/src/gpu_lde.rs @@ -1219,11 +1219,16 @@ where let inv_denoms_len = n.checked_mul(3).expect("inv_denoms u64 len overflow"); let inv_denoms_raw: &[u64] = unsafe { from_raw_parts(inv_denoms_host.as_ptr() as *const u64, inv_denoms_len) }; - let sums_raw = - match math_cuda::barycentric::barycentric_ext3_on_device(handle, row_stride, points_raw, inv_denoms_raw, n) { - Ok(v) => v, - Err(_) => return None, - }; + let sums_raw = match math_cuda::barycentric::barycentric_ext3_on_device( + handle, + row_stride, + points_raw, + inv_denoms_raw, + n, + ) { + Ok(v) => v, + Err(_) => return None, + }; GPU_BARY_CALLS.fetch_add(1, Ordering::Relaxed); let scalar = ood_ext3_scalar::(coset_offset_pow_n, n_inv, g_n_inv, z_pow_n); Some(apply_ext3_scalar::(&sums_raw, scalar, num_cols)) From f786e2f1b1de7bca2a6dd36bb171780e553534ec Mon Sep 17 00:00:00 2001 From: jotabulacios Date: Tue, 21 Jul 2026 16:43:32 -0300 Subject: [PATCH 05/11] Round2: explicit Option host composition LDE --- crypto/stark/src/prover.rs | 54 ++++++++++++++++++++++++++++---------- 1 file changed, 40 insertions(+), 14 deletions(-) diff --git a/crypto/stark/src/prover.rs b/crypto/stark/src/prover.rs index 0ff39c1d6..972ab5bb6 100644 --- a/crypto/stark/src/prover.rs +++ b/crypto/stark/src/prover.rs @@ -512,8 +512,15 @@ where F: IsField, FieldElement: AsBytes, { - /// Evaluations of the composition polynomial parts over the LDE domain. - pub(crate) lde_composition_poly_evaluations: Vec>>, + /// Host-side evaluations of the composition polynomial parts over the LDE + /// domain. `Some` on the CPU path and while the composition LDE is retained + /// on host (Steps A–C2, and C4 under `GPU_RETAIN_COMPOSITION_HOST`); `None` + /// when the composition LDE is device-only (C4), so every host consumer must + /// handle the device-only case explicitly instead of indexing empty data. + pub(crate) host_composition_lde: Option>>>, + /// Number of composition parts, always available even when the host LDE is + /// dropped (device-only). Equals `air.composition_poly_degree_bound / trace_length`. + pub(crate) num_composition_parts: usize, /// The Merkle tree built to compute the commitment to the composition polynomial parts. pub(crate) composition_poly_merkle_tree: BatchedMerkleTree, /// The commitment to the composition polynomial parts. @@ -526,6 +533,23 @@ where pub(crate) gpu_composition_tree: Option, } +impl Round2 +where + F: IsField, + FieldElement: AsBytes, +{ + /// The host composition-LDE evaluations. Panics if the composition LDE is + /// device-only (`host_composition_lde` is `None`): every caller of this must + /// be on a path that retained the host copy. The device-only paths read the + /// resident handle instead and never call this. + pub(crate) fn host_evals(&self) -> &[Vec>] { + self.host_composition_lde + .as_ref() + .expect("host composition LDE present (device-only paths must read the resident handle)") + .as_slice() + } +} + /// A container for the results of the third round of the STARK Prove protocol. pub(crate) struct Round3 { /// Evaluations of the trace polynomials, main and auxiliary, at the out-of-domain challenge. @@ -1403,8 +1427,10 @@ pub trait IsStarkProver< } } + let num_composition_parts = lde_composition_poly_parts_evaluations.len(); Ok(Round2 { - lde_composition_poly_evaluations: lde_composition_poly_parts_evaluations, + host_composition_lde: Some(lde_composition_poly_parts_evaluations), + num_composition_parts, composition_poly_merkle_tree, composition_poly_root, #[cfg(feature = "cuda")] @@ -1424,7 +1450,7 @@ pub trait IsStarkProver< FieldElement: AsBytes, FieldElement: AsBytes, { - let num_parts = round_2_result.lde_composition_poly_evaluations.len(); + let num_parts = round_2_result.num_composition_parts; let z_power = z.pow(num_parts); let domain_size = domain.interpolation_domain_size; let blowup_factor = domain.blowup_factor; @@ -1437,7 +1463,7 @@ pub trait IsStarkProver< 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 + .host_evals() .iter() .map(|lde_evals| { // Extract trace-size evaluations (stride = blowup_factor) @@ -1541,7 +1567,7 @@ pub trait IsStarkProver< let gamma = transcript.sample_field_element(); - let n_terms_composition_poly = round_2_result.lde_composition_poly_evaluations.len(); + let n_terms_composition_poly = round_2_result.num_composition_parts; // g·z pruning: only the current-row block (all columns) plus the masked // next-row columns get an opening / DEEP coefficient. let layout = Self::ood_layout(air); @@ -1679,7 +1705,7 @@ pub trait IsStarkProver< FieldElement: AsBytes, FieldElement: AsBytes, { - let num_parts = round_2_result.lde_composition_poly_evaluations.len(); + let num_parts = round_2_result.num_composition_parts; let z_power = z.pow(num_parts); // pole for H terms // Number of evaluation points per trace column (= transition_offsets.len() * step_size) @@ -1728,7 +1754,7 @@ pub trait IsStarkProver< crate::gpu_lde::try_deep_composition_gpu::( lde_trace, lde_trace.gpu_composition_parts(), - &round_2_result.lde_composition_poly_evaluations, + round_2_result.host_evals(), h_ood, &trace_ood_columns, composition_poly_gammas, @@ -1764,7 +1790,7 @@ pub trait IsStarkProver< crate::gpu_lde::try_deep_composition_gpu::( lde_trace, lde_trace.gpu_composition_parts(), - &round_2_result.lde_composition_poly_evaluations, + round_2_result.host_evals(), h_ood, &trace_ood_columns, composition_poly_gammas, @@ -1830,7 +1856,7 @@ pub trait IsStarkProver< // H terms for j in 0..num_parts { - let h_j_val = &round_2_result.lde_composition_poly_evaluations[j][i]; + let h_j_val = &round_2_result.host_evals()[j][i]; let h_j_ood = &h_ood[j]; result += &composition_poly_gammas[j] * (h_j_val - h_j_ood) * &inv_h[i]; } @@ -2335,14 +2361,14 @@ pub trait IsStarkProver< if let Some(proofs) = &comp_dev_proofs { let host = Self::open_composition_poly_with_proof( proofs[qi].clone(), - &round_2_result.lde_composition_poly_evaluations, + round_2_result.host_evals(), *index, ); // C2 cross-check: the values gathered off the resident // handle must equal the host values (host authoritative). if let Some(dev_vals) = comp_dev_values.as_ref() { let num_parts_comp = - round_2_result.lde_composition_poly_evaluations.len(); + round_2_result.num_composition_parts; let (even, odd) = Self::device_row_pair::( dev_vals, qi, @@ -2361,7 +2387,7 @@ pub trait IsStarkProver< } else { Self::open_composition_poly( &round_2_result.composition_poly_merkle_tree, - &round_2_result.lde_composition_poly_evaluations, + round_2_result.host_evals(), *index, ) } @@ -2370,7 +2396,7 @@ pub trait IsStarkProver< { Self::open_composition_poly( &round_2_result.composition_poly_merkle_tree, - &round_2_result.lde_composition_poly_evaluations, + round_2_result.host_evals(), *index, ) } From 6fb74a59a9e6475ca3ca1da8d1a904397c9d7db3 Mon Sep 17 00:00:00 2001 From: jotabulacios Date: Tue, 21 Jul 2026 17:47:04 -0300 Subject: [PATCH 06/11] Drop composition D2H when DEEP inputs resident --- crypto/math-cuda/src/lde.rs | 92 ++++++++++---- crypto/stark/src/gpu_lde.rs | 157 ++++++++++++++++++----- crypto/stark/src/prover.rs | 173 +++++++++++++++----------- prover/tests/cuda_path_integration.rs | 9 +- 4 files changed, 304 insertions(+), 127 deletions(-) diff --git a/crypto/math-cuda/src/lde.rs b/crypto/math-cuda/src/lde.rs index 4bf2d5bdd..e55572146 100644 --- a/crypto/math-cuda/src/lde.rs +++ b/crypto/math-cuda/src/lde.rs @@ -1534,23 +1534,33 @@ pub fn coset_lde_batch_ext3_into( weights: &[u64], outputs: &mut [&mut [u64]], ) -> Result<()> { - coset_lde_batch_ext3_into_inner(columns, n, blowup_factor, weights, outputs, false)?; + coset_lde_batch_ext3_into_inner(columns, n, blowup_factor, weights, outputs, false, true)?; Ok(()) } /// Like [`coset_lde_batch_ext3_into`] but also retains the device buffer and /// returns it as a [`GpuLdeExt3`] handle (slab layout `(c*3+k)*lde_size`), so /// the composition Merkle commit and DEEP can reuse it without re-uploading it -/// from host. The host `outputs` are still filled (compat with the existing R4 -/// composition openings, which still read the host evals). +/// from host. When `retain_host_lde` is true the host `outputs` are also filled; +/// when false the large D2H and CPU unpack are skipped and `outputs` may be +/// empty. Device consumers order their reads through the handle's `ready` event. pub fn coset_lde_batch_ext3_into_keep( columns: &[&[u64]], n: usize, blowup_factor: usize, weights: &[u64], outputs: &mut [&mut [u64]], + retain_host_lde: bool, ) -> Result { - let opt = coset_lde_batch_ext3_into_inner(columns, n, blowup_factor, weights, outputs, true)?; + let opt = coset_lde_batch_ext3_into_inner( + columns, + n, + blowup_factor, + weights, + outputs, + true, + retain_host_lde, + )?; Ok(opt.expect("keep_device_buf=true must return Some")) } @@ -1561,12 +1571,15 @@ fn coset_lde_batch_ext3_into_inner( weights: &[u64], outputs: &mut [&mut [u64]], keep_device_buf: bool, + retain_host_lde: bool, ) -> Result> { if columns.is_empty() { return Ok(None); } let m = columns.len(); - assert_eq!(outputs.len(), m, "outputs must match columns count"); + if retain_host_lde { + assert_eq!(outputs.len(), m, "outputs must match columns count"); + } // Empty domain must short-circuit before the power-of-two assert // (is_power_of_two returns false for 0). if n == 0 { @@ -1582,8 +1595,10 @@ fn coset_lde_batch_ext3_into_inner( assert_eq!(c.len(), 3 * n, "each ext3 column must be 3*n u64s"); } let lde_size = n * blowup_factor; - for o in outputs.iter() { - assert_eq!(o.len(), 3 * lde_size, "each output must be 3*lde_size u64s"); + if retain_host_lde { + 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_into lde_size"); let log_n = n.trailing_zeros() as u64; @@ -1596,18 +1611,45 @@ fn coset_lde_batch_ext3_into_inner( let stream = be.next_stream(); let staging_slot = be.pinned_staging(); - let mut staging = staging_slot.lock().unwrap(); - staging.ensure_capacity(mb * lde_size, &be.ctx)?; - let pinned = unsafe { staging.as_mut_slice(mb * lde_size) }; + let mut staging = Some(staging_slot.lock().unwrap()); + staging + .as_mut() + .expect("staging present") + .ensure_capacity(mb * lde_size, &be.ctx)?; - pack_ext3_to_pinned_slabs(columns, pinned, n); + { + let pinned = unsafe { + staging + .as_mut() + .expect("staging present") + .as_mut_slice(mb * lde_size) + }; + pack_ext3_to_pinned_slabs(columns, pinned, n); + } // Allocate + zero-pad device buffer holding 3M slabs of `lde_size`. let mut buf = stream.alloc_zeros::(mb * lde_size)?; // H2D: slab by slab into the first N slots of each `lde_size`-slab. - for s in 0..mb { - let mut dst = buf.slice_mut(s * lde_size..s * lde_size + n); - stream.memcpy_htod(&pinned[s * n..s * n + n], &mut dst)?; + { + let pinned = unsafe { + staging + .as_mut() + .expect("staging present") + .as_mut_slice(mb * lde_size) + }; + for s in 0..mb { + let mut dst = buf.slice_mut(s * lde_size..s * lde_size + n); + stream.memcpy_htod(&pinned[s * n..s * n + n], &mut dst)?; + } + } + + // Device-only releases the shared pinned staging before the long NTT. The + // H2Ds are asynchronous, so wait for just those copies before another + // worker can reuse the host buffer. We intentionally do not synchronize + // after the NTT; the `ready` event orders downstream device consumers. + if !retain_host_lde { + stream.synchronize()?; + drop(staging.take()); } let inv_tw = be.inv_twiddles_for(log_n)?; @@ -1667,20 +1709,26 @@ fn coset_lde_batch_ext3_into_inner( mb_u32, )?; - stream.memcpy_dtoh(&buf, &mut pinned[..mb * lde_size])?; - stream.synchronize()?; + if retain_host_lde { + let pinned = unsafe { + staging + .as_mut() + .expect("host-retained path keeps staging") + .as_mut_slice(mb * lde_size) + }; + stream.memcpy_dtoh(&buf, &mut pinned[..mb * lde_size])?; + stream.synchronize()?; - // Unpack: for each output column, re-interleave 3 slabs back into the - // ext3-per-element layout. - unpack_pinned_slabs_to_ext3(pinned, outputs, lde_size); + // Re-interleave 3 base-field slabs per output ext3 column. + unpack_pinned_slabs_to_ext3(pinned, outputs, lde_size); + } drop(staging); if keep_device_buf { // Same slab layout the comp-poly Merkle kernel and DEEP expect: // `buf[(c*3 + k) * lde_size + r]`. // Record an "LDE ready" event on the producer stream so a consumer on - // another stream can order its reads of `buf` against the fill. Today a - // `synchronize()` above already orders it (belt-and-suspenders); it - // becomes the real barrier once Step C4 drops that host D2H+sync. + // another stream can order its reads of `buf` against the fill. On the + // device-only path this is the only post-NTT ordering barrier. let ready = Some(std::sync::Arc::new(stream.record_event(None)?)); Ok(Some(GpuLdeExt3 { buf: std::sync::Arc::new(buf), diff --git a/crypto/stark/src/gpu_lde.rs b/crypto/stark/src/gpu_lde.rs index 32c788dc7..01f1daa8c 100644 --- a/crypto/stark/src/gpu_lde.rs +++ b/crypto/stark/src/gpu_lde.rs @@ -80,6 +80,7 @@ pub fn reset_all_gpu_call_counters() { GPU_COMPOSITION_CALLS.store(0, Ordering::Relaxed); GPU_OPENING_GATHER_CALLS.store(0, Ordering::Relaxed); GPU_DEVICE_ONLY_CALLS.store(0, Ordering::Relaxed); + GPU_COMPOSITION_DEVICE_ONLY_CALLS.store(0, Ordering::Relaxed); } pub(crate) static GPU_EXTEND_HALVES_CALLS: AtomicU64 = AtomicU64::new(0); @@ -124,6 +125,13 @@ pub fn gpu_device_only_calls() -> u64 { GPU_DEVICE_ONLY_CALLS.load(Ordering::Relaxed) } +/// Tables whose composition LDE skipped its host D2H because composition, +/// main, and aux handles were all available for the device-only R3/R4 path. +pub(crate) static GPU_COMPOSITION_DEVICE_ONLY_CALLS: AtomicU64 = AtomicU64::new(0); +pub fn gpu_composition_device_only_calls() -> u64 { + GPU_COMPOSITION_DEVICE_ONLY_CALLS.load(Ordering::Relaxed) +} + /// Runtime override to force the GPU composition path off (→ CPU accumulation). /// An escape hatch, and the A/B toggle for benchmarking the path against the CPU /// baseline in one process (no rebuild). Default off (path enabled). @@ -490,6 +498,7 @@ pub(crate) fn try_extend_two_halves_gpu( h0: &[FieldElement], h1: &[FieldElement], domain: &Domain, + retain_host_lde: bool, ) -> Option<( Vec>, Vec>, @@ -543,42 +552,49 @@ where w *= &g_inv; } - // Pre-allocate outputs. - let mut lde_h0 = vec![FieldElement::::zero(); lde_size]; - let mut lde_h1 = vec![FieldElement::::zero(); lde_size]; - - // Two ext3 columns (h0 + h1), each composed of 3 base-field components. + // Device-only avoids allocating and filling the two large host outputs. + let mut lde_h0: Vec> = Vec::new(); + let mut lde_h1: Vec> = Vec::new(); const NUM_COLS: usize = 2; + let ext3_len = lde_size + .checked_mul(3) + .expect("ext3 output length overflow"); GPU_LDE_CALLS.fetch_add((NUM_COLS * 3) as u64, Ordering::Relaxed); - let handle = { - let inputs: [&[u64]; 2] = [&h0_raw, &h1_raw]; - // View each output Vec> as &mut [u64] of length 3*lde_size. - let out0_ptr = lde_h0.as_mut_ptr() as *mut u64; - let out1_ptr = lde_h1.as_mut_ptr() as *mut u64; + let inputs: [&[u64]; 2] = [&h0_raw, &h1_raw]; + let handle = if retain_host_lde { + lde_h0 = vec![FieldElement::::zero(); lde_size]; + lde_h1 = vec![FieldElement::::zero(); lde_size]; // SAFETY: ext3 FieldElement is [u64; 3] in memory, and the Vec has len // = lde_size so the backing is 3*lde_size u64s. - let ext3_len = lde_size - .checked_mul(3) - .expect("ext3 output length overflow"); - let out0_slice = unsafe { from_raw_parts_mut(out0_ptr, ext3_len) }; - let out1_slice = unsafe { from_raw_parts_mut(out1_ptr, ext3_len) }; + let out0_slice = unsafe { from_raw_parts_mut(lde_h0.as_mut_ptr() as *mut u64, ext3_len) }; + let out1_slice = unsafe { from_raw_parts_mut(lde_h1.as_mut_ptr() as *mut u64, ext3_len) }; let mut outputs: [&mut [u64]; 2] = [out0_slice, out1_slice]; - // `_keep` fills the host `outputs` (still consumed by R4 openings) AND - // retains the device buffer as a handle, so the Merkle commit and DEEP - // reuse it instead of re-uploading the parts. let h = match math_cuda::lde::coset_lde_batch_ext3_into_keep( &inputs, n, blowup, &weights_u64, &mut outputs, + true, ) { Ok(h) => h, Err(_) => return None, }; - // The host D2H the `_into`/`_keep` still performs for compatibility. math_cuda::stagebytes::add_comp_h01_lde_d2h((ext3_len * NUM_COLS) * 8); h + } else { + let mut outputs: [&mut [u64]; 0] = []; + match math_cuda::lde::coset_lde_batch_ext3_into_keep( + &inputs, + n, + blowup, + &weights_u64, + &mut outputs, + false, + ) { + Ok(h) => h, + Err(_) => return None, + } }; Some((lde_h0, lde_h1, handle)) @@ -979,6 +995,19 @@ pub fn merkle_resident_disabled() -> bool { }) } +/// Force the A+B behavior for measurement/debugging: retain the host copy of +/// the composition LDE even when every consumer can use the resident handle. +/// The pre-residency toggles also require the host copy for their fallbacks. +pub fn retain_composition_host() -> bool { + static V: std::sync::OnceLock = std::sync::OnceLock::new(); + let explicit = *V.get_or_init(|| { + std::env::var("GPU_RETAIN_COMPOSITION_HOST") + .map(|v| v == "1") + .unwrap_or(false) + }); + explicit || deep_resident_disabled() || merkle_resident_disabled() +} + /// R2 GPU dispatch: build the composition-polynomial Merkle tree directly from /// the device-resident extended-LDE handle (the `_keep` extend buffer), skipping /// the host re-upload that [`try_build_comp_poly_tree_gpu`] performs. Same slab @@ -1483,19 +1512,39 @@ where F: IsField + IsSubFieldOf + 'static, E: IsField + 'static, { + static DEEP_DIAG: OnceLock = OnceLock::new(); + let diag = *DEEP_DIAG.get_or_init(|| { + std::env::var("GPU_DEEP_DIAG") + .map(|v| v == "1") + .unwrap_or(false) + }); + macro_rules! reject { + ($reason:expr) => {{ + if diag { + eprintln!("GPU_DEEP_DIAG reject: {}", $reason); + } + return None; + }}; + } + if TypeId::of::() != TypeId::of::() { - return None; + reject!("base field is not Goldilocks"); } if TypeId::of::() != TypeId::of::() { - return None; + reject!("extension field is not Goldilocks ext3"); } - let main = lde_trace.gpu_main()?; + let Some(main) = lde_trace.gpu_main() else { + reject!("main LDE has no device handle"); + }; let lde_size = main.lde_size; if lde_size < gpu_lde_threshold() { - return None; + reject!(format!( + "lde_size {lde_size} is below threshold {}", + gpu_lde_threshold() + )); } if !lde_size.is_power_of_two() { - return None; + reject!(format!("lde_size {lde_size} is not a power of two")); } let num_main = main.m; let aux_handle = lde_trace.gpu_aux(); @@ -1503,19 +1552,28 @@ where let num_total_cols = num_main + num_aux; let num_parts = composition_poly_gammas.len(); if h_ood.len() != num_parts { - return None; + reject!(format!( + "h_ood length {} != num_parts {num_parts}", + h_ood.len() + )); } if trace_ood_columns.len() != num_total_cols || trace_ood_columns.iter().any(|c| c.len() != num_eval_points) { - return None; + reject!(format!( + "trace OOD shape mismatch: columns={}, expected_columns={num_total_cols}, eval_points={num_eval_points}", + trace_ood_columns.len() + )); } if trace_terms_gammas.len() != num_total_cols || trace_terms_gammas .iter() .any(|c| c.len() != num_eval_points) { - return None; + reject!(format!( + "trace gamma shape mismatch: columns={}, expected_columns={num_total_cols}, eval_points={num_eval_points}", + trace_terms_gammas.len() + )); } let expected_inv_denoms = lde_size.checked_mul(1 + num_eval_points)?; // The fully-resident `(Some(parts), Some(dev_inv))` arm ignores the @@ -1527,22 +1585,35 @@ where // slicing an empty host buffer). let arm_needs_host_inv = !(parts_dev.is_some() && inv_denoms_dev.is_some()); if arm_needs_host_inv && inv_denoms_host.len() != expected_inv_denoms { - return None; + reject!(format!( + "host inv_denoms length {} != expected {expected_inv_denoms} (parts_dev={}, inv_dev={})", + inv_denoms_host.len(), + parts_dev.is_some(), + inv_denoms_dev.is_some() + )); } // Validate the host parts when we don't have a device handle, since // the math-cuda call will assert these. if parts_dev.is_none() { if parts_host.len() != num_parts { - return None; + reject!(format!( + "host parts length {} != num_parts {num_parts}", + parts_host.len() + )); } if parts_host.iter().any(|p| p.len() != lde_size) { - return None; + reject!(format!( + "a host part length differs from lde_size {lde_size}" + )); } } else if let Some(p) = parts_dev && (p.m != num_parts || p.lde_size != lde_size) { - return None; + reject!(format!( + "device parts shape m={}, lde_size={} != expected m={num_parts}, lde_size={lde_size}", + p.m, p.lde_size + )); } // Pack host buffers. SAFETY for ext3 transmutes: E == Ext3 by TypeId check. @@ -1577,6 +1648,16 @@ where // 2. Parts on device, inv_denoms on host. // 3. Both on host (fallback when R2 keep + denom-invert both missed). let parts_host_packed: Vec; + let arm = match (parts_dev.is_some(), inv_denoms_dev.is_some()) { + (true, true) => "device-parts+device-inv", + (true, false) => "device-parts+host-inv", + (false, _) => "host-parts+host-inv", + }; + if diag { + eprintln!( + "GPU_DEEP_DIAG attempt: arm={arm} lde_size={lde_size} main={num_main} aux={num_aux} parts={num_parts} eval_points={num_eval_points}" + ); + } let result = match (parts_dev, inv_denoms_dev) { (Some(parts), Some((inv_dev, stream))) => { math_cuda::deep::deep_composition_ext3_with_dev_parts_and_inv_denoms( @@ -1665,8 +1746,20 @@ where let deep_raw = match result { Ok(v) => v, - Err(_) => return None, + Err(e) => { + if diag { + eprintln!( + "GPU_DEEP_DIAG cuda-error: arm={arm} lde_size={lde_size} main={num_main} aux={num_aux} parts={num_parts} eval_points={num_eval_points}: {e:?}" + ); + } + return None; + } }; + if diag { + eprintln!( + "GPU_DEEP_DIAG success: arm={arm} lde_size={lde_size} main={num_main} aux={num_aux} parts={num_parts} eval_points={num_eval_points}" + ); + } GPU_DEEP_CALLS.fetch_add(1, Ordering::Relaxed); debug_assert_eq!(deep_raw.len(), lde_size * 3); Some(u64_to_ext3_vec::(&deep_raw)) diff --git a/crypto/stark/src/prover.rs b/crypto/stark/src/prover.rs index 972ab5bb6..1900e1e25 100644 --- a/crypto/stark/src/prover.rs +++ b/crypto/stark/src/prover.rs @@ -545,7 +545,9 @@ where pub(crate) fn host_evals(&self) -> &[Vec>] { self.host_composition_lde .as_ref() - .expect("host composition LDE present (device-only paths must read the resident handle)") + .expect( + "host composition LDE present (device-only paths must read the resident handle)", + ) .as_slice() } } @@ -1169,6 +1171,7 @@ pub trait IsStarkProver< domain: &Domain, twiddles: &LdeTwiddles, #[cfg(feature = "cuda")] gpu_parts_out: &mut Option, + #[cfg(feature = "cuda")] retain_host_lde: bool, ) -> Vec>> where FieldElement: AsBytes + Sync + Send, @@ -1206,7 +1209,7 @@ pub trait IsStarkProver< // `cuda` feature and a qualifying size. Falls through to CPU when not. #[cfg(feature = "cuda")] if let Some((lde_h0, lde_h1, handle)) = - crate::gpu_lde::try_extend_two_halves_gpu(&h0_evals, &h1_evals, domain) + crate::gpu_lde::try_extend_two_halves_gpu(&h0_evals, &h1_evals, domain, retain_host_lde) { // Hand the resident extended-LDE buffer to the session so the // composition Merkle commit and DEEP reuse it (mirrors the d>2 path). @@ -1286,6 +1289,23 @@ pub trait IsStarkProver< let number_of_parts = air.composition_poly_degree_bound(trace_length) / trace_length; + // Dropping the composition host LDE is safe only when DEEP has every + // trace input resident too. A composition handle alone is insufficient: + // some tables qualify for the R2 LDE but not for the R1 main/aux path. + #[cfg(feature = "cuda")] + let retain_host_composition_lde = { + let expected_lde_size = domain.lde_roots_of_unity_coset.len(); + let main_resident = round_1_result.lde_trace.gpu_main().is_some_and(|h| { + h.m == round_1_result.lde_trace.num_main_cols() && h.lde_size == expected_lde_size + }); + let aux_resident = round_1_result.lde_trace.num_aux_cols() == 0 + || round_1_result.lde_trace.gpu_aux().is_some_and(|h| { + h.m == round_1_result.lde_trace.num_aux_cols() + && h.lde_size == expected_lde_size + }); + crate::gpu_lde::retain_composition_host() || !main_resident || !aux_resident + }; + #[cfg(feature = "instruments")] let t_sub = Instant::now(); #[cfg(feature = "cuda")] @@ -1303,6 +1323,7 @@ pub trait IsStarkProver< domain, twiddles, &mut gpu_composition_parts, + retain_host_composition_lde, ) } #[cfg(not(feature = "cuda"))] @@ -1416,8 +1437,21 @@ pub trait IsStarkProver< #[cfg(feature = "instruments")] crate::instruments::store_r2_sub(constraints_dur, fft_dur, merkle_dur); + let num_composition_parts = lde_composition_poly_parts_evaluations.len(); + #[cfg(feature = "cuda")] + let composition_device_only = gpu_composition_parts.is_some() + && num_composition_parts > 0 + && lde_composition_poly_parts_evaluations + .iter() + .all(Vec::is_empty); + #[cfg(feature = "cuda")] + if composition_device_only { + crate::gpu_lde::GPU_COMPOSITION_DEVICE_ONLY_CALLS + .fetch_add(1, std::sync::atomic::Ordering::Relaxed); + } + // Fold the R2 device composition parts handle into the session (resident - // R2 to R4). The host evaluations stay in `Round2` for R4 openings. + // R2 to R4). #[cfg(feature = "cuda")] if let Some(handle) = gpu_composition_parts { // GPU_NO_RESIDENT_DEEP=1 withholds the handle so DEEP re-uploads @@ -1427,9 +1461,17 @@ pub trait IsStarkProver< } } - let num_composition_parts = lde_composition_poly_parts_evaluations.len(); + #[cfg(feature = "cuda")] + let host_composition_lde = if composition_device_only { + None + } else { + Some(lde_composition_poly_parts_evaluations) + }; + #[cfg(not(feature = "cuda"))] + let host_composition_lde = Some(lde_composition_poly_parts_evaluations); + Ok(Round2 { - host_composition_lde: Some(lde_composition_poly_parts_evaluations), + host_composition_lde, num_composition_parts, composition_poly_merkle_tree, composition_poly_root, @@ -1462,33 +1504,32 @@ pub trait IsStarkProver< 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 - .host_evals() - .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(); + let cpu_comp_ood = || -> Vec> { + round_2_result + .host_evals() + .iter() + .map(|lde_evals| { + 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() + }; - // Step C1: compute the same composition-parts OOD on the resident device - // handle and cross-check it against the CPU result above (host stays - // authoritative this stage). Proves the device path before C4 makes it - // the sole source and the host composition LDE is dropped. #[cfg(feature = "cuda")] - if let Some(handle) = round_1_result.lde_trace.gpu_composition_parts() { - if let Some(gpu_ood) = + let composition_poly_parts_ood_evaluation = round_1_result + .lde_trace + .gpu_composition_parts() + .and_then(|handle| { crate::gpu_lde::try_barycentric_ext3_on_comp_handle::( handle, blowup_factor, @@ -1499,21 +1540,10 @@ pub trait IsStarkProver< &comp_z_pow_n, &comp_inv_denoms, ) - { - assert_eq!( - gpu_ood.len(), - composition_poly_parts_ood_evaluation.len(), - "C1 composition OOD: GPU/CPU part count mismatch" - ); - for (j, (g, c)) in gpu_ood - .iter() - .zip(composition_poly_parts_ood_evaluation.iter()) - .enumerate() - { - assert_eq!(g, c, "C1 composition OOD: GPU != CPU at part {j}"); - } - } - } + }) + .unwrap_or_else(cpu_comp_ood); + #[cfg(not(feature = "cuda"))] + let composition_poly_parts_ood_evaluation = cpu_comp_ood(); // === Trace polynomials: barycentric evaluation via LDE === let trace_ood_evaluations = crate::trace::get_trace_evaluations_from_lde( @@ -1754,7 +1784,10 @@ pub trait IsStarkProver< crate::gpu_lde::try_deep_composition_gpu::( lde_trace, lde_trace.gpu_composition_parts(), - round_2_result.host_evals(), + round_2_result + .host_composition_lde + .as_deref() + .unwrap_or(&[]), h_ood, &trace_ood_columns, composition_poly_gammas, @@ -1790,7 +1823,10 @@ pub trait IsStarkProver< crate::gpu_lde::try_deep_composition_gpu::( lde_trace, lde_trace.gpu_composition_parts(), - round_2_result.host_evals(), + round_2_result + .host_composition_lde + .as_deref() + .unwrap_or(&[]), h_ood, &trace_ood_columns, composition_poly_gammas, @@ -1813,6 +1849,11 @@ pub trait IsStarkProver< !lde_trace.host_trace_empty(), "R4 DEEP composition fell back to the host trace, but it is device-only (empty)" ); + #[cfg(feature = "cuda")] + assert!( + round_2_result.host_composition_lde.is_some(), + "R4 DEEP composition fell back to the host loop, but the composition LDE is device-only" + ); // OOD column compression (Plonky3-style): precompute one value per eval point, // ood_compressed_k = Σ_j gamma[j][k] * ood[j][k]. @@ -2359,31 +2400,21 @@ pub trait IsStarkProver< #[cfg(feature = "cuda")] { if let Some(proofs) = &comp_dev_proofs { - let host = Self::open_composition_poly_with_proof( - proofs[qi].clone(), - round_2_result.host_evals(), - *index, - ); - // C2 cross-check: the values gathered off the resident - // handle must equal the host values (host authoritative). - if let Some(dev_vals) = comp_dev_values.as_ref() { - let num_parts_comp = - round_2_result.num_composition_parts; - let (even, odd) = Self::device_row_pair::( - dev_vals, - qi, - num_parts_comp, - ); - assert_eq!( - even, host.evaluations, - "C2 composition open: GPU != CPU (even) query {qi}" - ); - assert_eq!( - odd, host.evaluations_sym, - "C2 composition open: GPU != CPU (sym) query {qi}" - ); + match comp_dev_values.as_ref() { + Some(dev_vals) => { + let (even, odd) = Self::device_row_pair::( + dev_vals, + qi, + round_2_result.num_composition_parts, + ); + Self::open_polys_from_values(proofs[qi].clone(), even, odd) + } + None => Self::open_composition_poly_with_proof( + proofs[qi].clone(), + round_2_result.host_evals(), + *index, + ), } - host } else { Self::open_composition_poly( &round_2_result.composition_poly_merkle_tree, diff --git a/prover/tests/cuda_path_integration.rs b/prover/tests/cuda_path_integration.rs index b60cb3a34..fdbdd15f9 100644 --- a/prover/tests/cuda_path_integration.rs +++ b/prover/tests/cuda_path_integration.rs @@ -12,8 +12,9 @@ use lambda_vm_prover::test_utils::asm_elf_bytes; use lambda_vm_prover::{prove, verify}; use stark::gpu_lde::{ gpu_bary_calls, gpu_batch_invert_calls, gpu_comp_poly_tree_calls, gpu_composition_calls, - gpu_deep_calls, gpu_device_only_calls, gpu_extend_halves_calls, gpu_fri_calls, gpu_lde_calls, - gpu_logup_calls, gpu_opening_gather_calls, gpu_parts_lde_calls, reset_all_gpu_call_counters, + gpu_composition_device_only_calls, gpu_deep_calls, gpu_device_only_calls, + gpu_extend_halves_calls, gpu_fri_calls, gpu_lde_calls, gpu_logup_calls, + gpu_opening_gather_calls, gpu_parts_lde_calls, reset_all_gpu_call_counters, }; /// The R2 GPU composition-poly path (fused `H = z·Σβᵢ·Cᵢ + boundary`) fires and @@ -94,6 +95,10 @@ fn gpu_path_fires_end_to_end() { // DEEP fires once per table that took the R1 GPU path. assert!(gpu_deep_calls() > 0, "R4 GPU DEEP composition did not fire"); + assert!( + gpu_composition_device_only_calls() > 0, + "composition LDE device-only path did not fire" + ); // FRI commit fires once per table (commit_phase_from_evaluations). assert!(gpu_fri_calls() > 0, "R4 GPU FRI commit did not fire"); From b05e6792d266f0a7c7acd933039c51d560d6f6d0 Mon Sep 17 00:00:00 2001 From: jotabulacios Date: Wed, 22 Jul 2026 10:39:41 -0300 Subject: [PATCH 07/11] Retain FRI layer evals on GPU, gather openings --- crypto/math-cuda/src/fri.rs | 44 +++++++++++++--- crypto/stark/src/fri/fri_commitment.rs | 10 ++++ crypto/stark/src/fri/mod.rs | 2 +- crypto/stark/src/gpu_lde.rs | 70 ++++++++++++++++++++++---- 4 files changed, 106 insertions(+), 20 deletions(-) diff --git a/crypto/math-cuda/src/fri.rs b/crypto/math-cuda/src/fri.rs index 6bb22ff58..817d7e916 100644 --- a/crypto/math-cuda/src/fri.rs +++ b/crypto/math-cuda/src/fri.rs @@ -40,6 +40,17 @@ fn check_fault_injection() -> Result<()> { Ok(()) } +/// A FRI layer's folded evaluations kept resident on device. `buf` is the +/// layer's ext3 evals **interleaved** (`3 * len` u64, `[a0,b0,c0, a1,b1,c1, …]`), +/// and `len` is the number of ext3 evals in the layer (`n_out`), carried +/// explicitly so the query phase has the metadata even once the host `Vec` is +/// dropped (Step F2). Freed when the `buf` Arc drops. +#[derive(Clone)] +pub struct GpuFriEvals { + pub buf: Arc>, + pub len: usize, +} + /// 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. @@ -98,7 +109,8 @@ impl FriCommitState { pub fn fold_and_commit_layer( &mut self, zeta_raw: [u64; 3], - ) -> Result<(Vec, crate::lde::GpuMerkleTree)> { + retain_evals: bool, + ) -> Result<(Vec, crate::lde::GpuMerkleTree, Option)> { #[cfg(feature = "test-faults")] check_fault_injection()?; let be = backend()?; @@ -206,15 +218,28 @@ impl FriCommitState { // Sync and D2H. self.stream.synchronize()?; - // Layer evals: 3 * n_out u64 from the output buffer. - let layer_evals: Vec = if self.a_is_input { - let view = self.evals_b.slice(0..3 * n_out); - self.stream.clone_dtoh(&view)? + // The folded output buffer for this layer (3 * n_out ext3 u64). + let out_view = if self.a_is_input { + self.evals_b.slice(0..3 * n_out) } else { - let view = self.evals_a.slice(0..3 * n_out); - self.stream.clone_dtoh(&view)? + self.evals_a.slice(0..3 * n_out) }; + // Host copy — still authoritative until Step F2 routes openings to device. + let layer_evals: Vec = self.stream.clone_dtoh(&out_view)?; crate::stagebytes::add_fri_layer_d2h(layer_evals.len() * 8); + // F0: for committed (queried) layers, retain the folded evals on device in + // a fresh buffer (the ping-pong scratch `evals_a`/`evals_b` is overwritten + // by later folds) so the query phase can gather opened values on device. + // The terminal fold is never queried, so `retain_evals` is false there and + // this D2D copy is skipped. + let retained_evals = if retain_evals { + Some(GpuFriEvals { + buf: Arc::new(self.stream.clone_dtod(&out_view)?), + len: n_out, + }) + } 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. @@ -231,6 +256,9 @@ impl FriCommitState { leaves_len: num_leaves, root, }; - Ok((layer_evals, tree)) + Ok((layer_evals, tree, retained_evals)) } + + // (retained_evals is `Some` only when `retain_evals` is passed for a + // committed layer; the terminal fold passes false and gets `None`.) } diff --git a/crypto/stark/src/fri/fri_commitment.rs b/crypto/stark/src/fri/fri_commitment.rs index 58c9eed77..56902e39a 100644 --- a/crypto/stark/src/fri/fri_commitment.rs +++ b/crypto/stark/src/fri/fri_commitment.rs @@ -18,6 +18,14 @@ where /// `merkle_tree` is a root only placeholder. `None` on the CPU path. #[cfg(feature = "cuda")] pub gpu_tree: Option, + /// The layer's folded evaluations kept resident on device (ext3 interleaved, + /// `3 * len` u64), so the query phase can gather the opened + /// `evaluation[value_pos]` values on device. `Some` on the GPU commit path + /// (during F1 it coexists with the still-authoritative host `evaluation`; + /// from F2 it becomes the sole source and `evaluation` is dropped); `None` on + /// the CPU path. + #[cfg(feature = "cuda")] + pub gpu_evals: Option, } impl FriLayer @@ -32,6 +40,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 8f1172524..86bd3e111 100644 --- a/crypto/stark/src/fri/mod.rs +++ b/crypto/stark/src/fri/mod.rs @@ -145,7 +145,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 01f1daa8c..3a9e0c2a6 100644 --- a/crypto/stark/src/gpu_lde.rs +++ b/crypto/stark/src/gpu_lde.rs @@ -2079,22 +2079,25 @@ 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, dev_tree, dev_evals) = + match state.fold_and_commit_layer(zeta_raw, true) { + 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`. + // paths from it via `gather_proofs_dev`. The folded evals also stay + // resident in `gpu_evals` (Step F) so opened values are gathered on device. let evaluation = u64_to_ext3_vec::(&layer_evals_u64); 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); + layer.gpu_evals = dev_evals; fri_layer_list.push(layer); // >>>> Send commitment: [p_k] @@ -2108,7 +2111,8 @@ where 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, _tree, _dev_evals) = match state.fold_and_commit_layer(zeta_raw, false) + { Ok(v) => v, Err(_) => { *transcript = transcript_snapshot; @@ -2148,7 +2152,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() { @@ -2190,6 +2194,9 @@ where } // Reassemble per-query decommitments, matching the host walk's order. + // For each (query, layer) the opened symmetric value is at + // value_pos = (iota >> l) ^ 1 + // (distinct from the proof position `(iota >> l) >> 1` gathered above). let decommits = iotas .iter() .enumerate() @@ -2198,7 +2205,34 @@ 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 value_pos = index ^ 1; + let host_val = layer.evaluation[value_pos].clone(); + // F1 cross-check: gather the same value off the resident device + // evals (ext3 interleaved: 3 u64 at `3 * value_pos`) and assert it + // equals the host value. Host stays authoritative this stage; F2 + // makes the device gather the sole source and drops the host Vec. + // Residency is all-or-nothing (asserted above), so a resident layer + // MUST carry its evals — `expect`, not a silent skip, or a missing + // handle would pass the test on the host path. + let dev = layer + .gpu_evals + .as_ref() + .expect("resident FRI layer missing gpu_evals"); + assert!( + value_pos < dev.len, + "FRI value_pos {value_pos} out of layer len {}", + dev.len + ); + let base = value_pos * 3; + let raw: Vec = stream + .clone_dtoh(&dev.buf.slice(base..base + 3)) + .expect("FRI device eval gather (resident, no host fallback)"); + let dev_val = fri_ext3_from_raw::(&raw); + assert_eq!( + dev_val, host_val, + "F1 FRI eval: GPU != host, layer {l} value_pos {value_pos}" + ); + layers_evaluations_sym.push(host_val); layers_auth_paths.push(per_layer_proofs[l][q].clone()); index >>= 1; } @@ -2210,3 +2244,17 @@ where .collect(); Some(decommits) } + +/// Build one ext3 `FieldElement` from 3 interleaved u64 limbs `[a, b, c]`. +/// `E` must be the degree-3 Goldilocks extension (holds on the GPU FRI path). +#[cfg(feature = "cuda")] +fn fri_ext3_from_raw(raw: &[u64]) -> FieldElement +where + E: IsField + 'static, +{ + debug_assert_eq!(raw.len(), 3); + u64_to_ext3_vec::(raw) + .into_iter() + .next() + .expect("one ext3 element from 3 limbs") +} From 47d08a85b7d1d34ebec43c7897f828bb1f568955 Mon Sep 17 00:00:00 2001 From: jotabulacios Date: Wed, 22 Jul 2026 11:25:33 -0300 Subject: [PATCH 08/11] Drop FRI host D2H; gather openings on device --- crypto/math-cuda/src/fri.rs | 15 ++++- crypto/stark/src/gpu_lde.rs | 126 +++++++++++++++++++++--------------- 2 files changed, 85 insertions(+), 56 deletions(-) diff --git a/crypto/math-cuda/src/fri.rs b/crypto/math-cuda/src/fri.rs index 817d7e916..4675f10e3 100644 --- a/crypto/math-cuda/src/fri.rs +++ b/crypto/math-cuda/src/fri.rs @@ -110,6 +110,7 @@ impl FriCommitState { &mut self, zeta_raw: [u64; 3], retain_evals: bool, + retain_host: bool, ) -> Result<(Vec, crate::lde::GpuMerkleTree, Option)> { #[cfg(feature = "test-faults")] check_fault_injection()?; @@ -224,9 +225,17 @@ impl FriCommitState { } else { self.evals_a.slice(0..3 * n_out) }; - // Host copy — still authoritative until Step F2 routes openings to device. - let layer_evals: Vec = self.stream.clone_dtoh(&out_view)?; - crate::stagebytes::add_fri_layer_d2h(layer_evals.len() * 8); + // Host copy of the folded evals. Under device-only (F2) `retain_host` is + // false for committed layers and this big D2H is skipped (query openings + // are gathered on device); kept when `retain_host` (GPU_RETAIN_FRI_HOST, + // and always the terminal fold whose evals become the final-poly coeffs). + let layer_evals: Vec = if retain_host { + let v = self.stream.clone_dtoh(&out_view)?; + crate::stagebytes::add_fri_layer_d2h(v.len() * 8); + v + } else { + Vec::new() + }; // F0: for committed (queried) layers, retain the folded evals on device in // a fresh buffer (the ping-pong scratch `evals_a`/`evals_b` is overwritten // by later folds) so the query phase can gather opened values on device. diff --git a/crypto/stark/src/gpu_lde.rs b/crypto/stark/src/gpu_lde.rs index 3a9e0c2a6..f4d1e5d87 100644 --- a/crypto/stark/src/gpu_lde.rs +++ b/crypto/stark/src/gpu_lde.rs @@ -1008,6 +1008,20 @@ pub fn retain_composition_host() -> bool { explicit || deep_resident_disabled() || merkle_resident_disabled() } +/// Whether to keep the host copy of the FRI layer evaluations (Step F2). Default +/// drops it (device-only): the query phase gathers opened values off the resident +/// per-layer eval buffers. `GPU_RETAIN_FRI_HOST=1` keeps the host copy (the F0+F1 +/// behavior, and the ABBA baseline). Committed layers honor this; the terminal +/// fold always retains (its evals become the final-poly coefficients). +pub fn retain_fri_host() -> bool { + static V: std::sync::OnceLock = std::sync::OnceLock::new(); + *V.get_or_init(|| { + std::env::var("GPU_RETAIN_FRI_HOST") + .map(|v| v == "1") + .unwrap_or(false) + }) +} + /// R2 GPU dispatch: build the composition-polynomial Merkle tree directly from /// the device-resident extended-LDE handle (the `_keep` extend buffer), skipping /// the host re-upload that [`try_build_comp_poly_tree_gpu`] performs. Same slab @@ -2080,7 +2094,7 @@ where let zeta_raw: [u64; 3] = unsafe { [*zeta_ptr, *zeta_ptr.add(1), *zeta_ptr.add(2)] }; let (layer_evals_u64, dev_tree, dev_evals) = - match state.fold_and_commit_layer(zeta_raw, true) { + match state.fold_and_commit_layer(zeta_raw, true, retain_fri_host()) { Ok(v) => v, Err(_) => { *transcript = transcript_snapshot.clone(); @@ -2111,14 +2125,14 @@ where 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, _dev_evals) = match state.fold_and_commit_layer(zeta_raw, false) - { - Ok(v) => v, - Err(_) => { - *transcript = transcript_snapshot; - return None; - } - }; + let (terminal_evals_u64, _tree, _dev_evals) = + match state.fold_and_commit_layer(zeta_raw, false, true) { + Ok(v) => v, + Err(_) => { + *transcript = transcript_snapshot; + return None; + } + }; debug_assert_eq!(terminal_evals_u64.len(), layout.terminal_len * 3); let terminal_codeword = u64_to_ext3_vec::(&terminal_evals_u64); @@ -2179,18 +2193,54 @@ where .next_stream(); let num_layers = fri_layers.len(); - // Batched gather: one call per layer over all queries. + // Batched gather per layer over all queries: + // * auth paths from the resident tree (proof_pos = (iota >> l) >> 1) + // * opened values from the resident evals (value_pos = (iota >> l) ^ 1), + // scattered into one contiguous scratch and D2H'd once per layer (a + // single small D2H of `queries * 3` u64 — NOT the full-layer D2H). let mut per_layer_proofs: Vec>> = Vec::with_capacity(num_layers); + let mut per_layer_values: Vec>> = Vec::with_capacity(num_layers); for (l, layer) in fri_layers.iter().enumerate() { let tree = layer .gpu_tree .as_ref() .expect("FRI layers are device-resident as a group"); - let positions: Vec = iotas.iter().map(|&iota| (iota >> l) >> 1).collect(); + let proof_positions: Vec = iotas.iter().map(|&iota| (iota >> l) >> 1).collect(); per_layer_proofs.push( - gather_proofs_dev(tree, &positions, &stream) + gather_proofs_dev(tree, &proof_positions, &stream) .expect("device FRI-layer gather failed; resident tree has no host fallback"), ); + + let dev = layer + .gpu_evals + .as_ref() + .expect("resident FRI layer missing gpu_evals"); + let nq = iotas.len(); + let mut scratch = stream + .alloc_zeros::(nq * 3) + .expect("FRI value-gather scratch alloc"); + for (q, &iota) in iotas.iter().enumerate() { + let value_pos = (iota >> l) ^ 1; + assert!( + value_pos < dev.len, + "FRI value_pos {value_pos} out of layer len {}", + dev.len + ); + let base = value_pos * 3; + stream + .memcpy_dtod( + &dev.buf.slice(base..base + 3), + &mut scratch.slice_mut(q * 3..q * 3 + 3), + ) + .expect("FRI value-gather D2D scatter"); + } + let raw: Vec = stream + .clone_dtoh(&scratch) + .expect("FRI value-gather D2H (resident, no host fallback)"); + // Count the batched query-value gather (it replaces the full-layer D2H); + // otherwise the FRI stage-byte counter would omit it and look like ~0. + math_cuda::stagebytes::add_query_gather(raw.len() * 8); + per_layer_values.push(u64_to_ext3_vec::(&raw)); } // Reassemble per-query decommitments, matching the host walk's order. @@ -2206,33 +2256,17 @@ where let mut index = iota; for (l, layer) in fri_layers.iter().enumerate() { let value_pos = index ^ 1; - let host_val = layer.evaluation[value_pos].clone(); - // F1 cross-check: gather the same value off the resident device - // evals (ext3 interleaved: 3 u64 at `3 * value_pos`) and assert it - // equals the host value. Host stays authoritative this stage; F2 - // makes the device gather the sole source and drops the host Vec. - // Residency is all-or-nothing (asserted above), so a resident layer - // MUST carry its evals — `expect`, not a silent skip, or a missing - // handle would pass the test on the host path. - let dev = layer - .gpu_evals - .as_ref() - .expect("resident FRI layer missing gpu_evals"); - assert!( - value_pos < dev.len, - "FRI value_pos {value_pos} out of layer len {}", - dev.len - ); - let base = value_pos * 3; - let raw: Vec = stream - .clone_dtoh(&dev.buf.slice(base..base + 3)) - .expect("FRI device eval gather (resident, no host fallback)"); - let dev_val = fri_ext3_from_raw::(&raw); - assert_eq!( - dev_val, host_val, - "F1 FRI eval: GPU != host, layer {l} value_pos {value_pos}" - ); - layers_evaluations_sym.push(host_val); + // Device-gathered opened value (authoritative; F2 drops the host + // Vec). When the host evals are still retained + // (`GPU_RETAIN_FRI_HOST`), cross-check the device value against them. + let dev_val = per_layer_values[l][q].clone(); + if !layer.evaluation.is_empty() { + assert_eq!( + dev_val, layer.evaluation[value_pos], + "FRI eval: GPU != host, layer {l} value_pos {value_pos}" + ); + } + layers_evaluations_sym.push(dev_val); layers_auth_paths.push(per_layer_proofs[l][q].clone()); index >>= 1; } @@ -2244,17 +2278,3 @@ where .collect(); Some(decommits) } - -/// Build one ext3 `FieldElement` from 3 interleaved u64 limbs `[a, b, c]`. -/// `E` must be the degree-3 Goldilocks extension (holds on the GPU FRI path). -#[cfg(feature = "cuda")] -fn fri_ext3_from_raw(raw: &[u64]) -> FieldElement -where - E: IsField + 'static, -{ - debug_assert_eq!(raw.len(), 3); - u64_to_ext3_vec::(raw) - .into_iter() - .next() - .expect("one ext3 element from 3 limbs") -} From 1ec452300f79882573faede251ee236f1563cd9d Mon Sep 17 00:00:00 2001 From: jotabulacios Date: Wed, 22 Jul 2026 17:23:49 -0300 Subject: [PATCH 09/11] Direct GPU DEEP-to-FRI resident handoff --- crypto/math-cuda/kernels/fri.cu | 34 ++++ crypto/math-cuda/src/deep.rs | 80 ++++++++- crypto/math-cuda/src/device.rs | 2 + crypto/math-cuda/src/fri.rs | 78 ++++++-- crypto/math-cuda/src/stagebytes.rs | 27 ++- crypto/math-cuda/tests/fri_deep_handoff.rs | 104 +++++++++++ crypto/stark/src/fri/mod.rs | 1 + crypto/stark/src/gpu_lde.rs | 76 ++++++-- crypto/stark/src/instruments.rs | 39 +++- crypto/stark/src/prover.rs | 200 +++++++++++++++++++-- prover/src/instruments.rs | 22 ++- 11 files changed, 605 insertions(+), 58 deletions(-) create mode 100644 crypto/math-cuda/tests/fri_deep_handoff.rs diff --git a/crypto/math-cuda/kernels/fri.cu b/crypto/math-cuda/kernels/fri.cu index 63d72cef1..a582ad00d 100644 --- a/crypto/math-cuda/kernels/fri.cu +++ b/crypto/math-cuda/kernels/fri.cu @@ -47,6 +47,40 @@ extern "C" __global__ void fri_fold_ext3( out_p[2] = res.c; } +// First DEEP->FRI fold when the DEEP producer hands over its natural-order +// codeword directly. The legacy host bridge bit-reverses the full codeword and +// then pairs entries (2*j, 2*j+1). Bit reversal is an involution, so reading +// natural input at br(2*j) / br(2*j+1) is byte-identical while avoiding both +// the CPU permutation and a separate device permutation pass. Output j is in +// the ordinary FRI layer order, so every later fold uses fri_fold_ext3. +extern "C" __global__ void fri_fold_ext3_from_natural( + const uint64_t *__restrict__ in, + uint64_t n_out, + uint64_t log_n_in, + const uint64_t *__restrict__ inv_tw, + const uint64_t *__restrict__ zeta, + uint64_t *__restrict__ out) { + uint64_t j = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x; + if (j >= n_out) return; + + uint64_t br_lo = __brevll(2 * j) >> (64 - log_n_in); + uint64_t br_hi = __brevll(2 * j + 1) >> (64 - log_n_in); + const uint64_t *lo_p = in + br_lo * 3; + const uint64_t *hi_p = in + br_hi * 3; + + ext3::Fe3 lo = ext3::make(lo_p[0], lo_p[1], lo_p[2]); + ext3::Fe3 hi = ext3::make(hi_p[0], hi_p[1], hi_p[2]); + ext3::Fe3 sum = ext3::add(lo, hi); + ext3::Fe3 diff = ext3::sub(lo, hi); + ext3::Fe3 z = ext3::make(zeta[0], zeta[1], zeta[2]); + ext3::Fe3 res = ext3::add(sum, ext3::mul_base(ext3::mul(z, diff), inv_tw[j])); + + uint64_t *out_p = out + j * 3; + out_p[0] = res.a; + out_p[1] = res.b; + out_p[2] = res.c; +} + // update_twiddles: tw_out[j] = tw_in[2j]^2 for j in 0..n_out. // Separate input/output buffers: thread j reads tw_in[2j] while thread 2j // writes tw_out[2j], so an in-place version would race across threads. diff --git a/crypto/math-cuda/src/deep.rs b/crypto/math-cuda/src/deep.rs index 79df6b5d9..182e446bc 100644 --- a/crypto/math-cuda/src/deep.rs +++ b/crypto/math-cuda/src/deep.rs @@ -16,6 +16,21 @@ use crate::Result; use crate::device::backend; use crate::lde::{GpuLdeBase, GpuLdeExt3}; +/// Natural-order DEEP codeword kept on the producer stream for direct FRI +/// consumption. `len` counts ext3 elements; `buf` contains `3 * len` u64s. +pub struct GpuDeepEvals { + pub buf: CudaSlice, + pub len: usize, + pub stream: Arc, +} + +/// Result of the retaining DEEP entry point. `host` is empty when the caller +/// elects device-only handoff; the device handle is always populated. +pub struct DeepCompositionKeep { + pub host: Vec, + pub device: GpuDeepEvals, +} + /// Compute deep-composition evaluations on device. /// /// `num_eval_points = trace_terms_gammas_interleaved.len() / ((num_main + @@ -143,6 +158,51 @@ pub fn deep_composition_ext3_with_dev_parts_and_inv_denoms( row_stride: usize, domain_size: usize, ) -> Result> { + Ok(deep_composition_ext3_with_dev_parts_and_inv_denoms_keep( + stream, + main_lde, + aux_lde, + h_parts_dev, + inv_denoms_dev, + h_ood, + trace_ood, + gammas_h, + gammas_tr, + num_parts, + num_main, + num_aux, + num_eval_points, + row_stride, + domain_size, + true, + )? + .host) +} + +/// Retaining variant of +/// [`deep_composition_ext3_with_dev_parts_and_inv_denoms`]. When +/// `retain_host` is false it does not synchronize or copy the codeword to the +/// host; the returned handle remains bound to the producer stream so FRI can +/// consume it in FIFO order. +#[allow(clippy::too_many_arguments)] +pub fn deep_composition_ext3_with_dev_parts_and_inv_denoms_keep( + stream: &Arc, + main_lde: &GpuLdeBase, + aux_lde: Option<&GpuLdeExt3>, + h_parts_dev: &GpuLdeExt3, + inv_denoms_dev: &CudaSlice, + h_ood: &[u64], + trace_ood: &[u64], + gammas_h: &[u64], + gammas_tr: &[u64], + num_parts: usize, + num_main: usize, + num_aux: usize, + num_eval_points: usize, + row_stride: usize, + domain_size: usize, + retain_host: bool, +) -> Result { assert_eq!(main_lde.m, num_main); assert_eq!(h_parts_dev.m, num_parts); assert_eq!(h_parts_dev.lde_size, main_lde.lde_size); @@ -243,9 +303,22 @@ pub fn deep_composition_ext3_with_dev_parts_and_inv_denoms( .launch(cfg)?; } - let out = stream.clone_dtoh(&deep_out)?; - stream.synchronize()?; - Ok(out) + let host = if retain_host { + let out = stream.clone_dtoh(&deep_out)?; + crate::stagebytes::add_deep_out_d2h(core::mem::size_of_val(out.as_slice())); + stream.synchronize()?; + out + } else { + Vec::new() + }; + Ok(DeepCompositionKeep { + host, + device: GpuDeepEvals { + buf: deep_out, + len: domain_size, + stream: Arc::clone(stream), + }, + }) } #[allow(clippy::too_many_arguments)] @@ -370,6 +443,7 @@ fn deep_composition_ext3_impl( } let out = stream.clone_dtoh(&deep_out)?; + crate::stagebytes::add_deep_out_d2h(core::mem::size_of_val(out.as_slice())); stream.synchronize()?; Ok(out) } diff --git a/crypto/math-cuda/src/device.rs b/crypto/math-cuda/src/device.rs index 2bebe2cc0..1cd4140e5 100644 --- a/crypto/math-cuda/src/device.rs +++ b/crypto/math-cuda/src/device.rs @@ -180,6 +180,7 @@ pub struct Backend { // fri.cubin pub fri_fold_ext3: CudaFunction, + pub fri_fold_ext3_from_natural: CudaFunction, pub fri_update_twiddles: CudaFunction, // inverse.cubin @@ -379,6 +380,7 @@ impl Backend { gather_rows_ext3: bary.load_function("gather_rows_ext3")?, deep_composition_ext3_row: deep.load_function("deep_composition_ext3_row")?, fri_fold_ext3: fri.load_function("fri_fold_ext3")?, + fri_fold_ext3_from_natural: fri.load_function("fri_fold_ext3_from_natural")?, fri_update_twiddles: fri.load_function("fri_update_twiddles")?, compute_denoms_ext3: inverse.load_function("compute_denoms_ext3")?, block_inclusive_scan_fwd_ext3: inverse diff --git a/crypto/math-cuda/src/fri.rs b/crypto/math-cuda/src/fri.rs index 4675f10e3..2c2ca3d61 100644 --- a/crypto/math-cuda/src/fri.rs +++ b/crypto/math-cuda/src/fri.rs @@ -3,8 +3,9 @@ //! The host loop (in the stark crate) samples each layer's `zeta` from the //! transcript and feeds it in. This module keeps the folded evaluations, //! twiddles, and per-layer Merkle trees on device, only D2H'ing each -//! layer's root (to append to the transcript), plus its full evals and -//! tree nodes (to plug into `FriLayer` for the query phase). +//! layer's root (to append to the transcript). Folded evals and Merkle nodes +//! remain resident for the query phase; the terminal codeword alone returns +//! to the host for coefficient extraction. //! //! Mirrors `commit_phase_from_evaluations` at //! `crypto/stark/src/fri/mod.rs`. @@ -68,6 +69,9 @@ pub struct FriCommitState { pub current_n: usize, /// Which buffer holds the current layer's input. Toggles each fold. a_is_input: bool, + /// The direct DEEP handoff arrives in natural LDE order. Only the first + /// fold needs bit-reversed reads; its output has the ordinary FRI layout. + first_input_natural: bool, } impl FriCommitState { @@ -87,6 +91,7 @@ impl FriCommitState { 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)?; + crate::stagebytes::add_fri_initial_h2d(core::mem::size_of_val(evals_host)); let inv_tw = stream.clone_htod(inv_tw_host)?; Ok(Self { @@ -96,6 +101,41 @@ impl FriCommitState { inv_tw, current_n: n0, a_is_input: true, + first_input_natural: false, + }) + } + + /// Start FRI by taking ownership of a natural-order DEEP codeword already + /// resident on device. No D2H, CPU bit-reverse, H2D, or D2D staging pass is + /// performed. The first fold applies the legacy bit-reverse mapping while + /// loading; subsequent folds use the normal contiguous-pair kernel. + pub fn from_deep( + deep: crate::deep::GpuDeepEvals, + inv_tw_host: &[u64], + n0: usize, + ) -> Result { + assert!(n0 >= 2 && n0.is_power_of_two()); + assert_eq!(deep.len, n0); + assert_eq!(deep.buf.len(), 3 * n0); + assert_eq!(inv_tw_host.len(), n0 / 2); + + let crate::deep::GpuDeepEvals { + buf: evals_a, + len: _, + stream, + } = deep; + // SAFETY: the first fold writes every output slot before evals_b is read. + let evals_b = unsafe { stream.alloc::(3 * n0) }?; + let inv_tw = stream.clone_htod(inv_tw_host)?; + + Ok(Self { + stream, + evals_a, + evals_b, + inv_tw, + current_n: n0, + a_is_input: true, + first_input_natural: true, }) } @@ -147,15 +187,31 @@ impl FriCommitState { } else { (&self.evals_b, &mut self.evals_a) }; - unsafe { - self.stream - .launch_builder(&be.fri_fold_ext3) - .arg(input_evals) - .arg(&n_out_u64) - .arg(&self.inv_tw) - .arg(&zeta_dev) - .arg(output_evals) - .launch(cfg)?; + if self.first_input_natural { + let log_n_in = n_in.trailing_zeros() as u64; + unsafe { + self.stream + .launch_builder(&be.fri_fold_ext3_from_natural) + .arg(input_evals) + .arg(&n_out_u64) + .arg(&log_n_in) + .arg(&self.inv_tw) + .arg(&zeta_dev) + .arg(output_evals) + .launch(cfg)?; + } + self.first_input_natural = false; + } else { + unsafe { + self.stream + .launch_builder(&be.fri_fold_ext3) + .arg(input_evals) + .arg(&n_out_u64) + .arg(&self.inv_tw) + .arg(&zeta_dev) + .arg(output_evals) + .launch(cfg)?; + } } // SAFETY: keccak_fri_leaves_ext3 writes the leaves [num_leaves-1, 2*num_leaves-1) diff --git a/crypto/math-cuda/src/stagebytes.rs b/crypto/math-cuda/src/stagebytes.rs index 2b438f726..3e29d2958 100644 --- a/crypto/math-cuda/src/stagebytes.rs +++ b/crypto/math-cuda/src/stagebytes.rs @@ -13,6 +13,8 @@ //! * `comp_h01_lde_d2h`— D2H of the extended H₀/H₁ LDE result back to host //! * `comp_merkle_h2d` — H2D re-upload of those parts for the Merkle commit //! * `comp_deep_h2d` — H2D re-upload of composition parts for DEEP fallback +//! * `deep_out_d2h` — D2H of the completed DEEP codeword before FRI +//! * `fri_initial_h2d` — H2D of that codeword when GPU FRI starts //! * `fri_layer_d2h` — D2H of every FRI layer's evaluations //! * `query_gather` — D2H of Merkle paths during query openings @@ -24,6 +26,8 @@ static COMP_H01_H2D: AtomicU64 = AtomicU64::new(0); static COMP_H01_LDE_D2H: AtomicU64 = AtomicU64::new(0); static COMP_MERKLE_H2D: AtomicU64 = AtomicU64::new(0); static COMP_DEEP_H2D: AtomicU64 = AtomicU64::new(0); +static DEEP_OUT_D2H: AtomicU64 = AtomicU64::new(0); +static FRI_INITIAL_H2D: AtomicU64 = AtomicU64::new(0); static FRI_LAYER_D2H: AtomicU64 = AtomicU64::new(0); static QUERY_GATHER: AtomicU64 = AtomicU64::new(0); @@ -59,6 +63,12 @@ pub fn add_comp_merkle_h2d(bytes: usize) { pub fn add_comp_deep_h2d(bytes: usize) { add(&COMP_DEEP_H2D, bytes); } +pub fn add_deep_out_d2h(bytes: usize) { + add(&DEEP_OUT_D2H, bytes); +} +pub fn add_fri_initial_h2d(bytes: usize) { + add(&FRI_INITIAL_H2D, bytes); +} pub fn add_fri_layer_d2h(bytes: usize) { add(&FRI_LAYER_D2H, bytes); } @@ -74,6 +84,8 @@ pub fn reset() { &COMP_H01_LDE_D2H, &COMP_MERKLE_H2D, &COMP_DEEP_H2D, + &DEEP_OUT_D2H, + &FRI_INITIAL_H2D, &FRI_LAYER_D2H, &QUERY_GATHER, ] { @@ -93,6 +105,7 @@ pub fn report() -> Option { + mb(&COMP_H01_LDE_D2H) + mb(&COMP_MERKLE_H2D) + mb(&COMP_DEEP_H2D); + let deep_fri_bridge = mb(&DEEP_OUT_D2H) + mb(&FRI_INITIAL_H2D); let fri = mb(&FRI_LAYER_D2H); let q = mb(&QUERY_GATHER); let mut s = String::from("GPU stage bytes (host<->device, MB):\n"); @@ -117,11 +130,23 @@ pub fn report() -> Option { mb(&COMP_DEEP_H2D) )); s.push_str(&format!(" composition SUBTOTAL {:>10.1}\n", comp)); + s.push_str(&format!( + " DEEP output D2H {:>10.1}\n", + mb(&DEEP_OUT_D2H) + )); + s.push_str(&format!( + " FRI initial codeword H2D {:>10.1}\n", + mb(&FRI_INITIAL_H2D) + )); + s.push_str(&format!( + " DEEP->FRI bridge TOTAL {:>10.1}\n", + deep_fri_bridge + )); s.push_str(&format!(" FRI layer evals D2H {:>10.1}\n", fri)); s.push_str(&format!(" query gather D2H {:>10.1}\n", q)); s.push_str(&format!( " TOTAL (counted) {:>10.1}\n", - comp + fri + q + comp + deep_fri_bridge + fri + q )); Some(s) } diff --git a/crypto/math-cuda/tests/fri_deep_handoff.rs b/crypto/math-cuda/tests/fri_deep_handoff.rs new file mode 100644 index 000000000..064eff77c --- /dev/null +++ b/crypto/math-cuda/tests/fri_deep_handoff.rs @@ -0,0 +1,104 @@ +//! Parity for the direct natural-order DEEP -> first FRI fold handoff. + +use math::fft::bit_reversing::in_place_bit_reverse_permute; +use math::field::element::FieldElement; +use math::field::extensions_goldilocks::Degree3GoldilocksExtensionField; +use math::field::goldilocks::GoldilocksField; +use math::field::traits::IsPrimeField; +use math_cuda::deep::GpuDeepEvals; +use math_cuda::device::backend; +use math_cuda::fri::FriCommitState; +use rand::{Rng, SeedableRng}; +use rand_chacha::ChaCha8Rng; + +type Fp = FieldElement; +type Fp3 = FieldElement; + +fn rand_fp(rng: &mut ChaCha8Rng) -> Fp { + Fp::from_raw(rng.r#gen::()) +} + +fn rand_fp3(rng: &mut ChaCha8Rng) -> Fp3 { + Fp3::new([rand_fp(rng), rand_fp(rng), rand_fp(rng)]) +} + +fn raw3(e: &Fp3) -> [u64; 3] { + [ + *e.value()[0].value(), + *e.value()[1].value(), + *e.value()[2].value(), + ] +} + +fn canonical3(e: &Fp3) -> [u64; 3] { + [ + GoldilocksField::canonical(e.value()[0].value()), + GoldilocksField::canonical(e.value()[1].value()), + GoldilocksField::canonical(e.value()[2].value()), + ] +} + +fn run_parity(log_n: u32, seed: u64) { + let n = 1usize << log_n; + let mut rng = ChaCha8Rng::seed_from_u64(seed); + let natural: Vec = (0..n).map(|_| rand_fp3(&mut rng)).collect(); + let inv_tw: Vec = (0..n / 2).map(|_| rand_fp(&mut rng)).collect(); + let zeta = rand_fp3(&mut rng); + + let mut expected = natural.clone(); + in_place_bit_reverse_permute(&mut expected); + for j in 0..n / 2 { + let lo = &expected[2 * j]; + let hi = &expected[2 * j + 1]; + let sum = lo + hi; + let diff = lo - hi; + expected[j] = &sum + &(&inv_tw[j] * &(&zeta * &diff)); + } + expected.truncate(n / 2); + + let mut natural_raw = Vec::with_capacity(3 * n); + for e in &natural { + natural_raw.extend_from_slice(&raw3(e)); + } + let inv_tw_raw: Vec = inv_tw.iter().map(|x| *x.value()).collect(); + + let be = backend().unwrap(); + let stream = be.next_stream(); + let deep_buf = stream.clone_htod(&natural_raw).unwrap(); + stream.synchronize().unwrap(); + let deep = GpuDeepEvals { + buf: deep_buf, + len: n, + stream, + }; + let mut state = FriCommitState::from_deep(deep, &inv_tw_raw, n).unwrap(); + let (got_raw, _tree, _retained) = state + .fold_and_commit_layer(raw3(&zeta), false, true) + .unwrap(); + + assert_eq!(got_raw.len(), expected.len() * 3); + for (i, want) in expected.iter().enumerate() { + let got = Fp3::new([ + Fp::from_raw(got_raw[i * 3]), + Fp::from_raw(got_raw[i * 3 + 1]), + Fp::from_raw(got_raw[i * 3 + 2]), + ]); + assert_eq!( + canonical3(&got), + canonical3(want), + "first direct fold mismatch at log_n={log_n}, row={i}" + ); + } +} + +#[test] +fn direct_deep_first_fold_small() { + for log_n in 3..=10 { + run_parity(log_n, 0xD33F_0000 + log_n as u64); + } +} + +#[test] +fn direct_deep_first_fold_medium() { + run_parity(16, 0xD33F_F001); +} diff --git a/crypto/stark/src/fri/mod.rs b/crypto/stark/src/fri/mod.rs index 86bd3e111..55e3686d0 100644 --- a/crypto/stark/src/fri/mod.rs +++ b/crypto/stark/src/fri/mod.rs @@ -62,6 +62,7 @@ where // the CPU path below then runs as if the GPU had never been tried. if let Some(result) = crate::gpu_lde::try_fri_commit_gpu::( &evals, + None, transcript, coset_offset, domain_size, diff --git a/crypto/stark/src/gpu_lde.rs b/crypto/stark/src/gpu_lde.rs index f4d1e5d87..0733fc03a 100644 --- a/crypto/stark/src/gpu_lde.rs +++ b/crypto/stark/src/gpu_lde.rs @@ -1022,6 +1022,31 @@ pub fn retain_fri_host() -> bool { }) } +/// Enable the direct DEEP->FRI device handoff. It is the production default; +/// `GPU_DEEP_FRI_RESIDENT=0` restores the legacy D2H -> CPU bit-reverse -> H2D +/// bridge and is retained as an A/B control. +pub fn deep_fri_resident_enabled() -> bool { + static V: std::sync::OnceLock = std::sync::OnceLock::new(); + *V.get_or_init(|| { + std::env::var("GPU_DEEP_FRI_RESIDENT") + .map(|v| v != "0") + .unwrap_or(true) + }) +} + +/// Run the legacy host bridge against a cloned transcript and require roots, +/// terminal coefficients, and transcript state to match the direct handoff. +/// Disabled by default after parity and end-to-end validation; set +/// `GPU_DEEP_FRI_CROSSCHECK=1` to run both paths and compare them. +pub fn deep_fri_crosscheck() -> bool { + static V: std::sync::OnceLock = std::sync::OnceLock::new(); + *V.get_or_init(|| { + std::env::var("GPU_DEEP_FRI_CROSSCHECK") + .map(|v| v == "1") + .unwrap_or(false) + }) +} + /// R2 GPU dispatch: build the composition-polynomial Merkle tree directly from /// the device-resident extended-LDE handle (the `_keep` extend buffer), skipping /// the host re-upload that [`try_build_comp_poly_tree_gpu`] performs. Same slab @@ -1499,12 +1524,20 @@ where out } +/// GPU DEEP output. The host copy remains authoritative while the direct FRI +/// handoff is cross-checked; `device` is present only for the fully-resident +/// parts + inverse-denominators arm. +pub(crate) struct DeepGpuOutput { + pub host: Vec>, + pub device: Option, +} + /// R4 GPU dispatch: per-row DEEP composition over the full LDE domain. /// Reuses the device-resident main + (optional) aux LDE handles from R1 /// and, when supplied, the device-resident composition-parts LDE handle /// from the R2 `_keep` path. /// -/// Returns the `lde_size` ext3 evaluations of the DEEP polynomial on +/// Returns the `lde_size` ext3 evaluations plus an optional resident handle on /// success, or `None` to let the caller run its existing CPU loop. The /// caller's `inv_denoms` must be `inv_denoms[0..lde_size]` for the H-term /// and `inv_denoms[(1+k)*lde_size..(2+k)*lde_size]` for trace term k @@ -1521,7 +1554,8 @@ pub(crate) fn try_deep_composition_gpu( inv_denoms_host: &[FieldElement], inv_denoms_dev: Option<(&CudaSlice, &Arc)>, num_eval_points: usize, -) -> Option>> + retain_fully_resident_host: bool, +) -> Option> where F: IsField + IsSubFieldOf + 'static, E: IsField + 'static, @@ -1674,7 +1708,11 @@ where } let result = match (parts_dev, inv_denoms_dev) { (Some(parts), Some((inv_dev, stream))) => { - math_cuda::deep::deep_composition_ext3_with_dev_parts_and_inv_denoms( + // Keep the host codeword for the committed path and while the + // direct handoff is being cross-checked. In measured device-only + // mode FRI consumes `deep_out` on this same stream, so the D2H and + // its synchronisation are unnecessary. + math_cuda::deep::deep_composition_ext3_with_dev_parts_and_inv_denoms_keep( stream, main, aux_handle, @@ -1690,7 +1728,9 @@ where num_eval_points, row_stride_kernel, domain_size_kernel, + retain_fully_resident_host, ) + .map(|out| (out.host, Some(out.device))) } (Some(parts), None) => { let inv_h_raw: &[u64] = @@ -1715,6 +1755,7 @@ where row_stride_kernel, domain_size_kernel, ) + .map(|host| (host, None)) } (None, _) => { // De-interleave each ext3 part column into 3 contiguous base-field @@ -1755,10 +1796,11 @@ where row_stride_kernel, domain_size_kernel, ) + .map(|host| (host, None)) } }; - let deep_raw = match result { + let (deep_raw, deep_device) = match result { Ok(v) => v, Err(e) => { if diag { @@ -1775,8 +1817,11 @@ where ); } GPU_DEEP_CALLS.fetch_add(1, Ordering::Relaxed); - debug_assert_eq!(deep_raw.len(), lde_size * 3); - Some(u64_to_ext3_vec::(&deep_raw)) + debug_assert!(deep_raw.is_empty() || deep_raw.len() == lde_size * 3); + Some(DeepGpuOutput { + host: u64_to_ext3_vec::(&deep_raw), + device: deep_device, + }) } /// Build `inv_denoms[k*n + i] = 1 / (lift(coset_base[i]) - z_scalars[k])` @@ -2007,6 +2052,7 @@ where #[allow(clippy::type_complexity)] pub(crate) fn try_fri_commit_gpu( evals: &[FieldElement], + deep_device: Option, transcript: &mut T, coset_offset: &FieldElement, domain_size: usize, @@ -2033,7 +2079,10 @@ where if TypeId::of::() != TypeId::of::() { return None; } - let n0 = evals.len(); + // A resident DEEP codeword deliberately has no host `evals` in the + // measured path. Its typed handle is therefore authoritative for both + // length and storage. + let n0 = deep_device.as_ref().map_or(evals.len(), |deep| deep.len); if n0 != domain_size || !n0.is_power_of_two() || n0 < 2 { return None; } @@ -2053,10 +2102,15 @@ where inv_tw_u64.push(v); } - // SAFETY: E == Ext3; FieldElement backing is [u64; 3]. - let evals_u64: &[u64] = unsafe { ext3_slice_to_u64::(evals) }; - - let mut state = match math_cuda::fri::FriCommitState::new(evals_u64, &inv_tw_u64, n0) { + let state_result = match deep_device { + Some(deep) => math_cuda::fri::FriCommitState::from_deep(deep, &inv_tw_u64, n0), + None => { + // SAFETY: E == Ext3; FieldElement backing is [u64; 3]. + let evals_u64: &[u64] = unsafe { ext3_slice_to_u64::(evals) }; + math_cuda::fri::FriCommitState::new(evals_u64, &inv_tw_u64, n0) + } + }; + let mut state = match state_result { Ok(s) => s, Err(_) => return None, }; diff --git a/crypto/stark/src/instruments.rs b/crypto/stark/src/instruments.rs index 96bf6ffae..a7cce3356 100644 --- a/crypto/stark/src/instruments.rs +++ b/crypto/stark/src/instruments.rs @@ -185,12 +185,16 @@ pub struct TableSubOps { pub ood: Duration, /// Round 4: compute_deep_composition_poly_evaluations pub deep_comp: Duration, - /// Round 4: interpolate_fft + evaluate_fft + /// Round 4: serial CPU bit-reverse between GPU DEEP and GPU FRI. pub deep_extend: Duration, /// fri::commit_phase_from_evaluations pub fri_commit: Duration, - /// Round 4: grinding + FRI query + Merkle openings - pub queries: Duration, + /// Round 4: proof-of-work nonce search. + pub grinding: Duration, + /// Round 4: query-index sampling and FRI-layer decommitment. + pub fri_query: Duration, + /// Round 4: trace/composition Merkle openings. + pub openings: Duration, } /// Sub-operation breakdown for Round 1 aux commit pass. @@ -239,12 +243,15 @@ static AUX_INVERT_US: AtomicU64 = AtomicU64::new(0); static AUX_TERM_US: AtomicU64 = AtomicU64::new(0); static AUX_ACCUM_US: AtomicU64 = AtomicU64::new(0); +type R4SubDurations = (Duration, Duration, Duration, Duration, Duration, Duration); + thread_local! { static TIMING_DATA: RefCell> = const { RefCell::new(None) }; /// Round 2 sub-timings: (constraints, fft, merkle) static R2_SUB: RefCell> = const { RefCell::new(None) }; - /// Round 4 sub-timings: (fft, merkle, deep_comp, queries) - static R4_SUB: RefCell> = const { RefCell::new(None) }; + /// Round 4 sub-timings: (bit_reverse, fri_commit, deep_comp, grinding, + /// fri_query, openings). + static R4_SUB: RefCell> = const { RefCell::new(None) }; /// Assembled sub-ops from prove_rounds_2_to_4 (without reconstruct_round1 LDE time). static ROUND_SUB_OPS: RefCell> = const { RefCell::new(None) }; } @@ -333,11 +340,27 @@ pub fn take_r2_sub() -> Option<(Duration, Duration, Duration)> { R2_SUB.with(|cell| cell.borrow_mut().take()) } -pub fn store_r4_sub(fft: Duration, merkle: Duration, deep_comp: Duration, queries: Duration) { - R4_SUB.with(|cell| *cell.borrow_mut() = Some((fft, merkle, deep_comp, queries))); +pub fn store_r4_sub( + bit_reverse: Duration, + fri_commit: Duration, + deep_comp: Duration, + grinding: Duration, + fri_query: Duration, + openings: Duration, +) { + R4_SUB.with(|cell| { + *cell.borrow_mut() = Some(( + bit_reverse, + fri_commit, + deep_comp, + grinding, + fri_query, + openings, + )); + }); } -pub fn take_r4_sub() -> Option<(Duration, Duration, Duration, Duration)> { +pub fn take_r4_sub() -> Option<(Duration, Duration, Duration, Duration, Duration, Duration)> { R4_SUB.with(|cell| cell.borrow_mut().take()) } diff --git a/crypto/stark/src/prover.rs b/crypto/stark/src/prover.rs index 1900e1e25..299c8f229 100644 --- a/crypto/stark/src/prover.rs +++ b/crypto/stark/src/prover.rs @@ -577,6 +577,15 @@ pub(crate) struct Round4, E: IsField> { nonce: Option, } +/// Natural-order DEEP evaluations plus the optional device handle produced by +/// the fully-resident GPU arm. The host vector remains present during the +/// cross-check stage and is removed only after the direct FRI path is proven. +struct DeepCompositionOutput { + host: Vec>, + #[cfg(feature = "cuda")] + device: Option, +} + /// Returns the evaluations of the polynomial `p` over the lde domain defined by the given /// `blowup_factor`, `domain_size` and `offset`. The number of evaluations returned is `domain_size /// * blowup_factor`. The domain generator used is the one given by the implementation of `F` as `IsFFTField`. @@ -1620,10 +1629,35 @@ pub trait IsStarkProver< // <<<< Receive challenges: 𝛾ⱼ, 𝛾ⱼ' let gammas = deep_composition_coefficients; + let domain_size = domain.lde_roots_of_unity_coset.len(); + + // A valid FRI configuration may intentionally perform no folds (for + // example, a tiny trace or a terminal degree clamped to the full + // codeword). Keep the host DEEP output in those cases: the direct + // handoff is only authoritative when the GPU FRI dispatcher can + // actually consume it. + #[cfg(feature = "cuda")] + let direct_handoff_configured = { + let layout = crate::fri::terminal::FriFoldLayout::new( + domain_size.trailing_zeros(), + domain.blowup_factor.trailing_zeros(), + air.options().fri_final_poly_log_degree as u32, + ); + crate::gpu_lde::deep_fri_resident_enabled() + && layout.total_folds > 0 + && layout.terminal_len >= 2 + }; + #[cfg(feature = "cuda")] + let crosscheck_direct = crate::gpu_lde::deep_fri_crosscheck(); + #[cfg(feature = "cuda")] + let retain_fully_resident_host = !direct_handoff_configured || crosscheck_direct; + #[cfg(not(feature = "cuda"))] + let retain_fully_resident_host = true; + // Compute p₀ (deep composition polynomial) as N evaluations on trace-size coset #[cfg(feature = "instruments")] let t_sub = Instant::now(); - let deep_evals = Self::compute_deep_composition_poly_evaluations( + let deep_output = Self::compute_deep_composition_poly_evaluations( &round_1_result.lde_trace, round_2_result, round_3_result, @@ -1632,23 +1666,113 @@ pub trait IsStarkProver< &domain.trace_primitive_root, &gammas, &trace_term_coeffs, + retain_fully_resident_host, ); #[cfg(feature = "instruments")] let other_dur_1 = t_sub.elapsed(); - // DEEP evaluations are already at 2N LDE points — just bit-reverse for FRI. - // No iFFT+FFT extension needed (Plonky3-style direct LDE computation). - let domain_size = domain.lde_roots_of_unity_coset.len(); + #[cfg(feature = "cuda")] + let deep_device = deep_output.device; + + // DEEP evaluations are already at 2N LDE points. The legacy bridge + // bit-reverses them on CPU; the resident bridge folds that permutation + // into the first GPU FRI kernel and never materialises a host codeword. #[cfg(feature = "instruments")] let t_sub = Instant::now(); - let mut lde_evals = deep_evals; - in_place_bit_reverse_permute(&mut lde_evals); + let mut lde_evals = deep_output.host; + #[cfg(feature = "cuda")] + let direct_requested = direct_handoff_configured && deep_device.is_some(); + + #[cfg(feature = "cuda")] + let needs_host_bridge = !direct_requested || crosscheck_direct; + #[cfg(not(feature = "cuda"))] + let needs_host_bridge = true; + + if needs_host_bridge { + assert_eq!( + lde_evals.len(), + domain_size, + "R4 DEEP host codeword missing while the host FRI bridge is required" + ); + in_place_bit_reverse_permute(&mut lde_evals); + } #[cfg(feature = "instruments")] let r4_fft_dur = t_sub.elapsed(); // FRI commit phase from pre-computed evaluations #[cfg(feature = "instruments")] let t_sub = Instant::now(); + #[cfg(feature = "cuda")] + let direct_fri = if direct_requested { + let deep = deep_device.expect("direct DEEP->FRI requested without a device handle"); + let transcript_before = transcript.clone(); + let result = crate::gpu_lde::try_fri_commit_gpu::( + &lde_evals, + Some(deep), + transcript, + &coset_offset, + domain_size, + domain.blowup_factor.trailing_zeros(), + air.options().fri_final_poly_log_degree as u32, + ); + Some((transcript_before, result)) + } else { + None + }; + + #[cfg(feature = "cuda")] + let direct_fri = direct_fri.and_then(|(transcript_before, result)| { + result.map(|direct| (transcript_before, direct)) + }); + + #[cfg(feature = "cuda")] + let (fri_final_poly_coeffs, fri_layers) = if let Some((transcript_before, direct)) = + direct_fri + { + if crosscheck_direct { + let mut reference_transcript = transcript_before; + let reference = fri::commit_phase_from_evaluations( + lde_evals.clone(), + &mut reference_transcript, + &coset_offset, + domain_size, + domain.blowup_factor.trailing_zeros(), + air.options().fri_final_poly_log_degree as u32, + ); + let direct_roots: Vec<_> = direct.1.iter().map(|l| l.merkle_tree.root).collect(); + let reference_roots: Vec<_> = + reference.1.iter().map(|l| l.merkle_tree.root).collect(); + assert_eq!( + direct.0, reference.0, + "DEEP->FRI terminal coefficients mismatch" + ); + assert_eq!( + direct_roots, reference_roots, + "DEEP->FRI layer roots mismatch" + ); + assert_eq!( + transcript.state(), + reference_transcript.state(), + "DEEP->FRI transcript state mismatch" + ); + } + direct + } else { + assert!( + !direct_requested || crosscheck_direct, + "device-only DEEP->FRI handoff failed; refusing a host fallback without a host codeword" + ); + fri::commit_phase_from_evaluations( + lde_evals, + transcript, + &coset_offset, + domain_size, + domain.blowup_factor.trailing_zeros(), + air.options().fri_final_poly_log_degree as u32, + ) + }; + + #[cfg(not(feature = "cuda"))] let (fri_final_poly_coeffs, fri_layers) = fri::commit_phase_from_evaluations( lde_evals, transcript, @@ -1671,12 +1795,20 @@ pub trait IsStarkProver< transcript.append_bytes(&nonce_value.to_be_bytes()); nonce = Some(nonce_value); } + #[cfg(feature = "instruments")] + let grinding_dur = t_sub.elapsed(); + #[cfg(feature = "instruments")] + let t_sub = Instant::now(); let number_of_queries = air.options().fri_number_of_queries; let iotas = Self::sample_query_indexes(number_of_queries, domain, transcript); let query_list = fri::query_phase(&fri_layers, &iotas); + #[cfg(feature = "instruments")] + let fri_query_dur = t_sub.elapsed(); + #[cfg(feature = "instruments")] + let t_sub = Instant::now(); let fri_layers_merkle_roots: Vec<_> = fri_layers .iter() .map(|layer| layer.merkle_tree.root) @@ -1687,8 +1819,15 @@ pub trait IsStarkProver< #[cfg(feature = "instruments")] { - let queries_dur = t_sub.elapsed(); - crate::instruments::store_r4_sub(r4_fft_dur, r4_merkle_dur, other_dur_1, queries_dur); + let openings_dur = t_sub.elapsed(); + crate::instruments::store_r4_sub( + r4_fft_dur, + r4_merkle_dur, + other_dur_1, + grinding_dur, + fri_query_dur, + openings_dur, + ); } Round4 { @@ -1730,11 +1869,15 @@ pub trait IsStarkProver< primitive_root: &FieldElement, composition_poly_gammas: &[FieldElement], trace_terms_gammas: &[Vec>], - ) -> Vec> + retain_fully_resident_host: bool, + ) -> DeepCompositionOutput where FieldElement: AsBytes, FieldElement: AsBytes, { + #[cfg(not(feature = "cuda"))] + let _ = retain_fully_resident_host; + let num_parts = round_2_result.num_composition_parts; let z_power = z.pow(num_parts); // pole for H terms @@ -1795,9 +1938,13 @@ pub trait IsStarkProver< &[], Some((&inv_dev, &stream)), num_eval_points, + retain_fully_resident_host, ) { - return deep_evals; + return DeepCompositionOutput { + host: deep_evals.host, + device: deep_evals.device, + }; } } @@ -1834,9 +1981,13 @@ pub trait IsStarkProver< &denoms, None, num_eval_points, + true, ) { - return deep_evals; + return DeepCompositionOutput { + host: deep_evals.host, + device: deep_evals.device, + }; } } @@ -1892,7 +2043,7 @@ pub trait IsStarkProver< }) .collect(); - crate::par::par_map_collect(0..lde_size, |i| { + let host = crate::par::par_map_collect(0..lde_size, |i| { let mut result = FieldElement::::zero(); // H terms @@ -1917,7 +2068,12 @@ pub trait IsStarkProver< } result - }) + }); + DeepCompositionOutput { + host, + #[cfg(feature = "cuda")] + device: None, + } } /// Computes values and validity proofs of the evaluations of the composition polynomial parts @@ -3336,17 +3492,25 @@ pub trait IsStarkProver< let zero = Duration::ZERO; let (r2_constraints, r2_fft, r2_merkle) = crate::instruments::take_r2_sub().unwrap_or((zero, zero, zero)); - let (r4_fft, r4_merkle, r4_deep_comp, r4_queries) = - crate::instruments::take_r4_sub().unwrap_or((zero, zero, zero, zero)); + let ( + r4_bit_reverse, + r4_fri_commit, + r4_deep_comp, + r4_grinding, + r4_fri_query, + r4_openings, + ) = crate::instruments::take_r4_sub().unwrap_or((zero, zero, zero, zero, zero, zero)); crate::instruments::store_round_sub_ops(crate::instruments::TableSubOps { constraints: r2_constraints, comp_decompose: r2_fft, comp_commit: r2_merkle, ood: round_3_dur, deep_comp: r4_deep_comp, - deep_extend: r4_fft, - fri_commit: r4_merkle, - queries: r4_queries, + deep_extend: r4_bit_reverse, + fri_commit: r4_fri_commit, + grinding: r4_grinding, + fri_query: r4_fri_query, + openings: r4_openings, }); } diff --git a/prover/src/instruments.rs b/prover/src/instruments.rs index 0ea28273b..afae7ea18 100644 --- a/prover/src/instruments.rs +++ b/prover/src/instruments.rs @@ -140,7 +140,9 @@ pub fn print_report( entry.sub_ops.deep_comp += sub_ops.deep_comp; entry.sub_ops.deep_extend += sub_ops.deep_extend; entry.sub_ops.fri_commit += sub_ops.fri_commit; - entry.sub_ops.queries += sub_ops.queries; + entry.sub_ops.grinding += sub_ops.grinding; + entry.sub_ops.fri_query += sub_ops.fri_query; + entry.sub_ops.openings += sub_ops.openings; } let mut sorted: Vec<_> = merged.into_iter().collect(); @@ -177,7 +179,9 @@ pub fn print_report( let mut total_deep_comp = Duration::ZERO; let mut total_deep_extend = Duration::ZERO; let mut total_fri_commit = Duration::ZERO; - let mut total_queries = Duration::ZERO; + let mut total_grinding = Duration::ZERO; + let mut total_fri_query = Duration::ZERO; + let mut total_openings = Duration::ZERO; for (_, t) in &sorted { total_constraints += t.sub_ops.constraints; total_comp_decompose += t.sub_ops.comp_decompose; @@ -186,7 +190,9 @@ pub fn print_report( total_deep_comp += t.sub_ops.deep_comp; total_deep_extend += t.sub_ops.deep_extend; total_fri_commit += t.sub_ops.fri_commit; - total_queries += t.sub_ops.queries; + total_grinding += t.sub_ops.grinding; + total_fri_query += t.sub_ops.fri_query; + total_openings += t.sub_ops.openings; } let sub_ops_sum = total_constraints @@ -196,7 +202,9 @@ pub fn print_report( + total_deep_comp + total_deep_extend + total_fri_commit - + total_queries; + + total_grinding + + total_fri_query + + total_openings; if sub_ops_sum > Duration::ZERO { let mut sub_ops: Vec<(&str, Duration)> = vec![ ("R2 evaluate", total_constraints), @@ -204,9 +212,11 @@ pub fn print_report( ("R2 commit_bit_reversed (comp-poly)", total_comp_commit), ("R3 OOD evaluation", total_ood), ("R4 deep_composition_poly_evals", total_deep_comp), - ("R4 interpolate+evaluate_fft", total_deep_extend), + ("R4 DEEP->FRI CPU bit-reverse", total_deep_extend), ("R4 fri::commit_phase", total_fri_commit), - ("R4 queries & openings", total_queries), + ("R4 grinding", total_grinding), + ("R4 FRI query", total_fri_query), + ("R4 trace/composition openings", total_openings), ]; sub_ops.sort_by(|a, b| b.1.cmp(&a.1)); eprintln!( From 506909f3df8b96555ea11eda04d9e9462798cbeb Mon Sep 17 00:00:00 2001 From: jotabulacios Date: Fri, 24 Jul 2026 11:09:21 -0300 Subject: [PATCH 10/11] Fix cuda-test call sites for make lint --- crypto/stark/src/tests/prover_tests.rs | 15 +++++++++++++++ prover/tests/gpu_constraint_interp_real.rs | 1 + 2 files changed, 16 insertions(+) diff --git a/crypto/stark/src/tests/prover_tests.rs b/crypto/stark/src/tests/prover_tests.rs index a536a206a..4c4f215de 100644 --- a/crypto/stark/src/tests/prover_tests.rs +++ b/crypto/stark/src/tests/prover_tests.rs @@ -272,11 +272,26 @@ fn test_decompose_and_extend_d2_matches_original() { // --- New path: algebraic decomposition --- let twiddles = LdeTwiddles::new(&domain); assert!(!twiddles.has_composition_cache()); + #[cfg(not(feature = "cuda"))] let new_result = Prover::::decompose_and_extend_d2( &constraint_evaluations, &domain, &twiddles, ); + #[cfg(feature = "cuda")] + let new_result = { + // The cuda signature adds a resident-handle out-param and a + // retain-host-LDE flag; retain the host evals so the parity comparison + // below has them (falls back to the CPU extend without a GPU present). + let mut gpu_parts: Option = None; + Prover::::decompose_and_extend_d2( + &constraint_evaluations, + &domain, + &twiddles, + &mut gpu_parts, + true, + ) + }; #[cfg(not(feature = "cuda"))] assert!(twiddles.has_composition_cache()); diff --git a/prover/tests/gpu_constraint_interp_real.rs b/prover/tests/gpu_constraint_interp_real.rs index e464f2556..f59a7fb14 100644 --- a/prover/tests/gpu_constraint_interp_real.rs +++ b/prover/tests/gpu_constraint_interp_real.rs @@ -157,6 +157,7 @@ fn check_air(air: &dyn AIR, m: aux_cols, lde_size, tree: None, + ready: None, }; let gpu = try_eval_program_gpu( From d51680eff16617294fd849b75a3dbf8b1406c7e3 Mon Sep 17 00:00:00 2001 From: jotabulacios Date: Fri, 24 Jul 2026 12:17:07 -0300 Subject: [PATCH 11/11] Recover CPU FRI on device-only DEEP-FRI fault --- crypto/stark/src/fri/mod.rs | 9 ++- crypto/stark/src/prover.rs | 118 +++++++++++++++++----------- crypto/stark/src/tests/fri_tests.rs | 1 + 3 files changed, 81 insertions(+), 47 deletions(-) diff --git a/crypto/stark/src/fri/mod.rs b/crypto/stark/src/fri/mod.rs index 55e3686d0..5bd5742d4 100644 --- a/crypto/stark/src/fri/mod.rs +++ b/crypto/stark/src/fri/mod.rs @@ -37,6 +37,10 @@ pub fn commit_phase_from_evaluations< domain_size: usize, blowup_log: u32, final_poly_log_degree: u32, + // When true, skip the GPU fast path and commit on CPU. Used by the R4 + // GPU-fault recovery: once a device FRI attempt has failed, retrying the GPU + // would silently re-run it (masking the fallback), so recovery forces CPU. + force_cpu: bool, ) -> ( Vec>, Vec>>, @@ -52,8 +56,11 @@ where // snapshots the transcript before mutating it so a mid-loop cudarc // error restores state and lets the CPU loop below run as if the GPU // had never been tried. + #[cfg(not(feature = "cuda"))] + let _ = force_cpu; + #[cfg(feature = "cuda")] - { + if !force_cpu { // Try the GPU early-termination FRI commit first. `try_fri_commit_gpu` // drives the same commit phase on-device (Goldilocks + Ext3, above the // LDE size threshold, and only when folding actually happens) and returns diff --git a/crypto/stark/src/prover.rs b/crypto/stark/src/prover.rs index 299c8f229..3577073d8 100644 --- a/crypto/stark/src/prover.rs +++ b/crypto/stark/src/prover.rs @@ -1703,7 +1703,7 @@ pub trait IsStarkProver< #[cfg(feature = "instruments")] let t_sub = Instant::now(); #[cfg(feature = "cuda")] - let direct_fri = if direct_requested { + let (fri_final_poly_coeffs, fri_layers) = if direct_requested { let deep = deep_device.expect("direct DEEP->FRI requested without a device handle"); let transcript_before = transcript.clone(); let result = crate::gpu_lde::try_fri_commit_gpu::( @@ -1715,53 +1715,77 @@ pub trait IsStarkProver< domain.blowup_factor.trailing_zeros(), air.options().fri_final_poly_log_degree as u32, ); - Some((transcript_before, result)) - } else { - None - }; - - #[cfg(feature = "cuda")] - let direct_fri = direct_fri.and_then(|(transcript_before, result)| { - result.map(|direct| (transcript_before, direct)) - }); - - #[cfg(feature = "cuda")] - let (fri_final_poly_coeffs, fri_layers) = if let Some((transcript_before, direct)) = - direct_fri - { - if crosscheck_direct { - let mut reference_transcript = transcript_before; - let reference = fri::commit_phase_from_evaluations( - lde_evals.clone(), - &mut reference_transcript, - &coset_offset, - domain_size, - domain.blowup_factor.trailing_zeros(), - air.options().fri_final_poly_log_degree as u32, - ); - let direct_roots: Vec<_> = direct.1.iter().map(|l| l.merkle_tree.root).collect(); - let reference_roots: Vec<_> = - reference.1.iter().map(|l| l.merkle_tree.root).collect(); - assert_eq!( - direct.0, reference.0, - "DEEP->FRI terminal coefficients mismatch" - ); - assert_eq!( - direct_roots, reference_roots, - "DEEP->FRI layer roots mismatch" - ); - assert_eq!( - transcript.state(), - reference_transcript.state(), - "DEEP->FRI transcript state mismatch" - ); + match result { + Some(direct) => { + if crosscheck_direct { + let mut reference_transcript = transcript_before; + let reference = fri::commit_phase_from_evaluations( + lde_evals.clone(), + &mut reference_transcript, + &coset_offset, + domain_size, + domain.blowup_factor.trailing_zeros(), + air.options().fri_final_poly_log_degree as u32, + false, + ); + let direct_roots: Vec<_> = + direct.1.iter().map(|l| l.merkle_tree.root).collect(); + let reference_roots: Vec<_> = + reference.1.iter().map(|l| l.merkle_tree.root).collect(); + assert_eq!( + direct.0, reference.0, + "DEEP->FRI terminal coefficients mismatch" + ); + assert_eq!( + direct_roots, reference_roots, + "DEEP->FRI layer roots mismatch" + ); + assert_eq!( + transcript.state(), + reference_transcript.state(), + "DEEP->FRI transcript state mismatch" + ); + } + direct + } + // Device-only GPU FRI failed (e.g. a GPU fault mid-fold). The host + // codeword was dropped for residency, so recover it instead of + // aborting: restore the transcript to before the failed attempt, + // recompute the DEEP codeword on host, and run the CPU FRI. The + // recompute only happens on this rare failure path, so the happy + // path keeps the residency win (no host D2H). + None => { + *transcript = transcript_before; + let mut recovered = Self::compute_deep_composition_poly_evaluations( + &round_1_result.lde_trace, + round_2_result, + round_3_result, + z, + domain, + &domain.trace_primitive_root, + &gammas, + &trace_term_coeffs, + true, + ) + .host; + assert_eq!( + recovered.len(), + domain_size, + "recovered DEEP codeword size mismatch on GPU-FRI fallback" + ); + in_place_bit_reverse_permute(&mut recovered); + fri::commit_phase_from_evaluations( + recovered, + transcript, + &coset_offset, + domain_size, + domain.blowup_factor.trailing_zeros(), + air.options().fri_final_poly_log_degree as u32, + true, // force CPU: the GPU FRI already failed for this table + ) + } } - direct } else { - assert!( - !direct_requested || crosscheck_direct, - "device-only DEEP->FRI handoff failed; refusing a host fallback without a host codeword" - ); fri::commit_phase_from_evaluations( lde_evals, transcript, @@ -1769,6 +1793,7 @@ pub trait IsStarkProver< domain_size, domain.blowup_factor.trailing_zeros(), air.options().fri_final_poly_log_degree as u32, + false, ) }; @@ -1780,6 +1805,7 @@ pub trait IsStarkProver< domain_size, domain.blowup_factor.trailing_zeros(), air.options().fri_final_poly_log_degree as u32, + false, ); #[cfg(feature = "instruments")] let r4_merkle_dur = t_sub.elapsed(); diff --git a/crypto/stark/src/tests/fri_tests.rs b/crypto/stark/src/tests/fri_tests.rs index 10b34afbb..7360e7faa 100644 --- a/crypto/stark/src/tests/fri_tests.rs +++ b/crypto/stark/src/tests/fri_tests.rs @@ -181,6 +181,7 @@ fn test_commit_phase_early_termination_roundtrip() { initial_len, blowup_log, final_poly_log_degree, + false, ); assert_eq!(