diff --git a/crypto/math-cuda/build.rs b/crypto/math-cuda/build.rs index 316a9c7ed..dcfc65c93 100644 --- a/crypto/math-cuda/build.rs +++ b/crypto/math-cuda/build.rs @@ -121,5 +121,6 @@ fn main() { compile_ptx("fri.cu", "fri.ptx", have_nvcc); compile_ptx("inverse.cu", "inverse.ptx", have_nvcc); compile_ptx("logup.cu", "logup.ptx", have_nvcc); + compile_ptx("trace_cpu.cu", "trace_cpu.ptx", have_nvcc); compile_ptx("constraint_interp.cu", "constraint_interp.ptx", have_nvcc); } diff --git a/crypto/math-cuda/kernels/trace_cpu.cu b/crypto/math-cuda/kernels/trace_cpu.cu new file mode 100644 index 000000000..0f33e2f1d --- /dev/null +++ b/crypto/math-cuda/kernels/trace_cpu.cu @@ -0,0 +1,503 @@ +// On-GPU CPU trace-table fill: one thread per row, row-major output +// `out[row*NCOLS + col]` (the layout the device-input LDE seam consumes). +// +// Mirrors `prover/src/tables/cpu.rs::generate_cpu_trace`. The `CpuOperation` +// already carries the resolved values (rv1/rv2/arg2/res/next_pc/branch_cond), +// so the fill is pure bit-slicing — every column limb is < 2^32 (DWordWL lo/hi, +// DWordHL 16-bit halves, bytes, bools), so NO Goldilocks reduction is needed. +// +// Packed input, stride STRIDE u64 per op: +// [0]=timestamp [1]=pc [2]=imm [3]=next_pc [4]=rvd [5]=rv1 [6]=rv2 [7]=arg2 +// [8]=res [9]=flags [10]=bytes +// flags bits: 0 word_instr, 1 read_register1(raw), 2 read_register2(raw), +// 3 write_register(raw), 4 alu, 5 add, 6 sub, 7 memory, 8 branch, 9 ecall, +// 10 branch_cond. +// bytes: b0 rs1, b1 rs2, b2 rd, b3 half_instruction_length, b4 alu_flags, +// b5 mem_flags. +// +// The output buffer MUST be pre-zeroed (alloc_zeros); the kernel writes only the +// non-zero cells of real rows, and (TIMESTAMP, PC_0=1, NEXT_PC_0=1) for padding. + +#include + +#define NCOLS 38u +#define STRIDE 11u + +extern "C" __global__ void trace_cpu_fill(const uint64_t *ops, // n * STRIDE + uint64_t n, // real op count + uint64_t num_rows, // padded pow2 + uint64_t last_ts, // last real ts (0 if n==0) + uint64_t *out) // num_rows * NCOLS, zeroed +{ + uint64_t r = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x; + if (r >= num_rows) return; + uint64_t base = r * NCOLS; + + // Padding rows: continue the +4 timestamp cadence; PC_0 = NEXT_PC_0 = 1 + // (odd, unreachable); every other column stays zero. + if (r >= n) { + uint64_t j = r - n + 1; + out[base + 0] = last_ts + 4 * j; // TIMESTAMP + out[base + 1] = 1; // PC_0 + out[base + 21] = 1; // NEXT_PC_0 + return; + } + + const uint64_t *op = ops + r * STRIDE; + uint64_t ts = op[0]; + uint64_t pc = op[1]; + uint64_t imm = op[2]; + uint64_t npc = op[3]; + uint64_t rvd = op[4]; + uint64_t rv1 = op[5]; + uint64_t rv2 = op[6]; + uint64_t arg2 = op[7]; + uint64_t res = op[8]; + uint64_t fl = op[9]; + uint64_t by = op[10]; + + uint64_t word = fl & 1u; + uint64_t rr1 = (fl >> 1) & 1u; + uint64_t rr2 = (fl >> 2) & 1u; + uint64_t wr = (fl >> 3) & 1u; + uint64_t alu = (fl >> 4) & 1u; + uint64_t add = (fl >> 5) & 1u; + uint64_t sub = (fl >> 6) & 1u; + uint64_t mem = (fl >> 7) & 1u; + uint64_t br = (fl >> 8) & 1u; + uint64_t ec = (fl >> 9) & 1u; + uint64_t bcond = (fl >> 10) & 1u; + + uint64_t rs1 = by & 0xFFu; + uint64_t rs2 = (by >> 8) & 0xFFu; + uint64_t rd = (by >> 16) & 0xFFu; + uint64_t hil = (by >> 24) & 0xFFu; + uint64_t aluf = (by >> 32) & 0xFFu; + uint64_t memf = (by >> 40) & 0xFFu; + + // `effective(flag) = !word && flag`. Word-delegate rows suppress operational + // data (CPU32 owns it) — only the PC-advancing columns are set. + uint64_t nw = 1u - word; // !word + uint64_t rs1c = word ? 0u : rs1; + uint64_t rs2c = word ? 0u : rs2; + uint64_t rdc = word ? 0u : rd; + uint64_t immc = word ? 0u : imm; + uint64_t rvdc = word ? 0u : rvd; + uint64_t rv1c = word ? 0u : rv1; + uint64_t rv2c = word ? 0u : rv2; + uint64_t arg2c = word ? 0u : arg2; + uint64_t resc = word ? 0u : res; + + out[base + 0] = ts; // TIMESTAMP + out[base + 1] = pc & 0xFFFFFFFFu; // PC_0 + out[base + 2] = pc >> 32; // PC_1 + out[base + 3] = rs1c; // RS1 + out[base + 4] = rs2c; // RS2 + out[base + 5] = rdc; // RD + out[base + 6] = nw & rr1 & (rs1 != 0u); // READ_REGISTER1 = eff(rr1 && rs1!=0) + out[base + 7] = nw & rr2 & (rs2 != 0u); // READ_REGISTER2 + out[base + 8] = nw & wr & (rd != 0u); // WRITE_REGISTER + out[base + 9] = immc & 0xFFFFFFFFu; // IMM_0 + out[base + 10] = immc >> 32; // IMM_1 + out[base + 11] = hil; // HALF_INSTRUCTION_LENGTH (unmasked) + out[base + 12] = word; // WORD_INSTR + out[base + 13] = nw & alu; // ALU + out[base + 14] = word ? 0u : aluf; // ALU_FLAGS + out[base + 15] = nw & add; // ADD + out[base + 16] = nw & sub; // SUB + out[base + 17] = nw & mem; // MEMORY + out[base + 18] = word ? 0u : memf; // MEM_FLAGS + out[base + 19] = nw & br; // BRANCH + out[base + 20] = nw & ec; // ECALL + out[base + 21] = npc & 0xFFFFFFFFu; // NEXT_PC_0 + out[base + 22] = npc >> 32; // NEXT_PC_1 + out[base + 23] = rvdc & 0xFFFFFFFFu; // RVD_0 + out[base + 24] = rvdc >> 32; // RVD_1 + + // Inline-PC coordination columns. + uint64_t pcdr = nw & rr1 & (rs1 == 255u); // PC_DOUBLE_READ + uint64_t ts_lo = ts & 0xFFFFFFFFu; + uint64_t borrow = (pcdr == 0u && ts_lo < 3u) ? 1u : 0u; // PREV_PC_TIMESTAMP_BORROW + out[base + 25] = borrow; + out[base + 26] = pcdr; + out[base + 27] = rv1c & 0xFFFFFFFFu; // RV1_0 + out[base + 28] = rv1c >> 32; // RV1_1 + out[base + 29] = rv2c & 0xFFFFFFFFu; // RV2_0 + out[base + 30] = rv2c >> 32; // RV2_1 + out[base + 31] = arg2c & 0xFFFFFFFFu; // ARG2_0 + out[base + 32] = arg2c >> 32; // ARG2_1 + out[base + 33] = resc & 0xFFFFu; // RES_0 + out[base + 34] = (resc >> 16) & 0xFFFFu; // RES_1 + out[base + 35] = (resc >> 32) & 0xFFFFu; // RES_2 + out[base + 36] = (resc >> 48) & 0xFFFFu; // RES_3 + out[base + 37] = bcond; // BRANCH_COND (unmasked) +} + +// On-GPU MEMW_A (aligned memory) trace fill: one thread per row, row-major +// `out[row*MEMW_A_NCOLS + col]`. Mirrors +// `prover/src/tables/memw_aligned.rs::generate_memw_aligned_trace`. The op is +// already walked (old_value/old_timestamp filled by the memory-model walk), so +// this is pure bit-slicing — every limb is < 2^32 (halves/words/bytes/bools), no +// Goldilocks reduction. Padding rows (r >= n) stay all-zero (pre-zeroed buffer). +// +// Packed input, stride MEMW_A_STRIDE u64 per op: +// [0]=flags (bit0 is_register, bit1 is_read, bits8..16 width) +// [1]=base_address [2]=timestamp [3]=old_timestamp[0] +// [4..8]=value[0..8] packed 2×u32/u64 (value[2i] | value[2i+1]<<32) +// [8..12]=old[0..8] packed 2×u32/u64 +#define MEMW_A_NCOLS 29u +#define MEMW_A_STRIDE 12u + +extern "C" __global__ void memw_aligned_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_A_NCOLS; + if (r >= n) + return; // padding rows all-zero + + const uint64_t *op = ops + r * MEMW_A_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]; + uint64_t old_ts = op[3]; + + out[base + 0] = is_register; // IS_REGISTER + out[base + 1] = addr & 0xFFFFu; // BASE_ADDRESS[0] (low half) + out[base + 2] = (addr >> 16) & 0xFFFFu; // BASE_ADDRESS[1] (mid half) + out[base + 3] = (addr >> 32) & 0xFFFFFFFFu; // BASE_ADDRESS[2] (high word) + // VALUE[0..8] at cols 4..11. + for (int i = 0; i < 4; ++i) { + uint64_t p = op[4 + i]; + out[base + 4 + 2 * i] = p & 0xFFFFFFFFu; + out[base + 4 + 2 * i + 1] = p >> 32; + } + out[base + 12] = ts & 0xFFFFFFFFu; // TIMESTAMP_0 + out[base + 13] = ts >> 32; // TIMESTAMP_1 + out[base + 14] = (width == 2u) ? 1u : 0u; // WRITE2 + out[base + 15] = (width == 4u) ? 1u : 0u; // WRITE4 + out[base + 16] = (width == 8u) ? 1u : 0u; // WRITE8 + // OLD[0..8] at cols 17..24. + for (int i = 0; i < 4; ++i) { + uint64_t p = op[8 + i]; + out[base + 17 + 2 * i] = p & 0xFFFFFFFFu; + out[base + 17 + 2 * i + 1] = p >> 32; + } + out[base + 25] = old_ts & 0xFFFFFFFFu; // OLD_TIMESTAMP_0 + out[base + 26] = old_ts >> 32; // OLD_TIMESTAMP_1 + out[base + 27] = is_read; // MU_READ + out[base + 28] = 1u - is_read; // MU_WRITE +} + +// On-GPU LOAD trace fill (18 cols). Mirrors +// `prover/src/tables/load.rs::generate_load_trace`. Packed stride LOAD_STRIDE: +// [0]=flags (bit0 signed, bits8..16 width) [1]=base_address [2]=timestamp +// [3..7]=res[0..8] packed 2×u32/u64. Padding rows (r>=n) all-zero. +#define LOAD_NCOLS 18u +#define LOAD_STRIDE 7u + +extern "C" __global__ void load_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 * LOAD_NCOLS; + if (r >= n) + return; + + const uint64_t *op = ops + r * LOAD_STRIDE; + uint64_t fl = op[0]; + uint64_t is_signed = fl & 1u; + uint64_t width = (fl >> 8) & 0xFFu; + uint64_t addr = op[1]; + uint64_t ts = op[2]; + + out[base + 0] = addr & 0xFFFFFFFFu; // BASE_ADDRESS_0 + out[base + 1] = addr >> 32; // BASE_ADDRESS_1 + out[base + 2] = ts & 0xFFFFFFFFu; // TIMESTAMP_0 + out[base + 3] = ts >> 32; // TIMESTAMP_1 + out[base + 4] = (width == 2u) ? 1u : 0u; // READ2 + out[base + 5] = (width == 4u) ? 1u : 0u; // READ4 + out[base + 6] = (width == 8u) ? 1u : 0u; // READ8 + out[base + 7] = is_signed; // SIGNED + + uint64_t res[8]; + for (int i = 0; i < 4; ++i) { + uint64_t p = op[3 + i]; + res[2 * i] = p & 0xFFFFFFFFu; + res[2 * i + 1] = p >> 32; + } + for (int i = 0; i < 8; ++i) + out[base + 8 + i] = res[i]; // RES[0..8] + + int bidx = (width == 8u) ? 7 : (width == 4u) ? 3 : (width == 2u) ? 1 : 0; + out[base + 16] = (res[bidx] >> 7) & 1u; // SIGN_BIT + out[base + 17] = 1u; // MU (active row) +} + +// On-GPU STORE trace fill (16 cols). Mirrors +// `prover/src/tables/store.rs::generate_store_trace`. Packed stride STORE_STRIDE: +// [0]=flags (bit0 write2, bit1 write4, bit2 write8) [1]=base_address +// [2]=timestamp [3]=value. Padding rows (r>=n) all-zero. +#define STORE_NCOLS 16u +#define STORE_STRIDE 4u + +extern "C" __global__ void store_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 * STORE_NCOLS; + if (r >= n) + return; + + const uint64_t *op = ops + r * STORE_STRIDE; + uint64_t fl = op[0]; + uint64_t addr = op[1]; + uint64_t ts = op[2]; + uint64_t val = op[3]; + + out[base + 0] = addr & 0xFFFFFFFFu; // BASE_ADDRESS_0 + out[base + 1] = addr >> 32; // BASE_ADDRESS_1 + out[base + 2] = ts & 0xFFFFFFFFu; // TIMESTAMP_0 + out[base + 3] = ts >> 32; // TIMESTAMP_1 + out[base + 4] = fl & 1u; // WRITE2 + out[base + 5] = (fl >> 1) & 1u; // WRITE4 + out[base + 6] = (fl >> 2) & 1u; // WRITE8 + for (int i = 0; i < 8; ++i) + out[base + 7 + i] = (val >> (i * 8)) & 0xFFu; // VALUE[0..8] (DWordBL) + out[base + 15] = 1u; // MU +} + +// On-GPU SHIFT trace fill (29 cols). Mirrors +// `prover/src/tables/shift.rs::{generate_shift_trace, compute_aux}` — the fill +// RECOMPUTES the aux (bit_shift, zbs, x/y half decomposition, limb_shift one-hot, +// shifted DWordHL) from the compact input, so only 3 u64/op are uploaded. No +// dedup: μ = 1 per op. Padding rows (r >= n) set ZBS = 1 (all else 0). +// +// Packed stride SHIFT_STRIDE: [0]=value (4×u16 in_halves) [1]=shift_amount +// [2]=flags (bit0 direction(=right), bit1 signed, bit2 word_instr). + +// HWSL: (halfword << z) & 0xFFFF (z in [0,15]). +__device__ __forceinline__ uint16_t hwsl_(uint16_t h, uint32_t z) { + return z == 0u ? h : (uint16_t)((uint32_t)h << z); +} +// HWSL carry: halfword >> (16 - z). +__device__ __forceinline__ uint16_t hwslc_(uint16_t h, uint32_t z) { + return z == 0u ? (uint16_t)0 : (uint16_t)(h >> (16u - z)); +} + +#define SHIFT_NCOLS 29u +#define SHIFT_STRIDE 3u + +extern "C" __global__ void shift_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 * SHIFT_NCOLS; + if (r >= n) { + out[base + 12] = 1u; // ZBS = 1 on padding rows + return; + } + + const uint64_t *op = ops + r * SHIFT_STRIDE; + uint64_t value = op[0]; + uint64_t shift_amount = op[1]; + uint64_t fl = op[2]; + uint32_t direction = (uint32_t)(fl & 1u); // right + uint32_t is_signed = (uint32_t)((fl >> 1) & 1u); + uint32_t word_instr = (uint32_t)((fl >> 2) & 1u); + + uint16_t in_h[4]; + in_h[0] = (uint16_t)(value & 0xFFFFu); + in_h[1] = (uint16_t)((value >> 16) & 0xFFFFu); + in_h[2] = (uint16_t)((value >> 32) & 0xFFFFu); + in_h[3] = (uint16_t)((value >> 48) & 0xFFFFu); + uint8_t shift = (uint8_t)(shift_amount & 0xFFu); + uint32_t left = 1u - direction; + uint32_t right = direction; + + uint32_t is_negative = (is_signed && ((in_h[3] >> 15) & 1u)) ? 1u : 0u; + uint16_t extension = is_negative ? (uint16_t)0xFFFF : (uint16_t)0; + + uint8_t bit_shift; + if (left) + bit_shift = (uint8_t)(shift & 15u); + else + bit_shift = (uint8_t)((256u - (uint32_t)shift) & 15u); + uint32_t zbs = (bit_shift == 0u) ? 1u : 0u; + + uint16_t x[5] = {0, 0, 0, 0, 0}; + uint16_t y[4] = {0, 0, 0, 0}; + if (zbs) { + for (int i = 0; i < 4; ++i) { + if (left) + x[i] = in_h[i]; + else + y[i] = in_h[i]; + } + x[4] = 0; + } else { + for (int i = 0; i < 4; ++i) { + x[i] = hwsl_(in_h[i], bit_shift); + y[i] = hwslc_(in_h[i], bit_shift); + } + x[4] = hwsl_(extension, bit_shift); + } + + // limb_shift: one-hot of (shift >> 4) & mask. + uint32_t limb_idx = word_instr ? (uint32_t)((shift >> 4) & 1u) + : (uint32_t)((shift >> 4) & 3u); + uint32_t ls[4] = {0, 0, 0, 0}; + ls[limb_idx] = 1u; + + // shifted[i] (DWordHL). intra_left(k)=x[0] if k==0 else x[k]+y[k-1]; + // intra_right(k)=y[k]+x[k+1]. + uint16_t shifted[4]; + for (int i = 0; i < 4; ++i) { + uint16_t v = 0; + if (left) { + for (int j = 0; j <= i; ++j) + if (ls[j]) { + int k = i - j; + v = (uint16_t)(v + (k == 0 ? x[0] : (uint16_t)(x[k] + y[k - 1]))); + } + } + if (right) { + for (int j = 0; j <= 3 - i; ++j) + if (ls[j]) { + int k = i + j; + v = (uint16_t)(v + (uint16_t)(y[k] + x[k + 1])); + } + for (int j = 4 - i; j < 4; ++j) + if (ls[j]) + v = (uint16_t)(v + extension); + } + shifted[i] = v; + } + uint32_t out0 = (uint32_t)shifted[0] | ((uint32_t)shifted[1] << 16); + uint32_t out1 = (uint32_t)shifted[2] | ((uint32_t)shifted[3] << 16); + + out[base + 0] = in_h[0]; // IN_0 + out[base + 1] = in_h[1]; + out[base + 2] = in_h[2]; + out[base + 3] = in_h[3]; + out[base + 4] = shift; // SHIFT_AMOUNT + out[base + 5] = direction; // DIRECTION + out[base + 6] = is_signed; // SIGNED + out[base + 7] = word_instr; // WORD_INSTR + out[base + 8] = out0; // OUT_0 + out[base + 9] = out1; // OUT_1 + out[base + 10] = is_negative; // IS_NEGATIVE + out[base + 11] = bit_shift; // BIT_SHIFT + out[base + 12] = zbs; // ZBS + out[base + 13] = x[0]; // X_0..X_4 + out[base + 14] = x[1]; + out[base + 15] = x[2]; + out[base + 16] = x[3]; + out[base + 17] = x[4]; + out[base + 18] = y[0]; // Y_0..Y_3 + out[base + 19] = y[1]; + out[base + 20] = y[2]; + out[base + 21] = y[3]; + out[base + 22] = ls[0]; // LIMB_SHIFT_RAW_0..2 + out[base + 23] = ls[1]; + out[base + 24] = ls[2]; + out[base + 25] = 1u; // MU + out[base + 26] = (shift_amount >> 8) & 0xFFu; // SHIFT_B1 + out[base + 27] = (shift_amount >> 16) & 0xFFFFu; // SHIFT_H1 + out[base + 28] = (shift_amount >> 32) & 0xFFFFFFFFu; // SHIFT_HIGH +} + +// On-GPU LT trace fill (17 cols). Mirrors the per-row compute of +// `prover/src/tables/lt.rs::generate_lt_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 lt/out/sub/msbs. +// Order-independent (LogUp ALU bus) → validated by multiset/prove, not byte order. +// Padding rows (r >= n) stay all-zero. +// +// Packed stride LT_STRIDE: [0]=lhs [1]=rhs [2]=flags (bit0 signed, bit1 invert) +// [3]=multiplicity. +#define LT_NCOLS 17u +#define LT_STRIDE 4u + +extern "C" __global__ void lt_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 * LT_NCOLS; + if (r >= n) + return; + + const uint64_t *op = ops + r * LT_STRIDE; + uint64_t lhs = op[0]; + uint64_t rhs = op[1]; + uint64_t fl = op[2]; + uint64_t mult = op[3]; + uint64_t is_signed = fl & 1u; + uint64_t invert = (fl >> 1) & 1u; + + uint64_t lt = is_signed ? ((long long)lhs < (long long)rhs ? 1u : 0u) + : (lhs < rhs ? 1u : 0u); + uint64_t out_bit = lt ^ invert; + uint64_t sub = lhs - rhs; // wrapping + + out[base + 0] = lhs & 0xFFFFFFFFu; // LHS_0 (word) + out[base + 1] = (lhs >> 32) & 0xFFFFu; // LHS_1 (half) + out[base + 2] = (lhs >> 48) & 0xFFFFu; // LHS_2 (half) + out[base + 3] = rhs & 0xFFFFFFFFu; // RHS_0 + out[base + 4] = (rhs >> 32) & 0xFFFFu; // RHS_1 + out[base + 5] = (rhs >> 48) & 0xFFFFu; // RHS_2 + out[base + 6] = is_signed; // SIGNED + out[base + 7] = lt; // LT + out[base + 8] = sub & 0xFFFFu; // LHS_SUB_RHS_0 (DWordHL halves) + out[base + 9] = (sub >> 16) & 0xFFFFu; + out[base + 10] = (sub >> 32) & 0xFFFFu; + out[base + 11] = (sub >> 48) & 0xFFFFu; + out[base + 12] = (lhs >> 63) & 1u; // LHS_MSB + out[base + 13] = (rhs >> 63) & 1u; // RHS_MSB + out[base + 14] = invert; // INVERT + out[base + 15] = out_bit; // OUT + out[base + 16] = mult; // MU +} + +// 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 +// `prover/src/tables/memw_register.rs::generate_memw_register_trace_from_rows`: +// 0 ADDRESS=reg_addr/2, 1 TS0=ts&0xffffffff, 2 TS1=ts>>32, 3 VAL0, 4 VAL1, +// 5 OLD0, 6 OLD1, 7 OLD_TS_LO=old_ts&0xffffffff, 8 MU_READ=is_read, 9 MU_WRITE=!is_read. +// All limbs < 2^32 (no Goldilocks reduction). Padding rows are pre-zeroed. +extern "C" __global__ void memw_register_fill( + uint64_t n_acc, const uint32_t *reg_addr, const uint64_t *ts, + const uint64_t *value, const uint8_t *is_read, const long long *row_index, + const uint64_t *old_value, const uint64_t *old_ts, uint32_t ncols, + uint64_t *buf) { + uint64_t i = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x; + if (i >= n_acc) + return; + long long row = row_index[i]; + if (row < 0) + return; + uint64_t base = (uint64_t)row * ncols; + uint64_t v = value[i]; + uint64_t ov = old_value[i]; + uint64_t t = ts[i]; + uint64_t ot = old_ts[i]; + buf[base + 0] = (uint64_t)(reg_addr[i] / 2u); + buf[base + 1] = t & 0xFFFFFFFFull; + buf[base + 2] = t >> 32; + buf[base + 3] = v & 0xFFFFFFFFull; + buf[base + 4] = v >> 32; + buf[base + 5] = ov & 0xFFFFFFFFull; + buf[base + 6] = ov >> 32; + buf[base + 7] = ot & 0xFFFFFFFFull; + buf[base + 8] = is_read[i] ? 1ull : 0ull; + buf[base + 9] = is_read[i] ? 0ull : 1ull; +} diff --git a/crypto/math-cuda/src/device.rs b/crypto/math-cuda/src/device.rs index be0440d5e..24720eb8a 100644 --- a/crypto/math-cuda/src/device.rs +++ b/crypto/math-cuda/src/device.rs @@ -98,6 +98,7 @@ const DEEP_PTX: &str = include_str!(concat!(env!("OUT_DIR"), "/deep.ptx")); const FRI_PTX: &str = include_str!(concat!(env!("OUT_DIR"), "/fri.ptx")); const INVERSE_PTX: &str = include_str!(concat!(env!("OUT_DIR"), "/inverse.ptx")); const LOGUP_PTX: &str = include_str!(concat!(env!("OUT_DIR"), "/logup.ptx")); +const TRACE_CPU_PTX: &str = include_str!(concat!(env!("OUT_DIR"), "/trace_cpu.ptx")); const CONSTRAINT_INTERP_PTX: &str = include_str!(concat!(env!("OUT_DIR"), "/constraint_interp.ptx")); @@ -190,6 +191,14 @@ pub struct Backend { pub logup_finalize_accum_ext3: CudaFunction, pub logup_assemble_aux_ext3: CudaFunction, + // On-GPU trace generation. + pub trace_cpu_fill: CudaFunction, + pub memw_aligned_fill: CudaFunction, + pub load_fill: CudaFunction, + pub store_fill: CudaFunction, + pub shift_fill: CudaFunction, + pub lt_fill: CudaFunction, + pub memw_register_fill: CudaFunction, // constraint_interp.ptx pub constraint_interp_kernel: CudaFunction, pub constraint_composition_kernel: CudaFunction, @@ -297,6 +306,7 @@ impl Backend { let fri = ctx.load_module(Ptx::from_src(FRI_PTX))?; let inverse = ctx.load_module(Ptx::from_src(INVERSE_PTX))?; let logup = ctx.load_module(Ptx::from_src(LOGUP_PTX))?; + let trace_cpu = ctx.load_module(Ptx::from_src(TRACE_CPU_PTX))?; let constraint_interp = ctx.load_module(Ptx::from_src(CONSTRAINT_INTERP_PTX))?; let mut streams = Vec::with_capacity(STREAM_POOL_SIZE); @@ -385,6 +395,13 @@ impl Backend { logup_apply_offsets_add_ext3: logup.load_function("logup_apply_offsets_add_ext3")?, logup_finalize_accum_ext3: logup.load_function("logup_finalize_accum_ext3")?, logup_assemble_aux_ext3: logup.load_function("logup_assemble_aux_ext3")?, + trace_cpu_fill: trace_cpu.load_function("trace_cpu_fill")?, + memw_aligned_fill: trace_cpu.load_function("memw_aligned_fill")?, + load_fill: trace_cpu.load_function("load_fill")?, + store_fill: trace_cpu.load_function("store_fill")?, + shift_fill: trace_cpu.load_function("shift_fill")?, + lt_fill: trace_cpu.load_function("lt_fill")?, + memw_register_fill: trace_cpu.load_function("memw_register_fill")?, constraint_interp_kernel: constraint_interp .load_function("constraint_interp_kernel")?, constraint_composition_kernel: constraint_interp diff --git a/crypto/math-cuda/src/inverse.rs b/crypto/math-cuda/src/inverse.rs index 485e005f8..36762b478 100644 --- a/crypto/math-cuda/src/inverse.rs +++ b/crypto/math-cuda/src/inverse.rs @@ -84,12 +84,15 @@ pub fn batch_inverse_ext3_dev( stream: &Arc, ) -> Result> { assert!(n >= 1, "batch_inverse_ext3_dev requires n >= 1"); - // Runtime guard (not debug_assert): a u32 grid_dim is truncated past - // u32::MAX / BLOCK_SIZE, which would silently launch too few blocks - // and leave a tail uninverted. Reachable on LDE size 2^23+ × multi- - // eval-point R4. Returning Err lets the dispatcher's Err(_) => None - // route the caller to the CPU `inplace_batch_inverse` fallback. - if n > u32::MAX as usize / BLOCK_SIZE as usize { + // Runtime guard (not debug_assert): `n` must fit in u32 so neither `n as u32` + // below nor the kernel's flat index truncates. The block count + // `n.div_ceil(BLOCK_SIZE)` is well within CUDA's grid.x bound (2^31-1) for any + // such `n` — grid.x is NOT the 65535 cap that applies to grid.y/z, so the old + // `u32::MAX / BLOCK_SIZE` bound was 256× too strict and spuriously declined + // legitimate launches (e.g. logup batch-invert of 18 interactions × 2^20 rows + // = ~18.9M), silently routing to a wrong fallback. Returning Err lets the + // dispatcher's Err(_) => None route the caller to the CPU fallback. + if n > u32::MAX as usize { return Err(cudarc::driver::DriverError( cudarc::driver::sys::CUresult::CUDA_ERROR_INVALID_VALUE, )); @@ -182,10 +185,11 @@ pub fn compute_and_invert_denoms_ext3_dev( let total = k_scalars .checked_mul(n) .expect("compute_and_invert_denoms_ext3_dev: k_scalars * n overflow"); - // See `batch_inverse_ext3_dev` for the rationale: runtime Err, not - // debug_assert, so release builds also route past the silent-truncation - // hazard via the caller's CPU fallback. - if total > u32::MAX as usize / BLOCK_SIZE as usize { + // See `batch_inverse_ext3_dev` for the rationale: guard on `u32::MAX` (the + // cast/index bound), NOT `u32::MAX / BLOCK_SIZE` (which is the grid.y/z cap, not + // grid.x, and 256× too strict). Runtime Err, not debug_assert, so release + // builds route past the truncation hazard via the caller's CPU fallback. + if total > u32::MAX as usize { return Err(cudarc::driver::DriverError( cudarc::driver::sys::CUresult::CUDA_ERROR_INVALID_VALUE, )); diff --git a/crypto/math-cuda/src/lde.rs b/crypto/math-cuda/src/lde.rs index 30e524c2c..fd5c529f2 100644 --- a/crypto/math-cuda/src/lde.rs +++ b/crypto/math-cuda/src/lde.rs @@ -583,6 +583,40 @@ pub fn coset_lde_row_major_with_merkle_tree_keep( Ok((handle, lde_out)) } +/// Like [`coset_lde_row_major_with_merkle_tree_keep`] but the input is an +/// already-resident device buffer (`n * m` base-field elements, row-major +/// `[row*m + col]`). No PCIe upload: the buffer is copied device-to-device into +/// the LDE scratch. Used by the on-GPU trace generator so a device-born main +/// trace feeds the LDE without a host round-trip. Retains the trace-domain +/// column-major snapshot exactly like the host keep path, so the resident LogUp +/// aux fingerprint (#762) still works unchanged. +pub fn coset_lde_row_major_with_merkle_tree_keep_dev( + input_dev: &CudaSlice, + n: usize, + m: usize, + blowup_factor: usize, + weights: &[u64], +) -> Result<(GpuLdeBase, Vec)> { + let (tree, col_major_dev, lde_out, trace_col_major) = coset_lde_row_major_inner( + InnerInput::Dev(input_dev), + n, + m, + blowup_factor, + weights, + "coset_lde_row_major_dev lde_size", + true, + )?; + let handle = GpuLdeBase { + buf: Arc::new(col_major_dev), + m, + lde_size: n * blowup_factor, + tree: Some(tree), + trace_dev: trace_col_major.map(Arc::new), + trace_rows: n, + }; + Ok((handle, lde_out)) +} + /// 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/lib.rs b/crypto/math-cuda/src/lib.rs index 493480991..e60baa213 100644 --- a/crypto/math-cuda/src/lib.rs +++ b/crypto/math-cuda/src/lib.rs @@ -15,6 +15,7 @@ pub mod lde; pub mod logup; pub mod merkle; pub mod ntt; +pub mod trace_cpu; // Re-exported for downstream crates so they can refer to CUDA primitive // types without depending on cudarc directly. diff --git a/crypto/math-cuda/src/logup.rs b/crypto/math-cuda/src/logup.rs index e9d120ee1..a022728bb 100644 --- a/crypto/math-cuda/src/logup.rs +++ b/crypto/math-cuda/src/logup.rs @@ -44,12 +44,16 @@ pub struct LogupDescriptor<'a> { } fn cfg(total: usize) -> Result { - // See `batch_inverse_ext3_dev` for the rationale: a u32 grid_dim is - // truncated past u32::MAX / BLOCK_SIZE, which would silently launch too - // few blocks and leave a tail of the (uninitialized) output unwritten. - // Runtime Err, not debug_assert, so release builds also route to the - // caller's CPU fallback. - if total > u32::MAX as usize / BLOCK_SIZE as usize { + // One thread per element: `total` must fit in u32 so neither the `total as u32` + // cast below nor the kernel's `blockIdx.x*blockDim.x + threadIdx.x` index + // truncates. The block count `total.div_ceil(BLOCK_SIZE)` is well within CUDA's + // grid.x bound (2^31-1) for any such `total` — grid.x is NOT capped at the + // 65535 that applies to grid.y/z. (The previous bound `u32::MAX / BLOCK_SIZE` + // conflated the two and spuriously declined large single-chunk tables — e.g. + // SHIFT at 2^20 rows × 18 interactions = ~18.9M > u32::MAX/256 — silently + // falling back to a wrong, zeroed-host aux.) Runtime Err (not debug_assert) so + // release builds route to the caller's CPU fallback rather than truncate. + if total > u32::MAX as usize { return Err(cudarc::driver::DriverError( cudarc::driver::sys::CUresult::CUDA_ERROR_INVALID_VALUE, )); diff --git a/crypto/math-cuda/src/trace_cpu.rs b/crypto/math-cuda/src/trace_cpu.rs new file mode 100644 index 000000000..1b5ca6420 --- /dev/null +++ b/crypto/math-cuda/src/trace_cpu.rs @@ -0,0 +1,501 @@ +//! On-GPU CPU trace-table generation. Uploads packed `CpuOperation` fields and +//! fills the CPU table row-major on device (see `kernels/trace_cpu.cu`), leaving +//! the result resident so it feeds `coset_lde_row_major_with_merkle_tree_keep_dev` +//! with no host round-trip. + +use std::sync::Arc; + +use cudarc::driver::{CudaFunction, CudaSlice, CudaStream, LaunchConfig, PushKernelArg}; + +use crate::Result; +use crate::device::{Backend, backend}; + +/// CPU table width (`prover::tables::cpu::cols::NUM_COLUMNS`). +pub const CPU_NCOLS: usize = 38; +/// Packed input stride, u64 per op (must match `trace_cpu.cu`). +pub const CPU_OP_STRIDE: usize = 11; + +/// Build one CPU trace-table chunk on device from packed ops. +/// +/// `packed_ops` is `n * CPU_OP_STRIDE` u64s (see `trace_cpu.cu` for the field +/// layout). `num_rows` is the padded power-of-two row count and `last_ts` the +/// timestamp of the last real op (0 when `n == 0`). Returns the row-major device +/// buffer `[row*CPU_NCOLS + col]` (`num_rows * CPU_NCOLS` u64), ready for the +/// device-input LDE. +pub fn gpu_build_cpu_trace( + packed_ops: &[u64], + n: usize, + num_rows: usize, + last_ts: u64, +) -> Result> { + debug_assert_eq!(packed_ops.len(), n * CPU_OP_STRIDE); + let be = backend()?; + let stream = be.next_stream(); + + // `clone_htod` rejects empty slices; a 1-element dummy is never read because + // every row is a padding row when `n == 0`. + let ops_dev = if packed_ops.is_empty() { + stream.alloc_zeros::(CPU_OP_STRIDE)? + } else { + stream.clone_htod(packed_ops)? + }; + + // Zero-initialised: the kernel only writes non-zero cells. + let mut out = stream.alloc_zeros::(num_rows * CPU_NCOLS)?; + + unsafe { + stream + .launch_builder(&be.trace_cpu_fill) + .arg(&ops_dev) + .arg(&(n as u64)) + .arg(&(num_rows as u64)) + .arg(&last_ts) + .arg(&mut out) + .launch(LaunchConfig::for_num_elems(num_rows as u32))?; + } + stream.synchronize()?; + Ok(out) +} + +/// Host-returning wrapper over [`gpu_build_cpu_trace`] for byte-parity tests: +/// builds the row-major CPU table on device and copies it back on the same stream. +pub fn gpu_build_cpu_trace_host( + packed_ops: &[u64], + n: usize, + num_rows: usize, + last_ts: u64, +) -> Result> { + debug_assert_eq!(packed_ops.len(), n * CPU_OP_STRIDE); + let be = backend()?; + let stream = be.next_stream(); + + let ops_dev = if packed_ops.is_empty() { + stream.alloc_zeros::(CPU_OP_STRIDE)? + } else { + stream.clone_htod(packed_ops)? + }; + let mut out = stream.alloc_zeros::(num_rows * CPU_NCOLS)?; + + unsafe { + stream + .launch_builder(&be.trace_cpu_fill) + .arg(&ops_dev) + .arg(&(n as u64)) + .arg(&(num_rows as u64)) + .arg(&last_ts) + .arg(&mut out) + .launch(LaunchConfig::for_num_elems(num_rows as u32))?; + } + let host = stream.clone_dtoh(&out)?; + stream.synchronize()?; + Ok(host) +} + +/// MEMW_A table width (`prover::tables::memw_aligned::cols::NUM_COLUMNS`). +pub const MEMW_ALIGNED_NCOLS: usize = 29; +/// Packed MEMW_A input stride, u64 per op (must match `memw_aligned_fill` in `trace_cpu.cu`). +pub const MEMW_ALIGNED_STRIDE: usize = 12; + +/// Build one MEMW_A (aligned memory) trace-table chunk on device from packed ops. +/// +/// `packed_ops` is `n * MEMW_ALIGNED_STRIDE` u64s (see `trace_cpu.cu` for the field +/// layout). `num_rows` is the padded power-of-two row count. Returns the row-major +/// device buffer `[row*MEMW_ALIGNED_NCOLS + col]` (`num_rows * MEMW_ALIGNED_NCOLS` +/// u64), left resident so it feeds the device-input LDE with no full-column upload. +pub fn gpu_build_memw_aligned_trace( + packed_ops: &[u64], + n: usize, + num_rows: usize, +) -> Result> { + debug_assert_eq!(packed_ops.len(), n * MEMW_ALIGNED_STRIDE); + let be = backend()?; + let stream = be.next_stream(); + + let mut out = stream.alloc_zeros::(num_rows * MEMW_ALIGNED_NCOLS)?; + + // `clone_htod` rejects empty slices; a 1-op dummy is never read (all padding). + let ops_dev = if packed_ops.is_empty() { + stream.alloc_zeros::(MEMW_ALIGNED_STRIDE)? + } else { + stream.clone_htod(packed_ops)? + }; + + unsafe { + stream + .launch_builder(&be.memw_aligned_fill) + .arg(&ops_dev) + .arg(&(n as u64)) + .arg(&(num_rows as u64)) + .arg(&mut out) + .launch(LaunchConfig::for_num_elems(num_rows as u32))?; + } + stream.synchronize()?; + Ok(out) +} + +// ----------------------------------------------------------------------------- +// LOAD / STORE (per-row-map memory tables — same shape as MEMW_A). A shared +// interleaved-fill launcher keeps the per-table code to a packing convention + +// column/stride constants. +// ----------------------------------------------------------------------------- + +/// LOAD table width / packed input stride (must match `load_fill` in `trace_cpu.cu`). +pub const LOAD_NCOLS: usize = 18; +pub const LOAD_STRIDE: usize = 7; +/// STORE table width / packed input stride (must match `store_fill`). +pub const STORE_NCOLS: usize = 16; +pub const STORE_STRIDE: usize = 4; +/// SHIFT table width / packed input stride (must match `shift_fill`). +pub const SHIFT_NCOLS: usize = 29; +pub const SHIFT_STRIDE: usize = 3; +/// LT table width / packed input stride (must match `lt_fill`). Input is already +/// host-deduplicated: one unique op + summed multiplicity per row. +pub const LT_NCOLS: usize = 17; +pub const LT_STRIDE: usize = 4; + +/// 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 +/// (unsynchronized) device buffer. +#[allow(clippy::too_many_arguments)] +fn build_interleaved_on( + stream: &Arc, + kernel: &CudaFunction, + packed_ops: &[u64], + n: usize, + num_rows: usize, + ncols: usize, + stride: usize, +) -> Result> { + debug_assert_eq!(packed_ops.len(), n * stride); + let mut out = stream.alloc_zeros::(num_rows * ncols)?; + let ops_dev = if packed_ops.is_empty() { + stream.alloc_zeros::(stride)? + } else { + stream.clone_htod(packed_ops)? + }; + unsafe { + stream + .launch_builder(kernel) + .arg(&ops_dev) + .arg(&(n as u64)) + .arg(&(num_rows as u64)) + .arg(&mut out) + .launch(LaunchConfig::for_num_elems(num_rows as u32))?; + } + Ok(out) +} + +fn load_kernel(be: &Backend) -> &CudaFunction { + &be.load_fill +} +fn store_kernel(be: &Backend) -> &CudaFunction { + &be.store_fill +} +fn shift_kernel(be: &Backend) -> &CudaFunction { + &be.shift_fill +} +fn lt_kernel(be: &Backend) -> &CudaFunction { + &be.lt_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`). +pub fn gpu_build_lt_trace(packed_ops: &[u64], n: usize, num_rows: usize) -> Result> { + let be = backend()?; + let stream = be.next_stream(); + let out = build_interleaved_on( + &stream, + lt_kernel(be), + packed_ops, + n, + num_rows, + LT_NCOLS, + LT_STRIDE, + )?; + stream.synchronize()?; + Ok(out) +} + +/// Host-returning LT build for multiset-equality tests. +pub fn gpu_build_lt_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, + lt_kernel(be), + packed_ops, + n, + num_rows, + LT_NCOLS, + LT_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( + packed_ops: &[u64], + n: usize, + num_rows: usize, +) -> Result> { + let be = backend()?; + let stream = be.next_stream(); + let out = build_interleaved_on( + &stream, + shift_kernel(be), + packed_ops, + n, + num_rows, + SHIFT_NCOLS, + SHIFT_STRIDE, + )?; + stream.synchronize()?; + Ok(out) +} + +/// Host-returning SHIFT build for byte-parity tests. +pub fn gpu_build_shift_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, + shift_kernel(be), + packed_ops, + n, + num_rows, + SHIFT_NCOLS, + SHIFT_STRIDE, + )?; + let host = stream.clone_dtoh(&out)?; + stream.synchronize()?; + Ok(host) +} + +/// Build one LOAD trace-table chunk on device (residency-ready row-major buffer). +pub fn gpu_build_load_trace( + packed_ops: &[u64], + n: usize, + num_rows: usize, +) -> Result> { + let be = backend()?; + let stream = be.next_stream(); + let out = build_interleaved_on( + &stream, + load_kernel(be), + packed_ops, + n, + num_rows, + LOAD_NCOLS, + LOAD_STRIDE, + )?; + stream.synchronize()?; + Ok(out) +} + +/// Host-returning LOAD build for byte-parity tests. +pub fn gpu_build_load_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, + load_kernel(be), + packed_ops, + n, + num_rows, + LOAD_NCOLS, + LOAD_STRIDE, + )?; + let host = stream.clone_dtoh(&out)?; + stream.synchronize()?; + Ok(host) +} + +/// Build one STORE trace-table chunk on device (residency-ready row-major buffer). +pub fn gpu_build_store_trace( + packed_ops: &[u64], + n: usize, + num_rows: usize, +) -> Result> { + let be = backend()?; + let stream = be.next_stream(); + let out = build_interleaved_on( + &stream, + store_kernel(be), + packed_ops, + n, + num_rows, + STORE_NCOLS, + STORE_STRIDE, + )?; + stream.synchronize()?; + Ok(out) +} + +/// Host-returning STORE build for byte-parity tests. +pub fn gpu_build_store_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, + store_kernel(be), + packed_ops, + n, + num_rows, + STORE_NCOLS, + STORE_STRIDE, + )?; + let host = stream.clone_dtoh(&out)?; + stream.synchronize()?; + Ok(host) +} + +/// Host-returning wrapper over [`gpu_build_memw_aligned_trace`] for byte-parity +/// tests: builds the row-major MEMW_A buffer and copies it back on the same stream. +pub fn gpu_build_memw_aligned_trace_host( + packed_ops: &[u64], + n: usize, + num_rows: usize, +) -> Result> { + debug_assert_eq!(packed_ops.len(), n * MEMW_ALIGNED_STRIDE); + let be = backend()?; + let stream = be.next_stream(); + + let mut out = stream.alloc_zeros::(num_rows * MEMW_ALIGNED_NCOLS)?; + let ops_dev = if packed_ops.is_empty() { + stream.alloc_zeros::(MEMW_ALIGNED_STRIDE)? + } else { + stream.clone_htod(packed_ops)? + }; + unsafe { + stream + .launch_builder(&be.memw_aligned_fill) + .arg(&ops_dev) + .arg(&(n as u64)) + .arg(&(num_rows as u64)) + .arg(&mut out) + .launch(LaunchConfig::for_num_elems(num_rows as u32))?; + } + let host = stream.clone_dtoh(&out)?; + stream.synchronize()?; + Ok(host) +} + +// ----------------------------------------------------------------------------- +// MEMW_R (register fast-path) fill. Registers are pre-walked on the host (the +// sequential memory model recovers old_value/old_ts); this fills the 10 MEMW_R +// columns row-major on device from the walked rows, leaving the matrix resident +// for the device-input LDE (removes MEMW_R's full-column H2D). See the +// `memw_register_fill` kernel in `trace_cpu.cu`. +// ----------------------------------------------------------------------------- + +/// MEMW_R table width (`prover::tables::memw_register::cols::NUM_COLUMNS`). +pub const MEMW_REGISTER_NCOLS: usize = 10; + +/// Build the MEMW_R trace table on device from already-walked rows (the +/// `old_value`/`old_ts` recovered by the host walk are uploaded here). Returns the +/// residency-ready row-major `[row*NCOLS+col]` buffer. `row_index` is the identity +/// (every input row is a real MEMW_R row). +#[allow(clippy::too_many_arguments)] +pub fn gpu_fill_memw_register( + reg_addr: &[u32], + ts: &[u64], + value: &[u64], + is_read: &[u8], + old_value: &[u64], + old_ts: &[u64], + num_rows: usize, +) -> Result> { + let be = backend()?; + let stream = be.next_stream(); + let buf = fill_memw_register_on( + &stream, reg_addr, ts, value, is_read, old_value, old_ts, num_rows, + )?; + stream.synchronize()?; + Ok(buf) +} + +/// Host-returning wrapper over [`gpu_fill_memw_register`] for byte-parity tests. +#[allow(clippy::too_many_arguments)] +pub fn gpu_fill_memw_register_host( + reg_addr: &[u32], + ts: &[u64], + value: &[u64], + is_read: &[u8], + old_value: &[u64], + old_ts: &[u64], + num_rows: usize, +) -> Result> { + let be = backend()?; + let stream = be.next_stream(); + let buf = fill_memw_register_on( + &stream, reg_addr, ts, value, is_read, old_value, old_ts, num_rows, + )?; + let host = stream.clone_dtoh(&buf)?; + stream.synchronize()?; + Ok(host) +} + +#[allow(clippy::too_many_arguments)] +fn fill_memw_register_on( + stream: &std::sync::Arc, + reg_addr: &[u32], + ts: &[u64], + value: &[u64], + is_read: &[u8], + old_value: &[u64], + old_ts: &[u64], + num_rows: usize, +) -> Result> { + let n = reg_addr.len(); + debug_assert_eq!(ts.len(), n); + debug_assert_eq!(value.len(), n); + debug_assert_eq!(is_read.len(), n); + debug_assert_eq!(old_value.len(), n); + debug_assert_eq!(old_ts.len(), n); + let be = backend()?; + let mut buf = stream.alloc_zeros::(num_rows * MEMW_REGISTER_NCOLS)?; + if n == 0 { + return Ok(buf); + } + let keys_d = stream.clone_htod(reg_addr)?; + let ts_d = stream.clone_htod(ts)?; + let value_d = stream.clone_htod(value)?; + let is_read_d = stream.clone_htod(is_read)?; + let old_value_d = stream.clone_htod(old_value)?; + let old_ts_d = stream.clone_htod(old_ts)?; + let row_index: Vec = (0..n as i64).collect(); + let row_index_d = stream.clone_htod(&row_index)?; + let n_u64 = n as u64; + let ncols_u32 = MEMW_REGISTER_NCOLS as u32; + unsafe { + stream + .launch_builder(&be.memw_register_fill) + .arg(&n_u64) + .arg(&keys_d) + .arg(&ts_d) + .arg(&value_d) + .arg(&is_read_d) + .arg(&row_index_d) + .arg(&old_value_d) + .arg(&old_ts_d) + .arg(&ncols_u32) + .arg(&mut buf) + .launch(LaunchConfig::for_num_elems(n as u32))?; + } + Ok(buf) +} diff --git a/crypto/math-cuda/tests/lde_dev_parity.rs b/crypto/math-cuda/tests/lde_dev_parity.rs new file mode 100644 index 000000000..064653760 --- /dev/null +++ b/crypto/math-cuda/tests/lde_dev_parity.rs @@ -0,0 +1,86 @@ +//! Parity for the on-GPU trace generator's LDE seam: the device-input keep path +//! [`coset_lde_row_major_with_merkle_tree_keep_dev`] must produce a +//! byte-identical Merkle root and row-major LDE as the host-input keep path +//! [`coset_lde_row_major_with_merkle_tree_keep`] for the same matrix. This is +//! the "Step 1a" isolation test: it validates the seam without any trace-fill +//! kernel — upload a host matrix, run both paths, compare. +//! +//! Requires a GPU (skips cleanly if the CUDA backend is unavailable). + +use math::field::element::FieldElement; +use math::field::goldilocks::GoldilocksField; +use math::field::traits::{IsField, IsPrimeField}; +use math_cuda::device::backend; +use rand::{Rng, SeedableRng}; +use rand_chacha::ChaCha8Rng; + +/// Coset weights `[1/N, g/N, ...]` — the layout `crypto/stark/src/prover.rs` uses. +fn coset_weights(n: usize, coset_offset: u64) -> Vec { + let inv_n = FieldElement::::from(n as u64) + .inv() + .expect("n non-zero"); + let mut w = Vec::with_capacity(n); + let mut cur = *inv_n.value(); + for _ in 0..n { + w.push(cur); + cur = GoldilocksField::mul(&cur, &coset_offset); + } + w +} + +fn assert_dev_matches_host(log_n: u64, m: usize, blowup: usize, seed: u64) { + // Skip cleanly when no GPU/CUDA backend is present. + if backend().is_err() { + eprintln!("skipping lde_dev_parity: no CUDA backend"); + return; + } + + let n = 1usize << log_n; + let mut rng = ChaCha8Rng::seed_from_u64(seed); + // Row-major [row*m + col], the layout the trace generator emits. + let row_major: Vec = (0..n * m).map(|_| rng.r#gen::()).collect(); + + let coset_offset = 7u64; + let weights = coset_weights(n, coset_offset); + + // Host path: uploads `row_major` internally. + let (h_handle, h_lde) = math_cuda::lde::coset_lde_row_major_with_merkle_tree_keep( + &row_major, n, m, blowup, &weights, + ) + .expect("host keep"); + + // Device path: pre-upload the SAME matrix, then run the device-input LDE. + let be = backend().unwrap(); + let stream = be.next_stream(); + let dev = stream.clone_htod(&row_major).expect("upload matrix"); + stream.synchronize().expect("sync upload"); + let (d_handle, d_lde) = + math_cuda::lde::coset_lde_row_major_with_merkle_tree_keep_dev(&dev, n, m, blowup, &weights) + .expect("dev keep"); + + assert_eq!( + h_handle.tree.as_ref().unwrap().root, + d_handle.tree.as_ref().unwrap().root, + "Merkle root mismatch (n={n}, m={m}, blowup={blowup})" + ); + let hc: Vec = h_lde.iter().map(GoldilocksField::canonical).collect(); + let dc: Vec = d_lde.iter().map(GoldilocksField::canonical).collect(); + assert_eq!( + hc, dc, + "row-major LDE mismatch (n={n}, m={m}, blowup={blowup})" + ); +} + +#[test] +fn dev_matches_host_cpu_table_width() { + // CPU table is 38 columns — the P1 target. + assert_dev_matches_host(10, 38, 2, 0xC0DE); + assert_dev_matches_host(12, 38, 4, 0xBEEF); +} + +#[test] +fn dev_matches_host_various_widths() { + assert_dev_matches_host(12, 10, 2, 0x11); // memw_register width + assert_dev_matches_host(14, 16, 2, 0x22); // store width + assert_dev_matches_host(13, 49, 2, 0x33); // memw width +} diff --git a/crypto/stark/src/gpu_lde.rs b/crypto/stark/src/gpu_lde.rs index b23adb1ac..b89132d7f 100644 --- a/crypto/stark/src/gpu_lde.rs +++ b/crypto/stark/src/gpu_lde.rs @@ -554,6 +554,74 @@ where Some((tree, handle, lde_out)) } +/// Device-input analog of [`try_expand_leaf_and_tree_row_major_keep`]: the main +/// trace is already resident on device (row-major `[row*m + col]`, produced by +/// the on-GPU trace generator), so there is NO host upload — the LDE copies it +/// device-to-device. Returns the same (root-only host tree, resident LDE handle +/// with the trace-domain snapshot for the aux build, row-major host LDE) triple. +pub(crate) fn try_expand_leaf_and_tree_row_major_keep_dev( + input_dev: &math_cuda::CudaSlice, + n: usize, + m: usize, + blowup_factor: usize, + weights: &[FieldElement], +) -> Option<( + MerkleTree, + math_cuda::lde::GpuLdeBase, + Vec>, +)> +where + F: IsField + 'static, + E: IsField + 'static, + B: IsMerkleTreeBackend, +{ + // NOTE: no `lde_size < gpu_lde_threshold()` decline here. The threshold is a + // host-side heuristic (small tables are faster to LDE on the CPU, avoiding the + // H2D). For a DEVICE-resident input there is no valid CPU fallback — the host + // trace is a zeroed placeholder — so declining would silently commit zeros + // (bus imbalance / wrong proof). The input is already on-device, so we always + // run the GPU LDE regardless of size. + let _lde_size = n.saturating_mul(blowup_factor); + if TypeId::of::() != TypeId::of::() { + return None; + } + if TypeId::of::() != TypeId::of::() { + return None; + } + if input_dev.len() != n * m || m == 0 || n == 0 { + return None; + } + + let weights_u64 = unsafe { weights_to_u64::(weights) }; + + GPU_LDE_CALLS.fetch_add(m as u64, Ordering::Relaxed); + GPU_LEAF_HASH_CALLS.fetch_add(1, Ordering::Relaxed); + GPU_MERKLE_TREE_CALLS.fetch_add(1, Ordering::Relaxed); + + let (handle, lde_u64) = math_cuda::lde::coset_lde_row_major_with_merkle_tree_keep_dev( + input_dev, + 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(), + ) + }; + + let root = handle.tree.as_ref()?.root; + let tree = MerkleTree::::from_root(root); + Some((tree, handle, 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/instruments.rs b/crypto/stark/src/instruments.rs index 96bf6ffae..d0dee46e9 100644 --- a/crypto/stark/src/instruments.rs +++ b/crypto/stark/src/instruments.rs @@ -238,6 +238,14 @@ static AUX_FINGERPRINT_US: AtomicU64 = AtomicU64::new(0); static AUX_INVERT_US: AtomicU64 = AtomicU64::new(0); static AUX_TERM_US: AtomicU64 = AtomicU64::new(0); static AUX_ACCUM_US: AtomicU64 = AtomicU64::new(0); +// GPU trace-generation transfer accounting (the bytes we aim to eliminate). +// `MAIN_H2D_BYTES` sums every host->device upload of a main-trace matrix in +// `commit_main_trace`; `MAIN_DEV_BUILDS` counts tables that instead fed the LDE +// from an already-device-resident buffer (no upload) — 0 until the on-GPU +// trace generator lands, then the target is "all main tables". +static MAIN_H2D_BYTES: AtomicU64 = AtomicU64::new(0); +static MAIN_H2D_CALLS: AtomicU64 = AtomicU64::new(0); +static MAIN_DEV_BUILDS: AtomicU64 = AtomicU64::new(0); thread_local! { static TIMING_DATA: RefCell> = const { RefCell::new(None) }; @@ -281,6 +289,27 @@ pub fn accum_aux_accumulate(d: Duration) { AUX_ACCUM_US.fetch_add(d.as_micros() as u64, Ordering::Relaxed); } +/// Record a host->device upload of a main-trace matrix (`bytes`) at the LDE +/// commit boundary — the transfer the on-GPU trace generator removes. +pub fn accum_main_h2d(bytes: usize) { + MAIN_H2D_BYTES.fetch_add(bytes as u64, Ordering::Relaxed); + MAIN_H2D_CALLS.fetch_add(1, Ordering::Relaxed); +} + +/// Record a main table whose LDE input was already device-resident (no upload). +pub fn accum_main_dev_build() { + MAIN_DEV_BUILDS.fetch_add(1, Ordering::Relaxed); +} + +/// `(uploaded_bytes, upload_calls, device_resident_builds)` since the last reset. +pub fn take_main_transfer() -> (u64, u64, u64) { + ( + MAIN_H2D_BYTES.swap(0, Ordering::Relaxed), + MAIN_H2D_CALLS.swap(0, Ordering::Relaxed), + MAIN_DEV_BUILDS.swap(0, Ordering::Relaxed), + ) +} + pub fn take_r1_sub() -> Round1SubOps { Round1SubOps { main_lde: Duration::from_micros(R1_MAIN_LDE_US.swap(0, Ordering::Relaxed)), @@ -311,6 +340,9 @@ pub fn reset_all() { AUX_INVERT_US.store(0, Ordering::Relaxed); AUX_TERM_US.store(0, Ordering::Relaxed); AUX_ACCUM_US.store(0, Ordering::Relaxed); + MAIN_H2D_BYTES.store(0, Ordering::Relaxed); + MAIN_H2D_CALLS.store(0, Ordering::Relaxed); + MAIN_DEV_BUILDS.store(0, Ordering::Relaxed); TIMING_DATA.with(|cell| { cell.borrow_mut().take(); }); diff --git a/crypto/stark/src/logup_gpu.rs b/crypto/stark/src/logup_gpu.rs index bc9e88302..19da57766 100644 --- a/crypto/stark/src/logup_gpu.rs +++ b/crypto/stark/src/logup_gpu.rs @@ -431,7 +431,15 @@ where { return None; } - if trace_len < GPU_LOGUP_MIN_ROWS || main_cols.is_empty() || interactions.is_empty() { + // The `trace_len < GPU_LOGUP_MIN_ROWS` decline is a host-side heuristic (small + // tables are cheaper on the CPU aux path). It must NOT apply when the main is + // device-resident (`main_dev.is_some()`): that table's host trace is a zeroed + // placeholder, so a CPU fallback would build the aux from zeros → LogUp bus + // imbalance. Device-resident tables always build the aux on-device here. + if (trace_len < GPU_LOGUP_MIN_ROWS && main_dev.is_none()) + || main_cols.is_empty() + || interactions.is_empty() + { return None; } if std::env::var_os("LAMBDA_VM_NO_GPU_LOGUP").is_some() { diff --git a/crypto/stark/src/prover.rs b/crypto/stark/src/prover.rs index 5c03292da..c501e080e 100644 --- a/crypto/stark/src/prover.rs +++ b/crypto/stark/src/prover.rs @@ -777,6 +777,46 @@ pub trait IsStarkProver< // Fused GPU path (cuda only): row-major NTT — single H2D from the // already-row-major trace, no column extraction, no transpose. // Falls back to CPU if GPU path returns None. + // Device-born main trace (on-GPU trace generator): feed the LDE from the + // resident buffer directly — no host upload. Falls through to the host + // path if the handle is absent or the GPU path declines. + #[cfg(feature = "cuda")] + if precomputed.is_none() + && let Some(dev) = trace.main_input_dev() + { + let num_cols = trace.num_main_columns; + let n = trace.num_rows(); + #[cfg(feature = "instruments")] + let t_sub = Instant::now(); + if let Some((tree, handle, main_data)) = + crate::gpu_lde::try_expand_leaf_and_tree_row_major_keep_dev::< + Field, + Field, + BatchedMerkleTreeBackend, + >( + dev, + n, + num_cols, + domain.blowup_factor, + &twiddles.coset_weights, + ) + { + #[cfg(feature = "instruments")] + crate::instruments::accum_r1_main(t_sub.elapsed(), std::time::Duration::ZERO); + // No host->device upload happened for this table. + #[cfg(feature = "instruments")] + crate::instruments::accum_main_dev_build(); + let root = tree.root; + return Ok(( + TableCommit::plain(tree, root), + (main_data, num_cols), + Some(handle), + )); + } + } + + // Fused GPU host path (cuda only): row-major NTT — single H2D from the + // already-row-major trace, no column extraction, no transpose. #[cfg(feature = "cuda")] if precomputed.is_none() { let (trace_slice, num_cols) = trace.main_data_row_major(); @@ -805,6 +845,10 @@ pub trait IsStarkProver< let root = tree.root; #[cfg(feature = "instruments")] crate::instruments::accum_r1_main(main_lde_dur, std::time::Duration::ZERO); + // This table's main matrix was uploaded host->device for the LDE. + // Counts the transfer the on-GPU trace generator will remove. + #[cfg(feature = "instruments")] + crate::instruments::accum_main_h2d(trace_slice.len() * std::mem::size_of::()); return Ok(( TableCommit::plain(tree, root), (main_data, num_cols), @@ -2191,6 +2235,16 @@ pub trait IsStarkProver< heap_snaps.push(s); } + // Free the device-born main input buffers now that every main trace is + // committed. They are only the D2D source for `commit_main_trace`; the + // resident LDE handle (`main_gpu_handles`) carries what R2-4 needs. Holding + // them through the aux/DEEP/FRI VRAM peak makes large blocks OOM where the + // CPU-trace stream-and-free path fits. + #[cfg(feature = "cuda")] + for (_, trace, _) in air_trace_pairs.iter_mut() { + trace.clear_main_input_dev(); + } + // ===================================================================== // Round 1, Phase B: Sample shared LogUp challenges // ===================================================================== diff --git a/crypto/stark/src/trace.rs b/crypto/stark/src/trace.rs index 6d40425b7..b4647bad4 100644 --- a/crypto/stark/src/trace.rs +++ b/crypto/stark/src/trace.rs @@ -45,8 +45,41 @@ where /// LDE did not run for this table. #[cfg(feature = "cuda")] pub(crate) main_trace_dev: Option, + /// Device-born main trace produced by the on-GPU trace generator, row-major + /// `[row*num_main_columns + col]`, `num_rows * num_main_columns` u64s. When + /// present, `commit_main_trace` feeds the LDE from this buffer directly + /// (device-to-device, no host upload). None on the CPU trace-gen path. + #[cfg(feature = "cuda")] + pub(crate) main_input_dev: Option, +} + +/// Device-born main trace (row-major `[row*num_main_columns + col]`) produced by +/// the on-GPU trace generator, feeding `commit_main_trace` directly. GPU-only and +/// transient; excluded from logical trace equality and opaque in `Debug`, +/// matching `ResidentMainTrace`. +#[cfg(feature = "cuda")] +#[derive(Clone)] +pub(crate) struct DeviceMainInput { + pub(crate) buf: std::sync::Arc>, +} + +#[cfg(feature = "cuda")] +impl core::fmt::Debug for DeviceMainInput { + fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { + f.debug_struct("DeviceMainInput").finish_non_exhaustive() + } +} + +#[cfg(feature = "cuda")] +impl PartialEq for DeviceMainInput { + fn eq(&self, _other: &Self) -> bool { + true + } } +#[cfg(feature = "cuda")] +impl Eq for DeviceMainInput {} + /// Device-resident trace-domain main columns (column-major `[col*rows + row]`), /// retained from the R1 main LDE for the aux fingerprint kernel. GPU-only and /// transient; the device buffer is excluded from logical trace equality (only @@ -105,6 +138,8 @@ where resident_aux_ok: true, #[cfg(feature = "cuda")] main_trace_dev: None, + #[cfg(feature = "cuda")] + main_input_dev: None, } } @@ -133,6 +168,8 @@ where resident_aux_ok: true, #[cfg(feature = "cuda")] main_trace_dev: None, + #[cfg(feature = "cuda")] + main_input_dev: None, } } @@ -154,6 +191,8 @@ where resident_aux_ok: true, #[cfg(feature = "cuda")] main_trace_dev: None, + #[cfg(feature = "cuda")] + main_input_dev: None, } } @@ -213,6 +252,32 @@ where self.main_trace_dev = None; } + /// Attach a device-born main trace (row-major `[row*num_main_columns + col]`, + /// `num_rows * num_main_columns` u64s) from the on-GPU trace generator, so + /// `commit_main_trace` feeds the LDE from it directly (no host upload). + #[cfg(feature = "cuda")] + pub fn set_main_input_dev(&mut self, buf: std::sync::Arc>) { + self.main_input_dev = Some(DeviceMainInput { buf }); + } + + /// The device-born main trace buffer, if the on-GPU trace generator produced + /// one for this table. + #[cfg(feature = "cuda")] + pub fn main_input_dev(&self) -> Option<&math_cuda::CudaSlice> { + self.main_input_dev.as_ref().map(|b| b.buf.as_ref()) + } + + /// Drop the device-born main input buffer. It is only read by `commit_main_trace` + /// (the D2D source for the LDE); once the main commit is done nothing else uses + /// it, so freeing it after Round-1 Phase A reclaims the main-trace-sized device + /// buffers (≈ the whole main trace) before the aux LDE / DEEP / FRI VRAM peak — + /// without this the on-GPU trace generator holds them to end-of-proof and large + /// blocks (e.g. ethrex 10-tx) OOM where the CPU-trace (stream-and-free) path fits. + #[cfg(feature = "cuda")] + pub fn clear_main_input_dev(&mut self) { + self.main_input_dev = None; + } + pub fn num_steps(&self) -> usize { debug_assert!(self.main_table.height.is_multiple_of(self.step_size)); self.main_table.height / self.step_size diff --git a/prover/Cargo.toml b/prover/Cargo.toml index 57e8114d0..36c4b6423 100644 --- a/prover/Cargo.toml +++ b/prover/Cargo.toml @@ -7,7 +7,7 @@ license.workspace = true [features] default = ["parallel"] parallel = ["stark/parallel", "math/parallel", "crypto/parallel", "dep:rayon"] -cuda = ["stark/cuda"] +cuda = ["stark/cuda", "dep:math-cuda"] test-cuda-faults = ["cuda", "stark/test-cuda-faults"] debug-checks = ["stark/debug-checks"] instruments = ["stark/instruments"] @@ -20,6 +20,7 @@ crypto = { path = "../crypto/crypto" } math = { path = "../crypto/math" } executor = { path = "../executor" } ecsm = { path = "../crypto/ecsm" } +math-cuda = { path = "../crypto/math-cuda", optional = true } serde = { version = "1.0", features = ["derive"] } rayon = { version = "1.8.0", optional = true } sysinfo = { version = "0.31", default-features = false, features = ["system"] } diff --git a/prover/src/lib.rs b/prover/src/lib.rs index bd72a3670..588813f62 100644 --- a/prover/src/lib.rs +++ b/prover/src/lib.rs @@ -967,6 +967,16 @@ pub fn prove_with_options_and_inputs( drop(__root); let spans = stark::instruments::take_timeline(); print!("{}", stark::instruments::format_timeline(&spans)); + // Main-trace CPU->GPU transfer accounting: the bytes the on-GPU trace + // generator aims to eliminate, and how many tables already skipped the + // upload by feeding the LDE a device-resident buffer. + let (h2d_bytes, h2d_calls, dev_builds) = stark::instruments::take_main_transfer(); + if h2d_calls > 0 || dev_builds > 0 { + println!( + "[transfer] main-trace H2D: {:.1} MiB over {h2d_calls} table(s); device-resident builds: {dev_builds}", + h2d_bytes as f64 / (1024.0 * 1024.0), + ); + } if let Ok(path) = std::env::var("LAMBDA_VM_TIMELINE_JSON") { let _ = std::fs::write(&path, stark::instruments::timeline_json(&spans)); println!("[timeline] wrote {path}"); diff --git a/prover/src/tables/gpu_trace.rs b/prover/src/tables/gpu_trace.rs new file mode 100644 index 000000000..c93c2b523 --- /dev/null +++ b/prover/src/tables/gpu_trace.rs @@ -0,0 +1,430 @@ +//! On-GPU trace generation: build trace tables directly in device memory so +//! they feed the already-GPU LDE/commit without a host round-trip. +//! +//! This module is compiled only under the `cuda` feature. It hosts the +//! device-build dispatch (added table-by-table) plus the kill-switch used to +//! A/B the GPU path against the CPU trace generator. +//! +//! Design: `reports/tracegen/GPU-TRACEGEN-DESIGN-V2.md`. +#![cfg(feature = "cuda")] + +use std::sync::{Arc, OnceLock}; + +use stark::trace::TraceTable; + +use std::collections::HashMap; + +use super::cpu::{self, CpuOperation}; +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::shift::{self, ShiftOperation}; +use super::store::{self, StoreOperation}; +use super::types::{GoldilocksExtension, GoldilocksField}; + +/// When set (`LAMBDA_VM_CPU_TRACE=1`), all GPU trace-build dispatchers return +/// `None` so callers fall back to the CPU trace generator. This is the one-flag +/// A/B switch: same binary, `LAMBDA_VM_CPU_TRACE=1` runs the CPU baseline, +/// unset runs the GPU path. Read once and cached. +pub(crate) fn gpu_trace_disabled() -> bool { + static DISABLED: OnceLock = OnceLock::new(); + *DISABLED.get_or_init(|| { + std::env::var("LAMBDA_VM_CPU_TRACE") + .map(|v| v != "0" && !v.is_empty()) + .unwrap_or(false) + }) +} + +/// 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) +// ============================================================================= + +/// Marshal one chunk of `CpuOperation`s into the packed layout the `trace_cpu` +/// kernel consumes (stride `CPU_OP_STRIDE` u64/op). The kernel does the same +/// bit-slicing as `cpu::generate_cpu_trace`, so this only copies fields — no +/// per-column encoding on the host. +pub(crate) fn pack_cpu_ops(chunk: &[CpuOperation]) -> Vec { + let stride = math_cuda::trace_cpu::CPU_OP_STRIDE; + let mut packed = vec![0u64; chunk.len() * stride]; + for (i, op) in chunk.iter().enumerate() { + let f = &op.decode.fields; + let flags = (f.word_instr as u64) + | ((f.read_register1 as u64) << 1) + | ((f.read_register2 as u64) << 2) + | ((f.write_register as u64) << 3) + | ((f.alu as u64) << 4) + | ((f.add as u64) << 5) + | ((f.sub as u64) << 6) + | ((f.memory as u64) << 7) + | ((f.branch as u64) << 8) + | ((f.ecall as u64) << 9) + | ((op.branch_cond as u64) << 10); + let bytes = (f.rs1 as u64) + | ((f.rs2 as u64) << 8) + | ((f.rd as u64) << 16) + | ((f.half_instruction_length as u64) << 24) + | ((f.alu_flags as u64) << 32) + | ((f.mem_flags as u64) << 40); + let b = i * stride; + packed[b] = op.timestamp; + packed[b + 1] = op.decode.pc; + packed[b + 2] = op.decode.imm; + packed[b + 3] = op.next_pc; + packed[b + 4] = op.rvd; + packed[b + 5] = op.rv1; + packed[b + 6] = op.rv2; + packed[b + 7] = op.arg2; + packed[b + 8] = op.res; + packed[b + 9] = flags; + packed[b + 10] = bytes; + } + packed +} + +/// Build one CPU trace-table chunk on device: pack ops → GPU fill → a +/// `TraceTable` whose main matrix is resident on device (fed to the LDE with no +/// upload). The host main table is a zeroed placeholder sized for the correct +/// `num_rows`; it is never read on the GPU commit path (commit consumes the +/// device buffer, the aux build reads the resident snapshot, queries gather from +/// the device tree). Returns `None` if the GPU build fails, so the caller can +/// fall back to the CPU generator. +fn build_cpu_chunk( + chunk: &[CpuOperation], +) -> Option> { + let n = chunk.len(); + let num_rows = n.next_power_of_two().max(4); + let last_ts = chunk.last().map(|op| op.timestamp).unwrap_or(0); + let packed = pack_cpu_ops(chunk); + let dev = math_cuda::trace_cpu::gpu_build_cpu_trace(&packed, n, num_rows, last_ts).ok()?; + let mut trace = TraceTable::new_main( + crate::tables::types::zeroed_fe_vec(num_rows * cpu::cols::NUM_COLUMNS), + cpu::cols::NUM_COLUMNS, + 1, + ); + trace.set_main_input_dev(Arc::new(dev)); + Some(trace) +} + +/// Build all CPU trace-table chunks on device, mirroring `chunk_and_generate`'s +/// chunking (`max_rows`, one empty chunk when there are no ops). Returns `None` +/// when the kill-switch is set or any chunk fails to build, so the caller falls +/// back to the CPU generator. +pub(crate) fn gpu_build_cpu_trace_tables( + cpu_ops: &[CpuOperation], + max_rows: usize, +) -> Option>> { + gpu_build_tables(cpu_ops, max_rows, build_cpu_chunk) +} + +// ============================================================================= +// MEMW_R (register fast path — the biggest table, ~15M rows on ethrex) +// ============================================================================= + +/// Build one MEMW_R trace-table chunk on device: marshal the walked `RegRow`s into +/// the SoA the `memw_register_fill` kernel consumes, fill the 10 columns row-major +/// on device, and leave the matrix RESIDENT (fed to the LDE with no full-column +/// upload). The host main table is a zeroed placeholder sized to `num_rows`; it is +/// never read on the GPU commit path. The `old_*` come from the (correct, +/// precompile-inclusive) sequential walk — so this is program-agnostic — with only +/// the compact `RegRow` fields uploaded, not the full column matrix. Returns `None` +/// on GPU failure so the caller can fall back to the CPU fill. +fn build_memw_register_chunk( + chunk: &[RegRow], +) -> Option> { + let n = chunk.len(); + let num_rows = n.next_power_of_two().max(4); + + let mut reg_addr = Vec::with_capacity(n); + let mut ts = Vec::with_capacity(n); + let mut value = Vec::with_capacity(n); + let mut is_read = Vec::with_capacity(n); + let mut old_value = Vec::with_capacity(n); + let mut old_ts = Vec::with_capacity(n); + for r in chunk { + 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_ts.push(ot); + } + + let dev = math_cuda::trace_cpu::gpu_fill_memw_register( + ®_addr, &ts, &value, &is_read, &old_value, &old_ts, num_rows, + ) + .ok()?; + + let mut trace = TraceTable::new_main( + crate::tables::types::zeroed_fe_vec(num_rows * memw_register::cols::NUM_COLUMNS), + memw_register::cols::NUM_COLUMNS, + 1, + ); + trace.set_main_input_dev(Arc::new(dev)); + Some(trace) +} + +/// Build all MEMW_R trace-table chunks on device, mirroring `chunk_and_generate`'s +/// chunking (`max_rows`, one empty chunk when there are no rows). Returns `None` +/// when the kill-switch is set or any chunk fails to build, so the caller falls +/// back to the CPU fill. +pub(crate) fn gpu_build_memw_register_tables( + rows: &[RegRow], + max_rows: usize, +) -> Option>> { + gpu_build_tables(rows, max_rows, build_memw_register_chunk) +} + +// ============================================================================= +// MEMW_A (aligned memory — the biggest remaining uploader, ~2M rows on ethrex) +// ============================================================================= + +/// Pack one aligned `MemwOperation` into the stride-`MEMW_ALIGNED_STRIDE` layout +/// the `memw_aligned_fill` kernel consumes (see `trace_cpu.cu`). The op is already +/// walked (old_value/old_timestamp filled), so this only copies fields; value/old +/// (`[u32; 8]` each) pack two-per-u64. +pub(crate) fn pack_memw_aligned_op( + op: &MemwOperation, +) -> [u64; math_cuda::trace_cpu::MEMW_ALIGNED_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; + [ + flags, + op.base_address, + op.timestamp, + op.old_timestamp[0], + 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), + ] +} + +/// Build one MEMW_A trace-table chunk on device: pack the walked ops → GPU fill → +/// a resident matrix fed to the LDE with no full-column upload (only the compact +/// packed ops are H2D'd). Returns `None` on GPU failure so the caller falls back. +fn build_memw_aligned_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_ALIGNED_STRIDE); + for op in chunk { + packed.extend_from_slice(&pack_memw_aligned_op(op)); + } + let dev = math_cuda::trace_cpu::gpu_build_memw_aligned_trace(&packed, n, num_rows).ok()?; + let mut trace = TraceTable::new_main( + crate::tables::types::zeroed_fe_vec(num_rows * memw_aligned::cols::NUM_COLUMNS), + memw_aligned::cols::NUM_COLUMNS, + 1, + ); + trace.set_main_input_dev(Arc::new(dev)); + Some(trace) +} + +/// Build all MEMW_A trace-table chunks on device, mirroring `chunk_and_generate`'s +/// chunking. Returns `None` when the kill-switch is set or any chunk fails to build. +pub(crate) fn gpu_build_memw_aligned_tables( + ops: &[MemwOperation], + max_rows: usize, +) -> Option>> { + gpu_build_tables(ops, max_rows, build_memw_aligned_chunk) +} + +// ============================================================================= +// LOAD / STORE (per-row-map memory tables — same shape as MEMW_A) +// ============================================================================= + +/// Pack one `LoadOperation` into the `load_fill` stride (see `trace_cpu.cu`). +pub(crate) fn pack_load_op(op: &LoadOperation) -> [u64; math_cuda::trace_cpu::LOAD_STRIDE] { + let flags = (op.signed as u64) | ((op.width as u64) << 8); + let r = &op.res; + [ + flags, + op.base_address, + op.timestamp, + r[0] | (r[1] << 32), + r[2] | (r[3] << 32), + r[4] | (r[5] << 32), + r[6] | (r[7] << 32), + ] +} + +/// Pack one `StoreOperation` into the `store_fill` stride (see `trace_cpu.cu`). +pub(crate) fn pack_store_op(op: &StoreOperation) -> [u64; math_cuda::trace_cpu::STORE_STRIDE] { + let flags = (op.write2 as u64) | ((op.write4 as u64) << 1) | ((op.write8 as u64) << 2); + [flags, op.base_address, op.timestamp, op.value] +} + +fn build_load_chunk( + chunk: &[LoadOperation], +) -> 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::LOAD_STRIDE); + for op in chunk { + packed.extend_from_slice(&pack_load_op(op)); + } + let dev = math_cuda::trace_cpu::gpu_build_load_trace(&packed, n, num_rows).ok()?; + let mut trace = TraceTable::new_main( + crate::tables::types::zeroed_fe_vec(num_rows * load::cols::NUM_COLUMNS), + load::cols::NUM_COLUMNS, + 1, + ); + trace.set_main_input_dev(Arc::new(dev)); + Some(trace) +} + +pub(crate) fn gpu_build_load_tables( + ops: &[LoadOperation], + max_rows: usize, +) -> Option>> { + gpu_build_tables(ops, max_rows, build_load_chunk) +} + +fn build_store_chunk( + chunk: &[StoreOperation], +) -> 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::STORE_STRIDE); + for op in chunk { + packed.extend_from_slice(&pack_store_op(op)); + } + let dev = math_cuda::trace_cpu::gpu_build_store_trace(&packed, n, num_rows).ok()?; + let mut trace = TraceTable::new_main( + crate::tables::types::zeroed_fe_vec(num_rows * store::cols::NUM_COLUMNS), + store::cols::NUM_COLUMNS, + 1, + ); + trace.set_main_input_dev(Arc::new(dev)); + Some(trace) +} + +pub(crate) fn gpu_build_store_tables( + ops: &[StoreOperation], + max_rows: usize, +) -> Option>> { + gpu_build_tables(ops, max_rows, build_store_chunk) +} + +// ============================================================================= +// SHIFT (ALU table, no dedup — the kernel recomputes the shift aux on device) +// ============================================================================= + +/// Pack one `ShiftOperation` into the `shift_fill` stride (see `trace_cpu.cu`): +/// value (4×u16 in_halves), full shift_amount, and the flag bits. The kernel +/// recomputes bit_shift/zbs/x/y/limb_shift/out, so only 3 u64/op upload. +pub(crate) fn pack_shift_op(op: &ShiftOperation) -> [u64; math_cuda::trace_cpu::SHIFT_STRIDE] { + let h = &op.in_halves; + let value = + (h[0] as u64) | ((h[1] as u64) << 16) | ((h[2] as u64) << 32) | ((h[3] as u64) << 48); + let flags = (op.direction as u64) | ((op.signed as u64) << 1) | ((op.word_instr as u64) << 2); + [value, op.shift_amount, flags] +} + +fn build_shift_chunk( + chunk: &[ShiftOperation], +) -> 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::SHIFT_STRIDE); + for op in chunk { + packed.extend_from_slice(&pack_shift_op(op)); + } + let dev = math_cuda::trace_cpu::gpu_build_shift_trace(&packed, n, num_rows).ok()?; + let mut trace = TraceTable::new_main( + crate::tables::types::zeroed_fe_vec(num_rows * shift::cols::NUM_COLUMNS), + shift::cols::NUM_COLUMNS, + 1, + ); + trace.set_main_input_dev(Arc::new(dev)); + Some(trace) +} + +pub(crate) fn gpu_build_shift_tables( + ops: &[ShiftOperation], + max_rows: usize, +) -> Option>> { + gpu_build_tables(ops, max_rows, build_shift_chunk) +} + +// ============================================================================= +// LT (ALU dedup table): host per-chunk HashMap dedup → device fill (compute) +// ============================================================================= + +/// Pack one unique `LtOperation` + its multiplicity into the `lt_fill` stride. +pub(crate) fn pack_lt_op(op: &LtOperation, mult: u64) -> [u64; math_cuda::trace_cpu::LT_STRIDE] { + let flags = (op.signed as u64) | ((op.invert as u64) << 1); + [op.lhs, op.rhs, flags, mult] +} + +/// Build one LT trace-table chunk on device. Dedup happens HERE on the host (the +/// same per-chunk HashMap `generate_lt_trace` uses), then the unique ops + summed +/// multiplicities are filled on device. LT rides the permutation-invariant ALU +/// bus, so any row order is valid (validated by multiset/prove, not byte order). +fn build_lt_chunk( + chunk: &[LtOperation], +) -> Option> { + let mut map: HashMap = HashMap::new(); + for op in chunk { + *map.entry(op.clone()).or_insert(0) += 1; + } + let unique: Vec<(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(&pack_lt_op(op, *mult)); + } + let dev = math_cuda::trace_cpu::gpu_build_lt_trace(&packed, n, num_rows).ok()?; + let mut trace = TraceTable::new_main( + crate::tables::types::zeroed_fe_vec(num_rows * lt::cols::NUM_COLUMNS), + lt::cols::NUM_COLUMNS, + 1, + ); + trace.set_main_input_dev(Arc::new(dev)); + Some(trace) +} + +pub(crate) fn gpu_build_lt_tables( + ops: &[LtOperation], + max_rows: usize, +) -> Option>> { + // Each chunk dedups independently (matching `generate_lt_trace` per chunk). + gpu_build_tables(ops, max_rows, build_lt_chunk) +} diff --git a/prover/src/tables/memw.rs b/prover/src/tables/memw.rs index 338b85467..d786c2839 100644 --- a/prover/src/tables/memw.rs +++ b/prover/src/tables/memw.rs @@ -101,7 +101,7 @@ pub mod cols { // ========================================================================= /// A single MEMW operation to be added to the trace. -#[derive(Debug, Clone)] +#[derive(Debug, Clone, PartialEq, Eq)] pub struct MemwOperation { /// Whether this is a register access (true) or memory access (false) pub is_register: bool, diff --git a/prover/src/tables/memw_register.rs b/prover/src/tables/memw_register.rs index 590a55100..12f4c0d32 100644 --- a/prover/src/tables/memw_register.rs +++ b/prover/src/tables/memw_register.rs @@ -103,7 +103,7 @@ pub mod cols { /// - `old0/old1` = `old[0]`/`old[1]` /// - `old_ts_lo` = `old_timestamp[0] & 0xFFFF_FFFF` (the two words share old_timestamp, /// enforced by `is_register_op`; the upper limb is TIMESTAMP_1 = timestamp>>32) -#[derive(Debug, Clone, Copy)] +#[derive(Debug, Clone, Copy, PartialEq, Eq)] pub(crate) struct RegRow { /// Register index 0..=255 (`base_address / 2`); u16 keeps the struct at /// 32 bytes — it is the largest persisted array of the walk. @@ -158,6 +158,25 @@ impl RegRow { } } + /// Marshal to the SoA the on-device MEMW_R fill (`memw_register_fill`) + /// consumes: `(reg_addr = 2*address, timestamp, value, is_read, old_value, + /// old_ts)`. The old_timestamp upper limb is shared with TIMESTAMP_1 + /// (`timestamp >> 32`), matching the column encoding. + #[cfg(feature = "cuda")] + pub(crate) fn fill_soa(&self) -> (u32, u64, u64, u8, u64, u64) { + let value = (self.val0 as u64) | ((self.val1 as u64) << 32); + let old_value = (self.old0 as u64) | ((self.old1 as u64) << 32); + let old_ts = (self.old_ts_lo as u64) | ((self.timestamp >> 32) << 32); + ( + 2 * self.address as u32, + self.timestamp, + value, + self.is_read as u8, + old_value, + old_ts, + ) + } + /// Build a `RegRow` from a fully-formed register `MemwOperation`. Used on the /// precompile / commit / keccak / halt paths, which construct a `MemwOperation` /// first and only convert to the compact row once the op is known to route to diff --git a/prover/src/tables/mod.rs b/prover/src/tables/mod.rs index 0a86e4149..95a77a087 100644 --- a/prover/src/tables/mod.rs +++ b/prover/src/tables/mod.rs @@ -33,6 +33,8 @@ pub mod ecdas; pub mod ecsm; pub mod eq; pub mod global_memory; +#[cfg(feature = "cuda")] +pub mod gpu_trace; pub mod halt; pub mod keccak; pub mod keccak_rc; diff --git a/prover/src/tables/trace_builder.rs b/prover/src/tables/trace_builder.rs index 5c6b3085e..1c864dfc8 100644 --- a/prover/src/tables/trace_builder.rs +++ b/prover/src/tables/trace_builder.rs @@ -3187,6 +3187,29 @@ fn build_traces( // generate→spill order keeps trace memory bounded. let cpu_ops_ref = &cpu_ops; let gen_cpus = || { + // On-GPU CPU-table build (cuda, kill-switch off): the matrix is born + // resident on device and feeds the LDE with no host upload. Falls back + // to the CPU generator otherwise. Disk-spill needs the host matrix to + // spill, so it always uses the CPU path. + #[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_cpu_trace_tables(cpu_ops_ref, max_rows.cpu) + { + return Ok(tables); + } + } chunk_and_generate( cpu_ops_ref, max_rows.cpu, @@ -3205,6 +3228,31 @@ fn build_traces( ) }; let gen_memw_aligneds = || { + // On-GPU MEMW_A build (cuda, kill-switch off): the biggest remaining + // uploader (~2M rows on ethrex). Fill columns on device from the walked + // ops and leave the matrix resident (only the compact packed ops upload). + // Falls back to the CPU fill otherwise (and for disk-spill). + #[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_aligned_tables( + &memw_aligned_ops, + max_rows.memw_aligned, + ) + { + return Ok(tables); + } + } chunk_and_generate( &memw_aligned_ops, max_rows.memw_aligned, @@ -3214,6 +3262,31 @@ fn build_traces( ) }; let gen_memw_registers = || { + // On-GPU MEMW_R build (cuda, kill-switch off): fill the columns on device + // from the walked RegRows and leave the matrix resident so it feeds the LDE + // with no full-column upload. Falls back to the CPU fill otherwise (and for + // disk-spill, which needs the host matrix to spill). + #[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_register_tables( + &memw_register_rows, + max_rows.memw_register, + ) + { + return Ok(tables); + } + } // Direct-to-column fill from compact RegRows — the register fast path never // materializes a `Vec`. chunk_and_generate( @@ -3225,6 +3298,25 @@ fn build_traces( ) }; let gen_loads = || { + #[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_load_tables(&load_ops, max_rows.load) + { + return Ok(tables); + } + } chunk_and_generate( &load_ops, max_rows.load, @@ -3234,6 +3326,25 @@ fn build_traces( ) }; let gen_lts = || { + #[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_lt_tables(<_ops, max_rows.lt) + { + return Ok(tables); + } + } chunk_and_generate( <_ops, max_rows.lt, @@ -3243,6 +3354,25 @@ fn build_traces( ) }; let gen_shifts = || { + #[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_shift_tables(&shift_ops, max_rows.shift) + { + return Ok(tables); + } + } chunk_and_generate( &shift_ops, max_rows.shift, @@ -3300,6 +3430,25 @@ fn build_traces( ) }; let gen_stores = || { + #[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_store_tables(&store_ops, max_rows.store) + { + return Ok(tables); + } + } chunk_and_generate::( &store_ops, max_rows.store, diff --git a/prover/src/tests/gpu_fill_tests.rs b/prover/src/tests/gpu_fill_tests.rs new file mode 100644 index 000000000..ddc11ff34 --- /dev/null +++ b/prover/src/tests/gpu_fill_tests.rs @@ -0,0 +1,465 @@ +//! 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::{cpu, load, lt, memw_aligned, memw_register, 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" + ); +} diff --git a/prover/src/tests/mod.rs b/prover/src/tests/mod.rs index 4085f3bc5..8b6307243 100644 --- a/prover/src/tests/mod.rs +++ b/prover/src/tests/mod.rs @@ -46,6 +46,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)]