diff --git a/crypto/math-cuda/kernels/barycentric.cu b/crypto/math-cuda/kernels/barycentric.cu index 5c18bcb88..f76db471a 100644 --- a/crypto/math-cuda/kernels/barycentric.cu +++ b/crypto/math-cuda/kernels/barycentric.cu @@ -190,3 +190,45 @@ extern "C" __global__ void barycentric_ext3_batched_strided( out_ext3_int[col * 3 + 2] = sum.c; } } + +// Gather full rows from a device-resident base-field LDE (`buf[col*col_stride + +// row]`). One block per gathered row, threads stride over columns. Output is +// row-major `out[q*num_cols + col]` for gathered-row slot `q` — directly the +// concatenation of `gather_main_row(rows[q])` for each q. `rows` are the LDE row +// indices to gather (already the reversed query rows on the host side). +extern "C" __global__ void gather_rows_base( + const uint64_t *__restrict__ columns, + uint64_t col_stride, + uint64_t num_cols, + const uint32_t *__restrict__ rows, + uint64_t num_rows, + uint64_t *__restrict__ out +) { + uint64_t q = blockIdx.x; + if (q >= num_rows) return; + uint64_t row = rows[q]; + for (uint64_t col = threadIdx.x; col < num_cols; col += blockDim.x) { + out[q * num_cols + col] = columns[col * col_stride + row]; + } +} + +// Ext3 variant: M ext3 columns as 3M base slabs, `columns[(col*3+k)*col_stride + +// row]`. Output interleaved ext3: `out[(q*num_cols + col)*3 + k]`. +extern "C" __global__ void gather_rows_ext3( + const uint64_t *__restrict__ columns, + uint64_t col_stride, + uint64_t num_cols, + const uint32_t *__restrict__ rows, + uint64_t num_rows, + uint64_t *__restrict__ out +) { + uint64_t q = blockIdx.x; + if (q >= num_rows) return; + uint64_t row = rows[q]; + for (uint64_t col = threadIdx.x; col < num_cols; col += blockDim.x) { + uint64_t o = (q * num_cols + col) * 3; + out[o + 0] = columns[(col * 3 + 0) * col_stride + row]; + out[o + 1] = columns[(col * 3 + 1) * col_stride + row]; + out[o + 2] = columns[(col * 3 + 2) * col_stride + row]; + } +} diff --git a/crypto/math-cuda/src/barycentric.rs b/crypto/math-cuda/src/barycentric.rs index f299d1839..8fdea14f1 100644 --- a/crypto/math-cuda/src/barycentric.rs +++ b/crypto/math-cuda/src/barycentric.rs @@ -337,3 +337,81 @@ pub fn barycentric_ext3_on_device_with_dev_inv_denoms( stream.synchronize()?; Ok(out) } + +/// Gather full rows from a device-resident base-field LDE handle. `rows` are LDE +/// row indices; returns their column values row-major (`rows.len() * main.m` +/// u64, `out[q*num_cols + col]`) — i.e. the concatenation of +/// `gather_main_row(rows[q])` for each `q`. Runs on the caller's stream. +pub fn gather_rows_base_on_device( + main: &GpuLdeBase, + rows: &[u32], + stream: &Arc, +) -> Result> { + let num_cols = main.m; + if num_cols == 0 || rows.is_empty() { + return Ok(Vec::new()); + } + let be = backend()?; + let rows_dev = stream.clone_htod(rows)?; + let mut out = stream.alloc_zeros::(rows.len() * num_cols)?; + let col_stride = main.lde_size as u64; + let num_cols_u64 = num_cols as u64; + let num_rows_u64 = rows.len() as u64; + let cfg = LaunchConfig { + grid_dim: (rows.len() as u32, 1, 1), + block_dim: (BLOCK_DIM.min(num_cols as u32).max(1), 1, 1), + shared_mem_bytes: 0, + }; + unsafe { + stream + .launch_builder(&be.gather_rows_base) + .arg(main.buf.as_ref()) + .arg(&col_stride) + .arg(&num_cols_u64) + .arg(&rows_dev) + .arg(&num_rows_u64) + .arg(&mut out) + .launch(cfg)?; + } + let host = stream.clone_dtoh(&out)?; + stream.synchronize()?; + Ok(host) +} + +/// Ext3 sibling of [`gather_rows_base_on_device`]: returns `rows.len() * aux.m * +/// 3` u64, interleaved ext3 (`out[(q*num_cols + col)*3 + k]`). +pub fn gather_rows_ext3_on_device( + aux: &GpuLdeExt3, + rows: &[u32], + stream: &Arc, +) -> Result> { + let num_cols = aux.m; + if num_cols == 0 || rows.is_empty() { + return Ok(Vec::new()); + } + let be = backend()?; + 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; + let num_cols_u64 = num_cols as u64; + let num_rows_u64 = rows.len() as u64; + let cfg = LaunchConfig { + grid_dim: (rows.len() as u32, 1, 1), + block_dim: (BLOCK_DIM.min(num_cols as u32).max(1), 1, 1), + shared_mem_bytes: 0, + }; + unsafe { + stream + .launch_builder(&be.gather_rows_ext3) + .arg(aux.buf.as_ref()) + .arg(&col_stride) + .arg(&num_cols_u64) + .arg(&rows_dev) + .arg(&num_rows_u64) + .arg(&mut out) + .launch(cfg)?; + } + let host = stream.clone_dtoh(&out)?; + stream.synchronize()?; + Ok(host) +} diff --git a/crypto/math-cuda/src/device.rs b/crypto/math-cuda/src/device.rs index be0440d5e..bf75b696e 100644 --- a/crypto/math-cuda/src/device.rs +++ b/crypto/math-cuda/src/device.rs @@ -167,6 +167,8 @@ pub struct Backend { pub barycentric_ext3_batched: CudaFunction, pub barycentric_base_batched_strided: CudaFunction, pub barycentric_ext3_batched_strided: CudaFunction, + pub gather_rows_base: CudaFunction, + pub gather_rows_ext3: CudaFunction, // deep.ptx pub deep_composition_ext3_row: CudaFunction, @@ -367,6 +369,8 @@ impl Backend { .load_function("barycentric_base_batched_strided")?, barycentric_ext3_batched_strided: bary .load_function("barycentric_ext3_batched_strided")?, + gather_rows_base: bary.load_function("gather_rows_base")?, + 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_update_twiddles: fri.load_function("fri_update_twiddles")?, diff --git a/crypto/math-cuda/src/lde.rs b/crypto/math-cuda/src/lde.rs index 30e524c2c..06c1a02cd 100644 --- a/crypto/math-cuda/src/lde.rs +++ b/crypto/math-cuda/src/lde.rs @@ -400,6 +400,7 @@ enum InnerInput<'a> { /// required by the downstream GPU kernels DEEP/barycentric); callers wrap it in /// the appropriate LDE handle. #[allow(clippy::type_complexity)] +#[allow(clippy::too_many_arguments)] fn coset_lde_row_major_inner( input: InnerInput, n: usize, @@ -408,6 +409,7 @@ fn coset_lde_row_major_inner( weights: &[u64], what: &str, retain_trace_col_major: bool, + retain_host_lde: bool, ) -> Result<( GpuMerkleTree, CudaSlice, @@ -514,7 +516,10 @@ fn coset_lde_row_major_inner( // D2H the row-major LDE first (before the handle transpose). Release the // staging lock before the Merkle nodes transfer to minimise lock contention. - let lde_out = { + // Skipped entirely when `retain_host_lde` is false (the caller keeps the LDE + // device-only): this is the round-1 trace D2H the full-residency path + // eliminates — the big transfer/alloc win — so we return an empty host Vec. + let lde_out = if retain_host_lde { let staging_slot = be.pinned_staging(); let mut staging = staging_slot.lock().unwrap(); staging.ensure_capacity(lde_size * total_cols, &be.ctx)?; @@ -524,6 +529,8 @@ fn coset_lde_row_major_inner( let out = pinned[..lde_size * total_cols].to_vec(); drop(staging); out + } else { + Vec::new() }; // Keep the Merkle tree resident on device; copy only the 32 byte root so the @@ -562,6 +569,7 @@ pub fn coset_lde_row_major_with_merkle_tree_keep( m: usize, blowup_factor: usize, weights: &[u64], + retain_host_lde: bool, ) -> Result<(GpuLdeBase, Vec)> { let (tree, col_major_dev, lde_out, trace_col_major) = coset_lde_row_major_inner( InnerInput::Host(row_major), @@ -571,6 +579,7 @@ pub fn coset_lde_row_major_with_merkle_tree_keep( weights, "coset_lde_row_major lde_size", true, + retain_host_lde, )?; let handle = GpuLdeBase { buf: Arc::new(col_major_dev), @@ -599,6 +608,7 @@ pub fn coset_lde_ext3_row_major_with_merkle_tree_keep( m: usize, blowup_factor: usize, weights: &[u64], + retain_host_lde: bool, ) -> Result<(GpuLdeExt3, Vec)> { let (tree, col_major_dev, lde_out, _) = coset_lde_row_major_inner( InnerInput::Host(row_major), @@ -608,6 +618,7 @@ pub fn coset_lde_ext3_row_major_with_merkle_tree_keep( weights, "coset_lde_ext3_row_major lde_size", false, + retain_host_lde, )?; let handle = GpuLdeExt3 { buf: Arc::new(col_major_dev), @@ -628,6 +639,7 @@ pub fn coset_lde_ext3_row_major_with_merkle_tree_keep_dev( m: usize, blowup_factor: usize, weights: &[u64], + retain_host_lde: bool, ) -> Result<(GpuLdeExt3, Vec)> { let (tree, col_major_dev, lde_out, _) = coset_lde_row_major_inner( InnerInput::Dev(input_dev), @@ -637,6 +649,7 @@ pub fn coset_lde_ext3_row_major_with_merkle_tree_keep_dev( weights, "coset_lde_ext3_row_major_dev lde_size", false, + retain_host_lde, )?; let handle = GpuLdeExt3 { buf: Arc::new(col_major_dev), diff --git a/crypto/math-cuda/tests/gather_rows.rs b/crypto/math-cuda/tests/gather_rows.rs new file mode 100644 index 000000000..8b7f1b4a3 --- /dev/null +++ b/crypto/math-cuda/tests/gather_rows.rs @@ -0,0 +1,95 @@ +//! Parity: the device row-gather (`gather_rows_base` / `gather_rows_ext3`, +//! used by R4 query openings to read opened column values from the resident LDE +//! instead of the host trace) returns exactly the column values obtained by +//! directly indexing the column-major device buffer at each requested row. + +use std::sync::Arc; + +use math_cuda::barycentric::{gather_rows_base_on_device, gather_rows_ext3_on_device}; +use math_cuda::device::backend; +use math_cuda::lde::{GpuLdeBase, GpuLdeExt3}; +use rand::{Rng, SeedableRng}; +use rand_chacha::ChaCha8Rng; + +fn run_base(lde_size: usize, num_cols: usize, seed: u64) { + let mut rng = ChaCha8Rng::seed_from_u64(seed); + // Column-major base LDE: buf[col*lde_size + row]. + let buf: Vec = (0..num_cols * lde_size) + .map(|_| rng.r#gen::()) + .collect(); + + let be = backend().unwrap(); + let stream = be.next_stream(); + let dev = stream.clone_htod(&buf).unwrap(); + stream.synchronize().unwrap(); + let handle = GpuLdeBase { + buf: Arc::new(dev), + m: num_cols, + lde_size, + tree: None, + trace_dev: None, + trace_rows: 0, + }; + + let rows: Vec = (0..9).map(|_| rng.gen_range(0..lde_size) as u32).collect(); + let got = gather_rows_base_on_device(&handle, &rows, &stream).unwrap(); + assert_eq!(got.len(), rows.len() * num_cols, "base gather shape"); + for (q, &row) in rows.iter().enumerate() { + for col in 0..num_cols { + assert_eq!( + got[q * num_cols + col], + buf[col * lde_size + row as usize], + "base gather mismatch: row {row}, col {col}" + ); + } + } +} + +fn run_ext3(lde_size: usize, num_cols: usize, seed: u64) { + let mut rng = ChaCha8Rng::seed_from_u64(seed); + // De-interleaved ext3 LDE: buf[(col*3 + k)*lde_size + row]. + let buf: Vec = (0..num_cols * 3 * lde_size) + .map(|_| rng.r#gen::()) + .collect(); + + let be = backend().unwrap(); + let stream = be.next_stream(); + let dev = stream.clone_htod(&buf).unwrap(); + stream.synchronize().unwrap(); + let handle = GpuLdeExt3 { + buf: Arc::new(dev), + m: num_cols, + lde_size, + tree: None, + }; + + let rows: Vec = (0..9).map(|_| rng.gen_range(0..lde_size) as u32).collect(); + let got = gather_rows_ext3_on_device(&handle, &rows, &stream).unwrap(); + assert_eq!(got.len(), rows.len() * num_cols * 3, "ext3 gather shape"); + for (q, &row) in rows.iter().enumerate() { + for col in 0..num_cols { + let o = (q * num_cols + col) * 3; + for k in 0..3 { + assert_eq!( + got[o + k], + buf[(col * 3 + k) * lde_size + row as usize], + "ext3 gather mismatch: row {row}, col {col}, comp {k}" + ); + } + } + } +} + +#[test] +fn gather_rows_base_matches_direct_indexing() { + for (log_size, cols) in [(6u32, 3usize), (12, 20), (16, 8)] { + run_base(1usize << log_size, cols, 100 + log_size as u64); + } +} + +#[test] +fn gather_rows_ext3_matches_direct_indexing() { + for (log_size, cols) in [(6u32, 2usize), (12, 5), (16, 3)] { + run_ext3(1usize << log_size, cols, 200 + log_size as u64); + } +} diff --git a/crypto/math-cuda/tests/merkle_root_parity.rs b/crypto/math-cuda/tests/merkle_root_parity.rs index fcc9d226e..208353d95 100644 --- a/crypto/math-cuda/tests/merkle_root_parity.rs +++ b/crypto/math-cuda/tests/merkle_root_parity.rs @@ -305,6 +305,7 @@ fn new_row_major_pipeline_base_root_matches_cpu() { num_cols, blowup, &weights_u64, + true, ) .expect("new row-major GPU pipeline"); let gpu_root = handle.tree.as_ref().expect("resident merkle tree").root; @@ -368,6 +369,7 @@ fn new_row_major_pipeline_ext3_root_matches_cpu() { num_cols, blowup, &weights_u64, + true, ) .expect("new ext3 row-major GPU pipeline"); let gpu_root = handle.tree.as_ref().expect("resident merkle tree").root; diff --git a/crypto/stark/src/constraint_ir/gpu_interp.rs b/crypto/stark/src/constraint_ir/gpu_interp.rs index 8540ce81e..446844ac4 100644 --- a/crypto/stark/src/constraint_ir/gpu_interp.rs +++ b/crypto/stark/src/constraint_ir/gpu_interp.rs @@ -65,6 +65,39 @@ unsafe fn base_slice_to_u64(xs: &[FieldElement]) -> Vec { raw.to_vec() } +/// Lift raw base-field limbs (one canonical-Goldilocks `u64` per element, as +/// produced by the device row-gather kernels) back into owned `FieldElement`s. +/// Returns `None` unless `F == GoldilocksField`. The inverse of +/// [`base_slice_to_u64`]; used to feed device-gathered LDE rows into the generic +/// prover openings. +pub fn base_u64_to_field(raw: &[u64]) -> Option>> { + if TypeId::of::() != TypeId::of::() { + return None; + } + // SAFETY: the gate established `F == GoldilocksField`, whose `FieldElement` + // is `#[repr(transparent)]` over `u64`; `raw` (a `*const u64`) is 8-aligned. + let fe = + unsafe { std::slice::from_raw_parts(raw.as_ptr() as *const FieldElement, raw.len()) }; + Some(fe.to_vec()) +} + +/// Ext3 sibling of [`base_u64_to_field`]: `raw` holds `3` interleaved limbs per +/// element (`[c0, c1, c2]`). Returns `None` unless `E` is the degree-3 +/// Goldilocks extension. Inverse of [`ext3_slice_to_u64`]. +pub fn ext3_u64_to_field(raw: &[u64]) -> Option>> { + if TypeId::of::() != TypeId::of::() { + return None; + } + debug_assert_eq!(raw.len() % 3, 0, "ext3 limbs must come in triples"); + // SAFETY: the gate established the degree-3 Goldilocks extension, whose + // `FieldElement` is `#[repr(transparent)]` over `[u64; 3]`; `raw` (a + // `*const u64`) is 8-aligned, matching `[u64; 3]`'s alignment. + let fe = unsafe { + std::slice::from_raw_parts(raw.as_ptr() as *const FieldElement, raw.len() / 3) + }; + Some(fe.to_vec()) +} + /// Per-proof accumulation inputs (in `FieldElement` form) for /// [`try_eval_composition_gpu`], mirroring the CPU accumulation in /// `constraints::evaluator`. diff --git a/crypto/stark/src/constraints/evaluator.rs b/crypto/stark/src/constraints/evaluator.rs index acd4fef99..efeb10f8c 100644 --- a/crypto/stark/src/constraints/evaluator.rs +++ b/crypto/stark/src/constraints/evaluator.rs @@ -233,6 +233,17 @@ where } } + // Reaching here means the GPU composition path fell through to the CPU + // boundary + transition evaluation below, both of which read the host + // trace (`get_main`/`get_aux`, `RowFrame::from_lde`). Under the + // device-only gate that trace is empty, so a fall-through is a mis-gate + // or an unexpected GPU failure: hard-abort rather than read empty buffers. + #[cfg(feature = "cuda")] + assert!( + !lde_trace.host_trace_empty(), + "R2 composition fell back to the host trace, but it is device-only (empty)" + ); + // Fused boundary evaluation: compute (trace[col] - value) on-the-fly // instead of pre-computing all boundary_polys_evaluations. // This eliminates N_constraints × LDE_size intermediate allocations. diff --git a/crypto/stark/src/gpu_lde.rs b/crypto/stark/src/gpu_lde.rs index c812a3c1e..90dabf34c 100644 --- a/crypto/stark/src/gpu_lde.rs +++ b/crypto/stark/src/gpu_lde.rs @@ -78,6 +78,8 @@ pub fn reset_all_gpu_call_counters() { GPU_BATCH_INVERT_CALLS.store(0, Ordering::Relaxed); GPU_LOGUP_CALLS.store(0, Ordering::Relaxed); GPU_COMPOSITION_CALLS.store(0, Ordering::Relaxed); + GPU_OPENING_GATHER_CALLS.store(0, Ordering::Relaxed); + GPU_DEVICE_ONLY_CALLS.store(0, Ordering::Relaxed); } pub(crate) static GPU_EXTEND_HALVES_CALLS: AtomicU64 = AtomicU64::new(0); @@ -101,6 +103,27 @@ pub fn gpu_composition_calls() -> u64 { GPU_COMPOSITION_CALLS.load(Ordering::Relaxed) } +/// Successful device-resident-LDE opening-value gathers in +/// `open_deep_composition_poly` — one per main/aux trace whose R4 query rows +/// were read straight off the device LDE instead of the host trace (a +/// non-resident tree or non-Goldilocks tower falls back to the host gather and +/// is not counted). Guards against a silent regression where Stage-2 openings +/// quietly revert to the host path. +pub(crate) static GPU_OPENING_GATHER_CALLS: AtomicU64 = AtomicU64::new(0); +pub fn gpu_opening_gather_calls() -> u64 { + GPU_OPENING_GATHER_CALLS.load(Ordering::Relaxed) +} + +/// Tables whose round-1 LDE was kept device-only (host trace D2H skipped) — the +/// Stage-3 full-residency win. Incremented once per main trace that took the +/// `device_only` path. Zero means every table kept its host copy (gate never +/// engaged), so a residency regression drops this to 0 while proofs still +/// verify. +pub(crate) static GPU_DEVICE_ONLY_CALLS: AtomicU64 = AtomicU64::new(0); +pub fn gpu_device_only_calls() -> u64 { + GPU_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). @@ -123,6 +146,66 @@ pub(crate) fn gpu_composition_disabled() -> bool { }) } +/// Runtime override to force the Stage-3 device-only path off (keeps the round-1 +/// host D2H). Independent of the composition toggle, so the residency win can be +/// A/B-benched with the GPU composition path left on. Default off (path enabled). +static DEVICE_ONLY_DISABLED: std::sync::atomic::AtomicBool = + std::sync::atomic::AtomicBool::new(false); +pub fn set_device_only_disabled(v: bool) { + DEVICE_ONLY_DISABLED.store(v, Ordering::Relaxed); +} +pub(crate) fn device_only_disabled() -> bool { + if DEVICE_ONLY_DISABLED.load(Ordering::Relaxed) { + return true; + } + // Env fallback (cached): `LAMBDA_VM_DISABLE_DEVICE_ONLY=1`. + static ENV_DISABLED: OnceLock = OnceLock::new(); + *ENV_DISABLED.get_or_init(|| { + std::env::var("LAMBDA_VM_DISABLE_DEVICE_ONLY") + .map(|v| v == "1") + .unwrap_or(false) + }) +} + +/// Stage-3 device-only gate: `true` when a table's round-1 LDE can be left +/// device-resident (host D2H skipped) because every downstream round is +/// guaranteed to take its GPU path. A strict AND of all preconditions that +/// imply the R2 composition, R3 barycentric, R4 DEEP, and R4 opening GPU paths +/// all fire and read the device LDE. The per-round `host_trace_empty` +/// hard-abort guards are the safety net: if any precondition is nonetheless +/// violated at runtime (mis-gate or transient GPU error), the prove aborts +/// loudly rather than reading the empty host trace. +/// +/// `zerofier_uniform` must be the R1-derived conservative form (all constraints +/// share `end_exemptions == 0`), which implies `ZerofierEvaluations::is_uniform` +/// (a single cyclic group) — the condition the GPU composition kernel needs. +pub(crate) fn device_only_gate( + lde_size: usize, + n: usize, + is_preprocessed: bool, + offsets_contiguous: bool, + zerofier_uniform: bool, +) -> bool +where + F: 'static, + E: 'static, +{ + // debug-checks reconstruct the LDE from the host trace — keep it resident. + if cfg!(feature = "debug-checks") { + return false; + } + TypeId::of::() == TypeId::of::() + && TypeId::of::() == TypeId::of::() + && !device_only_disabled() + && !gpu_composition_disabled() + && lde_size.is_power_of_two() + && lde_size >= gpu_lde_threshold() + && n >= gpu_bary_threshold() + && !is_preprocessed + && offsets_contiguous + && zerofier_uniform +} + // ============================================================================ // Shared dispatch helpers // ============================================================================ @@ -481,6 +564,7 @@ pub(crate) fn try_expand_leaf_and_tree_row_major_keep( m: usize, blowup_factor: usize, weights: &[FieldElement], + retain_host_lde: bool, ) -> Option<( MerkleTree, math_cuda::lde::GpuLdeBase, @@ -513,12 +597,14 @@ where GPU_MERKLE_TREE_CALLS.fetch_add(1, Ordering::Relaxed); // The keep path keeps the Merkle tree resident on device (in `handle.tree`). + // `retain_host_lde=false` additionally skips the row-major D2H (device-only). let (handle, lde_u64) = math_cuda::lde::coset_lde_row_major_with_merkle_tree_keep( raw, n, m, blowup_factor, &weights_u64, + retain_host_lde, ) .ok()?; @@ -548,6 +634,7 @@ pub(crate) fn try_expand_leaf_and_tree_ext3_row_major_keep( m: usize, blowup_factor: usize, weights: &[FieldElement], + retain_host_lde: bool, ) -> Option<( MerkleTree, math_cuda::lde::GpuLdeExt3, @@ -582,12 +669,14 @@ where GPU_MERKLE_TREE_CALLS.fetch_add(1, Ordering::Relaxed); // The keep path keeps the Merkle tree resident on device (in `handle.tree`). + // `retain_host_lde=false` additionally skips the row-major D2H (device-only). let (handle, lde_u64) = math_cuda::lde::coset_lde_ext3_row_major_with_merkle_tree_keep( raw, n, m, blowup_factor, &weights_u64, + retain_host_lde, ) .ok()?; @@ -1149,6 +1238,7 @@ pub(crate) fn try_expand_leaf_and_tree_ext3_row_major_keep_dev( ra: &math_cuda::logup::ResidentAux, blowup_factor: usize, weights: &[FieldElement], + retain_host_lde: bool, ) -> Option<( MerkleTree, math_cuda::lde::GpuLdeExt3, @@ -1176,6 +1266,7 @@ where ra.num_aux_cols, blowup_factor, &weights_u64, + retain_host_lde, ) .ok()?; diff --git a/crypto/stark/src/prover.rs b/crypto/stark/src/prover.rs index 5c03292da..27ff172d2 100644 --- a/crypto/stark/src/prover.rs +++ b/crypto/stark/src/prover.rs @@ -254,6 +254,27 @@ where ) -> Round1 { let (main_data, num_main_cols) = lde.main; let (aux_data, num_aux_cols) = lde.aux; + + // Stage-3 device-only detection, inferred from the ACTUAL buffer state + // (not the gate's intent): a table whose round-1 D2H was skipped has an + // empty host buffer where it should have data. Using the real state is a + // safety property — if the `device_only` gate held but the GPU keep path + // fell back to CPU, the buffer is populated and this stays false, so the + // proof runs on the host trace as normal. A mixed state (one buffer + // empty, the other full) is treated as device-only so any host read + // hard-aborts rather than indexing an empty buffer. + #[cfg(feature = "cuda")] + let main_empty = num_main_cols > 0 && main_data.is_empty(); + #[cfg(feature = "cuda")] + let host_trace_empty = + main_empty || (num_aux_cols > 0 && aux_data.is_empty() && lde.gpu_aux.is_some()); + #[cfg(feature = "cuda")] + let device_num_rows = lde + .gpu_main + .as_ref() + .map(|h| h.lde_size) + .or_else(|| lde.gpu_aux.as_ref().map(|h| h.lde_size)); + #[allow(unused_mut)] let mut lde_trace = LDETraceTable::from_row_major( main_data, @@ -265,6 +286,20 @@ where ); #[cfg(feature = "cuda")] { + if host_trace_empty { + // Recover the LDE row count from the resident device handle + // whenever any host buffer is empty. `from_row_major` derives + // `num_rows` from `main_data` (or `aux_data` when there are no + // main columns); if that buffer was skipped it reads 0, so we + // overwrite from the handle's `lde_size` (the true row count). + // Idempotent when `from_row_major` already got it right, and it + // covers the aux-only (`num_main_cols == 0`) device-only case + // that a `main_empty`-only guard missed. + if let Some(n) = device_num_rows { + lde_trace.set_num_rows(n); + } + lde_trace.set_host_trace_empty(true); + } if let Some(h) = lde.gpu_main { lde_trace.set_gpu_main(h); } @@ -761,11 +796,48 @@ pub trait IsStarkProver< /// as a separate Merkle tree (the precomputed split for preprocessed /// tables) and the root is checked against the AIR-hardcoded commitment. #[allow(clippy::type_complexity)] + /// Stage-3 device-only gate for one table (see + /// [`crate::gpu_lde::device_only_gate`]). Derived purely from the AIR + + /// domain so the round-1 main-commit and aux-commit closures compute the + /// identical value and skip both host D2Hs consistently — the per-table + /// `host_trace_empty` flag covers both the main and aux buffers, so they + /// must be left empty together. + #[cfg(feature = "cuda")] + fn device_only_for( + air: &dyn AIR, + domain: &Domain, + ) -> bool { + // Preconditions the downstream GPU paths require that the numeric gate + // below does not capture. A table missing either would pass the gate, + // skip its host D2H, then hard-abort in round 2: + // - R2 composition unconditionally needs a device aux handle + // (`gpu_aux()?`), so the table must declare an aux trace. + // - The composition path needs a uniform zerofier with ≥1 group. An + // empty constraint set makes `all(end_exemptions == 0)` vacuously + // true here but `is_uniform()` false downstream (0 groups). + if !air.has_aux_trace() || air.constraints_meta().is_empty() { + return false; + } + let lde_size = domain.interpolation_domain_size * domain.blowup_factor; + let n = domain.interpolation_domain_size; + let offsets = &air.context().transition_offsets; + let offsets_contiguous = offsets.iter().enumerate().all(|(i, &o)| o == i); + let zerofier_uniform = air.constraints_meta().iter().all(|m| m.end_exemptions == 0); + crate::gpu_lde::device_only_gate::( + lde_size, + n, + air.is_preprocessed(), + offsets_contiguous, + zerofier_uniform, + ) + } + fn commit_main_trace( trace: &TraceTable, domain: &Domain, twiddles: &LdeTwiddles, precomputed: Option<(Commitment, usize)>, + #[cfg(feature = "cuda")] device_only: bool, #[cfg(feature = "disk-spill")] storage_mode: StorageMode, ) -> Result, ProvingError> where @@ -798,6 +870,7 @@ pub trait IsStarkProver< num_cols, domain.blowup_factor, &twiddles.coset_weights, + !device_only, ) { #[cfg(feature = "instruments")] @@ -805,6 +878,13 @@ pub trait IsStarkProver< let root = tree.root; #[cfg(feature = "instruments")] crate::instruments::accum_r1_main(main_lde_dur, std::time::Duration::ZERO); + // Count a device-only main commit only once the GPU keep path + // actually fired (handle produced + host trace intentionally + // empty), so the counter reflects real residency, not the gate. + if device_only { + crate::gpu_lde::GPU_DEVICE_ONLY_CALLS + .fetch_add(1, std::sync::atomic::Ordering::Relaxed); + } return Ok(( TableCommit::plain(tree, root), (main_data, num_cols), @@ -1611,6 +1691,16 @@ pub trait IsStarkProver< } } + // Reaching here means both GPU DEEP arms fell through to the host loop + // below (which reads `get_main`/`get_aux`). Under the device-only gate + // the host trace is empty, so a fall-through is a mis-gate or an + // unexpected GPU failure: hard-abort rather than read empty buffers. + #[cfg(feature = "cuda")] + assert!( + !lde_trace.host_trace_empty(), + "R4 DEEP composition fell back to the host trace, but it is device-only (empty)" + ); + // OOD column compression (Plonky3-style): precompute one value per eval point, // ood_compressed_k = Σ_j gamma[j][k] * ood[j][k]. // The per-LDE-point trace column sums are NOT precomputed — they are fused @@ -1809,6 +1899,36 @@ pub trait IsStarkProver< } } + /// Build a [`PolynomialOpenings`] from a device Merkle proof and a pair of + /// row-value vectors already gathered off the resident device LDE — the + /// fully device-sourced counterpart of [`Self::open_polys_with_proofs`], + /// which still reads its evaluations from the host LDE. + #[cfg(feature = "cuda")] + fn open_polys_from_values( + proof: Proof, + evaluations: Vec>, + evaluations_sym: Vec>, + ) -> PolynomialOpenings { + PolynomialOpenings { + proof, + evaluations, + evaluations_sym, + } + } + + /// Slice out query `qi`'s even/odd row (each `ncols` field elements) from the + /// row-major device gather `[even(q0), odd(q0), even(q1), odd(q1), ...]`. + #[cfg(feature = "cuda")] + fn device_row_pair( + vals: &[FieldElement], + qi: usize, + ncols: usize, + ) -> (Vec>, Vec>) { + let even = vals[(2 * qi) * ncols..(2 * qi + 1) * ncols].to_vec(); + let odd = vals[(2 * qi + 1) * ncols..(2 * qi + 2) * ncols].to_vec(); + (even, odd) + } + /// Open the deep composition polynomial on a list of indexes and their symmetric elements. fn open_deep_composition_poly( domain: &Domain, @@ -1828,6 +1948,23 @@ pub trait IsStarkProver< let num_precomputed_cols = main_commit.num_precomputed_cols; let total_cols = lde_trace.num_main_cols(); + // Row-pair LDE positions for every query, `[even(q0), odd(q0), ...]`. + // Each query opens the leaf at `challenge`, which pairs LDE rows + // `reverse_index(2·challenge)` (the queried point) and + // `reverse_index(2·challenge+1)` (its symmetric `-x` point). + #[cfg(feature = "cuda")] + let domain_size = domain.lde_roots_of_unity_coset.len() as u64; + #[cfg(feature = "cuda")] + let query_rows: Vec = indexes_to_open + .iter() + .flat_map(|&c| { + [ + reverse_index(c * 2, domain_size) as u32, + reverse_index(c * 2 + 1, domain_size) as u32, + ] + }) + .collect(); + // R4 trace proofs from the resident device trees, gathered in one batch // over all query positions instead of walking the host trees (byte // identical to the host proofs, guarded by the `merkle_gather` test). @@ -1882,6 +2019,83 @@ pub trait IsStarkProver< ) }); + // Full-residency Stage 2: gather each query's row-pair straight off the + // resident device LDE (a small D2H of only the queried rows), instead of + // indexing the full host LDE trace. The host trace is still resident this + // stage and every device row is cross-checked against it in the loop + // below; Stage 3 drops the host copy once this path is proven. `None` + // when the LDE is not device resident or the tower is not Goldilocks (→ + // host gather). Row-major: `q`-th row's columns at `[q*ncols ..]`. + // Gate the value gathers on the corresponding device-tree proofs so the + // two device arms stay aligned (`*_dev_proofs.is_some() ⇔ + // *_dev_values.is_some()` on the Goldilocks path) and we never gather + // rows for a tree that is not device resident. + #[cfg(feature = "cuda")] + let main_dev_values: Option>> = + main_dev_proofs.as_ref().and_then(|_| { + lde_trace.gpu_main().and_then(|h| { + let stream = lde_trace + .bound_stream() + .expect("bound stream for device-resident main-row gather"); + let raw = match math_cuda::barycentric::gather_rows_base_on_device( + h, + &query_rows, + &stream, + ) { + Ok(v) => v, + // A gather failure is only fatal under device-only (no host + // trace to fall back to). Otherwise return `None` so the host + // gather arms below serve the openings from the resident LDE. + Err(e) => { + assert!( + !lde_trace.host_trace_empty(), + "device main-row gather failed and the trace is device-only \ + (no host fallback): {e:?}" + ); + return None; + } + }; + crate::constraint_ir::gpu_interp::base_u64_to_field::(&raw) + }) + }); + #[cfg(feature = "cuda")] + if main_dev_values.is_some() { + crate::gpu_lde::GPU_OPENING_GATHER_CALLS + .fetch_add(1, std::sync::atomic::Ordering::Relaxed); + } + + #[cfg(feature = "cuda")] + let aux_dev_values: Option>> = + aux_dev_proofs.as_ref().and_then(|_| { + lde_trace.gpu_aux().and_then(|h| { + let stream = lde_trace + .bound_stream() + .expect("bound stream for device-resident aux-row gather"); + let raw = match math_cuda::barycentric::gather_rows_ext3_on_device( + h, + &query_rows, + &stream, + ) { + Ok(v) => v, + // Fatal only under device-only; otherwise fall back to host. + Err(e) => { + assert!( + !lde_trace.host_trace_empty(), + "device aux-row gather failed and the trace is device-only \ + (no host fallback): {e:?}" + ); + return None; + } + }; + crate::constraint_ir::gpu_interp::ext3_u64_to_field::(&raw) + }) + }); + #[cfg(feature = "cuda")] + if aux_dev_values.is_some() { + crate::gpu_lde::GPU_OPENING_GATHER_CALLS + .fetch_add(1, std::sync::atomic::Ordering::Relaxed); + } + for (qi, index) in indexes_to_open.iter().enumerate() { #[cfg(not(feature = "cuda"))] let _ = qi; @@ -1895,10 +2109,44 @@ pub trait IsStarkProver< #[cfg(feature = "cuda")] { if let Some(proofs) = &main_dev_proofs { - Self::open_polys_with_proofs(domain, proofs[qi].clone(), *index, |row| { - lde_trace.gather_main_row(row) - }) + let proof = proofs[qi].clone(); + if let Some(dev_vals) = &main_dev_values { + let (even, odd) = Self::device_row_pair(dev_vals, qi, total_cols); + // Cross-check the device gather against the host LDE. + // Skipped under device-only (host trace empty): the + // gather was proven bit-identical while the host copy + // was resident, and there is nothing to check against. + if !lde_trace.host_trace_empty() { + let r_even = reverse_index(*index * 2, domain_size); + let r_odd = reverse_index(*index * 2 + 1, domain_size); + assert_eq!( + even, + lde_trace.gather_main_row(r_even), + "device main-row gather mismatch (even), query {qi}" + ); + assert_eq!( + odd, + lde_trace.gather_main_row(r_odd), + "device main-row gather mismatch (odd), query {qi}" + ); + } + Self::open_polys_from_values(proof, even, odd) + } else { + // Device tree resident but the value gather is absent: + // reading the host gather is invalid under device-only. + assert!( + !lde_trace.host_trace_empty(), + "R4 main opening fell back to the host gather, but it is device-only (empty)" + ); + Self::open_polys_with_proofs(domain, proof, *index, |row| { + lde_trace.gather_main_row(row) + }) + } } else { + assert!( + !lde_trace.host_trace_empty(), + "R4 main opening fell back to the host tree, but it is device-only (empty)" + ); Self::open_polys_with(domain, &main_commit.tree, *index, |row| { lde_trace.gather_main_row(row) }) @@ -1950,10 +2198,40 @@ pub trait IsStarkProver< #[cfg(feature = "cuda")] { if let Some(proofs) = &aux_dev_proofs { - Self::open_polys_with_proofs(domain, proofs[qi].clone(), *index, |row| { - lde_trace.gather_aux_row(row) - }) + let proof = proofs[qi].clone(); + if let Some(dev_vals) = &aux_dev_values { + let ncols = lde_trace.num_aux_cols(); + let (even, odd) = Self::device_row_pair(dev_vals, qi, ncols); + // Cross-check skipped under device-only (host empty). + if !lde_trace.host_trace_empty() { + let r_even = reverse_index(*index * 2, domain_size); + let r_odd = reverse_index(*index * 2 + 1, domain_size); + assert_eq!( + even, + lde_trace.gather_aux_row(r_even), + "device aux-row gather mismatch (even), query {qi}" + ); + assert_eq!( + odd, + lde_trace.gather_aux_row(r_odd), + "device aux-row gather mismatch (odd), query {qi}" + ); + } + Self::open_polys_from_values(proof, even, odd) + } else { + assert!( + !lde_trace.host_trace_empty(), + "R4 aux opening fell back to the host gather, but it is device-only (empty)" + ); + Self::open_polys_with_proofs(domain, proof, *index, |row| { + lde_trace.gather_aux_row(row) + }) + } } else { + assert!( + !lde_trace.host_trace_empty(), + "R4 aux opening fell back to the host tree, but it is device-only (empty)" + ); Self::open_polys_with(domain, &aux.tree, *index, |row| { lde_trace.gather_aux_row(row) }) @@ -2155,11 +2433,19 @@ pub trait IsStarkProver< let precomputed = air .is_preprocessed() .then(|| (air.precomputed_commitment(), air.num_precomputed_columns())); + + // Stage-3 device-only gate: when it holds, `commit_main_trace` + // keeps the R1 LDE device-resident and skips the host D2H. + #[cfg(feature = "cuda")] + let device_only = Self::device_only_for(*air, domain); + Self::commit_main_trace( *trace, domain, twiddles, precomputed, + #[cfg(feature = "cuda")] + device_only, #[cfg(feature = "disk-spill")] storage_mode, ) @@ -2363,6 +2649,12 @@ pub trait IsStarkProver< if air.has_aux_trace() { let lde_size = domain.interpolation_domain_size * domain.blowup_factor; + // Same gate as the main commit (Phase A): skip the aux + // host D2H when device-only, so both buffers are left + // empty together for this table. + #[cfg(feature = "cuda")] + let device_only = Self::device_only_for(*air, domain); + // Resident GPU path: aux columns already on device (from // the resident LogUp aux build) — LDE straight from device // memory, no upload, no host column extraction. When the @@ -2379,7 +2671,10 @@ pub trait IsStarkProver< FieldExtension, BatchedMerkleTreeBackend, >( - ra, domain.blowup_factor, &twiddles.coset_weights + ra, + domain.blowup_factor, + &twiddles.coset_weights, + !device_only, ) .ok_or_else(|| { ProvingError::Fft( @@ -2421,6 +2716,7 @@ pub trait IsStarkProver< num_cols, domain.blowup_factor, &twiddles.coset_weights, + !device_only, ) { #[cfg(feature = "instruments")] diff --git a/crypto/stark/src/trace.rs b/crypto/stark/src/trace.rs index 6d40425b7..280aabb95 100644 --- a/crypto/stark/src/trace.rs +++ b/crypto/stark/src/trace.rs @@ -327,6 +327,14 @@ where pub(crate) num_rows: usize, pub(crate) lde_step_size: usize, pub(crate) blowup_factor: usize, + /// Full-residency (Stage 3): when true the round-1 D2H was intentionally + /// skipped and `main_data`/`aux_data` are empty — every round reads the LDE + /// off the device instead. Any code path that would read the host trace must + /// hard-abort on this flag rather than index an empty buffer, so a mis-gate + /// or an unexpected GPU fallback fails loudly instead of producing a wrong + /// proof. Always false today until the `device_only` gate is wired. + #[cfg(feature = "cuda")] + pub(crate) host_trace_empty: bool, /// Per table GPU residency session: owns this table's device LDE buffers /// and bound stream. Threaded R1 to R4. Empty on the CPU path. #[cfg(feature = "cuda")] @@ -459,6 +467,8 @@ where lde_step_size, blowup_factor, #[cfg(feature = "cuda")] + host_trace_empty: false, + #[cfg(feature = "cuda")] gpu_session: GpuTableSession::new(), } } @@ -494,6 +504,8 @@ where lde_step_size, blowup_factor, #[cfg(feature = "cuda")] + host_trace_empty: false, + #[cfg(feature = "cuda")] gpu_session: GpuTableSession::new(), } } @@ -511,6 +523,31 @@ where self.gpu_session.aux_lde = Some(h); } + /// Mark this table's host LDE trace as intentionally empty (Stage-3 + /// device-only path): the round-1 D2H was skipped and every host-trace read + /// must hard-abort instead of indexing the empty buffers. + #[cfg(feature = "cuda")] + pub fn set_host_trace_empty(&mut self, empty: bool) { + self.host_trace_empty = empty; + } + + /// Override the LDE row count. Needed on the device-only path: the host + /// buffers are empty, so `from_row_major` cannot infer `num_rows` from + /// `main_data.len()` — the caller supplies it from the device handle's + /// `lde_size` instead. + #[cfg(feature = "cuda")] + pub fn set_num_rows(&mut self, num_rows: usize) { + self.num_rows = num_rows; + } + + /// Whether the host LDE trace was intentionally left empty (see + /// [`Self::set_host_trace_empty`]). Guards on every host-read fallback check + /// this before touching `main_data`/`aux_data`. + #[cfg(feature = "cuda")] + pub fn host_trace_empty(&self) -> bool { + self.host_trace_empty + } + #[cfg(feature = "cuda")] pub fn gpu_main(&self) -> Option<&math_cuda::lde::GpuLdeBase> { self.gpu_session.main_lde.as_ref() @@ -742,6 +779,13 @@ where let main_evals: Vec> = if let Some(v) = main_gpu { v } else { + // Device-only tables have no host trace; a GPU fall-through here would + // read empty `main_data`. Hard-abort instead of a wrong OOD eval. + #[cfg(feature = "cuda")] + assert!( + !lde_trace.host_trace_empty(), + "R3 barycentric (main) fell back to the host trace, but it is device-only (empty)" + ); let inv_denoms_v = inv_denoms.get_or_insert_with(|| barycentric_inv_denoms(eval_point, &dc.points)); let col_scale = col_scale.get_or_insert_with(|| { @@ -793,6 +837,13 @@ where let aux_evals: Vec> = if let Some(v) = aux_gpu { v } else { + // Device-only tables have no host trace; a GPU fall-through here would + // read empty `aux_data`. Hard-abort instead of a wrong OOD eval. + #[cfg(feature = "cuda")] + assert!( + !lde_trace.host_trace_empty(), + "R3 barycentric (aux) fell back to the host trace, but it is device-only (empty)" + ); let inv_denoms_v = inv_denoms.get_or_insert_with(|| barycentric_inv_denoms(eval_point, &dc.points)); let col_scale = col_scale.get_or_insert_with(|| { diff --git a/prover/tests/cuda_path_integration.rs b/prover/tests/cuda_path_integration.rs index 7081baadf..b60cb3a34 100644 --- a/prover/tests/cuda_path_integration.rs +++ b/prover/tests/cuda_path_integration.rs @@ -12,8 +12,8 @@ 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_extend_halves_calls, gpu_fri_calls, gpu_lde_calls, gpu_logup_calls, - gpu_parts_lde_calls, reset_all_gpu_call_counters, + 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 @@ -153,3 +153,49 @@ fn gpu_proof_verifies_row_pair_commitment() { "GPU-produced proof (row-pair commitment) failed verification" ); } + +/// The full-residency Stage-2 R4 opening path fires: query row values are +/// gathered straight off the device LDE (not the host trace) and, guarded by the +/// in-prover cross-check against the host gather, still yield a verifying proof. +/// Guards a silent regression where the openings quietly revert to the host +/// path (which would still verify but drop the data-residency win). The +/// cross-check `assert_eq!`s inside `open_deep_composition_poly` are +/// release-active, so a divergent device gather would panic the prove here. +#[test] +#[ignore = "requires GPU; run with --ignored --nocapture"] +fn gpu_opening_gather_fires_and_verifies() { + let elf = asm_elf_bytes("fib_iterative_1M"); + reset_all_gpu_call_counters(); + let proof = prove(&elf).expect("prove"); + assert!( + gpu_opening_gather_calls() > 0, + "device-resident opening gather did not fire (openings fell back to the host LDE)" + ); + assert!( + verify(&proof, &elf).expect("verify"), + "GPU-produced proof (device-gathered openings) failed verification" + ); +} + +/// The full-residency Stage-3 device-only path fires: at least one table keeps +/// its round-1 LDE device-resident (the host D2H is skipped), and the proof +/// still verifies. This exercises every `host_trace_empty` hard-abort guard on +/// the happy path (none may fire) plus the GPU-only R2/R3/R4 paths reading the +/// device LDE with no host trace behind them. A regression that silently +/// reverts to the host D2H drops the counter to 0 (while the proof would still +/// verify), and a mis-gate that forces a host fallback panics one of the guards. +#[test] +#[ignore = "requires GPU; run with --ignored --nocapture"] +fn gpu_device_only_residency_fires_and_verifies() { + let elf = asm_elf_bytes("fib_iterative_1M"); + reset_all_gpu_call_counters(); + let proof = prove(&elf).expect("prove"); + assert!( + gpu_device_only_calls() > 0, + "device-only residency path did not fire (every table kept its host trace)" + ); + assert!( + verify(&proof, &elf).expect("verify"), + "GPU-produced proof (device-only residency) failed verification" + ); +} diff --git a/scripts/bench_abba.sh b/scripts/bench_abba.sh index 9acc7ae86..fd0e7ad3e 100755 --- a/scripts/bench_abba.sh +++ b/scripts/bench_abba.sh @@ -30,6 +30,11 @@ # Env: REBUILD=1 forces a rebuild even if cached binaries exist. # BENCH_FEATURES= cargo features for the cli build (default: jemalloc-stats). # The GPU ABBA workflow passes "jemalloc-stats,prover/cuda" to bench the GPU path. +# CONTINUATIONS=1 proves with --continuations (epochs; flat peak memory) on +# both sides — needed for large workloads that OOM monolithically. +# EPOCH_SIZE_LOG2= continuation epoch size (default 20; min 18). +# TX_COUNT= ethrex transfer fixture to prove (default 5; use 20 for a +# large continuation trace where GPU-residency wins are visible). # # Sizing (ethrex pair-noise sd ~1.2%, 80% power): ~12 pairs for a 1% effect, # ~18 for 0.8%, ~32 for 0.6%. Default 20 -> solid on 0.8-1%, ~60% power at 0.6% @@ -51,9 +56,23 @@ N_PAIRS="${3:-20}" # cli build features. Default matches the CPU bench; the GPU ABBA workflow overrides # with "jemalloc-stats,prover/cuda" to exercise the CUDA prover path. BENCH_FEATURES="${BENCH_FEATURES:-jemalloc-stats}" +# Continuation mode: split execution into epochs (flat peak memory) instead of a +# single monolithic prove. Required for large workloads (e.g. ethrex-20tx) that +# OOM monolithically, and the only way to bench a per-epoch GPU-residency change +# at a realistic trace size. When on, both binaries prove with +# `--continuations --epoch-size-log2 $EPOCH_SIZE_LOG2`. +CONTINUATIONS="${CONTINUATIONS:-0}" +EPOCH_SIZE_LOG2="${EPOCH_SIZE_LOG2:-20}" +# ethrex transfer-count fixture to prove (executor/tests/ethrex_${TX_COUNT}_transfers.bin). +TX_COUNT="${TX_COUNT:-5}" +if [ "$CONTINUATIONS" = "1" ]; then + CONT_ARGS="--continuations --epoch-size-log2 $EPOCH_SIZE_LOG2" +else + CONT_ARGS="" +fi ELF_REL="executor/program_artifacts/rust/ethrex.elf" -INPUT_REL="executor/tests/ethrex_5_transfers.bin" +INPUT_REL="executor/tests/ethrex_${TX_COUNT}_transfers.bin" WORK="/tmp/abba_run" WT="/tmp/abba_wt" PROOF="/tmp/abba_proof.bin" @@ -84,9 +103,9 @@ if [ ! -f "$ELF_REL" ]; then make "$ELF_REL" fi if [ ! -f "$INPUT_REL" ]; then - echo "==> Generating ethrex 5-transfer fixture (missing)" + echo "==> Generating ethrex ${TX_COUNT}-transfer fixture (missing)" ( cd tooling/ethrex-fixtures && cargo build --release ) - tooling/ethrex-fixtures/target/release/ethrex-fixtures 5 "$INPUT_REL" distinct + tooling/ethrex-fixtures/target/release/ethrex-fixtures "$TX_COUNT" "$INPUT_REL" distinct fi ELF="$(cd "$(dirname "$ELF_REL")" && pwd)/$(basename "$ELF_REL")" INPUT="$(cd "$(dirname "$INPUT_REL")" && pwd)/$(basename "$INPUT_REL")" @@ -143,7 +162,8 @@ fi # --- 3. Interleaved A/B/B/A measurement (fresh CSV -- pre-committed batch) --- run_prove() { # $1=binary -> echoes proving time (s) local out t - out="$("$1" prove "$ELF" --private-input "$INPUT" -o "$PROOF" --time 2>&1)" + # shellcheck disable=SC2086 # CONT_ARGS is intentionally word-split (0 or 2 args) + out="$("$1" prove "$ELF" --private-input "$INPUT" -o "$PROOF" --time $CONT_ARGS 2>&1)" rm -f "$PROOF" t="$(printf '%s\n' "$out" | grep -o 'Proving time: [0-9.]*' | awk '{print $3}')" if [ -z "$t" ]; then