diff --git a/crypto/math-cuda/build.rs b/crypto/math-cuda/build.rs index b2f61f9a2..8827e1d98 100644 --- a/crypto/math-cuda/build.rs +++ b/crypto/math-cuda/build.rs @@ -115,4 +115,5 @@ fn main() { compile_ptx("deep.cu", "deep.ptx", have_nvcc); compile_ptx("fri.cu", "fri.ptx", have_nvcc); compile_ptx("inverse.cu", "inverse.ptx", have_nvcc); + compile_ptx("logup.cu", "logup.ptx", have_nvcc); } diff --git a/crypto/math-cuda/kernels/logup.cu b/crypto/math-cuda/kernels/logup.cu new file mode 100644 index 000000000..8823a9ed2 --- /dev/null +++ b/crypto/math-cuda/kernels/logup.cu @@ -0,0 +1,238 @@ +// LogUp aux build: fingerprint kernel. +// +// One ext3 fingerprint per (interaction, row): +// lc = bus_id + sum_e alpha_powers[alpha_idx(e)] * base_e +// base_e = const_e + sum_t coef_t * main_col[col_t][row] (Goldilocks base) +// fp = z - lc +// Mirrors stark::logup_gpu::eval_fingerprint byte for byte. +// +// Layouts: +// main: column-major, main[col * num_rows + row]. +// descriptor (CSR): interactions -> elements -> terms. +// alpha_powers: ext3 interleaved, 3 limbs each. +// out: ext3 interleaved, out[(k*num_rows + row)*3 + {0,1,2}]. + +#include "ext3.cuh" + +using namespace ext3; + +extern "C" __global__ void logup_fingerprint_ext3( + const uint64_t *__restrict__ main, + uint32_t num_rows, + uint32_t num_interactions, + const uint64_t *__restrict__ bus_ids, + const uint32_t *__restrict__ elem_offsets, + const uint32_t *__restrict__ elem_alpha_idx, + const uint64_t *__restrict__ elem_const, + const uint32_t *__restrict__ term_offsets, + const uint64_t *__restrict__ term_coef, + const uint32_t *__restrict__ term_col, + const uint64_t *__restrict__ alpha_powers, + uint64_t z0, uint64_t z1, uint64_t z2, + uint64_t *__restrict__ out) { + uint64_t tid = blockIdx.x * (uint64_t)blockDim.x + threadIdx.x; + uint64_t total = (uint64_t)num_interactions * (uint64_t)num_rows; + if (tid >= total) + return; + + uint32_t k = (uint32_t)(tid / num_rows); + uint32_t row = (uint32_t)(tid % num_rows); + + Fe3 lc = make(bus_ids[k], 0, 0); + + uint32_t e_hi = elem_offsets[k + 1]; + for (uint32_t e = elem_offsets[k]; e < e_hi; ++e) { + uint64_t base = elem_const[e]; + uint32_t t_hi = term_offsets[e + 1]; + for (uint32_t t = term_offsets[e]; t < t_hi; ++t) { + uint64_t col_val = main[(uint64_t)term_col[t] * (uint64_t)num_rows + row]; + base = goldilocks::add(base, goldilocks::mul(term_coef[t], col_val)); + } + uint32_t ai = elem_alpha_idx[e]; + Fe3 a = make(alpha_powers[ai * 3 + 0], alpha_powers[ai * 3 + 1], + alpha_powers[ai * 3 + 2]); + lc = add(lc, mul_base(a, base)); + } + + Fe3 z = make(z0, z1, z2); + Fe3 fp = sub(z, lc); + uint64_t o = tid * 3; + out[o + 0] = fp.a; + out[o + 1] = fp.b; + out[o + 2] = fp.c; +} + +// Term combine: one ext3 per (output column, row): +// term = sum_{k in col} signed_mult_k(row) * reciprocal_k[row] +// signed_mult_k = mult_const[k] + sum_t mult_coef_t * main_col[col_t][row] +// (receiver sign already folded into the coefficients by the builder). +// reciprocals: ext3 interleaved, [(k*num_rows + row)*3 + limb]. +// out: ext3 interleaved, [(col*num_rows + row)*3 + limb]. +extern "C" __global__ void logup_term_ext3( + const uint64_t *__restrict__ main, + uint32_t num_rows, + const uint64_t *__restrict__ reciprocals, + uint32_t num_out_cols, + const uint32_t *__restrict__ out_col_offsets, + const uint32_t *__restrict__ out_col_interactions, + const uint64_t *__restrict__ mult_const, + const uint32_t *__restrict__ mult_term_offsets, + const uint64_t *__restrict__ mult_term_coef, + const uint32_t *__restrict__ mult_term_col, + uint64_t *__restrict__ out) { + uint64_t tid = blockIdx.x * (uint64_t)blockDim.x + threadIdx.x; + uint64_t total = (uint64_t)num_out_cols * (uint64_t)num_rows; + if (tid >= total) + return; + + uint32_t col = (uint32_t)(tid / num_rows); + uint32_t row = (uint32_t)(tid % num_rows); + + Fe3 term = zero(); + uint32_t ki_hi = out_col_offsets[col + 1]; + for (uint32_t ki = out_col_offsets[col]; ki < ki_hi; ++ki) { + uint32_t k = out_col_interactions[ki]; + + uint64_t m = mult_const[k]; + uint32_t t_hi = mult_term_offsets[k + 1]; + for (uint32_t t = mult_term_offsets[k]; t < t_hi; ++t) { + uint64_t col_val = main[(uint64_t)mult_term_col[t] * (uint64_t)num_rows + row]; + m = goldilocks::add(m, goldilocks::mul(mult_term_coef[t], col_val)); + } + + uint64_t ro = ((uint64_t)k * num_rows + row) * 3; + Fe3 r = make(reciprocals[ro], reciprocals[ro + 1], reciprocals[ro + 2]); + term = add(term, mul_base(r, m)); + } + + uint64_t o = tid * 3; + out[o + 0] = term.a; + out[o + 1] = term.b; + out[o + 2] = term.c; +} + +// =========================================================================== +// Accumulated column (K4): running sum of the term columns, on device. +// row_sum[i] = sum over all term columns of term[col][i] +// S = inclusive prefix scan of row_sum ; L = S[n-1] ; offset = L / N +// acc[i] = S[i] - (i+1) * offset (matches build_accumulated_column) +// Additive 3-phase Hillis-Steele scan (mirrors inverse.cu, add not mul). +// =========================================================================== + +#define LOGUP_BLK 256 + +// row_sum[i] = sum_c term[(c*num_rows + i)] over all num_cols term columns. +extern "C" __global__ void logup_row_sum_ext3( + const uint64_t *__restrict__ terms, uint32_t num_cols, uint32_t num_rows, + uint64_t *__restrict__ row_sum) { + uint64_t i = blockIdx.x * (uint64_t)blockDim.x + threadIdx.x; + if (i >= num_rows) + return; + Fe3 s = zero(); + for (uint32_t c = 0; c < num_cols; ++c) { + uint64_t o = ((uint64_t)c * num_rows + i) * 3; + s = add(s, make(terms[o], terms[o + 1], terms[o + 2])); + } + row_sum[i * 3] = s.a; + row_sum[i * 3 + 1] = s.b; + row_sum[i * 3 + 2] = s.c; +} + +// Per-block inclusive additive scan; writes block totals (last valid element). +extern "C" __global__ void logup_scan_block_add_ext3( + const uint64_t *__restrict__ input, uint64_t n, + uint64_t *__restrict__ scan_out, uint64_t *__restrict__ block_totals) { + __shared__ Fe3 sh[LOGUP_BLK]; + uint32_t tid = threadIdx.x; + uint64_t gid = blockIdx.x * (uint64_t)LOGUP_BLK + tid; + Fe3 v = (gid < n) ? make(input[gid * 3], input[gid * 3 + 1], input[gid * 3 + 2]) + : zero(); + sh[tid] = v; + __syncthreads(); + for (uint32_t off = 1; off < LOGUP_BLK; off <<= 1) { + Fe3 t = (tid >= off) ? sh[tid - off] : zero(); + __syncthreads(); + if (tid >= off) + sh[tid] = add(sh[tid], t); + __syncthreads(); + } + if (gid < n) { + uint64_t o = gid * 3; + scan_out[o] = sh[tid].a; + scan_out[o + 1] = sh[tid].b; + scan_out[o + 2] = sh[tid].c; + } + uint64_t block_end = (blockIdx.x + 1) * (uint64_t)LOGUP_BLK; + uint32_t last = (block_end <= n) + ? (LOGUP_BLK - 1) + : (uint32_t)(n - blockIdx.x * (uint64_t)LOGUP_BLK - 1); + if (tid == last) { + uint64_t b = blockIdx.x * 3; + block_totals[b] = sh[tid].a; + block_totals[b + 1] = sh[tid].b; + block_totals[b + 2] = sh[tid].c; + } +} + +// Phase 3: block b>0 adds the scanned prefix of preceding block totals. +extern "C" __global__ void logup_apply_offsets_add_ext3( + uint64_t *__restrict__ scan_inout, uint64_t n, + const uint64_t *__restrict__ block_totals_scanned) { + if (blockIdx.x == 0) + return; + uint64_t gid = blockIdx.x * (uint64_t)LOGUP_BLK + threadIdx.x; + if (gid >= n) + return; + uint64_t ob = (blockIdx.x - 1) * 3; + Fe3 off = make(block_totals_scanned[ob], block_totals_scanned[ob + 1], + block_totals_scanned[ob + 2]); + uint64_t o = gid * 3; + Fe3 v = add(make(scan_inout[o], scan_inout[o + 1], scan_inout[o + 2]), off); + scan_inout[o] = v.a; + scan_inout[o + 1] = v.b; + scan_inout[o + 2] = v.c; +} + +// acc[i] = scan[i] - (i+1) * (L * inv_N), L = scan[n-1]. inv_N is ext3 (1/N). +extern "C" __global__ void logup_finalize_accum_ext3( + const uint64_t *__restrict__ scan, uint64_t n, uint64_t inv0, uint64_t inv1, + uint64_t inv2, uint64_t *__restrict__ acc) { + uint64_t i = blockIdx.x * (uint64_t)blockDim.x + threadIdx.x; + if (i >= n) + return; + uint64_t lo = (n - 1) * 3; + Fe3 L = make(scan[lo], scan[lo + 1], scan[lo + 2]); + Fe3 offset = mul(L, make(inv0, inv1, inv2)); + Fe3 s = make(scan[i * 3], scan[i * 3 + 1], scan[i * 3 + 2]); + Fe3 a = sub(s, mul_base(offset, i + 1)); + acc[i * 3] = a.a; + acc[i * 3 + 1] = a.b; + acc[i * 3 + 2] = a.c; +} + +// Assemble the row-major aux trace buffer from the resident committed term +// columns + the accumulated column: +// aux[row * num_aux_cols + col] = committed[col][row] (col < num_committed) +// = accumulated[row] (col == num_committed) +// terms layout is [col][row] (column-major); aux is row-major [row][col]. +extern "C" __global__ void logup_assemble_aux_ext3( + const uint64_t *__restrict__ committed, uint32_t num_committed, + const uint64_t *__restrict__ accumulated, uint32_t num_rows, + uint64_t *__restrict__ aux) { + uint64_t i = blockIdx.x * (uint64_t)blockDim.x + threadIdx.x; + if (i >= num_rows) + return; + uint32_t num_aux_cols = num_committed + 1; + for (uint32_t col = 0; col < num_committed; ++col) { + uint64_t src = ((uint64_t)col * num_rows + i) * 3; + uint64_t dst = ((uint64_t)i * num_aux_cols + col) * 3; + aux[dst] = committed[src]; + aux[dst + 1] = committed[src + 1]; + aux[dst + 2] = committed[src + 2]; + } + uint64_t asrc = i * 3; + uint64_t adst = ((uint64_t)i * num_aux_cols + num_committed) * 3; + aux[adst] = accumulated[asrc]; + aux[adst + 1] = accumulated[asrc + 1]; + aux[adst + 2] = accumulated[asrc + 2]; +} diff --git a/crypto/math-cuda/src/device.rs b/crypto/math-cuda/src/device.rs index 3a149a83f..4d25184ef 100644 --- a/crypto/math-cuda/src/device.rs +++ b/crypto/math-cuda/src/device.rs @@ -97,6 +97,7 @@ const BARY_PTX: &str = include_str!(concat!(env!("OUT_DIR"), "/barycentric.ptx") 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")); /// Number of CUDA streams in the pool. Larger pools let many rayon-parallel /// callers overlap on the GPU without serializing on stream ownership. The @@ -179,6 +180,13 @@ pub struct Backend { pub block_inclusive_scan_rev_ext3: CudaFunction, pub apply_block_offsets_rev_ext3: CudaFunction, pub batch_inverse_combine_ext3: CudaFunction, + pub logup_fingerprint_ext3: CudaFunction, + pub logup_term_ext3: CudaFunction, + pub logup_row_sum_ext3: CudaFunction, + pub logup_scan_block_add_ext3: CudaFunction, + pub logup_apply_offsets_add_ext3: CudaFunction, + pub logup_finalize_accum_ext3: CudaFunction, + pub logup_assemble_aux_ext3: CudaFunction, // Twiddle caches keyed by log_n. fwd_twiddles: Mutex>>>>, @@ -280,6 +288,7 @@ impl Backend { let deep = ctx.load_module(Ptx::from_src(DEEP_PTX))?; 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 mut streams = Vec::with_capacity(STREAM_POOL_SIZE); for _ in 0..STREAM_POOL_SIZE { @@ -360,6 +369,13 @@ impl Backend { .load_function("block_inclusive_scan_rev_ext3")?, apply_block_offsets_rev_ext3: inverse.load_function("apply_block_offsets_rev_ext3")?, batch_inverse_combine_ext3: inverse.load_function("batch_inverse_combine_ext3")?, + logup_fingerprint_ext3: logup.load_function("logup_fingerprint_ext3")?, + logup_term_ext3: logup.load_function("logup_term_ext3")?, + logup_row_sum_ext3: logup.load_function("logup_row_sum_ext3")?, + logup_scan_block_add_ext3: logup.load_function("logup_scan_block_add_ext3")?, + 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")?, fwd_twiddles: Mutex::new(vec![None; max_log]), inv_twiddles: Mutex::new(vec![None; max_log]), ctx, diff --git a/crypto/math-cuda/src/lde.rs b/crypto/math-cuda/src/lde.rs index 427d84351..9a36196bf 100644 --- a/crypto/math-cuda/src/lde.rs +++ b/crypto/math-cuda/src/lde.rs @@ -390,15 +390,33 @@ fn launch_row_to_col_major( /// Returns (merkle_nodes, column-major device buffer, row-major LDE Vec). The /// buffer is transposed to column-major (as required by the downstream GPU /// kernels DEEP/barycentric); callers wrap it in the appropriate LDE handle. +/// Row-major LDE input: either a host slice (uploaded) or an already-resident +/// device buffer (copied device-to-device, no PCIe upload). +enum InnerInput<'a> { + Host(&'a [u64]), + Dev(&'a CudaSlice), +} + +#[allow(clippy::type_complexity)] fn coset_lde_row_major_inner( - row_major: &[u64], + input: InnerInput, n: usize, total_cols: usize, blowup_factor: usize, weights: &[u64], what: &str, -) -> Result<(GpuMerkleTree, CudaSlice, Vec)> { - assert_eq!(row_major.len(), n * total_cols); + retain_trace_col_major: bool, +) -> Result<( + GpuMerkleTree, + CudaSlice, + Vec, + Option>, +)> { + let input_len = match &input { + InnerInput::Host(h) => h.len(), + InnerInput::Dev(d) => d.len(), + }; + assert_eq!(input_len, n * total_cols); assert!(n.is_power_of_two()); assert_eq!(weights.len(), n); assert!(blowup_factor.is_power_of_two()); @@ -420,10 +438,27 @@ fn coset_lde_row_major_inner( let be = backend()?; let stream = be.next_stream(); - // H2D into a zeroed lde_size*total_cols buffer; only the first n*total_cols - // rows carry data, the remainder are already zero (zero-padding for LDE). + // Fill a zeroed lde_size*total_cols buffer; only the first n*total_cols rows + // carry data, the remainder are already zero (zero-padding for LDE). Host + // input uploads (H2D); device input copies in place (D2D, no PCIe upload). let mut buf = stream.alloc_zeros::(lde_size * total_cols)?; - stream.memcpy_htod(row_major, &mut buf.slice_mut(0..n * total_cols))?; + match input { + InnerInput::Host(h) => stream.memcpy_htod(h, &mut buf.slice_mut(0..n * total_cols))?, + InnerInput::Dev(d) => stream.memcpy_dtod(d, &mut buf.slice_mut(0..n * total_cols))?, + } + + // Snapshot the trace-domain input (column-major) before the iNTT overwrites + // it in place. The LogUp aux fingerprint kernel reads the main trace in + // place from this buffer, so R1 aux build skips the ~3 GB main re-upload. + // Transpose is a plain row->col transpose on the first n rows (not yet + // bit-reversed): dst[col*n + row] = buf[row*total_cols + col]. + let trace_col_major = if retain_trace_col_major { + Some(launch_row_to_col_major( + &stream, be, &buf, n, total_cols, n as u64, + )?) + } else { + None + }; let inv_tw = be.inv_twiddles_for(log_n)?; let fwd_tw = be.fwd_twiddles_for(log_lde)?; @@ -508,7 +543,7 @@ fn coset_lde_row_major_inner( leaves_len: num_leaves, root, }; - Ok((tree, col_major_dev, lde_out)) + Ok((tree, col_major_dev, lde_out, trace_col_major)) } /// Row-major LDE + Keccak + Merkle, all on-device, keeping the Merkle tree @@ -526,19 +561,22 @@ pub fn coset_lde_row_major_with_merkle_tree_keep( blowup_factor: usize, weights: &[u64], ) -> Result<(GpuLdeBase, Vec)> { - let (tree, col_major_dev, lde_out) = coset_lde_row_major_inner( - row_major, + let (tree, col_major_dev, lde_out, trace_col_major) = coset_lde_row_major_inner( + InnerInput::Host(row_major), n, m, blowup_factor, weights, "coset_lde_row_major lde_size", + true, )?; let handle = GpuLdeBase { buf: Arc::new(col_major_dev), m, lde_size: n * blowup_factor, tree: Some(tree), + trace_dev: trace_col_major.map(Arc::new), + trace_rows: n, }; Ok((handle, lde_out)) } @@ -560,13 +598,43 @@ pub fn coset_lde_ext3_row_major_with_merkle_tree_keep( blowup_factor: usize, weights: &[u64], ) -> Result<(GpuLdeExt3, Vec)> { - let (tree, col_major_dev, lde_out) = coset_lde_row_major_inner( - row_major, + let (tree, col_major_dev, lde_out, _) = coset_lde_row_major_inner( + InnerInput::Host(row_major), n, m * 3, blowup_factor, weights, "coset_lde_ext3_row_major lde_size", + false, + )?; + let handle = GpuLdeExt3 { + buf: Arc::new(col_major_dev), + m, + lde_size: n * blowup_factor, + tree: Some(tree), + }; + Ok((handle, lde_out)) +} + +/// Like [`coset_lde_ext3_row_major_with_merkle_tree_keep`] but the input is an +/// already-resident device buffer (`n * m` ext3 elements, row-major, `n*m*3` +/// u64s). No PCIe upload: the buffer is copied device-to-device into the LDE +/// scratch. Used by the resident LogUp aux path. +pub fn coset_lde_ext3_row_major_with_merkle_tree_keep_dev( + input_dev: &CudaSlice, + n: usize, + m: usize, + blowup_factor: usize, + weights: &[u64], +) -> Result<(GpuLdeExt3, Vec)> { + let (tree, col_major_dev, lde_out, _) = coset_lde_row_major_inner( + InnerInput::Dev(input_dev), + n, + m * 3, + blowup_factor, + weights, + "coset_lde_ext3_row_major_dev lde_size", + false, )?; let handle = GpuLdeExt3 { buf: Arc::new(col_major_dev), @@ -590,6 +658,12 @@ pub struct GpuLdeBase { pub m: usize, pub lde_size: usize, pub tree: Option, + /// Trace-domain main columns, column-major `[col*trace_rows + row]`, kept + /// resident from the R1 main LDE so the LogUp aux fingerprint kernel reads + /// them in place (no re-upload). None unless the base keep path retained it. + pub trace_dev: Option>>, + /// Row count (n) of `trace_dev`; 0 when `trace_dev` is None. + pub trace_rows: usize, } /// Handle to an ext3 LDE kept live on device, de-interleaved into 3 base @@ -1177,6 +1251,8 @@ fn coset_lde_batch_base_into_with_merkle_tree_inner( m, lde_size, tree: None, + trace_dev: None, + trace_rows: 0, })) } else { drop(buf); diff --git a/crypto/math-cuda/src/lib.rs b/crypto/math-cuda/src/lib.rs index 37f4bc2b7..14eead415 100644 --- a/crypto/math-cuda/src/lib.rs +++ b/crypto/math-cuda/src/lib.rs @@ -10,6 +10,7 @@ pub mod device; pub mod fri; pub mod inverse; pub mod lde; +pub mod logup; pub mod merkle; pub mod ntt; diff --git a/crypto/math-cuda/src/logup.rs b/crypto/math-cuda/src/logup.rs new file mode 100644 index 000000000..4fe7e9cd3 --- /dev/null +++ b/crypto/math-cuda/src/logup.rs @@ -0,0 +1,439 @@ +//! GPU LogUp aux build kernels. +//! +//! Two stages, mirroring `stark::logup_gpu`: +//! 1. `logup_fingerprints_dev`: one ext3 fingerprint per (interaction, row). +//! 2. `logup_term_columns`: fingerprints -> batch inverse -> per-output-column +//! signed-multiplicity combine, producing the committed + virtual term +//! columns. +//! +//! The descriptor is passed as plain array slices ([`LogupDescriptor`]) so this +//! crate stays independent of the stark types. + +use std::sync::Arc; + +use cudarc::driver::{CudaSlice, CudaStream, LaunchConfig, PushKernelArg}; + +use crate::Result; +use crate::device::backend; +use crate::inverse::batch_inverse_ext3_dev; + +const BLOCK_SIZE: u32 = 256; + +/// Flat LogUp descriptor for one table (CSR arrays, canonical Goldilocks). Built +/// by `stark::logup_gpu::build_fingerprint_descriptor`. +pub struct LogupDescriptor<'a> { + pub num_interactions: usize, + // fingerprint + pub bus_ids: &'a [u64], + pub elem_offsets: &'a [u32], + pub elem_alpha_idx: &'a [u32], + pub elem_const: &'a [u64], + pub term_offsets: &'a [u32], + pub term_coef: &'a [u64], + pub term_col: &'a [u32], + // term combine + pub num_out_cols: usize, + pub out_col_offsets: &'a [u32], + pub out_col_interactions: &'a [u32], + pub mult_const: &'a [u64], + pub mult_term_offsets: &'a [u32], + pub mult_term_coef: &'a [u64], + pub mult_term_col: &'a [u32], +} + +fn cfg(total: usize) -> LaunchConfig { + LaunchConfig { + grid_dim: ((total as u32).div_ceil(BLOCK_SIZE), 1, 1), + block_dim: (BLOCK_SIZE, 1, 1), + shared_mem_bytes: 0, + } +} + +/// Fingerprint kernel over a device-resident main trace. Returns the ext3 fp +/// buffer (`num_interactions * num_rows * 3`, layout `[(k*num_rows+row)*3+limb]`). +fn fingerprints_into_dev( + main_dev: &CudaSlice, + num_rows: usize, + d: &LogupDescriptor, + alpha_powers: &[u64], + z: [u64; 3], + stream: &Arc, +) -> Result> { + let total = d.num_interactions * num_rows; + let mut out = unsafe { stream.alloc::(total * 3) }?; + if total == 0 { + return Ok(out); + } + let be = backend()?; + let bus_ids = stream.clone_htod(d.bus_ids)?; + let elem_offsets = stream.clone_htod(d.elem_offsets)?; + let elem_alpha_idx = stream.clone_htod(d.elem_alpha_idx)?; + let elem_const = stream.clone_htod(d.elem_const)?; + let term_offsets = stream.clone_htod(d.term_offsets)?; + let term_coef = stream.clone_htod(d.term_coef)?; + let term_col = stream.clone_htod(d.term_col)?; + let alpha = stream.clone_htod(alpha_powers)?; + let num_rows_u32 = num_rows as u32; + let num_int_u32 = d.num_interactions as u32; + let (z0, z1, z2) = (z[0], z[1], z[2]); + unsafe { + stream + .launch_builder(&be.logup_fingerprint_ext3) + .arg(main_dev) + .arg(&num_rows_u32) + .arg(&num_int_u32) + .arg(&bus_ids) + .arg(&elem_offsets) + .arg(&elem_alpha_idx) + .arg(&elem_const) + .arg(&term_offsets) + .arg(&term_coef) + .arg(&term_col) + .arg(&alpha) + .arg(&z0) + .arg(&z1) + .arg(&z2) + .arg(&mut out) + .launch(cfg(total))?; + } + Ok(out) +} + +/// Compute fingerprints from a host main trace (column-major, `num_cols*num_rows`), +/// returning the resident ext3 buffer. The stream is synchronised before return. +pub fn logup_fingerprints_dev( + main_cols: &[u64], + num_rows: usize, + d: &LogupDescriptor, + alpha_powers: &[u64], + z: [u64; 3], + stream: &Arc, +) -> Result> { + let main_dev = stream.clone_htod(main_cols)?; + let out = fingerprints_into_dev(&main_dev, num_rows, d, alpha_powers, z, stream)?; + stream.synchronize()?; + Ok(out) +} + +/// Full term-column pipeline: fingerprints -> batch inverse -> term combine. +/// Returns the host term columns (`num_out_cols * num_rows * 3`, ext3 +/// interleaved, layout `[(col*num_rows+row)*3+limb]`). +pub fn logup_term_columns( + main_cols: &[u64], + num_rows: usize, + d: &LogupDescriptor, + alpha_powers: &[u64], + z: [u64; 3], +) -> Result> { + let be = backend()?; + let stream = be.next_stream(); + let timing = std::env::var_os("LAMBDA_VM_LOGUP_TIMING").is_some(); + let t0 = std::time::Instant::now(); + let main_dev = stream.clone_htod(main_cols)?; + if timing { + stream.synchronize()?; + } + let t1 = std::time::Instant::now(); + + let fp = fingerprints_into_dev(&main_dev, num_rows, d, alpha_powers, z, &stream)?; + let n = d.num_interactions * num_rows; + let recip = batch_inverse_ext3_dev(&fp, n, &stream)?; + + let total = d.num_out_cols * num_rows; + let mut out = unsafe { stream.alloc::(total * 3) }?; + if total == 0 { + stream.synchronize()?; + return Ok(Vec::new()); + } + + let out_col_offsets = stream.clone_htod(d.out_col_offsets)?; + let out_col_interactions = stream.clone_htod(d.out_col_interactions)?; + let mult_const = stream.clone_htod(d.mult_const)?; + let mult_term_offsets = stream.clone_htod(d.mult_term_offsets)?; + let mult_term_coef = stream.clone_htod(d.mult_term_coef)?; + let mult_term_col = stream.clone_htod(d.mult_term_col)?; + let num_rows_u32 = num_rows as u32; + let num_out_u32 = d.num_out_cols as u32; + unsafe { + stream + .launch_builder(&be.logup_term_ext3) + .arg(&main_dev) + .arg(&num_rows_u32) + .arg(&recip) + .arg(&num_out_u32) + .arg(&out_col_offsets) + .arg(&out_col_interactions) + .arg(&mult_const) + .arg(&mult_term_offsets) + .arg(&mult_term_coef) + .arg(&mult_term_col) + .arg(&mut out) + .launch(cfg(total))?; + } + if timing { + stream.synchronize()?; + } + let t2 = std::time::Instant::now(); + let host = stream.clone_dtoh(&out)?; + stream.synchronize()?; + let t3 = std::time::Instant::now(); + if timing { + eprintln!( + "LOGUP_GPU rows={} cols={} h2d_main={:?} compute={:?} d2h_terms={:?}", + num_rows, + main_cols.len() / num_rows, + t1 - t0, + t2 - t1, + t3 - t2, + ); + } + Ok(host) +} + +// Additive multi-block inclusive scan (mirrors inverse::scan_into_fwd, add). +fn scan_add_inplace( + stream: &Arc, + be: &crate::device::Backend, + buf: &mut CudaSlice, + n: usize, +) -> Result<()> { + if n <= 1 { + return Ok(()); + } + let k = (n as u32).div_ceil(BLOCK_SIZE); + let mut scan_out = unsafe { stream.alloc::(3 * n) }?; + let mut block_totals = unsafe { stream.alloc::(3 * k as usize) }?; + let n_u64 = n as u64; + let phase = LaunchConfig { + grid_dim: (k, 1, 1), + block_dim: (BLOCK_SIZE, 1, 1), + shared_mem_bytes: 0, + }; + unsafe { + stream + .launch_builder(&be.logup_scan_block_add_ext3) + .arg(&*buf) + .arg(&n_u64) + .arg(&mut scan_out) + .arg(&mut block_totals) + .launch(phase)?; + } + if k > 1 { + scan_add_inplace(stream, be, &mut block_totals, k as usize)?; + unsafe { + stream + .launch_builder(&be.logup_apply_offsets_add_ext3) + .arg(&mut scan_out) + .arg(&n_u64) + .arg(&block_totals) + .launch(phase)?; + } + } + stream.memcpy_dtod(&scan_out, buf)?; + Ok(()) +} + +/// The aux trace produced entirely on device: the row-major ext3 aux columns +/// resident on the GPU (fed straight to the aux LDE, no host round-trip), the +/// column count, and the host-side table contribution `L`. +#[derive(Clone)] +pub struct ResidentAux { + /// Row-major ext3 aux columns `[row * num_aux_cols + col]` (`committed + 1`). + pub buf: Arc>, + pub num_aux_cols: usize, + pub num_rows: usize, + /// LogUp table contribution (`L`), for the bus public inputs. + pub table_contribution: [u64; 3], +} + +// Debug/PartialEq/Eq compare only the host-side metadata (the device buffer is +// not comparable and never differs when the metadata matches); these exist so a +// `TraceTable` holding an optional `ResidentAux` can keep its derives. +impl std::fmt::Debug for ResidentAux { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("ResidentAux") + .field("num_aux_cols", &self.num_aux_cols) + .field("num_rows", &self.num_rows) + .finish() + } +} +impl PartialEq for ResidentAux { + fn eq(&self, other: &Self) -> bool { + self.num_aux_cols == other.num_aux_cols + && self.num_rows == other.num_rows + && self.table_contribution == other.table_contribution + } +} +impl Eq for ResidentAux {} + +/// Full aux build on device: fingerprints → invert → term columns → accumulate +/// scan → assemble the row-major aux trace buffer, all resident. `inv_n` is +/// `1/num_rows` embedded in ext3. The stream is synchronised before return. +/// Main trace input for the resident aux build: either a host column-major +/// buffer to upload, or an already-resident device buffer (from the R1 main +/// LDE) to read in place. The device form skips the ~3 GB main re-upload. +#[derive(Clone, Copy)] +pub enum ResidentMain<'a> { + Host(&'a [u64]), + Dev(&'a CudaSlice), +} + +#[allow(clippy::too_many_arguments)] +pub fn logup_aux_resident( + main: ResidentMain, + num_rows: usize, + d: &LogupDescriptor, + alpha_powers: &[u64], + z: [u64; 3], + inv_n: [u64; 3], + stream: &Arc, +) -> Result { + let be = backend()?; + // Per-phase timing (env LAMBDA_VM_LOGUP_TIMING): sync between phases so wall + // time is attributed correctly. Off = no extra syncs, production path. + let timing = std::env::var_os("LAMBDA_VM_LOGUP_TIMING").is_some(); + let sync_if = |on: bool| -> Result<()> { + if on { + stream.synchronize()?; + } + Ok(()) + }; + let t0 = std::time::Instant::now(); + + // Resident device main = zero upload; host main = one H2D. `uploaded` owns + // the staged buffer for the function scope so `main_dev` can borrow it. + let uploaded: Option> = match main { + ResidentMain::Dev(_) => None, + ResidentMain::Host(h) => Some(stream.clone_htod(h)?), + }; + let main_dev: &CudaSlice = match (main, &uploaded) { + (ResidentMain::Dev(d), _) => d, + (ResidentMain::Host(_), Some(up)) => up, + _ => unreachable!(), + }; + let main_len = main_dev.len(); + sync_if(timing)?; + let t_h2d = std::time::Instant::now(); + + let fp = fingerprints_into_dev(main_dev, num_rows, d, alpha_powers, z, stream)?; + sync_if(timing)?; + let t_fp = std::time::Instant::now(); + + let n = d.num_interactions * num_rows; + let recip = batch_inverse_ext3_dev(&fp, n, stream)?; + sync_if(timing)?; + let t_inv = std::time::Instant::now(); + + // Term columns (committed + virtual), resident, layout [col][row]. + // num_out is always >= 1 (the accumulated column); num_committed = num_out - 1. + debug_assert!( + d.num_out_cols >= 1, + "logup descriptor has no output columns" + ); + let num_out = d.num_out_cols; + let mut terms = unsafe { stream.alloc::(num_out * num_rows * 3) }?; + let out_col_offsets = stream.clone_htod(d.out_col_offsets)?; + let out_col_interactions = stream.clone_htod(d.out_col_interactions)?; + let mult_const = stream.clone_htod(d.mult_const)?; + let mult_term_offsets = stream.clone_htod(d.mult_term_offsets)?; + let mult_term_coef = stream.clone_htod(d.mult_term_coef)?; + let mult_term_col = stream.clone_htod(d.mult_term_col)?; + sync_if(timing)?; + let t_desc = std::time::Instant::now(); + let num_rows_u32 = num_rows as u32; + let num_out_u32 = num_out as u32; + unsafe { + stream + .launch_builder(&be.logup_term_ext3) + .arg(main_dev) + .arg(&num_rows_u32) + .arg(&recip) + .arg(&num_out_u32) + .arg(&out_col_offsets) + .arg(&out_col_interactions) + .arg(&mult_const) + .arg(&mult_term_offsets) + .arg(&mult_term_coef) + .arg(&mult_term_col) + .arg(&mut terms) + .launch(cfg(num_out * num_rows))?; + } + sync_if(timing)?; + let t_term = std::time::Instant::now(); + + // row_sum over all term columns → additive scan → accumulated column. + let mut row_sum = unsafe { stream.alloc::(num_rows * 3) }?; + unsafe { + stream + .launch_builder(&be.logup_row_sum_ext3) + .arg(&terms) + .arg(&num_out_u32) + .arg(&num_rows_u32) + .arg(&mut row_sum) + .launch(cfg(num_rows))?; + } + scan_add_inplace(stream, be, &mut row_sum, num_rows)?; // row_sum now holds S + let (i0, i1, i2) = (inv_n[0], inv_n[1], inv_n[2]); + let mut accumulated = unsafe { stream.alloc::(num_rows * 3) }?; + let n_u64 = num_rows as u64; + unsafe { + stream + .launch_builder(&be.logup_finalize_accum_ext3) + .arg(&row_sum) + .arg(&n_u64) + .arg(&i0) + .arg(&i1) + .arg(&i2) + .arg(&mut accumulated) + .launch(cfg(num_rows))?; + } + + // Assemble row-major aux buffer: committed (num_out-1) cols + accumulated. + let num_committed = num_out - 1; + let num_aux_cols = num_committed + 1; + let mut aux = unsafe { stream.alloc::(num_aux_cols * num_rows * 3) }?; + let num_committed_u32 = num_committed as u32; + unsafe { + stream + .launch_builder(&be.logup_assemble_aux_ext3) + .arg(&terms) + .arg(&num_committed_u32) + .arg(&accumulated) + .arg(&num_rows_u32) + .arg(&mut aux) + .launch(cfg(num_rows))?; + } + sync_if(timing)?; + let t_accum_done = std::time::Instant::now(); + + // L = table_contribution = S[n-1] (sum of all term columns, all rows). + let l_host: Vec = stream.clone_dtoh(&row_sum.slice((num_rows - 1) * 3..num_rows * 3))?; + stream.synchronize()?; + if timing { + let t_end = std::time::Instant::now(); + let ms = |a: std::time::Instant, b: std::time::Instant| (b - a).as_secs_f64() * 1e3; + let main_mb = (main_len * 8) as f64 / 1e6; + eprintln!( + "LOGUP_RESIDENT rows={} out_cols={} interactions={} main={:.0}MB | \ + h2d_main={:.2} fp={:.2} inv={:.2} desc_up={:.2} term={:.2} accum={:.2} l_dtoh={:.2} total={:.2} ms", + num_rows, + num_out, + d.num_interactions, + main_mb, + ms(t0, t_h2d), + ms(t_h2d, t_fp), + ms(t_fp, t_inv), + ms(t_inv, t_desc), + ms(t_desc, t_term), + ms(t_term, t_accum_done), + ms(t_accum_done, t_end), + ms(t0, t_end), + ); + } + Ok(ResidentAux { + buf: Arc::new(aux), + num_aux_cols, + num_rows, + table_contribution: [l_host[0], l_host[1], l_host[2]], + }) +} diff --git a/crypto/math-cuda/tests/barycentric_strided.rs b/crypto/math-cuda/tests/barycentric_strided.rs index 377a2b531..d96f7128b 100644 --- a/crypto/math-cuda/tests/barycentric_strided.rs +++ b/crypto/math-cuda/tests/barycentric_strided.rs @@ -50,6 +50,8 @@ fn run_base(log_trace: u32, blowup: usize, num_cols: usize, seed: u64) { m: num_cols, lde_size, tree: None, + trace_dev: None, + trace_rows: 0, }; // Pre-strided buffer for non-strided reference: trace-size picks of each col. diff --git a/crypto/math-cuda/tests/deep.rs b/crypto/math-cuda/tests/deep.rs index 6ab63be10..b7c027914 100644 --- a/crypto/math-cuda/tests/deep.rs +++ b/crypto/math-cuda/tests/deep.rs @@ -178,6 +178,8 @@ fn run_parity( m: num_main, lde_size, tree: None, + trace_dev: None, + trace_rows: 0, }; let aux_handle = if num_aux > 0 { Some(GpuLdeExt3 { diff --git a/crypto/stark/src/gpu_lde.rs b/crypto/stark/src/gpu_lde.rs index 3f1d81846..5cfdbd8a6 100644 --- a/crypto/stark/src/gpu_lde.rs +++ b/crypto/stark/src/gpu_lde.rs @@ -76,6 +76,7 @@ pub fn reset_all_gpu_call_counters() { GPU_DEEP_CALLS.store(0, Ordering::Relaxed); GPU_FRI_CALLS.store(0, Ordering::Relaxed); GPU_BATCH_INVERT_CALLS.store(0, Ordering::Relaxed); + GPU_LOGUP_CALLS.store(0, Ordering::Relaxed); } pub(crate) static GPU_EXTEND_HALVES_CALLS: AtomicU64 = AtomicU64::new(0); @@ -83,6 +84,12 @@ pub fn gpu_extend_halves_calls() -> u64 { GPU_EXTEND_HALVES_CALLS.load(Ordering::Relaxed) } +/// LogUp aux-build term-column GPU path attempts (one per table that took it). +pub(crate) static GPU_LOGUP_CALLS: AtomicU64 = AtomicU64::new(0); +pub fn gpu_logup_calls() -> u64 { + GPU_LOGUP_CALLS.load(Ordering::Relaxed) +} + // ============================================================================ // Shared dispatch helpers // ============================================================================ @@ -1104,7 +1111,57 @@ unsafe fn ext3_slice_to_u64(col: &[FieldElement]) -> &[u64] { /// Convert ext3 evals (3*n u64s, interleaved) into a freshly allocated /// `Vec>` of length `n`. Caller must have established /// `E == Ext3`. -fn u64_to_ext3_vec(raw: &[u64]) -> Vec> +/// Like [`try_expand_leaf_and_tree_ext3_row_major_keep`] but the aux columns are +/// already resident on device (from the GPU LogUp aux build) — no host upload. +/// Consumes the resident buffer via the device-input LDE. +pub(crate) fn try_expand_leaf_and_tree_ext3_row_major_keep_dev( + ra: &math_cuda::logup::ResidentAux, + blowup_factor: usize, + weights: &[FieldElement], +) -> Option<( + MerkleTree, + math_cuda::lde::GpuLdeExt3, + Vec>, +)> +where + F: IsField + 'static, + E: IsField + 'static, + B: IsMerkleTreeBackend, +{ + if TypeId::of::() != TypeId::of::() + || TypeId::of::() != TypeId::of::() + { + return None; + } + let weights_u64 = unsafe { weights_to_u64::(weights) }; + + GPU_LDE_CALLS.fetch_add((ra.num_aux_cols * 3) as u64, Ordering::Relaxed); + GPU_LEAF_HASH_CALLS.fetch_add(1, Ordering::Relaxed); + GPU_MERKLE_TREE_CALLS.fetch_add(1, Ordering::Relaxed); + + let (handle, lde_u64) = math_cuda::lde::coset_lde_ext3_row_major_with_merkle_tree_keep_dev( + &ra.buf, + ra.num_rows, + ra.num_aux_cols, + blowup_factor, + &weights_u64, + ) + .ok()?; + + let lde_out: Vec> = unsafe { + let mut v = std::mem::ManuallyDrop::new(lde_u64); + Vec::from_raw_parts( + v.as_mut_ptr() as *mut FieldElement, + v.len() / 3, + v.capacity() / 3, + ) + }; + let root = handle.tree.as_ref()?.root; + let tree = MerkleTree::::from_root(root); + Some((tree, handle, lde_out)) +} + +pub(crate) fn u64_to_ext3_vec(raw: &[u64]) -> Vec> where E: IsField + 'static, { diff --git a/crypto/stark/src/instruments.rs b/crypto/stark/src/instruments.rs index f263558aa..96bf6ffae 100644 --- a/crypto/stark/src/instruments.rs +++ b/crypto/stark/src/instruments.rs @@ -204,6 +204,14 @@ pub struct Round1SubOps { pub aux_lde: Duration, /// Aux trace: commit_bit_reversed (Merkle) pub aux_merkle: Duration, + /// Aux build: LogUp fingerprint computation (CPU). + pub aux_fingerprint: Duration, + /// Aux build: fingerprint batch inverse (CPU). + pub aux_invert: Duration, + /// Aux build: term combine (CPU). + pub aux_term: Duration, + /// Aux build: accumulated-column running sum (CPU). + pub aux_accumulate: Duration, } /// Timing data collected inside `multi_prove`. @@ -225,6 +233,11 @@ static R1_MAIN_LDE_US: AtomicU64 = AtomicU64::new(0); static R1_MAIN_MERKLE_US: AtomicU64 = AtomicU64::new(0); static R1_AUX_LDE_US: AtomicU64 = AtomicU64::new(0); static R1_AUX_MERKLE_US: AtomicU64 = AtomicU64::new(0); +// Aux build (LogUp) sub-phases, CPU time accumulated across tables/chunks. +static AUX_FINGERPRINT_US: AtomicU64 = AtomicU64::new(0); +static AUX_INVERT_US: AtomicU64 = AtomicU64::new(0); +static AUX_TERM_US: AtomicU64 = AtomicU64::new(0); +static AUX_ACCUM_US: AtomicU64 = AtomicU64::new(0); thread_local! { static TIMING_DATA: RefCell> = const { RefCell::new(None) }; @@ -256,12 +269,28 @@ pub fn accum_r1_aux(lde: Duration, merkle: Duration) { R1_AUX_MERKLE_US.fetch_add(merkle.as_micros() as u64, Ordering::Relaxed); } +/// Aux build (LogUp term column) sub-phase CPU times, summed across chunks. +pub fn accum_aux_term(fingerprint: Duration, invert: Duration, term: Duration) { + AUX_FINGERPRINT_US.fetch_add(fingerprint.as_micros() as u64, Ordering::Relaxed); + AUX_INVERT_US.fetch_add(invert.as_micros() as u64, Ordering::Relaxed); + AUX_TERM_US.fetch_add(term.as_micros() as u64, Ordering::Relaxed); +} + +/// Aux build accumulated-column (running sum) CPU time. +pub fn accum_aux_accumulate(d: Duration) { + AUX_ACCUM_US.fetch_add(d.as_micros() as u64, Ordering::Relaxed); +} + pub fn take_r1_sub() -> Round1SubOps { Round1SubOps { main_lde: Duration::from_micros(R1_MAIN_LDE_US.swap(0, Ordering::Relaxed)), main_merkle: Duration::from_micros(R1_MAIN_MERKLE_US.swap(0, Ordering::Relaxed)), aux_lde: Duration::from_micros(R1_AUX_LDE_US.swap(0, Ordering::Relaxed)), aux_merkle: Duration::from_micros(R1_AUX_MERKLE_US.swap(0, Ordering::Relaxed)), + aux_fingerprint: Duration::from_micros(AUX_FINGERPRINT_US.swap(0, Ordering::Relaxed)), + aux_invert: Duration::from_micros(AUX_INVERT_US.swap(0, Ordering::Relaxed)), + aux_term: Duration::from_micros(AUX_TERM_US.swap(0, Ordering::Relaxed)), + aux_accumulate: Duration::from_micros(AUX_ACCUM_US.swap(0, Ordering::Relaxed)), } } @@ -278,6 +307,10 @@ pub fn reset_all() { R1_MAIN_MERKLE_US.store(0, Ordering::Relaxed); R1_AUX_LDE_US.store(0, Ordering::Relaxed); R1_AUX_MERKLE_US.store(0, Ordering::Relaxed); + AUX_FINGERPRINT_US.store(0, Ordering::Relaxed); + AUX_INVERT_US.store(0, Ordering::Relaxed); + AUX_TERM_US.store(0, Ordering::Relaxed); + AUX_ACCUM_US.store(0, Ordering::Relaxed); TIMING_DATA.with(|cell| { cell.borrow_mut().take(); }); diff --git a/crypto/stark/src/lib.rs b/crypto/stark/src/lib.rs index 92a7a9697..caa1c73a0 100644 --- a/crypto/stark/src/lib.rs +++ b/crypto/stark/src/lib.rs @@ -20,6 +20,8 @@ pub mod gpu_lde; pub mod grinding; #[cfg(feature = "instruments")] pub mod instruments; +#[cfg(feature = "cuda")] +pub mod logup_gpu; pub mod lookup; pub(crate) mod par; pub mod profile_markers; diff --git a/crypto/stark/src/logup_gpu.rs b/crypto/stark/src/logup_gpu.rs new file mode 100644 index 000000000..10f2d990c --- /dev/null +++ b/crypto/stark/src/logup_gpu.rs @@ -0,0 +1,1102 @@ +//! GPU LogUp aux build: compile a table's bus interactions into a flat +//! descriptor the device fingerprint kernel can walk, plus a CPU evaluator that +//! mirrors the kernel exactly (the parity test pins them together). +//! +//! Fingerprint per interaction k at row i: +//! lc = bus_id + Σ_e α^{alpha_idx(e)} · e(i) +//! fp = z - lc +//! where each bus element e is `const + Σ_t coef_t · col_t[i]` in the base field +//! (Goldilocks), matching `BusValue::accumulate_fingerprint` / +//! `Packing::accumulate_fingerprint_with`. + +use std::any::TypeId; + +use crate::lookup::{ + BusInteraction, BusValue, LOGUP_CHALLENGE_ALPHA, LinearTerm, Multiplicity, Packing, + compute_alpha_powers, +}; +use math::field::element::FieldElement; +use math::field::extensions_goldilocks::Degree3GoldilocksExtensionField; +use math::field::goldilocks::GoldilocksField; +use math::field::traits::{IsField, IsSubFieldOf}; + +/// Minimum trace rows for the GPU aux-build path. Below this the CPU build wins +/// (dispatch + upload overhead). Correctness is unaffected: the fallback is +/// byte identical. +const GPU_LOGUP_MIN_ROWS: usize = 1 << 10; + +/// Goldilocks modulus 2^64 - 2^32 + 1. Coefficients are stored canonical so the +/// device path does plain Goldilocks arithmetic. +const GOLDILOCKS_P: u64 = 0xFFFF_FFFF_0000_0001; + +// Packing shift constants (powers of two), canonical Goldilocks. +const SHIFT_8: u64 = 1 << 8; +const SHIFT_16: u64 = 1 << 16; +const SHIFT_24: u64 = 1 << 24; + +/// Reduce a signed coefficient into canonical Goldilocks, matching +/// `FieldElement::::from(i64)`. |c| < 2^63 < p so no overflow. +fn i64_to_canonical(c: i64) -> u64 { + if c >= 0 { + c as u64 % GOLDILOCKS_P + } else { + GOLDILOCKS_P - (c.unsigned_abs() % GOLDILOCKS_P) + } +} + +/// Canonical Goldilocks negation. +fn neg_canonical(x: u64) -> u64 { + if x == 0 { 0 } else { GOLDILOCKS_P - x } +} + +/// Committed-pair / absorbed split, mirroring `lookup::split_interactions`. +fn split_interactions(num: usize) -> (usize, usize) { + if num <= 2 { + (0, num) + } else if num % 2 == 1 { + ((num - 1) / 2, 1) + } else { + ((num - 2) / 2, 2) + } +} + +/// Encode a multiplicity as a signed linear form `const + Σ coef·col` (sign +/// baked in: negated for receivers so `term = m'·recip` needs no extra sign). +/// Mirrors `Multiplicity::evaluate_with` for every variant. +fn encode_signed_multiplicity(m: &Multiplicity, is_sender: bool) -> (u64, Vec<(u64, u32)>) { + let mut cst: u128 = 0; + let mut terms: Vec<(u64, u32)> = Vec::new(); + match m { + Multiplicity::One => cst = 1, + Multiplicity::Column(c) => terms.push((1, *c as u32)), + Multiplicity::Sum(a, b) => { + terms.push((1, *a as u32)); + terms.push((1, *b as u32)); + } + Multiplicity::Negated(c) => { + cst = 1; + terms.push((neg_canonical(1), *c as u32)); + } + Multiplicity::Diff(a, b) => { + terms.push((1, *a as u32)); + terms.push((neg_canonical(1), *b as u32)); + } + Multiplicity::Sum3(a, b, c) => { + terms.push((1, *a as u32)); + terms.push((1, *b as u32)); + terms.push((1, *c as u32)); + } + Multiplicity::Linear(ts) => { + for t in ts { + match *t { + LinearTerm::Column { + coefficient, + column, + } => terms.push((i64_to_canonical(coefficient), column as u32)), + LinearTerm::ColumnUnsigned { + coefficient, + column, + } => terms.push((coefficient % GOLDILOCKS_P, column as u32)), + LinearTerm::Constant(v) => cst += i64_to_canonical(v) as u128, + } + } + } + } + let mut cst = (cst % GOLDILOCKS_P as u128) as u64; + if !is_sender { + cst = neg_canonical(cst); + for t in terms.iter_mut() { + t.0 = neg_canonical(t.0); + } + } + (cst, terms) +} + +/// Flat descriptor for one table's fingerprints. CSR layout: interactions index +/// into elements, elements index into terms. All coefficients canonical +/// Goldilocks. Ready to upload to the device fingerprint kernel. +#[derive(Clone, Debug, Default)] +pub struct FingerprintDescriptor { + pub num_interactions: usize, + /// `alpha_powers` must hold this many powers `[1, α, ... α^{len-1}]`. + pub alpha_powers_len: usize, + /// Per interaction: the α^0 (bus id) constant. + pub bus_ids: Vec, + /// Per interaction CSR offsets into the element arrays (len + 1). + pub elem_offsets: Vec, + /// Per element: the α power index (>= 1). + pub elem_alpha_idx: Vec, + /// Per element: additive constant (canonical; 0 for packings). + pub elem_const: Vec, + /// Per element CSR offsets into the term arrays (len + 1). + pub term_offsets: Vec, + /// Per term: coefficient (canonical Goldilocks). + pub term_coef: Vec, + /// Per term: main column index. + pub term_col: Vec, + + // --- term-combine (K3) data --- + /// Number of output term columns = committed pairs + 1 virtual. + pub num_out_cols: usize, + /// Per interaction: signed multiplicity constant (negated for receivers). + pub mult_const: Vec, + /// Per interaction CSR offsets into the multiplicity term arrays (len + 1). + pub mult_term_offsets: Vec, + /// Per multiplicity term: coefficient (signed, canonical Goldilocks). + pub mult_term_coef: Vec, + /// Per multiplicity term: main column index. + pub mult_term_col: Vec, + /// Per output column CSR offsets into `out_col_interactions` (len + 1). + pub out_col_offsets: Vec, + /// Interaction indices grouped per output column. + pub out_col_interactions: Vec, +} + +impl FingerprintDescriptor { + fn push_element(&mut self, alpha_idx: u32, const_val: u64, terms: &[(u64, u32)]) { + self.elem_alpha_idx.push(alpha_idx); + self.elem_const.push(const_val); + for &(coef, col) in terms { + self.term_coef.push(coef); + self.term_col.push(col); + } + self.term_offsets.push(self.term_coef.len() as u32); + } + + /// Expand one `BusValue` into elements starting at `alpha_off`; return the + /// number of bus elements (alpha powers) consumed. Mirrors + /// `BusValue::accumulate_fingerprint` exactly. + fn push_bus_value(&mut self, bv: &BusValue, alpha_off: u32) -> u32 { + match bv { + BusValue::Packed { + start_column, + packing, + } => { + let c = *start_column as u32; + match packing { + Packing::Direct => self.push_element(alpha_off, 0, &[(1, c)]), + Packing::Word2L => { + self.push_element(alpha_off, 0, &[(1, c), (SHIFT_16, c + 1)]) + } + Packing::Word4L => self.push_element( + alpha_off, + 0, + &[ + (1, c), + (SHIFT_8, c + 1), + (SHIFT_16, c + 2), + (SHIFT_24, c + 3), + ], + ), + Packing::DWordWL => { + self.push_element(alpha_off, 0, &[(1, c)]); + self.push_element(alpha_off + 1, 0, &[(1, c + 1)]); + } + Packing::DWordHHW => { + self.push_element(alpha_off, 0, &[(1, c)]); + self.push_element(alpha_off + 1, 0, &[(1, c + 1), (SHIFT_16, c + 2)]); + } + Packing::DWordWHH => { + self.push_element(alpha_off, 0, &[(1, c), (SHIFT_16, c + 1)]); + self.push_element(alpha_off + 1, 0, &[(1, c + 2)]); + } + Packing::DWordHL => { + self.push_element(alpha_off, 0, &[(1, c), (SHIFT_16, c + 1)]); + self.push_element(alpha_off + 1, 0, &[(1, c + 2), (SHIFT_16, c + 3)]); + } + Packing::DWordBL => { + self.push_element( + alpha_off, + 0, + &[ + (1, c), + (SHIFT_8, c + 1), + (SHIFT_16, c + 2), + (SHIFT_24, c + 3), + ], + ); + self.push_element( + alpha_off + 1, + 0, + &[ + (1, c + 4), + (SHIFT_8, c + 5), + (SHIFT_16, c + 6), + (SHIFT_24, c + 7), + ], + ); + } + Packing::QuadHL => { + for i in 0..4u32 { + let cc = c + i * 2; + self.push_element(alpha_off + i, 0, &[(1, cc), (SHIFT_16, cc + 1)]); + } + } + Packing::QuadWL => { + for i in 0..4u32 { + self.push_element(alpha_off + i, 0, &[(1, c + i)]); + } + } + } + packing.num_bus_elements() as u32 + } + BusValue::Linear(terms) => { + let mut const_val: u128 = 0; + let mut t: Vec<(u64, u32)> = Vec::new(); + for term in terms { + match *term { + LinearTerm::Column { + coefficient, + column, + } => t.push((i64_to_canonical(coefficient), column as u32)), + LinearTerm::ColumnUnsigned { + coefficient, + column, + } => t.push((coefficient % GOLDILOCKS_P, column as u32)), + LinearTerm::Constant(value) => { + const_val += i64_to_canonical(value) as u128; + } + } + } + self.push_element(alpha_off, (const_val % GOLDILOCKS_P as u128) as u64, &t); + 1 + } + } + } +} + +/// Compile a table's interactions into a [`FingerprintDescriptor`]. +pub fn build_fingerprint_descriptor(interactions: &[BusInteraction]) -> FingerprintDescriptor { + let mut d = FingerprintDescriptor { + num_interactions: interactions.len(), + ..Default::default() + }; + d.elem_offsets.push(0); + d.term_offsets.push(0); + d.mult_term_offsets.push(0); + let mut max_bus_elements = 0usize; + for it in interactions { + d.bus_ids.push(it.bus_id % GOLDILOCKS_P); + max_bus_elements = max_bus_elements.max(it.num_bus_elements()); + let mut alpha_off = 1u32; + for bv in &it.values { + alpha_off += d.push_bus_value(bv, alpha_off); + } + d.elem_offsets.push(d.elem_alpha_idx.len() as u32); + + // Signed multiplicity for the term combine. + let (cst, terms) = encode_signed_multiplicity(&it.multiplicity, it.is_sender); + d.mult_const.push(cst); + for (coef, col) in terms { + d.mult_term_coef.push(coef); + d.mult_term_col.push(col); + } + d.mult_term_offsets.push(d.mult_term_coef.len() as u32); + } + d.alpha_powers_len = max_bus_elements; + + // Output term columns: committed pair p = {2p, 2p+1}; the trailing 1-2 + // absorbed interactions form one virtual column. + let (committed_pairs, absorbed) = split_interactions(interactions.len()); + d.out_col_offsets.push(0); + for p in 0..committed_pairs { + d.out_col_interactions.push(2 * p as u32); + d.out_col_interactions.push(2 * p as u32 + 1); + d.out_col_offsets.push(d.out_col_interactions.len() as u32); + } + for k in (interactions.len() - absorbed)..interactions.len() { + d.out_col_interactions.push(k as u32); + } + d.out_col_offsets.push(d.out_col_interactions.len() as u32); + d.num_out_cols = committed_pairs + 1; + d +} + +impl FingerprintDescriptor { + /// Borrow the static arrays as the math-cuda flat descriptor (challenges + /// `alpha_powers`/`z` are passed separately at call time). + pub fn as_cuda(&self) -> math_cuda::logup::LogupDescriptor<'_> { + math_cuda::logup::LogupDescriptor { + num_interactions: self.num_interactions, + bus_ids: &self.bus_ids, + elem_offsets: &self.elem_offsets, + elem_alpha_idx: &self.elem_alpha_idx, + elem_const: &self.elem_const, + term_offsets: &self.term_offsets, + term_coef: &self.term_coef, + term_col: &self.term_col, + num_out_cols: self.num_out_cols, + out_col_offsets: &self.out_col_offsets, + out_col_interactions: &self.out_col_interactions, + mult_const: &self.mult_const, + mult_term_offsets: &self.mult_term_offsets, + mult_term_coef: &self.mult_term_coef, + mult_term_col: &self.mult_term_col, + } + } +} + +/// GPU aux-build term columns. Returns `(committed_columns, virtual_column)` +/// byte identical to the CPU path, or `None` to fall back (non Goldilocks, +/// below threshold, no GPU, or a GPU error). The committed columns are written +/// to the aux trace; the virtual column feeds the accumulated column. +#[allow(clippy::type_complexity)] +pub fn try_build_term_columns_gpu( + interactions: &[BusInteraction], + main_cols: &[Vec>], + trace_len: usize, + challenges: &[FieldElement], +) -> Option<(Vec>>, Vec>)> +where + F: IsField + 'static, + E: IsField + 'static, +{ + if TypeId::of::() != TypeId::of::() + || TypeId::of::() != TypeId::of::() + { + return None; + } + if trace_len < GPU_LOGUP_MIN_ROWS || main_cols.is_empty() || interactions.is_empty() { + return None; + } + // Escape hatch for A/B measurement: force the CPU aux build. + if std::env::var_os("LAMBDA_VM_NO_GPU_LOGUP").is_some() { + return None; + } + + let desc = build_fingerprint_descriptor(interactions); + if desc.num_out_cols == 0 { + return None; + } + + // main trace -> column-major u64. SAFETY: F == Goldilocks (repr(u64)). + let num_cols = main_cols.len(); + let mut main_flat = vec![0u64; num_cols * trace_len]; + for (c, col) in main_cols.iter().enumerate() { + for (r, e) in col.iter().enumerate() { + main_flat[c * trace_len + r] = unsafe { *(e.value() as *const _ as *const u64) }; + } + } + + // z + alpha powers. SAFETY: E == ext3 (repr [u64; 3]). + let z_arr = unsafe { *(challenges[0].value() as *const _ as *const [u64; 3]) }; + let alpha = &challenges[LOGUP_CHALLENGE_ALPHA]; + let alpha_powers = compute_alpha_powers(alpha, desc.alpha_powers_len); + let mut alpha_flat = vec![0u64; alpha_powers.len() * 3]; + for (i, p) in alpha_powers.iter().enumerate() { + let l = unsafe { *(p.value() as *const _ as *const [u64; 3]) }; + alpha_flat[i * 3..i * 3 + 3].copy_from_slice(&l); + } + + let md = desc.as_cuda(); + let term_flat = + math_cuda::logup::logup_term_columns(&main_flat, trace_len, &md, &alpha_flat, z_arr) + .ok()?; + crate::gpu_lde::GPU_LOGUP_CALLS.fetch_add(1, std::sync::atomic::Ordering::Relaxed); + + // term_flat layout [(col*trace_len + row)*3 + limb]; last column is virtual. + let mut cols: Vec>> = Vec::with_capacity(desc.num_out_cols); + for col in 0..desc.num_out_cols { + let lo = col * trace_len * 3; + cols.push(crate::gpu_lde::u64_to_ext3_vec::( + &term_flat[lo..lo + trace_len * 3], + )); + } + let virtual_column = cols.pop().unwrap(); + Some((cols, virtual_column)) +} + +/// GPU-resident aux build: produces the row-major aux columns on device (fed +/// straight to the aux LDE, no host round-trip) + the table contribution `L`. +/// Returns `None` to fall back (non Goldilocks, below threshold, no GPU, GPU +/// error). This is the residency path that avoids the term-column download. +pub fn try_build_aux_resident_gpu( + interactions: &[BusInteraction], + main_cols: &[Vec>], + main_dev: Option<(&math_cuda::CudaSlice, usize)>, + trace_len: usize, + challenges: &[FieldElement], +) -> Option +where + F: IsField + 'static, + E: IsField + 'static, +{ + if TypeId::of::() != TypeId::of::() + || TypeId::of::() != TypeId::of::() + { + return None; + } + if trace_len < GPU_LOGUP_MIN_ROWS || main_cols.is_empty() || interactions.is_empty() { + return None; + } + if std::env::var_os("LAMBDA_VM_NO_GPU_LOGUP").is_some() { + return None; + } + let desc = build_fingerprint_descriptor(interactions); + if desc.num_out_cols == 0 { + return None; + } + + let num_cols = main_cols.len(); + // Reuse the resident main trace from the R1 main LDE (column-major + // `[col*trace_len + row]`, same column order as `main_cols`) when it matches + // this table exactly; otherwise flatten + upload the host columns. The + // resident buffer skips the ~3 GB main re-upload. + let resident_main = + main_dev.filter(|&(buf, rows)| rows == trace_len && buf.len() == num_cols * trace_len); + let mut main_flat = Vec::new(); + if resident_main.is_none() { + main_flat = vec![0u64; num_cols * trace_len]; + for (c, col) in main_cols.iter().enumerate() { + for (r, e) in col.iter().enumerate() { + main_flat[c * trace_len + r] = unsafe { *(e.value() as *const _ as *const u64) }; + } + } + } + let z_arr = unsafe { *(challenges[0].value() as *const _ as *const [u64; 3]) }; + let alpha = &challenges[LOGUP_CHALLENGE_ALPHA]; + let alpha_powers = compute_alpha_powers(alpha, desc.alpha_powers_len); + let mut alpha_flat = vec![0u64; alpha_powers.len() * 3]; + for (i, p) in alpha_powers.iter().enumerate() { + let l = unsafe { *(p.value() as *const _ as *const [u64; 3]) }; + alpha_flat[i * 3..i * 3 + 3].copy_from_slice(&l); + } + // 1/N embedded in ext3 (matches the CPU offset = L * FieldElement::::from(N).inv()). + let inv_n_e = FieldElement::::from(trace_len as u64).inv().ok()?; + let inv_n = unsafe { *(inv_n_e.value() as *const _ as *const [u64; 3]) }; + + let be = math_cuda::device::backend().ok()?; + let stream = be.next_stream(); + let md = desc.as_cuda(); + let main = match resident_main { + Some((buf, _)) => math_cuda::logup::ResidentMain::Dev(buf), + None => math_cuda::logup::ResidentMain::Host(&main_flat), + }; + let ra = math_cuda::logup::logup_aux_resident( + main, + trace_len, + &md, + &alpha_flat, + z_arr, + inv_n, + &stream, + ) + .ok()?; + crate::gpu_lde::GPU_LOGUP_CALLS.fetch_add(1, std::sync::atomic::Ordering::Relaxed); + Some(ra) +} + +/// CPU reference: evaluate the fingerprint of interaction `k` at one row from +/// the descriptor. The device kernel performs the identical computation. Used by +/// the parity test and as the spec the kernel mirrors. +pub fn eval_fingerprint<'a, F, E>( + d: &FingerprintDescriptor, + k: usize, + get_col: impl Fn(usize) -> &'a FieldElement, + alpha_powers: &[FieldElement], + z: &FieldElement, +) -> FieldElement +where + F: IsField + IsSubFieldOf + 'a, + E: IsField, +{ + let mut lc = FieldElement::::from(d.bus_ids[k]); + let e_lo = d.elem_offsets[k] as usize; + let e_hi = d.elem_offsets[k + 1] as usize; + for e in e_lo..e_hi { + let mut base = FieldElement::::from(d.elem_const[e]); + let t_lo = d.term_offsets[e] as usize; + let t_hi = d.term_offsets[e + 1] as usize; + for t in t_lo..t_hi { + let coef = FieldElement::::from(d.term_coef[t]); + base += &coef * get_col(d.term_col[t] as usize); + } + lc += &base * &alpha_powers[d.elem_alpha_idx[e] as usize]; + } + z - &lc +} + +/// CPU reference: term column `out_col` at `row` = Σ over the column's +/// interactions of `signed_multiplicity · reciprocal`. `reciprocals` is laid out +/// `[k * num_rows + row]` (batch inverse of the fingerprints). Mirrors the K3 +/// kernel and the production term/accumulate combine. +pub fn eval_term<'a, F, E>( + d: &FingerprintDescriptor, + out_col: usize, + row: usize, + num_rows: usize, + get_col: impl Fn(usize) -> &'a FieldElement, + reciprocals: &[FieldElement], +) -> FieldElement +where + F: IsField + IsSubFieldOf + 'a, + E: IsField, +{ + let mut term = FieldElement::::zero(); + let lo = d.out_col_offsets[out_col] as usize; + let hi = d.out_col_offsets[out_col + 1] as usize; + for ki in lo..hi { + let k = d.out_col_interactions[ki] as usize; + let mut m = FieldElement::::from(d.mult_const[k]); + let t_lo = d.mult_term_offsets[k] as usize; + let t_hi = d.mult_term_offsets[k + 1] as usize; + for t in t_lo..t_hi { + m += &FieldElement::::from(d.mult_term_coef[t]) + * get_col(d.mult_term_col[t] as usize); + } + term += &m * &reciprocals[k * num_rows + row]; + } + term +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::lookup::{PackingShifts, compute_alpha_powers}; + use math::field::extensions_goldilocks::Degree3GoldilocksExtensionField; + use math::field::goldilocks::GoldilocksField; + + type F = GoldilocksField; + type E = Degree3GoldilocksExtensionField; + + // Reference fingerprint via the production accumulate path (source of truth). + fn reference_fp( + it: &BusInteraction, + main: &[Vec>], + row: usize, + alpha_powers: &[FieldElement], + z: &FieldElement, + shifts: &PackingShifts, + ) -> FieldElement { + let mut lc = FieldElement::::from(it.bus_id); + let mut off = 1usize; + for bv in &it.values { + off += bv.accumulate_fingerprint(main, row, alpha_powers, off, &mut lc, shifts); + } + z - &lc + } + + fn lcg(state: &mut u64) -> u64 { + *state = state + .wrapping_mul(6364136223846793005) + .wrapping_add(1442695040888963407); + *state + } + + #[test] + fn descriptor_fingerprint_matches_accumulate_path() { + use crate::lookup::Multiplicity::One; + let interactions = vec![ + BusInteraction::sender(0u64, One, Packing::Direct.columns(&[0])), + BusInteraction::sender(1u64, One, Packing::Word2L.columns(&[1])), + BusInteraction::sender(2u64, One, Packing::Word4L.columns(&[1])), + BusInteraction::sender(3u64, One, Packing::DWordWL.columns(&[2])), + BusInteraction::sender(4u64, One, Packing::DWordHHW.columns(&[0])), + BusInteraction::sender(5u64, One, Packing::DWordWHH.columns(&[0])), + BusInteraction::sender(6u64, One, Packing::DWordHL.columns(&[0])), + BusInteraction::sender(7u64, One, Packing::QuadWL.columns(&[0])), + BusInteraction::sender(8u64, One, Packing::QuadHL.columns(&[0])), + BusInteraction::sender( + 9u64, + One, + vec![ + BusValue::linear(vec![ + LinearTerm::Column { + coefficient: 3, + column: 1, + }, + LinearTerm::Column { + coefficient: -2, + column: 4, + }, + LinearTerm::Constant(42), + ]), + BusValue::column(5), + ], + ), + ]; + + let num_cols = 8; + let num_rows = 16; + let mut st = 0x1234_5678_9abc_def0u64; + let main: Vec>> = (0..num_cols) + .map(|_| { + (0..num_rows) + .map(|_| FieldElement::::from(lcg(&mut st))) + .collect() + }) + .collect(); + + // Base-embedded random alpha/z: distinct powers exercise every coef/index. + let alpha = FieldElement::::from(lcg(&mut st)); + let z = FieldElement::::from(lcg(&mut st)); + let shifts = PackingShifts::::new(); + + let desc = build_fingerprint_descriptor(&interactions); + let max_be = interactions + .iter() + .map(|i| i.num_bus_elements()) + .max() + .unwrap(); + assert_eq!(desc.alpha_powers_len, max_be); + let alpha_powers = compute_alpha_powers(&alpha, max_be); + + for (k, it) in interactions.iter().enumerate() { + for row in 0..num_rows { + let got = eval_fingerprint::(&desc, k, |c| &main[c][row], &alpha_powers, &z); + let want = reference_fp(it, &main, row, &alpha_powers, &z, &shifts); + assert_eq!(got, want, "fingerprint mismatch interaction {k} row {row}"); + } + } + } + + fn mk_ext3(st: &mut u64) -> FieldElement { + FieldElement::::new([ + FieldElement::::from(lcg(st)), + FieldElement::::from(lcg(st)), + FieldElement::::from(lcg(st)), + ]) + } + + fn limbs(e: &FieldElement) -> [u64; 3] { + let v = e.value(); + [*v[0].value(), *v[1].value(), *v[2].value()] + } + + // GPU fingerprint kernel vs the CPU evaluator, byte for byte (full ext3 + // alpha/z so mul_base is exercised). Runs on the GPU box. + #[test] + #[ignore = "requires GPU; run with --ignored"] + fn gpu_fingerprints_match_cpu() { + use crate::lookup::Multiplicity::One; + let interactions = vec![ + BusInteraction::sender(0u64, One, Packing::Direct.columns(&[0])), + BusInteraction::sender(1u64, One, Packing::Word4L.columns(&[0])), + BusInteraction::sender(2u64, One, Packing::DWordHL.columns(&[0])), + BusInteraction::sender(3u64, One, Packing::QuadHL.columns(&[0])), + BusInteraction::sender( + 4u64, + One, + vec![ + BusValue::linear(vec![ + LinearTerm::Column { + coefficient: 3, + column: 1, + }, + LinearTerm::Column { + coefficient: -2, + column: 2, + }, + LinearTerm::Constant(7), + ]), + BusValue::column(3), + ], + ), + ]; + + let num_cols = 8; + let num_rows = 64; + let mut st = 0xabcd_ef01_2345_6789u64; + let main: Vec>> = (0..num_cols) + .map(|_| { + (0..num_rows) + .map(|_| FieldElement::::from(lcg(&mut st))) + .collect() + }) + .collect(); + let alpha = mk_ext3(&mut st); + let z = mk_ext3(&mut st); + + let desc = build_fingerprint_descriptor(&interactions); + let alpha_powers = compute_alpha_powers(&alpha, desc.alpha_powers_len); + + // CPU reference, layout [(k*num_rows + row)*3 + limb]. + let mut cpu = vec![0u64; interactions.len() * num_rows * 3]; + for k in 0..interactions.len() { + for row in 0..num_rows { + let fp = eval_fingerprint::(&desc, k, |c| &main[c][row], &alpha_powers, &z); + let o = (k * num_rows + row) * 3; + cpu[o..o + 3].copy_from_slice(&limbs(&fp)); + } + } + + // Flatten GPU inputs. + let mut main_flat = vec![0u64; num_cols * num_rows]; + for c in 0..num_cols { + for r in 0..num_rows { + main_flat[c * num_rows + r] = *main[c][r].value(); + } + } + let mut alpha_flat = vec![0u64; alpha_powers.len() * 3]; + for (i, p) in alpha_powers.iter().enumerate() { + alpha_flat[i * 3..i * 3 + 3].copy_from_slice(&limbs(p)); + } + + let be = math_cuda::device::backend().unwrap(); + let stream = be.next_stream(); + let md = desc.as_cuda(); + let out_dev = math_cuda::logup::logup_fingerprints_dev( + &main_flat, + num_rows, + &md, + &alpha_flat, + limbs(&z), + &stream, + ) + .unwrap(); + let gpu: Vec = stream.clone_dtoh(&out_dev).unwrap(); + stream.synchronize().unwrap(); + + assert_eq!(gpu, cpu, "GPU fingerprints mismatch CPU evaluator"); + } + + // Faithful reference for one term column: fingerprint every interaction, + // batch invert, then Σ ±(multiplicity·recip). Mirrors compute_logup_term_column. + fn reference_term_column( + ints: &[&BusInteraction], + main: &[Vec>], + num_rows: usize, + alpha_powers: &[FieldElement], + z: &FieldElement, + shifts: &PackingShifts, + ) -> Vec> { + let mut fps: Vec> = Vec::with_capacity(ints.len() * num_rows); + for it in ints { + for row in 0..num_rows { + fps.push(reference_fp(it, main, row, alpha_powers, z, shifts)); + } + } + FieldElement::inplace_batch_inverse(&mut fps).unwrap(); + let mut out = vec![FieldElement::::zero(); num_rows]; + for (row, slot) in out.iter_mut().enumerate() { + let mut acc = FieldElement::::zero(); + for (k, it) in ints.iter().enumerate() { + let m = it.multiplicity.evaluate_at_row(main, row); + let t = &m * &fps[k * num_rows + row]; + acc = acc + if it.is_sender { t } else { -t }; + } + *slot = acc; + } + out + } + + // Interaction set exercising committed pairs + virtual and several + // multiplicity forms (5 interactions -> 2 pairs + 1 virtual, absorbed=1). + fn term_test_interactions() -> Vec { + use crate::lookup::Multiplicity; + vec![ + BusInteraction::sender(0u64, Multiplicity::Column(4), Packing::Direct.columns(&[0])), + BusInteraction::receiver(1u64, Multiplicity::One, Packing::Word4L.columns(&[0])), + BusInteraction::sender( + 2u64, + Multiplicity::Sum(4, 5), + Packing::DWordHL.columns(&[0]), + ), + BusInteraction::receiver( + 3u64, + Multiplicity::Negated(6), + Packing::QuadHL.columns(&[0]), + ), + BusInteraction::sender( + 4u64, + Multiplicity::Linear(vec![ + LinearTerm::Column { + coefficient: 1, + column: 4, + }, + LinearTerm::Column { + coefficient: -1, + column: 5, + }, + ]), + vec![BusValue::column(1), BusValue::column(2)], + ), + ] + } + + // CPU-only: descriptor term combine (eval_term over host-inverted + // eval_fingerprint) matches the reference. De-risks the multiplicity + // descriptor + output grouping without a GPU. + #[test] + fn descriptor_term_matches_reference_cpu() { + let interactions = term_test_interactions(); + let num_cols = 8; + let num_rows = 32; + let mut st = 0x9e37_79b9_7f4a_7c15u64; + let main: Vec>> = (0..num_cols) + .map(|_| { + (0..num_rows) + .map(|_| FieldElement::::from(lcg(&mut st) % 251)) + .collect() + }) + .collect(); + let alpha = FieldElement::::from(lcg(&mut st)); + let z = FieldElement::::from(lcg(&mut st)); + let shifts = PackingShifts::::new(); + + let desc = build_fingerprint_descriptor(&interactions); + let alpha_powers = compute_alpha_powers(&alpha, desc.alpha_powers_len); + + // Reciprocals of every interaction's fingerprint, laid out [k*num_rows+row]. + let mut recips: Vec> = Vec::with_capacity(interactions.len() * num_rows); + for k in 0..interactions.len() { + for row in 0..num_rows { + recips.push(eval_fingerprint::( + &desc, + k, + |c| &main[c][row], + &alpha_powers, + &z, + )); + } + } + FieldElement::inplace_batch_inverse(&mut recips).unwrap(); + + let groups: [Vec<&BusInteraction>; 3] = [ + vec![&interactions[0], &interactions[1]], + vec![&interactions[2], &interactions[3]], + vec![&interactions[4]], + ]; + assert_eq!(desc.num_out_cols, 3); + for (col, g) in groups.iter().enumerate() { + let want = reference_term_column(g, &main, num_rows, &alpha_powers, &z, &shifts); + for row in 0..num_rows { + let got = eval_term::(&desc, col, row, num_rows, |c| &main[c][row], &recips); + assert_eq!(got, want[row], "term mismatch col {col} row {row}"); + } + } + } + + // Full GPU term pipeline (fingerprint -> batch invert -> term) vs the CPU + // reference, byte for byte. Covers committed pairs + the virtual column. + #[test] + #[ignore = "requires GPU; run with --ignored"] + fn gpu_term_columns_match_cpu() { + // 5 interactions -> 2 committed pairs + 1 virtual (odd, absorbed=1). + let interactions = term_test_interactions(); + + let num_cols = 8; + let num_rows = 64; + let mut st = 0x5151_2323_9797_0e0eu64; + // Small column values so multiplicities like Negated (0/1) stay meaningful. + let main: Vec>> = (0..num_cols) + .map(|_| { + (0..num_rows) + .map(|_| FieldElement::::from(lcg(&mut st) % 251)) + .collect() + }) + .collect(); + let alpha = mk_ext3(&mut st); + let z = mk_ext3(&mut st); + let shifts = PackingShifts::::new(); + + let desc = build_fingerprint_descriptor(&interactions); + let alpha_powers = compute_alpha_powers(&alpha, desc.alpha_powers_len); + + // CPU reference term columns: 2 committed pairs + virtual (last 1). + let mut cpu = vec![0u64; desc.num_out_cols * num_rows * 3]; + let mut ref_cols: Vec>> = Vec::new(); + ref_cols.push(reference_term_column( + &[&interactions[0], &interactions[1]], + &main, + num_rows, + &alpha_powers, + &z, + &shifts, + )); + ref_cols.push(reference_term_column( + &[&interactions[2], &interactions[3]], + &main, + num_rows, + &alpha_powers, + &z, + &shifts, + )); + ref_cols.push(reference_term_column( + &[&interactions[4]], + &main, + num_rows, + &alpha_powers, + &z, + &shifts, + )); + for (col, rc) in ref_cols.iter().enumerate() { + for (row, v) in rc.iter().enumerate() { + let o = (col * num_rows + row) * 3; + cpu[o..o + 3].copy_from_slice(&limbs(v)); + } + } + + // GPU pipeline. + let mut main_flat = vec![0u64; num_cols * num_rows]; + for c in 0..num_cols { + for r in 0..num_rows { + main_flat[c * num_rows + r] = *main[c][r].value(); + } + } + let mut alpha_flat = vec![0u64; alpha_powers.len() * 3]; + for (i, p) in alpha_powers.iter().enumerate() { + alpha_flat[i * 3..i * 3 + 3].copy_from_slice(&limbs(p)); + } + let md = desc.as_cuda(); + let gpu = + math_cuda::logup::logup_term_columns(&main_flat, num_rows, &md, &alpha_flat, limbs(&z)) + .unwrap(); + + assert_eq!(desc.num_out_cols, 3); + assert_eq!(gpu, cpu, "GPU term columns mismatch CPU reference"); + } + + // Reference accumulated column, mirroring build_accumulated_column_from_terms. + fn reference_accumulate( + cols: &[Vec>], + num_rows: usize, + ) -> (Vec>, FieldElement) { + let mut total = FieldElement::::zero(); + for row in 0..num_rows { + for c in cols { + total = &total + &c[row]; + } + } + let n = FieldElement::::from(num_rows as u64); + let offset = &total * n.inv().unwrap(); + let mut acc = FieldElement::::zero(); + let mut out = Vec::with_capacity(num_rows); + for row in 0..num_rows { + let mut rs = FieldElement::::zero(); + for c in cols { + rs = &rs + &c[row]; + } + acc = &acc + &rs - &offset; + out.push(acc.clone()); + } + (out, total) + } + + // Full resident aux pipeline (fingerprint → invert → term → scan → assemble) + // vs the CPU reference, byte for byte: the row-major aux buffer (committed + + // accumulated) and the table_contribution L. Runs on the GPU box. + #[test] + #[ignore = "requires GPU; run with --ignored"] + fn gpu_aux_resident_matches_cpu() { + let interactions = term_test_interactions(); // 2 committed pairs + 1 virtual + let num_cols = 8; + // > BLOCK_SIZE (256) so the grid wide scan recurses across multiple blocks. + let num_rows = 1024; + let mut st = 0x243f_6a88_85a3_08d3u64; + let main: Vec>> = (0..num_cols) + .map(|_| { + (0..num_rows) + .map(|_| FieldElement::::from(lcg(&mut st) % 251)) + .collect() + }) + .collect(); + let alpha = mk_ext3(&mut st); + let z = mk_ext3(&mut st); + let shifts = PackingShifts::::new(); + let desc = build_fingerprint_descriptor(&interactions); + let alpha_powers = compute_alpha_powers(&alpha, desc.alpha_powers_len); + + let committed = vec![ + reference_term_column( + &[&interactions[0], &interactions[1]], + &main, + num_rows, + &alpha_powers, + &z, + &shifts, + ), + reference_term_column( + &[&interactions[2], &interactions[3]], + &main, + num_rows, + &alpha_powers, + &z, + &shifts, + ), + ]; + let virtual_col = reference_term_column( + &[&interactions[4]], + &main, + num_rows, + &alpha_powers, + &z, + &shifts, + ); + let mut all = committed.clone(); + all.push(virtual_col); + let (acc, total) = reference_accumulate(&all, num_rows); + + let num_aux = committed.len() + 1; + let mut expected = vec![0u64; num_aux * num_rows * 3]; + for row in 0..num_rows { + for (col, c) in committed.iter().enumerate() { + let o = (row * num_aux + col) * 3; + expected[o..o + 3].copy_from_slice(&limbs(&c[row])); + } + let o = (row * num_aux + committed.len()) * 3; + expected[o..o + 3].copy_from_slice(&limbs(&acc[row])); + } + + let mut main_flat = vec![0u64; num_cols * num_rows]; + for c in 0..num_cols { + for r in 0..num_rows { + main_flat[c * num_rows + r] = *main[c][r].value(); + } + } + let mut alpha_flat = vec![0u64; alpha_powers.len() * 3]; + for (i, p) in alpha_powers.iter().enumerate() { + alpha_flat[i * 3..i * 3 + 3].copy_from_slice(&limbs(p)); + } + let inv_n = limbs(&FieldElement::::from(num_rows as u64).inv().unwrap()); + + let be = math_cuda::device::backend().unwrap(); + let stream = be.next_stream(); + let md = desc.as_cuda(); + let ra = math_cuda::logup::logup_aux_resident( + math_cuda::logup::ResidentMain::Host(&main_flat), + num_rows, + &md, + &alpha_flat, + limbs(&z), + inv_n, + &stream, + ) + .unwrap(); + assert_eq!(ra.num_aux_cols, num_aux); + let gpu: Vec = stream.clone_dtoh(&*ra.buf).unwrap(); + stream.synchronize().unwrap(); + + assert_eq!( + ra.table_contribution, + limbs(&total), + "table_contribution L mismatch" + ); + assert_eq!(gpu, expected, "resident aux buffer mismatch CPU reference"); + + // Resident-main path: upload main col-major to device, then build via + // ResidentMain::Dev (no host upload). Must be byte-identical to Host. + let main_dev = stream.clone_htod(&main_flat).unwrap(); + stream.synchronize().unwrap(); + let ra_dev = math_cuda::logup::logup_aux_resident( + math_cuda::logup::ResidentMain::Dev(&main_dev), + num_rows, + &md, + &alpha_flat, + limbs(&z), + inv_n, + &stream, + ) + .unwrap(); + let gpu_dev: Vec = stream.clone_dtoh(&*ra_dev.buf).unwrap(); + stream.synchronize().unwrap(); + assert_eq!( + ra_dev.table_contribution, ra.table_contribution, + "resident-main L mismatch vs host-upload path" + ); + assert_eq!( + gpu_dev, expected, + "resident-main aux buffer mismatch CPU reference" + ); + } +} diff --git a/crypto/stark/src/lookup.rs b/crypto/stark/src/lookup.rs index cd41ac15a..0c5e5cc22 100644 --- a/crypto/stark/src/lookup.rs +++ b/crypto/stark/src/lookup.rs @@ -971,8 +971,8 @@ impl< impl crate::traits::AIR for AirWithBuses where - F: IsFFTField + IsSubFieldOf + IsPrimeField + Send + Sync, - E: IsField + Send + Sync, + F: IsFFTField + IsSubFieldOf + IsPrimeField + Send + Sync + 'static, + E: IsField + Send + Sync + 'static, B: BoundaryConstraintBuilder, PI: Send + Sync, CS: ConstraintSet, @@ -1109,6 +1109,13 @@ where let trace_len = trace.num_rows(); let _table_name = self.name.as_deref().unwrap_or("UNKNOWN"); + // Device-resident trace-domain main columns from the R1 main LDE, cloned + // (Arc, cheap) into a local so no borrow of `trace` is held across the + // `set_aux_resident` mutable borrow below. When present, the resident aux + // build reads them in place and skips the ~3 GB main re-upload. + #[cfg(all(feature = "cuda", not(feature = "debug-checks")))] + let resident_main = trace.main_trace_dev.clone(); + // Split interactions: committed pairs get term columns, last 1-2 are absorbed (virtual) let (num_committed_pairs, absorbed_count) = split_interactions(num_interactions); @@ -1121,49 +1128,91 @@ where // tables with many interactions. // Without `parallel`: sequential over pairs, sequential over rows. let interactions = &self.auxiliary_trace_build_data.interactions; - let build_pair = |i: usize| { - compute_logup_term_column( - &[&interactions[i * 2], &interactions[i * 2 + 1]], - &main_segment_cols, - trace_len, - challenges, - _table_name, - ) - }; - #[cfg(feature = "parallel")] - let committed_columns: Vec>> = if trace_len <= LOGUP_CHUNK_SIZE { - (0..num_committed_pairs) - .into_par_iter() - .map(build_pair) - .collect() - } else { - (0..num_committed_pairs).map(build_pair).collect() - }; - #[cfg(not(feature = "parallel"))] - let committed_columns: Vec>> = - (0..num_committed_pairs).map(build_pair).collect(); - - // Virtual column for absorbed interactions (NOT written to trace). - let virtual_column = if absorbed_count == 2 { - compute_logup_term_column( - &[ - &interactions[num_interactions - 2], - &interactions[num_interactions - 1], - ], + // GPU-resident aux build (Goldilocks + ext3, not disk-spill, not + // debug-checks): build the aux columns on device and keep them resident + // for the aux LDE (no term-column download). Returns the table + // contribution; the host set_aux + CPU accumulate below are skipped. + #[cfg(all(feature = "cuda", not(feature = "debug-checks")))] + if trace.resident_aux_ok() + && let Some(ra) = crate::logup_gpu::try_build_aux_resident_gpu::( + interactions, &main_segment_cols, + resident_main.as_ref().map(|r| (r.buf.as_ref(), r.rows)), trace_len, challenges, - _table_name, - ) - } else { - compute_logup_term_column( - &[&interactions[num_interactions - 1]], - &main_segment_cols, - trace_len, - challenges, - _table_name, ) + { + let table_contribution = crate::gpu_lde::u64_to_ext3_vec::(&ra.table_contribution) + .pop() + .expect("one ext3 element"); + trace.set_aux_resident(ra); + return Some(BusPublicInputs { table_contribution }); + } + + // GPU aux build (Goldilocks + ext3 + above threshold) computes all term + // columns on device, byte identical, and falls back to the CPU build. + #[cfg(feature = "cuda")] + let gpu_term_cols = crate::logup_gpu::try_build_term_columns_gpu::( + interactions, + &main_segment_cols, + trace_len, + challenges, + ); + #[cfg(not(feature = "cuda"))] + #[allow(clippy::type_complexity)] + let gpu_term_cols: Option<(Vec>>, Vec>)> = None; + + let (committed_columns, virtual_column) = match gpu_term_cols { + Some(cols) => cols, + None => { + let build_pair = |i: usize| { + compute_logup_term_column( + &[&interactions[i * 2], &interactions[i * 2 + 1]], + &main_segment_cols, + trace_len, + challenges, + _table_name, + ) + }; + + #[cfg(feature = "parallel")] + let committed_columns: Vec>> = if trace_len <= LOGUP_CHUNK_SIZE + { + (0..num_committed_pairs) + .into_par_iter() + .map(build_pair) + .collect() + } else { + (0..num_committed_pairs).map(build_pair).collect() + }; + #[cfg(not(feature = "parallel"))] + let committed_columns: Vec>> = + (0..num_committed_pairs).map(build_pair).collect(); + + // Virtual column for absorbed interactions (NOT written to trace). + let virtual_column = if absorbed_count == 2 { + compute_logup_term_column( + &[ + &interactions[num_interactions - 2], + &interactions[num_interactions - 1], + ], + &main_segment_cols, + trace_len, + challenges, + _table_name, + ) + } else { + compute_logup_term_column( + &[&interactions[num_interactions - 1]], + &main_segment_cols, + trace_len, + challenges, + _table_name, + ) + }; + (committed_columns, virtual_column) + } }; // Write only committed columns to trace @@ -1338,7 +1387,7 @@ impl Multiplicity { /// Evaluate the multiplicity for a single row of column-major main data. #[inline] - fn evaluate_at_row( + pub(crate) fn evaluate_at_row( &self, main_segment_cols: &[Vec>], row: usize, @@ -1544,6 +1593,8 @@ where let process_chunk = |chunk_start: usize, result_chunk: &mut [FieldElement]| { let chunk_len = result_chunk.len(); + #[cfg(feature = "instruments")] + let _t0 = std::time::Instant::now(); // Phase 1 — fingerprints, laid out as [int_0 rows…, int_1 rows…]. // fp[k*chunk_len + i] = interaction k at row chunk_start+i. @@ -1595,10 +1646,14 @@ where } } + #[cfg(feature = "instruments")] + let _t1 = std::time::Instant::now(); // Phase 2: batch invert FieldElement::inplace_batch_inverse(&mut fingerprints) .expect("fingerprint is zero - probability of sampling zero is negligible"); + #[cfg(feature = "instruments")] + let _t2 = std::time::Instant::now(); // Phase 3: Compute terms for (i, result_elem) in result_chunk.iter_mut().enumerate() { let row = chunk_start + i; @@ -1612,6 +1667,8 @@ where } *result_elem = acc; } + #[cfg(feature = "instruments")] + crate::instruments::accum_aux_term(_t1 - _t0, _t2 - _t1, std::time::Instant::now() - _t2); }; #[cfg(feature = "parallel")] @@ -1646,6 +1703,8 @@ where return FieldElement::zero(); } let trace_len = term_columns[0].len(); + #[cfg(feature = "instruments")] + let _t_acc = std::time::Instant::now(); // Compute L = sum of all terms across all rows let mut table_contribution = FieldElement::::zero(); @@ -1670,6 +1729,8 @@ where trace.set_aux(row, acc_column_idx, accumulated.clone()); } + #[cfg(feature = "instruments")] + crate::instruments::accum_aux_accumulate(std::time::Instant::now() - _t_acc); table_contribution } diff --git a/crypto/stark/src/prover.rs b/crypto/stark/src/prover.rs index cdf1cd1b2..f47b837d1 100644 --- a/crypto/stark/src/prover.rs +++ b/crypto/stark/src/prover.rs @@ -2220,6 +2220,28 @@ pub trait IsStarkProver< #[cfg(feature = "instruments")] let __sp = crate::instruments::span("r1_aux_build"); + // Disk-spill needs the aux columns in the host trace to spill them, so + // disable the GPU-resident aux build (it would keep them device-only). + #[cfg(all(feature = "cuda", feature = "disk-spill"))] + if storage_mode == StorageMode::Disk { + for (_, trace, _) in air_trace_pairs.iter_mut() { + trace.set_resident_aux_ok(false); + } + } + + // Thread each table's device-resident trace-domain main columns (kept by + // the R1 main LDE) onto its trace so the LogUp aux fingerprint kernel + // reads them in place instead of re-uploading ~3 GB. Tables without a GPU + // main handle (CPU LDE, preprocessed) fall back to the host upload path. + #[cfg(all(feature = "cuda", not(feature = "debug-checks")))] + for ((_, trace, _), gpu_main) in air_trace_pairs.iter_mut().zip(main_gpu_handles.iter()) { + if let Some(handle) = gpu_main + && let Some(td) = &handle.trace_dev + { + trace.set_main_trace_dev(std::sync::Arc::clone(td), handle.trace_rows); + } + } + #[cfg(feature = "parallel")] let aux_iter = air_trace_pairs.par_iter_mut(); #[cfg(not(feature = "parallel"))] @@ -2323,6 +2345,41 @@ pub trait IsStarkProver< if air.has_aux_trace() { let lde_size = domain.interpolation_domain_size * domain.blowup_factor; + // Resident GPU path: aux columns already on device (from + // the resident LogUp aux build) — LDE straight from device + // memory, no upload, no host column extraction. When the + // resident build fired the host aux trace is empty, so a + // device LDE failure is a hard abort, not a fall through to + // the host path below (which would commit a zero aux trace). + #[cfg(feature = "cuda")] + if let Some(ra) = trace.aux_resident() { + #[cfg(feature = "instruments")] + let t_sub = Instant::now(); + let (tree, handle, aux_data) = + crate::gpu_lde::try_expand_leaf_and_tree_ext3_row_major_keep_dev::< + Field, + FieldExtension, + BatchedMerkleTreeBackend, + >( + ra, domain.blowup_factor, &twiddles.coset_weights + ) + .ok_or_else(|| { + ProvingError::Fft( + "resident aux LDE failed; host aux trace is empty" + .to_string(), + ) + })?; + let num_cols = ra.num_aux_cols; + #[cfg(feature = "instruments")] + crate::instruments::accum_r1_aux(t_sub.elapsed(), Duration::ZERO); + let root = tree.root; + return Ok(( + Some(TableCommit::plain(tree, root)), + (aux_data, num_cols), + Some(handle), + )); + } + // Fused GPU path (cuda only): row-major ext3 NTT — single // H2D, no column extraction, no CPU transpose. #[cfg(feature = "cuda")] diff --git a/crypto/stark/src/trace.rs b/crypto/stark/src/trace.rs index 831b95284..41e1eb675 100644 --- a/crypto/stark/src/trace.rs +++ b/crypto/stark/src/trace.rs @@ -30,8 +30,53 @@ where pub num_main_columns: usize, pub num_aux_columns: usize, pub step_size: usize, + /// LogUp aux columns built resident on device (pre-LDE), threaded from the + /// R1 aux build to the R1 aux commit so they feed the aux LDE without a host + /// round-trip. None on the CPU / download path. + #[cfg(feature = "cuda")] + pub(crate) aux_resident: Option, + /// Whether the GPU-resident aux build is allowed (false under disk-spill, + /// which needs the aux columns in the host trace to spill them). + #[cfg(feature = "cuda")] + pub(crate) resident_aux_ok: bool, + /// Trace-domain main columns kept resident on device from the R1 main LDE + /// (column-major `[col*rows + row]`), so the R1 LogUp aux fingerprint kernel + /// reads them in place instead of re-uploading ~3 GB. None when the GPU main + /// LDE did not run for this table. + #[cfg(feature = "cuda")] + pub(crate) main_trace_dev: Option, +} + +/// Device-resident trace-domain main columns (column-major `[col*rows + row]`), +/// retained from the R1 main LDE for the aux fingerprint kernel. GPU-only and +/// transient; the device buffer is excluded from logical trace equality (only +/// `rows` participates) and opaque in `Debug`, matching `ResidentAux`. +#[cfg(feature = "cuda")] +#[derive(Clone)] +pub(crate) struct ResidentMainTrace { + pub(crate) buf: std::sync::Arc>, + pub(crate) rows: usize, +} + +#[cfg(feature = "cuda")] +impl core::fmt::Debug for ResidentMainTrace { + fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { + f.debug_struct("ResidentMainTrace") + .field("rows", &self.rows) + .finish_non_exhaustive() + } +} + +#[cfg(feature = "cuda")] +impl PartialEq for ResidentMainTrace { + fn eq(&self, other: &Self) -> bool { + self.rows == other.rows + } } +#[cfg(feature = "cuda")] +impl Eq for ResidentMainTrace {} + impl TraceTable where E: IsField, @@ -54,6 +99,12 @@ where num_main_columns, num_aux_columns, step_size, + #[cfg(feature = "cuda")] + aux_resident: None, + #[cfg(feature = "cuda")] + resident_aux_ok: true, + #[cfg(feature = "cuda")] + main_trace_dev: None, } } @@ -76,6 +127,12 @@ where num_main_columns, num_aux_columns, step_size, + #[cfg(feature = "cuda")] + aux_resident: None, + #[cfg(feature = "cuda")] + resident_aux_ok: true, + #[cfg(feature = "cuda")] + main_trace_dev: None, } } @@ -91,6 +148,12 @@ where num_main_columns, num_aux_columns, step_size, + #[cfg(feature = "cuda")] + aux_resident: None, + #[cfg(feature = "cuda")] + resident_aux_ok: true, + #[cfg(feature = "cuda")] + main_trace_dev: None, } } @@ -98,6 +161,50 @@ where self.main_table.height } + /// Store the resident (pre-LDE) LogUp aux columns, threaded to the aux commit. + #[cfg(feature = "cuda")] + pub fn set_aux_resident(&mut self, ra: math_cuda::logup::ResidentAux) { + self.aux_resident = Some(ra); + } + + /// Borrow the resident aux columns (read by the aux commit for the LDE). + #[cfg(feature = "cuda")] + pub fn aux_resident(&self) -> Option<&math_cuda::logup::ResidentAux> { + self.aux_resident.as_ref() + } + + /// Whether the GPU-resident aux build is allowed (false under disk-spill). + #[cfg(feature = "cuda")] + pub fn resident_aux_ok(&self) -> bool { + self.resident_aux_ok + } + + /// Disable the GPU-resident aux build (host trace needed, e.g. disk-spill). + #[cfg(feature = "cuda")] + pub fn set_resident_aux_ok(&mut self, ok: bool) { + self.resident_aux_ok = ok; + } + + /// Stash the device-resident trace-domain main columns from the R1 main LDE + /// (column-major `[col*rows + row]`) so the aux fingerprint kernel reads them + /// in place. + #[cfg(feature = "cuda")] + pub fn set_main_trace_dev( + &mut self, + buf: std::sync::Arc>, + rows: usize, + ) { + self.main_trace_dev = Some(ResidentMainTrace { buf, rows }); + } + + /// The device-resident main trace `(buffer, rows)`, if retained by R1. + #[cfg(feature = "cuda")] + pub fn main_trace_dev(&self) -> Option<(&math_cuda::CudaSlice, usize)> { + self.main_trace_dev + .as_ref() + .map(|r| (r.buf.as_ref(), r.rows)) + } + pub fn num_steps(&self) -> usize { debug_assert!(self.main_table.height.is_multiple_of(self.step_size)); self.main_table.height / self.step_size diff --git a/prover/src/instruments.rs b/prover/src/instruments.rs index aa5d1caa4..0ea28273b 100644 --- a/prover/src/instruments.rs +++ b/prover/src/instruments.rs @@ -87,6 +87,26 @@ pub fn print_report( total, ); row_sub(" Aux trace build (parallel)", mp.aux_build, total); + row_sub( + " LogUp fingerprint (CPU)", + mp.round1_sub.aux_fingerprint, + total, + ); + row_sub( + " LogUp batch invert (CPU)", + mp.round1_sub.aux_invert, + total, + ); + row_sub( + " LogUp term combine (CPU)", + mp.round1_sub.aux_term, + total, + ); + row_sub( + " LogUp accumulate scan (CPU)", + mp.round1_sub.aux_accumulate, + total, + ); row_sub(" Aux trace commit", mp.aux_commit, total); row_sub( " Aux LDE (fused GPU: LDE+Keccak+Merkle / CPU: LDE only)", diff --git a/prover/tests/cuda_path_integration.rs b/prover/tests/cuda_path_integration.rs index 8033828bf..c78b16e25 100644 --- a/prover/tests/cuda_path_integration.rs +++ b/prover/tests/cuda_path_integration.rs @@ -12,10 +12,27 @@ 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_parts_lde_calls, + gpu_extend_halves_calls, gpu_fri_calls, gpu_lde_calls, gpu_logup_calls, gpu_parts_lde_calls, reset_all_gpu_call_counters, }; +/// The GPU LogUp aux-build path fires and still yields a verifying proof. +#[test] +#[ignore = "requires GPU; run with --ignored --nocapture"] +fn gpu_logup_aux_build_fires_and_verifies() { + let elf = asm_elf_bytes("fib_iterative_1M"); + reset_all_gpu_call_counters(); + let proof = prove(&elf).expect("prove"); + assert!( + gpu_logup_calls() > 0, + "GPU LogUp aux-build path did not fire (tables below threshold or fell back)" + ); + assert!( + verify(&proof, &elf).expect("verify"), + "proof failed to verify" + ); +} + #[test] #[ignore = "requires GPU; run with --ignored --nocapture"] fn gpu_path_fires_end_to_end() {