diff --git a/Cargo.lock b/Cargo.lock index d2f699b80..5bfab63d4 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1687,6 +1687,7 @@ dependencies = [ "executor", "log", "math", + "math-cuda", "postcard", "rayon", "serde", diff --git a/crypto/math-cuda/build.rs b/crypto/math-cuda/build.rs index 8fc568e2b..316a9c7ed 100644 --- a/crypto/math-cuda/build.rs +++ b/crypto/math-cuda/build.rs @@ -121,4 +121,5 @@ fn main() { compile_ptx("fri.cu", "fri.ptx", have_nvcc); compile_ptx("inverse.cu", "inverse.ptx", have_nvcc); compile_ptx("logup.cu", "logup.ptx", have_nvcc); + compile_ptx("constraint_interp.cu", "constraint_interp.ptx", have_nvcc); } diff --git a/crypto/math-cuda/kernels/constraint_interp.cu b/crypto/math-cuda/kernels/constraint_interp.cu new file mode 100644 index 000000000..07ac0f9eb --- /dev/null +++ b/crypto/math-cuda/kernels/constraint_interp.cu @@ -0,0 +1,311 @@ +// Transition-constraint interpreter kernel. +// +// Evaluates a captured `ConstraintProgram` (lowered to the flat device blob by +// `crypto/stark/src/constraint_ir/device.rs`) over every row of a +// device-resident LDE, producing the per-constraint evaluations. It is a +// transliteration of the CPU walker `eval_device_program` (same module), with +// `FieldElement` arithmetic replaced by `goldilocks.cuh` / `ext3.cuh` — the two +// are asserted bit-for-bit equal by the pre-GPU parity test, so this kernel's +// output equals the compiled prover folder. +// +// Design (v1, "stripped" — mirrors OpenVM's GLOBAL=true quotient kernel): +// * One thread per LDE row, grid-stride over all rows (fixed launch, any size). +// * Per-thread value array in GLOBAL memory, strided by thread for coalescing: +// node `i` for this thread lives at `d_values[task_offset + i*num_threads]`. +// * ALL values are carried as ext3 `Fe3` (a base value `x` is the embedding +// `{x,0,0}`). Because embedding is a ring homomorphism and the device field +// arithmetic is bit-identical to the CPU path, doing every op in ext3 yields +// outputs bit-identical to the dim-split CPU walk — the per-node `dim` tag +// is therefore unused here and reserved for the base/ext-split optimization +// (Phase 6). `Op::Embed` is consequently the identity. +// +// Output is the per-constraint eval matrix `d_evals[c*num_rows + row]` (Fe3; +// base-rooted constraints carry their value in `.a`). Fusing the +// `z*Σ(Cᵢ·βᵢ) + boundary` accumulation into this kernel (to avoid a D2H of the +// matrix — the actual data-residency win) is the pipeline-integration follow-on. +// +// Op tags and the `Var` packing MUST stay in sync with +// `crypto/stark/src/constraint_ir/device.rs`. + +#include "goldilocks.cuh" +#include "ext3.cuh" + +using ext3::Fe3; + +// -- op tags (mirror device.rs OP_*) -- +#define OP_CONST_BASE 0u +#define OP_CONST_EXT 1u +#define OP_VAR 2u +#define OP_RAP_CHALLENGE 3u +#define OP_ALPHA_POW 4u +#define OP_TABLE_OFFSET 5u +#define OP_ADD 6u +#define OP_SUB 7u +#define OP_MUL 8u +#define OP_NEG 9u +#define OP_EMBED 10u + +// A flat IR node. Packed into two u64 words for a pure-u64 upload (matching the +// crate's device-buffer convention): +// word0 = op | (a << 32) ; word1 = b | (dim << 32) +// This mirrors the `#[repr(C)] DeviceNode { op, a, b, dim: u32 }` payload. +struct Node { + uint32_t op, a, b, dim; +}; + +__device__ __forceinline__ Node load_node(const uint64_t *d_nodes, uint64_t i) { + uint64_t w0 = d_nodes[2 * i]; + uint64_t w1 = d_nodes[2 * i + 1]; + Node n; + n.op = (uint32_t)(w0 & 0xFFFFFFFFull); + n.a = (uint32_t)(w0 >> 32); + n.b = (uint32_t)(w1 & 0xFFFFFFFFull); + n.dim = (uint32_t)(w1 >> 32); + return n; +} + +// Resolve an `Op::Var` leaf against the device-resident LDE columns. +// a = col (low 16 bits); b = main<<16 | offset<<8 | row (see device.rs pack_var) +// Base (main) columns are column-major `d_main[col*main_stride + r]`; ext (aux) +// columns store component k at `d_aux[(col*3 + k)*aux_stride + r]` (GpuLdeExt3). +// The frame `offset` selects row `r = (row + offset*next_step) mod num_rows`. +__device__ __forceinline__ Fe3 read_var(uint32_t a, uint32_t b, uint64_t row, uint64_t next_step, + uint64_t num_rows, const uint64_t *d_main, + uint64_t main_stride, const uint64_t *d_aux, + uint64_t aux_stride) { + uint32_t col = a & 0xFFFFu; + bool is_main = ((b >> 16) & 1u) != 0u; + uint32_t offset = (b >> 8) & 0xFFu; + + uint64_t r = row + (uint64_t)offset * next_step; + if (r >= num_rows) { + r -= num_rows; // wrap; offset*next_step < num_rows by construction + } + + if (is_main) { + uint64_t x = d_main[(uint64_t)col * main_stride + r]; + return ext3::make(x, 0, 0); + } + uint64_t base = (uint64_t)col * 3; + uint64_t a0 = d_aux[(base + 0) * aux_stride + r]; + uint64_t a1 = d_aux[(base + 1) * aux_stride + r]; + uint64_t a2 = d_aux[(base + 2) * aux_stride + r]; + return ext3::make(a0, a1, a2); +} + +extern "C" __global__ void constraint_interp_kernel( + // output: per-constraint eval matrix, constraint-major [num_roots * num_rows] + Fe3 *__restrict__ d_evals, + // program (flat blob) + const uint64_t *__restrict__ d_nodes, // 2 u64 per node + uint64_t num_nodes, + const uint64_t *__restrict__ d_base_consts, + const Fe3 *__restrict__ d_ext_consts, + const uint64_t *__restrict__ d_roots, + uint64_t num_roots, + // per-proof uniforms + const Fe3 *__restrict__ d_rap_challenges, + const Fe3 *__restrict__ d_alpha_powers, + const Fe3 *__restrict__ d_table_offset, // single element + // device-resident LDE + const uint64_t *__restrict__ d_main, + uint64_t main_stride, + const uint64_t *__restrict__ d_aux, + uint64_t aux_stride, + uint64_t next_step, + // sizing + uint64_t num_rows, + // scratch: per-thread value array, [num_nodes * num_threads] + Fe3 *__restrict__ d_values) { + uint64_t task_offset = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x; + uint64_t num_threads = (uint64_t)gridDim.x * blockDim.x; + + Fe3 *vals = d_values + task_offset; + uint64_t vstride = num_threads; + Fe3 table_offset = *d_table_offset; + + for (uint64_t row = task_offset; row < num_rows; row += num_threads) { + // Forward pass: id `i` references only nodes `< i`. + for (uint64_t i = 0; i < num_nodes; i++) { + Node nd = load_node(d_nodes, i); + Fe3 v; + switch (nd.op) { + case OP_CONST_BASE: + v = ext3::make(d_base_consts[nd.a], 0, 0); + break; + case OP_CONST_EXT: + v = d_ext_consts[nd.a]; + break; + case OP_VAR: + v = read_var(nd.a, nd.b, row, next_step, num_rows, d_main, main_stride, d_aux, + aux_stride); + break; + case OP_RAP_CHALLENGE: + v = d_rap_challenges[nd.a]; + break; + case OP_ALPHA_POW: + v = d_alpha_powers[nd.a]; + break; + case OP_TABLE_OFFSET: + v = table_offset; + break; + case OP_ADD: + v = ext3::add(vals[(uint64_t)nd.a * vstride], vals[(uint64_t)nd.b * vstride]); + break; + case OP_SUB: + v = ext3::sub(vals[(uint64_t)nd.a * vstride], vals[(uint64_t)nd.b * vstride]); + break; + case OP_MUL: + v = ext3::mul(vals[(uint64_t)nd.a * vstride], vals[(uint64_t)nd.b * vstride]); + break; + case OP_NEG: + v = ext3::neg(vals[(uint64_t)nd.a * vstride]); + break; + case OP_EMBED: + // All-ext representation: the base operand is already {x,0,0}. + v = vals[(uint64_t)nd.a * vstride]; + break; + default: + v = ext3::zero(); + break; + } + vals[i * vstride] = v; + } + + // Emit each constraint root. + for (uint64_t c = 0; c < num_roots; c++) { + uint64_t root = d_roots[c]; + d_evals[c * num_rows + row] = vals[root * vstride]; + } + } +} + +// Fused composition-polynomial kernel: same node walk as +// `constraint_interp_kernel`, but instead of emitting the per-constraint matrix +// it accumulates the composition-poly evaluation H(row) on-device — no matrix +// materialization, no D2H. Mirrors the CPU accumulation in +// `crypto/stark/src/constraints/evaluator.rs` (uniform-zerofier case): +// +// H(row) = z_inv[row % z_len] * Σ_c beta_trans[c] * C_c(row) (transition) +// + Σ_b z_b_inv[b*num_rows + row] * beta_bnd[b] * (trace_b - value_b) +// +// where a base-rooted C_c is carried as the embedding {C,0,0} (all-ext, exactly +// as the interpreter), z_inv is the cyclic base transition-zerofier inverse, and +// the boundary term reads the resident trace at column `b_col[b]` (main or aux). +extern "C" __global__ void constraint_composition_kernel( + // output: one H(row) per LDE row + Fe3 *__restrict__ d_h, + // program (flat blob) — identical to the interpreter kernel + const uint64_t *__restrict__ d_nodes, + uint64_t num_nodes, + const uint64_t *__restrict__ d_base_consts, + const Fe3 *__restrict__ d_ext_consts, + const uint64_t *__restrict__ d_roots, + uint64_t num_roots, + // per-proof uniforms + const Fe3 *__restrict__ d_rap_challenges, + const Fe3 *__restrict__ d_alpha_powers, + const Fe3 *__restrict__ d_table_offset, + // device-resident LDE + const uint64_t *__restrict__ d_main, + uint64_t main_stride, + const uint64_t *__restrict__ d_aux, + uint64_t aux_stride, + uint64_t next_step, + uint64_t num_rows, + // transition accumulation + const Fe3 *__restrict__ d_beta_trans, // [num_roots] + const uint64_t *__restrict__ d_z_inv, // [z_len], cyclic + uint64_t z_len, + // boundary accumulation + uint64_t num_boundary, + const uint64_t *__restrict__ d_b_col, // [num_boundary] + const uint64_t *__restrict__ d_b_is_aux, // [num_boundary] (0/1) + const Fe3 *__restrict__ d_b_value, // [num_boundary] + const Fe3 *__restrict__ d_b_beta, // [num_boundary] + const uint64_t *__restrict__ d_b_z_inv, // [num_boundary * num_rows] + // scratch value array, [num_nodes * num_threads] + Fe3 *__restrict__ d_values) { + uint64_t task_offset = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x; + uint64_t num_threads = (uint64_t)gridDim.x * blockDim.x; + + Fe3 *vals = d_values + task_offset; + uint64_t vstride = num_threads; + Fe3 table_offset = *d_table_offset; + + for (uint64_t row = task_offset; row < num_rows; row += num_threads) { + // Forward pass (identical to the interpreter kernel). + for (uint64_t i = 0; i < num_nodes; i++) { + Node nd = load_node(d_nodes, i); + Fe3 v; + switch (nd.op) { + case OP_CONST_BASE: + v = ext3::make(d_base_consts[nd.a], 0, 0); + break; + case OP_CONST_EXT: + v = d_ext_consts[nd.a]; + break; + case OP_VAR: + v = read_var(nd.a, nd.b, row, next_step, num_rows, d_main, main_stride, d_aux, + aux_stride); + break; + case OP_RAP_CHALLENGE: + v = d_rap_challenges[nd.a]; + break; + case OP_ALPHA_POW: + v = d_alpha_powers[nd.a]; + break; + case OP_TABLE_OFFSET: + v = table_offset; + break; + case OP_ADD: + v = ext3::add(vals[(uint64_t)nd.a * vstride], vals[(uint64_t)nd.b * vstride]); + break; + case OP_SUB: + v = ext3::sub(vals[(uint64_t)nd.a * vstride], vals[(uint64_t)nd.b * vstride]); + break; + case OP_MUL: + v = ext3::mul(vals[(uint64_t)nd.a * vstride], vals[(uint64_t)nd.b * vstride]); + break; + case OP_NEG: + v = ext3::neg(vals[(uint64_t)nd.a * vstride]); + break; + case OP_EMBED: + v = vals[(uint64_t)nd.a * vstride]; + break; + default: + v = ext3::zero(); + break; + } + vals[i * vstride] = v; + } + + // Transition: z_inv * Σ_c beta_c * C_c. + Fe3 sum = ext3::zero(); + for (uint64_t c = 0; c < num_roots; c++) { + Fe3 cval = vals[(uint64_t)d_roots[c] * vstride]; + sum = ext3::add(sum, ext3::mul(d_beta_trans[c], cval)); + } + Fe3 h = ext3::mul_base(sum, d_z_inv[row % z_len]); + + // Boundary: Σ_b z_b_inv[row] * beta_b * (trace[col_b] - value_b). + for (uint64_t b = 0; b < num_boundary; b++) { + uint64_t col = d_b_col[b]; + Fe3 tcell; + if (d_b_is_aux[b] != 0) { + uint64_t base = col * 3; + tcell = ext3::make(d_aux[(base + 0) * aux_stride + row], + d_aux[(base + 1) * aux_stride + row], + d_aux[(base + 2) * aux_stride + row]); + } else { + tcell = ext3::make(d_main[col * main_stride + row], 0, 0); + } + Fe3 bp = ext3::sub(tcell, d_b_value[b]); + // (z_b_inv * beta_b) * bp — matches the CPU op order. + Fe3 zb = ext3::mul_base(d_b_beta[b], d_b_z_inv[b * num_rows + row]); + h = ext3::add(h, ext3::mul(zb, bp)); + } + + d_h[row] = h; + } +} diff --git a/crypto/math-cuda/src/constraint_interp.rs b/crypto/math-cuda/src/constraint_interp.rs new file mode 100644 index 000000000..ba6baf8ea --- /dev/null +++ b/crypto/math-cuda/src/constraint_interp.rs @@ -0,0 +1,270 @@ +//! Host wrapper for the transition-constraint interpreter kernel +//! (`kernels/constraint_interp.cu`). +//! +//! Takes a constraint program already lowered to flat `u64` device arrays (by +//! `stark::constraint_ir::device::DeviceProgram`) plus the device-resident LDE +//! handles, uploads the program + per-proof uniforms, launches the interpreter +//! over every LDE row, and returns the per-constraint eval matrix. +//! +//! Layering note: this crate cannot see `stark`'s `DeviceProgram` type (stark +//! depends on math-cuda, not the reverse), so the caller flattens the program +//! into the raw `u64` slices below. The stark-side dispatch +//! (`stark::constraint_ir::gpu_interp`) owns that flattening + the TypeId gate. + +use cudarc::driver::{LaunchConfig, PushKernelArg}; + +use crate::Result; +use crate::device::backend; +use crate::lde::{GpuLdeBase, GpuLdeExt3}; + +const BLOCK_DIM: u32 = 256; +/// Cap on total threads (grid × block). Each thread owns `num_nodes` ext3 slots +/// in the global value scratch (`num_nodes × MAX_THREADS × 24 B`), so a fixed +/// cap bounds that buffer regardless of LDE size; threads grid-stride over the +/// remaining rows. 65536 mirrors OpenVM's quotient `TASK_SIZE`. +const MAX_THREADS: u32 = 1 << 16; + +/// Evaluate every constraint of a lowered program over the device-resident LDE. +/// +/// Returns the per-constraint eval matrix as raw ext3 limbs, constraint-major: +/// constraint `c`, row `r`, component `k` at `out[(c * num_rows + r) * 3 + k]`. +/// Base-rooted constraints carry their value in component 0. +/// +/// Inputs (all raw limbs, matching the crate's u64 device convention): +/// - `nodes`: 2 `u64` per IR node (`op | a<<32`, then `b | dim<<32`). +/// - `base_consts`: one `u64` per base constant. +/// - `ext_consts`, `rap_challenges`, `alpha_powers`: 3 `u64` per element. +/// - `table_offset`: exactly 3 `u64`. +/// - `roots`: one `u64` node id per constraint. +/// - `main`/`aux`: device-resident LDE handles; `next_step` is the LDE row +/// stride for a frame-offset step; `num_rows` is the number of LDE rows. +#[allow(clippy::too_many_arguments)] +pub fn eval_constraints_on_device( + nodes: &[u64], + num_nodes: usize, + base_consts: &[u64], + ext_consts: &[u64], + roots: &[u64], + rap_challenges: &[u64], + alpha_powers: &[u64], + table_offset: &[u64], + main: &GpuLdeBase, + aux: &GpuLdeExt3, + next_step: usize, + num_rows: usize, +) -> Result> { + let num_roots = roots.len(); + if num_rows == 0 || num_roots == 0 || num_nodes == 0 { + return Ok(vec![0u64; num_roots * num_rows * 3]); + } + debug_assert_eq!(nodes.len(), 2 * num_nodes, "2 u64 per node"); + debug_assert_eq!(table_offset.len(), 3, "table_offset is one ext3 element"); + + let be = backend()?; + let stream = be.next_stream(); + + // Upload the program + uniforms (the column data never crosses PCIe — it is + // already resident in `main.buf` / `aux.buf`). + let d_nodes = stream.clone_htod(nodes)?; + let d_base_consts = stream.clone_htod(base_consts)?; + let d_ext_consts = stream.clone_htod(ext_consts)?; + let d_roots = stream.clone_htod(roots)?; + let d_rap = stream.clone_htod(rap_challenges)?; + let d_alpha = stream.clone_htod(alpha_powers)?; + let d_offset = stream.clone_htod(table_offset)?; + + // Fixed thread count, grid-stride over rows. + let max_grid = MAX_THREADS / BLOCK_DIM; + let grid = (num_rows as u32).div_ceil(BLOCK_DIM).clamp(1, max_grid); + let num_threads = (grid as usize) * (BLOCK_DIM as usize); + + // Per-thread value scratch ([num_nodes * num_threads] ext3) and output. + let mut d_values = stream.alloc_zeros::(num_nodes * num_threads * 3)?; + let mut d_evals = stream.alloc_zeros::(num_roots * num_rows * 3)?; + + let num_nodes_u64 = num_nodes as u64; + let num_roots_u64 = num_roots as u64; + let main_stride = main.lde_size as u64; + let aux_stride = aux.lde_size as u64; + let next_step_u64 = next_step as u64; + let num_rows_u64 = num_rows as u64; + + let cfg = LaunchConfig { + grid_dim: (grid, 1, 1), + block_dim: (BLOCK_DIM, 1, 1), + shared_mem_bytes: 0, + }; + unsafe { + stream + .launch_builder(&be.constraint_interp_kernel) + .arg(&mut d_evals) + .arg(&d_nodes) + .arg(&num_nodes_u64) + .arg(&d_base_consts) + .arg(&d_ext_consts) + .arg(&d_roots) + .arg(&num_roots_u64) + .arg(&d_rap) + .arg(&d_alpha) + .arg(&d_offset) + .arg(main.buf.as_ref()) + .arg(&main_stride) + .arg(aux.buf.as_ref()) + .arg(&aux_stride) + .arg(&next_step_u64) + .arg(&num_rows_u64) + .arg(&mut d_values) + .launch(cfg)?; + } + let out = stream.clone_dtoh(&d_evals)?; + stream.synchronize()?; + Ok(out) +} + +/// The per-proof accumulation inputs that turn per-constraint evals into the +/// composition-poly evaluation `H(row)` (all raw limbs; see +/// [`eval_composition_on_device`]). +pub struct CompositionAccum<'a> { + /// Transition combination coefficients β, one ext3 per constraint root + /// (`num_roots * 3` u64). + pub beta_trans: &'a [u64], + /// Cyclic transition-zerofier inverse, base field (`z_len` u64), indexed + /// `row % z_len`. + pub z_inv: &'a [u64], + /// Boundary constraint columns (`num_boundary` u64). + pub b_col: &'a [u64], + /// Boundary main/aux selector, 0 = main / 1 = aux (`num_boundary` u64). + pub b_is_aux: &'a [u64], + /// Boundary target values (`num_boundary * 3` u64, ext3). + pub b_value: &'a [u64], + /// Boundary combination coefficients β_b (`num_boundary * 3` u64, ext3). + pub b_beta: &'a [u64], + /// Boundary zerofier inverses, base field (`num_boundary * num_rows` u64), + /// indexed `b * num_rows + row`. + pub b_z_inv: &'a [u64], +} + +/// Evaluate the constraints AND fuse the composition accumulation on-device: +/// `H(row) = z_inv[row]·Σ βᵢ·Cᵢ + Σ_b z_b_inv[row]·β_b·(trace_b − value_b)`, +/// returning `H` as raw ext3 limbs (`num_rows * 3` u64, `out[row*3 + k]`). No +/// per-constraint matrix is materialized. +/// +/// Uniform-zerofier case only (the VM has no end-exemptions); the caller gates +/// on `is_uniform` and falls back to CPU otherwise. +#[allow(clippy::too_many_arguments)] +pub fn eval_composition_on_device( + nodes: &[u64], + num_nodes: usize, + base_consts: &[u64], + ext_consts: &[u64], + roots: &[u64], + rap_challenges: &[u64], + alpha_powers: &[u64], + table_offset: &[u64], + main: &GpuLdeBase, + aux: &GpuLdeExt3, + next_step: usize, + num_rows: usize, + accum: &CompositionAccum, +) -> Result> { + let num_roots = roots.len(); + if num_rows == 0 { + return Ok(Vec::new()); + } + debug_assert_eq!(nodes.len(), 2 * num_nodes, "2 u64 per node"); + debug_assert_eq!(accum.beta_trans.len(), num_roots * 3, "β per root"); + let num_boundary = accum.b_col.len(); + debug_assert_eq!( + accum.b_z_inv.len(), + num_boundary * num_rows, + "z_b_inv shape" + ); + // The kernel indexes these by `num_boundary`; a caller mismatch would be an + // OOB device read rather than a clean panic, so pin all boundary shapes. + debug_assert_eq!(accum.b_is_aux.len(), num_boundary, "b_is_aux per boundary"); + debug_assert_eq!( + accum.b_value.len(), + num_boundary * 3, + "b_value ext3 per boundary" + ); + debug_assert_eq!( + accum.b_beta.len(), + num_boundary * 3, + "b_beta ext3 per boundary" + ); + + let be = backend()?; + let stream = be.next_stream(); + + let d_nodes = stream.clone_htod(nodes)?; + let d_base_consts = stream.clone_htod(base_consts)?; + let d_ext_consts = stream.clone_htod(ext_consts)?; + let d_roots = stream.clone_htod(roots)?; + let d_rap = stream.clone_htod(rap_challenges)?; + let d_alpha = stream.clone_htod(alpha_powers)?; + let d_offset = stream.clone_htod(table_offset)?; + + let d_beta_trans = stream.clone_htod(accum.beta_trans)?; + let d_z_inv = stream.clone_htod(accum.z_inv)?; + let d_b_col = stream.clone_htod(accum.b_col)?; + let d_b_is_aux = stream.clone_htod(accum.b_is_aux)?; + let d_b_value = stream.clone_htod(accum.b_value)?; + let d_b_beta = stream.clone_htod(accum.b_beta)?; + let d_b_z_inv = stream.clone_htod(accum.b_z_inv)?; + + let max_grid = MAX_THREADS / BLOCK_DIM; + let grid = (num_rows as u32).div_ceil(BLOCK_DIM).clamp(1, max_grid); + let num_threads = (grid as usize) * (BLOCK_DIM as usize); + + let mut d_values = stream.alloc_zeros::(num_nodes * num_threads * 3)?; + let mut d_h = stream.alloc_zeros::(num_rows * 3)?; + + let num_nodes_u64 = num_nodes as u64; + let num_roots_u64 = num_roots as u64; + let main_stride = main.lde_size as u64; + let aux_stride = aux.lde_size as u64; + let next_step_u64 = next_step as u64; + let num_rows_u64 = num_rows as u64; + let z_len_u64 = (accum.z_inv.len() as u64).max(1); + let num_boundary_u64 = num_boundary as u64; + + let cfg = LaunchConfig { + grid_dim: (grid, 1, 1), + block_dim: (BLOCK_DIM, 1, 1), + shared_mem_bytes: 0, + }; + unsafe { + stream + .launch_builder(&be.constraint_composition_kernel) + .arg(&mut d_h) + .arg(&d_nodes) + .arg(&num_nodes_u64) + .arg(&d_base_consts) + .arg(&d_ext_consts) + .arg(&d_roots) + .arg(&num_roots_u64) + .arg(&d_rap) + .arg(&d_alpha) + .arg(&d_offset) + .arg(main.buf.as_ref()) + .arg(&main_stride) + .arg(aux.buf.as_ref()) + .arg(&aux_stride) + .arg(&next_step_u64) + .arg(&num_rows_u64) + .arg(&d_beta_trans) + .arg(&d_z_inv) + .arg(&z_len_u64) + .arg(&num_boundary_u64) + .arg(&d_b_col) + .arg(&d_b_is_aux) + .arg(&d_b_value) + .arg(&d_b_beta) + .arg(&d_b_z_inv) + .arg(&mut d_values) + .launch(cfg)?; + } + let out = stream.clone_dtoh(&d_h)?; + stream.synchronize()?; + Ok(out) +} diff --git a/crypto/math-cuda/src/device.rs b/crypto/math-cuda/src/device.rs index b1d6ca489..be0440d5e 100644 --- a/crypto/math-cuda/src/device.rs +++ b/crypto/math-cuda/src/device.rs @@ -98,6 +98,8 @@ const DEEP_PTX: &str = include_str!(concat!(env!("OUT_DIR"), "/deep.ptx")); const FRI_PTX: &str = include_str!(concat!(env!("OUT_DIR"), "/fri.ptx")); const INVERSE_PTX: &str = include_str!(concat!(env!("OUT_DIR"), "/inverse.ptx")); const LOGUP_PTX: &str = include_str!(concat!(env!("OUT_DIR"), "/logup.ptx")); +const CONSTRAINT_INTERP_PTX: &str = + include_str!(concat!(env!("OUT_DIR"), "/constraint_interp.ptx")); /// Number of CUDA streams in the pool. Larger pools let many rayon-parallel /// callers overlap on the GPU without serializing on stream ownership. The @@ -188,6 +190,10 @@ pub struct Backend { pub logup_finalize_accum_ext3: CudaFunction, pub logup_assemble_aux_ext3: CudaFunction, + // constraint_interp.ptx + pub constraint_interp_kernel: CudaFunction, + pub constraint_composition_kernel: CudaFunction, + // Twiddle caches keyed by log_n. fwd_twiddles: Mutex>>>>, inv_twiddles: Mutex>>>>, @@ -291,6 +297,7 @@ impl Backend { let fri = ctx.load_module(Ptx::from_src(FRI_PTX))?; let inverse = ctx.load_module(Ptx::from_src(INVERSE_PTX))?; let logup = ctx.load_module(Ptx::from_src(LOGUP_PTX))?; + let constraint_interp = ctx.load_module(Ptx::from_src(CONSTRAINT_INTERP_PTX))?; let mut streams = Vec::with_capacity(STREAM_POOL_SIZE); for _ in 0..STREAM_POOL_SIZE { @@ -378,6 +385,10 @@ impl Backend { logup_apply_offsets_add_ext3: logup.load_function("logup_apply_offsets_add_ext3")?, logup_finalize_accum_ext3: logup.load_function("logup_finalize_accum_ext3")?, logup_assemble_aux_ext3: logup.load_function("logup_assemble_aux_ext3")?, + constraint_interp_kernel: constraint_interp + .load_function("constraint_interp_kernel")?, + constraint_composition_kernel: constraint_interp + .load_function("constraint_composition_kernel")?, fwd_twiddles: Mutex::new(vec![None; max_log]), inv_twiddles: Mutex::new(vec![None; max_log]), ctx, diff --git a/crypto/math-cuda/src/lib.rs b/crypto/math-cuda/src/lib.rs index abc487750..493480991 100644 --- a/crypto/math-cuda/src/lib.rs +++ b/crypto/math-cuda/src/lib.rs @@ -6,6 +6,7 @@ //! pipelines or used by the parity test suite. pub mod barycentric; +pub mod constraint_interp; pub mod deep; pub mod device; pub mod fri; diff --git a/crypto/stark/src/constraint_ir/device.rs b/crypto/stark/src/constraint_ir/device.rs new file mode 100644 index 000000000..6c522d103 --- /dev/null +++ b/crypto/stark/src/constraint_ir/device.rs @@ -0,0 +1,477 @@ +//! Device (GPU) lowering of a Goldilocks [`ConstraintProgram`] to a flat, +//! `#[repr(C)]` blob a CUDA interpreter kernel can walk directly — plus a CPU +//! reference walker over that same blob. +//! +//! This is the *one* concrete lowering point: the rest of the constraint IR is +//! field-generic (``), but genericity does not cross to CUDA, so here we +//! commit to the Goldilocks base field and its degree-3 extension. Any other +//! field tower never reaches this module — it stays on the generic +//! [`interp`](super::interp) path. +//! +//! Two things live here: +//! +//! - [`DeviceProgram::lower`] — serialize a `ConstraintProgram` into flat, device-uploadable arrays: a `#[repr(C)]` +//! [`DeviceNode`] list, `u64` / `[u64; 3]` constant tables, `roots`, and +//! `num_base`. Field constants become raw limbs via `FieldElement::to_raw`, +//! which is byte-identical to how [`crate::gpu_lde`] already hands Goldilocks +//! elements to the device (a `#[repr(transparent)]` `u64` / `[u64; 3]`). +//! +//! - [`eval_device_program`] — a CPU forward pass over the *flat* node array +//! (not the [`Op`] enum), decoding leaves from the raw limb tables and +//! reproducing the exact [`interp::run`](super::interp) semantics +//! (per-node [`Dim`] drives base-vs-extension arithmetic, mixed operands +//! auto-embed). It is the model of the GPU kernel's per-thread walk, so a +//! bit-for-bit match against [`eval_program`](super::interp::eval_program) +//! pins the on-device layout and control flow *before* any CUDA exists. The +//! kernel is then a transliteration of this walk with the `FieldElement` +//! arithmetic swapped for `goldilocks.cuh` / `ext3.cuh`. + +use math::field::element::FieldElement; +use math::field::extensions_goldilocks::Degree3GoldilocksExtensionField as GoldilocksExtension; +use math::field::goldilocks::GoldilocksField; + +use super::ir::{ConstraintProgram, Dim, Op}; + +type FpE = FieldElement; +type Ext3E = FieldElement; + +// ------------------------------------------------------------------------- +// Wire tags — MUST match the CUDA kernel's `switch (op)` and dim checks. +// ------------------------------------------------------------------------- + +/// `a` = index into `base_consts`. +pub const OP_CONST_BASE: u32 = 0; +/// `a` = index into `ext_consts`. +pub const OP_CONST_EXT: u32 = 1; +/// Trace-cell read; `a`/`b` pack the [`Op::Var`] fields (see [`pack_var`]). +pub const OP_VAR: u32 = 2; +/// `a` = index into the per-proof `rap_challenges` uniform buffer. +pub const OP_RAP_CHALLENGE: u32 = 3; +/// `a` = index into the per-proof `logup_alpha_powers` uniform buffer. +pub const OP_ALPHA_POW: u32 = 4; +/// The per-proof LogUp table offset uniform; no operands. +pub const OP_TABLE_OFFSET: u32 = 5; +/// `a`, `b` = node ids. +pub const OP_ADD: u32 = 6; +/// `a`, `b` = node ids. +pub const OP_SUB: u32 = 7; +/// `a`, `b` = node ids. +pub const OP_MUL: u32 = 8; +/// `a` = node id. +pub const OP_NEG: u32 = 9; +/// `a` = node id (base → extension embed). +pub const OP_EMBED: u32 = 10; + +/// Node result is a base-field value. +pub const DIM_BASE: u32 = 0; +/// Node result is an extension-field value. +pub const DIM_EXT: u32 = 1; + +/// One flattened IR instruction: 16 bytes, `#[repr(C)]` for a 1:1 device +/// upload. `op` is an `OP_*` tag; the meaning of `a`/`b` depends on `op` (node +/// ids for arithmetic, table indices for constants/uniforms, packed [`Op::Var`] +/// fields for [`OP_VAR`]); `dim` is [`DIM_BASE`] or [`DIM_EXT`]. +#[repr(C)] +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct DeviceNode { + pub op: u32, + pub a: u32, + pub b: u32, + pub dim: u32, +} + +/// Pack an [`Op::Var`]'s fields into the `(a, b)` operand words: +/// `a` holds `col`; `b` holds `main` (bit 16), `offset` (bits 8..16), `row` +/// (bits 0..8). The CUDA kernel unpacks with the same layout. +#[inline] +pub fn pack_var(main: bool, offset: u8, row: u8, col: u16) -> (u32, u32) { + let a = col as u32; + let b = ((main as u32) << 16) | ((offset as u32) << 8) | (row as u32); + (a, b) +} + +/// Inverse of [`pack_var`]: `(main, offset, row, col)`. +#[inline] +pub fn unpack_var(a: u32, b: u32) -> (bool, u8, u8, u16) { + let col = (a & 0xFFFF) as u16; + let main = ((b >> 16) & 1) != 0; + let offset = ((b >> 8) & 0xFF) as u8; + let row = (b & 0xFF) as u8; + (main, offset, row, col) +} + +/// A [`ConstraintProgram`] lowered to flat, device-uploadable arrays. Constants +/// are canonical raw limbs (`u64` base / `[u64; 3]` extension), matching the +/// `#[repr(transparent)]` layout the GPU trace buffers already use. +#[derive(Clone, Debug)] +pub struct DeviceProgram { + /// Topologically ordered instruction list (id `i` references only `< i`). + pub nodes: Vec, + /// Base-field constant table, indexed by [`OP_CONST_BASE`]. + pub base_consts: Vec, + /// Extension-field constant table, indexed by [`OP_CONST_EXT`]. + pub ext_consts: Vec<[u64; 3]>, + /// Per-constraint root node ids, indexed by `constraint_idx`. + pub roots: Vec, + /// Number of leading ([`DIM_BASE`]-rooted) constraints written to + /// `base_evals`; the rest go to `ext_evals`. + pub num_base: u32, +} + +impl DeviceProgram { + /// Lower a concrete-Goldilocks [`ConstraintProgram`] to its flat device + /// form. Pure serialization — no field arithmetic, no device access. + pub fn lower(prog: &ConstraintProgram) -> Self { + let nodes = prog + .nodes + .iter() + .zip(prog.dims.iter()) + .map(|(op, dim)| { + let dim = match dim { + Dim::Base => DIM_BASE, + Dim::Ext => DIM_EXT, + }; + let (op, a, b) = match *op { + Op::ConstBase(idx) => (OP_CONST_BASE, idx, 0), + Op::ConstExt(idx) => (OP_CONST_EXT, idx, 0), + Op::Var { + main, + offset, + row, + col, + } => { + let (a, b) = pack_var(main, offset, row, col); + (OP_VAR, a, b) + } + Op::RapChallenge { idx } => (OP_RAP_CHALLENGE, idx as u32, 0), + Op::AlphaPow { idx } => (OP_ALPHA_POW, idx as u32, 0), + Op::TableOffset => (OP_TABLE_OFFSET, 0, 0), + Op::Add(a, b) => (OP_ADD, a, b), + Op::Sub(a, b) => (OP_SUB, a, b), + Op::Mul(a, b) => (OP_MUL, a, b), + Op::Neg(a) => (OP_NEG, a, 0), + Op::Embed(a) => (OP_EMBED, a, 0), + }; + DeviceNode { op, a, b, dim } + }) + .collect(); + + let base_consts = prog.base_consts.iter().map(|c| *c.value()).collect(); + let ext_consts = prog.ext_consts.iter().map(encode_ext).collect(); + + DeviceProgram { + nodes, + base_consts, + ext_consts, + roots: prog.roots.clone(), + num_base: prog.num_base as u32, + } + } +} + +/// A node's computed value during the walk: base or extension field element. +#[derive(Clone)] +enum Value { + Base(FpE), + Ext(Ext3E), +} + +impl Value { + /// Promote to the extension field, embedding a base value if needed. + fn to_ext(&self) -> Ext3E { + match self { + Value::Base(x) => (*x).to_extension::(), + Value::Ext(x) => *x, + } + } + + fn as_base(&self) -> &FpE { + match self { + Value::Base(x) => x, + Value::Ext(_) => panic!("expected a base value but found an extension value"), + } + } +} + +/// `[u64; 3]` → extension element. +#[inline] +fn decode_ext(limbs: [u64; 3]) -> Ext3E { + Ext3E::from_raw([ + FpE::from_raw(limbs[0]), + FpE::from_raw(limbs[1]), + FpE::from_raw(limbs[2]), + ]) +} + +/// Extension element → `[u64; 3]`. +#[inline] +fn encode_ext(x: &Ext3E) -> [u64; 3] { + let limbs = x.value(); + [*limbs[0].value(), *limbs[1].value(), *limbs[2].value()] +} + +/// Apply a binary op, auto-embedding to the extension when the result dimension +/// is [`DIM_EXT`] (or either operand is already an extension value) — the exact +/// rule of [`interp::binop`](super::interp). +#[inline] +fn binop( + values: &[Value], + a: u32, + b: u32, + dim: u32, + base_op: impl Fn(FpE, FpE) -> FpE, + ext_op: impl Fn(Ext3E, Ext3E) -> Ext3E, +) -> Value { + let va = &values[a as usize]; + let vb = &values[b as usize]; + if dim == DIM_BASE + && let (Value::Base(x), Value::Base(y)) = (va, vb) + { + return Value::Base(base_op(*x, *y)); + } + Value::Ext(ext_op(va.to_ext(), vb.to_ext())) +} + +/// Full prover-shaped forward pass over the *flat* device blob, in raw limbs — +/// the CPU model of the GPU kernel. Mirrors +/// [`eval_program`](super::interp::eval_program): base-rooted constraints +/// (`c < num_base`) land in `base_evals`, the rest in `ext_evals`, with the +/// same auto-embed semantics. +/// +/// `main[offset][col]` / `aux[offset][col]` are the frame's trace cells; +/// `rap_challenges` / `alpha_powers` / `table_offset` are the per-proof +/// uniforms. All values are the same raw `u64` / `[u64; 3]` limbs the device +/// buffers carry. +#[allow(clippy::too_many_arguments)] +pub fn eval_device_program( + dev: &DeviceProgram, + main: &[Vec], + aux: &[Vec<[u64; 3]>], + rap_challenges: &[[u64; 3]], + alpha_powers: &[[u64; 3]], + table_offset: [u64; 3], + base_evals: &mut [u64], + ext_evals: &mut [[u64; 3]], +) { + let mut values: Vec = Vec::with_capacity(dev.nodes.len()); + + for node in &dev.nodes { + let v = match node.op { + OP_CONST_BASE => Value::Base(FpE::from_raw(dev.base_consts[node.a as usize])), + OP_CONST_EXT => Value::Ext(decode_ext(dev.ext_consts[node.a as usize])), + OP_VAR => { + let (is_main, offset, _row, col) = unpack_var(node.a, node.b); + if is_main { + Value::Base(FpE::from_raw(main[offset as usize][col as usize])) + } else { + Value::Ext(decode_ext(aux[offset as usize][col as usize])) + } + } + OP_RAP_CHALLENGE => Value::Ext(decode_ext(rap_challenges[node.a as usize])), + OP_ALPHA_POW => Value::Ext(decode_ext(alpha_powers[node.a as usize])), + OP_TABLE_OFFSET => Value::Ext(decode_ext(table_offset)), + OP_ADD => binop( + &values, + node.a, + node.b, + node.dim, + |x, y| x + y, + |x, y| x + y, + ), + OP_SUB => binop( + &values, + node.a, + node.b, + node.dim, + |x, y| x - y, + |x, y| x - y, + ), + OP_MUL => binop( + &values, + node.a, + node.b, + node.dim, + |x, y| x * y, + |x, y| x * y, + ), + OP_NEG => { + let val = &values[node.a as usize]; + if node.dim == DIM_BASE { + match val { + Value::Base(x) => Value::Base(-x), + // Dim/value mismatch: keep it in the extension, as interp does. + Value::Ext(x) => Value::Ext(-*x), + } + } else { + Value::Ext(-val.to_ext()) + } + } + OP_EMBED => Value::Ext(values[node.a as usize].to_ext()), + other => panic!("unknown device op tag {other}"), + }; + values.push(v); + } + + for (c, &root) in dev.roots.iter().enumerate() { + let v = &values[root as usize]; + if (c as u32) < dev.num_base { + base_evals[c] = *v.as_base().value(); + } else { + ext_evals[c] = encode_ext(&v.to_ext()); + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::constraint_ir::builder::IrBuilder; + use crate::constraint_ir::interp::eval_program; + use crate::frame::Frame; + use crate::table::TableView; + use crate::traits::TransitionEvaluationContext; + + type Gl = GoldilocksField; + type Ext = GoldilocksExtension; + + fn fp(v: u64) -> FpE { + FpE::from(v) + } + fn ext3(a: u64, b: u64, c: u64) -> Ext3E { + Ext3E::from_raw([fp(a), fp(b), fp(c)]) + } + + struct SplitMix64(u64); + impl SplitMix64 { + fn next_u64(&mut self) -> u64 { + self.0 = self.0.wrapping_add(0x9E37_79B9_7F4A_7C15); + let mut z = self.0; + z = (z ^ (z >> 30)).wrapping_mul(0xBF58_476D_1CE4_E5B9); + z = (z ^ (z >> 27)).wrapping_mul(0x94D0_49BB_1331_11EB); + z ^ (z >> 31) + } + fn ext(&mut self) -> Ext3E { + ext3(self.next_u64(), self.next_u64(), self.next_u64()) + } + } + + #[test] + fn pack_var_roundtrips() { + for &(m, off, row, col) in &[ + (true, 0u8, 0u8, 0u16), + (false, 1, 0, 7), + (true, 3, 0, 65535), + (false, 255, 255, 12345), + ] { + let (a, b) = pack_var(m, off, row, col); + assert_eq!(unpack_var(a, b), (m, off, row, col)); + } + } + + /// A program that exercises every `Op` variant and both dims, with a + /// base-rooted constraint and extension (LogUp-shaped) roots, next-row + /// reads, and mixed base×ext arithmetic. Roots: 0 (base) is a pure base + /// expression; 1 and 2 (ext) touch challenges, alpha powers, table offset, + /// aux, embed and negation. + fn all_ops_program() -> ConstraintProgram { + let mut b = IrBuilder::::new(); + + // Root 0 (base): (m0 + m1) * 2 - m0_next , all base, incl. next-row. + let m0 = b.main(0, 0); + let m1 = b.main(0, 1); + let m0n = b.main(1, 0); + let two = b.const_base(2); + let sum = b.add(m0, m1); + let scaled = b.mul(sum, two); + let base_root = b.sub(scaled, m0n); + b.emit(0, base_root); + + // Root 1 (ext): m0 * challenge(0) + alpha_pow(1) * aux(0,0) - table_offset + let ch = b.challenge(0); + let ap = b.alpha_power(1); + let au = b.aux(0, 0); + let off = b.table_offset(); + let t1 = b.mul(m0, ch); // base × ext → ext (auto-embed) + let t2 = b.mul(ap, au); // ext × ext + let s = b.add(t1, t2); + let ext_root = b.sub(s, off); + b.emit(1, ext_root); + + // Root 2 (ext): embed(m1) + (-aux(0,1)) + const_ext + let em = b.embed(m1); + let au1 = b.aux(0, 1); + let nau1 = b.neg(au1); // ext negation + let ce = b.const_ext(ext3(9, 8, 7)); + let s2 = b.add(em, nau1); + let ext_root2 = b.add(s2, ce); + b.emit(2, ext_root2); + + b.finish(1) // 1 base root, 2 ext roots + } + + #[test] + fn device_walk_matches_interp_all_ops() { + let prog = all_ops_program(); + let dev = DeviceProgram::lower(&prog); + + let mut rng = SplitMix64(0x0123_4567_89AB_CDEF); + for _ in 0..1000 { + // Two frame steps (offset 0 and 1), 2 main cols + 2 aux cols each. + let main_vals: Vec> = (0..2) + .map(|_| vec![fp(rng.next_u64()), fp(rng.next_u64())]) + .collect(); + let aux_vals: Vec> = (0..2).map(|_| vec![rng.ext(), rng.ext()]).collect(); + let rap = vec![rng.ext(), rng.ext()]; + let alpha = vec![rng.ext(), rng.ext()]; + let offset = rng.ext(); + + // Reference: the generic interpreter over the ConstraintProgram. + let steps: Vec> = main_vals + .iter() + .zip(aux_vals.iter()) + .map(|(m, a)| TableView::::new(vec![m.clone()], vec![a.clone()])) + .collect(); + let frame = Frame::::new(steps); + let ctx = TransitionEvaluationContext::new_prover( + frame.as_row_frame(), + &rap, + &alpha, + &offset, + ); + let mut base_ref = vec![FpE::zero(); 1]; + let mut ext_ref = vec![Ext3E::zero(); 3]; + eval_program(&prog, &ctx, &mut base_ref, &mut ext_ref); + + // Device walk over the flat blob, in raw limbs. + let main_raw: Vec> = main_vals + .iter() + .map(|r| r.iter().map(|x| *x.value()).collect()) + .collect(); + let aux_raw: Vec> = aux_vals + .iter() + .map(|r| r.iter().map(encode_ext).collect()) + .collect(); + let rap_raw: Vec<[u64; 3]> = rap.iter().map(encode_ext).collect(); + let alpha_raw: Vec<[u64; 3]> = alpha.iter().map(encode_ext).collect(); + let mut base_dev = vec![0u64; 1]; + let mut ext_dev = vec![[0u64; 3]; 3]; + eval_device_program( + &dev, + &main_raw, + &aux_raw, + &rap_raw, + &alpha_raw, + encode_ext(&offset), + &mut base_dev, + &mut ext_dev, + ); + + // Base root 0. + assert_eq!(base_dev[0], *base_ref[0].value()); + // Ext roots 1 and 2 (absolute indices; slot 0 unused for ext). + assert_eq!(ext_dev[1], encode_ext(&ext_ref[1])); + assert_eq!(ext_dev[2], encode_ext(&ext_ref[2])); + } + } +} diff --git a/crypto/stark/src/constraint_ir/gpu_interp.rs b/crypto/stark/src/constraint_ir/gpu_interp.rs new file mode 100644 index 000000000..8540ce81e --- /dev/null +++ b/crypto/stark/src/constraint_ir/gpu_interp.rs @@ -0,0 +1,269 @@ +//! GPU dispatch for the constraint interpreter (the device edge). +//! +//! Lowers a captured [`ConstraintProgram`] to its flat device blob +//! ([`DeviceProgram`]), flattens it plus the per-proof uniforms into the raw +//! `u64` slices the kernel reads, and launches +//! [`math_cuda::constraint_interp::eval_constraints_on_device`] over the +//! device-resident LDE. Returns the per-constraint eval matrix, or `None` to +//! signal the caller to fall back to the CPU path. +//! +//! This is the *one* concrete-Goldilocks lowering point: the IR is field-generic +//! everywhere else, and genericity does not cross to CUDA. A `TypeId` gate +//! establishes `F = GoldilocksField` / `E = Degree3GoldilocksExtensionField` +//! before a single `unsafe` reinterpret to the concrete program — the same +//! device-edge seam `crate::gpu_lde` uses for the LDE. +//! +//! The whole module is `#[cfg(feature = "cuda")]`; without the feature the +//! caller only ever has the CPU interpreter. + +use std::any::TypeId; + +use math::field::element::FieldElement; +use math::field::extensions_goldilocks::Degree3GoldilocksExtensionField as GoldilocksExtension; +use math::field::goldilocks::GoldilocksField; +use math::field::traits::IsField; + +use math_cuda::lde::{GpuLdeBase, GpuLdeExt3}; + +use super::device::DeviceProgram; +use super::ir::ConstraintProgram; + +/// Pack the lowered node list into 2 `u64` per node (`op | a<<32`, `b | dim<<32`), +/// the encoding the kernel's `load_node` decodes. +fn pack_nodes(dev: &DeviceProgram) -> Vec { + let mut out = Vec::with_capacity(dev.nodes.len() * 2); + for n in &dev.nodes { + out.push(n.op as u64 | ((n.a as u64) << 32)); + out.push(n.b as u64 | ((n.dim as u64) << 32)); + } + out +} + +/// Flatten `[[u64; 3]]` ext3 limbs to a contiguous `u64` slice. +fn flatten_ext3(xs: &[[u64; 3]]) -> Vec { + xs.iter().flat_map(|e| e.iter().copied()).collect() +} + +/// Reinterpret a slice of ext3 field elements as flat `u64` (3 per element). +/// +/// # Safety +/// The caller must have established `E == Degree3GoldilocksExtensionField`, +/// whose `FieldElement` is `#[repr(transparent)]` over `[u64; 3]` — the same +/// invariant `crate::gpu_lde::columns_to_u64_ext3` relies on. +unsafe fn ext3_slice_to_u64(xs: &[FieldElement]) -> Vec { + let raw = unsafe { std::slice::from_raw_parts(xs.as_ptr() as *const u64, xs.len() * 3) }; + raw.to_vec() +} + +/// Reinterpret a slice of base field elements as flat `u64` (1 per element). +/// +/// # Safety +/// The caller must have established `F == GoldilocksField`, whose `FieldElement` +/// is `#[repr(transparent)]` over `u64`. +unsafe fn base_slice_to_u64(xs: &[FieldElement]) -> Vec { + let raw = unsafe { std::slice::from_raw_parts(xs.as_ptr() as *const u64, xs.len()) }; + raw.to_vec() +} + +/// Per-proof accumulation inputs (in `FieldElement` form) for +/// [`try_eval_composition_gpu`], mirroring the CPU accumulation in +/// `constraints::evaluator`. +pub struct CompositionInputs<'a, F: IsField, E: IsField> { + /// Transition coefficients β, one per constraint root. + pub beta_trans: &'a [FieldElement], + /// Cyclic transition-zerofier inverse (base field, `blowup`-length). + pub z_inv: &'a [FieldElement], + /// Boundary constraint columns. + pub b_col: &'a [usize], + /// Boundary main/aux selector. + pub b_is_aux: &'a [bool], + /// Boundary target values. + pub b_value: &'a [FieldElement], + /// Boundary coefficients β_b. + pub b_beta: &'a [FieldElement], + /// Boundary zerofier inverses (base field), laid out `b * num_rows + row`. + pub b_z_inv: &'a [FieldElement], +} + +/// The lowered device program plus the packed per-proof uniforms shared by both +/// GPU dispatch entry points. Produced by [`lower_and_pack`]. +struct LoweredCall { + dev: DeviceProgram, + nodes: Vec, + ext_consts: Vec, + roots: Vec, + rap: Vec, + alpha: Vec, + offset: Vec, +} + +/// The single concrete-Goldilocks lowering seam shared by +/// [`try_eval_composition_gpu`] and [`try_eval_program_gpu`]: gate on the +/// Goldilocks tower, reinterpret the generic program once, lower it to the flat +/// device blob, and pack the three ext3 uniforms. Returns `None` (→ CPU +/// fallback) for any other field tower. Factoring this keeps the sole `unsafe` +/// program reinterpret and the TypeId gate in one place instead of two. +fn lower_and_pack( + prog: &ConstraintProgram, + rap_challenges: &[FieldElement], + alpha_powers: &[FieldElement], + table_offset: &FieldElement, +) -> Option +where + F: IsField + 'static, + E: IsField + 'static, +{ + if TypeId::of::() != TypeId::of::() + || TypeId::of::() != TypeId::of::() + { + return None; + } + // SAFETY: the TypeId gate established `F = GoldilocksField` and + // `E = Degree3GoldilocksExtensionField`; the generic program has the exact + // layout of the concrete one (constants are `#[repr(transparent)]` over + // `u64` / `[u64; 3]`). + let prog: &ConstraintProgram = + unsafe { &*(prog as *const _ as *const _) }; + + let dev = DeviceProgram::lower(prog); + let nodes = pack_nodes(&dev); + let ext_consts = flatten_ext3(&dev.ext_consts); + let roots: Vec = dev.roots.iter().map(|&r| r as u64).collect(); + + // SAFETY: `E` is the ext3 tower (gated above). + let rap = unsafe { ext3_slice_to_u64(rap_challenges) }; + let alpha = unsafe { ext3_slice_to_u64(alpha_powers) }; + let offset = unsafe { ext3_slice_to_u64(std::slice::from_ref(table_offset)) }; + + Some(LoweredCall { + dev, + nodes, + ext_consts, + roots, + rap, + alpha, + offset, + }) +} + +/// Fused composition-poly evaluation on the GPU: returns `H(row)` as raw ext3 +/// limbs (`num_rows * 3` u64), or `None` for non-Goldilocks towers (→ CPU +/// fallback). `H(row) = z_inv·Σβᵢ·Cᵢ + Σ_b z_b_inv·β_b·(trace_b − value_b)`, +/// the uniform-zerofier accumulation of `evaluator::evaluate`. +#[allow(clippy::too_many_arguments)] +pub fn try_eval_composition_gpu( + prog: &ConstraintProgram, + main: &GpuLdeBase, + aux: &GpuLdeExt3, + rap_challenges: &[FieldElement], + alpha_powers: &[FieldElement], + table_offset: &FieldElement, + next_step: usize, + num_rows: usize, + inputs: &CompositionInputs, +) -> Option> +where + F: IsField + 'static, + E: IsField + 'static, +{ + let LoweredCall { + dev, + nodes, + ext_consts, + roots, + rap, + alpha, + offset, + } = lower_and_pack(prog, rap_challenges, alpha_powers, table_offset)?; + + // SAFETY: `E`/`F` are the Goldilocks tower (established in `lower_and_pack`). + let beta_trans = unsafe { ext3_slice_to_u64(inputs.beta_trans) }; + let z_inv = unsafe { base_slice_to_u64(inputs.z_inv) }; + let b_value = unsafe { ext3_slice_to_u64(inputs.b_value) }; + let b_beta = unsafe { ext3_slice_to_u64(inputs.b_beta) }; + let b_z_inv = unsafe { base_slice_to_u64(inputs.b_z_inv) }; + let b_col: Vec = inputs.b_col.iter().map(|&c| c as u64).collect(); + let b_is_aux: Vec = inputs.b_is_aux.iter().map(|&a| a as u64).collect(); + + let accum = math_cuda::constraint_interp::CompositionAccum { + beta_trans: &beta_trans, + z_inv: &z_inv, + b_col: &b_col, + b_is_aux: &b_is_aux, + b_value: &b_value, + b_beta: &b_beta, + b_z_inv: &b_z_inv, + }; + + let result = math_cuda::constraint_interp::eval_composition_on_device( + &nodes, + dev.nodes.len(), + &dev.base_consts, + &ext_consts, + &roots, + &rap, + &alpha, + &offset, + main, + aux, + next_step, + num_rows, + &accum, + ); + if result.is_ok() { + crate::gpu_lde::GPU_COMPOSITION_CALLS.fetch_add(1, std::sync::atomic::Ordering::Relaxed); + } + result.ok() +} + +/// Evaluate a captured program on the GPU, returning the per-constraint eval +/// matrix as raw ext3 limbs (constraint-major: constraint `c`, row `r`, +/// component `k` at `out[(c * num_rows + r) * 3 + k]`), or `None` if the field +/// tower is not the Goldilocks/degree-3 pair (→ CPU fallback). +/// +/// `main`/`aux` are the device-resident LDE handles; `rap_challenges`, +/// `alpha_powers`, `table_offset` are the per-proof uniforms; `next_step` is the +/// LDE row stride for a frame-offset step; `num_rows` is the LDE row count. +#[allow(clippy::too_many_arguments)] +pub fn try_eval_program_gpu( + prog: &ConstraintProgram, + main: &GpuLdeBase, + aux: &GpuLdeExt3, + rap_challenges: &[FieldElement], + alpha_powers: &[FieldElement], + table_offset: &FieldElement, + next_step: usize, + num_rows: usize, +) -> Option> +where + F: IsField + 'static, + E: IsField + 'static, +{ + let LoweredCall { + dev, + nodes, + ext_consts, + roots, + rap, + alpha, + offset, + } = lower_and_pack(prog, rap_challenges, alpha_powers, table_offset)?; + + let result = math_cuda::constraint_interp::eval_constraints_on_device( + &nodes, + dev.nodes.len(), + &dev.base_consts, + &ext_consts, + &roots, + &rap, + &alpha, + &offset, + main, + aux, + next_step, + num_rows, + ); + + // Any device error is mapped to a CPU fallback, never propagated. + result.ok() +} diff --git a/crypto/stark/src/constraint_ir/mod.rs b/crypto/stark/src/constraint_ir/mod.rs index 258aa23b7..380f32d86 100644 --- a/crypto/stark/src/constraint_ir/mod.rs +++ b/crypto/stark/src/constraint_ir/mod.rs @@ -16,12 +16,19 @@ //! - [`ir`]: the IR data structures ([`ConstraintProgram`], [`Op`], [`Dim`]). //! - [`builder`]: the [`IrBuilder`] and [`Expr`] capture API. //! - [`interp`]: a CPU forward-pass interpreter over the IR. +//! - [`device`]: the concrete-Goldilocks flat lowering ([`DeviceProgram`]) for +//! the GPU kernel, plus a CPU walker over that flat blob (the pre-GPU parity +//! oracle). //! //! [`ConstraintProgram`]: ir::ConstraintProgram //! [`Op`]: ir::Op //! [`Dim`]: ir::Dim +//! [`DeviceProgram`]: device::DeviceProgram pub mod builder; +pub mod device; +#[cfg(feature = "cuda")] +pub mod gpu_interp; pub mod interp; pub mod ir; @@ -29,5 +36,6 @@ pub mod ir; mod tests; pub use builder::{Expr, IrBuilder}; +pub use device::{DeviceNode, DeviceProgram, eval_device_program}; pub use interp::{eval_program, eval_program_base, eval_program_verifier}; pub use ir::{ConstraintProgram, Dim, Op}; diff --git a/crypto/stark/src/constraints/evaluator.rs b/crypto/stark/src/constraints/evaluator.rs index 48434b6e1..acd4fef99 100644 --- a/crypto/stark/src/constraints/evaluator.rs +++ b/crypto/stark/src/constraints/evaluator.rs @@ -182,7 +182,11 @@ where transition_coefficients: &[FieldElement], boundary_coefficients: &[FieldElement], rap_challenges: &[FieldElement], - ) -> Vec> { + ) -> Vec> + where + Field: 'static, + FieldExtension: 'static, + { let boundary_constraints = &self.boundary_constraints; let mut boundary_step_points: Vec<(usize, FieldElement)> = Vec::new(); let boundary_zerofiers_inverse_evaluations: Vec>> = @@ -208,6 +212,27 @@ where }) .collect::>>>(); + let zerofier_data = air.transition_zerofier_evaluations_grouped(domain); + + // GPU composition path: fuse H(row) = z_inv·Σβᵢ·Cᵢ + boundary on-device + // (no CPU trace read, no per-constraint matrix). Falls through to the CPU + // path below when the GPU LDE is absent, the field is not Goldilocks, the + // zerofier is non-uniform, or the transition offsets are non-contiguous. + #[cfg(feature = "cuda")] + { + if let Some(h) = self.try_evaluate_composition_gpu( + air, + lde_trace, + rap_challenges, + transition_coefficients, + boundary_coefficients, + &zerofier_data, + &boundary_zerofiers_inverse_evaluations, + ) { + return h; + } + } + // Fused boundary evaluation: compute (trace[col] - value) on-the-fly // instead of pre-computing all boundary_polys_evaluations. // This eliminates N_constraints × LDE_size intermediate allocations. @@ -237,8 +262,6 @@ where }) .collect(); - let zerofier_data = air.transition_zerofier_evaluations_grouped(domain); - // Iterate over all LDE domain and compute the part of the composition polynomial // related to the transition constraints and add it to the already computed part of the // boundary constraints. @@ -258,4 +281,110 @@ where &self.logup_table_offset, ) } + + /// GPU composition path: produce `H(row)` on-device (transition + boundary + /// fused, no CPU trace read, no per-constraint matrix), returning `None` to + /// fall back to the CPU path when the GPU LDE is absent, the tower is not + /// Goldilocks/degree-3, the zerofier is non-uniform (end-exemptions), or the + /// transition offsets are non-contiguous. The result feeds the existing + /// decompose + composition commit exactly like the CPU `H` vector. + #[cfg(feature = "cuda")] + #[allow(clippy::too_many_arguments)] + fn try_evaluate_composition_gpu( + &self, + air: &dyn AIR, + lde_trace: &LDETraceTable, + rap_challenges: &[FieldElement], + transition_coefficients: &[FieldElement], + boundary_coefficients: &[FieldElement], + zerofier_data: &ZerofierEvaluations, + boundary_z_inv: &[Vec>], + ) -> Option>> + where + Field: 'static, + FieldExtension: 'static, + { + use math::field::extensions_goldilocks::Degree3GoldilocksExtensionField; + use math::field::goldilocks::GoldilocksField; + use std::any::TypeId; + + if TypeId::of::() != TypeId::of::() + || TypeId::of::() != TypeId::of::() + { + return None; + } + if crate::gpu_lde::gpu_composition_disabled() { + return None; + } + if !zerofier_data.is_uniform() { + return None; + } + // The kernel's row math assumes Var offsets index a contiguous [0..n) + // frame (offset·step). The VM uses [0, 1]; anything else → CPU. + let offsets = &air.context().transition_offsets; + if !offsets.iter().enumerate().all(|(i, &o)| o == i) { + return None; + } + let main = lde_trace.gpu_main()?; + let aux = lde_trace.gpu_aux()?; + + let prog = air.constraint_program(); + + // LogUp alpha powers, exactly as `evaluate_transitions` derives them. + 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() + }; + + // Boundary spec (aligned with `boundary_coefficients` / `boundary_z_inv`). + let bcs = &self.boundary_constraints.constraints; + let b_col: Vec = bcs.iter().map(|c| c.col).collect(); + let b_is_aux: Vec = bcs.iter().map(|c| c.is_aux).collect(); + let b_value: Vec> = + bcs.iter().map(|c| c.value.clone()).collect(); + let b_z_inv_flat: Vec> = + boundary_z_inv.iter().flatten().cloned().collect(); + + let inputs = crate::constraint_ir::gpu_interp::CompositionInputs { + beta_trans: transition_coefficients, + z_inv: &zerofier_data.groups[0], + b_col: &b_col, + b_is_aux: &b_is_aux, + b_value: &b_value, + b_beta: boundary_coefficients, + b_z_inv: &b_z_inv_flat, + }; + + let next_step = lde_trace.lde_step_size; // == blowup_factor (single-row steps) + let num_rows = lde_trace.num_rows(); + + let raw = crate::constraint_ir::gpu_interp::try_eval_composition_gpu( + prog, + main, + aux, + rap_challenges, + &logup_alpha_powers, + &self.logup_table_offset, + next_step, + num_rows, + &inputs, + )?; + + // SAFETY: the TypeId gate established `FieldExtension == + // Degree3GoldilocksExtensionField`, which is `#[repr(transparent)]` over + // `[u64; 3]`; `raw.len() == num_rows * 3`. + let h: Vec> = unsafe { + std::slice::from_raw_parts( + raw.as_ptr() as *const FieldElement, + raw.len() / 3, + ) + } + .to_vec(); + Some(h) + } } diff --git a/crypto/stark/src/gpu_lde.rs b/crypto/stark/src/gpu_lde.rs index 6a81162a5..c812a3c1e 100644 --- a/crypto/stark/src/gpu_lde.rs +++ b/crypto/stark/src/gpu_lde.rs @@ -77,6 +77,7 @@ pub fn reset_all_gpu_call_counters() { GPU_FRI_CALLS.store(0, Ordering::Relaxed); GPU_BATCH_INVERT_CALLS.store(0, Ordering::Relaxed); GPU_LOGUP_CALLS.store(0, Ordering::Relaxed); + GPU_COMPOSITION_CALLS.store(0, Ordering::Relaxed); } pub(crate) static GPU_EXTEND_HALVES_CALLS: AtomicU64 = AtomicU64::new(0); @@ -92,6 +93,36 @@ pub fn gpu_logup_calls() -> u64 { GPU_LOGUP_CALLS.load(Ordering::Relaxed) } +/// Successful GPU composition-poly (`H(row)`) dispatches — one per table whose +/// round-2 constraint evaluation took the fused on-device path (a failed attempt +/// or a gate miss falls back to the CPU accumulation and is not counted). +pub(crate) static GPU_COMPOSITION_CALLS: AtomicU64 = AtomicU64::new(0); +pub fn gpu_composition_calls() -> u64 { + GPU_COMPOSITION_CALLS.load(Ordering::Relaxed) +} + +/// Runtime override to force the GPU composition path off (→ CPU accumulation). +/// An escape hatch, and the A/B toggle for benchmarking the path against the CPU +/// baseline in one process (no rebuild). Default off (path enabled). +static GPU_COMPOSITION_DISABLED: std::sync::atomic::AtomicBool = + std::sync::atomic::AtomicBool::new(false); +pub fn set_gpu_composition_disabled(v: bool) { + GPU_COMPOSITION_DISABLED.store(v, Ordering::Relaxed); +} +pub(crate) fn gpu_composition_disabled() -> bool { + if GPU_COMPOSITION_DISABLED.load(Ordering::Relaxed) { + return true; + } + // Env fallback (cached), so an unmodified prove binary can A/B the path: + // `LAMBDA_VM_DISABLE_GPU_COMPOSITION=1`. + static ENV_DISABLED: OnceLock = OnceLock::new(); + *ENV_DISABLED.get_or_init(|| { + std::env::var("LAMBDA_VM_DISABLE_GPU_COMPOSITION") + .map(|v| v == "1") + .unwrap_or(false) + }) +} + // ============================================================================ // Shared dispatch helpers // ============================================================================ diff --git a/crypto/stark/tests/gpu_constraint_interp.rs b/crypto/stark/tests/gpu_constraint_interp.rs new file mode 100644 index 000000000..bc6512f84 --- /dev/null +++ b/crypto/stark/tests/gpu_constraint_interp.rs @@ -0,0 +1,512 @@ +//! GPU↔CPU parity for the transition-constraint interpreter kernel +//! (`crypto/math-cuda/kernels/constraint_interp.cu`). +//! +//! The kernel — driven through `gpu_interp::try_eval_program_gpu` — must produce +//! the per-constraint eval matrix bit-for-bit identical to the CPU reference +//! oracle [`eval_device_program`] (the flat-blob forward walk in +//! `constraint_ir::device`). That oracle is itself pinned bit-for-bit to the +//! production folder across all 26 tables by +//! `lambda_vm_prover::tests::constraint_program_device_tests`, so GPU == oracle +//! closes the chain GPU == compiled prover folder without needing a GPU there. +//! +//! Layouts (must match the kernel + the host wrapper): +//! * base LDE column-major `buf[col * lde_size + row]` (`GpuLdeBase`) +//! * ext3 LDE de-interleaved `buf[(col*3 + k) * lde_size + row]` (`GpuLdeExt3`) +//! * `lde_size` is the row stride; here `lde_size == num_rows`, `next_step = 1`. +//! +//! CRITICAL — same-cell reads: the kernel resolves an `Op::Var{offset}` leaf at +//! LDE row `(row + offset * next_step) mod num_rows`. So for kernel row `r` the +//! oracle is fed `main[o][c] = base_lde[c][(r + o) mod num_rows]` and +//! `aux[o][c] = aux_lde[c][(r + o) mod num_rows]` — reproducing the real +//! next-row wrap exactly. +//! +//! Requires the `cuda` feature and a visible GPU. + +#![cfg(feature = "cuda")] + +use std::sync::Arc; + +use math::field::element::FieldElement; +use math::field::extensions_goldilocks::Degree3GoldilocksExtensionField as Ext; +use math::field::goldilocks::GoldilocksField as Gl; + +use math_cuda::device::backend; +use math_cuda::lde::{GpuLdeBase, GpuLdeExt3}; + +use stark::constraint_ir::device::{ + DeviceProgram, OP_ALPHA_POW, OP_RAP_CHALLENGE, OP_VAR, eval_device_program, unpack_var, +}; +use stark::constraint_ir::{ConstraintProgram, IrBuilder}; + +type Fp = FieldElement; +type Fp3 = FieldElement; + +/// Deterministic SplitMix64 (no `rand` needed; matches the device.rs oracle test). +struct SplitMix64(u64); +impl SplitMix64 { + fn next_u64(&mut self) -> u64 { + self.0 = self.0.wrapping_add(0x9E37_79B9_7F4A_7C15); + let mut z = self.0; + z = (z ^ (z >> 30)).wrapping_mul(0xBF58_476D_1CE4_E5B9); + z = (z ^ (z >> 27)).wrapping_mul(0x94D0_49BB_1331_11EB); + z ^ (z >> 31) + } + fn fp3(&mut self) -> Fp3 { + // Raw (possibly non-canonical) limbs — the stronger test, and exactly + // what a real LDE column carries. + Fp3::from_raw([ + Fp::from_raw(self.next_u64()), + Fp::from_raw(self.next_u64()), + Fp::from_raw(self.next_u64()), + ]) + } +} + +fn fp(v: u64) -> Fp { + Fp::from(v) +} +fn ext3(a: u64, b: u64, c: u64) -> Fp3 { + Fp3::from_raw([fp(a), fp(b), fp(c)]) +} + +/// Extension element → raw `[u64; 3]` limbs (the device representation). +fn enc(x: &Fp3) -> [u64; 3] { + let l = x.value(); + [*l[0].value(), *l[1].value(), *l[2].value()] +} + +/// The all-ops synthetic program (mirrors `device.rs`'s own oracle test): every +/// `Op` variant, both dims, a base-rooted constraint plus two ext (LogUp-shaped) +/// roots, next-row reads, and mixed base×ext arithmetic. +fn all_ops_program() -> ConstraintProgram { + let mut b = IrBuilder::::new(); + + // Root 0 (base): (m0 + m1) * 2 - m0_next, all base, incl. next-row. + let m0 = b.main(0, 0); + let m1 = b.main(0, 1); + let m0n = b.main(1, 0); + let two = b.const_base(2); + let sum = b.add(m0, m1); + let scaled = b.mul(sum, two); + let base_root = b.sub(scaled, m0n); + b.emit(0, base_root); + + // Root 1 (ext): m0 * challenge(0) + alpha_pow(1) * aux(0,0) - table_offset + let ch = b.challenge(0); + let ap = b.alpha_power(1); + let au = b.aux(0, 0); + let off = b.table_offset(); + let t1 = b.mul(m0, ch); // base × ext → ext (auto-embed) + let t2 = b.mul(ap, au); // ext × ext + let s = b.add(t1, t2); + let ext_root = b.sub(s, off); + b.emit(1, ext_root); + + // Root 2 (ext): embed(m1) + (-aux(0,1)) + const_ext + let em = b.embed(m1); + let au1 = b.aux(0, 1); + let nau1 = b.neg(au1); // ext negation + let ce = b.const_ext(ext3(9, 8, 7)); + let s2 = b.add(em, nau1); + let ext_root2 = b.add(s2, ce); + b.emit(2, ext_root2); + + b.finish(1) // 1 base root, 2 ext roots +} + +/// Derive the trace/uniform footprint the program actually touches, so the +/// harness works for any program (synthetic or real): #main cols, #aux cols, +/// #rap challenges, #alpha powers, and the max frame offset. +fn program_footprint(dev: &DeviceProgram) -> (usize, usize, usize, usize, usize) { + let (mut main_cols, mut aux_cols, mut rap_len, mut alpha_len, mut max_off) = (0, 0, 0, 0, 0); + for n in &dev.nodes { + match n.op { + OP_VAR => { + let (is_main, offset, _row, col) = unpack_var(n.a, n.b); + let col = col as usize + 1; + if is_main { + main_cols = main_cols.max(col); + } else { + aux_cols = aux_cols.max(col); + } + max_off = max_off.max(offset as usize); + } + OP_RAP_CHALLENGE => rap_len = rap_len.max(n.a as usize + 1), + OP_ALPHA_POW => alpha_len = alpha_len.max(n.a as usize + 1), + _ => {} + } + } + (main_cols, aux_cols, rap_len, alpha_len, max_off) +} + +/// Full GPU↔CPU differential for one program over an 8-row random LDE. +fn check_program(prog: &ConstraintProgram, label: &str, seed: u64) { + const NUM_ROWS: usize = 8; + const NEXT_STEP: usize = 1; + let lde_size = NUM_ROWS; + + let dev = DeviceProgram::lower(prog); + let (main_cols, aux_cols, rap_len, alpha_len, max_off) = program_footprint(&dev); + let n_off = max_off + 1; + assert!( + n_off <= NUM_ROWS, + "[{label}] program frame span {n_off} exceeds NUM_ROWS {NUM_ROWS}" + ); + let n = dev.roots.len(); + let num_base = dev.num_base as usize; + + let mut rng = SplitMix64(seed); + + // Host-side random LDE, kept column-major so we can both upload it and feed + // the oracle the exact same cells. + let base_host: Vec> = (0..main_cols) + .map(|_| (0..NUM_ROWS).map(|_| rng.next_u64()).collect()) + .collect(); + let aux_host: Vec> = (0..aux_cols) + .map(|_| (0..NUM_ROWS).map(|_| enc(&rng.fp3())).collect()) + .collect(); + + // Per-proof uniforms. + let rap: Vec = (0..rap_len.max(1)).map(|_| rng.fp3()).collect(); + let alpha: Vec = (0..alpha_len.max(1)).map(|_| rng.fp3()).collect(); + let offset = rng.fp3(); + + // Pack the LDE into the device buffer layouts and upload. + let mut base_flat = vec![0u64; main_cols * lde_size]; + for (c, col) in base_host.iter().enumerate() { + for (r, v) in col.iter().enumerate() { + base_flat[c * lde_size + r] = *v; + } + } + let mut aux_flat = vec![0u64; aux_cols * 3 * lde_size]; + for (c, col) in aux_host.iter().enumerate() { + for (r, v) in col.iter().enumerate() { + aux_flat[(c * 3) * lde_size + r] = v[0]; + aux_flat[(c * 3 + 1) * lde_size + r] = v[1]; + aux_flat[(c * 3 + 2) * lde_size + r] = v[2]; + } + } + + let be = backend().expect("cuda backend"); + let stream = be.next_stream(); + let base_dev = stream.clone_htod(&base_flat).expect("upload base LDE"); + let aux_dev = stream.clone_htod(&aux_flat).expect("upload aux LDE"); + stream.synchronize().expect("sync uploads"); + + let main = GpuLdeBase { + buf: Arc::new(base_dev), + m: main_cols, + lde_size, + tree: None, + trace_dev: None, + trace_rows: 0, + }; + let aux = GpuLdeExt3 { + buf: Arc::new(aux_dev), + m: aux_cols, + lde_size, + tree: None, + }; + + // GPU: launch the interpreter over every row. + let gpu = stark::constraint_ir::gpu_interp::try_eval_program_gpu( + prog, &main, &aux, &rap, &alpha, &offset, NEXT_STEP, NUM_ROWS, + ) + .unwrap_or_else(|| panic!("[{label}] GPU path (Goldilocks ext3) must engage")); + assert_eq!(gpu.len(), n * NUM_ROWS * 3, "[{label}] eval matrix shape"); + + // CPU oracle, row by row, reading the SAME wrapped LDE cells. + let rap_raw: Vec<[u64; 3]> = rap.iter().map(enc).collect(); + let alpha_raw: Vec<[u64; 3]> = alpha.iter().map(enc).collect(); + let off_raw = enc(&offset); + + for r in 0..NUM_ROWS { + let main_raw: Vec> = (0..n_off) + .map(|o| { + (0..main_cols) + .map(|c| base_host[c][(r + o) % NUM_ROWS]) + .collect() + }) + .collect(); + let aux_raw: Vec> = (0..n_off) + .map(|o| { + (0..aux_cols) + .map(|c| aux_host[c][(r + o) % NUM_ROWS]) + .collect() + }) + .collect(); + + let mut base_o = vec![0u64; num_base]; + let mut ext_o = vec![[0u64; 3]; n]; + eval_device_program( + &dev, + &main_raw, + &aux_raw, + &rap_raw, + &alpha_raw, + off_raw, + &mut base_o, + &mut ext_o, + ); + + for c in 0..n { + let g = |k: usize| gpu[(c * NUM_ROWS + r) * 3 + k]; + if c < num_base { + assert_eq!( + g(0), + base_o[c], + "[{label}] base constraint {c}, row {r}: GPU {} vs CPU {}", + g(0), + base_o[c] + ); + // A base-rooted constraint carries its value in component 0; + // the embedding pads components 1 and 2 with zero. + assert_eq!( + g(1), + 0, + "[{label}] base constraint {c}, row {r}: comp1 != 0" + ); + assert_eq!( + g(2), + 0, + "[{label}] base constraint {c}, row {r}: comp2 != 0" + ); + } else { + let got = [g(0), g(1), g(2)]; + assert_eq!( + got, ext_o[c], + "[{label}] ext constraint {c}, row {r}: GPU {got:?} vs CPU {:?}", + ext_o[c] + ); + } + } + } +} + +#[test] +fn gpu_matches_cpu_oracle_all_ops() { + // A few seeds to exercise the reduce/overflow paths with different limbs. + for seed in [0x0123_4567_89AB_CDEF, 0xDEAD_BEEF_CAFE_F00D, 1, 42] { + check_program(&all_ops_program(), "ALL_OPS", seed); + } +} + +// ------------------------------------------------------------------------ +// Fused composition-poly kernel (`constraint_composition_kernel`): the GPU +// H(row) must match the CPU accumulation of `constraints::evaluator` applied +// to the same per-constraint evals — z_inv·Σβᵢ·Cᵢ + Σ_b z_b_inv·β_b·(trace−val). +// ------------------------------------------------------------------------ + +use stark::constraint_ir::gpu_interp::{CompositionInputs, try_eval_composition_gpu}; + +/// CPU reference for one row: mirror `evaluator.rs` (uniform case) exactly, +/// consuming `eval_device_program`'s per-constraint evals. +#[allow(clippy::too_many_arguments)] +fn composition_oracle_row( + base_evals: &[u64], + ext_evals: &[[u64; 3]], + num_base: usize, + beta_trans: &[Fp3], + z_inv_row: Fp, + b_terms: &[(bool, usize, Fp3, Fp3, Fp)], // (is_aux, col, value, beta, z_inv_row) + base_row: &[u64], + aux_row: &[[u64; 3]], +) -> Fp3 { + let mut sum = Fp3::zero(); + for (c, beta) in beta_trans.iter().enumerate() { + // eval * beta, base×ext for base constraints (matches evaluator.rs:89/92). + if c < num_base { + sum += Fp::from_raw(base_evals[c]) * *beta; + } else { + let e = ext_evals[c]; + sum += + Fp3::from_raw([Fp::from_raw(e[0]), Fp::from_raw(e[1]), Fp::from_raw(e[2])]) * *beta; + } + } + let mut h = z_inv_row * sum; // z * sum (base×ext) + for &(is_aux, col, value, beta, zinv) in b_terms { + let tcell = if is_aux { + let a = aux_row[col]; + Fp3::from_raw([Fp::from_raw(a[0]), Fp::from_raw(a[1]), Fp::from_raw(a[2])]) + } else { + Fp::from_raw(base_row[col]).to_extension::() + }; + let bp = tcell - value; + h += zinv * beta * bp; // (base×ext)×ext, matches evaluator.rs:234 + } + h +} + +fn check_composition(prog: &ConstraintProgram, label: &str, seed: u64) { + const NUM_ROWS: usize = 8; + const NEXT_STEP: usize = 1; + const Z_LEN: usize = 2; // blowup-length cyclic transition zerofier inverse + let lde_size = NUM_ROWS; + + let dev = DeviceProgram::lower(prog); + let (main_cols, aux_cols, rap_len, alpha_len, max_off) = program_footprint(&dev); + assert!(max_off < NUM_ROWS); + let n = dev.roots.len(); + let num_base = dev.num_base as usize; + + let mut rng = SplitMix64(seed); + + let base_host: Vec> = (0..main_cols) + .map(|_| (0..NUM_ROWS).map(|_| rng.next_u64()).collect()) + .collect(); + let aux_host: Vec> = (0..aux_cols) + .map(|_| (0..NUM_ROWS).map(|_| enc(&rng.fp3())).collect()) + .collect(); + let rap: Vec = (0..rap_len.max(1)).map(|_| rng.fp3()).collect(); + let alpha: Vec = (0..alpha_len.max(1)).map(|_| rng.fp3()).collect(); + let offset = rng.fp3(); + + // Accumulation inputs (synthetic but shaped exactly like the real ones). + let beta_trans: Vec = (0..n).map(|_| rng.fp3()).collect(); + let z_inv: Vec = (0..Z_LEN).map(|_| fp(rng.next_u64())).collect(); + + // Two boundary constraints: one main (col 0), one aux (last aux col). + let b_defs: Vec<(bool, usize)> = { + let mut v = Vec::new(); + if main_cols > 0 { + v.push((false, 0)); + } + if aux_cols > 0 { + v.push((true, aux_cols - 1)); + } + v + }; + let num_boundary = b_defs.len(); + let b_col: Vec = b_defs.iter().map(|&(_, c)| c).collect(); + let b_is_aux: Vec = b_defs.iter().map(|&(a, _)| a).collect(); + let b_value: Vec = (0..num_boundary).map(|_| rng.fp3()).collect(); + let b_beta: Vec = (0..num_boundary).map(|_| rng.fp3()).collect(); + // b_z_inv laid out b*num_rows + row. + let b_z_inv: Vec = (0..num_boundary * NUM_ROWS) + .map(|_| fp(rng.next_u64())) + .collect(); + + // Upload the LDE and build handles. + let mut base_flat = vec![0u64; main_cols.max(1) * lde_size]; + for (c, col) in base_host.iter().enumerate() { + for (r, v) in col.iter().enumerate() { + base_flat[c * lde_size + r] = *v; + } + } + let mut aux_flat = vec![0u64; aux_cols.max(1) * 3 * lde_size]; + for (c, col) in aux_host.iter().enumerate() { + for (r, v) in col.iter().enumerate() { + aux_flat[(c * 3) * lde_size + r] = v[0]; + aux_flat[(c * 3 + 1) * lde_size + r] = v[1]; + aux_flat[(c * 3 + 2) * lde_size + r] = v[2]; + } + } + let be = backend().expect("cuda backend"); + let stream = be.next_stream(); + let base_dev = stream.clone_htod(&base_flat).expect("upload base"); + let aux_dev = stream.clone_htod(&aux_flat).expect("upload aux"); + stream.synchronize().expect("sync"); + let main = GpuLdeBase { + buf: Arc::new(base_dev), + m: main_cols, + lde_size, + tree: None, + trace_dev: None, + trace_rows: 0, + }; + let aux = GpuLdeExt3 { + buf: Arc::new(aux_dev), + m: aux_cols, + lde_size, + tree: None, + }; + + let inputs = CompositionInputs { + beta_trans: &beta_trans, + z_inv: &z_inv, + b_col: &b_col, + b_is_aux: &b_is_aux, + b_value: &b_value, + b_beta: &b_beta, + b_z_inv: &b_z_inv, + }; + let gpu = try_eval_composition_gpu( + prog, &main, &aux, &rap, &alpha, &offset, NEXT_STEP, NUM_ROWS, &inputs, + ) + .unwrap_or_else(|| panic!("[{label}] GPU composition path must engage")); + assert_eq!(gpu.len(), NUM_ROWS * 3, "[{label}] H shape"); + + // CPU oracle, row by row. + let rap_raw: Vec<[u64; 3]> = rap.iter().map(enc).collect(); + let alpha_raw: Vec<[u64; 3]> = alpha.iter().map(enc).collect(); + let off_raw = enc(&offset); + let n_off = max_off + 1; + + for r in 0..NUM_ROWS { + let main_raw: Vec> = (0..n_off) + .map(|o| { + (0..main_cols) + .map(|c| base_host[c][(r + o) % NUM_ROWS]) + .collect() + }) + .collect(); + let aux_raw: Vec> = (0..n_off) + .map(|o| { + (0..aux_cols) + .map(|c| aux_host[c][(r + o) % NUM_ROWS]) + .collect() + }) + .collect(); + let mut base_o = vec![0u64; num_base]; + let mut ext_o = vec![[0u64; 3]; n]; + eval_device_program( + &dev, + &main_raw, + &aux_raw, + &rap_raw, + &alpha_raw, + off_raw, + &mut base_o, + &mut ext_o, + ); + + // Boundary terms read the current row (offset 0). + let b_terms: Vec<(bool, usize, Fp3, Fp3, Fp)> = (0..num_boundary) + .map(|b| { + ( + b_is_aux[b], + b_col[b], + b_value[b], + b_beta[b], + b_z_inv[b * NUM_ROWS + r], + ) + }) + .collect(); + + let h_cpu = composition_oracle_row( + &base_o, + &ext_o, + num_base, + &beta_trans, + z_inv[r % Z_LEN], + &b_terms, + &main_raw[0], + &aux_raw[0], + ); + + let h_gpu = [gpu[r * 3], gpu[r * 3 + 1], gpu[r * 3 + 2]]; + assert_eq!( + h_gpu, + enc(&h_cpu), + "[{label}] H mismatch row {r} seed {seed:#x}: GPU {h_gpu:?} vs CPU {:?}", + enc(&h_cpu) + ); + } +} + +#[test] +fn gpu_composition_matches_cpu_oracle_all_ops() { + for seed in [0x0123_4567_89AB_CDEF, 0xDEAD_BEEF_CAFE_F00D, 7] { + check_composition(&all_ops_program(), "ALL_OPS_COMP", seed); + } +} diff --git a/prover/Cargo.toml b/prover/Cargo.toml index 3695689d6..57e8114d0 100644 --- a/prover/Cargo.toml +++ b/prover/Cargo.toml @@ -37,6 +37,11 @@ tiny-keccak = { version = "2.0", features = ["keccak"] } # Enable stark's test-utils so cross-crate tests can reach # `compute_precomputed_commitment_for_testing`. Only active under cargo test/bench. stark = { path = "../crypto/stark", features = ["test-utils"] } +# Device-resident LDE handles for the cuda-gated GPU constraint-interp parity +# test (`tests/gpu_constraint_interp_real.rs`). Its build.rs stubs empty PTX when +# nvcc is absent, so this compiles on CPU-only hosts too; the test itself is +# `#[cfg(feature = "cuda")]` and only runs with a GPU. +math-cuda = { path = "../crypto/math-cuda" } [[bench]] name = "vm_prover_benchmark" diff --git a/prover/src/tests/constraint_program_device_tests.rs b/prover/src/tests/constraint_program_device_tests.rs new file mode 100644 index 000000000..9062df681 --- /dev/null +++ b/prover/src/tests/constraint_program_device_tests.rs @@ -0,0 +1,185 @@ +//! Device-lowering differential tests: every production table's captured +//! constraint program is lowered to its flat GPU form +//! ([`stark::constraint_ir::DeviceProgram`]) and the CPU walk over that flat +//! blob ([`eval_device_program`]) is compared bit-for-bit against the compiled +//! prover folder on random off-trace frames. +//! +//! This is the pre-GPU parity oracle: the CUDA kernel is a transliteration of +//! [`eval_device_program`], so a bit-for-bit match here pins the on-device +//! layout and control flow (op tags, `Var` packing, const side-tables, the +//! base/ext output split) against the production path — across every real +//! bus-interaction layout, `Multiplicity` variant, and LogUp suffix — *before* +//! any CUDA exists. It complements the synthetic all-ops coverage in +//! `stark::constraint_ir::device`'s own unit test. +//! +//! The folder (`compute_transition_prover`) is the oracle: it is the production +//! prove path, independently pinned by the prove→verify suites and +//! cross-version verification. + +use math::field::element::FieldElement; +use stark::constraint_ir::{DeviceProgram, eval_device_program}; +use stark::frame::Frame; +use stark::proof::options::GoldilocksCubicProofOptions; +use stark::table::TableView; +use stark::traits::{AIR, TransitionEvaluationContext}; + +use crate::tables::types::{GoldilocksExtension, GoldilocksField}; +use crate::test_utils::*; + +type Gl = GoldilocksField; +type Ext3 = GoldilocksExtension; +type Fp = FieldElement; +type Fp3 = FieldElement; + +const TRIALS: usize = 200; + +/// Deterministic SplitMix64 (no `rand` dependency). +struct SplitMix64(u64); +impl SplitMix64 { + fn next_u64(&mut self) -> u64 { + self.0 = self.0.wrapping_add(0x9E37_79B9_7F4A_7C15); + let mut z = self.0; + z = (z ^ (z >> 30)).wrapping_mul(0xBF58_476D_1CE4_E5B9); + z = (z ^ (z >> 27)).wrapping_mul(0x94D0_49BB_1331_11EB); + z ^ (z >> 31) + } + fn fp3(&mut self) -> Fp3 { + Fp3::new([ + Fp::from(self.next_u64()), + Fp::from(self.next_u64()), + Fp::from(self.next_u64()), + ]) + } +} + +/// Extension element → raw `[u64; 3]` limbs (the device representation). +fn enc(x: &Fp3) -> [u64; 3] { + let limbs = x.value(); + [*limbs[0].value(), *limbs[1].value(), *limbs[2].value()] +} + +/// Device differential for one production AIR: lower the captured program once, +/// then assert on random two-step prover frames that the flat-blob walk matches +/// the compiled folder, base and extension constraints alike. +fn check_air_device( + air: &dyn AIR, + label: &str, +) { + let n = air.context().num_transition_constraints; + let num_base = air.num_base_transition_constraints(); + let (n_main, n_aux) = air.trace_layout(); + + let prog = air.constraint_program(); + let dev = DeviceProgram::lower(prog); + assert_eq!(dev.roots.len(), n, "[{label}] one root per constraint"); + assert_eq!( + dev.num_base as usize, num_base, + "[{label}] num_base preserved" + ); + + let mut rng = SplitMix64(0xDECA_FBAD ^ label.len() as u64); + for trial in 0..TRIALS { + let mk_step = |rng: &mut SplitMix64| { + let main: Vec = (0..n_main).map(|_| Fp::from(rng.next_u64())).collect(); + let aux: Vec = (0..n_aux).map(|_| rng.fp3()).collect(); + TableView::new(vec![main], vec![aux]) + }; + let frame = Frame::::new(vec![mk_step(&mut rng), mk_step(&mut rng)]); + let challenges = vec![rng.fp3(), rng.fp3()]; // [z, alpha] + let alphas: Vec = (0..air.max_bus_elements() + 2).map(|_| rng.fp3()).collect(); + let offset = rng.fp3(); + + // Raw-limb inputs for the device walk, extracted from the frame. + let main_raw: Vec> = (0..2) + .map(|off| { + let step = frame.get_evaluation_step(off); + (0..n_main) + .map(|c| *step.get_main_evaluation_element(0, c).value()) + .collect() + }) + .collect(); + let aux_raw: Vec> = (0..2) + .map(|off| { + let step = frame.get_evaluation_step(off); + (0..n_aux) + .map(|c| enc(step.get_aux_evaluation_element(0, c))) + .collect() + }) + .collect(); + let rap_raw: Vec<[u64; 3]> = challenges.iter().map(enc).collect(); + let alpha_raw: Vec<[u64; 3]> = alphas.iter().map(enc).collect(); + let off_raw = enc(&offset); + + // Oracle: the compiled prover folder. + let ctx = TransitionEvaluationContext::new_prover( + frame.as_row_frame(), + &challenges, + &alphas, + &offset, + ); + let mut f_base = vec![Fp::zero(); num_base]; + let mut f_ext = vec![Fp3::zero(); n]; + air.compute_transition_prover(&ctx, &mut f_base, &mut f_ext); + + // Device walk over the flat blob. + let mut d_base = vec![0u64; num_base]; + let mut d_ext = vec![[0u64; 3]; n]; + eval_device_program( + &dev, + &main_raw, + &aux_raw, + &rap_raw, + &alpha_raw, + off_raw, + &mut d_base, + &mut d_ext, + ); + + for c in 0..num_base { + assert_eq!( + d_base[c], + *f_base[c].value(), + "[{label}] folder vs device, base constraint {c}, trial {trial}" + ); + } + for c in num_base..n { + assert_eq!( + d_ext[c], + enc(&f_ext[c]), + "[{label}] folder vs device, ext constraint {c}, trial {trial}" + ); + } + } +} + +#[test] +fn all_table_programs_lower_and_match_folders() { + let opts = GoldilocksCubicProofOptions::with_blowup(2).expect("blowup=2 valid"); + + check_air_device(&create_cpu_air(&opts), "CPU"); + check_air_device(&create_bitwise_air(&opts), "BITWISE"); + check_air_device(&create_lt_air(&opts), "LT"); + check_air_device(&create_shift_air(&opts), "SHIFT"); + check_air_device(&create_eq_air(&opts), "EQ"); + check_air_device(&create_bytewise_air(&opts), "BYTEWISE"); + check_air_device(&create_store_air(&opts), "STORE"); + check_air_device(&create_cpu32_air(&opts), "CPU32"); + check_air_device(&create_memw_air(&opts), "MEMW"); + check_air_device(&create_memw_aligned_air(&opts), "MEMW_A"); + check_air_device(&create_memw_register_air(&opts), "MEMW_R"); + check_air_device(&create_load_air(&opts), "LOAD"); + check_air_device(&create_decode_air(&opts), "DECODE"); + check_air_device(&create_mul_air(&opts), "MUL"); + check_air_device(&create_dvrm_air(&opts), "DVRM"); + check_air_device(&create_branch_air(&opts), "BRANCH"); + check_air_device(&create_halt_air(&opts), "HALT"); + check_air_device(&create_commit_air(&opts), "COMMIT"); + check_air_device(&create_page_air(&opts, 0x1000), "PAGE"); + check_air_device(&create_register_air(&opts), "REGISTER"); + check_air_device(&create_keccak_air(&opts), "KECCAK"); + check_air_device(&create_keccak_rnd_air(&opts), "KECCAK_RND"); + check_air_device(&create_keccak_rc_air(&opts), "KECCAK_RC"); + check_air_device(&create_ecsm_air(&opts), "ECSM"); + check_air_device(&create_ec_scalar_air(&opts), "EC_SCALAR"); + check_air_device(&create_ecdas_air(&opts), "ECDAS"); +} diff --git a/prover/src/tests/mod.rs b/prover/src/tests/mod.rs index 2c06d7fcf..f7cc6b77c 100644 --- a/prover/src/tests/mod.rs +++ b/prover/src/tests/mod.rs @@ -17,6 +17,8 @@ pub mod compute_commit_bus_offset_tests; #[cfg(test)] pub mod constraint_emit_tests; #[cfg(test)] +pub mod constraint_program_device_tests; +#[cfg(test)] pub mod constraint_program_tests; #[cfg(test)] pub mod constraint_set_tests_a; diff --git a/prover/tests/cuda_path_integration.rs b/prover/tests/cuda_path_integration.rs index e0a587b88..7081baadf 100644 --- a/prover/tests/cuda_path_integration.rs +++ b/prover/tests/cuda_path_integration.rs @@ -11,11 +11,30 @@ use lambda_vm_prover::test_utils::asm_elf_bytes; use lambda_vm_prover::{prove, verify}; use stark::gpu_lde::{ - gpu_bary_calls, gpu_batch_invert_calls, gpu_comp_poly_tree_calls, gpu_deep_calls, - gpu_extend_halves_calls, gpu_fri_calls, gpu_lde_calls, gpu_logup_calls, gpu_parts_lde_calls, - reset_all_gpu_call_counters, + gpu_bary_calls, gpu_batch_invert_calls, gpu_comp_poly_tree_calls, gpu_composition_calls, + gpu_deep_calls, gpu_extend_halves_calls, gpu_fri_calls, gpu_lde_calls, gpu_logup_calls, + gpu_parts_lde_calls, reset_all_gpu_call_counters, }; +/// The R2 GPU composition-poly path (fused `H = z·Σβᵢ·Cᵢ + boundary`) fires and +/// yields a verifying proof. Guards both a silent CPU fallback (counter == 0) +/// and a bad-`H` regression (fires but the proof fails verification). +#[test] +#[ignore = "requires GPU; run with --ignored --nocapture"] +fn gpu_composition_path_fires_and_verifies() { + let elf = asm_elf_bytes("fib_iterative_1M"); + reset_all_gpu_call_counters(); + let proof = prove(&elf).expect("prove"); + assert!( + gpu_composition_calls() > 0, + "GPU composition path did not fire (tables below threshold or gate fell back to CPU)" + ); + assert!( + verify(&proof, &elf).expect("verify"), + "GPU-produced proof (fused composition) failed verification" + ); +} + /// The GPU LogUp aux-build path fires and still yields a verifying proof. #[test] #[ignore = "requires GPU; run with --ignored --nocapture"] diff --git a/prover/tests/gpu_constraint_interp_real.rs b/prover/tests/gpu_constraint_interp_real.rs new file mode 100644 index 000000000..586e14e3f --- /dev/null +++ b/prover/tests/gpu_constraint_interp_real.rs @@ -0,0 +1,256 @@ +//! GPU↔CPU parity for the transition-constraint interpreter kernel on the *real* +//! production programs. +//! +//! For every production table, lower its captured `ConstraintProgram` to the +//! flat device blob, run `constraint_interp.cu` (via +//! `gpu_interp::try_eval_program_gpu`) over an 8-row random LDE, and assert the +//! per-constraint eval matrix is bit-for-bit identical to the CPU reference +//! oracle `eval_device_program`. The oracle is pinned bit-for-bit to the +//! compiled prover folder by `tests::constraint_program_device_tests`, so +//! GPU == oracle closes the chain GPU == folder across all 26 tables. +//! +//! Complements the synthetic all-ops coverage in +//! `stark`'s `tests/gpu_constraint_interp.rs`: this exercises the kernel on the +//! real node graphs (hundreds of nodes, many roots, every bus/LogUp shape). +//! +//! Requires the `cuda` feature and a visible GPU. Run with: +//! ```text +//! cargo test -p lambda-vm-prover --features cuda --test gpu_constraint_interp_real +//! ``` + +#![cfg(feature = "cuda")] + +use std::sync::Arc; + +use math::field::element::FieldElement; +use math::field::extensions_goldilocks::Degree3GoldilocksExtensionField as Ext; +use math::field::goldilocks::GoldilocksField as Gl; + +use math_cuda::device::backend; +use math_cuda::lde::{GpuLdeBase, GpuLdeExt3}; + +use stark::constraint_ir::device::{ + DeviceProgram, OP_ALPHA_POW, OP_RAP_CHALLENGE, OP_VAR, eval_device_program, unpack_var, +}; +use stark::constraint_ir::gpu_interp::try_eval_program_gpu; +use stark::proof::options::GoldilocksCubicProofOptions; +use stark::traits::AIR; + +use lambda_vm_prover::test_utils::*; + +type Fp = FieldElement; +type Fp3 = FieldElement; + +struct SplitMix64(u64); +impl SplitMix64 { + fn next_u64(&mut self) -> u64 { + self.0 = self.0.wrapping_add(0x9E37_79B9_7F4A_7C15); + let mut z = self.0; + z = (z ^ (z >> 30)).wrapping_mul(0xBF58_476D_1CE4_E5B9); + z = (z ^ (z >> 27)).wrapping_mul(0x94D0_49BB_1331_11EB); + z ^ (z >> 31) + } + fn fp3(&mut self) -> Fp3 { + Fp3::from_raw([ + Fp::from_raw(self.next_u64()), + Fp::from_raw(self.next_u64()), + Fp::from_raw(self.next_u64()), + ]) + } +} + +fn enc(x: &Fp3) -> [u64; 3] { + let l = x.value(); + [*l[0].value(), *l[1].value(), *l[2].value()] +} + +/// #main cols, #aux cols, #rap challenges, #alpha powers, max frame offset the +/// program actually references. +fn program_footprint(dev: &DeviceProgram) -> (usize, usize, usize, usize, usize) { + let (mut main_cols, mut aux_cols, mut rap_len, mut alpha_len, mut max_off) = (0, 0, 0, 0, 0); + for n in &dev.nodes { + match n.op { + OP_VAR => { + let (is_main, offset, _row, col) = unpack_var(n.a, n.b); + let col = col as usize + 1; + if is_main { + main_cols = main_cols.max(col); + } else { + aux_cols = aux_cols.max(col); + } + max_off = max_off.max(offset as usize); + } + OP_RAP_CHALLENGE => rap_len = rap_len.max(n.a as usize + 1), + OP_ALPHA_POW => alpha_len = alpha_len.max(n.a as usize + 1), + _ => {} + } + } + (main_cols, aux_cols, rap_len, alpha_len, max_off) +} + +fn check_air(air: &dyn AIR, label: &str) { + const NUM_ROWS: usize = 8; + const NEXT_STEP: usize = 1; + let lde_size = NUM_ROWS; + + let prog = air.constraint_program(); + let dev = DeviceProgram::lower(prog); + let (main_cols, aux_cols, rap_len, alpha_len, max_off) = program_footprint(&dev); + let n_off = max_off + 1; + assert!( + n_off <= NUM_ROWS, + "[{label}] frame span {n_off} exceeds NUM_ROWS {NUM_ROWS}" + ); + let n = dev.roots.len(); + let num_base = dev.num_base as usize; + assert_eq!(n, air.num_transition_constraints(), "[{label}] root count"); + assert_eq!( + num_base, + air.num_base_transition_constraints(), + "[{label}] num_base" + ); + + for seed in [0xABCD_0000u64 ^ label.len() as u64, 0x5EED_1234] { + let mut rng = SplitMix64(seed); + + let base_host: Vec> = (0..main_cols) + .map(|_| (0..NUM_ROWS).map(|_| rng.next_u64()).collect()) + .collect(); + let aux_host: Vec> = (0..aux_cols) + .map(|_| (0..NUM_ROWS).map(|_| enc(&rng.fp3())).collect()) + .collect(); + let rap: Vec = (0..rap_len.max(1)).map(|_| rng.fp3()).collect(); + let alpha: Vec = (0..alpha_len.max(1)).map(|_| rng.fp3()).collect(); + let offset = rng.fp3(); + + let mut base_flat = vec![0u64; main_cols.max(1) * lde_size]; + for (c, col) in base_host.iter().enumerate() { + for (r, v) in col.iter().enumerate() { + base_flat[c * lde_size + r] = *v; + } + } + let mut aux_flat = vec![0u64; aux_cols.max(1) * 3 * lde_size]; + for (c, col) in aux_host.iter().enumerate() { + for (r, v) in col.iter().enumerate() { + aux_flat[(c * 3) * lde_size + r] = v[0]; + aux_flat[(c * 3 + 1) * lde_size + r] = v[1]; + aux_flat[(c * 3 + 2) * lde_size + r] = v[2]; + } + } + + let be = backend().expect("cuda backend"); + let stream = be.next_stream(); + let base_dev = stream.clone_htod(&base_flat).expect("upload base LDE"); + let aux_dev = stream.clone_htod(&aux_flat).expect("upload aux LDE"); + stream.synchronize().expect("sync uploads"); + + let main = GpuLdeBase { + buf: Arc::new(base_dev), + m: main_cols, + lde_size, + tree: None, + trace_dev: None, + trace_rows: 0, + }; + let aux = GpuLdeExt3 { + buf: Arc::new(aux_dev), + m: aux_cols, + lde_size, + tree: None, + }; + + let gpu = try_eval_program_gpu( + prog, &main, &aux, &rap, &alpha, &offset, NEXT_STEP, NUM_ROWS, + ) + .unwrap_or_else(|| panic!("[{label}] GPU path must engage")); + assert_eq!(gpu.len(), n * NUM_ROWS * 3, "[{label}] matrix shape"); + + let rap_raw: Vec<[u64; 3]> = rap.iter().map(enc).collect(); + let alpha_raw: Vec<[u64; 3]> = alpha.iter().map(enc).collect(); + let off_raw = enc(&offset); + + for r in 0..NUM_ROWS { + let main_raw: Vec> = (0..n_off) + .map(|o| { + (0..main_cols) + .map(|c| base_host[c][(r + o) % NUM_ROWS]) + .collect() + }) + .collect(); + let aux_raw: Vec> = (0..n_off) + .map(|o| { + (0..aux_cols) + .map(|c| aux_host[c][(r + o) % NUM_ROWS]) + .collect() + }) + .collect(); + + let mut base_o = vec![0u64; num_base]; + let mut ext_o = vec![[0u64; 3]; n]; + eval_device_program( + &dev, + &main_raw, + &aux_raw, + &rap_raw, + &alpha_raw, + off_raw, + &mut base_o, + &mut ext_o, + ); + + for c in 0..n { + let g = |k: usize| gpu[(c * NUM_ROWS + r) * 3 + k]; + if c < num_base { + assert_eq!( + g(0), + base_o[c], + "[{label}] base c{c} r{r} seed{seed:#x}: GPU {} vs CPU {}", + g(0), + base_o[c] + ); + assert_eq!(g(1), 0, "[{label}] base c{c} r{r}: comp1 != 0"); + assert_eq!(g(2), 0, "[{label}] base c{c} r{r}: comp2 != 0"); + } else { + let got = [g(0), g(1), g(2)]; + assert_eq!( + got, ext_o[c], + "[{label}] ext c{c} r{r} seed{seed:#x}: GPU {got:?} vs CPU {:?}", + ext_o[c] + ); + } + } + } + } +} + +#[test] +fn all_table_programs_gpu_match_cpu_oracle() { + let opts = GoldilocksCubicProofOptions::with_blowup(2).expect("blowup=2 valid"); + + check_air(&create_cpu_air(&opts), "CPU"); + check_air(&create_bitwise_air(&opts), "BITWISE"); + check_air(&create_lt_air(&opts), "LT"); + check_air(&create_shift_air(&opts), "SHIFT"); + check_air(&create_eq_air(&opts), "EQ"); + check_air(&create_bytewise_air(&opts), "BYTEWISE"); + check_air(&create_store_air(&opts), "STORE"); + check_air(&create_cpu32_air(&opts), "CPU32"); + check_air(&create_memw_air(&opts), "MEMW"); + check_air(&create_memw_aligned_air(&opts), "MEMW_A"); + check_air(&create_memw_register_air(&opts), "MEMW_R"); + check_air(&create_load_air(&opts), "LOAD"); + check_air(&create_decode_air(&opts), "DECODE"); + check_air(&create_mul_air(&opts), "MUL"); + check_air(&create_dvrm_air(&opts), "DVRM"); + check_air(&create_branch_air(&opts), "BRANCH"); + check_air(&create_halt_air(&opts), "HALT"); + check_air(&create_commit_air(&opts), "COMMIT"); + check_air(&create_page_air(&opts, 0x1000), "PAGE"); + check_air(&create_register_air(&opts), "REGISTER"); + check_air(&create_keccak_air(&opts), "KECCAK"); + check_air(&create_keccak_rnd_air(&opts), "KECCAK_RND"); + check_air(&create_keccak_rc_air(&opts), "KECCAK_RC"); + check_air(&create_ecsm_air(&opts), "ECSM"); + check_air(&create_ec_scalar_air(&opts), "EC_SCALAR"); + check_air(&create_ecdas_air(&opts), "ECDAS"); +} diff --git a/scripts/bench_abba_gpu.sh b/scripts/bench_abba_gpu.sh new file mode 100755 index 000000000..b54753d9a --- /dev/null +++ b/scripts/bench_abba_gpu.sh @@ -0,0 +1,120 @@ +#!/usr/bin/env bash +# +# bench_abba_gpu.sh — interleaved A/B/B/A of the GPU composition path (ON vs OFF) +# on the ethrex workload, using ONE cuda `cli` binary and the +# LAMBDA_VM_DISABLE_GPU_COMPOSITION runtime toggle. +# +# WHY a toggle instead of two refs (as in bench_abba.sh): the change under test +# is a single feature-gated branch (round-2 constraint eval on GPU vs CPU). One +# binary flipped by an env var isolates exactly that branch with zero build/ref +# differences — no compiler/codegen drift between the two sides. +# +# A = GPU composition path ON (LAMBDA_VM_DISABLE_GPU_COMPOSITION unset) +# B = GPU composition path OFF (=1 -> CPU per-row accumulation) +# +# CONVENTION: reported % = (A - B)/B = (GPU_on - CPU)/CPU. NEGATIVE = GPU faster. +# +# USAGE: scripts/bench_abba_gpu.sh [N_PAIRS=10] [TX_COUNT=20] +# Env: REBUILD=1 force a cli rebuild +# CUDARC_PIN= pin math-cuda's cudarc to a CUDA version (rented-box +# driver may lack cudarc-latest symbols) +# BENCH_FEATURES cli features (default: jemalloc-stats,prover/cuda) + +set -euo pipefail +N_PAIRS="${1:-10}" +TX_COUNT="${2:-20}" +BENCH_FEATURES="${BENCH_FEATURES:-jemalloc-stats,prover/cuda}" + +ELF_REL="executor/program_artifacts/rust/ethrex.elf" +INPUT_REL="executor/tests/ethrex_${TX_COUNT}_transfers.bin" +WORK="/tmp/abba_gpu_run" +PROOF="/tmp/abba_gpu_proof.bin" + +ROOT="$(git rev-parse --show-toplevel)" +cd "$ROOT" +command -v python3 >/dev/null 2>&1 || { echo "ERROR: python3 required." >&2; exit 1; } +mkdir -p "$WORK" + +[ -f "$ELF_REL" ] || { echo "ERROR: $ELF_REL missing (copy the ethrex guest ELF in)." >&2; exit 1; } +if [ ! -f "$INPUT_REL" ]; then + echo "==> Generating ethrex ${TX_COUNT}-transfer fixture" + ( cd tooling/ethrex-fixtures && cargo build --release ) + tooling/ethrex-fixtures/target/release/ethrex-fixtures "$TX_COUNT" "$INPUT_REL" distinct +fi +ELF="$ROOT/$ELF_REL" +INPUT="$ROOT/$INPUT_REL" + +if [ "${REBUILD:-0}" = "1" ] || [ ! -x "$WORK/cli" ]; then + if [ -n "${CUDARC_PIN:-}" ]; then + sed -i "s/\"cuda-version-from-build-system\"/\"${CUDARC_PIN}\"/; /\"fallback-latest\"/d" \ + crypto/math-cuda/Cargo.toml + echo " cudarc pinned to ${CUDARC_PIN}" + fi + echo "==> Building cli (features: $BENCH_FEATURES)" + cargo build --release -p cli --features "$BENCH_FEATURES" + cp target/release/cli "$WORK/cli" + [ -n "${CUDARC_PIN:-}" ] && git checkout -- crypto/math-cuda/Cargo.toml +else + echo "==> Reusing cached cli (REBUILD=1 to force)" +fi + +run_prove() { # $1 = 0|1 (disable-flag) -> proving time (s) + local out t + out="$(LAMBDA_VM_DISABLE_GPU_COMPOSITION="$1" "$WORK/cli" prove "$ELF" \ + --private-input "$INPUT" -o "$PROOF" --time 2>&1)" + rm -f "$PROOF" + t="$(printf '%s\n' "$out" | grep -o 'Proving time: [0-9.]*' | awk '{print $3}')" + if [ -z "$t" ]; then + echo "ERROR: could not parse 'Proving time':" >&2; printf '%s\n' "$out" >&2; exit 1 + fi + echo "$t" +} + +# Warm-up (PTX load, pools, pinned alloc) so pair 1 isn't an outlier. +echo "==> Warm-up prove" +run_prove 0 >/dev/null + +echo "==> $N_PAIRS pairs, ethrex ${TX_COUNT} txs (A=GPU-comp ON, B=OFF; - = GPU faster)" +printf 'pair,a_time,b_time\n' > "$WORK/pairs.csv" +for i in $(seq 1 "$N_PAIRS"); do + if [ $((i % 2)) -eq 1 ]; then + a="$(run_prove 0)"; b="$(run_prove 1)" # odd: A then B + else + b="$(run_prove 1)"; a="$(run_prove 0)" # even: B then A (ABBA) + fi + printf '%d,%s,%s\n' "$i" "$a" "$b" >> "$WORK/pairs.csv" + printf ' pair %2d/%d A(GPU)=%ss B(CPU)=%ss %+.2f%%\n' \ + "$i" "$N_PAIRS" "$a" "$b" "$(awk "BEGIN{print ($a-$b)/$b*100}")" +done + +python3 - "$WORK/pairs.csv" <<'PY' +import sys, csv, math +rows = list(csv.DictReader(open(sys.argv[1]))) +A = [float(r['a_time']) for r in rows] # GPU composition ON +B = [float(r['b_time']) for r in rows] # GPU composition OFF (CPU) +n = len(A) +d = [(a - b) / b * 100.0 for a, b in zip(A, B)] +mean = sum(d) / n +var = sum((x - mean) ** 2 for x in d) / (n - 1) if n > 1 else 0.0 +sd = math.sqrt(var); se = sd / math.sqrt(n) if n else float('inf') +TT = {1:12.706,2:4.303,3:3.182,4:2.776,5:2.571,6:2.447,7:2.365,8:2.306,9:2.262, + 10:2.228,11:2.201,12:2.179,13:2.160,14:2.145,15:2.131,16:2.120,17:2.110, + 18:2.101,19:2.093,20:2.086,25:2.060,30:2.042,40:2.021,50:2.009,60:2.000} +df = n - 1 +tc = TT.get(df) or (1.96 if df > 120 else TT[min(TT, key=lambda k: abs(k - df))]) +lo, hi = mean - tc * se, mean + tc * se +def median(xs): + s = sorted(xs); m = len(s) + return s[m // 2] if m % 2 else (s[m // 2 - 1] + s[m // 2]) / 2 +med = median(d) +print("\n=== GPU-composition ABBA (A=GPU ON, B=CPU; - = GPU faster) ===") +print(f" pairs: {n} mean A (GPU): {sum(A)/n:.3f}s mean B (CPU): {sum(B)/n:.3f}s") +print(f" paired-t mean {mean:+.2f}% sd {sd:.2f}% se {se:.2f}% 95% CI [{lo:+.2f}%, {hi:+.2f}%]") +print(f" median {med:+.2f}%") +if hi < 0: + print(f" => GPU path faster by ~{-mean:.2f}% (CI below 0)") +elif lo > 0: + print(f" => GPU path slower by ~{mean:.2f}% (CI above 0)") +else: + print(f" => inconclusive at n={n} (CI straddles 0); point ~{med:+.2f}%") +PY