From c690d3554004a931278c2102d5e141bee59c5faa Mon Sep 17 00:00:00 2001 From: Joaquin Carletti Date: Mon, 13 Jul 2026 16:12:50 -0300 Subject: [PATCH 1/3] add new tables --- crypto/math-cuda/kernels/trace_cpu.cu | 265 +++++++++++++++++ crypto/math-cuda/src/device.rs | 8 + crypto/math-cuda/src/trace_cpu.rs | 192 ++++++++++++ prover/src/tables/gpu_trace.rs | 272 +++++++++++++++++ prover/src/tables/trace_builder.rs | 402 ++++++++++++++++++++++++++ 5 files changed, 1139 insertions(+) diff --git a/crypto/math-cuda/kernels/trace_cpu.cu b/crypto/math-cuda/kernels/trace_cpu.cu index 0f33e2f1d..e5c236465 100644 --- a/crypto/math-cuda/kernels/trace_cpu.cu +++ b/crypto/math-cuda/kernels/trace_cpu.cu @@ -467,6 +467,271 @@ extern "C" __global__ void lt_fill(const uint64_t *ops, uint64_t n, out[base + 16] = mult; // MU } +// On-GPU EQ trace fill (12 cols). Mirrors the per-row compute of +// `prover/src/tables/eq.rs::generate_eq_trace`. Dedup is done on the HOST (the +// same per-chunk HashMap the CPU uses); this kernel receives already-unique ops +// with their summed multiplicity, one per row, and recomputes diff/eq/res. +// Order-independent (LogUp ALU bus) → validated by multiset/prove, not byte order. +// Padding rows (r >= n) stay all-zero. +// +// Packed stride EQ_STRIDE: [0]=a [1]=b [2]=flags (bit0 invert) [3]=multiplicity. +#define EQ_NCOLS 12u +#define EQ_STRIDE 4u + +extern "C" __global__ void eq_fill(const uint64_t *ops, uint64_t n, + uint64_t num_rows, uint64_t *out) { + uint64_t r = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x; + if (r >= num_rows) + return; + uint64_t base = r * EQ_NCOLS; + if (r >= n) + return; + + const uint64_t *op = ops + r * EQ_STRIDE; + uint64_t a = op[0]; + uint64_t b = op[1]; + uint64_t fl = op[2]; + uint64_t mult = op[3]; + uint64_t invert = fl & 1u; + + uint64_t eq = (a == b) ? 1u : 0u; + uint64_t res = eq ^ invert; + uint64_t diff = a - b; // wrapping + + out[base + 0] = a & 0xFFFFFFFFu; // A_0 (DWordWL lo) + out[base + 1] = a >> 32; // A_1 (DWordWL hi) + out[base + 2] = b & 0xFFFFFFFFu; // B_0 + out[base + 3] = b >> 32; // B_1 + out[base + 4] = invert; // INVERT + out[base + 5] = res; // RES = eq ^ invert + out[base + 6] = diff & 0xFFFFu; // DIFF_0 (DWordHL halves) + out[base + 7] = (diff >> 16) & 0xFFFFu; // DIFF_1 + out[base + 8] = (diff >> 32) & 0xFFFFu; // DIFF_2 + out[base + 9] = (diff >> 48) & 0xFFFFu; // DIFF_3 + out[base + 10] = eq; // EQ = (a == b) + out[base + 11] = mult; // MU +} + +// On-GPU BYTEWISE trace fill (26 cols). Mirrors the per-row compute of +// `prover/src/tables/bytewise.rs::generate_bytewise_trace`. Dedup is done on the +// HOST (the same per-chunk HashMap the CPU uses); this kernel receives unique ops +// + summed multiplicity, one per row, and recomputes res = a AND/OR/XOR b, then +// byte-splits a/b/res (DWordBL). Order-independent (LogUp ALU bus) → validated by +// multiset/prove, not byte order. Padding rows (r >= n) stay all-zero. +// +// Packed stride BYTEWISE_STRIDE: [0]=a [1]=b [2]=op (alu_op: 0 AND, 1 OR, 2 XOR) +// [3]=multiplicity. +#define BYTEWISE_NCOLS 26u +#define BYTEWISE_STRIDE 4u + +extern "C" __global__ void bytewise_fill(const uint64_t *ops, uint64_t n, + uint64_t num_rows, uint64_t *out) { + uint64_t r = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x; + if (r >= num_rows) + return; + uint64_t base = r * BYTEWISE_NCOLS; + if (r >= n) + return; + + const uint64_t *op = ops + r * BYTEWISE_STRIDE; + uint64_t a = op[0]; + uint64_t b = op[1]; + uint64_t opc = op[2]; + uint64_t mult = op[3]; + + // op: 0 AND, 1 OR, 2 XOR (only these reach BYTEWISE). + uint64_t res = (opc == 0u) ? (a & b) : (opc == 1u) ? (a | b) : (a ^ b); + + for (int i = 0; i < 8; ++i) { + out[base + 0 + i] = (a >> (i * 8)) & 0xFFu; // A[0..8] (DWordBL) + out[base + 8 + i] = (b >> (i * 8)) & 0xFFu; // B[0..8] + out[base + 17 + i] = (res >> (i * 8)) & 0xFFu; // RES[0..8] + } + out[base + 16] = opc; // OP + out[base + 25] = mult; // MU +} + +// On-GPU MUL trace fill (26 cols). Mirrors the per-row compute of +// `prover/src/tables/mul.rs::generate_mul_trace` — the full 128-bit signed/unsigned +// product plus the sign-extended convolution `raw_product[0..4]`. Dedup is done on +// the HOST (the same per-chunk HashMap the CPU uses, keyed by op with split +// mu_lo/mu_hi from `wants_hi`); this kernel receives unique ops + both +// multiplicities, one per row. Order-independent (LogUp ALU bus) → validated by +// multiset/prove, not byte order. Padding rows (r >= n) stay all-zero. +// +// Packed stride MUL_STRIDE: [0]=lhs [1]=rhs [2]=flags (bit0 lhs_signed, bit1 +// rhs_signed) [3]=mu_lo [4]=mu_hi. +#define MUL_NCOLS 26u +#define MUL_STRIDE 5u + +extern "C" __global__ void mul_fill(const uint64_t *ops, uint64_t n, + uint64_t num_rows, uint64_t *out) { + uint64_t r = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x; + if (r >= num_rows) + return; + uint64_t base = r * MUL_NCOLS; + if (r >= n) + return; + + const uint64_t *op = ops + r * MUL_STRIDE; + uint64_t lhs = op[0]; + uint64_t rhs = op[1]; + uint64_t fl = op[2]; + uint64_t mu_lo = op[3]; + uint64_t mu_hi = op[4]; + uint64_t lhs_signed = fl & 1u; + uint64_t rhs_signed = (fl >> 1) & 1u; + + // Full 128-bit product. Signed operands sign-extend (int64 -> int128); + // unsigned zero-extend (uint64 -> int128, value-preserving since < 2^64). + __int128 a = lhs_signed ? (__int128)(int64_t)lhs : (__int128)lhs; + __int128 b = rhs_signed ? (__int128)(int64_t)rhs : (__int128)rhs; + __int128 product = a * b; // wrapping + uint64_t lo = (uint64_t)product; + uint64_t hi = (uint64_t)((unsigned __int128)product >> 64); + + uint64_t lhs_is_neg = (lhs_signed && ((int64_t)lhs < 0)) ? 1u : 0u; + uint64_t rhs_is_neg = (rhs_signed && ((int64_t)rhs < 0)) ? 1u : 0u; + + // Sign-extended halfword arrays: [0..4] = halfwords, [4..8] = 0xFFFF*is_neg. + uint64_t lfill = lhs_is_neg ? 0xFFFFull : 0ull; + uint64_t rfill = rhs_is_neg ? 0xFFFFull : 0ull; + uint64_t lhs_ext[8]; + uint64_t rhs_ext[8]; + for (int j = 0; j < 4; ++j) { + lhs_ext[j] = (lhs >> (16 * j)) & 0xFFFFu; + rhs_ext[j] = (rhs >> (16 * j)) & 0xFFFFu; + } + for (int j = 4; j < 8; ++j) { + lhs_ext[j] = lfill; + rhs_ext[j] = rfill; + } + + // raw_product[i] = Σ_k 2^(16k) Σ_j lhs_ext[j]*rhs_ext[idx-j], idx = 2i+k. + uint64_t raw[4]; + for (int i = 0; i < 4; ++i) { + unsigned __int128 sum = 0; + for (int k = 0; k <= 1; ++k) { + int idx = 2 * i + k; + if (idx < 8) { + for (int j = 0; j <= idx; ++j) { + if (j < 8 && (idx - j) < 8) { + unsigned __int128 term = + (unsigned __int128)lhs_ext[j] * (unsigned __int128)rhs_ext[idx - j]; + sum += term << (16 * k); + } + } + } + } + raw[i] = (uint64_t)sum; + } + + for (int j = 0; j < 4; ++j) { + out[base + 0 + j] = (lhs >> (16 * j)) & 0xFFFFu; // LHS_0..3 (DWordHL) + out[base + 5 + j] = (rhs >> (16 * j)) & 0xFFFFu; // RHS_0..3 + out[base + 10 + j] = (lo >> (16 * j)) & 0xFFFFu; // LO_0..3 + out[base + 14 + j] = (hi >> (16 * j)) & 0xFFFFu; // HI_0..3 + } + out[base + 4] = lhs_signed; // LHS_SIGNED + out[base + 9] = rhs_signed; // RHS_SIGNED + out[base + 18] = lhs_is_neg; // LHS_IS_NEGATIVE + out[base + 19] = rhs_is_neg; // RHS_IS_NEGATIVE + out[base + 20] = raw[0]; // RAW_PRODUCT_0..3 + out[base + 21] = raw[1]; + out[base + 22] = raw[2]; + out[base + 23] = raw[3]; + out[base + 24] = mu_lo; // MU_LO + out[base + 25] = mu_hi; // MU_HI +} + +// On-GPU DVRM trace fill (34 cols). Mirrors the per-row compute of +// `prover/src/tables/dvrm.rs::generate_dvrm_trace` — RISC-V signed/unsigned +// division & remainder with the div-by-zero and MIN/-1 overflow special cases, +// plus the abs/sign aux columns and n_sub_r. Dedup is done on the HOST (per-chunk +// HashMap keyed by op with split mu_q/mu_r from `wants_remainder`); this kernel +// receives unique ops + both multiplicities, one per row. Order-independent +// (LogUp ALU bus) → validated by multiset/prove, not byte order. Padding rows +// (r >= n) stay all-zero. +// +// Packed stride DVRM_STRIDE: [0]=n [1]=d [2]=flags (bit0 signed) [3]=mu_q [4]=mu_r. +#define DVRM_NCOLS 34u +#define DVRM_STRIDE 5u + +extern "C" __global__ void dvrm_fill(const uint64_t *ops, uint64_t n_ops, + uint64_t num_rows, uint64_t *out) { + uint64_t ri = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x; + if (ri >= num_rows) + return; + uint64_t base = ri * DVRM_NCOLS; + if (ri >= n_ops) + return; + + const uint64_t *op = ops + ri * DVRM_STRIDE; + uint64_t n = op[0]; + uint64_t d = op[1]; + uint64_t fl = op[2]; + uint64_t mu_q = op[3]; + uint64_t mu_r = op[4]; + uint64_t is_signed = fl & 1u; + + uint64_t div_by_zero = (d == 0ull) ? 1u : 0u; + uint64_t overflow = + (is_signed && n == 0x8000000000000000ull && d == 0xFFFFFFFFFFFFFFFFull) ? 1u : 0u; + + // Branch order matches the Rust: div-by-zero and overflow are handled before + // the divide, so the signed path never hits INT64_MIN / -1 (UB) or /0. + uint64_t q, rem; + if (div_by_zero) { + q = 0xFFFFFFFFFFFFFFFFull; + rem = n; + } else if (overflow) { + q = n; // i64::MIN + rem = 0ull; + } else if (is_signed) { + int64_t ni = (int64_t)n; + int64_t di = (int64_t)d; + q = (uint64_t)(ni / di); // truncates toward zero, == wrapping_div + rem = (uint64_t)(ni % di); // == wrapping_rem (sign follows dividend) + } else { + q = n / d; + rem = n % d; + } + + uint64_t sign_n = (is_signed && (n >> 63)) ? 1u : 0u; + uint64_t sign_d = (is_signed && (d >> 63)) ? 1u : 0u; + uint64_t sign_q = (is_signed && !overflow) ? 1u : 0u; + uint64_t sign_r = (is_signed && (rem >> 63)) ? 1u : 0u; + uint64_t n_sub_r = n - rem; // wrapping + uint64_t sign_n_sub_r = (is_signed && (n_sub_r >> 63)) ? 1u : 0u; + + // abs(x) for a two's-complement negative is 0 - x (mod 2^64); correct for + // i64::MIN too (matches Rust `unsigned_abs`). Non-negative passes through. + uint64_t abs_r = sign_r ? (0ull - rem) : rem; + uint64_t abs_d = sign_d ? (0ull - d) : d; + + for (int j = 0; j < 4; ++j) { + out[base + 0 + j] = (n >> (16 * j)) & 0xFFFFu; // N_0..3 (DWordHL) + out[base + 4 + j] = (d >> (16 * j)) & 0xFFFFu; // D_0..3 + out[base + 9 + j] = (q >> (16 * j)) & 0xFFFFu; // Q_0..3 + out[base + 13 + j] = (rem >> (16 * j)) & 0xFFFFu; // R_0..3 + out[base + 23 + j] = (n_sub_r >> (16 * j)) & 0xFFFFu; // N_SUB_R_0..3 + } + out[base + 8] = is_signed; // SIGNED + out[base + 17] = div_by_zero; // DIV_BY_ZERO + out[base + 18] = overflow; // OVERFLOW + out[base + 19] = abs_r & 0xFFFFFFFFu; // ABS_R_0 (DWordWL) + out[base + 20] = abs_r >> 32; // ABS_R_1 + out[base + 21] = abs_d & 0xFFFFFFFFu; // ABS_D_0 + out[base + 22] = abs_d >> 32; // ABS_D_1 + out[base + 27] = sign_n_sub_r; // SIGN_N_SUB_R + out[base + 28] = sign_n; // SIGN_N + out[base + 29] = sign_d; // SIGN_D + out[base + 30] = sign_q; // SIGN_Q + out[base + 31] = sign_r; // SIGN_R + out[base + 32] = mu_q; // MU_Q + out[base + 33] = mu_r; // MU_R +} + // On-GPU MEMW_R (register fast-path) fill: write the 10 MEMW_R columns ROW-MAJOR // from the host-walked rows. One thread per row (row_index is the identity for // the fill-from-walked-rows path). Columns mirror diff --git a/crypto/math-cuda/src/device.rs b/crypto/math-cuda/src/device.rs index 7a4ab90af..857a37bbe 100644 --- a/crypto/math-cuda/src/device.rs +++ b/crypto/math-cuda/src/device.rs @@ -196,6 +196,10 @@ pub struct Backend { pub store_fill: CudaFunction, pub shift_fill: CudaFunction, pub lt_fill: CudaFunction, + pub eq_fill: CudaFunction, + pub bytewise_fill: CudaFunction, + pub mul_fill: CudaFunction, + pub dvrm_fill: CudaFunction, pub memw_register_fill: CudaFunction, // Twiddle caches keyed by log_n. @@ -395,6 +399,10 @@ impl Backend { store_fill: trace_cpu.load_function("store_fill")?, shift_fill: trace_cpu.load_function("shift_fill")?, lt_fill: trace_cpu.load_function("lt_fill")?, + eq_fill: trace_cpu.load_function("eq_fill")?, + bytewise_fill: trace_cpu.load_function("bytewise_fill")?, + mul_fill: trace_cpu.load_function("mul_fill")?, + dvrm_fill: trace_cpu.load_function("dvrm_fill")?, memw_register_fill: trace_cpu.load_function("memw_register_fill")?, fwd_twiddles: Mutex::new(vec![None; max_log]), inv_twiddles: Mutex::new(vec![None; max_log]), diff --git a/crypto/math-cuda/src/trace_cpu.rs b/crypto/math-cuda/src/trace_cpu.rs index 1b5ca6420..ce960740e 100644 --- a/crypto/math-cuda/src/trace_cpu.rs +++ b/crypto/math-cuda/src/trace_cpu.rs @@ -152,6 +152,22 @@ pub const SHIFT_STRIDE: usize = 3; /// host-deduplicated: one unique op + summed multiplicity per row. pub const LT_NCOLS: usize = 17; pub const LT_STRIDE: usize = 4; +/// EQ table width / packed input stride (must match `eq_fill`). Like LT, input is +/// host-deduplicated: one unique op + summed multiplicity per row. +pub const EQ_NCOLS: usize = 12; +pub const EQ_STRIDE: usize = 4; +/// BYTEWISE table width / packed input stride (must match `bytewise_fill`). Like +/// LT, input is host-deduplicated: one unique op + summed multiplicity per row. +pub const BYTEWISE_NCOLS: usize = 26; +pub const BYTEWISE_STRIDE: usize = 4; +/// MUL table width / packed input stride (must match `mul_fill`). Input is +/// host-deduplicated: one unique op + split (mu_lo, mu_hi) multiplicities per row. +pub const MUL_NCOLS: usize = 26; +pub const MUL_STRIDE: usize = 5; +/// DVRM table width / packed input stride (must match `dvrm_fill`). Input is +/// host-deduplicated: one unique op + split (mu_q, mu_r) multiplicities per row. +pub const DVRM_NCOLS: usize = 34; +pub const DVRM_STRIDE: usize = 5; /// Shared interleaved fill on `stream`: upload `packed` (n × stride u64) and run /// `kernel(ops, n, num_rows, out)` into a zeroed row-major buffer. Returns the @@ -197,6 +213,18 @@ fn shift_kernel(be: &Backend) -> &CudaFunction { fn lt_kernel(be: &Backend) -> &CudaFunction { &be.lt_fill } +fn eq_kernel(be: &Backend) -> &CudaFunction { + &be.eq_fill +} +fn bytewise_kernel(be: &Backend) -> &CudaFunction { + &be.bytewise_fill +} +fn mul_kernel(be: &Backend) -> &CudaFunction { + &be.mul_fill +} +fn dvrm_kernel(be: &Backend) -> &CudaFunction { + &be.dvrm_fill +} /// Build one LT trace-table chunk on device from HOST-DEDUPLICATED ops (one unique /// op + summed multiplicity per row; see `lt_fill` in `trace_cpu.cu`). @@ -234,6 +262,170 @@ pub fn gpu_build_lt_trace_host(packed_ops: &[u64], n: usize, num_rows: usize) -> Ok(host) } +/// Build one EQ trace-table chunk on device from HOST-DEDUPLICATED ops (one unique +/// op + summed multiplicity per row; see `eq_fill` in `trace_cpu.cu`). +pub fn gpu_build_eq_trace(packed_ops: &[u64], n: usize, num_rows: usize) -> Result> { + let be = backend()?; + let stream = be.next_stream(); + let out = build_interleaved_on( + &stream, + eq_kernel(be), + packed_ops, + n, + num_rows, + EQ_NCOLS, + EQ_STRIDE, + )?; + stream.synchronize()?; + Ok(out) +} + +/// Host-returning EQ build for multiset-equality tests. +pub fn gpu_build_eq_trace_host(packed_ops: &[u64], n: usize, num_rows: usize) -> Result> { + let be = backend()?; + let stream = be.next_stream(); + let out = build_interleaved_on( + &stream, + eq_kernel(be), + packed_ops, + n, + num_rows, + EQ_NCOLS, + EQ_STRIDE, + )?; + let host = stream.clone_dtoh(&out)?; + stream.synchronize()?; + Ok(host) +} + +/// Build one BYTEWISE trace-table chunk on device from HOST-DEDUPLICATED ops (one +/// unique op + summed multiplicity per row; see `bytewise_fill` in `trace_cpu.cu`). +pub fn gpu_build_bytewise_trace( + packed_ops: &[u64], + n: usize, + num_rows: usize, +) -> Result> { + let be = backend()?; + let stream = be.next_stream(); + let out = build_interleaved_on( + &stream, + bytewise_kernel(be), + packed_ops, + n, + num_rows, + BYTEWISE_NCOLS, + BYTEWISE_STRIDE, + )?; + stream.synchronize()?; + Ok(out) +} + +/// Host-returning BYTEWISE build for multiset-equality tests. +pub fn gpu_build_bytewise_trace_host( + packed_ops: &[u64], + n: usize, + num_rows: usize, +) -> Result> { + let be = backend()?; + let stream = be.next_stream(); + let out = build_interleaved_on( + &stream, + bytewise_kernel(be), + packed_ops, + n, + num_rows, + BYTEWISE_NCOLS, + BYTEWISE_STRIDE, + )?; + let host = stream.clone_dtoh(&out)?; + stream.synchronize()?; + Ok(host) +} + +/// Build one MUL trace-table chunk on device from HOST-DEDUPLICATED ops (one unique +/// op + split mu_lo/mu_hi multiplicities per row; see `mul_fill` in `trace_cpu.cu`). +pub fn gpu_build_mul_trace( + packed_ops: &[u64], + n: usize, + num_rows: usize, +) -> Result> { + let be = backend()?; + let stream = be.next_stream(); + let out = build_interleaved_on( + &stream, + mul_kernel(be), + packed_ops, + n, + num_rows, + MUL_NCOLS, + MUL_STRIDE, + )?; + stream.synchronize()?; + Ok(out) +} + +/// Host-returning MUL build for multiset-equality tests. +pub fn gpu_build_mul_trace_host(packed_ops: &[u64], n: usize, num_rows: usize) -> Result> { + let be = backend()?; + let stream = be.next_stream(); + let out = build_interleaved_on( + &stream, + mul_kernel(be), + packed_ops, + n, + num_rows, + MUL_NCOLS, + MUL_STRIDE, + )?; + let host = stream.clone_dtoh(&out)?; + stream.synchronize()?; + Ok(host) +} + +/// Build one DVRM trace-table chunk on device from HOST-DEDUPLICATED ops (one unique +/// op + split mu_q/mu_r multiplicities per row; see `dvrm_fill` in `trace_cpu.cu`). +pub fn gpu_build_dvrm_trace( + packed_ops: &[u64], + n: usize, + num_rows: usize, +) -> Result> { + let be = backend()?; + let stream = be.next_stream(); + let out = build_interleaved_on( + &stream, + dvrm_kernel(be), + packed_ops, + n, + num_rows, + DVRM_NCOLS, + DVRM_STRIDE, + )?; + stream.synchronize()?; + Ok(out) +} + +/// Host-returning DVRM build for multiset-equality tests. +pub fn gpu_build_dvrm_trace_host( + packed_ops: &[u64], + n: usize, + num_rows: usize, +) -> Result> { + let be = backend()?; + let stream = be.next_stream(); + let out = build_interleaved_on( + &stream, + dvrm_kernel(be), + packed_ops, + n, + num_rows, + DVRM_NCOLS, + DVRM_STRIDE, + )?; + let host = stream.clone_dtoh(&out)?; + stream.synchronize()?; + Ok(host) +} + /// Build one SHIFT trace-table chunk on device (residency-ready row-major buffer). /// The kernel recomputes the shift aux from the packed inputs (see `trace_cpu.cu`). pub fn gpu_build_shift_trace( diff --git a/prover/src/tables/gpu_trace.rs b/prover/src/tables/gpu_trace.rs index e587423ee..c469372c0 100644 --- a/prover/src/tables/gpu_trace.rs +++ b/prover/src/tables/gpu_trace.rs @@ -14,12 +14,16 @@ use stark::trace::TraceTable; use std::collections::HashMap; +use super::bytewise::{self, BytewiseOperation}; use super::cpu::{self, CpuOperation}; +use super::dvrm::{self, DvrmOperation}; +use super::eq::{self, EqOperation}; use super::load::{self, LoadOperation}; use super::lt::{self, LtOperation}; use super::memw::MemwOperation; use super::memw_aligned; use super::memw_register::{self, RegRow}; +use super::mul::{self, MulOperation}; use super::shift::{self, ShiftOperation}; use super::store::{self, StoreOperation}; use super::types::{GoldilocksExtension, GoldilocksField}; @@ -486,3 +490,271 @@ pub(crate) fn gpu_build_lt_tables( } Some(tables) } + +// ============================================================================= +// EQ (ALU dedup table): host per-chunk HashMap dedup → device fill (compute) +// ============================================================================= + +/// Pack one unique `EqOperation` + its multiplicity into the `eq_fill` stride. +pub(crate) fn pack_eq_op(op: &EqOperation, mult: u64) -> [u64; math_cuda::trace_cpu::EQ_STRIDE] { + let flags = op.invert as u64; + [op.a, op.b, flags, mult] +} + +/// Build one EQ trace-table chunk on device. Dedup happens HERE on the host (the +/// same per-chunk HashMap `generate_eq_trace` uses), then the unique ops + summed +/// multiplicities are filled on device. EQ rides the permutation-invariant ALU +/// bus, so any row order is valid (validated by multiset/prove, not byte order). +fn build_eq_chunk( + chunk: &[EqOperation], +) -> Option> { + let mut map: HashMap = HashMap::new(); + for op in chunk { + *map.entry(op.clone()).or_insert(0) += 1; + } + let unique: Vec<(EqOperation, u64)> = map.into_iter().collect(); + let n = unique.len(); + let num_rows = n.next_power_of_two().max(4); + let mut packed = Vec::with_capacity(n * math_cuda::trace_cpu::EQ_STRIDE); + for (op, mult) in &unique { + packed.extend_from_slice(&pack_eq_op(op, *mult)); + } + let dev = math_cuda::trace_cpu::gpu_build_eq_trace(&packed, n, num_rows).ok()?; + let mut trace = TraceTable::new_main( + crate::tables::types::zeroed_fe_vec(num_rows * eq::cols::NUM_COLUMNS), + eq::cols::NUM_COLUMNS, + 1, + ); + trace.set_main_input_dev(Arc::new(dev)); + Some(trace) +} + +pub(crate) fn gpu_build_eq_tables( + ops: &[EqOperation], + max_rows: usize, +) -> Option>> { + if gpu_trace_disabled() { + return None; + } + // Chunk the RAW ops exactly like `chunk_and_generate`; each chunk dedups + // independently (matching `generate_eq_trace` per chunk). + let chunks: Vec<&[EqOperation]> = if ops.is_empty() { + vec![&[][..]] + } else { + ops.chunks(max_rows).collect() + }; + let mut tables = Vec::with_capacity(chunks.len()); + for chunk in chunks { + tables.push(build_eq_chunk(chunk)?); + } + Some(tables) +} + +// ============================================================================= +// BYTEWISE (ALU dedup table): host per-chunk HashMap dedup → device fill (compute) +// ============================================================================= + +/// Pack one unique `BytewiseOperation` + its multiplicity into the `bytewise_fill` +/// stride. +pub(crate) fn pack_bytewise_op( + op: &BytewiseOperation, + mult: u64, +) -> [u64; math_cuda::trace_cpu::BYTEWISE_STRIDE] { + [op.a, op.b, op.op as u64, mult] +} + +/// Build one BYTEWISE trace-table chunk on device. Dedup happens HERE on the host +/// (the same per-chunk HashMap `generate_bytewise_trace` uses), then the unique ops +/// with summed multiplicities are filled on device. BYTEWISE rides the +/// permutation-invariant ALU bus, so any row order is valid (validated by +/// multiset/prove, not byte order). +fn build_bytewise_chunk( + chunk: &[BytewiseOperation], +) -> Option> { + let mut map: HashMap = HashMap::new(); + for op in chunk { + *map.entry(op.clone()).or_insert(0) += 1; + } + let unique: Vec<(BytewiseOperation, u64)> = map.into_iter().collect(); + let n = unique.len(); + let num_rows = n.next_power_of_two().max(4); + let mut packed = Vec::with_capacity(n * math_cuda::trace_cpu::BYTEWISE_STRIDE); + for (op, mult) in &unique { + packed.extend_from_slice(&pack_bytewise_op(op, *mult)); + } + let dev = math_cuda::trace_cpu::gpu_build_bytewise_trace(&packed, n, num_rows).ok()?; + let mut trace = TraceTable::new_main( + crate::tables::types::zeroed_fe_vec(num_rows * bytewise::cols::NUM_COLUMNS), + bytewise::cols::NUM_COLUMNS, + 1, + ); + trace.set_main_input_dev(Arc::new(dev)); + Some(trace) +} + +pub(crate) fn gpu_build_bytewise_tables( + ops: &[BytewiseOperation], + max_rows: usize, +) -> Option>> { + if gpu_trace_disabled() { + return None; + } + // Chunk the RAW ops exactly like `chunk_and_generate`; each chunk dedups + // independently (matching `generate_bytewise_trace` per chunk). + let chunks: Vec<&[BytewiseOperation]> = if ops.is_empty() { + vec![&[][..]] + } else { + ops.chunks(max_rows).collect() + }; + let mut tables = Vec::with_capacity(chunks.len()); + for chunk in chunks { + tables.push(build_bytewise_chunk(chunk)?); + } + Some(tables) +} + +// ============================================================================= +// MUL (ALU dedup table): host per-chunk HashMap dedup → device fill (128-bit +// product + sign-extended convolution recomputed on device) +// ============================================================================= + +/// Pack one unique `MulOperation` + its split multiplicities into the `mul_fill` +/// stride. +pub(crate) fn pack_mul_op( + op: &MulOperation, + mu_lo: u64, + mu_hi: u64, +) -> [u64; math_cuda::trace_cpu::MUL_STRIDE] { + let flags = (op.lhs_signed as u64) | ((op.rhs_signed as u64) << 1); + [op.lhs, op.rhs, flags, mu_lo, mu_hi] +} + +/// Build one MUL trace-table chunk on device. Dedup happens HERE on the host (the +/// same per-chunk HashMap `generate_mul_trace` uses, keyed by op with mu_lo/mu_hi +/// accumulated from each `wants_hi`), then the unique ops + both multiplicities are +/// filled on device (the kernel recomputes the 128-bit product and raw_products). +/// MUL rides the permutation-invariant ALU bus, so any row order is valid +/// (validated by multiset/prove, not byte order). +fn build_mul_chunk( + chunk: &[(MulOperation, bool)], +) -> Option> { + // (mu_lo, mu_hi) per unique op, matching `MulMultiplicities`. + let mut map: HashMap = HashMap::new(); + for (op, wants_hi) in chunk { + let e = map.entry(op.clone()).or_insert((0, 0)); + if *wants_hi { + e.1 += 1; + } else { + e.0 += 1; + } + } + let unique: Vec<(MulOperation, (u64, u64))> = map.into_iter().collect(); + let n = unique.len(); + let num_rows = n.next_power_of_two().max(4); + let mut packed = Vec::with_capacity(n * math_cuda::trace_cpu::MUL_STRIDE); + for (op, (mu_lo, mu_hi)) in &unique { + packed.extend_from_slice(&pack_mul_op(op, *mu_lo, *mu_hi)); + } + let dev = math_cuda::trace_cpu::gpu_build_mul_trace(&packed, n, num_rows).ok()?; + let mut trace = TraceTable::new_main( + crate::tables::types::zeroed_fe_vec(num_rows * mul::cols::NUM_COLUMNS), + mul::cols::NUM_COLUMNS, + 1, + ); + trace.set_main_input_dev(Arc::new(dev)); + Some(trace) +} + +pub(crate) fn gpu_build_mul_tables( + ops: &[(MulOperation, bool)], + max_rows: usize, +) -> Option>> { + if gpu_trace_disabled() { + return None; + } + // Chunk the RAW (op, wants_hi) pairs exactly like `chunk_and_generate`; each + // chunk dedups independently (matching `generate_mul_trace` per chunk). + let chunks: Vec<&[(MulOperation, bool)]> = if ops.is_empty() { + vec![&[][..]] + } else { + ops.chunks(max_rows).collect() + }; + let mut tables = Vec::with_capacity(chunks.len()); + for chunk in chunks { + tables.push(build_mul_chunk(chunk)?); + } + Some(tables) +} + +// ============================================================================= +// DVRM (ALU dedup table): host per-chunk HashMap dedup → device fill (signed/ +// unsigned division & remainder + abs/sign aux recomputed on device) +// ============================================================================= + +/// Pack one unique `DvrmOperation` + its split multiplicities into the `dvrm_fill` +/// stride. +pub(crate) fn pack_dvrm_op( + op: &DvrmOperation, + mu_q: u64, + mu_r: u64, +) -> [u64; math_cuda::trace_cpu::DVRM_STRIDE] { + let flags = op.signed as u64; + [op.n, op.d, flags, mu_q, mu_r] +} + +/// Build one DVRM trace-table chunk on device. Dedup happens HERE on the host (the +/// same per-chunk HashMap `generate_dvrm_trace` uses, keyed by op with mu_q/mu_r +/// accumulated from each `wants_remainder`), then the unique ops + both +/// multiplicities are filled on device (the kernel recomputes the quotient, +/// remainder, and abs/sign aux). DVRM rides the permutation-invariant ALU bus, so +/// any row order is valid (validated by multiset/prove, not byte order). +fn build_dvrm_chunk( + chunk: &[(DvrmOperation, bool)], +) -> Option> { + // (mu_q, mu_r) per unique op, matching `DvrmMultiplicities`. + let mut map: HashMap = HashMap::new(); + for (op, wants_remainder) in chunk { + let e = map.entry(op.clone()).or_insert((0, 0)); + if *wants_remainder { + e.1 += 1; + } else { + e.0 += 1; + } + } + let unique: Vec<(DvrmOperation, (u64, u64))> = map.into_iter().collect(); + let n = unique.len(); + let num_rows = n.next_power_of_two().max(4); + let mut packed = Vec::with_capacity(n * math_cuda::trace_cpu::DVRM_STRIDE); + for (op, (mu_q, mu_r)) in &unique { + packed.extend_from_slice(&pack_dvrm_op(op, *mu_q, *mu_r)); + } + let dev = math_cuda::trace_cpu::gpu_build_dvrm_trace(&packed, n, num_rows).ok()?; + let mut trace = TraceTable::new_main( + crate::tables::types::zeroed_fe_vec(num_rows * dvrm::cols::NUM_COLUMNS), + dvrm::cols::NUM_COLUMNS, + 1, + ); + trace.set_main_input_dev(Arc::new(dev)); + Some(trace) +} + +pub(crate) fn gpu_build_dvrm_tables( + ops: &[(DvrmOperation, bool)], + max_rows: usize, +) -> Option>> { + if gpu_trace_disabled() { + return None; + } + // Chunk the RAW (op, wants_remainder) pairs exactly like `chunk_and_generate`; + // each chunk dedups independently (matching `generate_dvrm_trace` per chunk). + let chunks: Vec<&[(DvrmOperation, bool)]> = if ops.is_empty() { + vec![&[][..]] + } else { + ops.chunks(max_rows).collect() + }; + let mut tables = Vec::with_capacity(chunks.len()); + for chunk in chunks { + tables.push(build_dvrm_chunk(chunk)?); + } + Some(tables) +} diff --git a/prover/src/tables/trace_builder.rs b/prover/src/tables/trace_builder.rs index 7b90d6fa0..b2227ae87 100644 --- a/prover/src/tables/trace_builder.rs +++ b/prover/src/tables/trace_builder.rs @@ -3382,6 +3382,25 @@ fn build_traces( ) }; let gen_muls = || { + #[cfg(feature = "cuda")] + { + let use_gpu = { + #[cfg(feature = "disk-spill")] + { + storage_mode != StorageMode::Disk + } + #[cfg(not(feature = "disk-spill"))] + { + true + } + }; + if use_gpu + && let Some(tables) = + crate::tables::gpu_trace::gpu_build_mul_tables(&mul_ops, max_rows.mul) + { + return Ok(tables); + } + } chunk_and_generate( &mul_ops, max_rows.mul, @@ -3391,6 +3410,25 @@ fn build_traces( ) }; let gen_dvrms = || { + #[cfg(feature = "cuda")] + { + let use_gpu = { + #[cfg(feature = "disk-spill")] + { + storage_mode != StorageMode::Disk + } + #[cfg(not(feature = "disk-spill"))] + { + true + } + }; + if use_gpu + && let Some(tables) = + crate::tables::gpu_trace::gpu_build_dvrm_tables(&dvrm_ops, max_rows.dvrm) + { + return Ok(tables); + } + } chunk_and_generate( &dvrm_ops, max_rows.dvrm, @@ -3412,6 +3450,25 @@ fn build_traces( // dispatch, so they are generated empty — one padded (μ=0) chunk each, which // contributes nothing to any bus. let gen_eqs = || { + #[cfg(feature = "cuda")] + { + let use_gpu = { + #[cfg(feature = "disk-spill")] + { + storage_mode != StorageMode::Disk + } + #[cfg(not(feature = "disk-spill"))] + { + true + } + }; + if use_gpu + && let Some(tables) = + crate::tables::gpu_trace::gpu_build_eq_tables(&eq_ops, max_rows.eq) + { + return Ok(tables); + } + } chunk_and_generate::( &eq_ops, max_rows.eq, @@ -3421,6 +3478,27 @@ fn build_traces( ) }; let gen_bytewises = || { + #[cfg(feature = "cuda")] + { + let use_gpu = { + #[cfg(feature = "disk-spill")] + { + storage_mode != StorageMode::Disk + } + #[cfg(not(feature = "disk-spill"))] + { + true + } + }; + if use_gpu + && let Some(tables) = crate::tables::gpu_trace::gpu_build_bytewise_tables( + &bytewise_ops, + max_rows.bytewise, + ) + { + return Ok(tables); + } + } chunk_and_generate::( &bytewise_ops, max_rows.bytewise, @@ -4999,4 +5077,328 @@ mod gpu_fill_tests { "device LT rows must match the CPU fill as a multiset" ); } + + /// EQ device fill: like LT, dedup rides the permutation-invariant ALU bus, so + /// the row ORDER is non-deterministic (HashMap iteration). Validate as a + /// MULTISET — the real rows (μ>0), incl. summed multiplicities, must match the + /// CPU `generate_eq_trace`. Covers a==b and a!=b, invert, high words, and + /// duplicates (μ>1). Skips cleanly with no GPU. + #[cfg(feature = "cuda")] + #[test] + fn gpu_eq_fill_matches_cpu_multiset() { + if math_cuda::device::backend().is_err() { + eprintln!("skipping gpu_eq_fill_matches_cpu_multiset: no CUDA backend"); + return; + } + // Raw ops with equal/unequal operands, high words, and deliberate duplicates. + let mut raw = Vec::new(); + for i in 0..800u64 { + let a = i.wrapping_mul(0x9E37_79B9_7F4A_7C15) ^ (i << 40); + // Every 5th op has a == b to exercise the eq=1 path. + let b = if i % 5 == 0 { + a + } else { + i.wrapping_mul(0x1234_5678_9ABC_DEF1).rotate_left(11) + }; + let invert = i % 3 == 0; + let op = eq::EqOperation::new(a, b, invert); + raw.push(op.clone()); + if i % 4 == 0 { + raw.push(op); // duplicate → multiplicity 2 + } + } + + let ncols = math_cuda::trace_cpu::EQ_NCOLS; + // Extract the real (μ>0) rows of a row-major u64 buffer as a sorted multiset. + let real_rows = |flat: &[u64]| -> Vec> { + let mut rows: Vec> = flat + .chunks(ncols) + .filter(|row| row[eq::cols::MU] > 0) + .map(|row| row.to_vec()) + .collect(); + rows.sort(); + rows + }; + + let cpu_table = eq::generate_eq_trace(&raw); + let (cpu_fe, w) = cpu_table.main_data_row_major(); + assert_eq!(w, ncols); + let cpu_u64: Vec = cpu_fe + .iter() + .map(|e| unsafe { *(e.value() as *const u64) }) + .collect(); + + // Dedup on the host exactly like `gpu_build_eq_tables`, then device-fill. + let mut map: std::collections::HashMap = HashMap::new(); + for op in &raw { + *map.entry(op.clone()).or_insert(0) += 1; + } + let unique: Vec<(eq::EqOperation, u64)> = map.into_iter().collect(); + let n = unique.len(); + let num_rows = n.next_power_of_two().max(4); + let mut packed = Vec::with_capacity(n * math_cuda::trace_cpu::EQ_STRIDE); + for (op, mult) in &unique { + packed.extend_from_slice(&crate::tables::gpu_trace::pack_eq_op(op, *mult)); + } + let gpu_u64 = math_cuda::trace_cpu::gpu_build_eq_trace_host(&packed, n, num_rows) + .expect("device EQ build must run on a box with a CUDA backend"); + + assert_eq!( + real_rows(&gpu_u64), + real_rows(&cpu_u64), + "device EQ rows must match the CPU fill as a multiset" + ); + } + + /// BYTEWISE device fill: like LT/EQ, dedup rides the permutation-invariant ALU + /// bus, so the row ORDER is non-deterministic. Validate as a MULTISET — the real + /// rows (μ>0), incl. summed multiplicities, must match `generate_bytewise_trace`. + /// Covers AND/OR/XOR, full 64-bit operands, and duplicates (μ>1). Skips cleanly + /// with no GPU. + #[cfg(feature = "cuda")] + #[test] + fn gpu_bytewise_fill_matches_cpu_multiset() { + use crate::tables::types::alu_op; + if math_cuda::device::backend().is_err() { + eprintln!("skipping gpu_bytewise_fill_matches_cpu_multiset: no CUDA backend"); + return; + } + // Raw ops cycling AND/OR/XOR, full-word operands, with deliberate duplicates. + let mut raw = Vec::new(); + for i in 0..900u64 { + let a = i.wrapping_mul(0x9E37_79B9_7F4A_7C15) ^ (i << 33); + let b = i.wrapping_mul(0x1234_5678_9ABC_DEF1).rotate_left(19); + let op = [alu_op::AND, alu_op::OR, alu_op::XOR][(i % 3) as usize]; + let bw = bytewise::BytewiseOperation::new(a, b, op); + raw.push(bw.clone()); + if i % 4 == 0 { + raw.push(bw); // duplicate → multiplicity 2 + } + } + + let ncols = math_cuda::trace_cpu::BYTEWISE_NCOLS; + // Extract the real (μ>0) rows of a row-major u64 buffer as a sorted multiset. + let real_rows = |flat: &[u64]| -> Vec> { + let mut rows: Vec> = flat + .chunks(ncols) + .filter(|row| row[bytewise::cols::MU] > 0) + .map(|row| row.to_vec()) + .collect(); + rows.sort(); + rows + }; + + let cpu_table = bytewise::generate_bytewise_trace(&raw); + let (cpu_fe, w) = cpu_table.main_data_row_major(); + assert_eq!(w, ncols); + let cpu_u64: Vec = cpu_fe + .iter() + .map(|e| unsafe { *(e.value() as *const u64) }) + .collect(); + + // Dedup on the host exactly like `gpu_build_bytewise_tables`, then device-fill. + let mut map: std::collections::HashMap = HashMap::new(); + for op in &raw { + *map.entry(op.clone()).or_insert(0) += 1; + } + let unique: Vec<(bytewise::BytewiseOperation, u64)> = map.into_iter().collect(); + let n = unique.len(); + let num_rows = n.next_power_of_two().max(4); + let mut packed = Vec::with_capacity(n * math_cuda::trace_cpu::BYTEWISE_STRIDE); + for (op, mult) in &unique { + packed.extend_from_slice(&crate::tables::gpu_trace::pack_bytewise_op(op, *mult)); + } + let gpu_u64 = math_cuda::trace_cpu::gpu_build_bytewise_trace_host(&packed, n, num_rows) + .expect("device BYTEWISE build must run on a box with a CUDA backend"); + + assert_eq!( + real_rows(&gpu_u64), + real_rows(&cpu_u64), + "device BYTEWISE rows must match the CPU fill as a multiset" + ); + } + + /// MUL device fill: like the other ALU tables, dedup rides the + /// permutation-invariant ALU bus, so the row ORDER is non-deterministic. + /// Validate as a MULTISET — the real rows (mu_lo>0 or mu_hi>0), incl. the split + /// multiplicities, must match `generate_mul_trace`. Covers all four + /// signed/unsigned combos, negative operands, the 128-bit product + raw_product + /// convolution, and lo/hi (wants_hi) requests with duplicates. Skips cleanly + /// with no GPU. + #[cfg(feature = "cuda")] + #[test] + fn gpu_mul_fill_matches_cpu_multiset() { + if math_cuda::device::backend().is_err() { + eprintln!("skipping gpu_mul_fill_matches_cpu_multiset: no CUDA backend"); + return; + } + // Raw (op, wants_hi) pairs: varied operands, all sign combos, both lo and hi + // requests, with deliberate duplicates so mu_lo/mu_hi accumulate. + let mut raw = Vec::new(); + for i in 0..800u64 { + let lhs = i.wrapping_mul(0x9E37_79B9_7F4A_7C15) ^ (i << 40); + let rhs = i.wrapping_mul(0x1234_5678_9ABC_DEF1).rotate_left(23); + let lhs_signed = i % 2 == 0; + let rhs_signed = i % 3 == 0; + let op = mul::MulOperation::new(lhs, lhs_signed, rhs, rhs_signed); + raw.push((op.clone(), i % 2 == 0)); + if i % 3 == 0 { + raw.push((op.clone(), true)); // extra hi request + } + if i % 5 == 0 { + raw.push((op, false)); // extra lo request (duplicate) + } + } + // Explicit sign edge cases: i64::MIN, -1, treated signed and unsigned. + for &(a, b) in &[ + (0x8000_0000_0000_0000u64, 0xFFFF_FFFF_FFFF_FFFFu64), + (0xFFFF_FFFF_FFFF_FFFFu64, 0x8000_0000_0000_0000u64), + ] { + raw.push((mul::MulOperation::new(a, true, b, true), false)); + raw.push((mul::MulOperation::new(a, false, b, false), true)); + } + + let ncols = math_cuda::trace_cpu::MUL_NCOLS; + // Real rows: mu_lo>0 or mu_hi>0. + let real_rows = |flat: &[u64]| -> Vec> { + let mut rows: Vec> = flat + .chunks(ncols) + .filter(|row| row[mul::cols::MU_LO] > 0 || row[mul::cols::MU_HI] > 0) + .map(|row| row.to_vec()) + .collect(); + rows.sort(); + rows + }; + + let cpu_table = mul::generate_mul_trace(&raw); + let (cpu_fe, w) = cpu_table.main_data_row_major(); + assert_eq!(w, ncols); + let cpu_u64: Vec = cpu_fe + .iter() + .map(|e| unsafe { *(e.value() as *const u64) }) + .collect(); + + // Dedup on the host exactly like `gpu_build_mul_tables`, then device-fill. + let mut map: std::collections::HashMap = HashMap::new(); + for (op, wants_hi) in &raw { + let e = map.entry(op.clone()).or_insert((0, 0)); + if *wants_hi { + e.1 += 1; + } else { + e.0 += 1; + } + } + let unique: Vec<(mul::MulOperation, (u64, u64))> = map.into_iter().collect(); + let n = unique.len(); + let num_rows = n.next_power_of_two().max(4); + let mut packed = Vec::with_capacity(n * math_cuda::trace_cpu::MUL_STRIDE); + for (op, (mu_lo, mu_hi)) in &unique { + packed.extend_from_slice(&crate::tables::gpu_trace::pack_mul_op(op, *mu_lo, *mu_hi)); + } + let gpu_u64 = math_cuda::trace_cpu::gpu_build_mul_trace_host(&packed, n, num_rows) + .expect("device MUL build must run on a box with a CUDA backend"); + + assert_eq!( + real_rows(&gpu_u64), + real_rows(&cpu_u64), + "device MUL rows must match the CPU fill as a multiset" + ); + } + + /// DVRM device fill: like the other ALU tables, dedup rides the + /// permutation-invariant ALU bus, so the row ORDER is non-deterministic. + /// Validate as a MULTISET — the real rows (mu_q>0 or mu_r>0), incl. the split + /// multiplicities, must match `generate_dvrm_trace`. Covers signed/unsigned, + /// division-by-zero, the MIN/-1 signed overflow, negative operands, and + /// quotient/remainder (wants_remainder) requests with duplicates. Skips cleanly + /// with no GPU. + #[cfg(feature = "cuda")] + #[test] + fn gpu_dvrm_fill_matches_cpu_multiset() { + if math_cuda::device::backend().is_err() { + eprintln!("skipping gpu_dvrm_fill_matches_cpu_multiset: no CUDA backend"); + return; + } + // Raw (op, wants_remainder) pairs: varied operands (incl. periodic d==0), + // both signednesses, q and r requests, with deliberate duplicates. + let mut raw = Vec::new(); + for i in 0..800u64 { + let n = i.wrapping_mul(0x9E37_79B9_7F4A_7C15) ^ (i << 40); + let d = if i % 11 == 0 { + 0 + } else { + i.wrapping_mul(0x1234_5678_9ABC_DEF1).rotate_left(29) + }; + let signed = i % 2 == 0; + let op = dvrm::DvrmOperation::new(n, d, signed); + raw.push((op.clone(), i % 2 == 0)); + if i % 3 == 0 { + raw.push((op.clone(), true)); // extra remainder request + } + if i % 5 == 0 { + raw.push((op, false)); // extra quotient request (duplicate) + } + } + // Explicit edge cases: signed overflow (MIN/-1), div-by-zero, MIN numerator, + // -1 denominator. + let min = 0x8000_0000_0000_0000u64; + let neg1 = 0xFFFF_FFFF_FFFF_FFFFu64; + for &(n, d, s) in &[ + (min, neg1, true), // signed overflow + (min, neg1, false), // unsigned: /(2^64-1), no overflow + (123u64, 0u64, true), // div-by-zero, signed + (123u64, 0u64, false), // div-by-zero, unsigned + (min, 7u64, true), // negative numerator + (100u64, neg1, true), // negative denominator (n != MIN) + ] { + raw.push((dvrm::DvrmOperation::new(n, d, s), false)); + raw.push((dvrm::DvrmOperation::new(n, d, s), true)); + } + + let ncols = math_cuda::trace_cpu::DVRM_NCOLS; + // Real rows: mu_q>0 or mu_r>0. + let real_rows = |flat: &[u64]| -> Vec> { + let mut rows: Vec> = flat + .chunks(ncols) + .filter(|row| row[dvrm::cols::MU_Q] > 0 || row[dvrm::cols::MU_R] > 0) + .map(|row| row.to_vec()) + .collect(); + rows.sort(); + rows + }; + + let cpu_table = dvrm::generate_dvrm_trace(&raw); + let (cpu_fe, w) = cpu_table.main_data_row_major(); + assert_eq!(w, ncols); + let cpu_u64: Vec = cpu_fe + .iter() + .map(|e| unsafe { *(e.value() as *const u64) }) + .collect(); + + // Dedup on the host exactly like `gpu_build_dvrm_tables`, then device-fill. + let mut map: std::collections::HashMap = HashMap::new(); + for (op, wants_remainder) in &raw { + let e = map.entry(op.clone()).or_insert((0, 0)); + if *wants_remainder { + e.1 += 1; + } else { + e.0 += 1; + } + } + let unique: Vec<(dvrm::DvrmOperation, (u64, u64))> = map.into_iter().collect(); + let n = unique.len(); + let num_rows = n.next_power_of_two().max(4); + let mut packed = Vec::with_capacity(n * math_cuda::trace_cpu::DVRM_STRIDE); + for (op, (mu_q, mu_r)) in &unique { + packed.extend_from_slice(&crate::tables::gpu_trace::pack_dvrm_op(op, *mu_q, *mu_r)); + } + let gpu_u64 = math_cuda::trace_cpu::gpu_build_dvrm_trace_host(&packed, n, num_rows) + .expect("device DVRM build must run on a box with a CUDA backend"); + + assert_eq!( + real_rows(&gpu_u64), + real_rows(&cpu_u64), + "device DVRM rows must match the CPU fill as a multiset" + ); + } } From a31f153c005c1dd57a0f01c57849030e98351ee2 Mon Sep 17 00:00:00 2001 From: Joaquin Carletti Date: Mon, 13 Jul 2026 17:24:08 -0300 Subject: [PATCH 2/3] move tables to gpu --- crypto/math-cuda/kernels/trace_cpu.cu | 218 ++++++++++++++++++++ crypto/math-cuda/src/device.rs | 6 + crypto/math-cuda/src/lde.rs | 69 +++++++ crypto/math-cuda/src/trace_cpu.rs | 153 ++++++++++++++ crypto/stark/src/gpu_lde.rs | 57 ++++++ crypto/stark/src/prover.rs | 56 ++++-- prover/src/tables/gpu_trace.rs | 216 +++++++++++++++++++- prover/src/tables/trace_builder.rs | 276 ++++++++++++++++++++++++++ 8 files changed, 1038 insertions(+), 13 deletions(-) diff --git a/crypto/math-cuda/kernels/trace_cpu.cu b/crypto/math-cuda/kernels/trace_cpu.cu index e5c236465..6e12a0a03 100644 --- a/crypto/math-cuda/kernels/trace_cpu.cu +++ b/crypto/math-cuda/kernels/trace_cpu.cu @@ -732,6 +732,224 @@ extern "C" __global__ void dvrm_fill(const uint64_t *ops, uint64_t n_ops, out[base + 33] = mu_r; // MU_R } +// On-GPU BRANCH trace fill (14 cols). Mirrors the per-row compute of +// `prover/src/tables/branch.rs::generate_branch_trace`: next_pc = (base + offset) +// & ~1, where base = jalr ? register : pc, split into 3 high halfwords + 2 low +// bytes (LSB masked) plus the unmasked low byte. Dedup is done on the HOST (the +// same per-chunk HashMap the CPU uses); this kernel receives unique ops + summed +// multiplicity, one per row. Order-independent (LogUp lookup bus) → validated by +// multiset/prove, not byte order. Padding rows (r >= n) stay all-zero. +// +// Packed stride BRANCH_STRIDE: [0]=pc [1]=offset [2]=register [3]=flags (bit0 +// jalr) [4]=multiplicity. +#define BRANCH_NCOLS 14u +#define BRANCH_STRIDE 5u + +extern "C" __global__ void branch_fill(const uint64_t *ops, uint64_t n, + uint64_t num_rows, uint64_t *out) { + uint64_t r = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x; + if (r >= num_rows) + return; + uint64_t base = r * BRANCH_NCOLS; + if (r >= n) + return; + + const uint64_t *op = ops + r * BRANCH_STRIDE; + uint64_t pc = op[0]; + uint64_t offset = op[1]; + uint64_t reg = op[2]; + uint64_t fl = op[3]; + uint64_t mult = op[4]; + uint64_t jalr = fl & 1u; + + uint64_t b = jalr ? reg : pc; + uint64_t unmasked = b + offset; // wrapping + uint64_t next_pc = unmasked & ~1ull; + + out[base + 0] = pc & 0xFFFFFFFFu; // PC_0 (DWordWL) + out[base + 1] = pc >> 32; // PC_1 + out[base + 2] = offset & 0xFFFFFFFFu; // OFFSET_0 + out[base + 3] = offset >> 32; // OFFSET_1 + out[base + 4] = reg & 0xFFFFFFFFu; // REGISTER_0 + out[base + 5] = reg >> 32; // REGISTER_1 + out[base + 6] = jalr; // JALR + out[base + 7] = (next_pc >> 16) & 0xFFFFu; // NEXT_PC_HIGH_0 (halves) + out[base + 8] = (next_pc >> 32) & 0xFFFFu; // NEXT_PC_HIGH_1 + out[base + 9] = (next_pc >> 48) & 0xFFFFu; // NEXT_PC_HIGH_2 + out[base + 10] = next_pc & 0xFFu; // NEXT_PC_LOW_0 (masked LSB) + out[base + 11] = (next_pc >> 8) & 0xFFu; // NEXT_PC_LOW_1 + out[base + 12] = unmasked & 0xFFu; // UNMASKED_LOW_BYTE + out[base + 13] = mult; // MU +} + +// On-GPU CPU32 trace fill (38 cols). Mirrors the per-row compute of +// `prover/src/tables/cpu32.rs::{generate_cpu32_trace, compute_aux}` — the delegated +// `*W` (32-bit word) instructions. Sign/zero-extends rv1/rv2 to arg1/arg2 and +// sign-extends the 32-bit result to rvd, per RV64 `*W` semantics. Per-row (μ=1, no +// dedup) → byte-identical to the CPU fill. Padding rows (r >= n) stay all-zero. +// +// Packed stride CPU32_STRIDE: [0]=timestamp [1]=pc [2]=rv1 [3]=rv2 [4]=imm [5]=res +// [6]=flags (bit0 read_register1, bit1 read_register2, bit2 write_register, +// bit3 alu, bit4 add, bit5 sub) [7]=bytes (b0 rs1, b1 rs2, b2 rd, b3 alu_flags, +// b4 half_instruction_length). +#define CPU32_NCOLS 38u +#define CPU32_STRIDE 8u + +extern "C" __global__ void cpu32_fill(const uint64_t *ops, uint64_t n, + uint64_t num_rows, uint64_t *out) { + uint64_t r = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x; + if (r >= num_rows) + return; + uint64_t base = r * CPU32_NCOLS; + if (r >= n) + return; + + const uint64_t *op = ops + r * CPU32_STRIDE; + uint64_t ts = op[0]; + uint64_t pc = op[1]; + uint64_t rv1 = op[2]; + uint64_t rv2 = op[3]; + uint64_t imm = op[4]; + uint64_t res = op[5]; + uint64_t fl = op[6]; + uint64_t by = op[7]; + + uint64_t rr1 = fl & 1u; + uint64_t rr2 = (fl >> 1) & 1u; + uint64_t wr = (fl >> 2) & 1u; + uint64_t alu = (fl >> 3) & 1u; + uint64_t add = (fl >> 4) & 1u; + uint64_t sub = (fl >> 5) & 1u; + + uint64_t rs1 = by & 0xFFu; + uint64_t rs2 = (by >> 8) & 0xFFu; + uint64_t rd = (by >> 16) & 0xFFu; + uint64_t aluf = (by >> 24) & 0xFFu; + uint64_t hil = (by >> 32) & 0xFFu; + + // signed = alu_flags bit 5 (ALU_FLAGS_SIGNED). rv1/rv2 sign bits are gated by + // `signed`; the result is always sign-extended for *W. + uint64_t is_signed = (aluf >> 5) & 1u; + uint64_t rv1_sign = (is_signed && ((rv1 >> 31) & 1u)) ? 1u : 0u; + uint64_t rv2_sign = (is_signed && ((rv2 >> 31) & 1u)) ? 1u : 0u; + uint64_t res_sign = ((res >> 31) & 1u) ? 1u : 0u; + + uint64_t arg1_hi = rv1_sign ? 0xFFFFFFFFull : 0ull; + uint64_t arg1 = (rv1 & 0xFFFFFFFFull) | (arg1_hi << 32); + + // By the decoding assumption exactly one of rv2 / imm is nonzero, so the + // per-word sums do not overflow; the masks mirror the Rust exactly regardless. + uint64_t arg2_lo = (rv2 & 0xFFFFFFFFull) + (imm & 0xFFFFFFFFull); + uint64_t arg2_hi = (rv2_sign ? 0xFFFFFFFFull : 0ull) + (imm >> 32); + uint64_t arg2 = (arg2_lo & 0xFFFFFFFFull) | (arg2_hi << 32); + + uint64_t rvd_hi = res_sign ? 0xFFFFFFFFull : 0ull; + uint64_t rvd = (res & 0xFFFFFFFFull) | (rvd_hi << 32); + + out[base + 0] = ts & 0xFFFFFFFFu; // TIMESTAMP_0 (DWordWL) + out[base + 1] = ts >> 32; // TIMESTAMP_1 + out[base + 2] = pc & 0xFFFFFFFFu; // PC_0 + out[base + 3] = pc >> 32; // PC_1 + out[base + 4] = rs1; // RS1 + out[base + 5] = rr1; // READ_REGISTER1 + out[base + 6] = rv1 & 0xFFFFu; // RV1_0 (DWordWHH: half) + out[base + 7] = (rv1 >> 16) & 0xFFFFu; // RV1_1 (half) + out[base + 8] = (rv1 >> 32) & 0xFFFFFFFFu; // RV1_2 (word) + out[base + 9] = rv1_sign; // RV1_SIGN + out[base + 10] = arg1 & 0xFFFFFFFFu; // ARG1_0 (DWordWL) + out[base + 11] = arg1 >> 32; // ARG1_1 + out[base + 12] = rs2; // RS2 + out[base + 13] = rr2; // READ_REGISTER2 + out[base + 14] = rv2 & 0xFFFFu; // RV2_0 + out[base + 15] = (rv2 >> 16) & 0xFFFFu; // RV2_1 + out[base + 16] = (rv2 >> 32) & 0xFFFFFFFFu; // RV2_2 + out[base + 17] = rv2_sign; // RV2_SIGN + out[base + 18] = imm & 0xFFFFFFFFu; // IMM_0 + out[base + 19] = imm >> 32; // IMM_1 + out[base + 20] = arg2 & 0xFFFFFFFFu; // ARG2_0 + out[base + 21] = arg2 >> 32; // ARG2_1 + out[base + 22] = res & 0xFFFFu; // RES_0 (DWordHL: 4 halves) + out[base + 23] = (res >> 16) & 0xFFFFu; // RES_1 + out[base + 24] = (res >> 32) & 0xFFFFu; // RES_2 + out[base + 25] = (res >> 48) & 0xFFFFu; // RES_3 + out[base + 26] = res_sign; // RES_SIGN + out[base + 27] = rd; // RD + out[base + 28] = wr; // WRITE_REGISTER + out[base + 29] = rvd & 0xFFFFFFFFu; // RVD_0 + out[base + 30] = rvd >> 32; // RVD_1 + out[base + 31] = alu; // ALU + out[base + 32] = aluf; // ALU_FLAGS + out[base + 33] = add; // ADD + out[base + 34] = sub; // SUB + out[base + 35] = hil; // HALF_INSTRUCTION_LENGTH + out[base + 36] = is_signed; // SIGNED + out[base + 37] = 1u; // MU (active row) +} + +// On-GPU MEMW (general / unaligned / split-timestamp) trace fill (49 cols). Mirrors +// `prover/src/tables/memw.rs::generate_memw_trace`. The op is already walked +// (old/old_timestamp filled), so this is bit-slicing + the carry[i] aux +// (base_addr_lo + (i+1) >= 2^32). Per-row (no dedup) → byte-identical to the CPU +// fill. Padding rows (r >= n) stay all-zero. +// +// Packed stride MEMW_STRIDE: [0]=flags (bit0 is_register, bit1 is_read, bits8..16 +// width) [1]=base_address [2]=timestamp [3..7]=value[0..8] packed 2×u32/u64 +// [7..11]=old[0..8] packed 2×u32/u64 [11..19]=old_timestamp[0..8]. +#define MEMW_NCOLS 49u +#define MEMW_STRIDE 19u + +extern "C" __global__ void memw_fill(const uint64_t *ops, uint64_t n, + uint64_t num_rows, uint64_t *out) { + uint64_t r = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x; + if (r >= num_rows) + return; + uint64_t base = r * MEMW_NCOLS; + if (r >= n) + return; + + const uint64_t *op = ops + r * MEMW_STRIDE; + uint64_t fl = op[0]; + uint64_t is_register = fl & 1u; + uint64_t is_read = (fl >> 1) & 1u; + uint64_t width = (fl >> 8) & 0xFFu; + uint64_t addr = op[1]; + uint64_t ts = op[2]; + + out[base + 0] = is_register; // IS_REGISTER + out[base + 1] = addr & 0xFFFFFFFFu; // BASE_ADDRESS_0 (DWordWL) + out[base + 2] = addr >> 32; // BASE_ADDRESS_1 + // VALUE[0..8] at cols 3..11 (each column holds one full value limb). + for (int i = 0; i < 4; ++i) { + uint64_t p = op[3 + i]; + out[base + 3 + 2 * i] = p & 0xFFFFFFFFu; + out[base + 3 + 2 * i + 1] = p >> 32; + } + out[base + 11] = ts & 0xFFFFFFFFu; // TIMESTAMP_0 + out[base + 12] = ts >> 32; // TIMESTAMP_1 + out[base + 13] = (width == 2u) ? 1u : 0u; // WRITE2 + out[base + 14] = (width == 4u) ? 1u : 0u; // WRITE4 + out[base + 15] = (width == 8u) ? 1u : 0u; // WRITE8 + // OLD[0..8] at cols 16..24. + for (int i = 0; i < 4; ++i) { + uint64_t p = op[7 + i]; + out[base + 16 + 2 * i] = p & 0xFFFFFFFFu; + out[base + 16 + 2 * i + 1] = p >> 32; + } + // CARRY[0..7] at cols 24..31: carry when adding (i+1) to base_address low word. + uint64_t base_lo = addr & 0xFFFFFFFFu; + for (int i = 0; i < 7; ++i) { + out[base + 24 + i] = (base_lo + (uint64_t)(i + 1) >= (1ull << 32)) ? 1u : 0u; + } + // OLD_TIMESTAMP[i] as DWordWL at cols 31 + 2i (8 timestamps → cols 31..47). + for (int i = 0; i < 8; ++i) { + uint64_t ot = op[11 + i]; + out[base + 31 + 2 * i] = ot & 0xFFFFFFFFu; + out[base + 31 + 2 * i + 1] = ot >> 32; + } + out[base + 47] = is_read; // MU_READ + out[base + 48] = 1u - is_read; // MU_WRITE +} + // On-GPU MEMW_R (register fast-path) fill: write the 10 MEMW_R columns ROW-MAJOR // from the host-walked rows. One thread per row (row_index is the identity for // the fill-from-walked-rows path). Columns mirror diff --git a/crypto/math-cuda/src/device.rs b/crypto/math-cuda/src/device.rs index 857a37bbe..4b64c4a35 100644 --- a/crypto/math-cuda/src/device.rs +++ b/crypto/math-cuda/src/device.rs @@ -200,6 +200,9 @@ pub struct Backend { pub bytewise_fill: CudaFunction, pub mul_fill: CudaFunction, pub dvrm_fill: CudaFunction, + pub branch_fill: CudaFunction, + pub cpu32_fill: CudaFunction, + pub memw_fill: CudaFunction, pub memw_register_fill: CudaFunction, // Twiddle caches keyed by log_n. @@ -403,6 +406,9 @@ impl Backend { bytewise_fill: trace_cpu.load_function("bytewise_fill")?, mul_fill: trace_cpu.load_function("mul_fill")?, dvrm_fill: trace_cpu.load_function("dvrm_fill")?, + branch_fill: trace_cpu.load_function("branch_fill")?, + cpu32_fill: trace_cpu.load_function("cpu32_fill")?, + memw_fill: trace_cpu.load_function("memw_fill")?, memw_register_fill: trace_cpu.load_function("memw_register_fill")?, fwd_twiddles: Mutex::new(vec![None; max_log]), inv_twiddles: Mutex::new(vec![None; max_log]), diff --git a/crypto/math-cuda/src/lde.rs b/crypto/math-cuda/src/lde.rs index fd5c529f2..dd6fa4534 100644 --- a/crypto/math-cuda/src/lde.rs +++ b/crypto/math-cuda/src/lde.rs @@ -617,6 +617,75 @@ pub fn coset_lde_row_major_with_merkle_tree_keep_dev( Ok((handle, lde_out)) } +/// Row-major coset LDE ONLY — no Merkle, no column-major transpose. Returns the +/// row-major LDE `[row*total_cols + col]` (`lde_size * total_cols` u64s), i.e. the +/// exact buffer the CPU `Polynomial::coset_lde_full_expand_row_major` produces. +/// +/// Used by the **preprocessed-table** commit path: those tables commit two *subset* +/// Merkle trees (static columns + multiplicity columns) that the device +/// full-column Merkle cannot represent, so the Merkle stays on the host — but the +/// LDE (the dominant cost) still runs on GPU. Same LDE math as +/// [`coset_lde_row_major_inner`] (single H2D → iNTT → coset → NTT → single D2H), +/// only the commit/transpose steps are dropped. +pub fn coset_lde_row_major_no_merkle( + row_major: &[u64], + n: usize, + total_cols: usize, + blowup_factor: usize, + weights: &[u64], +) -> Result> { + assert_eq!(row_major.len(), n * total_cols); + assert!(n.is_power_of_two()); + assert_eq!(weights.len(), n); + assert!(blowup_factor.is_power_of_two()); + let lde_size = n * blowup_factor; + assert_u32_domain(lde_size, "coset_lde_row_major_no_merkle lde_size"); + let log_n = n.trailing_zeros() as u64; + let log_lde = lde_size.trailing_zeros() as u64; + let n_u64 = n as u64; + let lde_u64 = lde_size as u64; + let cols_u64 = total_cols as u64; + + let be = backend()?; + let stream = be.next_stream(); + + // Zero-padded lde_size*total_cols buffer; first n*total_cols rows carry data. + let mut buf = stream.alloc_zeros::(lde_size * total_cols)?; + stream.memcpy_htod(row_major, &mut buf.slice_mut(0..n * total_cols))?; + + let inv_tw = be.inv_twiddles_for(log_n)?; + let fwd_tw = be.fwd_twiddles_for(log_lde)?; + let weights_dev = stream.clone_htod(weights)?; + + // iNTT → coset weights (incl. 1/N) → forward NTT at lde_size. Identical to the + // committed device path's expansion, so the row-major output is bit-identical. + launch_bit_reverse_row_major(stream.as_ref(), be, &mut buf, n_u64, log_n, cols_u64)?; + run_row_major_ntt_body( + stream.as_ref(), + be, + &mut buf, + inv_tw.as_ref(), + n_u64, + log_n, + cols_u64, + )?; + launch_pointwise_mul_row_major(stream.as_ref(), be, &mut buf, &weights_dev, n_u64, cols_u64)?; + launch_bit_reverse_row_major(stream.as_ref(), be, &mut buf, lde_u64, log_lde, cols_u64)?; + run_row_major_ntt_body( + stream.as_ref(), + be, + &mut buf, + fwd_tw.as_ref(), + lde_u64, + log_lde, + cols_u64, + )?; + + let out = stream.clone_dtoh(&buf)?; + stream.synchronize()?; + Ok(out) +} + /// Row-major ext3 LDE + Keccak + Merkle, all on-device. /// /// `Fp3` is `[u64; 3]` in memory, so row-major ext3 with `m` ext3 columns is diff --git a/crypto/math-cuda/src/trace_cpu.rs b/crypto/math-cuda/src/trace_cpu.rs index ce960740e..b4131cead 100644 --- a/crypto/math-cuda/src/trace_cpu.rs +++ b/crypto/math-cuda/src/trace_cpu.rs @@ -168,6 +168,18 @@ pub const MUL_STRIDE: usize = 5; /// host-deduplicated: one unique op + split (mu_q, mu_r) multiplicities per row. pub const DVRM_NCOLS: usize = 34; pub const DVRM_STRIDE: usize = 5; +/// BRANCH table width / packed input stride (must match `branch_fill`). Input is +/// host-deduplicated: one unique op + summed multiplicity per row. +pub const BRANCH_NCOLS: usize = 14; +pub const BRANCH_STRIDE: usize = 5; +/// CPU32 table width / packed input stride (must match `cpu32_fill`). Per-row +/// (μ=1, no dedup): one op per row. +pub const CPU32_NCOLS: usize = 38; +pub const CPU32_STRIDE: usize = 8; +/// MEMW (general/unaligned) table width / packed input stride (must match +/// `memw_fill`). Per-row (no dedup): one walked op per row. +pub const MEMW_NCOLS: usize = 49; +pub const MEMW_STRIDE: usize = 19; /// Shared interleaved fill on `stream`: upload `packed` (n × stride u64) and run /// `kernel(ops, n, num_rows, out)` into a zeroed row-major buffer. Returns the @@ -225,6 +237,15 @@ fn mul_kernel(be: &Backend) -> &CudaFunction { fn dvrm_kernel(be: &Backend) -> &CudaFunction { &be.dvrm_fill } +fn branch_kernel(be: &Backend) -> &CudaFunction { + &be.branch_fill +} +fn cpu32_kernel(be: &Backend) -> &CudaFunction { + &be.cpu32_fill +} +fn memw_kernel(be: &Backend) -> &CudaFunction { + &be.memw_fill +} /// Build one LT trace-table chunk on device from HOST-DEDUPLICATED ops (one unique /// op + summed multiplicity per row; see `lt_fill` in `trace_cpu.cu`). @@ -426,6 +447,138 @@ pub fn gpu_build_dvrm_trace_host( Ok(host) } +/// Build one BRANCH trace-table chunk on device from HOST-DEDUPLICATED ops (one +/// unique op + summed multiplicity per row; see `branch_fill` in `trace_cpu.cu`). +pub fn gpu_build_branch_trace( + packed_ops: &[u64], + n: usize, + num_rows: usize, +) -> Result> { + let be = backend()?; + let stream = be.next_stream(); + let out = build_interleaved_on( + &stream, + branch_kernel(be), + packed_ops, + n, + num_rows, + BRANCH_NCOLS, + BRANCH_STRIDE, + )?; + stream.synchronize()?; + Ok(out) +} + +/// Host-returning BRANCH build for multiset-equality tests. +pub fn gpu_build_branch_trace_host( + packed_ops: &[u64], + n: usize, + num_rows: usize, +) -> Result> { + let be = backend()?; + let stream = be.next_stream(); + let out = build_interleaved_on( + &stream, + branch_kernel(be), + packed_ops, + n, + num_rows, + BRANCH_NCOLS, + BRANCH_STRIDE, + )?; + let host = stream.clone_dtoh(&out)?; + stream.synchronize()?; + Ok(host) +} + +/// Build one CPU32 trace-table chunk on device (per-row `*W` fill; see `cpu32_fill` +/// in `trace_cpu.cu`). +pub fn gpu_build_cpu32_trace( + packed_ops: &[u64], + n: usize, + num_rows: usize, +) -> Result> { + let be = backend()?; + let stream = be.next_stream(); + let out = build_interleaved_on( + &stream, + cpu32_kernel(be), + packed_ops, + n, + num_rows, + CPU32_NCOLS, + CPU32_STRIDE, + )?; + stream.synchronize()?; + Ok(out) +} + +/// Host-returning CPU32 build for byte-parity tests. +pub fn gpu_build_cpu32_trace_host( + packed_ops: &[u64], + n: usize, + num_rows: usize, +) -> Result> { + let be = backend()?; + let stream = be.next_stream(); + let out = build_interleaved_on( + &stream, + cpu32_kernel(be), + packed_ops, + n, + num_rows, + CPU32_NCOLS, + CPU32_STRIDE, + )?; + let host = stream.clone_dtoh(&out)?; + stream.synchronize()?; + Ok(host) +} + +/// Build one MEMW (general/unaligned) trace-table chunk on device from walked ops +/// (see `memw_fill` in `trace_cpu.cu`). +pub fn gpu_build_memw_trace( + packed_ops: &[u64], + n: usize, + num_rows: usize, +) -> Result> { + let be = backend()?; + let stream = be.next_stream(); + let out = build_interleaved_on( + &stream, + memw_kernel(be), + packed_ops, + n, + num_rows, + MEMW_NCOLS, + MEMW_STRIDE, + )?; + stream.synchronize()?; + Ok(out) +} + +/// Host-returning MEMW (general) build for byte-parity tests. +pub fn gpu_build_memw_trace_host( + packed_ops: &[u64], + n: usize, + num_rows: usize, +) -> Result> { + let be = backend()?; + let stream = be.next_stream(); + let out = build_interleaved_on( + &stream, + memw_kernel(be), + packed_ops, + n, + num_rows, + MEMW_NCOLS, + MEMW_STRIDE, + )?; + let host = stream.clone_dtoh(&out)?; + stream.synchronize()?; + Ok(host) +} + /// Build one SHIFT trace-table chunk on device (residency-ready row-major buffer). /// The kernel recomputes the shift aux from the packed inputs (see `trace_cpu.cu`). pub fn gpu_build_shift_trace( diff --git a/crypto/stark/src/gpu_lde.rs b/crypto/stark/src/gpu_lde.rs index 923b2457d..5dd09c8b9 100644 --- a/crypto/stark/src/gpu_lde.rs +++ b/crypto/stark/src/gpu_lde.rs @@ -576,6 +576,63 @@ where Some((tree, handle, lde_out)) } +/// GPU coset LDE with NO Merkle, for the **preprocessed-table** commit path +/// (DECODE, BITWISE, PAGE, REGISTER, KECCAK_RC). LDE the row-major host trace on +/// device and return the row-major LDE as `Vec>`; the caller builds +/// the two *subset* Merkle trees (static columns + multiplicity columns) on the +/// host, because the device full-column Merkle cannot represent that split. +/// +/// Returns `None` (→ caller LDEs on CPU) for non-Goldilocks fields, tables below +/// the GPU size threshold (kernel launch overhead dominates — e.g. the 32-row +/// KECCAK_RC), empty/mismatched input, or GPU failure. Unlike the device-resident +/// seam there IS a valid CPU fallback here (the host trace is real), so the size +/// threshold applies. +pub(crate) fn try_lde_row_major_no_merkle( + row_major: &[FieldElement], + n: usize, + m: usize, + blowup_factor: usize, + weights: &[FieldElement], +) -> Option>> +where + F: IsField + 'static, + E: IsField + 'static, +{ + let lde_size = n.saturating_mul(blowup_factor); + if lde_size < gpu_lde_threshold() { + return None; + } + if TypeId::of::() != TypeId::of::() { + return None; + } + if TypeId::of::() != TypeId::of::() { + return None; + } + if row_major.len() != n * m || m == 0 || n == 0 { + return None; + } + + let raw: &[u64] = unsafe { from_raw_parts(row_major.as_ptr() as *const u64, n * m) }; + let weights_u64 = unsafe { weights_to_u64::(weights) }; + + GPU_LDE_CALLS.fetch_add(m as u64, Ordering::Relaxed); + + let lde_u64 = + math_cuda::lde::coset_lde_row_major_no_merkle(raw, n, m, blowup_factor, &weights_u64) + .ok()?; + + // Transmute Vec → Vec> (zero-copy, E == GoldilocksField). + 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(), + v.capacity(), + ) + }; + Some(lde_out) +} + /// Row-major ext3 GPU path: single H2D → row-major NTT (m*3 base-field cols) → /// row-major Keccak → Merkle → single D2H → transpose to GpuLdeExt3 handle. /// Same optimization as the base-field path: no extract_columns, no CPU transpose. diff --git a/crypto/stark/src/prover.rs b/crypto/stark/src/prover.rs index c501e080e..13f3cba5c 100644 --- a/crypto/stark/src/prover.rs +++ b/crypto/stark/src/prover.rs @@ -866,24 +866,56 @@ pub trait IsStarkProver< #[cfg(feature = "instruments")] let t_sub = Instant::now(); - let mut main_data: Vec> = Vec::with_capacity(lde_size * total_cols); - main_data.extend_from_slice(trace_data); + // Preprocessed tables (`precomputed.is_some()`) commit two *host* subset + // Merkle trees (static columns + multiplicity columns), so they cannot use + // the device full-column Merkle fast paths above — but the LDE, the dominant + // cost, still runs on GPU here; only the subset Merkles stay on host. A + // non-preprocessed table that fell through to this path already declined the + // GPU above, so it stays on CPU. Same row-major LDE math either way, so the + // subset roots are identical (and the static root is re-checked below). + #[cfg(feature = "cuda")] + let gpu_lde: Option>> = if precomputed.is_some() { + let rows = if total_cols > 0 { + trace_data.len() / total_cols + } else { + 0 + }; + crate::gpu_lde::try_lde_row_major_no_merkle::( + trace_data, + rows, + total_cols, + domain.blowup_factor, + &twiddles.coset_weights, + ) + } else { + None + }; + #[cfg(not(feature = "cuda"))] + let gpu_lde: Option>> = None; + + let main_data: Vec> = match gpu_lde { + Some(lde) => lde, + None => { + let mut cpu = Vec::with_capacity(lde_size * total_cols); + cpu.extend_from_slice(trace_data); + Polynomial::>::coset_lde_full_expand_row_major::( + &mut cpu, + total_cols, + domain.blowup_factor, + &twiddles.coset_weights, + &twiddles.two_half_inv, + &twiddles.two_half_fwd, + ) + .expect("row-major coset LDE expansion"); + cpu + } + }; #[cfg(feature = "disk-spill")] if storage_mode == StorageMode::Disk { trace.main_table.advise_drop_cache(); } - Polynomial::>::coset_lde_full_expand_row_major::( - &mut main_data, - total_cols, - domain.blowup_factor, - &twiddles.coset_weights, - &twiddles.two_half_inv, - &twiddles.two_half_fwd, - ) - .expect("row-major coset LDE expansion"); - #[cfg(feature = "instruments")] let main_lde_dur = t_sub.elapsed(); diff --git a/prover/src/tables/gpu_trace.rs b/prover/src/tables/gpu_trace.rs index c469372c0..68893df8b 100644 --- a/prover/src/tables/gpu_trace.rs +++ b/prover/src/tables/gpu_trace.rs @@ -14,13 +14,15 @@ use stark::trace::TraceTable; use std::collections::HashMap; +use super::branch::{self, BranchOperation}; use super::bytewise::{self, BytewiseOperation}; use super::cpu::{self, CpuOperation}; +use super::cpu32::{self, Cpu32Operation}; use super::dvrm::{self, DvrmOperation}; use super::eq::{self, EqOperation}; use super::load::{self, LoadOperation}; use super::lt::{self, LtOperation}; -use super::memw::MemwOperation; +use super::memw::{self, MemwOperation}; use super::memw_aligned; use super::memw_register::{self, RegRow}; use super::mul::{self, MulOperation}; @@ -758,3 +760,215 @@ pub(crate) fn gpu_build_dvrm_tables( } Some(tables) } + +// ============================================================================= +// BRANCH (branch/jump target dedup table): host per-chunk HashMap dedup → device +// fill (next_pc = (base + offset) & ~1 recomputed on device) +// ============================================================================= + +/// Pack one unique `BranchOperation` + its multiplicity into the `branch_fill` +/// stride. +pub(crate) fn pack_branch_op( + op: &BranchOperation, + mult: u64, +) -> [u64; math_cuda::trace_cpu::BRANCH_STRIDE] { + let flags = op.jalr as u64; + [op.pc, op.offset, op.register, flags, mult] +} + +/// Build one BRANCH trace-table chunk on device. Dedup happens HERE on the host +/// (the same per-chunk HashMap `generate_branch_trace` uses), then the unique ops + +/// summed multiplicities are filled on device (the kernel recomputes next_pc and +/// its byte/half decomposition). BRANCH is a permutation-invariant lookup table, so +/// any row order is valid (validated by multiset/prove, not byte order). +fn build_branch_chunk( + chunk: &[BranchOperation], +) -> Option> { + let mut map: HashMap = HashMap::new(); + for op in chunk { + *map.entry(op.clone()).or_insert(0) += 1; + } + let unique: Vec<(BranchOperation, u64)> = map.into_iter().collect(); + let n = unique.len(); + let num_rows = n.next_power_of_two().max(4); + let mut packed = Vec::with_capacity(n * math_cuda::trace_cpu::BRANCH_STRIDE); + for (op, mult) in &unique { + packed.extend_from_slice(&pack_branch_op(op, *mult)); + } + let dev = math_cuda::trace_cpu::gpu_build_branch_trace(&packed, n, num_rows).ok()?; + let mut trace = TraceTable::new_main( + crate::tables::types::zeroed_fe_vec(num_rows * branch::cols::NUM_COLUMNS), + branch::cols::NUM_COLUMNS, + 1, + ); + trace.set_main_input_dev(Arc::new(dev)); + Some(trace) +} + +pub(crate) fn gpu_build_branch_tables( + ops: &[BranchOperation], + max_rows: usize, +) -> Option>> { + if gpu_trace_disabled() { + return None; + } + // Chunk the RAW ops exactly like `chunk_and_generate`; each chunk dedups + // independently (matching `generate_branch_trace` per chunk). + let chunks: Vec<&[BranchOperation]> = if ops.is_empty() { + vec![&[][..]] + } else { + ops.chunks(max_rows).collect() + }; + let mut tables = Vec::with_capacity(chunks.len()); + for chunk in chunks { + tables.push(build_branch_chunk(chunk)?); + } + Some(tables) +} + +// ============================================================================= +// CPU32 (delegated *W instructions — per-row, no dedup, like the CPU table) +// ============================================================================= + +/// Pack one `Cpu32Operation` into the `cpu32_fill` stride (see `trace_cpu.cu`). The +/// kernel recomputes the sign-extension aux (arg1/arg2/rvd, sign bits), so this +/// only copies the raw fields. +pub(crate) fn pack_cpu32_op(op: &Cpu32Operation) -> [u64; math_cuda::trace_cpu::CPU32_STRIDE] { + let flags = (op.read_register1 as u64) + | ((op.read_register2 as u64) << 1) + | ((op.write_register as u64) << 2) + | ((op.alu as u64) << 3) + | ((op.add as u64) << 4) + | ((op.sub as u64) << 5); + let bytes = (op.rs1 as u64) + | ((op.rs2 as u64) << 8) + | ((op.rd as u64) << 16) + | ((op.alu_flags as u64) << 24) + | ((op.half_instruction_length as u64) << 32); + [ + op.timestamp, + op.pc, + op.rv1, + op.rv2, + op.imm, + op.res, + flags, + bytes, + ] +} + +/// Build one CPU32 trace-table chunk on device: pack the ops → GPU fill → a resident +/// matrix fed to the LDE with no full-column upload. Per-row (μ=1, no dedup), so the +/// device fill is byte-identical to `generate_cpu32_trace`. +fn build_cpu32_chunk( + chunk: &[Cpu32Operation], +) -> Option> { + let n = chunk.len(); + let num_rows = n.next_power_of_two().max(4); + let mut packed = Vec::with_capacity(n * math_cuda::trace_cpu::CPU32_STRIDE); + for op in chunk { + packed.extend_from_slice(&pack_cpu32_op(op)); + } + let dev = math_cuda::trace_cpu::gpu_build_cpu32_trace(&packed, n, num_rows).ok()?; + let mut trace = TraceTable::new_main( + crate::tables::types::zeroed_fe_vec(num_rows * cpu32::cols::NUM_COLUMNS), + cpu32::cols::NUM_COLUMNS, + 1, + ); + trace.set_main_input_dev(Arc::new(dev)); + Some(trace) +} + +pub(crate) fn gpu_build_cpu32_tables( + ops: &[Cpu32Operation], + max_rows: usize, +) -> Option>> { + if gpu_trace_disabled() { + return None; + } + let chunks: Vec<&[Cpu32Operation]> = if ops.is_empty() { + vec![&[][..]] + } else { + ops.chunks(max_rows).collect() + }; + let mut tables = Vec::with_capacity(chunks.len()); + for chunk in chunks { + tables.push(build_cpu32_chunk(chunk)?); + } + Some(tables) +} + +// ============================================================================= +// MEMW (general / unaligned memory — per-op, same MemwOperation as MEMW_A) +// ============================================================================= + +/// Pack one walked `MemwOperation` into the `memw_fill` stride (see `trace_cpu.cu`). +/// value/old (`[u32; 8]` each) pack two-per-u64; the 8 old_timestamps are full u64. +pub(crate) fn pack_memw_op(op: &MemwOperation) -> [u64; math_cuda::trace_cpu::MEMW_STRIDE] { + let flags = (op.is_register as u64) | ((op.is_read as u64) << 1) | ((op.width as u64) << 8); + let v = &op.value; + let o = &op.old; + let ot = &op.old_timestamp; + [ + flags, + op.base_address, + op.timestamp, + v[0] as u64 | ((v[1] as u64) << 32), + v[2] as u64 | ((v[3] as u64) << 32), + v[4] as u64 | ((v[5] as u64) << 32), + v[6] as u64 | ((v[7] as u64) << 32), + o[0] as u64 | ((o[1] as u64) << 32), + o[2] as u64 | ((o[3] as u64) << 32), + o[4] as u64 | ((o[5] as u64) << 32), + o[6] as u64 | ((o[7] as u64) << 32), + ot[0], + ot[1], + ot[2], + ot[3], + ot[4], + ot[5], + ot[6], + ot[7], + ] +} + +/// Build one MEMW (general) trace-table chunk on device: pack the walked ops → GPU +/// fill → a resident matrix fed to the LDE with no full-column upload. Per-row (no +/// dedup), so the device fill is byte-identical to `generate_memw_trace`. +fn build_memw_chunk( + chunk: &[MemwOperation], +) -> Option> { + let n = chunk.len(); + let num_rows = n.next_power_of_two().max(4); + let mut packed = Vec::with_capacity(n * math_cuda::trace_cpu::MEMW_STRIDE); + for op in chunk { + packed.extend_from_slice(&pack_memw_op(op)); + } + let dev = math_cuda::trace_cpu::gpu_build_memw_trace(&packed, n, num_rows).ok()?; + let mut trace = TraceTable::new_main( + crate::tables::types::zeroed_fe_vec(num_rows * memw::cols::NUM_COLUMNS), + memw::cols::NUM_COLUMNS, + 1, + ); + trace.set_main_input_dev(Arc::new(dev)); + Some(trace) +} + +pub(crate) fn gpu_build_memw_tables( + ops: &[MemwOperation], + max_rows: usize, +) -> Option>> { + if gpu_trace_disabled() { + return None; + } + let chunks: Vec<&[MemwOperation]> = if ops.is_empty() { + vec![&[][..]] + } else { + ops.chunks(max_rows).collect() + }; + let mut tables = Vec::with_capacity(chunks.len()); + for chunk in chunks { + tables.push(build_memw_chunk(chunk)?); + } + Some(tables) +} diff --git a/prover/src/tables/trace_builder.rs b/prover/src/tables/trace_builder.rs index b2227ae87..0163ad6bc 100644 --- a/prover/src/tables/trace_builder.rs +++ b/prover/src/tables/trace_builder.rs @@ -3219,6 +3219,25 @@ fn build_traces( ) }; let gen_memws = || { + #[cfg(feature = "cuda")] + { + let use_gpu = { + #[cfg(feature = "disk-spill")] + { + storage_mode != StorageMode::Disk + } + #[cfg(not(feature = "disk-spill"))] + { + true + } + }; + if use_gpu + && let Some(tables) = + crate::tables::gpu_trace::gpu_build_memw_tables(&memw_ops, max_rows.memw) + { + return Ok(tables); + } + } chunk_and_generate( &memw_ops, max_rows.memw, @@ -3438,6 +3457,25 @@ fn build_traces( ) }; let gen_branches = || { + #[cfg(feature = "cuda")] + { + let use_gpu = { + #[cfg(feature = "disk-spill")] + { + storage_mode != StorageMode::Disk + } + #[cfg(not(feature = "disk-spill"))] + { + true + } + }; + if use_gpu + && let Some(tables) = + crate::tables::gpu_trace::gpu_build_branch_tables(&branch_ops, max_rows.branch) + { + return Ok(tables); + } + } chunk_and_generate( &branch_ops, max_rows.branch, @@ -3536,6 +3574,25 @@ fn build_traces( ) }; let gen_cpu32s = || { + #[cfg(feature = "cuda")] + { + let use_gpu = { + #[cfg(feature = "disk-spill")] + { + storage_mode != StorageMode::Disk + } + #[cfg(not(feature = "disk-spill"))] + { + true + } + }; + if use_gpu + && let Some(tables) = + crate::tables::gpu_trace::gpu_build_cpu32_tables(&cpu32_ops, max_rows.cpu32) + { + return Ok(tables); + } + } chunk_and_generate::( &cpu32_ops, max_rows.cpu32, @@ -5401,4 +5458,223 @@ mod gpu_fill_tests { "device DVRM rows must match the CPU fill as a multiset" ); } + + /// BRANCH device fill: a permutation-invariant lookup table (dedup + summed + /// multiplicity), so the row ORDER is non-deterministic. Validate as a MULTISET + /// — the real rows (μ>0) must match `generate_branch_trace`. Covers JALR (base = + /// register) vs PC-relative, wrapping add, odd offsets (so LSB masking differs + /// from the unmasked low byte), high address bits, and duplicates (μ>1). Skips + /// cleanly with no GPU. + #[cfg(feature = "cuda")] + #[test] + fn gpu_branch_fill_matches_cpu_multiset() { + if math_cuda::device::backend().is_err() { + eprintln!("skipping gpu_branch_fill_matches_cpu_multiset: no CUDA backend"); + return; + } + let mut raw = Vec::new(); + for i in 0..800u64 { + let pc = 0x1000u64.wrapping_add(i.wrapping_mul(4)) ^ (i << 34); + // Odd offsets so `next_pc` (LSB masked) differs from the unmasked low + // byte; high bits set to exercise the wrapping add. + let offset = i.wrapping_mul(0x9E37_79B9_7F4A_7C15) | 1; + let register = i.wrapping_mul(0x1234_5678_9ABC_DEF1).rotate_left(7); + let jalr = i % 2 == 0; + let op = branch::BranchOperation::new(pc, offset, register, jalr); + raw.push(op.clone()); + if i % 4 == 0 { + raw.push(op); // duplicate → μ=2 + } + } + + let ncols = math_cuda::trace_cpu::BRANCH_NCOLS; + let real_rows = |flat: &[u64]| -> Vec> { + let mut rows: Vec> = flat + .chunks(ncols) + .filter(|row| row[branch::cols::MU] > 0) + .map(|row| row.to_vec()) + .collect(); + rows.sort(); + rows + }; + + let cpu_table = branch::generate_branch_trace(&raw); + let (cpu_fe, w) = cpu_table.main_data_row_major(); + assert_eq!(w, ncols); + let cpu_u64: Vec = cpu_fe + .iter() + .map(|e| unsafe { *(e.value() as *const u64) }) + .collect(); + + // Dedup on the host exactly like `gpu_build_branch_tables`, then device-fill. + let mut map: std::collections::HashMap = HashMap::new(); + for op in &raw { + *map.entry(op.clone()).or_insert(0) += 1; + } + let unique: Vec<(branch::BranchOperation, u64)> = map.into_iter().collect(); + let n = unique.len(); + let num_rows = n.next_power_of_two().max(4); + let mut packed = Vec::with_capacity(n * math_cuda::trace_cpu::BRANCH_STRIDE); + for (op, mult) in &unique { + packed.extend_from_slice(&crate::tables::gpu_trace::pack_branch_op(op, *mult)); + } + let gpu_u64 = math_cuda::trace_cpu::gpu_build_branch_trace_host(&packed, n, num_rows) + .expect("device BRANCH build must run on a box with a CUDA backend"); + + assert_eq!( + real_rows(&gpu_u64), + real_rows(&cpu_u64), + "device BRANCH rows must match the CPU fill as a multiset" + ); + } + + /// CPU32 device fill must be byte-identical to `cpu32::generate_cpu32_trace`. + /// Per-row (μ=1, no dedup), so full byte parity holds. Covers signed/unsigned + /// (alu_flags bit 5), negative rv1/rv2/res (sign-extension into arg1/arg2/rvd), + /// imm-vs-rv2 operands, flag/register combinations, high words, and padding + /// (n=300 < 512). Skips cleanly with no GPU. + #[cfg(feature = "cuda")] + #[test] + fn gpu_cpu32_fill_matches_cpu() { + if math_cuda::device::backend().is_err() { + eprintln!("skipping gpu_cpu32_fill_matches_cpu: no CUDA backend"); + return; + } + let mut ops = Vec::new(); + for i in 0..300u64 { + let signed = i % 2 == 0; + // alu_flags: bit 5 = signed; low bits carry a varied opcode. + let alu_flags = ((i % 20) as u8) | if signed { 1 << 5 } else { 0 }; + // Alternate imm-driven vs rv2-driven arg2 (decode assumption: one nonzero). + let use_imm = i % 3 == 0; + let rv2 = if use_imm { + 0 + } else { + i.wrapping_mul(0xDEAD_BEEF).rotate_left(9) + }; + let imm = if use_imm { + i.wrapping_mul(0x9E37_79B9) | (i << 40) + } else { + 0 + }; + ops.push(cpu32::Cpu32Operation { + timestamp: 4 * i + 8 + (i << 34), + pc: 0x1000 + i * 4 + ((i % 3) << 33), + rs1: (i % 32) as u8, + read_register1: i % 3 != 0, + rv1: i.wrapping_mul(0x1234_5678_9ABC) ^ (i << 31), // exercise bit 31 + rs2: ((i + 5) % 32) as u8, + read_register2: !use_imm, + rv2, + imm, + res: i.wrapping_mul(0xABCD_1234) ^ (i << 31), + rd: ((i + 7) % 32) as u8, + write_register: i % 4 != 0, + alu: i % 5 != 0, + alu_flags, + add: i % 5 == 1, + sub: i % 5 == 2, + half_instruction_length: if i % 2 == 0 { 1 } else { 2 }, + }); + } + let n = ops.len(); + let num_rows = n.next_power_of_two().max(4); + + let cpu_table = cpu32::generate_cpu32_trace(&ops); + let (cpu_fe, w) = cpu_table.main_data_row_major(); + assert_eq!(w, math_cuda::trace_cpu::CPU32_NCOLS); + let cpu_u64: Vec = cpu_fe + .iter() + .map(|e| unsafe { *(e.value() as *const u64) }) + .collect(); + + let mut packed = Vec::with_capacity(n * math_cuda::trace_cpu::CPU32_STRIDE); + for op in &ops { + packed.extend_from_slice(&crate::tables::gpu_trace::pack_cpu32_op(op)); + } + let gpu_u64 = math_cuda::trace_cpu::gpu_build_cpu32_trace_host(&packed, n, num_rows) + .expect("device CPU32 build must run on a box with a CUDA backend"); + + assert_eq!(gpu_u64.len(), num_rows * math_cuda::trace_cpu::CPU32_NCOLS); + assert_eq!( + gpu_u64, cpu_u64, + "device CPU32 table must be byte-identical to the CPU fill" + ); + } + + /// MEMW (general/unaligned) device fill must be byte-identical to + /// `memw::generate_memw_trace`. Per-row (no dedup). Covers memory and register + /// accesses, widths 1/2/4/8, read/write, base addresses that straddle the 2^32 + /// boundary (so `carry[i]` fires), distinct per-byte old_timestamps (the + /// split-timestamp path), and full-u32 value/old limbs (exercises both halves of + /// the 2×u32/u64 packing). Skips cleanly with no GPU. + #[cfg(feature = "cuda")] + #[test] + fn gpu_memw_fill_matches_cpu() { + if math_cuda::device::backend().is_err() { + eprintln!("skipping gpu_memw_fill_matches_cpu: no CUDA backend"); + return; + } + let mut ops = Vec::new(); + for i in 0..400u64 { + let width = [1u8, 2, 4, 8][(i % 4) as usize]; + let is_read = i % 2 == 0; + let is_register = i % 7 == 0; + // Low word near 2^32 so `base_lo + (i+1)` overflows for some rows. + let base = (0x1_0000_0000u64 * (i % 3)) + (0xFFFF_FFF8u64 - (i % 16)) + i; + let value = [ + (i as u32).wrapping_mul(2_654_435_761), + (i as u32) ^ 0xABCD_1234, + i as u32, + 0xDEAD_0000 | (i as u32 & 0xFFFF), + 7, + 0, + (i as u32).wrapping_add(99), + (i as u32) & 0xFF, + ]; + let old = [ + (i as u32).wrapping_add(3), + (i as u32).wrapping_mul(17), + 0, + (i as u32) ^ 0x0BAD_F00D, + (i as u32).wrapping_add(5), + (i as u32).wrapping_mul(23), + 0, + (i as u32).wrapping_add(7), + ]; + let ts = 4 * i + 100; + // Distinct per-byte old_timestamps (the unaligned split-timestamp path). + let mut old_ts = [0u64; 8]; + for (j, o) in old_ts.iter_mut().enumerate() { + *o = (4 * i + 3 + j as u64) ^ ((j as u64) << 33); + } + ops.push( + memw::MemwOperation::new(is_register, base, value, ts, width, is_read) + .with_old(old, old_ts), + ); + } + let n = ops.len(); + let num_rows = n.next_power_of_two().max(4); + + let cpu_table = memw::generate_memw_trace(&ops); + let (cpu_fe, w) = cpu_table.main_data_row_major(); + assert_eq!(w, math_cuda::trace_cpu::MEMW_NCOLS); + let cpu_u64: Vec = cpu_fe + .iter() + .map(|e| unsafe { *(e.value() as *const u64) }) + .collect(); + + let mut packed = Vec::with_capacity(n * math_cuda::trace_cpu::MEMW_STRIDE); + for op in &ops { + packed.extend_from_slice(&crate::tables::gpu_trace::pack_memw_op(op)); + } + let gpu_u64 = math_cuda::trace_cpu::gpu_build_memw_trace_host(&packed, n, num_rows) + .expect("device MEMW build must run on a box with a CUDA backend"); + + assert_eq!(gpu_u64.len(), num_rows * math_cuda::trace_cpu::MEMW_NCOLS); + assert_eq!( + gpu_u64, cpu_u64, + "device MEMW table must be byte-identical to the CPU fill" + ); + } } From c8241583faad3d377d335df5c9fc562d03ca8075 Mon Sep 17 00:00:00 2001 From: Joaquin Carletti Date: Mon, 13 Jul 2026 18:11:52 -0300 Subject: [PATCH 3/3] move tests --- prover/src/tables/gpu_trace.rs | 241 ++----- prover/src/tables/trace_builder.rs | 1013 ---------------------------- prover/src/tests/gpu_fill_tests.rs | 1004 +++++++++++++++++++++++++++ prover/src/tests/mod.rs | 2 + 4 files changed, 1053 insertions(+), 1207 deletions(-) create mode 100644 prover/src/tests/gpu_fill_tests.rs diff --git a/prover/src/tables/gpu_trace.rs b/prover/src/tables/gpu_trace.rs index 68893df8b..6a43a8240 100644 --- a/prover/src/tables/gpu_trace.rs +++ b/prover/src/tables/gpu_trace.rs @@ -43,6 +43,33 @@ pub(crate) fn gpu_trace_disabled() -> bool { }) } +/// Shared GPU table-build dispatcher. Mirrors `chunk_and_generate`'s chunking +/// (`ops.chunks(max_rows)`, with one empty chunk when there are no ops so the +/// table still emits its padded minimum), builds each chunk on device via +/// `build_chunk`, and collects the resident tables. Returns `None` when the +/// kill-switch is set (`LAMBDA_VM_CPU_TRACE`) or any chunk fails to build, so the +/// caller falls back to the CPU generator. Every `gpu_build_*_tables` entry point +/// is a thin wrapper over this. +fn gpu_build_tables( + ops: &[T], + max_rows: usize, + build_chunk: impl Fn(&[T]) -> Option>, +) -> Option>> { + if gpu_trace_disabled() { + return None; + } + let chunks: Vec<&[T]> = if ops.is_empty() { + vec![&[][..]] + } else { + ops.chunks(max_rows).collect() + }; + let mut tables = Vec::with_capacity(chunks.len()); + for chunk in chunks { + tables.push(build_chunk(chunk)?); + } + Some(tables) +} + // ============================================================================= // CPU table (the first table built on device — see GPU-TRACEGEN-DESIGN-V2 §P1) // ============================================================================= @@ -121,19 +148,7 @@ pub(crate) fn gpu_build_cpu_trace_tables( cpu_ops: &[CpuOperation], max_rows: usize, ) -> Option>> { - if gpu_trace_disabled() { - return None; - } - let chunks: Vec<&[CpuOperation]> = if cpu_ops.is_empty() { - vec![&[][..]] - } else { - cpu_ops.chunks(max_rows).collect() - }; - let mut tables = Vec::with_capacity(chunks.len()); - for chunk in chunks { - tables.push(build_cpu_chunk(chunk)?); - } - Some(tables) + gpu_build_tables(cpu_ops, max_rows, build_cpu_chunk) } // ============================================================================= @@ -192,19 +207,7 @@ pub(crate) fn gpu_build_memw_register_tables( rows: &[RegRow], max_rows: usize, ) -> Option>> { - if gpu_trace_disabled() { - return None; - } - let chunks: Vec<&[RegRow]> = if rows.is_empty() { - vec![&[][..]] - } else { - rows.chunks(max_rows).collect() - }; - let mut tables = Vec::with_capacity(chunks.len()); - for chunk in chunks { - tables.push(build_memw_register_chunk(chunk)?); - } - Some(tables) + gpu_build_tables(rows, max_rows, build_memw_register_chunk) } // ============================================================================= @@ -265,19 +268,7 @@ pub(crate) fn gpu_build_memw_aligned_tables( ops: &[MemwOperation], max_rows: usize, ) -> Option>> { - if gpu_trace_disabled() { - return None; - } - let chunks: Vec<&[MemwOperation]> = if ops.is_empty() { - vec![&[][..]] - } else { - ops.chunks(max_rows).collect() - }; - let mut tables = Vec::with_capacity(chunks.len()); - for chunk in chunks { - tables.push(build_memw_aligned_chunk(chunk)?); - } - Some(tables) + gpu_build_tables(ops, max_rows, build_memw_aligned_chunk) } // ============================================================================= @@ -328,19 +319,7 @@ pub(crate) fn gpu_build_load_tables( ops: &[LoadOperation], max_rows: usize, ) -> Option>> { - if gpu_trace_disabled() { - return None; - } - let chunks: Vec<&[LoadOperation]> = if ops.is_empty() { - vec![&[][..]] - } else { - ops.chunks(max_rows).collect() - }; - let mut tables = Vec::with_capacity(chunks.len()); - for chunk in chunks { - tables.push(build_load_chunk(chunk)?); - } - Some(tables) + gpu_build_tables(ops, max_rows, build_load_chunk) } fn build_store_chunk( @@ -366,19 +345,7 @@ pub(crate) fn gpu_build_store_tables( ops: &[StoreOperation], max_rows: usize, ) -> Option>> { - if gpu_trace_disabled() { - return None; - } - let chunks: Vec<&[StoreOperation]> = if ops.is_empty() { - vec![&[][..]] - } else { - ops.chunks(max_rows).collect() - }; - let mut tables = Vec::with_capacity(chunks.len()); - for chunk in chunks { - tables.push(build_store_chunk(chunk)?); - } - Some(tables) + gpu_build_tables(ops, max_rows, build_store_chunk) } // ============================================================================= @@ -419,19 +386,7 @@ pub(crate) fn gpu_build_shift_tables( ops: &[ShiftOperation], max_rows: usize, ) -> Option>> { - if gpu_trace_disabled() { - return None; - } - let chunks: Vec<&[ShiftOperation]> = if ops.is_empty() { - vec![&[][..]] - } else { - ops.chunks(max_rows).collect() - }; - let mut tables = Vec::with_capacity(chunks.len()); - for chunk in chunks { - tables.push(build_shift_chunk(chunk)?); - } - Some(tables) + gpu_build_tables(ops, max_rows, build_shift_chunk) } // ============================================================================= @@ -476,21 +431,8 @@ pub(crate) fn gpu_build_lt_tables( ops: &[LtOperation], max_rows: usize, ) -> Option>> { - if gpu_trace_disabled() { - return None; - } - // Chunk the RAW ops exactly like `chunk_and_generate`; each chunk dedups - // independently (matching `generate_lt_trace` per chunk). - let chunks: Vec<&[LtOperation]> = if ops.is_empty() { - vec![&[][..]] - } else { - ops.chunks(max_rows).collect() - }; - let mut tables = Vec::with_capacity(chunks.len()); - for chunk in chunks { - tables.push(build_lt_chunk(chunk)?); - } - Some(tables) + // Each chunk dedups independently (matching `generate_lt_trace` per chunk). + gpu_build_tables(ops, max_rows, build_lt_chunk) } // ============================================================================= @@ -535,21 +477,8 @@ pub(crate) fn gpu_build_eq_tables( ops: &[EqOperation], max_rows: usize, ) -> Option>> { - if gpu_trace_disabled() { - return None; - } - // Chunk the RAW ops exactly like `chunk_and_generate`; each chunk dedups - // independently (matching `generate_eq_trace` per chunk). - let chunks: Vec<&[EqOperation]> = if ops.is_empty() { - vec![&[][..]] - } else { - ops.chunks(max_rows).collect() - }; - let mut tables = Vec::with_capacity(chunks.len()); - for chunk in chunks { - tables.push(build_eq_chunk(chunk)?); - } - Some(tables) + // Each chunk dedups independently (matching `generate_eq_trace` per chunk). + gpu_build_tables(ops, max_rows, build_eq_chunk) } // ============================================================================= @@ -598,21 +527,8 @@ pub(crate) fn gpu_build_bytewise_tables( ops: &[BytewiseOperation], max_rows: usize, ) -> Option>> { - if gpu_trace_disabled() { - return None; - } - // Chunk the RAW ops exactly like `chunk_and_generate`; each chunk dedups - // independently (matching `generate_bytewise_trace` per chunk). - let chunks: Vec<&[BytewiseOperation]> = if ops.is_empty() { - vec![&[][..]] - } else { - ops.chunks(max_rows).collect() - }; - let mut tables = Vec::with_capacity(chunks.len()); - for chunk in chunks { - tables.push(build_bytewise_chunk(chunk)?); - } - Some(tables) + // Each chunk dedups independently (matching `generate_bytewise_trace` per chunk). + gpu_build_tables(ops, max_rows, build_bytewise_chunk) } // ============================================================================= @@ -671,21 +587,8 @@ pub(crate) fn gpu_build_mul_tables( ops: &[(MulOperation, bool)], max_rows: usize, ) -> Option>> { - if gpu_trace_disabled() { - return None; - } - // Chunk the RAW (op, wants_hi) pairs exactly like `chunk_and_generate`; each - // chunk dedups independently (matching `generate_mul_trace` per chunk). - let chunks: Vec<&[(MulOperation, bool)]> = if ops.is_empty() { - vec![&[][..]] - } else { - ops.chunks(max_rows).collect() - }; - let mut tables = Vec::with_capacity(chunks.len()); - for chunk in chunks { - tables.push(build_mul_chunk(chunk)?); - } - Some(tables) + // Each chunk dedups independently (matching `generate_mul_trace` per chunk). + gpu_build_tables(ops, max_rows, build_mul_chunk) } // ============================================================================= @@ -744,21 +647,8 @@ pub(crate) fn gpu_build_dvrm_tables( ops: &[(DvrmOperation, bool)], max_rows: usize, ) -> Option>> { - if gpu_trace_disabled() { - return None; - } - // Chunk the RAW (op, wants_remainder) pairs exactly like `chunk_and_generate`; - // each chunk dedups independently (matching `generate_dvrm_trace` per chunk). - let chunks: Vec<&[(DvrmOperation, bool)]> = if ops.is_empty() { - vec![&[][..]] - } else { - ops.chunks(max_rows).collect() - }; - let mut tables = Vec::with_capacity(chunks.len()); - for chunk in chunks { - tables.push(build_dvrm_chunk(chunk)?); - } - Some(tables) + // Each chunk dedups independently (matching `generate_dvrm_trace` per chunk). + gpu_build_tables(ops, max_rows, build_dvrm_chunk) } // ============================================================================= @@ -809,21 +699,8 @@ pub(crate) fn gpu_build_branch_tables( ops: &[BranchOperation], max_rows: usize, ) -> Option>> { - if gpu_trace_disabled() { - return None; - } - // Chunk the RAW ops exactly like `chunk_and_generate`; each chunk dedups - // independently (matching `generate_branch_trace` per chunk). - let chunks: Vec<&[BranchOperation]> = if ops.is_empty() { - vec![&[][..]] - } else { - ops.chunks(max_rows).collect() - }; - let mut tables = Vec::with_capacity(chunks.len()); - for chunk in chunks { - tables.push(build_branch_chunk(chunk)?); - } - Some(tables) + // Each chunk dedups independently (matching `generate_branch_trace` per chunk). + gpu_build_tables(ops, max_rows, build_branch_chunk) } // ============================================================================= @@ -883,19 +760,7 @@ pub(crate) fn gpu_build_cpu32_tables( ops: &[Cpu32Operation], max_rows: usize, ) -> Option>> { - if gpu_trace_disabled() { - return None; - } - let chunks: Vec<&[Cpu32Operation]> = if ops.is_empty() { - vec![&[][..]] - } else { - ops.chunks(max_rows).collect() - }; - let mut tables = Vec::with_capacity(chunks.len()); - for chunk in chunks { - tables.push(build_cpu32_chunk(chunk)?); - } - Some(tables) + gpu_build_tables(ops, max_rows, build_cpu32_chunk) } // ============================================================================= @@ -958,17 +823,5 @@ pub(crate) fn gpu_build_memw_tables( ops: &[MemwOperation], max_rows: usize, ) -> Option>> { - if gpu_trace_disabled() { - return None; - } - let chunks: Vec<&[MemwOperation]> = if ops.is_empty() { - vec![&[][..]] - } else { - ops.chunks(max_rows).collect() - }; - let mut tables = Vec::with_capacity(chunks.len()); - for chunk in chunks { - tables.push(build_memw_chunk(chunk)?); - } - Some(tables) + gpu_build_tables(ops, max_rows, build_memw_chunk) } diff --git a/prover/src/tables/trace_builder.rs b/prover/src/tables/trace_builder.rs index 0163ad6bc..aaff3ad40 100644 --- a/prover/src/tables/trace_builder.rs +++ b/prover/src/tables/trace_builder.rs @@ -4665,1016 +4665,3 @@ impl Traces { ) } } - -#[cfg(all(test, feature = "cuda"))] -mod gpu_fill_tests { - //! Byte-parity (and multiset, for order-independent LT) between each table's - //! device fill and the CPU `generate_*_trace`, so the on-GPU trace tables are - //! bit-identical to the CPU path. Each test skips cleanly with no CUDA backend. - - use super::*; - - /// CPU-table device fill must be byte-identical to `cpu::generate_cpu_trace`. - /// The CPU kernel is the most intricate of the seven (word-delegate column - /// masking, `PC_DOUBLE_READ`/`PREV_PC_TIMESTAMP_BORROW`, and the `+4` padding - /// cadence with PC=1), so this guards it the same way its six siblings are - /// guarded. The fill is a pure function of the packed op fields, so the synthetic - /// ops need only be diverse (not a valid execution): word-delegate rows, x255 PC - /// reads (`pc_double_read`), `ts_lo < 3` rows (`prev_pc_timestamp_borrow`), x0 - /// registers, and high-word values (DWordWL/DWordHL splits). `n = 300 < 512` - /// forces padding rows, exercising the `+4` cadence off `last_ts`. Skips cleanly - /// with no GPU. - #[cfg(feature = "cuda")] - #[test] - fn gpu_cpu_fill_matches_cpu() { - use crate::tables::types::{DecodeEntry, ShrunkDecode}; - if math_cuda::device::backend().is_err() { - eprintln!("skipping gpu_cpu_fill_matches_cpu: no CUDA backend"); - return; - } - let mut ops = Vec::new(); - for i in 0..300u64 { - let word = i % 5 == 0; - // rs1 cycles through x255 (PC register), x0, and normal registers so the - // pc_double_read and register-zero-suppression paths are both hit. - let rs1 = match i % 4 { - 0 => 255u8, - 1 => 0, - _ => (i % 31 + 1) as u8, - }; - let fields = ShrunkDecode { - read_register1: i % 3 != 0, - read_register2: i % 2 == 0, - write_register: i % 3 == 0, - word_instr: word, - alu: i % 6 == 0, - add: i % 6 == 1, - sub: i % 6 == 2, - memory: i % 6 == 3, - branch: i % 6 == 4, - ecall: i % 6 == 5, - rs1, - rs2: (i % 32) as u8, - rd: ((i + 7) % 32) as u8, - half_instruction_length: if i % 2 == 0 { 1 } else { 2 }, - alu_flags: (i & 0xFF) as u8, - mem_flags: ((i >> 1) & 0xFF) as u8, - }; - // A few timestamps with ts_lo < 3 (0,1,2) trip prev_pc_timestamp_borrow; - // the rest are large and continue a +4 cadence. - let timestamp = match i { - 0 => 0, - 1 => 1, - 2 => 2, - _ => 4 * i + 100, - }; - // PC/imm/next_pc carry high words (>32 bits) to exercise the DWordWL split - // into PC_1/NEXT_PC_1/IMM_1. - let decode = DecodeEntry { - pc: 0x1000 + i * 4 + ((i % 3) << 33), - imm: i.wrapping_mul(0x9E37_79B9) | ((i % 4) << 40), - fields, - }; - ops.push(CpuOperation { - decode, - timestamp, - next_pc: 0x2000 + i * 4 + ((i % 2) << 34), - rvd: i.wrapping_mul(0x1234_5678_9ABC), - rv1: i.wrapping_mul(0xDEAD_BEEF).rotate_left(13), - rv2: (i ^ 0xFFFF_0000_1111) << 3, - arg2: i.wrapping_add(0x7777_8888_9999), - res: i.wrapping_mul(0xABCD_1234_5678) ^ (i << 48), - branch_cond: i % 3 == 1, - ..Default::default() - }); - } - let n = ops.len(); - let num_rows = n.next_power_of_two().max(4); - let last_ts = ops.last().map(|op| op.timestamp).unwrap_or(0); - - let cpu_table = cpu::generate_cpu_trace(&ops); - let (cpu_fe, w) = cpu_table.main_data_row_major(); - assert_eq!(w, math_cuda::trace_cpu::CPU_NCOLS); - let cpu_u64: Vec = cpu_fe - .iter() - .map(|e| unsafe { *(e.value() as *const u64) }) - .collect(); - - let packed = crate::tables::gpu_trace::pack_cpu_ops(&ops); - let gpu_u64 = math_cuda::trace_cpu::gpu_build_cpu_trace_host(&packed, n, num_rows, last_ts) - .expect("device CPU build must run on a box with a CUDA backend"); - - assert_eq!(gpu_u64.len(), num_rows * math_cuda::trace_cpu::CPU_NCOLS); - assert_eq!( - gpu_u64, cpu_u64, - "device CPU table must be byte-identical to the CPU fill" - ); - } - - /// MEMW_R (register fast-path) device fill must be byte-identical to the CPU - /// `generate_memw_register_trace_from_rows`. The rows are the walked register - /// rows; here we build synthetic `RegRow`s (read + write, varied values/old/ts) - /// and fill both ways. - #[cfg(feature = "cuda")] - #[test] - fn gpu_memw_register_fill_matches_cpu() { - if math_cuda::device::backend().is_err() { - eprintln!("skipping gpu_memw_register_fill_matches_cpu: no CUDA backend"); - return; - } - let mut rows = Vec::new(); - for i in 0..500u64 { - let reg_addr = 2 * (i % 40); // even word-address = 2*reg_index - let ts = 4 * i + 100; - let old_ts = 4 * i + 50; - let val = i.wrapping_mul(2_654_435_761); - let old = i.wrapping_mul(40_503) ^ 0xABCD_1234; - let is_read = i % 3 == 0; - rows.push(memw_register::RegRow::new( - reg_addr, - ts, - val as u32, - (val >> 32) as u32, - old as u32, - (old >> 32) as u32, - old_ts, - is_read, - )); - } - let num_rows = rows.len().next_power_of_two().max(4); - - let cpu_table = memw_register::generate_memw_register_trace_from_rows(&rows); - let (cpu_fe, w) = cpu_table.main_data_row_major(); - assert_eq!(w, math_cuda::trace_cpu::MEMW_REGISTER_NCOLS); - let cpu_u64: Vec = cpu_fe - .iter() - .map(|e| unsafe { *(e.value() as *const u64) }) - .collect(); - - let mut reg_addr = Vec::new(); - let mut ts = Vec::new(); - let mut value = Vec::new(); - let mut is_read = Vec::new(); - let mut old_value = Vec::new(); - let mut old_tsv = Vec::new(); - for r in &rows { - let (ra, t, v, ir, ov, ot) = r.fill_soa(); - reg_addr.push(ra); - ts.push(t); - value.push(v); - is_read.push(ir); - old_value.push(ov); - old_tsv.push(ot); - } - let gpu_u64 = math_cuda::trace_cpu::gpu_fill_memw_register_host( - ®_addr, &ts, &value, &is_read, &old_value, &old_tsv, num_rows, - ) - .expect("device MEMW_R fill must run on a box with a CUDA backend"); - - assert_eq!( - gpu_u64.len(), - num_rows * math_cuda::trace_cpu::MEMW_REGISTER_NCOLS - ); - assert_eq!( - gpu_u64, cpu_u64, - "device MEMW_R fill must be byte-identical to the CPU fill" - ); - } - - /// MEMW_A (aligned memory — the biggest remaining uploader) device fill must be - /// byte-identical to the CPU `generate_memw_aligned_trace`. Exercises memory - /// (widths 2/4/8), register fallback (`is_register`), high address bits, and - /// the 2×u32-per-u64 value/old packing. Skips cleanly with no GPU. - #[cfg(feature = "cuda")] - #[test] - fn gpu_memw_aligned_fill_matches_cpu() { - if math_cuda::device::backend().is_err() { - eprintln!("skipping gpu_memw_aligned_fill_matches_cpu: no CUDA backend"); - return; - } - let mut ops = Vec::new(); - for i in 0..500u64 { - let width = [2u8, 4, 8][(i % 3) as usize]; - let is_read = i % 2 == 0; - let is_register = i % 7 == 0; - // Vary the high address word too (whh split has a 32-bit high word). - let base = (0x1_0000_0000u64 * (i % 4)) + 0x8000 + i * 8; - let value = [ - (i as u32).wrapping_mul(2654435761), - (i as u32) ^ 0xABCD_1234, - i as u32, - 0, - 7, - 0, - 0, - (i as u32) & 0xFF, - ]; - let old = [ - (i as u32).wrapping_add(99), - (i as u32) ^ 0x0BAD_F00D, - 0, - 3, - 0, - 0, - 0, - 0, - ]; - let ts = 4 * i + 100; - let old_ts = [4 * i + 50; 8]; - ops.push( - MemwOperation::new(is_register, base, value, ts, width, is_read) - .with_old(old, old_ts), - ); - } - let n = ops.len(); - let num_rows = n.next_power_of_two().max(4); - - let cpu_table = memw_aligned::generate_memw_aligned_trace(&ops); - let (cpu_fe, width_cols) = cpu_table.main_data_row_major(); - assert_eq!(width_cols, math_cuda::trace_cpu::MEMW_ALIGNED_NCOLS); - let cpu_u64: Vec = cpu_fe - .iter() - .map(|e| unsafe { *(e.value() as *const u64) }) - .collect(); - - let mut packed = Vec::with_capacity(n * math_cuda::trace_cpu::MEMW_ALIGNED_STRIDE); - for op in &ops { - packed.extend_from_slice(&crate::tables::gpu_trace::pack_memw_aligned_op(op)); - } - let gpu_u64 = math_cuda::trace_cpu::gpu_build_memw_aligned_trace_host(&packed, n, num_rows) - .expect("device MEMW_A build must run on a box with a CUDA backend"); - - assert_eq!( - gpu_u64.len(), - num_rows * math_cuda::trace_cpu::MEMW_ALIGNED_NCOLS - ); - assert_eq!( - gpu_u64, cpu_u64, - "device MEMW_A table must be byte-identical to the CPU fill" - ); - } - - /// LOAD device fill byte-parity (widths 1/2/4/8, signed/unsigned, sign-bit, - /// high address bits, res-byte packing). Skips cleanly with no GPU. - #[cfg(feature = "cuda")] - #[test] - fn gpu_load_fill_matches_cpu() { - if math_cuda::device::backend().is_err() { - eprintln!("skipping gpu_load_fill_matches_cpu: no CUDA backend"); - return; - } - let mut ops = Vec::new(); - for i in 0..400u64 { - let width = [1u8, 2, 4, 8][(i % 4) as usize]; - let signed = i % 2 == 0; - let base = (0x1_0000_0000u64 * (i % 3)) + 0x400 + i * 8; - let ts = 4 * i + 7; - let mut res = [0u64; 8]; - for (j, r) in res.iter_mut().enumerate() { - *r = (i.wrapping_mul(31) + j as u64) & 0xFF; - } - ops.push(load::LoadOperation::new(base, ts, width, signed, res)); - } - let n = ops.len(); - let num_rows = n.next_power_of_two().max(4); - let cpu = load::generate_load_trace(&ops); - let (fe, w) = cpu.main_data_row_major(); - assert_eq!(w, math_cuda::trace_cpu::LOAD_NCOLS); - let cpu_u64: Vec = fe - .iter() - .map(|e| unsafe { *(e.value() as *const u64) }) - .collect(); - let mut packed = Vec::with_capacity(n * math_cuda::trace_cpu::LOAD_STRIDE); - for op in &ops { - packed.extend_from_slice(&crate::tables::gpu_trace::pack_load_op(op)); - } - let gpu = math_cuda::trace_cpu::gpu_build_load_trace_host(&packed, n, num_rows) - .expect("device LOAD build must run on a box with a CUDA backend"); - assert_eq!( - gpu, cpu_u64, - "device LOAD table must be byte-identical to the CPU fill" - ); - } - - /// STORE device fill byte-parity (widths 1/2/4/8 via `bytes`, full-value - /// DWordBL split, high address bits). Skips cleanly with no GPU. - #[cfg(feature = "cuda")] - #[test] - fn gpu_store_fill_matches_cpu() { - if math_cuda::device::backend().is_err() { - eprintln!("skipping gpu_store_fill_matches_cpu: no CUDA backend"); - return; - } - let mut ops = Vec::new(); - for i in 0..400u64 { - let bytes = [1u8, 2, 4, 8][(i % 4) as usize]; - let base = (0x1_0000_0000u64 * (i % 3)) + 0x800 + i * 8; - let ts = 4 * i + 9; - let value = i.wrapping_mul(0x9E37_79B9_7F4A_7C15); - ops.push(store::StoreOperation::new(base, ts, value, bytes)); - } - let n = ops.len(); - let num_rows = n.next_power_of_two().max(4); - let cpu = store::generate_store_trace(&ops); - let (fe, w) = cpu.main_data_row_major(); - assert_eq!(w, math_cuda::trace_cpu::STORE_NCOLS); - let cpu_u64: Vec = fe - .iter() - .map(|e| unsafe { *(e.value() as *const u64) }) - .collect(); - let mut packed = Vec::with_capacity(n * math_cuda::trace_cpu::STORE_STRIDE); - for op in &ops { - packed.extend_from_slice(&crate::tables::gpu_trace::pack_store_op(op)); - } - let gpu = math_cuda::trace_cpu::gpu_build_store_trace_host(&packed, n, num_rows) - .expect("device STORE build must run on a box with a CUDA backend"); - assert_eq!( - gpu, cpu_u64, - "device STORE table must be byte-identical to the CPU fill" - ); - } - - /// SHIFT device fill byte-parity: the kernel recomputes the full aux - /// (bit_shift/zbs/x/y/limb_shift/out) — covers left/right, signed/unsigned, - /// word/64-bit, shift amounts spanning 0 / >16 / >32 / >64, negative MSB inputs, - /// and the padding rows (ZBS=1). SHIFT is per-row (μ=1), so byte-parity holds. - #[cfg(feature = "cuda")] - #[test] - fn gpu_shift_fill_matches_cpu() { - if math_cuda::device::backend().is_err() { - eprintln!("skipping gpu_shift_fill_matches_cpu: no CUDA backend"); - return; - } - let mut ops = Vec::new(); - // Exhaustive over shift byte 0..256 × all flags × MSB-set / edge values - // (so the signed arithmetic-right-shift extension path is fully covered). - let values: [u64; 8] = [ - 0, - 1, - 0x8000_0000_0000_0000, - 0xFFFF_FFFF_FFFF_FFFF, - 0x1234_5678_9ABC_DEF0, - 0xFEDC_BA98_7654_3210, - 0x0000_0000_FFFF_FFFF, - 0xFFFF_FFFF_0000_0000, - ]; - for &value in &values { - for shift_amount in 0u64..256 { - for &direction in &[false, true] { - for &signed in &[false, true] { - for &word_instr in &[false, true] { - ops.push(shift::ShiftOperation::new( - value, - shift_amount, - direction, - signed, - word_instr, - )); - } - } - } - } - } - // Also a few large shift_amounts (high SHIFT_B1/H1/HIGH limbs) — real ops - // carry the full rv2 on the ALU bus. - for &sa in &[0x1_0000u64, 0xFFFF_FFFFu64, 0x1234_5678_9ABCu64, u64::MAX] { - ops.push(shift::ShiftOperation::new( - 0xDEAD_BEEF_CAFE_1234, - sa, - false, - true, - false, - )); - } - let n = ops.len(); - let num_rows = n.next_power_of_two().max(4); - let cpu = shift::generate_shift_trace(&ops); - let (fe, w) = cpu.main_data_row_major(); - assert_eq!(w, math_cuda::trace_cpu::SHIFT_NCOLS); - let cpu_u64: Vec = fe - .iter() - .map(|e| unsafe { *(e.value() as *const u64) }) - .collect(); - let mut packed = Vec::with_capacity(n * math_cuda::trace_cpu::SHIFT_STRIDE); - for op in &ops { - packed.extend_from_slice(&crate::tables::gpu_trace::pack_shift_op(op)); - } - let gpu = math_cuda::trace_cpu::gpu_build_shift_trace_host(&packed, n, num_rows) - .expect("device SHIFT build must run on a box with a CUDA backend"); - assert_eq!( - gpu, cpu_u64, - "device SHIFT table must be byte-identical to the CPU fill" - ); - } - - /// LT device fill: dedup rides the permutation-invariant ALU bus, so the row - /// ORDER is non-deterministic (HashMap iteration). Validate as a MULTISET — - /// the set of real rows (μ>0), incl. summed multiplicities, must match the CPU - /// `generate_lt_trace`. Covers signed/unsigned, invert, MSBs, and duplicates - /// (μ>1). Skips cleanly with no GPU. - #[cfg(feature = "cuda")] - #[test] - fn gpu_lt_fill_matches_cpu_multiset() { - if math_cuda::device::backend().is_err() { - eprintln!("skipping gpu_lt_fill_matches_cpu_multiset: no CUDA backend"); - return; - } - // Raw ops with deliberate duplicates (some pushed twice → μ=2). - let mut raw = Vec::new(); - for i in 0..800u64 { - let lhs = i.wrapping_mul(0x9E37_79B9_7F4A_7C15) ^ (i << 50); - let rhs = i.wrapping_mul(0x1234_5678_9ABC_DEF1).rotate_left(17); - let signed = i % 2 == 0; - let invert = i % 3 == 0; - let op = lt::LtOperation::new_with_invert(lhs, rhs, signed, invert); - raw.push(op.clone()); - if i % 4 == 0 { - raw.push(op); // duplicate → multiplicity 2 - } - } - - let ncols = math_cuda::trace_cpu::LT_NCOLS; - // Extract the real (μ>0) rows of a row-major u64 buffer as a sorted multiset. - let real_rows = |flat: &[u64]| -> Vec> { - let mut rows: Vec> = flat - .chunks(ncols) - .filter(|row| row[lt::cols::MU] > 0) - .map(|row| row.to_vec()) - .collect(); - rows.sort(); - rows - }; - - let cpu_table = lt::generate_lt_trace(&raw); - let (cpu_fe, w) = cpu_table.main_data_row_major(); - assert_eq!(w, ncols); - let cpu_u64: Vec = cpu_fe - .iter() - .map(|e| unsafe { *(e.value() as *const u64) }) - .collect(); - - // Dedup on the host exactly like `gpu_build_lt_tables`, then device-fill. - let mut map: std::collections::HashMap = HashMap::new(); - for op in &raw { - *map.entry(op.clone()).or_insert(0) += 1; - } - let unique: Vec<(lt::LtOperation, u64)> = map.into_iter().collect(); - let n = unique.len(); - let num_rows = n.next_power_of_two().max(4); - let mut packed = Vec::with_capacity(n * math_cuda::trace_cpu::LT_STRIDE); - for (op, mult) in &unique { - packed.extend_from_slice(&crate::tables::gpu_trace::pack_lt_op(op, *mult)); - } - let gpu_u64 = math_cuda::trace_cpu::gpu_build_lt_trace_host(&packed, n, num_rows) - .expect("device LT build must run on a box with a CUDA backend"); - - assert_eq!( - real_rows(&gpu_u64), - real_rows(&cpu_u64), - "device LT rows must match the CPU fill as a multiset" - ); - } - - /// EQ device fill: like LT, dedup rides the permutation-invariant ALU bus, so - /// the row ORDER is non-deterministic (HashMap iteration). Validate as a - /// MULTISET — the real rows (μ>0), incl. summed multiplicities, must match the - /// CPU `generate_eq_trace`. Covers a==b and a!=b, invert, high words, and - /// duplicates (μ>1). Skips cleanly with no GPU. - #[cfg(feature = "cuda")] - #[test] - fn gpu_eq_fill_matches_cpu_multiset() { - if math_cuda::device::backend().is_err() { - eprintln!("skipping gpu_eq_fill_matches_cpu_multiset: no CUDA backend"); - return; - } - // Raw ops with equal/unequal operands, high words, and deliberate duplicates. - let mut raw = Vec::new(); - for i in 0..800u64 { - let a = i.wrapping_mul(0x9E37_79B9_7F4A_7C15) ^ (i << 40); - // Every 5th op has a == b to exercise the eq=1 path. - let b = if i % 5 == 0 { - a - } else { - i.wrapping_mul(0x1234_5678_9ABC_DEF1).rotate_left(11) - }; - let invert = i % 3 == 0; - let op = eq::EqOperation::new(a, b, invert); - raw.push(op.clone()); - if i % 4 == 0 { - raw.push(op); // duplicate → multiplicity 2 - } - } - - let ncols = math_cuda::trace_cpu::EQ_NCOLS; - // Extract the real (μ>0) rows of a row-major u64 buffer as a sorted multiset. - let real_rows = |flat: &[u64]| -> Vec> { - let mut rows: Vec> = flat - .chunks(ncols) - .filter(|row| row[eq::cols::MU] > 0) - .map(|row| row.to_vec()) - .collect(); - rows.sort(); - rows - }; - - let cpu_table = eq::generate_eq_trace(&raw); - let (cpu_fe, w) = cpu_table.main_data_row_major(); - assert_eq!(w, ncols); - let cpu_u64: Vec = cpu_fe - .iter() - .map(|e| unsafe { *(e.value() as *const u64) }) - .collect(); - - // Dedup on the host exactly like `gpu_build_eq_tables`, then device-fill. - let mut map: std::collections::HashMap = HashMap::new(); - for op in &raw { - *map.entry(op.clone()).or_insert(0) += 1; - } - let unique: Vec<(eq::EqOperation, u64)> = map.into_iter().collect(); - let n = unique.len(); - let num_rows = n.next_power_of_two().max(4); - let mut packed = Vec::with_capacity(n * math_cuda::trace_cpu::EQ_STRIDE); - for (op, mult) in &unique { - packed.extend_from_slice(&crate::tables::gpu_trace::pack_eq_op(op, *mult)); - } - let gpu_u64 = math_cuda::trace_cpu::gpu_build_eq_trace_host(&packed, n, num_rows) - .expect("device EQ build must run on a box with a CUDA backend"); - - assert_eq!( - real_rows(&gpu_u64), - real_rows(&cpu_u64), - "device EQ rows must match the CPU fill as a multiset" - ); - } - - /// BYTEWISE device fill: like LT/EQ, dedup rides the permutation-invariant ALU - /// bus, so the row ORDER is non-deterministic. Validate as a MULTISET — the real - /// rows (μ>0), incl. summed multiplicities, must match `generate_bytewise_trace`. - /// Covers AND/OR/XOR, full 64-bit operands, and duplicates (μ>1). Skips cleanly - /// with no GPU. - #[cfg(feature = "cuda")] - #[test] - fn gpu_bytewise_fill_matches_cpu_multiset() { - use crate::tables::types::alu_op; - if math_cuda::device::backend().is_err() { - eprintln!("skipping gpu_bytewise_fill_matches_cpu_multiset: no CUDA backend"); - return; - } - // Raw ops cycling AND/OR/XOR, full-word operands, with deliberate duplicates. - let mut raw = Vec::new(); - for i in 0..900u64 { - let a = i.wrapping_mul(0x9E37_79B9_7F4A_7C15) ^ (i << 33); - let b = i.wrapping_mul(0x1234_5678_9ABC_DEF1).rotate_left(19); - let op = [alu_op::AND, alu_op::OR, alu_op::XOR][(i % 3) as usize]; - let bw = bytewise::BytewiseOperation::new(a, b, op); - raw.push(bw.clone()); - if i % 4 == 0 { - raw.push(bw); // duplicate → multiplicity 2 - } - } - - let ncols = math_cuda::trace_cpu::BYTEWISE_NCOLS; - // Extract the real (μ>0) rows of a row-major u64 buffer as a sorted multiset. - let real_rows = |flat: &[u64]| -> Vec> { - let mut rows: Vec> = flat - .chunks(ncols) - .filter(|row| row[bytewise::cols::MU] > 0) - .map(|row| row.to_vec()) - .collect(); - rows.sort(); - rows - }; - - let cpu_table = bytewise::generate_bytewise_trace(&raw); - let (cpu_fe, w) = cpu_table.main_data_row_major(); - assert_eq!(w, ncols); - let cpu_u64: Vec = cpu_fe - .iter() - .map(|e| unsafe { *(e.value() as *const u64) }) - .collect(); - - // Dedup on the host exactly like `gpu_build_bytewise_tables`, then device-fill. - let mut map: std::collections::HashMap = HashMap::new(); - for op in &raw { - *map.entry(op.clone()).or_insert(0) += 1; - } - let unique: Vec<(bytewise::BytewiseOperation, u64)> = map.into_iter().collect(); - let n = unique.len(); - let num_rows = n.next_power_of_two().max(4); - let mut packed = Vec::with_capacity(n * math_cuda::trace_cpu::BYTEWISE_STRIDE); - for (op, mult) in &unique { - packed.extend_from_slice(&crate::tables::gpu_trace::pack_bytewise_op(op, *mult)); - } - let gpu_u64 = math_cuda::trace_cpu::gpu_build_bytewise_trace_host(&packed, n, num_rows) - .expect("device BYTEWISE build must run on a box with a CUDA backend"); - - assert_eq!( - real_rows(&gpu_u64), - real_rows(&cpu_u64), - "device BYTEWISE rows must match the CPU fill as a multiset" - ); - } - - /// MUL device fill: like the other ALU tables, dedup rides the - /// permutation-invariant ALU bus, so the row ORDER is non-deterministic. - /// Validate as a MULTISET — the real rows (mu_lo>0 or mu_hi>0), incl. the split - /// multiplicities, must match `generate_mul_trace`. Covers all four - /// signed/unsigned combos, negative operands, the 128-bit product + raw_product - /// convolution, and lo/hi (wants_hi) requests with duplicates. Skips cleanly - /// with no GPU. - #[cfg(feature = "cuda")] - #[test] - fn gpu_mul_fill_matches_cpu_multiset() { - if math_cuda::device::backend().is_err() { - eprintln!("skipping gpu_mul_fill_matches_cpu_multiset: no CUDA backend"); - return; - } - // Raw (op, wants_hi) pairs: varied operands, all sign combos, both lo and hi - // requests, with deliberate duplicates so mu_lo/mu_hi accumulate. - let mut raw = Vec::new(); - for i in 0..800u64 { - let lhs = i.wrapping_mul(0x9E37_79B9_7F4A_7C15) ^ (i << 40); - let rhs = i.wrapping_mul(0x1234_5678_9ABC_DEF1).rotate_left(23); - let lhs_signed = i % 2 == 0; - let rhs_signed = i % 3 == 0; - let op = mul::MulOperation::new(lhs, lhs_signed, rhs, rhs_signed); - raw.push((op.clone(), i % 2 == 0)); - if i % 3 == 0 { - raw.push((op.clone(), true)); // extra hi request - } - if i % 5 == 0 { - raw.push((op, false)); // extra lo request (duplicate) - } - } - // Explicit sign edge cases: i64::MIN, -1, treated signed and unsigned. - for &(a, b) in &[ - (0x8000_0000_0000_0000u64, 0xFFFF_FFFF_FFFF_FFFFu64), - (0xFFFF_FFFF_FFFF_FFFFu64, 0x8000_0000_0000_0000u64), - ] { - raw.push((mul::MulOperation::new(a, true, b, true), false)); - raw.push((mul::MulOperation::new(a, false, b, false), true)); - } - - let ncols = math_cuda::trace_cpu::MUL_NCOLS; - // Real rows: mu_lo>0 or mu_hi>0. - let real_rows = |flat: &[u64]| -> Vec> { - let mut rows: Vec> = flat - .chunks(ncols) - .filter(|row| row[mul::cols::MU_LO] > 0 || row[mul::cols::MU_HI] > 0) - .map(|row| row.to_vec()) - .collect(); - rows.sort(); - rows - }; - - let cpu_table = mul::generate_mul_trace(&raw); - let (cpu_fe, w) = cpu_table.main_data_row_major(); - assert_eq!(w, ncols); - let cpu_u64: Vec = cpu_fe - .iter() - .map(|e| unsafe { *(e.value() as *const u64) }) - .collect(); - - // Dedup on the host exactly like `gpu_build_mul_tables`, then device-fill. - let mut map: std::collections::HashMap = HashMap::new(); - for (op, wants_hi) in &raw { - let e = map.entry(op.clone()).or_insert((0, 0)); - if *wants_hi { - e.1 += 1; - } else { - e.0 += 1; - } - } - let unique: Vec<(mul::MulOperation, (u64, u64))> = map.into_iter().collect(); - let n = unique.len(); - let num_rows = n.next_power_of_two().max(4); - let mut packed = Vec::with_capacity(n * math_cuda::trace_cpu::MUL_STRIDE); - for (op, (mu_lo, mu_hi)) in &unique { - packed.extend_from_slice(&crate::tables::gpu_trace::pack_mul_op(op, *mu_lo, *mu_hi)); - } - let gpu_u64 = math_cuda::trace_cpu::gpu_build_mul_trace_host(&packed, n, num_rows) - .expect("device MUL build must run on a box with a CUDA backend"); - - assert_eq!( - real_rows(&gpu_u64), - real_rows(&cpu_u64), - "device MUL rows must match the CPU fill as a multiset" - ); - } - - /// DVRM device fill: like the other ALU tables, dedup rides the - /// permutation-invariant ALU bus, so the row ORDER is non-deterministic. - /// Validate as a MULTISET — the real rows (mu_q>0 or mu_r>0), incl. the split - /// multiplicities, must match `generate_dvrm_trace`. Covers signed/unsigned, - /// division-by-zero, the MIN/-1 signed overflow, negative operands, and - /// quotient/remainder (wants_remainder) requests with duplicates. Skips cleanly - /// with no GPU. - #[cfg(feature = "cuda")] - #[test] - fn gpu_dvrm_fill_matches_cpu_multiset() { - if math_cuda::device::backend().is_err() { - eprintln!("skipping gpu_dvrm_fill_matches_cpu_multiset: no CUDA backend"); - return; - } - // Raw (op, wants_remainder) pairs: varied operands (incl. periodic d==0), - // both signednesses, q and r requests, with deliberate duplicates. - let mut raw = Vec::new(); - for i in 0..800u64 { - let n = i.wrapping_mul(0x9E37_79B9_7F4A_7C15) ^ (i << 40); - let d = if i % 11 == 0 { - 0 - } else { - i.wrapping_mul(0x1234_5678_9ABC_DEF1).rotate_left(29) - }; - let signed = i % 2 == 0; - let op = dvrm::DvrmOperation::new(n, d, signed); - raw.push((op.clone(), i % 2 == 0)); - if i % 3 == 0 { - raw.push((op.clone(), true)); // extra remainder request - } - if i % 5 == 0 { - raw.push((op, false)); // extra quotient request (duplicate) - } - } - // Explicit edge cases: signed overflow (MIN/-1), div-by-zero, MIN numerator, - // -1 denominator. - let min = 0x8000_0000_0000_0000u64; - let neg1 = 0xFFFF_FFFF_FFFF_FFFFu64; - for &(n, d, s) in &[ - (min, neg1, true), // signed overflow - (min, neg1, false), // unsigned: /(2^64-1), no overflow - (123u64, 0u64, true), // div-by-zero, signed - (123u64, 0u64, false), // div-by-zero, unsigned - (min, 7u64, true), // negative numerator - (100u64, neg1, true), // negative denominator (n != MIN) - ] { - raw.push((dvrm::DvrmOperation::new(n, d, s), false)); - raw.push((dvrm::DvrmOperation::new(n, d, s), true)); - } - - let ncols = math_cuda::trace_cpu::DVRM_NCOLS; - // Real rows: mu_q>0 or mu_r>0. - let real_rows = |flat: &[u64]| -> Vec> { - let mut rows: Vec> = flat - .chunks(ncols) - .filter(|row| row[dvrm::cols::MU_Q] > 0 || row[dvrm::cols::MU_R] > 0) - .map(|row| row.to_vec()) - .collect(); - rows.sort(); - rows - }; - - let cpu_table = dvrm::generate_dvrm_trace(&raw); - let (cpu_fe, w) = cpu_table.main_data_row_major(); - assert_eq!(w, ncols); - let cpu_u64: Vec = cpu_fe - .iter() - .map(|e| unsafe { *(e.value() as *const u64) }) - .collect(); - - // Dedup on the host exactly like `gpu_build_dvrm_tables`, then device-fill. - let mut map: std::collections::HashMap = HashMap::new(); - for (op, wants_remainder) in &raw { - let e = map.entry(op.clone()).or_insert((0, 0)); - if *wants_remainder { - e.1 += 1; - } else { - e.0 += 1; - } - } - let unique: Vec<(dvrm::DvrmOperation, (u64, u64))> = map.into_iter().collect(); - let n = unique.len(); - let num_rows = n.next_power_of_two().max(4); - let mut packed = Vec::with_capacity(n * math_cuda::trace_cpu::DVRM_STRIDE); - for (op, (mu_q, mu_r)) in &unique { - packed.extend_from_slice(&crate::tables::gpu_trace::pack_dvrm_op(op, *mu_q, *mu_r)); - } - let gpu_u64 = math_cuda::trace_cpu::gpu_build_dvrm_trace_host(&packed, n, num_rows) - .expect("device DVRM build must run on a box with a CUDA backend"); - - assert_eq!( - real_rows(&gpu_u64), - real_rows(&cpu_u64), - "device DVRM rows must match the CPU fill as a multiset" - ); - } - - /// BRANCH device fill: a permutation-invariant lookup table (dedup + summed - /// multiplicity), so the row ORDER is non-deterministic. Validate as a MULTISET - /// — the real rows (μ>0) must match `generate_branch_trace`. Covers JALR (base = - /// register) vs PC-relative, wrapping add, odd offsets (so LSB masking differs - /// from the unmasked low byte), high address bits, and duplicates (μ>1). Skips - /// cleanly with no GPU. - #[cfg(feature = "cuda")] - #[test] - fn gpu_branch_fill_matches_cpu_multiset() { - if math_cuda::device::backend().is_err() { - eprintln!("skipping gpu_branch_fill_matches_cpu_multiset: no CUDA backend"); - return; - } - let mut raw = Vec::new(); - for i in 0..800u64 { - let pc = 0x1000u64.wrapping_add(i.wrapping_mul(4)) ^ (i << 34); - // Odd offsets so `next_pc` (LSB masked) differs from the unmasked low - // byte; high bits set to exercise the wrapping add. - let offset = i.wrapping_mul(0x9E37_79B9_7F4A_7C15) | 1; - let register = i.wrapping_mul(0x1234_5678_9ABC_DEF1).rotate_left(7); - let jalr = i % 2 == 0; - let op = branch::BranchOperation::new(pc, offset, register, jalr); - raw.push(op.clone()); - if i % 4 == 0 { - raw.push(op); // duplicate → μ=2 - } - } - - let ncols = math_cuda::trace_cpu::BRANCH_NCOLS; - let real_rows = |flat: &[u64]| -> Vec> { - let mut rows: Vec> = flat - .chunks(ncols) - .filter(|row| row[branch::cols::MU] > 0) - .map(|row| row.to_vec()) - .collect(); - rows.sort(); - rows - }; - - let cpu_table = branch::generate_branch_trace(&raw); - let (cpu_fe, w) = cpu_table.main_data_row_major(); - assert_eq!(w, ncols); - let cpu_u64: Vec = cpu_fe - .iter() - .map(|e| unsafe { *(e.value() as *const u64) }) - .collect(); - - // Dedup on the host exactly like `gpu_build_branch_tables`, then device-fill. - let mut map: std::collections::HashMap = HashMap::new(); - for op in &raw { - *map.entry(op.clone()).or_insert(0) += 1; - } - let unique: Vec<(branch::BranchOperation, u64)> = map.into_iter().collect(); - let n = unique.len(); - let num_rows = n.next_power_of_two().max(4); - let mut packed = Vec::with_capacity(n * math_cuda::trace_cpu::BRANCH_STRIDE); - for (op, mult) in &unique { - packed.extend_from_slice(&crate::tables::gpu_trace::pack_branch_op(op, *mult)); - } - let gpu_u64 = math_cuda::trace_cpu::gpu_build_branch_trace_host(&packed, n, num_rows) - .expect("device BRANCH build must run on a box with a CUDA backend"); - - assert_eq!( - real_rows(&gpu_u64), - real_rows(&cpu_u64), - "device BRANCH rows must match the CPU fill as a multiset" - ); - } - - /// CPU32 device fill must be byte-identical to `cpu32::generate_cpu32_trace`. - /// Per-row (μ=1, no dedup), so full byte parity holds. Covers signed/unsigned - /// (alu_flags bit 5), negative rv1/rv2/res (sign-extension into arg1/arg2/rvd), - /// imm-vs-rv2 operands, flag/register combinations, high words, and padding - /// (n=300 < 512). Skips cleanly with no GPU. - #[cfg(feature = "cuda")] - #[test] - fn gpu_cpu32_fill_matches_cpu() { - if math_cuda::device::backend().is_err() { - eprintln!("skipping gpu_cpu32_fill_matches_cpu: no CUDA backend"); - return; - } - let mut ops = Vec::new(); - for i in 0..300u64 { - let signed = i % 2 == 0; - // alu_flags: bit 5 = signed; low bits carry a varied opcode. - let alu_flags = ((i % 20) as u8) | if signed { 1 << 5 } else { 0 }; - // Alternate imm-driven vs rv2-driven arg2 (decode assumption: one nonzero). - let use_imm = i % 3 == 0; - let rv2 = if use_imm { - 0 - } else { - i.wrapping_mul(0xDEAD_BEEF).rotate_left(9) - }; - let imm = if use_imm { - i.wrapping_mul(0x9E37_79B9) | (i << 40) - } else { - 0 - }; - ops.push(cpu32::Cpu32Operation { - timestamp: 4 * i + 8 + (i << 34), - pc: 0x1000 + i * 4 + ((i % 3) << 33), - rs1: (i % 32) as u8, - read_register1: i % 3 != 0, - rv1: i.wrapping_mul(0x1234_5678_9ABC) ^ (i << 31), // exercise bit 31 - rs2: ((i + 5) % 32) as u8, - read_register2: !use_imm, - rv2, - imm, - res: i.wrapping_mul(0xABCD_1234) ^ (i << 31), - rd: ((i + 7) % 32) as u8, - write_register: i % 4 != 0, - alu: i % 5 != 0, - alu_flags, - add: i % 5 == 1, - sub: i % 5 == 2, - half_instruction_length: if i % 2 == 0 { 1 } else { 2 }, - }); - } - let n = ops.len(); - let num_rows = n.next_power_of_two().max(4); - - let cpu_table = cpu32::generate_cpu32_trace(&ops); - let (cpu_fe, w) = cpu_table.main_data_row_major(); - assert_eq!(w, math_cuda::trace_cpu::CPU32_NCOLS); - let cpu_u64: Vec = cpu_fe - .iter() - .map(|e| unsafe { *(e.value() as *const u64) }) - .collect(); - - let mut packed = Vec::with_capacity(n * math_cuda::trace_cpu::CPU32_STRIDE); - for op in &ops { - packed.extend_from_slice(&crate::tables::gpu_trace::pack_cpu32_op(op)); - } - let gpu_u64 = math_cuda::trace_cpu::gpu_build_cpu32_trace_host(&packed, n, num_rows) - .expect("device CPU32 build must run on a box with a CUDA backend"); - - assert_eq!(gpu_u64.len(), num_rows * math_cuda::trace_cpu::CPU32_NCOLS); - assert_eq!( - gpu_u64, cpu_u64, - "device CPU32 table must be byte-identical to the CPU fill" - ); - } - - /// MEMW (general/unaligned) device fill must be byte-identical to - /// `memw::generate_memw_trace`. Per-row (no dedup). Covers memory and register - /// accesses, widths 1/2/4/8, read/write, base addresses that straddle the 2^32 - /// boundary (so `carry[i]` fires), distinct per-byte old_timestamps (the - /// split-timestamp path), and full-u32 value/old limbs (exercises both halves of - /// the 2×u32/u64 packing). Skips cleanly with no GPU. - #[cfg(feature = "cuda")] - #[test] - fn gpu_memw_fill_matches_cpu() { - if math_cuda::device::backend().is_err() { - eprintln!("skipping gpu_memw_fill_matches_cpu: no CUDA backend"); - return; - } - let mut ops = Vec::new(); - for i in 0..400u64 { - let width = [1u8, 2, 4, 8][(i % 4) as usize]; - let is_read = i % 2 == 0; - let is_register = i % 7 == 0; - // Low word near 2^32 so `base_lo + (i+1)` overflows for some rows. - let base = (0x1_0000_0000u64 * (i % 3)) + (0xFFFF_FFF8u64 - (i % 16)) + i; - let value = [ - (i as u32).wrapping_mul(2_654_435_761), - (i as u32) ^ 0xABCD_1234, - i as u32, - 0xDEAD_0000 | (i as u32 & 0xFFFF), - 7, - 0, - (i as u32).wrapping_add(99), - (i as u32) & 0xFF, - ]; - let old = [ - (i as u32).wrapping_add(3), - (i as u32).wrapping_mul(17), - 0, - (i as u32) ^ 0x0BAD_F00D, - (i as u32).wrapping_add(5), - (i as u32).wrapping_mul(23), - 0, - (i as u32).wrapping_add(7), - ]; - let ts = 4 * i + 100; - // Distinct per-byte old_timestamps (the unaligned split-timestamp path). - let mut old_ts = [0u64; 8]; - for (j, o) in old_ts.iter_mut().enumerate() { - *o = (4 * i + 3 + j as u64) ^ ((j as u64) << 33); - } - ops.push( - memw::MemwOperation::new(is_register, base, value, ts, width, is_read) - .with_old(old, old_ts), - ); - } - let n = ops.len(); - let num_rows = n.next_power_of_two().max(4); - - let cpu_table = memw::generate_memw_trace(&ops); - let (cpu_fe, w) = cpu_table.main_data_row_major(); - assert_eq!(w, math_cuda::trace_cpu::MEMW_NCOLS); - let cpu_u64: Vec = cpu_fe - .iter() - .map(|e| unsafe { *(e.value() as *const u64) }) - .collect(); - - let mut packed = Vec::with_capacity(n * math_cuda::trace_cpu::MEMW_STRIDE); - for op in &ops { - packed.extend_from_slice(&crate::tables::gpu_trace::pack_memw_op(op)); - } - let gpu_u64 = math_cuda::trace_cpu::gpu_build_memw_trace_host(&packed, n, num_rows) - .expect("device MEMW build must run on a box with a CUDA backend"); - - assert_eq!(gpu_u64.len(), num_rows * math_cuda::trace_cpu::MEMW_NCOLS); - assert_eq!( - gpu_u64, cpu_u64, - "device MEMW table must be byte-identical to the CPU fill" - ); - } -} diff --git a/prover/src/tests/gpu_fill_tests.rs b/prover/src/tests/gpu_fill_tests.rs new file mode 100644 index 000000000..68ae13c26 --- /dev/null +++ b/prover/src/tests/gpu_fill_tests.rs @@ -0,0 +1,1004 @@ +//! GPU device-fill parity tests: byte-parity (and multiset, for the +//! order-independent dedup ALU tables) between each table's on-device fill and its +//! CPU `generate_*_trace`, so the GPU trace tables are bit-identical to the CPU +//! path. Each test skips cleanly with no CUDA backend. Grouped here (rather than +//! inline in the table modules) so production code carries no test code; the whole +//! module is `cuda`-gated at its `mod.rs` registration. + +use std::collections::HashMap; + +use crate::tables::cpu::CpuOperation; +use crate::tables::memw::MemwOperation; +use crate::tables::{ + branch, bytewise, cpu, cpu32, dvrm, eq, load, lt, memw, memw_aligned, memw_register, mul, + shift, store, +}; + +/// CPU-table device fill must be byte-identical to `cpu::generate_cpu_trace`. +/// The CPU kernel is the most intricate of the seven (word-delegate column +/// masking, `PC_DOUBLE_READ`/`PREV_PC_TIMESTAMP_BORROW`, and the `+4` padding +/// cadence with PC=1), so this guards it the same way its six siblings are +/// guarded. The fill is a pure function of the packed op fields, so the synthetic +/// ops need only be diverse (not a valid execution): word-delegate rows, x255 PC +/// reads (`pc_double_read`), `ts_lo < 3` rows (`prev_pc_timestamp_borrow`), x0 +/// registers, and high-word values (DWordWL/DWordHL splits). `n = 300 < 512` +/// forces padding rows, exercising the `+4` cadence off `last_ts`. Skips cleanly +/// with no GPU. +#[test] +fn gpu_cpu_fill_matches_cpu() { + use crate::tables::types::{DecodeEntry, ShrunkDecode}; + if math_cuda::device::backend().is_err() { + eprintln!("skipping gpu_cpu_fill_matches_cpu: no CUDA backend"); + return; + } + let mut ops = Vec::new(); + for i in 0..300u64 { + let word = i % 5 == 0; + // rs1 cycles through x255 (PC register), x0, and normal registers so the + // pc_double_read and register-zero-suppression paths are both hit. + let rs1 = match i % 4 { + 0 => 255u8, + 1 => 0, + _ => (i % 31 + 1) as u8, + }; + let fields = ShrunkDecode { + read_register1: i % 3 != 0, + read_register2: i % 2 == 0, + write_register: i % 3 == 0, + word_instr: word, + alu: i % 6 == 0, + add: i % 6 == 1, + sub: i % 6 == 2, + memory: i % 6 == 3, + branch: i % 6 == 4, + ecall: i % 6 == 5, + rs1, + rs2: (i % 32) as u8, + rd: ((i + 7) % 32) as u8, + half_instruction_length: if i % 2 == 0 { 1 } else { 2 }, + alu_flags: (i & 0xFF) as u8, + mem_flags: ((i >> 1) & 0xFF) as u8, + }; + // A few timestamps with ts_lo < 3 (0,1,2) trip prev_pc_timestamp_borrow; + // the rest are large and continue a +4 cadence. + let timestamp = match i { + 0 => 0, + 1 => 1, + 2 => 2, + _ => 4 * i + 100, + }; + // PC/imm/next_pc carry high words (>32 bits) to exercise the DWordWL split + // into PC_1/NEXT_PC_1/IMM_1. + let decode = DecodeEntry { + pc: 0x1000 + i * 4 + ((i % 3) << 33), + imm: i.wrapping_mul(0x9E37_79B9) | ((i % 4) << 40), + fields, + }; + ops.push(CpuOperation { + decode, + timestamp, + next_pc: 0x2000 + i * 4 + ((i % 2) << 34), + rvd: i.wrapping_mul(0x1234_5678_9ABC), + rv1: i.wrapping_mul(0xDEAD_BEEF).rotate_left(13), + rv2: (i ^ 0xFFFF_0000_1111) << 3, + arg2: i.wrapping_add(0x7777_8888_9999), + res: i.wrapping_mul(0xABCD_1234_5678) ^ (i << 48), + branch_cond: i % 3 == 1, + ..Default::default() + }); + } + let n = ops.len(); + let num_rows = n.next_power_of_two().max(4); + let last_ts = ops.last().map(|op| op.timestamp).unwrap_or(0); + + let cpu_table = cpu::generate_cpu_trace(&ops); + let (cpu_fe, w) = cpu_table.main_data_row_major(); + assert_eq!(w, math_cuda::trace_cpu::CPU_NCOLS); + let cpu_u64: Vec = cpu_fe + .iter() + .map(|e| unsafe { *(e.value() as *const u64) }) + .collect(); + + let packed = crate::tables::gpu_trace::pack_cpu_ops(&ops); + let gpu_u64 = math_cuda::trace_cpu::gpu_build_cpu_trace_host(&packed, n, num_rows, last_ts) + .expect("device CPU build must run on a box with a CUDA backend"); + + assert_eq!(gpu_u64.len(), num_rows * math_cuda::trace_cpu::CPU_NCOLS); + assert_eq!( + gpu_u64, cpu_u64, + "device CPU table must be byte-identical to the CPU fill" + ); +} + +/// MEMW_R (register fast-path) device fill must be byte-identical to the CPU +/// `generate_memw_register_trace_from_rows`. The rows are the walked register +/// rows; here we build synthetic `RegRow`s (read + write, varied values/old/ts) +/// and fill both ways. +#[test] +fn gpu_memw_register_fill_matches_cpu() { + if math_cuda::device::backend().is_err() { + eprintln!("skipping gpu_memw_register_fill_matches_cpu: no CUDA backend"); + return; + } + let mut rows = Vec::new(); + for i in 0..500u64 { + let reg_addr = 2 * (i % 40); // even word-address = 2*reg_index + let ts = 4 * i + 100; + let old_ts = 4 * i + 50; + let val = i.wrapping_mul(2_654_435_761); + let old = i.wrapping_mul(40_503) ^ 0xABCD_1234; + let is_read = i % 3 == 0; + rows.push(memw_register::RegRow::new( + reg_addr, + ts, + val as u32, + (val >> 32) as u32, + old as u32, + (old >> 32) as u32, + old_ts, + is_read, + )); + } + let num_rows = rows.len().next_power_of_two().max(4); + + let cpu_table = memw_register::generate_memw_register_trace_from_rows(&rows); + let (cpu_fe, w) = cpu_table.main_data_row_major(); + assert_eq!(w, math_cuda::trace_cpu::MEMW_REGISTER_NCOLS); + let cpu_u64: Vec = cpu_fe + .iter() + .map(|e| unsafe { *(e.value() as *const u64) }) + .collect(); + + let mut reg_addr = Vec::new(); + let mut ts = Vec::new(); + let mut value = Vec::new(); + let mut is_read = Vec::new(); + let mut old_value = Vec::new(); + let mut old_tsv = Vec::new(); + for r in &rows { + let (ra, t, v, ir, ov, ot) = r.fill_soa(); + reg_addr.push(ra); + ts.push(t); + value.push(v); + is_read.push(ir); + old_value.push(ov); + old_tsv.push(ot); + } + let gpu_u64 = math_cuda::trace_cpu::gpu_fill_memw_register_host( + ®_addr, &ts, &value, &is_read, &old_value, &old_tsv, num_rows, + ) + .expect("device MEMW_R fill must run on a box with a CUDA backend"); + + assert_eq!( + gpu_u64.len(), + num_rows * math_cuda::trace_cpu::MEMW_REGISTER_NCOLS + ); + assert_eq!( + gpu_u64, cpu_u64, + "device MEMW_R fill must be byte-identical to the CPU fill" + ); +} + +/// MEMW_A (aligned memory — the biggest remaining uploader) device fill must be +/// byte-identical to the CPU `generate_memw_aligned_trace`. Exercises memory +/// (widths 2/4/8), register fallback (`is_register`), high address bits, and +/// the 2×u32-per-u64 value/old packing. Skips cleanly with no GPU. +#[test] +fn gpu_memw_aligned_fill_matches_cpu() { + if math_cuda::device::backend().is_err() { + eprintln!("skipping gpu_memw_aligned_fill_matches_cpu: no CUDA backend"); + return; + } + let mut ops = Vec::new(); + for i in 0..500u64 { + let width = [2u8, 4, 8][(i % 3) as usize]; + let is_read = i % 2 == 0; + let is_register = i % 7 == 0; + // Vary the high address word too (whh split has a 32-bit high word). + let base = (0x1_0000_0000u64 * (i % 4)) + 0x8000 + i * 8; + let value = [ + (i as u32).wrapping_mul(2654435761), + (i as u32) ^ 0xABCD_1234, + i as u32, + 0, + 7, + 0, + 0, + (i as u32) & 0xFF, + ]; + let old = [ + (i as u32).wrapping_add(99), + (i as u32) ^ 0x0BAD_F00D, + 0, + 3, + 0, + 0, + 0, + 0, + ]; + let ts = 4 * i + 100; + let old_ts = [4 * i + 50; 8]; + ops.push( + MemwOperation::new(is_register, base, value, ts, width, is_read).with_old(old, old_ts), + ); + } + let n = ops.len(); + let num_rows = n.next_power_of_two().max(4); + + let cpu_table = memw_aligned::generate_memw_aligned_trace(&ops); + let (cpu_fe, width_cols) = cpu_table.main_data_row_major(); + assert_eq!(width_cols, math_cuda::trace_cpu::MEMW_ALIGNED_NCOLS); + let cpu_u64: Vec = cpu_fe + .iter() + .map(|e| unsafe { *(e.value() as *const u64) }) + .collect(); + + let mut packed = Vec::with_capacity(n * math_cuda::trace_cpu::MEMW_ALIGNED_STRIDE); + for op in &ops { + packed.extend_from_slice(&crate::tables::gpu_trace::pack_memw_aligned_op(op)); + } + let gpu_u64 = math_cuda::trace_cpu::gpu_build_memw_aligned_trace_host(&packed, n, num_rows) + .expect("device MEMW_A build must run on a box with a CUDA backend"); + + assert_eq!( + gpu_u64.len(), + num_rows * math_cuda::trace_cpu::MEMW_ALIGNED_NCOLS + ); + assert_eq!( + gpu_u64, cpu_u64, + "device MEMW_A table must be byte-identical to the CPU fill" + ); +} + +/// LOAD device fill byte-parity (widths 1/2/4/8, signed/unsigned, sign-bit, +/// high address bits, res-byte packing). Skips cleanly with no GPU. +#[test] +fn gpu_load_fill_matches_cpu() { + if math_cuda::device::backend().is_err() { + eprintln!("skipping gpu_load_fill_matches_cpu: no CUDA backend"); + return; + } + let mut ops = Vec::new(); + for i in 0..400u64 { + let width = [1u8, 2, 4, 8][(i % 4) as usize]; + let signed = i % 2 == 0; + let base = (0x1_0000_0000u64 * (i % 3)) + 0x400 + i * 8; + let ts = 4 * i + 7; + let mut res = [0u64; 8]; + for (j, r) in res.iter_mut().enumerate() { + *r = (i.wrapping_mul(31) + j as u64) & 0xFF; + } + ops.push(load::LoadOperation::new(base, ts, width, signed, res)); + } + let n = ops.len(); + let num_rows = n.next_power_of_two().max(4); + let cpu = load::generate_load_trace(&ops); + let (fe, w) = cpu.main_data_row_major(); + assert_eq!(w, math_cuda::trace_cpu::LOAD_NCOLS); + let cpu_u64: Vec = fe + .iter() + .map(|e| unsafe { *(e.value() as *const u64) }) + .collect(); + let mut packed = Vec::with_capacity(n * math_cuda::trace_cpu::LOAD_STRIDE); + for op in &ops { + packed.extend_from_slice(&crate::tables::gpu_trace::pack_load_op(op)); + } + let gpu = math_cuda::trace_cpu::gpu_build_load_trace_host(&packed, n, num_rows) + .expect("device LOAD build must run on a box with a CUDA backend"); + assert_eq!( + gpu, cpu_u64, + "device LOAD table must be byte-identical to the CPU fill" + ); +} + +/// STORE device fill byte-parity (widths 1/2/4/8 via `bytes`, full-value +/// DWordBL split, high address bits). Skips cleanly with no GPU. +#[test] +fn gpu_store_fill_matches_cpu() { + if math_cuda::device::backend().is_err() { + eprintln!("skipping gpu_store_fill_matches_cpu: no CUDA backend"); + return; + } + let mut ops = Vec::new(); + for i in 0..400u64 { + let bytes = [1u8, 2, 4, 8][(i % 4) as usize]; + let base = (0x1_0000_0000u64 * (i % 3)) + 0x800 + i * 8; + let ts = 4 * i + 9; + let value = i.wrapping_mul(0x9E37_79B9_7F4A_7C15); + ops.push(store::StoreOperation::new(base, ts, value, bytes)); + } + let n = ops.len(); + let num_rows = n.next_power_of_two().max(4); + let cpu = store::generate_store_trace(&ops); + let (fe, w) = cpu.main_data_row_major(); + assert_eq!(w, math_cuda::trace_cpu::STORE_NCOLS); + let cpu_u64: Vec = fe + .iter() + .map(|e| unsafe { *(e.value() as *const u64) }) + .collect(); + let mut packed = Vec::with_capacity(n * math_cuda::trace_cpu::STORE_STRIDE); + for op in &ops { + packed.extend_from_slice(&crate::tables::gpu_trace::pack_store_op(op)); + } + let gpu = math_cuda::trace_cpu::gpu_build_store_trace_host(&packed, n, num_rows) + .expect("device STORE build must run on a box with a CUDA backend"); + assert_eq!( + gpu, cpu_u64, + "device STORE table must be byte-identical to the CPU fill" + ); +} + +/// SHIFT device fill byte-parity: the kernel recomputes the full aux +/// (bit_shift/zbs/x/y/limb_shift/out) — covers left/right, signed/unsigned, +/// word/64-bit, shift amounts spanning 0 / >16 / >32 / >64, negative MSB inputs, +/// and the padding rows (ZBS=1). SHIFT is per-row (μ=1), so byte-parity holds. +#[test] +fn gpu_shift_fill_matches_cpu() { + if math_cuda::device::backend().is_err() { + eprintln!("skipping gpu_shift_fill_matches_cpu: no CUDA backend"); + return; + } + let mut ops = Vec::new(); + // Exhaustive over shift byte 0..256 × all flags × MSB-set / edge values + // (so the signed arithmetic-right-shift extension path is fully covered). + let values: [u64; 8] = [ + 0, + 1, + 0x8000_0000_0000_0000, + 0xFFFF_FFFF_FFFF_FFFF, + 0x1234_5678_9ABC_DEF0, + 0xFEDC_BA98_7654_3210, + 0x0000_0000_FFFF_FFFF, + 0xFFFF_FFFF_0000_0000, + ]; + for &value in &values { + for shift_amount in 0u64..256 { + for &direction in &[false, true] { + for &signed in &[false, true] { + for &word_instr in &[false, true] { + ops.push(shift::ShiftOperation::new( + value, + shift_amount, + direction, + signed, + word_instr, + )); + } + } + } + } + } + // Also a few large shift_amounts (high SHIFT_B1/H1/HIGH limbs) — real ops + // carry the full rv2 on the ALU bus. + for &sa in &[0x1_0000u64, 0xFFFF_FFFFu64, 0x1234_5678_9ABCu64, u64::MAX] { + ops.push(shift::ShiftOperation::new( + 0xDEAD_BEEF_CAFE_1234, + sa, + false, + true, + false, + )); + } + let n = ops.len(); + let num_rows = n.next_power_of_two().max(4); + let cpu = shift::generate_shift_trace(&ops); + let (fe, w) = cpu.main_data_row_major(); + assert_eq!(w, math_cuda::trace_cpu::SHIFT_NCOLS); + let cpu_u64: Vec = fe + .iter() + .map(|e| unsafe { *(e.value() as *const u64) }) + .collect(); + let mut packed = Vec::with_capacity(n * math_cuda::trace_cpu::SHIFT_STRIDE); + for op in &ops { + packed.extend_from_slice(&crate::tables::gpu_trace::pack_shift_op(op)); + } + let gpu = math_cuda::trace_cpu::gpu_build_shift_trace_host(&packed, n, num_rows) + .expect("device SHIFT build must run on a box with a CUDA backend"); + assert_eq!( + gpu, cpu_u64, + "device SHIFT table must be byte-identical to the CPU fill" + ); +} + +/// LT device fill: dedup rides the permutation-invariant ALU bus, so the row +/// ORDER is non-deterministic (HashMap iteration). Validate as a MULTISET — +/// the set of real rows (μ>0), incl. summed multiplicities, must match the CPU +/// `generate_lt_trace`. Covers signed/unsigned, invert, MSBs, and duplicates +/// (μ>1). Skips cleanly with no GPU. +#[test] +fn gpu_lt_fill_matches_cpu_multiset() { + if math_cuda::device::backend().is_err() { + eprintln!("skipping gpu_lt_fill_matches_cpu_multiset: no CUDA backend"); + return; + } + // Raw ops with deliberate duplicates (some pushed twice → μ=2). + let mut raw = Vec::new(); + for i in 0..800u64 { + let lhs = i.wrapping_mul(0x9E37_79B9_7F4A_7C15) ^ (i << 50); + let rhs = i.wrapping_mul(0x1234_5678_9ABC_DEF1).rotate_left(17); + let signed = i % 2 == 0; + let invert = i % 3 == 0; + let op = lt::LtOperation::new_with_invert(lhs, rhs, signed, invert); + raw.push(op.clone()); + if i % 4 == 0 { + raw.push(op); // duplicate → multiplicity 2 + } + } + + let ncols = math_cuda::trace_cpu::LT_NCOLS; + // Extract the real (μ>0) rows of a row-major u64 buffer as a sorted multiset. + let real_rows = |flat: &[u64]| -> Vec> { + let mut rows: Vec> = flat + .chunks(ncols) + .filter(|row| row[lt::cols::MU] > 0) + .map(|row| row.to_vec()) + .collect(); + rows.sort(); + rows + }; + + let cpu_table = lt::generate_lt_trace(&raw); + let (cpu_fe, w) = cpu_table.main_data_row_major(); + assert_eq!(w, ncols); + let cpu_u64: Vec = cpu_fe + .iter() + .map(|e| unsafe { *(e.value() as *const u64) }) + .collect(); + + // Dedup on the host exactly like `gpu_build_lt_tables`, then device-fill. + let mut map: std::collections::HashMap = HashMap::new(); + for op in &raw { + *map.entry(op.clone()).or_insert(0) += 1; + } + let unique: Vec<(lt::LtOperation, u64)> = map.into_iter().collect(); + let n = unique.len(); + let num_rows = n.next_power_of_two().max(4); + let mut packed = Vec::with_capacity(n * math_cuda::trace_cpu::LT_STRIDE); + for (op, mult) in &unique { + packed.extend_from_slice(&crate::tables::gpu_trace::pack_lt_op(op, *mult)); + } + let gpu_u64 = math_cuda::trace_cpu::gpu_build_lt_trace_host(&packed, n, num_rows) + .expect("device LT build must run on a box with a CUDA backend"); + + assert_eq!( + real_rows(&gpu_u64), + real_rows(&cpu_u64), + "device LT rows must match the CPU fill as a multiset" + ); +} + +/// EQ device fill: like LT, dedup rides the permutation-invariant ALU bus, so +/// the row ORDER is non-deterministic (HashMap iteration). Validate as a +/// MULTISET — the real rows (μ>0), incl. summed multiplicities, must match the +/// CPU `generate_eq_trace`. Covers a==b and a!=b, invert, high words, and +/// duplicates (μ>1). Skips cleanly with no GPU. +#[test] +fn gpu_eq_fill_matches_cpu_multiset() { + if math_cuda::device::backend().is_err() { + eprintln!("skipping gpu_eq_fill_matches_cpu_multiset: no CUDA backend"); + return; + } + // Raw ops with equal/unequal operands, high words, and deliberate duplicates. + let mut raw = Vec::new(); + for i in 0..800u64 { + let a = i.wrapping_mul(0x9E37_79B9_7F4A_7C15) ^ (i << 40); + // Every 5th op has a == b to exercise the eq=1 path. + let b = if i % 5 == 0 { + a + } else { + i.wrapping_mul(0x1234_5678_9ABC_DEF1).rotate_left(11) + }; + let invert = i % 3 == 0; + let op = eq::EqOperation::new(a, b, invert); + raw.push(op.clone()); + if i % 4 == 0 { + raw.push(op); // duplicate → multiplicity 2 + } + } + + let ncols = math_cuda::trace_cpu::EQ_NCOLS; + // Extract the real (μ>0) rows of a row-major u64 buffer as a sorted multiset. + let real_rows = |flat: &[u64]| -> Vec> { + let mut rows: Vec> = flat + .chunks(ncols) + .filter(|row| row[eq::cols::MU] > 0) + .map(|row| row.to_vec()) + .collect(); + rows.sort(); + rows + }; + + let cpu_table = eq::generate_eq_trace(&raw); + let (cpu_fe, w) = cpu_table.main_data_row_major(); + assert_eq!(w, ncols); + let cpu_u64: Vec = cpu_fe + .iter() + .map(|e| unsafe { *(e.value() as *const u64) }) + .collect(); + + // Dedup on the host exactly like `gpu_build_eq_tables`, then device-fill. + let mut map: std::collections::HashMap = HashMap::new(); + for op in &raw { + *map.entry(op.clone()).or_insert(0) += 1; + } + let unique: Vec<(eq::EqOperation, u64)> = map.into_iter().collect(); + let n = unique.len(); + let num_rows = n.next_power_of_two().max(4); + let mut packed = Vec::with_capacity(n * math_cuda::trace_cpu::EQ_STRIDE); + for (op, mult) in &unique { + packed.extend_from_slice(&crate::tables::gpu_trace::pack_eq_op(op, *mult)); + } + let gpu_u64 = math_cuda::trace_cpu::gpu_build_eq_trace_host(&packed, n, num_rows) + .expect("device EQ build must run on a box with a CUDA backend"); + + assert_eq!( + real_rows(&gpu_u64), + real_rows(&cpu_u64), + "device EQ rows must match the CPU fill as a multiset" + ); +} + +/// BYTEWISE device fill: like LT/EQ, dedup rides the permutation-invariant ALU +/// bus, so the row ORDER is non-deterministic. Validate as a MULTISET — the real +/// rows (μ>0), incl. summed multiplicities, must match `generate_bytewise_trace`. +/// Covers AND/OR/XOR, full 64-bit operands, and duplicates (μ>1). Skips cleanly +/// with no GPU. +#[test] +fn gpu_bytewise_fill_matches_cpu_multiset() { + use crate::tables::types::alu_op; + if math_cuda::device::backend().is_err() { + eprintln!("skipping gpu_bytewise_fill_matches_cpu_multiset: no CUDA backend"); + return; + } + // Raw ops cycling AND/OR/XOR, full-word operands, with deliberate duplicates. + let mut raw = Vec::new(); + for i in 0..900u64 { + let a = i.wrapping_mul(0x9E37_79B9_7F4A_7C15) ^ (i << 33); + let b = i.wrapping_mul(0x1234_5678_9ABC_DEF1).rotate_left(19); + let op = [alu_op::AND, alu_op::OR, alu_op::XOR][(i % 3) as usize]; + let bw = bytewise::BytewiseOperation::new(a, b, op); + raw.push(bw.clone()); + if i % 4 == 0 { + raw.push(bw); // duplicate → multiplicity 2 + } + } + + let ncols = math_cuda::trace_cpu::BYTEWISE_NCOLS; + // Extract the real (μ>0) rows of a row-major u64 buffer as a sorted multiset. + let real_rows = |flat: &[u64]| -> Vec> { + let mut rows: Vec> = flat + .chunks(ncols) + .filter(|row| row[bytewise::cols::MU] > 0) + .map(|row| row.to_vec()) + .collect(); + rows.sort(); + rows + }; + + let cpu_table = bytewise::generate_bytewise_trace(&raw); + let (cpu_fe, w) = cpu_table.main_data_row_major(); + assert_eq!(w, ncols); + let cpu_u64: Vec = cpu_fe + .iter() + .map(|e| unsafe { *(e.value() as *const u64) }) + .collect(); + + // Dedup on the host exactly like `gpu_build_bytewise_tables`, then device-fill. + let mut map: std::collections::HashMap = HashMap::new(); + for op in &raw { + *map.entry(op.clone()).or_insert(0) += 1; + } + let unique: Vec<(bytewise::BytewiseOperation, u64)> = map.into_iter().collect(); + let n = unique.len(); + let num_rows = n.next_power_of_two().max(4); + let mut packed = Vec::with_capacity(n * math_cuda::trace_cpu::BYTEWISE_STRIDE); + for (op, mult) in &unique { + packed.extend_from_slice(&crate::tables::gpu_trace::pack_bytewise_op(op, *mult)); + } + let gpu_u64 = math_cuda::trace_cpu::gpu_build_bytewise_trace_host(&packed, n, num_rows) + .expect("device BYTEWISE build must run on a box with a CUDA backend"); + + assert_eq!( + real_rows(&gpu_u64), + real_rows(&cpu_u64), + "device BYTEWISE rows must match the CPU fill as a multiset" + ); +} + +/// MUL device fill: like the other ALU tables, dedup rides the +/// permutation-invariant ALU bus, so the row ORDER is non-deterministic. +/// Validate as a MULTISET — the real rows (mu_lo>0 or mu_hi>0), incl. the split +/// multiplicities, must match `generate_mul_trace`. Covers all four +/// signed/unsigned combos, negative operands, the 128-bit product + raw_product +/// convolution, and lo/hi (wants_hi) requests with duplicates. Skips cleanly +/// with no GPU. +#[test] +fn gpu_mul_fill_matches_cpu_multiset() { + if math_cuda::device::backend().is_err() { + eprintln!("skipping gpu_mul_fill_matches_cpu_multiset: no CUDA backend"); + return; + } + // Raw (op, wants_hi) pairs: varied operands, all sign combos, both lo and hi + // requests, with deliberate duplicates so mu_lo/mu_hi accumulate. + let mut raw = Vec::new(); + for i in 0..800u64 { + let lhs = i.wrapping_mul(0x9E37_79B9_7F4A_7C15) ^ (i << 40); + let rhs = i.wrapping_mul(0x1234_5678_9ABC_DEF1).rotate_left(23); + let lhs_signed = i % 2 == 0; + let rhs_signed = i % 3 == 0; + let op = mul::MulOperation::new(lhs, lhs_signed, rhs, rhs_signed); + raw.push((op.clone(), i % 2 == 0)); + if i % 3 == 0 { + raw.push((op.clone(), true)); // extra hi request + } + if i % 5 == 0 { + raw.push((op, false)); // extra lo request (duplicate) + } + } + // Explicit sign edge cases: i64::MIN, -1, treated signed and unsigned. + for &(a, b) in &[ + (0x8000_0000_0000_0000u64, 0xFFFF_FFFF_FFFF_FFFFu64), + (0xFFFF_FFFF_FFFF_FFFFu64, 0x8000_0000_0000_0000u64), + ] { + raw.push((mul::MulOperation::new(a, true, b, true), false)); + raw.push((mul::MulOperation::new(a, false, b, false), true)); + } + + let ncols = math_cuda::trace_cpu::MUL_NCOLS; + // Real rows: mu_lo>0 or mu_hi>0. + let real_rows = |flat: &[u64]| -> Vec> { + let mut rows: Vec> = flat + .chunks(ncols) + .filter(|row| row[mul::cols::MU_LO] > 0 || row[mul::cols::MU_HI] > 0) + .map(|row| row.to_vec()) + .collect(); + rows.sort(); + rows + }; + + let cpu_table = mul::generate_mul_trace(&raw); + let (cpu_fe, w) = cpu_table.main_data_row_major(); + assert_eq!(w, ncols); + let cpu_u64: Vec = cpu_fe + .iter() + .map(|e| unsafe { *(e.value() as *const u64) }) + .collect(); + + // Dedup on the host exactly like `gpu_build_mul_tables`, then device-fill. + let mut map: std::collections::HashMap = HashMap::new(); + for (op, wants_hi) in &raw { + let e = map.entry(op.clone()).or_insert((0, 0)); + if *wants_hi { + e.1 += 1; + } else { + e.0 += 1; + } + } + let unique: Vec<(mul::MulOperation, (u64, u64))> = map.into_iter().collect(); + let n = unique.len(); + let num_rows = n.next_power_of_two().max(4); + let mut packed = Vec::with_capacity(n * math_cuda::trace_cpu::MUL_STRIDE); + for (op, (mu_lo, mu_hi)) in &unique { + packed.extend_from_slice(&crate::tables::gpu_trace::pack_mul_op(op, *mu_lo, *mu_hi)); + } + let gpu_u64 = math_cuda::trace_cpu::gpu_build_mul_trace_host(&packed, n, num_rows) + .expect("device MUL build must run on a box with a CUDA backend"); + + assert_eq!( + real_rows(&gpu_u64), + real_rows(&cpu_u64), + "device MUL rows must match the CPU fill as a multiset" + ); +} + +/// DVRM device fill: like the other ALU tables, dedup rides the +/// permutation-invariant ALU bus, so the row ORDER is non-deterministic. +/// Validate as a MULTISET — the real rows (mu_q>0 or mu_r>0), incl. the split +/// multiplicities, must match `generate_dvrm_trace`. Covers signed/unsigned, +/// division-by-zero, the MIN/-1 signed overflow, negative operands, and +/// quotient/remainder (wants_remainder) requests with duplicates. Skips cleanly +/// with no GPU. +#[test] +fn gpu_dvrm_fill_matches_cpu_multiset() { + if math_cuda::device::backend().is_err() { + eprintln!("skipping gpu_dvrm_fill_matches_cpu_multiset: no CUDA backend"); + return; + } + // Raw (op, wants_remainder) pairs: varied operands (incl. periodic d==0), + // both signednesses, q and r requests, with deliberate duplicates. + let mut raw = Vec::new(); + for i in 0..800u64 { + let n = i.wrapping_mul(0x9E37_79B9_7F4A_7C15) ^ (i << 40); + let d = if i % 11 == 0 { + 0 + } else { + i.wrapping_mul(0x1234_5678_9ABC_DEF1).rotate_left(29) + }; + let signed = i % 2 == 0; + let op = dvrm::DvrmOperation::new(n, d, signed); + raw.push((op.clone(), i % 2 == 0)); + if i % 3 == 0 { + raw.push((op.clone(), true)); // extra remainder request + } + if i % 5 == 0 { + raw.push((op, false)); // extra quotient request (duplicate) + } + } + // Explicit edge cases: signed overflow (MIN/-1), div-by-zero, MIN numerator, + // -1 denominator. + let min = 0x8000_0000_0000_0000u64; + let neg1 = 0xFFFF_FFFF_FFFF_FFFFu64; + for &(n, d, s) in &[ + (min, neg1, true), // signed overflow + (min, neg1, false), // unsigned: /(2^64-1), no overflow + (123u64, 0u64, true), // div-by-zero, signed + (123u64, 0u64, false), // div-by-zero, unsigned + (min, 7u64, true), // negative numerator + (100u64, neg1, true), // negative denominator (n != MIN) + ] { + raw.push((dvrm::DvrmOperation::new(n, d, s), false)); + raw.push((dvrm::DvrmOperation::new(n, d, s), true)); + } + + let ncols = math_cuda::trace_cpu::DVRM_NCOLS; + // Real rows: mu_q>0 or mu_r>0. + let real_rows = |flat: &[u64]| -> Vec> { + let mut rows: Vec> = flat + .chunks(ncols) + .filter(|row| row[dvrm::cols::MU_Q] > 0 || row[dvrm::cols::MU_R] > 0) + .map(|row| row.to_vec()) + .collect(); + rows.sort(); + rows + }; + + let cpu_table = dvrm::generate_dvrm_trace(&raw); + let (cpu_fe, w) = cpu_table.main_data_row_major(); + assert_eq!(w, ncols); + let cpu_u64: Vec = cpu_fe + .iter() + .map(|e| unsafe { *(e.value() as *const u64) }) + .collect(); + + // Dedup on the host exactly like `gpu_build_dvrm_tables`, then device-fill. + let mut map: std::collections::HashMap = HashMap::new(); + for (op, wants_remainder) in &raw { + let e = map.entry(op.clone()).or_insert((0, 0)); + if *wants_remainder { + e.1 += 1; + } else { + e.0 += 1; + } + } + let unique: Vec<(dvrm::DvrmOperation, (u64, u64))> = map.into_iter().collect(); + let n = unique.len(); + let num_rows = n.next_power_of_two().max(4); + let mut packed = Vec::with_capacity(n * math_cuda::trace_cpu::DVRM_STRIDE); + for (op, (mu_q, mu_r)) in &unique { + packed.extend_from_slice(&crate::tables::gpu_trace::pack_dvrm_op(op, *mu_q, *mu_r)); + } + let gpu_u64 = math_cuda::trace_cpu::gpu_build_dvrm_trace_host(&packed, n, num_rows) + .expect("device DVRM build must run on a box with a CUDA backend"); + + assert_eq!( + real_rows(&gpu_u64), + real_rows(&cpu_u64), + "device DVRM rows must match the CPU fill as a multiset" + ); +} + +/// BRANCH device fill: a permutation-invariant lookup table (dedup + summed +/// multiplicity), so the row ORDER is non-deterministic. Validate as a MULTISET +/// — the real rows (μ>0) must match `generate_branch_trace`. Covers JALR (base = +/// register) vs PC-relative, wrapping add, odd offsets (so LSB masking differs +/// from the unmasked low byte), high address bits, and duplicates (μ>1). Skips +/// cleanly with no GPU. +#[test] +fn gpu_branch_fill_matches_cpu_multiset() { + if math_cuda::device::backend().is_err() { + eprintln!("skipping gpu_branch_fill_matches_cpu_multiset: no CUDA backend"); + return; + } + let mut raw = Vec::new(); + for i in 0..800u64 { + let pc = 0x1000u64.wrapping_add(i.wrapping_mul(4)) ^ (i << 34); + // Odd offsets so `next_pc` (LSB masked) differs from the unmasked low + // byte; high bits set to exercise the wrapping add. + let offset = i.wrapping_mul(0x9E37_79B9_7F4A_7C15) | 1; + let register = i.wrapping_mul(0x1234_5678_9ABC_DEF1).rotate_left(7); + let jalr = i % 2 == 0; + let op = branch::BranchOperation::new(pc, offset, register, jalr); + raw.push(op.clone()); + if i % 4 == 0 { + raw.push(op); // duplicate → μ=2 + } + } + + let ncols = math_cuda::trace_cpu::BRANCH_NCOLS; + let real_rows = |flat: &[u64]| -> Vec> { + let mut rows: Vec> = flat + .chunks(ncols) + .filter(|row| row[branch::cols::MU] > 0) + .map(|row| row.to_vec()) + .collect(); + rows.sort(); + rows + }; + + let cpu_table = branch::generate_branch_trace(&raw); + let (cpu_fe, w) = cpu_table.main_data_row_major(); + assert_eq!(w, ncols); + let cpu_u64: Vec = cpu_fe + .iter() + .map(|e| unsafe { *(e.value() as *const u64) }) + .collect(); + + // Dedup on the host exactly like `gpu_build_branch_tables`, then device-fill. + let mut map: std::collections::HashMap = HashMap::new(); + for op in &raw { + *map.entry(op.clone()).or_insert(0) += 1; + } + let unique: Vec<(branch::BranchOperation, u64)> = map.into_iter().collect(); + let n = unique.len(); + let num_rows = n.next_power_of_two().max(4); + let mut packed = Vec::with_capacity(n * math_cuda::trace_cpu::BRANCH_STRIDE); + for (op, mult) in &unique { + packed.extend_from_slice(&crate::tables::gpu_trace::pack_branch_op(op, *mult)); + } + let gpu_u64 = math_cuda::trace_cpu::gpu_build_branch_trace_host(&packed, n, num_rows) + .expect("device BRANCH build must run on a box with a CUDA backend"); + + assert_eq!( + real_rows(&gpu_u64), + real_rows(&cpu_u64), + "device BRANCH rows must match the CPU fill as a multiset" + ); +} + +/// CPU32 device fill must be byte-identical to `cpu32::generate_cpu32_trace`. +/// Per-row (μ=1, no dedup), so full byte parity holds. Covers signed/unsigned +/// (alu_flags bit 5), negative rv1/rv2/res (sign-extension into arg1/arg2/rvd), +/// imm-vs-rv2 operands, flag/register combinations, high words, and padding +/// (n=300 < 512). Skips cleanly with no GPU. +#[test] +fn gpu_cpu32_fill_matches_cpu() { + if math_cuda::device::backend().is_err() { + eprintln!("skipping gpu_cpu32_fill_matches_cpu: no CUDA backend"); + return; + } + let mut ops = Vec::new(); + for i in 0..300u64 { + let signed = i % 2 == 0; + // alu_flags: bit 5 = signed; low bits carry a varied opcode. + let alu_flags = ((i % 20) as u8) | if signed { 1 << 5 } else { 0 }; + // Alternate imm-driven vs rv2-driven arg2 (decode assumption: one nonzero). + let use_imm = i % 3 == 0; + let rv2 = if use_imm { + 0 + } else { + i.wrapping_mul(0xDEAD_BEEF).rotate_left(9) + }; + let imm = if use_imm { + i.wrapping_mul(0x9E37_79B9) | (i << 40) + } else { + 0 + }; + ops.push(cpu32::Cpu32Operation { + timestamp: 4 * i + 8 + (i << 34), + pc: 0x1000 + i * 4 + ((i % 3) << 33), + rs1: (i % 32) as u8, + read_register1: i % 3 != 0, + rv1: i.wrapping_mul(0x1234_5678_9ABC) ^ (i << 31), // exercise bit 31 + rs2: ((i + 5) % 32) as u8, + read_register2: !use_imm, + rv2, + imm, + res: i.wrapping_mul(0xABCD_1234) ^ (i << 31), + rd: ((i + 7) % 32) as u8, + write_register: i % 4 != 0, + alu: i % 5 != 0, + alu_flags, + add: i % 5 == 1, + sub: i % 5 == 2, + half_instruction_length: if i % 2 == 0 { 1 } else { 2 }, + }); + } + let n = ops.len(); + let num_rows = n.next_power_of_two().max(4); + + let cpu_table = cpu32::generate_cpu32_trace(&ops); + let (cpu_fe, w) = cpu_table.main_data_row_major(); + assert_eq!(w, math_cuda::trace_cpu::CPU32_NCOLS); + let cpu_u64: Vec = cpu_fe + .iter() + .map(|e| unsafe { *(e.value() as *const u64) }) + .collect(); + + let mut packed = Vec::with_capacity(n * math_cuda::trace_cpu::CPU32_STRIDE); + for op in &ops { + packed.extend_from_slice(&crate::tables::gpu_trace::pack_cpu32_op(op)); + } + let gpu_u64 = math_cuda::trace_cpu::gpu_build_cpu32_trace_host(&packed, n, num_rows) + .expect("device CPU32 build must run on a box with a CUDA backend"); + + assert_eq!(gpu_u64.len(), num_rows * math_cuda::trace_cpu::CPU32_NCOLS); + assert_eq!( + gpu_u64, cpu_u64, + "device CPU32 table must be byte-identical to the CPU fill" + ); +} + +/// MEMW (general/unaligned) device fill must be byte-identical to +/// `memw::generate_memw_trace`. Per-row (no dedup). Covers memory and register +/// accesses, widths 1/2/4/8, read/write, base addresses that straddle the 2^32 +/// boundary (so `carry[i]` fires), distinct per-byte old_timestamps (the +/// split-timestamp path), and full-u32 value/old limbs (exercises both halves of +/// the 2×u32/u64 packing). Skips cleanly with no GPU. +#[test] +fn gpu_memw_fill_matches_cpu() { + if math_cuda::device::backend().is_err() { + eprintln!("skipping gpu_memw_fill_matches_cpu: no CUDA backend"); + return; + } + let mut ops = Vec::new(); + for i in 0..400u64 { + let width = [1u8, 2, 4, 8][(i % 4) as usize]; + let is_read = i % 2 == 0; + let is_register = i % 7 == 0; + // Low word near 2^32 so `base_lo + (i+1)` overflows for some rows. + let base = (0x1_0000_0000u64 * (i % 3)) + (0xFFFF_FFF8u64 - (i % 16)) + i; + let value = [ + (i as u32).wrapping_mul(2_654_435_761), + (i as u32) ^ 0xABCD_1234, + i as u32, + 0xDEAD_0000 | (i as u32 & 0xFFFF), + 7, + 0, + (i as u32).wrapping_add(99), + (i as u32) & 0xFF, + ]; + let old = [ + (i as u32).wrapping_add(3), + (i as u32).wrapping_mul(17), + 0, + (i as u32) ^ 0x0BAD_F00D, + (i as u32).wrapping_add(5), + (i as u32).wrapping_mul(23), + 0, + (i as u32).wrapping_add(7), + ]; + let ts = 4 * i + 100; + // Distinct per-byte old_timestamps (the unaligned split-timestamp path). + let mut old_ts = [0u64; 8]; + for (j, o) in old_ts.iter_mut().enumerate() { + *o = (4 * i + 3 + j as u64) ^ ((j as u64) << 33); + } + ops.push( + memw::MemwOperation::new(is_register, base, value, ts, width, is_read) + .with_old(old, old_ts), + ); + } + let n = ops.len(); + let num_rows = n.next_power_of_two().max(4); + + let cpu_table = memw::generate_memw_trace(&ops); + let (cpu_fe, w) = cpu_table.main_data_row_major(); + assert_eq!(w, math_cuda::trace_cpu::MEMW_NCOLS); + let cpu_u64: Vec = cpu_fe + .iter() + .map(|e| unsafe { *(e.value() as *const u64) }) + .collect(); + + let mut packed = Vec::with_capacity(n * math_cuda::trace_cpu::MEMW_STRIDE); + for op in &ops { + packed.extend_from_slice(&crate::tables::gpu_trace::pack_memw_op(op)); + } + let gpu_u64 = math_cuda::trace_cpu::gpu_build_memw_trace_host(&packed, n, num_rows) + .expect("device MEMW build must run on a box with a CUDA backend"); + + assert_eq!(gpu_u64.len(), num_rows * math_cuda::trace_cpu::MEMW_NCOLS); + assert_eq!( + gpu_u64, cpu_u64, + "device MEMW table must be byte-identical to the CPU fill" + ); +} diff --git a/prover/src/tests/mod.rs b/prover/src/tests/mod.rs index c3600c359..e43337ec9 100644 --- a/prover/src/tests/mod.rs +++ b/prover/src/tests/mod.rs @@ -44,6 +44,8 @@ pub mod ecdas_tests; pub mod ecsm_tests; #[cfg(test)] pub mod eq_tests; +#[cfg(all(test, feature = "cuda"))] +pub mod gpu_fill_tests; #[cfg(test)] pub mod keccak_rnd_tests; #[cfg(test)]