diff --git a/Makefile b/Makefile index 25dce43de..40dc04b3d 100644 --- a/Makefile +++ b/Makefile @@ -565,6 +565,8 @@ test-cuda-integration: test-cuda-fallback: cargo test -p lambda-vm-prover --release --features test-cuda-faults \ --test cuda_fallback_tests -- --ignored --nocapture --test-threads=1 + cargo test -p lambda-vm-prover --release --features lambda-vm-prover/cuda \ + --test gpu_force_downgrade -- --ignored --nocapture --test-threads=1 # The prover/stark/crypto/ecsm test suite with the GPU (cuda) path enabled (requires NVIDIA # GPU + nvcc). The GPU CI counterpart of CPU CI's sharded prover tests. Single-threaded: the diff --git a/crypto/math-cuda/build.rs b/crypto/math-cuda/build.rs index fbd70eb5b..bbb9943b9 100644 --- a/crypto/math-cuda/build.rs +++ b/crypto/math-cuda/build.rs @@ -72,6 +72,13 @@ fn to_real_arch(arch: &str) -> String { } } +/// Single source for the barycentric multi-kernel eval-point cap. The CUDA +/// side sizes a per-thread accumulator array with it (`BARY_MAX_K`, passed via +/// `-D` below) and the Rust dispatch asserts against it (generated into +/// `bary_consts.rs`) — defining it twice invites stack corruption in the +/// kernel the day one side moves without the other. +const BARY_MAX_EVAL_POINTS: usize = 8; + fn compile_kernel(src: &str, out_name: &str, have_nvcc: bool) { let manifest_dir = PathBuf::from(env::var("CARGO_MANIFEST_DIR").unwrap()); let out_dir = PathBuf::from(env::var("OUT_DIR").unwrap()); @@ -118,6 +125,7 @@ fn compile_kernel(src: &str, out_name: &str, have_nvcc: bool) { let mut cmd = Command::new(nvcc_path()); cmd.args(["--cubin", "-O3", "-std=c++17", "-arch", &arch]); + cmd.arg(format!("-DBARY_MAX_K={BARY_MAX_EVAL_POINTS}")); // SASS→source line mapping for Nsight Compute. Unlike -G this does not // change codegen, but keep it opt-in so production cubins stay byte-stable. if env::var("LAMBDA_VM_NVCC_LINEINFO").is_ok_and(|v| v != "0" && !v.is_empty()) { @@ -136,6 +144,19 @@ fn compile_kernel(src: &str, out_name: &str, have_nvcc: bool) { } fn main() { + // Rust-side mirror of the kernel cap; see BARY_MAX_EVAL_POINTS above. + let out_dir = PathBuf::from(env::var("OUT_DIR").unwrap()); + fs::write( + out_dir.join("bary_consts.rs"), + format!( + "/// Compile-time cap of the multi kernels' per-thread accumulator array\n\ + /// (`BARY_MAX_K` in barycentric.cu — single-sourced from build.rs).\n\ + /// Callers with more evaluation points fall back to the per-point kernels.\n\ + pub const BARY_MAX_EVAL_POINTS: usize = {BARY_MAX_EVAL_POINTS};\n" + ), + ) + .expect("failed to write bary_consts.rs"); + // Headers aren't compiled, so emit rerun-if-changed to rebuild on // header edits. println!("cargo:rerun-if-changed=kernels/goldilocks.cuh"); diff --git a/crypto/math-cuda/kernels/barycentric.cu b/crypto/math-cuda/kernels/barycentric.cu index f76db471a..a9da64b23 100644 --- a/crypto/math-cuda/kernels/barycentric.cu +++ b/crypto/math-cuda/kernels/barycentric.cu @@ -191,6 +191,137 @@ extern "C" __global__ void barycentric_ext3_batched_strided( } } +// Multi-eval-point + row-chunked barycentric. Two fixes over the *_strided +// kernels above: (1) the LDE column data is read ONCE for all K evaluation +// points (K inv_denom blocks, K accumulators) instead of once per point, and +// (2) each column is split into `num_chunks` row ranges so the grid is +// `num_cols * num_chunks` blocks instead of `num_cols` — the single-block-per- +// column grid left most SMs idle at typical column counts. Blocks emit partial +// sums; `barycentric_combine_partials` folds the chunk axis. +// +// `inv_denoms` holds K contiguous blocks of 3N u64 (ext3 interleaved), one per +// evaluation point — the layout `compute_and_invert_denoms_ext3_dev` already +// produces. Partials layout: `[(k*num_cols + col)*num_chunks + chunk]` ext3 +// interleaved, so the combine pass reads each (k, col)'s chunks contiguously. +#ifndef BARY_MAX_K +#error "BARY_MAX_K must be passed by build.rs (-DBARY_MAX_K=...) — single-sourced there" +#endif + +extern "C" __global__ void barycentric_base_strided_multi( + const uint64_t *columns, + uint64_t col_stride, + uint64_t row_stride, + const uint64_t *coset_points, + const uint64_t *inv_denoms, + uint64_t n, + uint64_t k_points, + uint64_t num_chunks, + uint64_t *partials +) { + uint64_t col = blockIdx.x; + uint64_t chunk = blockIdx.y; + const uint64_t *col_data = columns + col * col_stride; + uint64_t chunk_len = (n + num_chunks - 1) / num_chunks; + uint64_t start = chunk * chunk_len; + uint64_t end = start + chunk_len < n ? start + chunk_len : n; + + ext3::Fe3 acc[BARY_MAX_K]; + for (uint32_t k = 0; k < k_points; ++k) acc[k] = ext3::zero(); + + for (uint64_t i = start + threadIdx.x; i < end; i += BARY_BLOCK_DIM) { + uint64_t eval = col_data[i * row_stride]; + uint64_t point = coset_points[i]; + uint64_t pe = goldilocks::mul(point, eval); + for (uint32_t k = 0; k < k_points; ++k) { + const uint64_t *inv = inv_denoms + (uint64_t)k * 3 * n + i * 3; + ext3::Fe3 inv_d = ext3::make(inv[0], inv[1], inv[2]); + acc[k] = ext3::add(acc[k], ext3::mul_base(inv_d, pe)); + } + } + + for (uint32_t k = 0; k < k_points; ++k) { + ext3::Fe3 sum = block_reduce_ext3(acc[k]); + if (threadIdx.x == 0) { + uint64_t o = ((k * gridDim.x + col) * num_chunks + chunk) * 3; + partials[o + 0] = sum.a; + partials[o + 1] = sum.b; + partials[o + 2] = sum.c; + } + // block_reduce_ext3 reuses its shared buffers: every thread must be + // done reading round k's result before round k+1 overwrites them. + __syncthreads(); + } +} + +extern "C" __global__ void barycentric_ext3_strided_multi( + const uint64_t *columns, + uint64_t col_stride, + uint64_t row_stride, + const uint64_t *coset_points, + const uint64_t *inv_denoms, + uint64_t n, + uint64_t k_points, + uint64_t num_chunks, + uint64_t *partials +) { + uint64_t col = blockIdx.x; + uint64_t chunk = blockIdx.y; + const uint64_t *slab_a = columns + (col * 3 + 0) * col_stride; + const uint64_t *slab_b = columns + (col * 3 + 1) * col_stride; + const uint64_t *slab_c = columns + (col * 3 + 2) * col_stride; + uint64_t chunk_len = (n + num_chunks - 1) / num_chunks; + uint64_t start = chunk * chunk_len; + uint64_t end = start + chunk_len < n ? start + chunk_len : n; + + ext3::Fe3 acc[BARY_MAX_K]; + for (uint32_t k = 0; k < k_points; ++k) acc[k] = ext3::zero(); + + for (uint64_t i = start + threadIdx.x; i < end; i += BARY_BLOCK_DIM) { + uint64_t lde_i = i * row_stride; + ext3::Fe3 eval = ext3::make(slab_a[lde_i], slab_b[lde_i], slab_c[lde_i]); + uint64_t point = coset_points[i]; + ext3::Fe3 pe = ext3::mul_base(eval, point); + for (uint32_t k = 0; k < k_points; ++k) { + const uint64_t *inv = inv_denoms + (uint64_t)k * 3 * n + i * 3; + ext3::Fe3 inv_d = ext3::make(inv[0], inv[1], inv[2]); + acc[k] = ext3::add(acc[k], ext3::mul(pe, inv_d)); + } + } + + for (uint32_t k = 0; k < k_points; ++k) { + ext3::Fe3 sum = block_reduce_ext3(acc[k]); + if (threadIdx.x == 0) { + uint64_t o = ((k * gridDim.x + col) * num_chunks + chunk) * 3; + partials[o + 0] = sum.a; + partials[o + 1] = sum.b; + partials[o + 2] = sum.c; + } + __syncthreads(); + } +} + +// Fold the chunk axis of the multi kernels' partials: one thread per +// (k, col) pair sums its `num_chunks` ext3 partials sequentially (the whole +// buffer is tiny — K * cols * chunks). Output `out_ext3_int[k*num_cols+col]`, +// same per-column layout as the single-point kernels, K blocks concatenated. +extern "C" __global__ void barycentric_combine_partials( + const uint64_t *partials, + uint64_t num_chunks, + uint64_t total, + uint64_t *out_ext3_int +) { + uint64_t idx = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x; + if (idx >= total) return; + const uint64_t *row = partials + idx * num_chunks * 3; + ext3::Fe3 acc = ext3::zero(); + for (uint64_t c = 0; c < num_chunks; ++c) { + acc = ext3::add(acc, ext3::make(row[c * 3 + 0], row[c * 3 + 1], row[c * 3 + 2])); + } + out_ext3_int[idx * 3 + 0] = acc.a; + out_ext3_int[idx * 3 + 1] = acc.b; + out_ext3_int[idx * 3 + 2] = acc.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 diff --git a/crypto/math-cuda/src/barycentric.rs b/crypto/math-cuda/src/barycentric.rs index e9aceaea2..cce002ac0 100644 --- a/crypto/math-cuda/src/barycentric.rs +++ b/crypto/math-cuda/src/barycentric.rs @@ -350,6 +350,163 @@ pub fn barycentric_ext3_on_device_with_dev_inv_denoms( Ok(out) } +include!(concat!(env!("OUT_DIR"), "/bary_consts.rs")); + +/// Row-chunk count for the multi kernels: enough `cols * chunks` blocks to +/// occupy the device, without shrinking a chunk's row range below the point +/// where launch + combine overhead dominates. +fn bary_num_chunks(num_cols: usize, n: usize) -> usize { + let by_occupancy = (2048 / num_cols.max(1)).max(1); + let by_rows = (n / 8192).max(1); + by_occupancy.min(by_rows).min(64) +} + +/// Multi-eval-point counterpart of +/// [`barycentric_base_on_device_with_dev_inv_denoms`]: one pass over the LDE +/// column data computes the barycentric sums for ALL `k_points` evaluation +/// points (their inv_denom blocks live contiguously in `inv_denoms_dev`, the +/// layout `compute_and_invert_denoms_ext3_dev` produces). Returns +/// `3 * k_points * num_cols` u64: `k_points` concatenated per-column blocks, +/// each in the same layout as the single-point kernels. +pub fn barycentric_base_multi_on_device( + stream: &Arc, + main_handle: &GpuLdeBase, + row_stride: usize, + coset_points_dev: &CudaSlice, + inv_denoms_dev: &CudaSlice, + n: usize, + k_points: usize, +) -> Result> { + main_handle.wait_ready_on(stream)?; + assert!((1..=BARY_MAX_EVAL_POINTS).contains(&k_points)); + assert!(coset_points_dev.len() >= n); + assert!(inv_denoms_dev.len() >= k_points * 3 * n); + let num_cols = main_handle.m; + if num_cols == 0 || n == 0 { + return Ok(vec![0; 3 * k_points * num_cols]); + } + let be = backend()?; + let num_chunks = bary_num_chunks(num_cols, n); + let total = k_points * num_cols; + let mut partials = stream.alloc_zeros::(total * num_chunks * 3)?; + let mut out_dev = stream.alloc_zeros::(3 * total)?; + let points_view = coset_points_dev.slice(0..n); + let inv_view = inv_denoms_dev.slice(0..k_points * 3 * n); + + let col_stride_u64 = main_handle.lde_size as u64; + let row_stride_u64 = row_stride as u64; + let n_u64 = n as u64; + let k_u64 = k_points as u64; + let chunks_u64 = num_chunks as u64; + let total_u64 = total as u64; + let cfg = LaunchConfig { + grid_dim: (num_cols as u32, num_chunks as u32, 1), + block_dim: (BLOCK_DIM, 1, 1), + shared_mem_bytes: 0, + }; + unsafe { + stream + .launch_builder(&be.barycentric_base_strided_multi) + .arg(main_handle.buf.as_ref()) + .arg(&col_stride_u64) + .arg(&row_stride_u64) + .arg(&points_view) + .arg(&inv_view) + .arg(&n_u64) + .arg(&k_u64) + .arg(&chunks_u64) + .arg(&mut partials) + .launch(cfg)?; + } + let combine_cfg = LaunchConfig { + grid_dim: (total.div_ceil(BLOCK_DIM as usize) as u32, 1, 1), + block_dim: (BLOCK_DIM, 1, 1), + shared_mem_bytes: 0, + }; + unsafe { + stream + .launch_builder(&be.barycentric_combine_partials) + .arg(&partials) + .arg(&chunks_u64) + .arg(&total_u64) + .arg(&mut out_dev) + .launch(combine_cfg)?; + } + let out = stream.clone_dtoh(&out_dev)?; + stream.synchronize()?; + Ok(out) +} + +/// Ext3 counterpart of [`barycentric_base_multi_on_device`]. +pub fn barycentric_ext3_multi_on_device( + stream: &Arc, + aux_handle: &GpuLdeExt3, + row_stride: usize, + coset_points_dev: &CudaSlice, + inv_denoms_dev: &CudaSlice, + n: usize, + k_points: usize, +) -> Result> { + aux_handle.wait_ready_on(stream)?; + assert!((1..=BARY_MAX_EVAL_POINTS).contains(&k_points)); + assert!(coset_points_dev.len() >= n); + assert!(inv_denoms_dev.len() >= k_points * 3 * n); + let num_cols = aux_handle.m; + if num_cols == 0 || n == 0 { + return Ok(vec![0; 3 * k_points * num_cols]); + } + let be = backend()?; + let num_chunks = bary_num_chunks(num_cols, n); + let total = k_points * num_cols; + let mut partials = stream.alloc_zeros::(total * num_chunks * 3)?; + let mut out_dev = stream.alloc_zeros::(3 * total)?; + let points_view = coset_points_dev.slice(0..n); + let inv_view = inv_denoms_dev.slice(0..k_points * 3 * n); + + let col_stride_u64 = aux_handle.lde_size as u64; + let row_stride_u64 = row_stride as u64; + let n_u64 = n as u64; + let k_u64 = k_points as u64; + let chunks_u64 = num_chunks as u64; + let total_u64 = total as u64; + let cfg = LaunchConfig { + grid_dim: (num_cols as u32, num_chunks as u32, 1), + block_dim: (BLOCK_DIM, 1, 1), + shared_mem_bytes: 0, + }; + unsafe { + stream + .launch_builder(&be.barycentric_ext3_strided_multi) + .arg(aux_handle.buf.as_ref()) + .arg(&col_stride_u64) + .arg(&row_stride_u64) + .arg(&points_view) + .arg(&inv_view) + .arg(&n_u64) + .arg(&k_u64) + .arg(&chunks_u64) + .arg(&mut partials) + .launch(cfg)?; + } + let combine_cfg = LaunchConfig { + grid_dim: (total.div_ceil(BLOCK_DIM as usize) as u32, 1, 1), + block_dim: (BLOCK_DIM, 1, 1), + shared_mem_bytes: 0, + }; + unsafe { + stream + .launch_builder(&be.barycentric_combine_partials) + .arg(&partials) + .arg(&chunks_u64) + .arg(&total_u64) + .arg(&mut out_dev) + .launch(combine_cfg)?; + } + let out = stream.clone_dtoh(&out_dev)?; + 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 diff --git a/crypto/math-cuda/src/device.rs b/crypto/math-cuda/src/device.rs index a7c129cc8..1ddb8927c 100644 --- a/crypto/math-cuda/src/device.rs +++ b/crypto/math-cuda/src/device.rs @@ -207,6 +207,9 @@ pub struct Backend { pub barycentric_ext3_batched: CudaFunction, pub barycentric_base_batched_strided: CudaFunction, pub barycentric_ext3_batched_strided: CudaFunction, + pub barycentric_base_strided_multi: CudaFunction, + pub barycentric_ext3_strided_multi: CudaFunction, + pub barycentric_combine_partials: CudaFunction, pub gather_rows_base: CudaFunction, pub gather_rows_ext3: CudaFunction, @@ -438,6 +441,9 @@ impl Backend { .load_function("barycentric_base_batched_strided")?, barycentric_ext3_batched_strided: bary .load_function("barycentric_ext3_batched_strided")?, + barycentric_base_strided_multi: bary.load_function("barycentric_base_strided_multi")?, + barycentric_ext3_strided_multi: bary.load_function("barycentric_ext3_strided_multi")?, + barycentric_combine_partials: bary.load_function("barycentric_combine_partials")?, 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")?, diff --git a/crypto/math-cuda/src/lde.rs b/crypto/math-cuda/src/lde.rs index 3d8bfa207..9bbd9958d 100644 --- a/crypto/math-cuda/src/lde.rs +++ b/crypto/math-cuda/src/lde.rs @@ -666,19 +666,26 @@ fn coset_lde_row_major_inner( /// the whole tree copy to host is eliminated; query openings gather paths from /// the device tree. /// -/// Input: `row_major` is a flat `n * m` slice in row-major order. Returns the -/// `GpuLdeBase` handle (column-major buf, plus the device tree) and the -/// row-major LDE Vec. +/// Input: `row_major` is a flat `n * m` slice in row-major order; when +/// `predev` carries the same data already on device (pre-uploaded off the +/// critical path), the expansion D2D-copies from it instead of a fresh H2D. +/// Returns the `GpuLdeBase` handle (column-major buf, plus the device tree) +/// and the row-major LDE Vec. pub fn coset_lde_row_major_with_merkle_tree_keep( row_major: &[u64], + predev: Option<&CudaSlice>, n: usize, m: usize, blowup_factor: usize, weights: &[u64], retain_host_lde: bool, ) -> Result<(GpuLdeBase, Vec)> { + let input = match predev { + Some(d) if d.len() == row_major.len() => InnerInput::Dev(d), + _ => InnerInput::Host(row_major), + }; let (tree, col_major_dev, lde_out, trace_col_major, ready) = coset_lde_row_major_inner( - InnerInput::Host(row_major), + input, n, m, blowup_factor, @@ -715,14 +722,17 @@ pub fn coset_lde_row_major_with_merkle_tree_keep( /// Returns `(precomputed_nodes, handle, row_major_lde)`. The handle also /// carries the column-major LDE + trace snapshot for downstream GPU rounds. #[allow(clippy::type_complexity)] +#[allow(clippy::too_many_arguments)] pub fn coset_lde_row_major_split_trees( row_major: &[u64], + predev: Option<&CudaSlice>, n: usize, m: usize, blowup_factor: usize, weights: &[u64], split_col: usize, build_precomputed: bool, + retain_host_lde: bool, ) -> Result<(Option>, GpuLdeBase, Vec)> { assert!(split_col > 0 && split_col < m, "split inside the row"); assert!(n.is_power_of_two(), "n must be a power of two"); @@ -744,16 +754,12 @@ pub fn coset_lde_row_major_split_trees( let be = backend()?; let stream = be.next_stream(); - let (buf, trace_col_major) = expand_row_major_on_stream( - &stream, - be, - InnerInput::Host(row_major), - n, - m, - blowup_factor, - weights, - true, - )?; + let input = match predev { + Some(d) if d.len() == row_major.len() => InnerInput::Dev(d), + _ => InnerInput::Host(row_major), + }; + let (buf, trace_col_major) = + expand_row_major_on_stream(&stream, be, input, n, m, blowup_factor, weights, true)?; // One subset tree per column range, built sequentially on the stream. let build_subset_tree_dev = |col_start: u64, col_end: u64| -> Result> { @@ -801,10 +807,13 @@ pub fn coset_lde_row_major_split_trees( } }; - // D2H the row-major LDE (preprocessed tables always keep the host copy — - // they are excluded from the device-only gate). - let lde_pending = - crate::device::async_dtoh_via(&stream, be.pinned_staging(), &be.ctx, &buf, lde_size * m)?; + // D2H the row-major LDE only when the caller keeps a host copy; under + // device-only every downstream consumer reads the handle. + let lde_pending = retain_host_lde + .then(|| { + crate::device::async_dtoh_via(&stream, be.pinned_staging(), &be.ctx, &buf, lde_size * m) + }) + .transpose()?; // Column-major handle for downstream GPU rounds (DEEP, barycentric, // constraint composition). @@ -812,10 +821,13 @@ pub fn coset_lde_row_major_split_trees( let ready = be.take_event()?; ready.event().record(&stream)?; - let lde_out = { - let mut out = vec![0u64; lde_size * m]; - lde_pending.wait_into_u64(&mut out)?; - out + let lde_out = match lde_pending { + Some(pending) => { + let mut out = vec![0u64; lde_size * m]; + pending.wait_into_u64(&mut out)?; + out + } + None => Vec::new(), }; let handle = GpuLdeBase { diff --git a/crypto/math-cuda/tests/barycentric_multi.rs b/crypto/math-cuda/tests/barycentric_multi.rs new file mode 100644 index 000000000..4f58f27df --- /dev/null +++ b/crypto/math-cuda/tests/barycentric_multi.rs @@ -0,0 +1,167 @@ +//! Parity: the multi-eval-point chunked barycentric kernels match K separate +//! single-point strided calls over the same device LDE handle. + +use std::sync::Arc; + +use math::field::element::FieldElement; +use math::field::goldilocks::GoldilocksField; +use math_cuda::barycentric::{ + barycentric_base_multi_on_device, barycentric_base_on_device, barycentric_ext3_multi_on_device, + barycentric_ext3_on_device, +}; +use math_cuda::device::backend; +use math_cuda::lde::{GpuLdeBase, GpuLdeExt3}; +use rand::{Rng, SeedableRng}; +use rand_chacha::ChaCha8Rng; + +type Fp = FieldElement; + +fn rand_fp(rng: &mut ChaCha8Rng) -> Fp { + Fp::from_raw(rng.r#gen::()) +} + +fn run_base(log_trace: u32, blowup: usize, num_cols: usize, k_points: usize, seed: u64) { + let n = 1usize << log_trace; + let lde_size = n * blowup; + let mut rng = ChaCha8Rng::seed_from_u64(seed); + let mut lde_flat = vec![0u64; num_cols * lde_size]; + for v in lde_flat.iter_mut() { + *v = *rand_fp(&mut rng).value(); + } + let coset_points: Vec = (0..n).map(|_| rng.r#gen::()).collect(); + // K contiguous inv_denom blocks of 3n, the R3DevContext layout. + let inv_denoms_all: Vec = (0..(k_points * n * 3)) + .map(|_| rng.r#gen::()) + .collect(); + + let be = backend().unwrap(); + let stream = be.next_stream(); + let lde_dev = stream.clone_htod(&lde_flat).unwrap(); + let points_dev = stream.clone_htod(&coset_points).unwrap(); + let inv_dev = stream.clone_htod(&inv_denoms_all).unwrap(); + stream.synchronize().unwrap(); + let handle = GpuLdeBase { + ready: None, + buf: Arc::new(lde_dev), + m: num_cols, + lde_size, + tree: None, + trace_dev: None, + trace_rows: 0, + }; + + let multi = barycentric_base_multi_on_device( + &stream, + &handle, + blowup, + &points_dev, + &inv_dev, + n, + k_points, + ) + .unwrap(); + assert_eq!(multi.len(), 3 * k_points * num_cols); + + for k in 0..k_points { + let single = barycentric_base_on_device( + &handle, + blowup, + &coset_points, + &inv_denoms_all[k * 3 * n..(k + 1) * 3 * n], + n, + ) + .unwrap(); + assert_eq!( + &multi[k * 3 * num_cols..(k + 1) * 3 * num_cols], + &single[..], + "base multi mismatch at k={k} (log_trace={log_trace}, blowup={blowup}, \ + cols={num_cols}, k_points={k_points})" + ); + } +} + +fn run_ext3(log_trace: u32, blowup: usize, num_cols: usize, k_points: usize, seed: u64) { + let n = 1usize << log_trace; + let lde_size = n * blowup; + let mut rng = ChaCha8Rng::seed_from_u64(seed); + let mut lde_flat = vec![0u64; num_cols * 3 * lde_size]; + for v in lde_flat.iter_mut() { + *v = *rand_fp(&mut rng).value(); + } + let coset_points: Vec = (0..n).map(|_| rng.r#gen::()).collect(); + let inv_denoms_all: Vec = (0..(k_points * n * 3)) + .map(|_| rng.r#gen::()) + .collect(); + + let be = backend().unwrap(); + let stream = be.next_stream(); + let lde_dev = stream.clone_htod(&lde_flat).unwrap(); + let points_dev = stream.clone_htod(&coset_points).unwrap(); + let inv_dev = stream.clone_htod(&inv_denoms_all).unwrap(); + stream.synchronize().unwrap(); + let handle = GpuLdeExt3 { + ready: None, + buf: Arc::new(lde_dev), + m: num_cols, + lde_size, + tree: None, + }; + + let multi = barycentric_ext3_multi_on_device( + &stream, + &handle, + blowup, + &points_dev, + &inv_dev, + n, + k_points, + ) + .unwrap(); + assert_eq!(multi.len(), 3 * k_points * num_cols); + + for k in 0..k_points { + let single = barycentric_ext3_on_device( + &handle, + blowup, + &coset_points, + &inv_denoms_all[k * 3 * n..(k + 1) * 3 * n], + n, + ) + .unwrap(); + assert_eq!( + &multi[k * 3 * num_cols..(k + 1) * 3 * num_cols], + &single[..], + "ext3 multi mismatch at k={k} (log_trace={log_trace}, blowup={blowup}, \ + cols={num_cols}, k_points={k_points})" + ); + } +} + +#[test] +fn bary_base_multi_matches_single_point() { + // Covers: k=1 degenerate, the production k=2, the kernel cap k=8, a + // single-chunk tiny n, and a column count that forces the chunk heuristic + // to its occupancy branch. + for (log_t, blowup, cols, k) in [ + (4u32, 2usize, 3usize, 1usize), + (8, 4, 10, 2), + (12, 2, 5, 3), + (14, 2, 100, 2), + (10, 2, 4, 8), + ] { + run_base(log_t, blowup, cols, k, 3000 + log_t as u64 + k as u64); + } +} + +#[test] +fn bary_ext3_multi_matches_single_point() { + for (log_t, blowup, cols, k) in [ + (4u32, 2usize, 2usize, 1usize), + (8, 4, 5, 2), + (10, 2, 3, 3), + (14, 2, 40, 2), + (10, 2, 4, 8), + ] { + run_ext3(log_t, blowup, cols, k, 4000 + log_t as u64 + k as u64); + } +} diff --git a/crypto/math-cuda/tests/merkle_root_parity.rs b/crypto/math-cuda/tests/merkle_root_parity.rs index 208353d95..410828268 100644 --- a/crypto/math-cuda/tests/merkle_root_parity.rs +++ b/crypto/math-cuda/tests/merkle_root_parity.rs @@ -301,6 +301,7 @@ fn new_row_major_pipeline_base_root_matches_cpu() { let (handle, _lde) = math_cuda::lde::coset_lde_row_major_with_merkle_tree_keep( &row_major, + None, n, num_cols, blowup, diff --git a/crypto/stark/src/gpu_lde.rs b/crypto/stark/src/gpu_lde.rs index 98830fcc7..3f435bad6 100644 --- a/crypto/stark/src/gpu_lde.rs +++ b/crypto/stark/src/gpu_lde.rs @@ -39,10 +39,12 @@ use crate::trace::LDETraceTable; /// check is on **lde size**, not trace length, because that's what /// determines the FFT workload. /// -/// 2^19 is a conservative default calibrated against a 46-core machine where -/// rayon-parallel CPU LDE is already fast. Override via env var for tuning -/// on smaller machines, see `crypto/math-cuda/tests/bench_quick.rs`. -const DEFAULT_GPU_LDE_THRESHOLD: usize = 1 << 19; +/// The commit itself is not the whole cost: a table committed on CPU has no +/// device handle, so every R2-R4 GPU dispatch re-uploads its LDE. 2^14 is the +/// measured sweep optimum on ethrex continuations (2^14 beats 2^15..2^19 and +/// also beats "everything on GPU", where sub-2^14 tables lose to launch +/// overhead). Override via env var for tuning. +const DEFAULT_GPU_LDE_THRESHOLD: usize = 1 << 14; fn gpu_lde_threshold() -> usize { static CACHED: OnceLock = OnceLock::new(); @@ -54,6 +56,59 @@ fn gpu_lde_threshold() -> usize { }) } +/// Minimum LDE size for the device-only envelope, decoupled from the commit +/// threshold above. Committing on GPU and keeping the handle resident pays +/// from small sizes (it kills the per-round re-uploads); dropping the HOST +/// copy is a much stronger contract — every downstream dispatch must take its +/// GPU path or the prove hard-aborts, and the gate cannot mirror kernel-side +/// eligibility (the LOCKSTEP note below). Keep device-only to the large-table +/// envelope where those paths are exercised; mid tables keep a host copy so a +/// dispatch decline degrades to CPU instead of aborting. +const DEFAULT_DEVICE_ONLY_MIN_LDE: usize = 1 << 19; + +fn gpu_device_only_threshold() -> usize { + static CACHED: OnceLock = OnceLock::new(); + *CACHED.get_or_init(|| { + std::env::var("LAMBDA_VM_GPU_DEVICE_ONLY_THRESHOLD") + .ok() + .and_then(|s| s.parse().ok()) + .unwrap_or(DEFAULT_DEVICE_ONLY_MIN_LDE) + }) +} + +/// Test hook: decline the device R2 path unconditionally so device-only +/// tables exercise the [`materialize_lde_trace_host`] recovery end to end. +pub(crate) fn gpu_force_downgrade() -> bool { + static CACHED: OnceLock = OnceLock::new(); + *CACHED.get_or_init(|| std::env::var("LAMBDA_VM_GPU_FORCE_DOWNGRADE").is_ok_and(|v| v != "0")) +} + +/// Diagnostic hook: recompute the R2 composition parts and the R3 OOD +/// evaluations on host after each device dispatch and panic (naming the table +/// and stage) on any mismatch. Localizes silent device-side corruption. +pub(crate) fn gpu_xcheck() -> bool { + static CACHED: OnceLock = OnceLock::new(); + *CACHED.get_or_init(|| std::env::var("LAMBDA_VM_GPU_XCHECK").is_ok_and(|v| v != "0")) +} + +/// Serialize the device R2 window (constraint eval + decompose) across +/// tables. Concurrent R2 windows under VRAM pressure can transiently corrupt +/// a whole H buffer (root mechanism unidentified; reruns on the same resident +/// inputs come out correct), yielding a proof that fails verification. +/// Serializing only this window eliminates it at negligible cost — the +/// windows rarely overlap. `LAMBDA_VM_GPU_SERIALIZE_R2=0` disables the lock +/// (e.g. to bisect or once the underlying race is fixed). +pub(crate) fn r2_serialize_guard() -> Option> { + static ENABLED: OnceLock = OnceLock::new(); + static LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(()); + if *ENABLED.get_or_init(|| !std::env::var("LAMBDA_VM_GPU_SERIALIZE_R2").is_ok_and(|v| v == "0")) + { + Some(LOCK.lock().unwrap()) + } else { + None + } +} + /// Incremented by the `try_expand_*` functions per base-field column handed to /// the GPU dispatch (an ext3 column counts as 3, one per base component), /// before the GPU call. A failed call returns without decrementing it, so it @@ -189,7 +244,6 @@ pub(crate) fn device_only_disabled() -> bool { pub(crate) fn device_only_gate( lde_size: usize, n: usize, - is_preprocessed: bool, offsets_contiguous: bool, zerofier_uniform: bool, ) -> bool @@ -205,9 +259,8 @@ where && !device_only_disabled() && !gpu_composition_disabled() && lde_size.is_power_of_two() - && lde_size >= gpu_lde_threshold() + && lde_size >= gpu_device_only_threshold() && n >= gpu_bary_threshold() - && !is_preprocessed && offsets_contiguous && zerofier_uniform } @@ -679,6 +732,7 @@ pub fn gpu_leaf_hash_calls() -> u64 { /// openings gather paths from the device tree via [`gather_proofs_dev`]. pub(crate) fn try_expand_leaf_and_tree_row_major_keep( row_major: &[FieldElement], + predev: Option<&math_cuda::CudaSlice>, n: usize, m: usize, blowup_factor: usize, @@ -719,6 +773,7 @@ where // `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, + predev, n, m, blowup_factor, @@ -774,16 +829,21 @@ where /// downstream GPU rounds. /// /// `build_precomputed=false` skips the precomputed tree (process-cache hit); -/// the first element is then `None`. +/// the first element is then `None`. With `want_host=false` the row-major LDE +/// D2H is skipped and the returned Vec is empty (device-only tables: every +/// consumer reads the handle). #[allow(clippy::type_complexity)] +#[allow(clippy::too_many_arguments)] pub(crate) fn try_expand_split_trees_row_major_keep( row_major: &[FieldElement], + predev: Option<&math_cuda::CudaSlice>, n: usize, m: usize, blowup_factor: usize, weights: &[FieldElement], split_col: usize, build_precomputed: bool, + want_host: bool, ) -> Option<( Option>, MerkleTree, @@ -821,12 +881,14 @@ where let (pre_nodes, handle, lde_u64) = math_cuda::lde::coset_lde_row_major_split_trees( raw, + predev, n, m, blowup_factor, &weights_u64, split_col, build_precomputed, + want_host, ) .ok()?; @@ -1275,6 +1337,146 @@ where Some(apply_ext3_scalar::(&sums_raw, scalar, num_cols)) } +/// Multi-eval-point variant of [`try_barycentric_base_on_handle`]: one kernel +/// pass over the main LDE computes the OOD sums for every evaluation point at +/// once (their inv_denom blocks are contiguous in the [`R3DevContext`] buffer), +/// instead of re-reading the column data per point. Returns one scaled eval Vec +/// per point, or `None` (→ per-point dispatch / CPU fallback) when the handle +/// is absent, thresholds miss, there are more points than the kernel's +/// accumulator cap, or the math-cuda call errs. +#[allow(clippy::too_many_arguments)] +pub(crate) fn try_barycentric_base_on_handle_multi( + lde_trace: &LDETraceTable, + row_stride: usize, + coset_points_len: usize, + coset_offset_pow_n: &FieldElement, + n_inv: &FieldElement, + g_n_inv: &FieldElement, + z_pows: &[FieldElement], + ctx: &R3DevContext, +) -> Option>>> +where + F: IsField + IsSubFieldOf + 'static, + E: IsField + 'static, +{ + if !is_goldilocks_ext3_tower::() { + return None; + } + let k_points = z_pows.len(); + if k_points == 0 || k_points > math_cuda::barycentric::BARY_MAX_EVAL_POINTS { + return None; + } + let main = lde_trace.gpu_main()?; + let num_cols = main.m; + if num_cols == 0 { + return Some(vec![Vec::new(); k_points]); + } + let n = coset_points_len; + if !n.is_power_of_two() || n < gpu_bary_threshold() { + return None; + } + if main.lde_size != n.checked_mul(row_stride)? { + return None; + } + if ctx.inv_denoms.len() < k_points * 3 * n { + return None; + } + + let sums_raw = math_cuda::barycentric::barycentric_base_multi_on_device( + &ctx.stream, + main, + row_stride, + &ctx.coset_points, + &ctx.inv_denoms, + n, + k_points, + ) + .ok()?; + GPU_BARY_CALLS.fetch_add(k_points as u64, Ordering::Relaxed); + + Some( + z_pows + .iter() + .enumerate() + .map(|(k, z_pow_n)| { + let scalar = ood_ext3_scalar::(coset_offset_pow_n, n_inv, g_n_inv, z_pow_n); + apply_ext3_scalar::( + &sums_raw[k * 3 * num_cols..(k + 1) * 3 * num_cols], + scalar, + num_cols, + ) + }) + .collect(), + ) +} + +/// Aux (ext3) counterpart of [`try_barycentric_base_on_handle_multi`]. +#[allow(clippy::too_many_arguments)] +pub(crate) fn try_barycentric_ext3_on_handle_multi( + lde_trace: &LDETraceTable, + row_stride: usize, + coset_points_len: usize, + coset_offset_pow_n: &FieldElement, + n_inv: &FieldElement, + g_n_inv: &FieldElement, + z_pows: &[FieldElement], + ctx: &R3DevContext, +) -> Option>>> +where + F: IsField + IsSubFieldOf + 'static, + E: IsField + 'static, +{ + if !is_goldilocks_ext3_tower::() { + return None; + } + let k_points = z_pows.len(); + if k_points == 0 || k_points > math_cuda::barycentric::BARY_MAX_EVAL_POINTS { + return None; + } + let aux = lde_trace.gpu_aux()?; + let num_cols = aux.m; + if num_cols == 0 { + return Some(vec![Vec::new(); k_points]); + } + let n = coset_points_len; + if !n.is_power_of_two() || n < gpu_bary_threshold() { + return None; + } + if aux.lde_size != n.checked_mul(row_stride)? { + return None; + } + if ctx.inv_denoms.len() < k_points * 3 * n { + return None; + } + + let sums_raw = math_cuda::barycentric::barycentric_ext3_multi_on_device( + &ctx.stream, + aux, + row_stride, + &ctx.coset_points, + &ctx.inv_denoms, + n, + k_points, + ) + .ok()?; + GPU_BARY_CALLS.fetch_add(k_points as u64, Ordering::Relaxed); + + Some( + z_pows + .iter() + .enumerate() + .map(|(k, z_pow_n)| { + let scalar = ood_ext3_scalar::(coset_offset_pow_n, n_inv, g_n_inv, z_pow_n); + apply_ext3_scalar::( + &sums_raw[k * 3 * num_cols..(k + 1) * 3 * num_cols], + scalar, + num_cols, + ) + }) + .collect(), + ) +} + /// Ext3 counterpart of [`try_barycentric_base_on_handle`] for the aux LDE. /// Reads `lde_trace.gpu_aux()` (the de-interleaved 3-slab device buffer). #[allow(clippy::too_many_arguments)] @@ -1413,6 +1615,234 @@ pub fn gpu_fri_calls() -> u64 { /// are counted here, so a single failed dispatch does not necessarily lower /// the total; R3's fallbacks are CPU-only, so a failure there does. pub(crate) static GPU_BATCH_INVERT_CALLS: AtomicU64 = AtomicU64::new(0); +/// Times a device-only table had to be downgraded back to a host trace +/// because a downstream device path missed at runtime (see +/// [`materialize_lde_trace_host`]). Nonzero values mean the device-only gate +/// admitted a table some dispatch later declined — correct but slower, and +/// worth mirroring the missing condition into the gate. +pub(crate) static GPU_DEVICE_ONLY_DOWNGRADES: AtomicU64 = AtomicU64::new(0); +pub fn gpu_device_only_downgrades() -> u64 { + GPU_DEVICE_ONLY_DOWNGRADES.load(Ordering::Relaxed) +} + +/// Recover a device-only table for the host path: download the resident main +/// and aux LDEs from their device handles into the host buffers and clear the +/// device-only flag. A side whose host buffer is already populated (a mixed +/// state: one commit fell back to CPU while the other stayed device-only) is +/// kept as is — only the missing side is downloaded. The class-level safety +/// net under the device-only gate — a static predicate can never mirror every +/// reason a dynamic dispatch might decline (kernel eligibility, transient +/// errors, shapes a new workload brings), so any miss lands here and degrades +/// to a slower-but-correct CPU round instead of a hard abort. Returns false +/// (→ the caller's abort) only when a missing side has no handle or a +/// download fails. +pub(crate) fn materialize_lde_trace_host( + lde_trace: &mut crate::trace::LDETraceTable, +) -> bool +where + F: IsField + IsSubFieldOf + 'static, + E: IsField + 'static, +{ + if !lde_trace.host_trace_empty() { + return true; + } + if !is_goldilocks_ext3_tower::() { + return false; + } + let Some(stream) = lde_trace.bound_stream() else { + return false; + }; + + // Main: column-major device buf -> row-major host Vec. An empty Vec tells + // `set_host_data` to keep the buffer that is already there. + let main_data: Vec> = + if lde_trace.num_main_cols() == 0 || !lde_trace.main_data.is_empty() { + Vec::new() + } else { + let Some(h) = lde_trace.gpu_main() else { + return false; + }; + if h.m != lde_trace.num_main_cols() || h.lde_size != lde_trace.num_rows() { + return false; + } + let Some(data) = download_main_lde_row_major::(h, &stream) else { + return false; + }; + data + }; + + // Aux: de-interleaved ext3 slabs -> row-major interleaved host Vec. + let aux_data: Vec> = + if lde_trace.num_aux_cols() == 0 || !lde_trace.aux_data.is_empty() { + Vec::new() + } else { + let Some(h) = lde_trace.gpu_aux() else { + return false; + }; + if h.m != lde_trace.num_aux_cols() || h.lde_size != lde_trace.num_rows() { + return false; + } + if h.wait_ready_on(&stream).is_err() { + return false; + } + let Ok(slabs) = stream.clone_dtoh(h.buf.as_ref()) else { + return false; + }; + if stream.synchronize().is_err() { + return false; + } + let (m, lde) = (h.m, h.lde_size); + let mut interleaved = vec![0u64; m * lde * 3]; + for c in 0..m { + for k in 0..3 { + let slab = &slabs[(c * 3 + k) * lde..(c * 3 + k + 1) * lde]; + for r in 0..lde { + interleaved[(r * m + c) * 3 + k] = slab[r]; + } + } + } + // SAFETY: E == Ext3 per the tower check; FieldElement backing + // is [u64; 3]. + unsafe { + let mut v = std::mem::ManuallyDrop::new(interleaved); + Vec::from_raw_parts( + v.as_mut_ptr() as *mut FieldElement, + v.len() / 3, + v.capacity() / 3, + ) + } + }; + + lde_trace.set_host_data(main_data, aux_data); + GPU_DEVICE_ONLY_DOWNGRADES.fetch_add(1, Ordering::Relaxed); + true +} + +/// Download a resident main LDE (column-major device buf) into the row-major +/// host Vec the CPU rounds read. Shared by the R1 and R2 downgrade paths. +pub(crate) fn download_main_lde_row_major( + h: &math_cuda::lde::GpuLdeBase, + stream: &std::sync::Arc, +) -> Option>> +where + F: IsField + 'static, +{ + if TypeId::of::() != TypeId::of::() { + return None; + } + h.wait_ready_on(stream).ok()?; + let col_major = stream.clone_dtoh(h.buf.as_ref()).ok()?; + stream.synchronize().ok()?; + let (m, lde) = (h.m, h.lde_size); + if col_major.len() != m * lde { + return None; + } + let mut row_major = vec![0u64; m * lde]; + for c in 0..m { + for r in 0..lde { + row_major[r * m + c] = col_major[c * lde + r]; + } + } + // SAFETY: F == Goldilocks (gated above); FieldElement is + // #[repr(transparent)] over u64. + Some(unsafe { + let mut v = std::mem::ManuallyDrop::new(row_major); + Vec::from_raw_parts( + v.as_mut_ptr() as *mut FieldElement, + v.len(), + v.capacity(), + ) + }) +} + +/// R1 counterpart of [`materialize_lde_trace_host`]: download the resident +/// aux trace (already row-major ext3, matching the host layout) into the +/// trace's aux table, so the aux commit continues on the host arms when the +/// device aux LDE declines at runtime. +pub(crate) fn materialize_aux_trace_host(trace: &mut crate::trace::TraceTable) -> bool +where + F: IsField + IsSubFieldOf + 'static, + E: IsField + 'static, +{ + if !is_goldilocks_ext3_tower::() { + return false; + } + let (buf, rows, cols) = match trace.aux_resident.as_ref() { + Some(ra) => (ra.buf.clone(), ra.num_rows, ra.num_aux_cols), + None => return false, + }; + let Ok(be) = math_cuda::device::backend() else { + return false; + }; + let stream = be.next_stream(); + let Ok(raw) = stream.clone_dtoh(buf.as_ref()) else { + return false; + }; + if stream.synchronize().is_err() || raw.len() != rows * cols * 3 { + return false; + } + let data = u64_to_ext3_vec::(&raw); + trace.aux_table = crate::table::Table::new(data, cols); + trace.num_aux_columns = cols; + // The declined device LDE attempt can leave kernels enqueued on another + // stream still reading this buffer; its owning stream is long idle, so + // dropping here would complete the stream-ordered free immediately and + // the pool could hand the memory to a concurrent table's allocation + // while those kernels run. Drain the device before the drop — this is a + // rare recovery path. + if be.ctx.synchronize().is_err() { + return false; + } + trace.aux_resident = None; + GPU_DEVICE_ONLY_DOWNGRADES.fetch_add(1, Ordering::Relaxed); + true +} + +/// Diagnostic: download a resident ext3 handle (3-slab layout) as per-column +/// host Vecs. Used by the xcheck post-mortem to compare the committed R2 +/// parts against a host recompute. +pub(crate) fn download_ext3_columns( + h: &math_cuda::lde::GpuLdeExt3, +) -> Option>>> +where + E: IsField + 'static, +{ + if TypeId::of::() != TypeId::of::() { + return None; + } + let be = math_cuda::device::backend().ok()?; + let stream = be.next_stream(); + h.wait_ready_on(&stream).ok()?; + let slabs = stream.clone_dtoh(h.buf.as_ref()).ok()?; + stream.synchronize().ok()?; + let (m, lde) = (h.m, h.lde_size); + if slabs.len() != m * lde * 3 { + return None; + } + let mut cols = Vec::with_capacity(m); + for c in 0..m { + let mut interleaved = vec![0u64; lde * 3]; + for k in 0..3 { + let slab = &slabs[(c * 3 + k) * lde..(c * 3 + k + 1) * lde]; + for r in 0..lde { + interleaved[r * 3 + k] = slab[r]; + } + } + cols.push(u64_to_ext3_vec::(&interleaved)); + } + Some(cols) +} + +/// The device's VRAM admission budget in bytes, if a CUDA backend is up. +/// Lets callers outside this crate (the epoch builder's trace pre-upload) +/// size their riding-ahead allocations relative to the same budget the +/// per-table scheduler admits against. +pub fn device_vram_budget_bytes() -> Option { + math_cuda::device::backend() + .ok() + .map(|be| be.vram_budget_bytes()) +} + pub fn gpu_batch_invert_calls() -> u64 { GPU_BATCH_INVERT_CALLS.load(Ordering::Relaxed) } @@ -1979,10 +2409,7 @@ where // SAFETY: F == Goldilocks per TypeId check; FieldElement is // #[repr(transparent)] over u64. let coset_u64: &[u64] = unsafe { from_raw_parts(coset_base.as_ptr() as *const u64, n) }; - let coset_dev = match stream.clone_htod(coset_u64) { - Ok(s) => s, - Err(_) => return None, - }; + let coset_dev = coset_points_device_handle(coset_u64, stream)?; // SAFETY: E == Ext3 per TypeId check. let z_u64: &[u64] = unsafe { ext3_slice_to_u64::(z_scalars) }; @@ -1999,6 +2426,66 @@ where } } +/// Device-resident coset point buffers, keyed by `(len, points[0], points[1])` +/// — a geometric coset is fully determined by its length and first two terms, +/// so the key needs no allocation pinning. R3 OOD and the R4 DEEP inv_denoms +/// build used to re-upload the SAME domain points per table per epoch (~19 GB +/// per 100tx prove measured); one upload per distinct coset now serves the +/// whole process (a handful of sizes, ~2-16 MiB each, never evicted — same +/// policy as the host-side domain caches). +#[allow(clippy::type_complexity)] +fn coset_points_device_cache() +-> &'static std::sync::Mutex>>> { + static CACHE: OnceLock< + std::sync::Mutex>>>, + > = OnceLock::new(); + CACHE.get_or_init(Default::default) +} + +/// Resolve a host coset-points slice to its device-resident copy, uploading +/// once per distinct coset. The first upload synchronizes its stream so the +/// buffer is safe to read from any other stream afterwards. Returns `None` on +/// upload failure (→ the caller's fallback). +fn coset_points_device_handle( + coset_u64: &[u64], + stream: &Arc, +) -> Option>> { + if coset_u64.len() < 2 { + return stream.clone_htod(coset_u64).ok().map(Arc::new); + } + let key = (coset_u64.len(), coset_u64[0], coset_u64[1]); + if let Some(h) = coset_points_device_cache().lock().unwrap().get(&key) { + return Some(h.clone()); + } + // The key only determines the full contents for a geometric sequence + // `p_i = p_0·w^i`: verify it at sampled indices so a non-coset caller + // trips here instead of silently aliasing another entry. Insert-only — + // a handful of times per process. + { + type Fp = FieldElement; + let p0 = Fp::from_raw(coset_u64[0]); + let w = Fp::from_raw(coset_u64[1]) + * p0.inv() + .expect("coset_points_device_handle: coset offset must be nonzero"); + for i in [2usize, coset_u64.len() / 2, coset_u64.len() - 1] { + assert_eq!( + Fp::from_raw(coset_u64[i]), + &p0 * &w.pow(i as u64), + "coset_points_device_handle: input is not a geometric coset" + ); + } + } + let buf = stream.clone_htod(coset_u64).ok()?; + // Settle the copy before publishing: consumers run on other streams. + stream.synchronize().ok()?; + let h = Arc::new(buf); + coset_points_device_cache() + .lock() + .unwrap() + .insert(key, h.clone()); + Some(h) +} + /// Convenience wrapper for prover callers that don't yet own a stream: /// acquires the math-cuda backend, allocates a fresh stream, and produces /// a device-resident `inv_denoms` buffer plus the stream that owns it. @@ -2084,7 +2571,7 @@ pub(crate) fn gather_proofs_dev( #[derive(Debug)] pub(crate) struct R3DevContext { pub inv_denoms: CudaSlice, - pub coset_points: CudaSlice, + pub coset_points: Arc>, pub stream: Arc, } @@ -2130,7 +2617,7 @@ where // SAFETY: F == Goldilocks per TypeId check; FieldElement is // #[repr(transparent)] over u64. let coset_u64: &[u64] = unsafe { from_raw_parts(coset_base.as_ptr() as *const u64, n) }; - let coset_points = stream.clone_htod(coset_u64).ok()?; + let coset_points = coset_points_device_handle(coset_u64, &stream)?; // SAFETY: E == Ext3 per TypeId check. let z_u64: &[u64] = unsafe { ext3_slice_to_u64::(z_scalars) }; @@ -2576,7 +3063,7 @@ mod split_tree_tests { let (pre_tree, mult_tree, handle, lde) = try_expand_split_trees_row_major_keep::>( - &data, n, m, blowup, &weights, split, true, + &data, None, n, m, blowup, &weights, split, true, true, ) .expect("GPU split path must engage above the threshold"); let pre_tree = pre_tree.expect("precomputed tree was requested"); diff --git a/crypto/stark/src/prover.rs b/crypto/stark/src/prover.rs index 4047458bc..ed32978fc 100644 --- a/crypto/stark/src/prover.rs +++ b/crypto/stark/src/prover.rs @@ -1031,15 +1031,20 @@ pub trait IsStarkProver< 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; + // The device-resident R2 path only exists for the d=2 quotient + // decomposition; any other part count skips it entirely and needs the + // host evaluator, which device-only leaves without data. + if air.composition_poly_degree_bound(n) / n != 2 { + return false; + } + let lde_size = domain.interpolation_domain_size * domain.blowup_factor; let offsets_contiguous = crate::gpu_lde::offsets_are_contiguous(&air.context().transition_offsets); 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, ) @@ -1088,6 +1093,7 @@ pub trait IsStarkProver< BatchedMerkleTreeBackend, >( trace_slice, + trace.main_rowmajor_dev(), n, num_cols, domain.blowup_factor, @@ -1148,16 +1154,22 @@ pub trait IsStarkProver< BatchedMerkleTreeBackend, >( trace_slice, + trace.main_rowmajor_dev(), n, num_cols, domain.blowup_factor, &twiddles.coset_weights, num_precomputed, cached_pre.is_none(), + !device_only, ) { #[cfg(feature = "instruments")] crate::instruments::accum_r1_main(t_sub.elapsed(), std::time::Duration::ZERO); + if device_only { + crate::gpu_lde::GPU_DEVICE_ONLY_CALLS + .fetch_add(1, std::sync::atomic::Ordering::Relaxed); + } let precomputed_tree = match cached_pre { // Cache key == the root a rebuild would be verified // against, so a hit needs no re-check. @@ -1589,32 +1601,40 @@ pub trait IsStarkProver< #[cfg(feature = "cuda")] let mut precomputed_parts: Option>>> = None; #[cfg(feature = "cuda")] - if number_of_parts == 2 - && let Some(h_dev) = evaluator.evaluate_dev( - air, - &round_1_result.lde_trace, - domain, - transition_coefficients, - boundary_coefficients, - &round_1_result.rap_challenges, - ) { - match crate::gpu_lde::try_decompose_extend_d2_dev::( - &h_dev, - twiddles.inv_2x(domain), - &twiddles.composition(domain).weights, - !round_1_result.lde_trace.host_trace_empty(), - ) { - Some((parts, handle)) => { - gpu_composition_parts = Some(handle); - precomputed_parts = Some(parts); - } - None => { - if let Some(h) = - crate::gpu_lde::download_comp_h_to_field::(&h_dev) - { - precomputed_parts = - Some(Self::decompose_and_extend_d2(&h, domain, twiddles)); + // Serializing this window across tables (device constraint eval + + // decompose, where H is born) eliminates a transient whole-buffer + // H corruption seen under concurrent R2 windows on VRAM pressure. + // The commit and every host arm run outside the lock. + let _r2_serial_guard = crate::gpu_lde::r2_serialize_guard(); + if number_of_parts == 2 + && !crate::gpu_lde::gpu_force_downgrade() + && let Some(h_dev) = evaluator.evaluate_dev( + air, + &round_1_result.lde_trace, + domain, + transition_coefficients, + boundary_coefficients, + &round_1_result.rap_challenges, + ) + { + match crate::gpu_lde::try_decompose_extend_d2_dev::( + &h_dev, + twiddles.inv_2x(domain), + &twiddles.composition(domain).weights, + !round_1_result.lde_trace.host_trace_empty(), + ) { + Some((parts, handle)) => { + gpu_composition_parts = Some(handle); + precomputed_parts = Some(parts); + } + None => { + if let Some(h) = + crate::gpu_lde::download_comp_h_to_field::(&h_dev) + { + precomputed_parts = + Some(Self::decompose_and_extend_d2(&h, domain, twiddles)); + } } } } @@ -1633,11 +1653,36 @@ pub trait IsStarkProver< // failed. Abort with the device-only contract's message rather than a // bare index-out-of-bounds from somewhere inside the evaluator. #[cfg(feature = "cuda")] - if precomputed_parts.is_none() { + if precomputed_parts.is_none() && round_1_result.lde_trace.host_trace_empty() { + // The device R2 path missed on a device-only table. The gate is a + // static predicate and cannot mirror every dynamic decline, so + // recover instead of aborting: download the resident LDEs from + // the device handles and continue on the host path — slower for + // this table, never wrong. The abort remains only for the case + // where the handles themselves cannot serve the data. + let recovered = + crate::gpu_lde::materialize_lde_trace_host(&mut round_1_result.lde_trace); + if recovered { + // Rare by design; the name tells which condition the gate is + // missing so it can be mirrored as an optimization. + eprintln!( + "[gpu] device-only downgrade: table={} n={} num_parts={} \ + (device R2 path declined; continuing on host)", + air.name(), + trace_length, + number_of_parts, + ); + } assert!( - !round_1_result.lde_trace.host_trace_empty(), - "R2 composition fell back to the host evaluator, but the trace \ - is device-only (empty)" + recovered, + "R2 composition fell back to the host evaluator on a device-only \ + trace and the resident handles could not be downloaded: \ + table={} n={} num_parts={} main_cols={} aux_cols={}", + air.name(), + trace_length, + number_of_parts, + round_1_result.lde_trace.num_main_cols(), + round_1_result.lde_trace.num_aux_cols(), ); } @@ -1723,6 +1768,7 @@ pub trait IsStarkProver< #[cfg(not(feature = "cuda"))] cpu_eval()? }; + #[cfg(feature = "instruments")] let fft_dur = t_sub.elapsed(); @@ -2591,9 +2637,12 @@ pub trait IsStarkProver< /// One query's trace-poly opening with the device-resident fast paths: /// device Merkle proof + device-gathered values when both are present, the /// device proof with a host gather when only the tree is resident, and the - /// full host walk otherwise. One body for the main and aux arms, so the - /// device↔host cross-check and the R4 `host_trace_empty` hard-abort guards - /// exist exactly once. + /// full host walk otherwise. One body for the main, aux and preprocessed + /// multiplicity arms, so the device↔host cross-check and the R4 + /// `host_trace_empty` hard-abort guards exist exactly once. The device + /// gather always pulls the full `ncols` row; `col_range` selects the + /// committed subset (the full row for plain arms, `[split, ncols)` for the + /// multiplicity subset) and must match what `gather` returns. #[cfg(feature = "cuda")] #[allow(clippy::too_many_arguments)] fn open_trace_polys_device( @@ -2605,6 +2654,7 @@ pub trait IsStarkProver< qi: usize, challenge: usize, ncols: usize, + col_range: std::ops::Range, what: &str, gather: G, ) -> PolynomialOpenings @@ -2618,6 +2668,13 @@ pub trait IsStarkProver< !lde_trace.host_trace_empty(), "R4 {what} opening fell back to the host tree, but it is device-only (empty)" ); + // A root-only host tree means the nodes are device-resident: the + // host walk would emit an empty path for position 0 instead of + // failing, so a broken proofs↔tree pairing must abort here. + assert!( + !tree.is_root_only(), + "R4 {what} opening fell back to a root-only host tree (nodes device-resident)" + ); return Self::open_polys_with(domain, tree, challenge, gather); }; let proof = proofs[qi].clone(); @@ -2631,6 +2688,7 @@ pub trait IsStarkProver< return Self::open_polys_with_proofs(domain, proof, challenge, gather); }; let (even, odd) = Self::device_row_pair(dev_vals, qi, ncols); + let (even, odd) = (even[col_range.clone()].to_vec(), odd[col_range].to_vec()); // 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 @@ -2700,8 +2758,8 @@ pub trait IsStarkProver< // is a hard abort. When the tree is not device resident the value is // `None` and the openings below walk the full host tree. // For preprocessed tables the resident tree is the multiplicity subset - // tree (the host `main_commit.tree` is root only); values still come - // from the host LDE range gather below. + // tree (the host `main_commit.tree` is root only); values come from the + // same device row gather as plain tables, sliced per subset below. #[cfg(feature = "cuda")] let main_dev_proofs: Option>> = lde_trace .gpu_main() @@ -2755,10 +2813,8 @@ pub trait IsStarkProver< // *_dev_values.is_some()` on the Goldilocks path) and we never gather // rows for a tree that is not device resident. #[cfg(feature = "cuda")] - let main_dev_values: Option>> = (!is_preprocessed) - .then_some(()) - .and(main_dev_proofs.as_ref()) - .and_then(|_| { + let main_dev_values: Option>> = + main_dev_proofs.as_ref().and_then(|_| { lde_trace.gpu_main().and_then(|h| { Self::gather_query_rows_device( lde_trace, @@ -2834,41 +2890,25 @@ pub trait IsStarkProver< // For preprocessed tables, open the main split (multiplicities only); // for normal tables, open all main columns. let main_trace_opening = if is_preprocessed { - // Multiplicity subset: device proof (resident subset tree) + - // host range gather for the values. + // Multiplicity subset: same device fast paths as the plain + // arm, sliced to the committed `[split, total)` column range. #[cfg(feature = "cuda")] { - match &main_dev_proofs { - Some(proofs) => Self::open_polys_with_proofs( - domain, - proofs[qi].clone(), - *index, - |row| { - lde_trace.gather_main_row_range( - row, - num_precomputed_cols, - total_cols, - ) - }, - ), - None => { - // A root-only host tree means the nodes are - // device-resident: this arm would emit an empty - // path for query position 0 instead of failing. - assert!( - !main_commit.tree.is_root_only(), - "preprocessed opening fell back to the host tree, \ - but it is root-only (nodes device-resident)" - ); - Self::open_polys_with(domain, &main_commit.tree, *index, |row| { - lde_trace.gather_main_row_range( - row, - num_precomputed_cols, - total_cols, - ) - }) - } - } + Self::open_trace_polys_device( + domain, + lde_trace, + main_dev_proofs.as_ref(), + main_dev_values.as_ref(), + &main_commit.tree, + qi, + *index, + total_cols, + num_precomputed_cols..total_cols, + "multiplicity", + |row| { + lde_trace.gather_main_row_range(row, num_precomputed_cols, total_cols) + }, + ) } #[cfg(not(feature = "cuda"))] Self::open_polys_with(domain, &main_commit.tree, *index, |row| { @@ -2886,6 +2926,7 @@ pub trait IsStarkProver< qi, *index, total_cols, + 0..total_cols, "main", |row| lde_trace.gather_main_row(row), ) @@ -2899,7 +2940,61 @@ pub trait IsStarkProver< }; // For preprocessed tables, also open the precomputed-columns tree. + // The tree is always a full host tree (process-wide cache), so the + // Merkle path comes from the host walk; the VALUES come from the + // device row gather when the LDE is resident (sliced to the + // `[0, split)` range), host range gather otherwise. let precomputed_trace_opening = main_commit.precomputed_tree.as_ref().map(|tree| { + #[cfg(feature = "cuda")] + { + match main_dev_values.as_ref() { + Some(vals) => { + let (even, odd) = Self::device_row_pair(vals, qi, total_cols); + let (even, odd) = ( + even[..num_precomputed_cols].to_vec(), + odd[..num_precomputed_cols].to_vec(), + ); + // Query 0 stays a release canary, same rationale + // as `open_trace_polys_device`. + if (cfg!(debug_assertions) || qi == 0) && !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_range( + r_even, + 0, + num_precomputed_cols + ), + "device precomputed-row gather mismatch (even), query {qi}" + ); + assert_eq!( + odd, + lde_trace.gather_main_row_range(r_odd, 0, num_precomputed_cols), + "device precomputed-row gather mismatch (odd), query {qi}" + ); + } + Self::open_polys_from_values( + tree.get_proof_by_pos(*index) + .expect("FRI query index in bounds"), + even, + odd, + ) + } + None => { + assert!( + !lde_trace.host_trace_empty(), + "R4 precomputed opening fell back to the host gather, \ + but it is device-only (empty)" + ); + Self::open_polys_with(domain, tree, *index, |row| { + lde_trace.gather_main_row_range(row, 0, num_precomputed_cols) + }) + } + } + } + #[cfg(not(feature = "cuda"))] Self::open_polys_with(domain, tree, *index, |row| { lde_trace.gather_main_row_range(row, 0, num_precomputed_cols) }) @@ -2986,6 +3081,7 @@ pub trait IsStarkProver< qi, *index, lde_trace.num_aux_cols(), + 0..lde_trace.num_aux_cols(), "aux", |row| lde_trace.gather_aux_row(row), ) @@ -3358,6 +3454,7 @@ pub trait IsStarkProver< #[cfg(feature = "cuda")] { trace.clear_main_trace_dev(); + trace.clear_main_rowmajor_dev(); if let Some(handle) = gpu_main_cells[idx].lock().unwrap().as_mut() { handle.trace_dev = None; handle.trace_rows = 0; @@ -3379,23 +3476,29 @@ pub trait IsStarkProver< if air.has_aux_trace() { let lde_size = domain.interpolation_domain_size * domain.blowup_factor; - // Same gate as the Round 1 main commit: skip the aux - // host D2H when device-only, so both buffers are left - // empty together for this table. + // Same gate as the Round 1 main commit — but only if + // that commit actually produced a device handle. If + // the GPU main commit declined and fell back to CPU, + // skipping the aux D2H here would mark the trace + // device-only with no main handle to serve it, turning + // a recoverable fallback into a hard abort downstream. #[cfg(feature = "cuda")] - let device_only = Self::device_only_for(*air, domain); + let mut device_only = Self::device_only_for(*air, domain) + && gpu_main_cells[idx].lock().unwrap().is_some(); // 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 // resident build fired the host aux trace is empty, so a - // device LDE failure is a hard abort, not a fall through to - // the host path below (which would commit a zero aux trace). + // device LDE failure downloads the resident aux trace and + // continues on the host arms below (falling through as-is + // would commit a zero aux trace). #[cfg(feature = "cuda")] - if let Some(ra) = trace.aux_resident() { + if trace.aux_resident().is_some() { #[cfg(feature = "instruments")] let t_sub = Instant::now(); - let (tree, handle, aux_data) = + let num_cols = trace.aux_resident().map_or(0, |ra| ra.num_aux_cols); + let expand = |ra: &math_cuda::logup::ResidentAux| { crate::gpu_lde::try_expand_leaf_and_tree_ext3_row_major_keep_dev::< Field, FieldExtension, @@ -3406,21 +3509,79 @@ pub trait IsStarkProver< &twiddles.coset_weights, !device_only, ) - .ok_or_else(|| { - ProvingError::Fft( - "resident aux LDE failed; host aux trace is empty" - .to_string(), - ) - })?; - let num_cols = ra.num_aux_cols; - #[cfg(feature = "instruments")] - crate::instruments::accum_r1_aux(t_sub.elapsed(), Duration::ZERO); - let root = tree.root; - return Ok(( - Some(TableCommit::plain(tree, root)), - (aux_data, num_cols), - Some(handle), - )); + }; + let mut expanded = expand(trace.aux_resident().expect("checked above")); + if expanded.is_none() + && let Ok(be) = math_cuda::device::backend() + && be.ctx.synchronize().is_ok() + { + // The decline is usually transient VRAM + // pressure from concurrent tables; a device + // drain releases those peaks, so one retry + // tends to keep the table fully resident + // instead of paying the host downgrade. + eprintln!( + "[gpu] resident aux LDE declined: table={} \ + (retrying after device drain)", + air.name(), + ); + expanded = expand(trace.aux_resident().expect("checked above")); + } + if let Some((tree, handle, aux_data)) = expanded { + #[cfg(feature = "instruments")] + crate::instruments::accum_r1_aux(t_sub.elapsed(), Duration::ZERO); + let root = tree.root; + return Ok(( + Some(TableCommit::plain(tree, root)), + (aux_data, num_cols), + Some(handle), + )); + } + // The device aux LDE declined at runtime (transient + // VRAM pressure, usually) and there is no host aux + // trace to fall back to. Same class as the R2 + // downgrade: download the resident aux trace — and + // the main LDE if this table was device-only — and + // continue fully host-backed on the arms below. + let mut recovered = crate::gpu_lde::materialize_aux_trace_host(*trace); + if recovered && device_only { + let mut cell = main_lde_cells[idx].lock().unwrap(); + if let Some((data, _)) = cell.as_mut() + && data.is_empty() + && trace.num_main_columns > 0 + { + recovered = match ( + gpu_main_cells[idx].lock().unwrap().as_ref(), + math_cuda::device::backend(), + ) { + (Some(h), Ok(be)) => { + match crate::gpu_lde::download_main_lde_row_major::( + h, + &be.next_stream(), + ) { + Some(v) => { + *data = v; + true + } + None => false, + } + } + _ => false, + }; + } + } + if !recovered { + return Err(ProvingError::Fft( + "resident aux LDE failed; host aux trace is empty".to_string(), + )); + } + eprintln!( + "[gpu] resident-aux downgrade: table={} rows={} \ + (device aux LDE declined; continuing on host)", + air.name(), + trace.num_rows(), + ); + device_only = false; } // Fused GPU path (cuda only): row-major ext3 NTT — single @@ -3718,6 +3879,370 @@ pub trait IsStarkProver< // TODO: propagate errors instead of unwrap() in open_deep_composition_poly and FRI operations /// Executes rounds 2-4 and generates a STARK proof for the trace `main_trace` with public inputs `pub_inputs`. /// Warning: the transcript must be safely initialized before passing it to this method. + /// Diagnostic (see `gpu_lde::gpu_xcheck`): the verifier's step-2 + /// composition consistency check run in-process on the freshly computed + /// R3 values — H(z) reconstructed from the trace OOD evaluations must + /// match the folded parts OOD. Near-zero cost (one constraint evaluation + /// at a single point), so it can run on every table without disturbing + /// the timing that provokes VRAM-pressure bugs. Mirrors + /// `step_2_verify_claimed_composition_polynomial` in `verifier.rs`. + #[cfg(feature = "cuda")] + #[allow(clippy::too_many_arguments)] + fn composition_ood_consistent( + air: &dyn AIR, + pub_inputs: &PI, + domain: &Domain, + rap_challenges: &[FieldElement], + bus_public_inputs: Option<&BusPublicInputs>, + transition_coefficients: &[FieldElement], + boundary_coefficients: &[FieldElement], + z: &FieldElement, + trace_ood: &Table, + parts_ood: &[FieldElement], + ) -> bool { + use crate::lookup::{LOGUP_CHALLENGE_ALPHA, compute_alpha_powers}; + use crate::traits::TransitionEvaluationContext; + + let trace_length = domain.interpolation_domain_size; + let boundary_constraints = + air.boundary_constraints(pub_inputs, rap_challenges, bus_public_inputs, trace_length); + let mut step_to_point: std::collections::HashMap> = + std::collections::HashMap::new(); + let boundary_points: Vec> = boundary_constraints + .constraints + .iter() + .map(|c| { + step_to_point + .entry(c.step) + .or_insert_with(|| domain.trace_primitive_root.pow(c.step as u64)) + .clone() + }) + .collect(); + + let main_trace_width = air.trace_layout().0; + let ood_row = trace_ood.get_row(0); + let (nums, mut dens): ( + Vec>, + Vec>, + ) = boundary_constraints + .constraints + .iter() + .zip(&boundary_points) + .map(|(c, point)| { + let column_idx = if c.is_aux { + main_trace_width + c.col + } else { + c.col + }; + (-&c.value + &ood_row[column_idx], -point + z) + }) + .unzip(); + if FieldElement::inplace_batch_inverse(&mut dens).is_err() { + return false; + } + let boundary_sum: FieldElement = nums + .iter() + .zip(&dens) + .zip(boundary_coefficients) + .map(|((num, den), beta)| num * den * beta) + .fold(FieldElement::zero(), |acc, x| acc + x); + + let Some(num_main_trace_columns) = + trace_ood.width.checked_sub(air.num_auxiliary_rap_columns()) + else { + return false; + }; + let logup_alpha_powers: Vec> = + if rap_challenges.len() > LOGUP_CHALLENGE_ALPHA { + compute_alpha_powers( + &rap_challenges[LOGUP_CHALLENGE_ALPHA], + air.max_bus_elements(), + ) + } else { + Vec::new() + }; + let logup_table_offset = match bus_public_inputs { + Some(bpi) => { + let n = FieldElement::::from(trace_length as u64); + match n.inv() { + Ok(n_inv) => n_inv * &bpi.table_contribution, + Err(_) => return false, + } + } + None => FieldElement::zero(), + }; + + // Frame over the OOD grid, mirroring `StarkTableView::into_frame` + // (that view carries rkyv bounds this generic context lacks). + let step_size = air.step_size(); + debug_assert!(trace_ood.height.is_multiple_of(step_size)); + let steps: Vec> = (0..trace_ood + .height) + .step_by(step_size) + .map(|initial| { + let mut main = Vec::new(); + let mut aux = Vec::new(); + for row_idx in initial..initial + step_size { + let row = trace_ood.get_row(row_idx); + main.push(row[..num_main_trace_columns].to_vec()); + aux.push(row[num_main_trace_columns..].to_vec()); + } + crate::table::TableView::new(main, aux) + }) + .collect(); + let ood_frame = crate::frame::Frame::new(steps); + let ctx = TransitionEvaluationContext::new_verifier( + &ood_frame, + rap_challenges, + &logup_alpha_powers, + &logup_table_offset, + ); + let transition_evals = air.compute_transition(&ctx); + + let mut denominators = + vec![FieldElement::::zero(); air.num_transition_constraints()]; + air.constraints_meta().iter().for_each(|m| { + denominators[m.constraint_idx] = crate::constraints::zerofier::evaluate_zerofier( + m, + z, + &domain.trace_primitive_root, + trace_length, + ); + }); + let transition_sum = transition_evals + .into_iter() + .zip(transition_coefficients) + .zip(denominators) + .fold(FieldElement::zero(), |acc, ((eval, beta), den)| { + acc + beta * eval * &den + }); + + let ood_evaluation = &boundary_sum + transition_sum; + let claimed = parts_ood + .iter() + .rev() + .fold(FieldElement::zero(), |acc, coeff| acc * z + coeff); + claimed == ood_evaluation + } + + /// Diagnostic follow-up when [`Self::composition_ood_consistent`] fails: + /// recompute each device-derived stage on host for THIS table only and + /// report which one diverges, then panic (the proof would not verify). + /// Runs after the corruption already happened, so the expensive host + /// recomputes cannot mask the failure they are diagnosing. + #[cfg(feature = "cuda")] + #[allow(clippy::too_many_arguments)] + fn xcheck_post_mortem( + air: &dyn AIR, + pub_inputs: &PI, + domain: &Domain, + twiddles: &LdeTwiddles, + round_1_result: &mut Round1, + transition_coefficients: &[FieldElement], + boundary_coefficients: &[FieldElement], + round_2_result: &Round2, + round_3_result: &Round3, + z: &FieldElement, + ) where + FieldElement: AsBytes, + FieldElement: AsBytes, + { + let name = air.name(); + let trace_length = domain.interpolation_domain_size; + eprintln!("[xcheck] FAIL composition consistency: table={name} n={trace_length}"); + + if round_1_result.lde_trace.host_trace_empty() + && !crate::gpu_lde::materialize_lde_trace_host(&mut round_1_result.lde_trace) + { + panic!("[xcheck] table={name}: cannot materialize host trace for post-mortem"); + } + + // Stage 1: R2 parts (device H + decompose) vs full host recompute. + let evaluator = ConstraintEvaluator::new( + air, + pub_inputs, + &round_1_result.rap_challenges, + round_1_result.bus_public_inputs.as_ref(), + trace_length, + ); + let host_h = evaluator.evaluate( + air, + &round_1_result.lde_trace, + domain, + transition_coefficients, + boundary_coefficients, + &round_1_result.rap_challenges, + ); + let host_parts = Self::decompose_and_extend_d2(&host_h, domain, twiddles); + let device_parts: Option>>> = if round_2_result + .lde_composition_poly_evaluations + .first() + .is_some_and(|p| !p.is_empty()) + { + Some(round_2_result.lde_composition_poly_evaluations.clone()) + } else { + round_1_result + .lde_trace + .gpu_composition_parts() + .and_then(crate::gpu_lde::download_ext3_columns::) + }; + let mut r2_verdict = "UNAVAILABLE (no device parts to compare)".to_string(); + if let Some(dev) = &device_parts { + r2_verdict = "ok".to_string(); + 'outer: for (pi, (hp, dp)) in host_parts.iter().zip(dev).enumerate() { + if hp.len() != dp.len() { + r2_verdict = + format!("LEN MISMATCH part={pi} host={} dev={}", hp.len(), dp.len()); + break; + } + for (ri, (x, y)) in hp.iter().zip(dp.iter()).enumerate() { + if x != y { + r2_verdict = format!("MISMATCH part={pi} row={ri} host={x:?} device={y:?}"); + break 'outer; + } + } + } + } + eprintln!("[xcheck] table={name} R2 parts: {r2_verdict}"); + + // Corruption shape: how much of each part differs, and where. A whole + // buffer points at H itself; a contiguous chunk at one kernel pass; a + // strided pattern at slab/component confusion. + if let Some(dev) = &device_parts { + for (pi, (hp, dp)) in host_parts.iter().zip(dev).enumerate() { + if hp.len() != dp.len() { + continue; + } + let mism: Vec = hp + .iter() + .zip(dp.iter()) + .enumerate() + .filter(|(_, (x, y))| x != y) + .map(|(i, _)| i) + .collect(); + if !mism.is_empty() { + eprintln!( + "[xcheck] table={name} part={pi}: {} of {} rows differ, first={} last={}", + mism.len(), + hp.len(), + mism[0], + mism[mism.len() - 1], + ); + } + } + } + + // Rerun the device R2 chain for this table now that the storm has + // passed: a correct rerun means a transient race during the original + // run; the same wrong values mean a persistently corrupted device + // input (zerofiers, IR buffers, resident LDEs). + let rerun: Option>>> = evaluator + .evaluate_dev( + air, + &round_1_result.lde_trace, + domain, + transition_coefficients, + boundary_coefficients, + &round_1_result.rap_challenges, + ) + .and_then(|h_dev| { + crate::gpu_lde::try_decompose_extend_d2_dev::( + &h_dev, + twiddles.inv_2x(domain), + &twiddles.composition(domain).weights, + true, + ) + .map(|(parts, _handle)| parts) + }); + let rerun_verdict = match &rerun { + None => "device rerun declined".to_string(), + Some(p2) if *p2 == host_parts => { + "rerun matches HOST (transient race in the original run)".to_string() + } + Some(p2) if device_parts.as_ref().is_some_and(|dp| p2 == dp) => { + "rerun matches ORIGINAL DEVICE (persistent corrupted device input)".to_string() + } + Some(_) => "rerun matches NEITHER".to_string(), + }; + eprintln!("[xcheck] table={name} R2 rerun: {rerun_verdict}"); + + // Stage 2: R3 trace OOD vs the host arms. + let dc = domain.ood_constants(); + let host_ood = crate::trace::with_r3_force_host(|| { + crate::trace::get_trace_evaluations_from_lde( + &round_1_result.lde_trace, + domain, + z, + &air.context().transition_offsets, + air.step_size(), + dc, + ) + }); + let got = &round_3_result.trace_ood_evaluations; + let mut r3_trace_verdict = "ok".to_string(); + if host_ood.width != got.width || host_ood.height != got.height { + r3_trace_verdict = "SHAPE MISMATCH".to_string(); + } else { + 'outer: for r in 0..host_ood.height { + for c in 0..host_ood.width { + if host_ood.get(r, c) != got.get(r, c) { + r3_trace_verdict = format!( + "MISMATCH row={r} col={c} host={:?} device={:?}", + host_ood.get(r, c), + got.get(r, c) + ); + break 'outer; + } + } + } + } + eprintln!("[xcheck] table={name} R3 trace_ood: {r3_trace_verdict}"); + + // Stage 3: R3 parts OOD vs the host arm over the HOST-recomputed parts + // (independent of the device H), and over the device parts when + // available (isolates barycentric vs upstream). + let num_parts = round_3_result.composition_poly_parts_ood_evaluation.len(); + let z_power = z.pow(num_parts); + let comp_z_pow_n = z_power.pow(trace_length); + let comp_inv_denoms = math::polynomial::barycentric_inv_denoms(&z_power, &dc.points); + let ood_of = + |parts: &[Vec>]| -> Vec> { + parts + .iter() + .map(|lde_evals| { + let evals: Vec> = (0..trace_length) + .map(|i| lde_evals[i * domain.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 host_parts_ood = ood_of(&host_parts); + eprintln!( + "[xcheck] table={name} R3 parts_ood: claimed={:?} host_from_host_parts={:?} host_from_device_parts={:?}", + round_3_result.composition_poly_parts_ood_evaluation, + host_parts_ood, + device_parts.as_deref().map(ood_of), + ); + + eprintln!( + "[xcheck] table={name}: composition OOD inconsistency (R2 parts: {r2_verdict}; \ + R2 rerun: {rerun_verdict}; R3 trace_ood: {r3_trace_verdict}); aborting" + ); + // abort() and not panic!: a panicking prover thread deadlocks the + // epoch pipeline (producer stuck in a bounded send), which would turn + // every diagnostic catch into a hung process. + std::process::abort(); + } + fn prove_rounds_2_to_4( air: &dyn AIR, pub_inputs: &PI, @@ -3796,6 +4321,38 @@ pub trait IsStarkProver< #[cfg(feature = "instruments")] let round_3_dur = t_r3.elapsed(); + // Diagnostic: verifier-equivalent composition consistency check, run + // per table at negligible cost; on failure, per-stage host recompute + // names where the corruption entered (then panics). + #[cfg(feature = "cuda")] + if crate::gpu_lde::gpu_xcheck() + && !Self::composition_ood_consistent( + air, + pub_inputs, + domain, + &round_1_result.rap_challenges, + round_1_result.bus_public_inputs.as_ref(), + &transition_coefficients, + &boundary_coefficients, + &z, + &round_3_result.trace_ood_evaluations, + &round_3_result.composition_poly_parts_ood_evaluation, + ) + { + Self::xcheck_post_mortem( + air, + pub_inputs, + domain, + twiddles, + round_1_result, + &transition_coefficients, + &boundary_coefficients, + &round_2_result, + &round_3_result, + &z, + ); + } + // >>>> Send values: tⱼ(zgᵏ). g·z pruning: split the full OOD table into // the current-row block (all columns) and the pruned next-row block // (masked columns only), and absorb only the surviving values — the diff --git a/crypto/stark/src/trace.rs b/crypto/stark/src/trace.rs index b34023ac3..bce6fffa9 100644 --- a/crypto/stark/src/trace.rs +++ b/crypto/stark/src/trace.rs @@ -45,6 +45,85 @@ where /// LDE did not run for this table. #[cfg(feature = "cuda")] pub(crate) main_trace_dev: Option, + /// Row-major main trace pre-uploaded to device off the prove critical path + /// (by the epoch pipeline's builder thread, which finishes ~1s before the + /// prover consumes the epoch). The R1 main commit D2D-copies from it + /// instead of paying the H2D inside its chain. + #[cfg(feature = "cuda")] + pub(crate) main_rowmajor_dev: Option, +} + +/// Device-resident row-major main trace, pre-uploaded ahead of the prove. +/// Excluded from logical trace equality and opaque in `Debug`, matching +/// [`ResidentMainTrace`]. +#[cfg(feature = "cuda")] +#[derive(Clone)] +pub(crate) struct PreUploadedMainTrace { + pub(crate) buf: std::sync::Arc>, +} + +#[cfg(feature = "cuda")] +impl core::fmt::Debug for PreUploadedMainTrace { + fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { + f.debug_struct("PreUploadedMainTrace") + .finish_non_exhaustive() + } +} + +#[cfg(feature = "cuda")] +impl PartialEq for PreUploadedMainTrace { + fn eq(&self, _other: &Self) -> bool { + true + } +} + +#[cfg(feature = "cuda")] +impl Eq for PreUploadedMainTrace {} + +// Separate impl: the `TypeId` tower check needs `'static`, which the main +// `TraceTable` impl does not require of its parameters. +#[cfg(feature = "cuda")] +impl TraceTable +where + E: IsField + 'static, + F: IsSubFieldOf + IsFFTField + 'static, +{ + /// Pre-upload the row-major main trace to device, off the prove critical + /// path (called from the epoch pipeline's builder thread). Returns the + /// bytes uploaded (0 = skipped: non-Goldilocks tower, empty, below the + /// size floor, or upload failure — the commit then does its own H2D). + /// The upload stream is synchronized before publishing, so any stream may + /// read the buffer afterwards. + pub fn preupload_main_to_device(&mut self, min_bytes: usize) -> usize { + use std::any::TypeId; + if self.main_rowmajor_dev.is_some() { + return 0; + } + if TypeId::of::() != TypeId::of::() { + return 0; + } + let (data, cols) = self.main_data_row_major(); + let bytes = std::mem::size_of_val(data); + if cols == 0 || data.is_empty() || bytes < min_bytes { + return 0; + } + let Ok(be) = math_cuda::device::backend() else { + return 0; + }; + let stream = be.next_stream(); + // SAFETY: F == Goldilocks per the TypeId check; FieldElement is + // #[repr(transparent)] over u64. + let raw: &[u64] = + unsafe { core::slice::from_raw_parts(data.as_ptr() as *const u64, data.len()) }; + let Ok(buf) = stream.clone_htod(raw) else { + return 0; + }; + if stream.synchronize().is_err() { + return 0; + } + self.main_rowmajor_dev = Some(PreUploadedMainTrace { buf: Arc::new(buf) }); + bytes + } } /// Device-resident trace-domain main columns (column-major `[col*rows + row]`), @@ -105,6 +184,8 @@ where resident_aux_ok: true, #[cfg(feature = "cuda")] main_trace_dev: None, + #[cfg(feature = "cuda")] + main_rowmajor_dev: None, } } @@ -133,6 +214,8 @@ where resident_aux_ok: true, #[cfg(feature = "cuda")] main_trace_dev: None, + #[cfg(feature = "cuda")] + main_rowmajor_dev: None, } } @@ -154,6 +237,8 @@ where resident_aux_ok: true, #[cfg(feature = "cuda")] main_trace_dev: None, + #[cfg(feature = "cuda")] + main_rowmajor_dev: None, } } @@ -213,6 +298,20 @@ where self.main_trace_dev = None; } + /// The pre-uploaded row-major main trace, if the builder produced one. + #[cfg(feature = "cuda")] + pub(crate) fn main_rowmajor_dev(&self) -> Option<&math_cuda::CudaSlice> { + self.main_rowmajor_dev.as_ref().map(|p| p.buf.as_ref()) + } + + /// Drop the pre-uploaded row-major trace. Its only consumer is the R1 main + /// commit, so the prover clears it alongside `clear_main_trace_dev` to + /// reclaim the VRAM before the aux-commit + DEEP/FRI peak. + #[cfg(feature = "cuda")] + pub fn clear_main_rowmajor_dev(&mut self) { + self.main_rowmajor_dev = None; + } + pub fn num_steps(&self) -> usize { debug_assert!(self.main_table.height.is_multiple_of(self.step_size)); self.main_table.height / self.step_size @@ -537,6 +636,27 @@ where /// `main_data.len()` — the caller supplies it from the device handle's /// `lde_size` instead. #[cfg(feature = "cuda")] + /// Install downloaded host buffers on a device-only table and clear the + /// flag: from here every host read is valid again. An empty Vec keeps + /// that side's existing buffer (either the side has no columns or it + /// already held a host copy in a mixed state). Only meaningful from + /// [`crate::gpu_lde::materialize_lde_trace_host`], which guarantees the + /// buffers match the device handles' layout. + #[cfg(feature = "cuda")] + pub(crate) fn set_host_data( + &mut self, + main_data: Vec>, + aux_data: Vec>, + ) { + if !main_data.is_empty() { + self.main_data = main_data; + } + if !aux_data.is_empty() { + self.aux_data = aux_data; + } + self.host_trace_empty = false; + } + pub fn set_num_rows(&mut self, num_rows: usize) { self.num_rows = num_rows; } @@ -657,6 +777,23 @@ where } } +// Diagnostic (see `gpu_lde::gpu_xcheck`): while set on the current thread, +// `get_trace_evaluations_from_lde` skips every GPU dispatch and runs the +// host arms, so a second call can cross-check the device results. +#[cfg(feature = "cuda")] +thread_local! { + static R3_FORCE_HOST: std::cell::Cell = const { std::cell::Cell::new(false) }; +} + +/// Run `f` with the R3 GPU dispatches disabled on this thread. +#[cfg(feature = "cuda")] +pub(crate) fn with_r3_force_host(f: impl FnOnce() -> R) -> R { + R3_FORCE_HOST.with(|c| c.set(true)); + let out = f(); + R3_FORCE_HOST.with(|c| c.set(false)); + out +} + /// Evaluates trace polynomials at OOD points using barycentric interpolation /// on the LDE evaluations, without needing coefficient-form polynomials. /// @@ -710,16 +847,56 @@ where // into a single device context. The barycentric kernels below read // both via offset, with no per-eval-point or per-{main,aux} H2D. #[cfg(feature = "cuda")] - let r3_ctx: Option = + let r3_force_host = R3_FORCE_HOST.with(|c| c.get()); + #[cfg(feature = "cuda")] + let r3_ctx: Option = if r3_force_host { + None + } else { crate::gpu_lde::try_prep_r3_dev_context::( &dc.points, &evaluation_points, lde_trace.bound_stream(), - ); + ) + }; #[allow(unused_variables)] #[cfg(not(feature = "cuda"))] let r3_ctx: Option<()> = None; + // Multi-eval-point GPU fast path: ONE kernel pass per {main, aux} computes + // the barycentric sums for every evaluation point (the per-point loop below + // then just consumes its slice). `None` (handle absent, too many points, + // kernel error) falls through to the per-point dispatch inside the loop, + // which preserves the original behavior arm by arm. + #[cfg(feature = "cuda")] + let (main_multi, aux_multi) = match r3_ctx.as_ref() { + Some(ctx) => { + let z_pows: Vec> = evaluation_points.iter().map(|p| p.pow(n)).collect(); + ( + crate::gpu_lde::try_barycentric_base_on_handle_multi::( + lde_trace, + bf, + n, + &dc.offset_pow_n, + &dc.size_inv, + &dc.offset_pow_n_inv, + &z_pows, + ctx, + ), + crate::gpu_lde::try_barycentric_ext3_on_handle_multi::( + lde_trace, + bf, + n, + &dc.offset_pow_n, + &dc.size_inv, + &dc.offset_pow_n_inv, + &z_pows, + ctx, + ), + ) + } + None => (None, None), + }; + #[cfg_attr(not(feature = "cuda"), allow(clippy::unused_enumerate_index))] for (eval_point_idx, eval_point) in evaluation_points.iter().enumerate() { // Silence unused warning under non-cuda where eval_point_idx is @@ -763,17 +940,26 @@ where #[cfg(feature = "cuda")] let r3_arg = r3_ctx.as_ref().map(|ctx| (ctx, eval_point_idx * 3 * n)); #[cfg(feature = "cuda")] - let main_gpu = crate::gpu_lde::try_barycentric_base_on_handle::( - lde_trace, - bf, - &dc.points, - &dc.offset_pow_n, - &dc.size_inv, - &dc.offset_pow_n_inv, - &z_pow_n, - inv_denoms.as_deref().unwrap_or(&[]), - r3_arg, - ); + let main_gpu = if r3_force_host { + None + } else { + main_multi + .as_ref() + .map(|per_point| per_point[eval_point_idx].clone()) + .or_else(|| { + crate::gpu_lde::try_barycentric_base_on_handle::( + lde_trace, + bf, + &dc.points, + &dc.offset_pow_n, + &dc.size_inv, + &dc.offset_pow_n_inv, + &z_pow_n, + inv_denoms.as_deref().unwrap_or(&[]), + r3_arg, + ) + }) + }; #[cfg(not(feature = "cuda"))] let main_gpu: Option>> = None; @@ -781,10 +967,12 @@ where 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. + // read empty `main_data`. Hard-abort instead of a wrong OOD eval. The + // check is on the buffer itself, not the table-wide flag: a mixed + // state can leave a valid host copy on one side only. #[cfg(feature = "cuda")] assert!( - !lde_trace.host_trace_empty(), + lde_trace.num_main_cols() == 0 || !lde_trace.main_data.is_empty(), "R3 barycentric (main) fell back to the host trace, but it is device-only (empty)" ); let inv_denoms_v = @@ -821,17 +1009,26 @@ where #[cfg(feature = "cuda")] let r3_arg_aux = r3_ctx.as_ref().map(|ctx| (ctx, eval_point_idx * 3 * n)); #[cfg(feature = "cuda")] - let aux_gpu = crate::gpu_lde::try_barycentric_ext3_on_handle::( - lde_trace, - bf, - &dc.points, - &dc.offset_pow_n, - &dc.size_inv, - &dc.offset_pow_n_inv, - &z_pow_n, - inv_denoms.as_deref().unwrap_or(&[]), - r3_arg_aux, - ); + let aux_gpu = if r3_force_host { + None + } else { + aux_multi + .as_ref() + .map(|per_point| per_point[eval_point_idx].clone()) + .or_else(|| { + crate::gpu_lde::try_barycentric_ext3_on_handle::( + lde_trace, + bf, + &dc.points, + &dc.offset_pow_n, + &dc.size_inv, + &dc.offset_pow_n_inv, + &z_pow_n, + inv_denoms.as_deref().unwrap_or(&[]), + r3_arg_aux, + ) + }) + }; #[cfg(not(feature = "cuda"))] let aux_gpu: Option>> = None; @@ -839,10 +1036,11 @@ where 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. + // read empty `aux_data`. Hard-abort instead of a wrong OOD eval. Same + // buffer-level check as the main arm: mixed states are valid here. #[cfg(feature = "cuda")] assert!( - !lde_trace.host_trace_empty(), + lde_trace.num_aux_cols() == 0 || !lde_trace.aux_data.is_empty(), "R3 barycentric (aux) fell back to the host trace, but it is device-only (empty)" ); let inv_denoms_v = diff --git a/prover/src/continuation.rs b/prover/src/continuation.rs index 8f3e68db4..664c101ee 100644 --- a/prover/src/continuation.rs +++ b/prover/src/continuation.rs @@ -1231,6 +1231,17 @@ pub fn prove_continuation( drop(__nvtx); match traces { Ok(traces) => { + // Pre-upload the big main traces from this builder thread + // (idle slack ahead of the prover), so the R1 main commits + // skip their H2D. + #[cfg(feature = "cuda")] + let traces = { + let mut traces = traces; + #[cfg(feature = "instruments")] + let __sp = stark::instruments::span("p6_trace_preupload"); + traces.preupload_main_traces(); + traces + }; let prepared = PreparedEpoch { index: job.index, register_init: job.register_init, @@ -1763,6 +1774,101 @@ mod tests { use super::*; use crate::test_utils::asm_elf_bytes; + // Diagnostic (not a regression test): structurally diff two continuation + // proof bundles of the same input. The prover is deterministic, so the + // first differing field per table names the round where a corrupt run + // diverged. Run with: + // PROOF_A= PROOF_B= \ + // cargo test -p prover --release proof_diff -- --ignored --nocapture + #[test] + #[ignore] + fn proof_diff() { + fn load(path: &str) -> ContinuationProof { + use std::os::unix::fs::FileExt; + let file = std::fs::File::open(path).unwrap(); + let len = file.metadata().unwrap().len() as usize; + let mut aligned = rkyv::util::AlignedVec::<16>::with_capacity(len); + aligned.resize(len, 0); + file.read_exact_at(&mut aligned, 0).unwrap(); + rkyv::from_bytes::(&aligned).unwrap() + } + fn table_eq(a: &stark::table::Table, b: &stark::table::Table) -> bool { + if a.width != b.width || a.height != b.height { + return false; + } + (0..a.height).all(|r| (0..a.width).all(|c| a.get(r, c) == b.get(r, c))) + } + fn diff_multi(label: &str, a: &MultiProof, b: &MultiProof) { + assert_eq!(a.proofs.len(), b.proofs.len(), "{label}: table count"); + for (t, (pa, pb)) in a.proofs.iter().zip(b.proofs.iter()).enumerate() { + let mut d = Vec::new(); + if pa.lde_trace_main_merkle_root != pb.lde_trace_main_merkle_root { + d.push("main_root"); + } + if pa.lde_trace_aux_merkle_root != pb.lde_trace_aux_merkle_root { + d.push("aux_root"); + } + if pa.lde_trace_precomputed_merkle_root != pb.lde_trace_precomputed_merkle_root { + d.push("preproc_root"); + } + if pa.bus_public_inputs.as_ref().map(|x| &x.table_contribution) + != pb.bus_public_inputs.as_ref().map(|x| &x.table_contribution) + { + d.push("bus_pi"); + } + if pa.composition_poly_root != pb.composition_poly_root { + d.push("comp_root"); + } + if !table_eq(&pa.trace_ood_evaluations, &pb.trace_ood_evaluations) { + d.push("trace_ood"); + } + if !table_eq( + &pa.trace_ood_next_evaluations, + &pb.trace_ood_next_evaluations, + ) { + d.push("trace_ood_next"); + } + if pa.composition_poly_parts_ood_evaluation + != pb.composition_poly_parts_ood_evaluation + { + d.push("parts_ood"); + } + if pa.fri_layers_merkle_roots != pb.fri_layers_merkle_roots { + d.push("fri_roots"); + } + if pa.fri_final_poly_coeffs != pb.fri_final_poly_coeffs { + d.push("fri_final"); + } + if pa.nonce != pb.nonce { + d.push("nonce"); + } + if !d.is_empty() { + println!( + "{label} table {t} (cols={} len={}): {d:?}", + pa.trace_ood_evaluations.width, pa.trace_length + ); + } + } + } + let a = load(&std::env::var("PROOF_A").unwrap()); + let b = load(&std::env::var("PROOF_B").unwrap()); + assert_eq!(a.epochs.len(), b.epochs.len(), "epoch count"); + for (e, (ea, eb)) in a.epochs.iter().zip(b.epochs.iter()).enumerate() { + diff_multi(&format!("epoch {e}"), &ea.proof, &eb.proof); + if ea.public_output != eb.public_output { + println!("epoch {e}: public_output differs"); + } + if ea.reg_fini != eb.reg_fini { + println!("epoch {e}: reg_fini differs"); + } + if ea.l2g_root != eb.l2g_root { + println!("epoch {e}: l2g_root differs"); + } + } + diff_multi("global", &a.global, &b.global); + println!("diff complete"); + } + // `test_commit_split` issues two Commit syscalls, one early and one late, so a // small epoch puts the second commit in a later epoch. That epoch starts with // x254 > 0 (the carried commit index), which exercises the cross-epoch commit diff --git a/prover/src/tables/bitwise.rs b/prover/src/tables/bitwise.rs index 45bddb636..c73e1e341 100644 --- a/prover/src/tables/bitwise.rs +++ b/prover/src/tables/bitwise.rs @@ -411,6 +411,10 @@ pub fn update_multiplicities( trace: &mut TraceTable, ops: &[BitwiseOperation], ) { + // A pre-uploaded device copy of the main trace would go stale with the + // in-place edits below; drop it so the commit re-uploads fresh data. + #[cfg(feature = "cuda")] + trace.clear_main_rowmajor_dev(); for op in ops { let row = row_index(op.x, op.y, op.z); let mu_col = mu_column(op.lookup_type); diff --git a/prover/src/tables/trace_builder.rs b/prover/src/tables/trace_builder.rs index 5ec9fa566..f067b45fd 100644 --- a/prover/src/tables/trace_builder.rs +++ b/prover/src/tables/trace_builder.rs @@ -3830,6 +3830,83 @@ pub fn count_table_lengths( } impl Traces { + /// Pre-upload the epoch's biggest main traces to device, called from the + /// epoch pipeline's builder thread (idle slack ahead of the prover) so the + /// R1 main commits D2D-copy instead of paying the H2D inside their chains. + /// Biggest tables first, bounded by `LAMBDA_VM_TRACE_PREUPLOAD_MB` (default + /// 4096) of VRAM riding ahead per epoch; tables that don't fit (or are + /// below the 8 MiB floor, or whose upload fails) keep the normal H2D path. + #[cfg(feature = "cuda")] + pub fn preupload_main_traces(&mut self) { + const MIN_BYTES: usize = 8 << 20; + static BUDGET_BYTES: std::sync::OnceLock = std::sync::OnceLock::new(); + let budget = *BUDGET_BYTES.get_or_init(|| { + // Default OFF: pre-uploading was wall-neutral on the 5090 (the + // scheduler already hides the H2D) and its riding-ahead buffers + // sit outside the VRAM admission gate — at epoch 2^22 they pushed + // the prove past the card's headroom. Opt in for PCIe-bound + // setups via the env var. + let env_cap = std::env::var("LAMBDA_VM_TRACE_PREUPLOAD_MB") + .ok() + .and_then(|s| s.parse::().ok()) + .unwrap_or(0) + << 20; + // These buffers ride ahead of the prover's own VRAM admission + // gate (they exist before their table is admitted), so cap them + // to a slice of the device budget rather than competing with the + // prove peak on small cards. + match stark::gpu_lde::device_vram_budget_bytes() { + Some(dev) => env_cap.min((dev / 4) as usize), + None => env_cap, + } + }); + if budget == 0 { + return; + } + + let mut tables: Vec<&mut TraceTable> = Vec::new(); + tables.extend(self.cpus.iter_mut()); + tables.extend(self.lts.iter_mut()); + tables.extend(self.shifts.iter_mut()); + tables.extend(self.memws.iter_mut()); + tables.extend(self.memw_aligneds.iter_mut()); + tables.extend(self.memw_registers.iter_mut()); + tables.extend(self.loads.iter_mut()); + tables.extend(self.muls.iter_mut()); + tables.extend(self.dvrms.iter_mut()); + tables.extend(self.pages.iter_mut()); + tables.extend(self.branches.iter_mut()); + tables.extend(self.eqs.iter_mut()); + tables.extend(self.bytewises.iter_mut()); + tables.extend(self.stores.iter_mut()); + tables.extend(self.cpu32s.iter_mut()); + // BITWISE is excluded: `prove_epoch` mutates its multiplicities in + // place (L2G range-check lookups) after the build, which would leave + // a stale device copy to be committed. + tables.push(&mut self.decode); + tables.push(&mut self.keccak); + tables.push(&mut self.keccak_rnd); + tables.push(&mut self.ecsm); + tables.push(&mut self.ecdas); + + let bytes_of = |t: &TraceTable| { + t.num_rows() * t.num_main_columns * 8 + }; + tables.sort_by_key(|t| std::cmp::Reverse(bytes_of(t))); + + let mut left = budget; + for t in tables { + let est = bytes_of(t); + if est < MIN_BYTES { + break; + } + if est > left { + continue; + } + left -= t.preupload_main_to_device(MIN_BYTES); + } + } + /// Returns the total number of main-trace field elements across all tables. /// /// Counts only the main (base-field) trace columns — equivalent to SP1's diff --git a/prover/tests/gpu_force_downgrade.rs b/prover/tests/gpu_force_downgrade.rs new file mode 100644 index 000000000..b1d8cc897 --- /dev/null +++ b/prover/tests/gpu_force_downgrade.rs @@ -0,0 +1,45 @@ +//! End-to-end exercise of the device-only downgrade recovery: with +//! `LAMBDA_VM_GPU_FORCE_DOWNGRADE` set, every device-only table declines its +//! device R2 path, downloads its resident LDEs back to host +//! (`materialize_lde_trace_host`) and finishes on the host evaluator — and +//! the proof must still verify. The device-only threshold is lowered so the +//! small fixture actually produces device-only tables. +//! +//! Lives in its own integration-test binary: the env hooks are cached in +//! process-wide `OnceLock`s, so they must be set before any other test's GPU +//! dispatch initializes them. +//! +//! Requires the `cuda` feature and a visible GPU. Run with: +//! +//! ```text +//! cargo test -p lambda-vm-prover --release --features cuda \ +//! --test gpu_force_downgrade -- --ignored --nocapture +//! ``` +#![cfg(feature = "cuda")] + +#[test] +#[ignore = "requires GPU; run with --ignored --nocapture"] +fn forced_downgrade_prove_verifies() { + // SAFETY: single test in this binary, set before any GPU dispatch. + unsafe { + std::env::set_var("LAMBDA_VM_GPU_FORCE_DOWNGRADE", "1"); + std::env::set_var("LAMBDA_VM_GPU_DEVICE_ONLY_THRESHOLD", "16384"); + } + let ws = std::path::PathBuf::from(env!("CARGO_MANIFEST_DIR")) + .parent() + .expect("workspace root") + .to_path_buf(); + let elf = std::fs::read(ws.join("executor/program_artifacts/rust/ethrex.elf")) + .expect("need ethrex.elf — run `make compile-programs-rust`"); + let input = std::fs::read(ws.join("executor/tests/ethrex_simple_tx.bin")).expect("fixture"); + + let proof = lambda_vm_prover::prove_with_inputs(&elf, &input).expect("prove"); + assert!( + stark::gpu_lde::gpu_device_only_downgrades() > 0, + "no table took the forced downgrade — the hook or the device-only gate moved" + ); + assert!( + lambda_vm_prover::verify(&proof, &elf).expect("verify"), + "downgraded proof must verify" + ); +} diff --git a/scripts/profiling/h2d_histo.py b/scripts/profiling/h2d_histo.py new file mode 100644 index 000000000..1317cc103 --- /dev/null +++ b/scripts/profiling/h2d_histo.py @@ -0,0 +1,79 @@ +#!/usr/bin/env python3 +"""H2D/D2H attribution histogram from an nsys sqlite export. + +Groups memcpys by (enclosing phase, innermost NVTX range, size) so the +dominant uploaders inside a phase are identifiable by name + size fingerprint. +Reuses the loaders from nsys_phase_busy.py (same directory). +""" + +import os +import sqlite3 +import sys +from collections import defaultdict + +sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) +from nsys_phase_busy import ( + base_name, + build_range_lookup, + load_api_calls, + load_gpu_rows, + load_nvtx, + load_strings, + tables, +) + + +def main(): + db = sys.argv[1] + con = sqlite3.connect(f"file:{db}?mode=ro", uri=True) + tset = tables(con) + strings = load_strings(con, tset) + nvtx = load_nvtx(con, tset, strings) + _, memcpys = load_gpu_rows(con, tset, strings) + api = load_api_calls(con, tset) + chain_at = build_range_lookup(nvtx) + + def chain_for(corr, start): + if corr in api: + api_start, tid = api[corr] + c = chain_at(tid, api_start) + if c: + return c + return [] + + def coarse_of(chain): + for name in reversed(chain): + if "[" not in name: + return name + return base_name(chain[0]) if chain else "(none)" + + def innermost(chain): + return base_name(chain[-1]) if chain else "(none)" + + # (direction, phase, inner, bytes) -> [count, total_bytes, total_ns] + hist = defaultdict(lambda: [0, 0, 0]) + for start, end, kind, nbytes, corr in memcpys: + if kind not in ("h2d", "d2h"): + continue + chain = chain_for(corr, start) + key = (kind, coarse_of(chain), innermost(chain), nbytes) + h = hist[key] + h[0] += 1 + h[1] += nbytes + h[2] += end - start + + for direction in ("h2d", "d2h"): + rows = [(k, v) for k, v in hist.items() if k[0] == direction] + rows.sort(key=lambda kv: -kv[1][1]) + total_gb = sum(v[1] for _, v in rows) / 2**30 + print(f"\n== {direction.upper()} total {total_gb:.1f} GiB — top 20 by bytes ==") + print(f"{'phase':<28} {'inner range':<28} {'size MiB':>9} {'count':>6} {'GiB':>7} {'ms':>8}") + for (_, phase, inner, nbytes), (cnt, tot, ns) in rows[:20]: + print( + f"{phase:<28} {inner:<28} {nbytes / 2**20:>9.2f} {cnt:>6} " + f"{tot / 2**30:>7.2f} {ns / 1e6:>8.1f}" + ) + + +if __name__ == "__main__": + main()