diff --git a/crypto/math-cuda/kernels/trace_cpu.cu b/crypto/math-cuda/kernels/trace_cpu.cu index 0f33e2f1d..6e12a0a03 100644 --- a/crypto/math-cuda/kernels/trace_cpu.cu +++ b/crypto/math-cuda/kernels/trace_cpu.cu @@ -467,6 +467,489 @@ 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 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 24720eb8a..554d3cd5b 100644 --- a/crypto/math-cuda/src/device.rs +++ b/crypto/math-cuda/src/device.rs @@ -198,6 +198,13 @@ 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 branch_fill: CudaFunction, + pub cpu32_fill: CudaFunction, + pub memw_fill: CudaFunction, pub memw_register_fill: CudaFunction, // constraint_interp.ptx pub constraint_interp_kernel: CudaFunction, @@ -401,6 +408,13 @@ 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")?, + 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")?, constraint_interp_kernel: constraint_interp .load_function("constraint_interp_kernel")?, 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 1b5ca6420..b4131cead 100644 --- a/crypto/math-cuda/src/trace_cpu.rs +++ b/crypto/math-cuda/src/trace_cpu.rs @@ -152,6 +152,34 @@ 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; +/// 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 @@ -197,6 +225,27 @@ 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 +} +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`). @@ -234,6 +283,302 @@ 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 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 b89132d7f..5b13ee6c2 100644 --- a/crypto/stark/src/gpu_lde.rs +++ b/crypto/stark/src/gpu_lde.rs @@ -622,6 +622,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 c93c2b523..6a43a8240 100644 --- a/prover/src/tables/gpu_trace.rs +++ b/prover/src/tables/gpu_trace.rs @@ -14,12 +14,18 @@ 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}; use super::shift::{self, ShiftOperation}; use super::store::{self, StoreOperation}; use super::types::{GoldilocksExtension, GoldilocksField}; @@ -428,3 +434,394 @@ pub(crate) fn gpu_build_lt_tables( // Each chunk dedups independently (matching `generate_lt_trace` per chunk). gpu_build_tables(ops, max_rows, build_lt_chunk) } + +// ============================================================================= +// 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>> { + // Each chunk dedups independently (matching `generate_eq_trace` per chunk). + gpu_build_tables(ops, max_rows, build_eq_chunk) +} + +// ============================================================================= +// 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>> { + // Each chunk dedups independently (matching `generate_bytewise_trace` per chunk). + gpu_build_tables(ops, max_rows, build_bytewise_chunk) +} + +// ============================================================================= +// 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>> { + // Each chunk dedups independently (matching `generate_mul_trace` per chunk). + gpu_build_tables(ops, max_rows, build_mul_chunk) +} + +// ============================================================================= +// 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>> { + // Each chunk dedups independently (matching `generate_dvrm_trace` per chunk). + gpu_build_tables(ops, max_rows, build_dvrm_chunk) +} + +// ============================================================================= +// 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>> { + // Each chunk dedups independently (matching `generate_branch_trace` per chunk). + gpu_build_tables(ops, max_rows, build_branch_chunk) +} + +// ============================================================================= +// 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>> { + gpu_build_tables(ops, max_rows, build_cpu32_chunk) +} + +// ============================================================================= +// 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>> { + 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 1c864dfc8..aaff3ad40 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, @@ -3382,6 +3401,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 +3429,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, @@ -3400,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, @@ -3412,6 +3488,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 +3516,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, @@ -3458,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, diff --git a/prover/src/tests/gpu_fill_tests.rs b/prover/src/tests/gpu_fill_tests.rs index ddc11ff34..68ae13c26 100644 --- a/prover/src/tests/gpu_fill_tests.rs +++ b/prover/src/tests/gpu_fill_tests.rs @@ -9,7 +9,10 @@ use std::collections::HashMap; use crate::tables::cpu::CpuOperation; use crate::tables::memw::MemwOperation; -use crate::tables::{cpu, load, lt, memw_aligned, memw_register, shift, store}; +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 @@ -463,3 +466,539 @@ fn gpu_lt_fill_matches_cpu_multiset() { "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" + ); +}