diff --git a/crypto/math-cuda/build.rs b/crypto/math-cuda/build.rs index 403b89430..61db35a6b 100644 --- a/crypto/math-cuda/build.rs +++ b/crypto/math-cuda/build.rs @@ -191,6 +191,10 @@ fn main() { compile_kernel("deep.cu", "deep.cubin", have_nvcc, &[]); compile_kernel("fri.cu", "fri.cubin", have_nvcc, &[]); compile_kernel("inverse.cu", "inverse.cubin", have_nvcc, &[]); + // RPX256 (XHash12) leaves and parents — the algebraic hash's device + // kernels. Pinned on the host by `tests/host_kat/rpx_host_kat.cpp`; the + // cubin needs no `-D`: RPX has no compile-time knob. + compile_kernel("rpx.cu", "rpx.cubin", have_nvcc, &[]); compile_kernel("logup.cu", "logup.cubin", have_nvcc, &[]); compile_kernel( "constraint_interp.cu", diff --git a/crypto/math-cuda/kernels/rpx.cu b/crypto/math-cuda/kernels/rpx.cu index cd5657d48..96553060f 100644 --- a/crypto/math-cuda/kernels/rpx.cu +++ b/crypto/math-cuda/kernels/rpx.cu @@ -1,8 +1,9 @@ // RPX256 (Rescue-Prime eXtended / XHash12) over Goldilocks at width 12 on // device — the permutation, the rate-8 overwrite-duplex leaf sponge and the -// Merkle parent. Phase 1 of the per-table GPU redo's lane K: arithmetic only. -// The leaf/tree kernels that stream table rows through `rpx::Sponge` and -// `rpx::compress` are phase 2 and follow `blake3.cu:338-620`'s shape. +// Merkle parent (lane K phase 1, arithmetic), then the leaf/tree kernels that +// stream table rows through `rpx::Sponge` and `rpx::compress` and the +// permutation probe (phase 2, the `extern "C"` surface at the end of the +// file, kernel for kernel the twin of `blake3.cu:338-620`). // // THE ORACLE is the Rust host implementation, byte for byte: // `prover/src/lfm/rpx.rs` `Rpx256::permute` (:280-316) — schedule FB E FB E FB E M, @@ -19,7 +20,7 @@ // PROVENANCE, layered exactly as the Rust module's own (rpx.rs "PROVENANCE"): // the FB round IS RPO's round with RPO's constants, and those are pinned by // nineteen EXTERNAL miden-crypto vectors, which `tests/host_kat/rpx_host_kat.cpp` -// replays through `fb_round` composed seven times. The E round (the cubic +// replays through `fb_round(s, r)` composed seven times. The E round (the cubic // extension) and the schedule have no published vector anywhere; they are // pinned to the Rust oracle's output (`prover/tests/rpx_host_kat_vectors.rs`) // and, independently, to naive polynomial arithmetic in the harness. @@ -63,6 +64,17 @@ #include "goldilocks.cuh" #include "ext3.cuh" +// `permute` is a REAL device function, never inlined (see its CODE SHAPE +// note). The host shim has no `__noinline__`; on the host the attribute only +// matters to the code-size probe, which asks for it explicitly. +#if defined(__CUDACC__) +#define RPX_NOINLINE __noinline__ +#elif defined(RPX_HOST_NOINLINE) +#define RPX_NOINLINE __attribute__((noinline)) +#else +#define RPX_NOINLINE +#endif + namespace rpx { enum : int { @@ -141,10 +153,13 @@ __device__ __constant__ uint64_t ARK2[NUM_ROUNDS][STATE_FELTS] = { 16460604813734957368ull, 9643968136937729763ull, 3611348709641382851ull, 18256379591337759196ull}, }; -// First ROW of the circulant MDS: `M[i][j] = MDS_CIRC_ROW[(j − i) mod 12]` -// (rpo.rs:107-114). Stored 32-bit so each MDS term is one 32×32→64 MAC. The -// row sums to 160, which is the bound `mds` rests on. -__device__ __constant__ uint32_t MDS_CIRC_ROW[STATE_FELTS] = {7, 23, 8, 26, 13, 10, 9, 7, 6, 22, 21, 8}; +// First ROW of the circulant MDS, `M[i][j] = ROW[(j − i) mod 12]` +// (rpo.rs:107-114), stored TWICE so that `MDS_CIRC_ROW2[j + 12 − i]` is the +// entry with no modulo: the output-lane loop in `mds` is rolled, so `i` is a +// runtime value there. 32-bit so each MDS term is one 32×32→64 MAC. The row +// sums to 160, which is the bound `mds` rests on. +__device__ __constant__ uint32_t MDS_CIRC_ROW2[2 * STATE_FELTS] = { + 7, 23, 8, 26, 13, 10, 9, 7, 6, 22, 21, 8, 7, 23, 8, 26, 13, 10, 9, 7, 6, 22, 21, 8}; // --------------------------------------------------------------------------- // Field-op forwarders. Under nvcc they are the `goldilocks.cuh` / `ext3.cuh` @@ -209,12 +224,15 @@ __device__ __forceinline__ void mds(uint64_t s[STATE_FELTS]) { hi32[j] = (uint32_t)(s[j] >> 32); } uint64_t out[STATE_FELTS]; -#pragma unroll + // Rolled over output lanes: twelve iterations of twenty-four MACs, one + // twelfth of the unrolled body's code for the same instruction count. +#pragma unroll 1 for (int i = 0; i < STATE_FELTS; ++i) { uint64_t acc_lo = 0, acc_hi = 0; // Σ c·l_j and Σ c·h_j, each < 2^40 + const int rot = STATE_FELTS - i; // MDS_CIRC_ROW2[j + rot] = ROW[(j − i) mod 12] #pragma unroll for (int j = 0; j < STATE_FELTS; ++j) { - const uint32_t c = MDS_CIRC_ROW[(j + STATE_FELTS - i) % STATE_FELTS]; + const uint32_t c = MDS_CIRC_ROW2[j + rot]; acc_lo += (uint64_t)c * (uint64_t)lo32[j]; acc_hi += (uint64_t)c * (uint64_t)hi32[j]; } @@ -243,7 +261,9 @@ __device__ __forceinline__ uint64_t sbox(uint64_t x) { template __device__ __forceinline__ uint64_t square_n(uint64_t x) { -#pragma unroll + // Rolled: the chain is serial anyway, and unrolled it is what made one + // permutation ~49k lines of PTX. The unroll factor here is a tuning knob. +#pragma unroll 1 for (int i = 0; i < N; ++i) x = fmul(x, x); return x; } @@ -310,33 +330,35 @@ __device__ __forceinline__ CubicExt ext_power7(const CubicExt &a) { } // --------------------------------------------------------------------------- -// Rounds. `R` is the round index into ARK1/ARK2 — a template parameter so the -// constant-bank offsets fold at compile time. +// Rounds. `r` is the round index into ARK1/ARK2 — a runtime value, so one copy +// of each round body serves every round; the constant-bank address is +// computed, which costs nothing next to the round's arithmetic. Every lane +// loop is rolled for the same reason (see `permute`'s CODE SHAPE note). // --------------------------------------------------------------------------- // FB: `MDS → +ARK1 → x^7 → MDS → +ARK2 → x^{1/7}` — RPO's round exactly // (rpo.rs:561-582, rpx.rs:283-295). RPX runs it at R = 0, 2, 4; RPO at 0..7. -template -__device__ __forceinline__ void fb_round(uint64_t s[STATE_FELTS]) { +__device__ __forceinline__ void fb_round(uint64_t s[STATE_FELTS], int r) { mds(s); -#pragma unroll - for (int i = 0; i < STATE_FELTS; ++i) s[i] = fadd(s[i], ARK1[R][i]); -#pragma unroll +#pragma unroll 1 + for (int i = 0; i < STATE_FELTS; ++i) s[i] = fadd(s[i], ARK1[r][i]); +#pragma unroll 1 for (int i = 0; i < STATE_FELTS; ++i) s[i] = sbox(s[i]); mds(s); -#pragma unroll - for (int i = 0; i < STATE_FELTS; ++i) s[i] = fadd(s[i], ARK2[R][i]); -#pragma unroll +#pragma unroll 1 + for (int i = 0; i < STATE_FELTS; ++i) s[i] = fadd(s[i], ARK2[r][i]); + // The twelve chains are independent; a GPU hides their latency with other + // warps, not by unrolling one thread's twelve chains into straight line. +#pragma unroll 1 for (int i = 0; i < STATE_FELTS; ++i) s[i] = inv_sbox(s[i]); } // E: `+ARK1 → x^7` in the cubic extension on four lane-triples, NO linear // layer (rpx.rs:296-307; the design, not an omission — rpx.rs:275-279). -template -__device__ __forceinline__ void ext_round(uint64_t s[STATE_FELTS]) { -#pragma unroll - for (int i = 0; i < STATE_FELTS; ++i) s[i] = fadd(s[i], ARK1[R][i]); -#pragma unroll +__device__ __forceinline__ void ext_round(uint64_t s[STATE_FELTS], int r) { +#pragma unroll 1 + for (int i = 0; i < STATE_FELTS; ++i) s[i] = fadd(s[i], ARK1[r][i]); +#pragma unroll 1 for (int e = 0; e < EXT_ELEMENTS; ++e) { const int base = e * EXT_DEGREE; CubicExt x; @@ -351,23 +373,32 @@ __device__ __forceinline__ void ext_round(uint64_t s[STATE_FELTS]) { } // M: `MDS → +ARK1`, a linear finish with no S-box (rpx.rs:308-313). -template -__device__ __forceinline__ void final_round(uint64_t s[STATE_FELTS]) { +__device__ __forceinline__ void final_round(uint64_t s[STATE_FELTS], int r) { mds(s); -#pragma unroll - for (int i = 0; i < STATE_FELTS; ++i) s[i] = fadd(s[i], ARK1[R][i]); +#pragma unroll 1 + for (int i = 0; i < STATE_FELTS; ++i) s[i] = fadd(s[i], ARK1[r][i]); } // The permutation: `FB E FB E FB E M` (rpx.rs:280-316), output CANONICAL. -__device__ void permute(uint64_t s[STATE_FELTS]) { - fb_round<0>(s); - ext_round<1>(s); - fb_round<2>(s); - ext_round<3>(s); - fb_round<4>(s); - ext_round<5>(s); - final_round<6>(s); -#pragma unroll +// +// ★ CODE SHAPE. A real (`RPX_NOINLINE`) function with rolled loops, on +// purpose. The first cubin build of the fully inlined, fully unrolled form ran +// 41 minutes and emitted 56 MB of PTX: one permutation was ~49k straight-line +// lines (the inverse S-box chain unrolled over twelve lanes, three times) and +// every leaf kernel carried one copy per `permute` call site — seven in the +// comp-poly kernel. Rolled and called, the whole file is a few thousand lines +// and every kernel shares one body. The price is loop overhead of order 10% of +// the permutation's instructions and the state living in local memory across +// the call; the `-Xptxas -v` report and the unroll factors of `square_n` and +// the lane loops are the tuning knobs, in that order. +RPX_NOINLINE __device__ void permute(uint64_t s[STATE_FELTS]) { +#pragma unroll 1 + for (int r = 0; r + 1 < NUM_ROUNDS; r += 2) { + fb_round(s, r); + ext_round(s, r + 1); + } + final_round(s, NUM_ROUNDS - 1); +#pragma unroll 1 for (int i = 0; i < STATE_FELTS; ++i) s[i] = goldilocks::canonical(s[i]); } @@ -449,3 +480,315 @@ __device__ __forceinline__ void compress(const uint64_t left[DIGEST_FELTS], } } // namespace rpx + +// =========================================================================== +// PHASE 2 — the device-facing surface: node bytes, leaf kernels, Merkle +// compressors and the permutation probe. Kernel for kernel the twin of +// `blake3.cu:338-620`, with the chain replaced by `rpx::Sponge` and the parent +// by `rpx::compress`. +// +// NODE BYTES. A node is four canonical felts, each stored as eight BIG-ENDIAN +// bytes — `digest_to_commitment` (algebraic_commit.rs:112-118) — so 32 bytes, +// the same slot width as a BLAKE3 or keccak node, and the device tree's bytes +// equal the host's. Digests leave `permute` canonical; a parent reads its +// children back with `commitment_to_digest`'s big-endian decoding. The device +// is little-endian, so both directions byte-swap (`bswap64`); the 32-byte node +// offsets inside a 256-byte-aligned `cuMemAlloc` buffer make the u64 accesses +// aligned, the same precondition the BLAKE3 u32 accesses rest on. +// +// A LEAF absorbs exactly the felt sequence the host leaf hashes: the same +// read pattern as the BLAKE3 kernel it twins (`leaves_bit_reversed_grouped`, +// commitment.rs:67 — bit-reversed rows, each column by column, an ext3 element +// as its three components), which is the sequence `felts_from_bytes` rebuilds +// from the leaf bytes, so `hash_bytes == hash_data` holds on device by +// construction. The felt count is known before the loop, as the overwrite +// duplex's padding flag needs it (A1). Raw `[0, 2^64)` storage is absorbed as +// is: the permutation is representation-independent, and the host +// canonicalises before serialising — same field value, same digest. +// =========================================================================== + +namespace rpx { + +// Byte-swap a u64: the device reads a host big-endian felt from a node and +// writes one back. Plain shifts so the host shim compiles it; nvcc lowers it +// to two PRMTs. +__device__ __forceinline__ uint64_t bswap64(uint64_t x) { + x = ((x & 0x00FF00FF00FF00FFull) << 8) | ((x >> 8) & 0x00FF00FF00FF00FFull); + x = ((x & 0x0000FFFF0000FFFFull) << 16) | ((x >> 16) & 0x0000FFFF0000FFFFull); + return (x << 32) | (x >> 32); +} + +// Four felts → one 32-byte node, `digest_to_commitment`'s layout. +__device__ __forceinline__ void store_digest_be(const uint64_t digest[DIGEST_FELTS], uint8_t *node) { + uint64_t *dst = reinterpret_cast(node); +#pragma unroll + for (int i = 0; i < DIGEST_FELTS; ++i) dst[i] = bswap64(digest[i]); +} + +// One 32-byte node → four felts, `commitment_to_digest`'s decoding. +__device__ __forceinline__ void load_digest_be(const uint8_t *node, uint64_t digest[DIGEST_FELTS]) { + const uint64_t *src = reinterpret_cast(node); +#pragma unroll + for (int i = 0; i < DIGEST_FELTS; ++i) digest[i] = bswap64(src[i]); +} + +// A Merkle parent in place in the node buffer — `parent` (algebraic_commit.rs +// :248-252): decode both children, `compress`, encode. Node buffer layout as +// `blake3.cu` / `keccak.cu` / the CPU `merkle.rs`: children at +// `nodes[parent_begin + n_pairs .. parent_begin + 3*n_pairs]`, parents at +// `nodes[parent_begin .. parent_begin + n_pairs]`, 32 bytes per node. +__device__ __forceinline__ void hash_merkle_parent(uint8_t *nodes, uint64_t parent_begin, + uint64_t n_pairs, uint64_t tid) { + uint64_t left[DIGEST_FELTS], right[DIGEST_FELTS], out[DIGEST_FELTS]; + load_digest_be(nodes + (parent_begin + n_pairs + 2 * tid) * 32, left); + load_digest_be(nodes + (parent_begin + n_pairs + 2 * tid + 1) * 32, right); + compress(left, right, out); + store_digest_be(out, nodes + (parent_begin + tid) * 32); +} + +} // namespace rpx + +// --------------------------------------------------------------------------- +// Leaf kernels. Twins of `blake3_leaves_*` / `blake3_comp_poly_leaves_ext3` / +// `blake3_fri_leaves_ext3`, argument for argument; one thread hashes one leaf. +// --------------------------------------------------------------------------- + +// Goldilocks BASE-FIELD leaf hashing, one leaf per bit-reversed row: column +// `c` of row `br` at `columns_base_ptr[c * col_stride + br]`. +// Twin of `blake3_leaves_base_batched` (`blake3.cu:346`). +extern "C" __global__ void rpx_leaves_base_batched( + const uint64_t *columns_base_ptr, + uint64_t col_stride, + uint64_t num_cols, + uint64_t num_rows, + uint64_t log_num_rows, + uint8_t *hashed_leaves_out) { + uint64_t tid = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x; + if (tid >= num_rows) return; + uint64_t br = __brevll(tid) >> (64 - log_num_rows); + + rpx::Sponge sp; + sp.init(num_cols); + for (uint64_t c = 0; c < num_cols; ++c) sp.absorb(columns_base_ptr[c * col_stride + br]); + uint64_t digest[rpx::DIGEST_FELTS]; + sp.finalize(digest); + rpx::store_digest_be(digest, hashed_leaves_out + tid * 32); +} + +// BASE-FIELD row-pair leaf hashing: leaf `tid` hashes bit-reversed rows +// `2*tid` and `2*tid+1`, each column by column, first row then second. +// `num_leaves = num_rows / 2`. Twin of `blake3_leaves_base_row_pair_batched`. +extern "C" __global__ void rpx_leaves_base_row_pair_batched( + const uint64_t *columns_base_ptr, + uint64_t col_stride, + uint64_t num_cols, + uint64_t num_rows, + uint64_t log_num_rows, + uint8_t *hashed_leaves_out) { + uint64_t tid = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x; + uint64_t num_leaves = num_rows >> 1; + if (tid >= num_leaves) return; + uint64_t br_0 = __brevll(2 * tid) >> (64 - log_num_rows); + uint64_t br_1 = __brevll(2 * tid + 1) >> (64 - log_num_rows); + + rpx::Sponge sp; + sp.init(2 * num_cols); + for (uint64_t c = 0; c < num_cols; ++c) sp.absorb(columns_base_ptr[c * col_stride + br_0]); + for (uint64_t c = 0; c < num_cols; ++c) sp.absorb(columns_base_ptr[c * col_stride + br_1]); + uint64_t digest[rpx::DIGEST_FELTS]; + sp.finalize(digest); + rpx::store_digest_be(digest, hashed_leaves_out + tid * 32); +} + +// EXT3 leaf hashing, one leaf per bit-reversed row, components in three +// separate base slabs: column `c` component `k` at +// `columns_base_ptr[(c*3 + k) * col_stride + br]`; an element is absorbed as +// `[comp0, comp1, comp2]`, matching `write_bytes_be`. +// Twin of `blake3_leaves_ext3_batched`. +extern "C" __global__ void rpx_leaves_ext3_batched( + const uint64_t *columns_base_ptr, + uint64_t col_stride, + uint64_t num_cols, // number of ext3 columns (NOT slabs) + uint64_t num_rows, + uint64_t log_num_rows, + uint8_t *hashed_leaves_out) { + uint64_t tid = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x; + if (tid >= num_rows) return; + uint64_t br = __brevll(tid) >> (64 - log_num_rows); + + rpx::Sponge sp; + sp.init(3 * num_cols); + for (uint64_t c = 0; c < num_cols; ++c) { +#pragma unroll + for (int k = 0; k < 3; ++k) { + sp.absorb(columns_base_ptr[(c * 3 + (uint64_t)k) * col_stride + br]); + } + } + uint64_t digest[rpx::DIGEST_FELTS]; + sp.finalize(digest); + rpx::store_digest_be(digest, hashed_leaves_out + tid * 32); +} + +// Composition-polynomial leaf hashing: each leaf absorbs `2 * num_parts` ext3 +// values from bit-reversed rows `2*tid` and `2*tid+1`, (row 0: parts) then +// (row 1: parts), three base components per value. +// Twin of `blake3_comp_poly_leaves_ext3`. +extern "C" __global__ void rpx_comp_poly_leaves_ext3( + const uint64_t *parts_base_ptr, + uint64_t col_stride, + uint64_t num_parts, + uint64_t num_rows, + uint64_t log_num_rows, + uint8_t *leaves_out) { + uint64_t tid = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x; + uint64_t num_leaves = num_rows >> 1; + if (tid >= num_leaves) return; + uint64_t br_0 = __brevll(2 * tid) >> (64 - log_num_rows); + uint64_t br_1 = __brevll(2 * tid + 1) >> (64 - log_num_rows); + + rpx::Sponge sp; + sp.init(2 * 3 * num_parts); + for (uint64_t p = 0; p < num_parts; ++p) { +#pragma unroll + for (int k = 0; k < 3; ++k) { + sp.absorb(parts_base_ptr[(p * 3 + (uint64_t)k) * col_stride + br_0]); + } + } + for (uint64_t p = 0; p < num_parts; ++p) { +#pragma unroll + for (int k = 0; k < 3; ++k) { + sp.absorb(parts_base_ptr[(p * 3 + (uint64_t)k) * col_stride + br_1]); + } + } + uint64_t digest[rpx::DIGEST_FELTS]; + sp.finalize(digest); + rpx::store_digest_be(digest, leaves_out + tid * 32); +} + +// FRI layer leaf hashing: each leaf absorbs two consecutive ext3 values from an +// interleaved eval vector `[a0,a1,a2,b0,b1,b2,...]` — six felts, so a single +// block, no padding flag (`6 mod 8 = 6` in capacity lane 8). No bit reversal. +// The host is `AlgebraicPairBackend::hash_data` (algebraic_commit.rs:318-329). +// Twin of `blake3_fri_leaves_ext3`. +extern "C" __global__ void rpx_fri_leaves_ext3( + const uint64_t *evals_interleaved, // 3 * num_evals u64s + uint64_t num_leaves, // = num_evals / 2 + uint8_t *leaves_out) { + uint64_t tid = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x; + if (tid >= num_leaves) return; + const uint64_t *pair = evals_interleaved + 2 * tid * 3; + + rpx::Sponge sp; + sp.init(6); +#pragma unroll + for (int i = 0; i < 6; ++i) sp.absorb(pair[i]); + uint64_t digest[rpx::DIGEST_FELTS]; + sp.finalize(digest); + rpx::store_digest_be(digest, leaves_out + tid * 32); +} + +// Row-major ROW-PAIR leaf hashing: leaf `tid` absorbs row `reverse_index(2*tid)` +// then row `reverse_index(2*tid+1)`, each `m` lanes read contiguously from +// `data + br * m`. `m` is the row stride in u64s: base trace = column count, +// ext3 trace = 3 * column count (an ext3 element's components are consecutive). +// Twin of `blake3_leaves_base_row_major_row_pair`; the fused LDE+commit +// pipeline's leaf kernel (`lde.rs` `coset_lde_row_major_inner`). +extern "C" __global__ void rpx_leaves_base_row_major_row_pair( + const uint64_t *data, + uint64_t m, + uint64_t num_rows, + uint64_t log_num_rows, + uint8_t *hashed_leaves_out) { + uint64_t tid = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x; + uint64_t num_leaves = num_rows >> 1; + if (tid >= num_leaves) return; + uint64_t br_0 = __brevll(2 * tid) >> (64 - log_num_rows); + uint64_t br_1 = __brevll(2 * tid + 1) >> (64 - log_num_rows); + const uint64_t *row_0 = data + br_0 * m; + const uint64_t *row_1 = data + br_1 * m; + + rpx::Sponge sp; + sp.init(2 * m); + for (uint64_t c = 0; c < m; ++c) sp.absorb(row_0[c]); + for (uint64_t c = 0; c < m; ++c) sp.absorb(row_1[c]); + uint64_t digest[rpx::DIGEST_FELTS]; + sp.finalize(digest); + rpx::store_digest_be(digest, hashed_leaves_out + tid * 32); +} + +// Column-range variant: each leaf absorbs only columns `[col_start, col_end)` +// of the two bit-reversed rows while `m` stays the full row stride — the CPU +// `commit_rows_bit_reversed_subset`, how preprocessed tables commit their +// precomputed and multiplicity column ranges to separate trees over one LDE. +// Twin of `blake3_leaves_base_row_major_row_pair_range`. +extern "C" __global__ void rpx_leaves_base_row_major_row_pair_range( + const uint64_t *data, + uint64_t m, + uint64_t col_start, + uint64_t col_end, + uint64_t num_rows, + uint64_t log_num_rows, + uint8_t *hashed_leaves_out) { + uint64_t tid = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x; + uint64_t num_leaves = num_rows >> 1; + if (tid >= num_leaves) return; + uint64_t br_0 = __brevll(2 * tid) >> (64 - log_num_rows); + uint64_t br_1 = __brevll(2 * tid + 1) >> (64 - log_num_rows); + const uint64_t *row_0 = data + br_0 * m; + const uint64_t *row_1 = data + br_1 * m; + + rpx::Sponge sp; + sp.init(2 * (col_end - col_start)); + for (uint64_t c = col_start; c < col_end; ++c) sp.absorb(row_0[c]); + for (uint64_t c = col_start; c < col_end; ++c) sp.absorb(row_1[c]); + uint64_t digest[rpx::DIGEST_FELTS]; + sp.finalize(digest); + rpx::store_digest_be(digest, hashed_leaves_out + tid * 32); +} + +// --------------------------------------------------------------------------- +// Merkle level / tail. Same launch split as BLAKE3's: one thread per pair per +// level while a level is wide, then ONE single-block launch that grid-strides +// every remaining level with a barrier between them. +// --------------------------------------------------------------------------- + +// One level of the inner tree: each thread compresses one child pair. +extern "C" __global__ void rpx_merkle_level(uint8_t *nodes, + uint64_t parent_begin, // in 32-byte nodes + uint64_t n_pairs) { + uint64_t tid = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x; + if (tid >= n_pairs) return; + rpx::hash_merkle_parent(nodes, parent_begin, n_pairs, tid); +} + +// Every remaining level from `level_begin` up to the root, in one block. +// Twin of `blake3_merkle_tail`. +extern "C" __global__ void rpx_merkle_tail(uint8_t *nodes, uint64_t level_begin) { + uint64_t lb = level_begin; + while (lb != 0) { + uint64_t nb = lb / 2; + uint64_t n_pairs = lb - nb; + for (uint64_t tid = threadIdx.x; tid < n_pairs; tid += blockDim.x) { + rpx::hash_merkle_parent(nodes, nb, n_pairs, tid); + } + __syncthreads(); + lb = nb; + } +} + +// --------------------------------------------------------------------------- +// Parity-harness entry point: `n` independent permutations, one thread each. +// The bare device permutation is otherwise unreachable from host code; this is +// what lets the GPU be checked against the host `Rpx256` (and the host-KAT's +// oracle tables) before any tree is built. Not on any production path. +// --------------------------------------------------------------------------- +extern "C" __global__ void rpx_permute_probe(const uint64_t *states, uint64_t n, uint64_t *out) { + uint64_t tid = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x; + if (tid >= n) return; + uint64_t s[rpx::STATE_FELTS]; +#pragma unroll + for (int i = 0; i < rpx::STATE_FELTS; ++i) s[i] = states[tid * rpx::STATE_FELTS + i]; + rpx::permute(s); +#pragma unroll + for (int i = 0; i < rpx::STATE_FELTS; ++i) out[tid * rpx::STATE_FELTS + i] = s[i]; +} diff --git a/crypto/math-cuda/src/device.rs b/crypto/math-cuda/src/device.rs index c62bf2855..cee3f304b 100644 --- a/crypto/math-cuda/src/device.rs +++ b/crypto/math-cuda/src/device.rs @@ -136,6 +136,7 @@ const LOGUP_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/logup.cubin const CONSTRAINT_INTERP_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/constraint_interp.cubin")); const BLAKE3_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/blake3.cubin")); +const RPX_CUBIN: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/rpx.cubin")); /// Number of CUDA streams in the pool. Larger pools let many rayon-parallel /// callers overlap on the GPU without serializing on stream ownership. The @@ -238,6 +239,22 @@ pub struct Backend { pub blake3_blocks_of_felts_probe: CudaFunction, pub blake3_chain_probe: CudaFunction, + // rpx.cubin — the RPX256 (XHash12) leaf kernels, Merkle level/tail + // compressors and the permutation probe (see `kernels/rpx.cu`). Twin for + // twin with the blake3 set above and in the same order; the probe is the + // only host-visible handle on the bare device permutation, which the parity + // tests check against the host `Rpx256`. + pub rpx_leaves_base_row_major_row_pair: CudaFunction, + pub rpx_leaves_base_row_major_row_pair_range: CudaFunction, + pub rpx_leaves_base_batched: CudaFunction, + pub rpx_leaves_base_row_pair_batched: CudaFunction, + pub rpx_leaves_ext3_batched: CudaFunction, + pub rpx_comp_poly_leaves_ext3: CudaFunction, + pub rpx_fri_leaves_ext3: CudaFunction, + pub rpx_merkle_level: CudaFunction, + pub rpx_merkle_tail: CudaFunction, + pub rpx_permute_probe: CudaFunction, + // barycentric.cubin pub barycentric_base_batched: CudaFunction, pub barycentric_ext3_batched: CudaFunction, @@ -448,6 +465,7 @@ impl Backend { let constraint_interp = ctx.load_module(Ptx::from_binary(CONSTRAINT_INTERP_CUBIN.to_vec()))?; let blake3 = ctx.load_module(Ptx::from_binary(BLAKE3_CUBIN.to_vec()))?; + let rpx = ctx.load_module(Ptx::from_binary(RPX_CUBIN.to_vec()))?; let mut streams = Vec::with_capacity(STREAM_POOL_SIZE); for _ in 0..STREAM_POOL_SIZE { @@ -560,6 +578,20 @@ impl Backend { blake3_serialize_felts_probe: blake3.load_function("blake3_serialize_felts_probe")?, blake3_blocks_of_felts_probe: blake3.load_function("blake3_blocks_of_felts_probe")?, blake3_chain_probe: blake3.load_function("blake3_chain_probe")?, + + rpx_leaves_base_row_major_row_pair: rpx + .load_function("rpx_leaves_base_row_major_row_pair")?, + rpx_leaves_base_row_major_row_pair_range: rpx + .load_function("rpx_leaves_base_row_major_row_pair_range")?, + rpx_leaves_base_batched: rpx.load_function("rpx_leaves_base_batched")?, + rpx_leaves_base_row_pair_batched: rpx + .load_function("rpx_leaves_base_row_pair_batched")?, + rpx_leaves_ext3_batched: rpx.load_function("rpx_leaves_ext3_batched")?, + rpx_comp_poly_leaves_ext3: rpx.load_function("rpx_comp_poly_leaves_ext3")?, + rpx_fri_leaves_ext3: rpx.load_function("rpx_fri_leaves_ext3")?, + rpx_merkle_level: rpx.load_function("rpx_merkle_level")?, + rpx_merkle_tail: rpx.load_function("rpx_merkle_tail")?, + rpx_permute_probe: rpx.load_function("rpx_permute_probe")?, barycentric_base_batched: bary.load_function("barycentric_base_batched")?, barycentric_ext3_batched: bary.load_function("barycentric_ext3_batched")?, barycentric_base_batched_strided: bary diff --git a/crypto/math-cuda/src/fri.rs b/crypto/math-cuda/src/fri.rs index 2615dc3d1..156fc3e11 100644 --- a/crypto/math-cuda/src/fri.rs +++ b/crypto/math-cuda/src/fri.rs @@ -205,7 +205,14 @@ impl FriCommitState { num_leaves_u64, &mut leaves_view, )?, - DeviceHash::Rpo256 | DeviceHash::Rpx256 | DeviceHash::Poseidon => unimplemented!( + DeviceHash::Rpx256 => crate::rpx::launch_fri_leaves_ext3( + self.stream.as_ref(), + be, + &out, + num_leaves_u64, + &mut leaves_view, + )?, + DeviceHash::Rpo256 | DeviceHash::Poseidon => unimplemented!( "{:?} device commit not yet ported (FRI layer ext3 leaves)", self.hash ), @@ -224,7 +231,13 @@ impl FriCommitState { &mut nodes_dev, num_leaves, )?, - DeviceHash::Rpo256 | DeviceHash::Rpx256 | DeviceHash::Poseidon => unimplemented!( + DeviceHash::Rpx256 => crate::rpx::build_inner_tree_levels( + self.stream.as_ref(), + be, + &mut nodes_dev, + num_leaves, + )?, + DeviceHash::Rpo256 | DeviceHash::Poseidon => unimplemented!( "{:?} device commit not yet ported (FRI layer inner tree levels)", self.hash ), diff --git a/crypto/math-cuda/src/lde.rs b/crypto/math-cuda/src/lde.rs index 33c45e557..5e91cea8f 100644 --- a/crypto/math-cuda/src/lde.rs +++ b/crypto/math-cuda/src/lde.rs @@ -1054,7 +1054,10 @@ fn build_inner_tree_levels_for( DeviceHash::Blake3 => { crate::blake3::build_inner_tree_levels(stream, be, nodes_dev, leaves_len) } - DeviceHash::Rpo256 | DeviceHash::Rpx256 | DeviceHash::Poseidon => { + DeviceHash::Rpx256 => { + crate::rpx::build_inner_tree_levels(stream, be, nodes_dev, leaves_len) + } + DeviceHash::Rpo256 | DeviceHash::Poseidon => { unimplemented!("{hash:?} device commit not yet ported (inner tree levels)") } } @@ -1141,7 +1144,16 @@ fn coset_lde_row_major_inner( log_lde, &mut leaves_view, )?, - DeviceHash::Rpo256 | DeviceHash::Rpx256 | DeviceHash::Poseidon => { + DeviceHash::Rpx256 => crate::rpx::launch_leaves_base_row_major_row_pair( + stream.as_ref(), + be, + &buf, + cols_u64, + lde_u64, + log_lde, + &mut leaves_view, + )?, + DeviceHash::Rpo256 | DeviceHash::Poseidon => { unimplemented!("{hash:?} device commit not yet ported (row-major row-pair leaves)") } } @@ -1339,7 +1351,18 @@ pub fn coset_lde_row_major_split_trees( log_lde, &mut leaves_view, )?, - DeviceHash::Rpo256 | DeviceHash::Rpx256 | DeviceHash::Poseidon => unimplemented!( + DeviceHash::Rpx256 => crate::rpx::launch_leaves_base_row_major_row_pair_range( + stream.as_ref(), + be, + &buf, + cols_u64, + col_start, + col_end, + lde_u64, + log_lde, + &mut leaves_view, + )?, + DeviceHash::Rpo256 | DeviceHash::Poseidon => unimplemented!( "{hash:?} device commit not yet ported (row-major row-pair leaves, column range)" ), } @@ -2154,7 +2177,23 @@ fn coset_lde_batch_base_into_with_merkle_tree_inner( lde_u64, &mut leaves_view, )?, - (DeviceHash::Rpo256 | DeviceHash::Rpx256 | DeviceHash::Poseidon, _) => { + (DeviceHash::Rpx256, true) => crate::rpx::launch_leaves_base_row_pair( + stream.as_ref(), + &buf, + col_stride_u64, + m as u64, + lde_u64, + &mut leaves_view, + )?, + (DeviceHash::Rpx256, false) => crate::rpx::launch_leaves_base( + stream.as_ref(), + &buf, + col_stride_u64, + m as u64, + lde_u64, + &mut leaves_view, + )?, + (DeviceHash::Rpo256 | DeviceHash::Poseidon, _) => { unimplemented!("{hash:?} device commit not yet ported (column-major base leaves)") } } @@ -2399,7 +2438,17 @@ fn evaluate_poly_coset_batch_ext3_into_inner( log_num_rows, &mut leaves_view, )?, - DeviceHash::Rpo256 | DeviceHash::Rpx256 | DeviceHash::Poseidon => { + DeviceHash::Rpx256 => crate::rpx::launch_comp_poly_leaves_ext3( + stream.as_ref(), + be, + &buf, + col_stride_u64, + num_parts_u64, + lde_u64, + log_num_rows, + &mut leaves_view, + )?, + DeviceHash::Rpo256 | DeviceHash::Poseidon => { unimplemented!("{hash:?} device commit not yet ported (comp-poly ext3 leaves)") } } diff --git a/crypto/math-cuda/src/lib.rs b/crypto/math-cuda/src/lib.rs index 6470d5bb3..e0af3781f 100644 --- a/crypto/math-cuda/src/lib.rs +++ b/crypto/math-cuda/src/lib.rs @@ -21,6 +21,7 @@ pub mod merkle; pub mod mmcs; pub mod ntt; pub mod nvtx; +pub mod rpx; // Re-exported for downstream crates so they can refer to CUDA primitive // types without depending on cudarc directly. @@ -42,8 +43,9 @@ pub type Result = std::result::Result; /// (keccak-256, or `Blake3Chain` at the compiled round count), exactly as on /// the host. /// -/// ★ The three ALGEBRAIC keys name hashes whose device kernels are not yet -/// ported. Every dispatch site in this crate carries an arm for them that +/// ★ Of the three ALGEBRAIC keys, [`DeviceHash::Rpx256`] is ported +/// ([`rpx`]); RPO256 and Poseidon name hashes whose device kernels are not. +/// Every dispatch site in this crate carries an arm for the unported keys that /// aborts with `unimplemented!` naming the hash — never an arm that launches a /// byte-hash kernel in its place. The keys exist ahead of their kernels so the /// host side (`stark::config::DeviceTreeBackend`) can name every commitment @@ -59,8 +61,8 @@ pub enum DeviceHash { /// RPO256 leaves and parents. No device kernels yet: every dispatch site /// aborts loudly on this key. Rpo256, - /// RPX256 (XHash12) leaves and parents. No device kernels yet: every - /// dispatch site aborts loudly on this key. + /// RPX256 (XHash12) leaves and parents — [`rpx`]'s kernels, the + /// algebraic family's first device port. Rpx256, /// ⚠ Poseidon-original — UNSHIPPABLE on the host side too; present so the /// key set mirrors `CommitmentHash` one-to-one. No device kernels. diff --git a/crypto/math-cuda/src/rpx.rs b/crypto/math-cuda/src/rpx.rs index 3468b1e29..ba8be5dab 100644 --- a/crypto/math-cuda/src/rpx.rs +++ b/crypto/math-cuda/src/rpx.rs @@ -1,21 +1,757 @@ -//! RPX256 (Rescue-Prime eXtended, width 12) device launch code — PHASE 2. +//! GPU RPX256 (XHash12) for Merkle commits — the leaf kernels, the parent/level +//! compressors, and the permutation probe that is the host's only handle on the +//! bare device permutation. //! -//! Phase 1 (lane K) ships the kernel SOURCE, `kernels/rpx.cu`, pinned by the -//! host known-answer harness `tests/host_kat/rpx_host_kat.cpp` (run with -//! `make test-rpx-host-kat`): the permutation, the rate-8 overwrite-duplex -//! leaf sponge and the Merkle parent, compiled on the host through -//! `cuda_host_shim.h` and checked against miden-crypto's RPO vectors and the -//! Rust oracle `prover/src/lfm/rpx.rs`. +//! Twin of [`crate::blake3`], launcher for launcher, so the two read against +//! each other; production dispatch reaches this module through +//! [`crate::DeviceHash::Rpx256`] from the fused LDE+commit pipelines +//! ([`crate::lde`]), the comp-poly tree builders (`stark::gpu_lde`) and the FRI +//! layer commits ([`crate::fri`]). //! -//! This module is the placeholder for the phase-2 launch wrappers — the leaf -//! kernels mirroring `blake3.rs`'s, `merkle_level` / `merkle_tail` over -//! `rpx::compress`, and the third arm in every `match hash` — and it is -//! deliberately NOT declared in `lib.rs` yet: nothing compiles it. Phase 2 -//! adds `pub mod rpx;` to `lib.rs` and -//! `compile_kernel("rpx.cu", "rpx.cubin", have_nvcc, &[])` to `build.rs` -//! (both lane D's files, requested through the coordinator). +//! # What a parent is //! -//! Digest layout contract for that work: a digest is four CANONICAL Goldilocks -//! felts (the kernel canonicalises every permutation output), serialised as -//! `digest_to_commitment` does — each felt's eight big-endian bytes, 32 bytes -//! per node, the same slot width as a BLAKE3 digest. +//! `hash_new_parent(left, right)` is ONE permutation of `[left ‖ right ‖ 0⁴]` +//! with the digest read from lanes 0..4 — `algebraic_commit::parent`, which is +//! `Rpx256::merge` in miden's terms. The children are decoded from their node +//! bytes as `commitment_to_digest` does (four big-endian u64s) and the parent +//! is encoded back as `digest_to_commitment` does, so device nodes are the +//! host's bytes. +//! +//! # What a leaf is +//! +//! The rate-8 OVERWRITE duplex `algebraic_commit::sponge_leaf` over the felt +//! sequence the host leaf hashes — the same read pattern as the BLAKE3 kernel +//! each leaf kernel twins (`leaves_bit_reversed_grouped`), which is exactly the +//! sequence `felts_from_bytes` rebuilds from the leaf bytes, so the +//! `hash_bytes == hash_data` contract holds on device by construction. +//! +//! # Coverage +//! +//! All seven leaf kernels, both tree compressors and the wrapper twins are +//! here; the `launch_*` functions are what the dispatch sites call. The +//! permutation itself is pinned without a GPU by +//! `tests/host_kat/rpx_host_kat.cpp` (`make test-rpx-host-kat`); the device +//! build is pinned against the host by `prover/tests/rpx_device_parity.rs`. + +use cudarc::driver::{CudaSlice, CudaStream, CudaViewMut, LaunchConfig, PushKernelArg}; +use std::sync::Arc; + +use crate::Result; +use crate::device::{Backend, backend}; +use crate::lde::pack_ext3_to_pinned_slabs; + +/// Felts in one permutation state. +pub const STATE_FELTS: usize = 12; + +/// Threads per block for the RPX kernels. +/// +/// [`crate::merkle`]'s 128 rather than BLAKE3's 256: a thread carries a +/// twelve-lane u64 state plus the inverse S-box's live temporaries, a register +/// footprint closer to keccak's 25 u64 lanes than to BLAKE3's 32 u32 words, and +/// 128 is the Blackwell register-file limit that keccak already runs at. To be +/// re-measured with `-Xptxas -v` (phase-2 gate); this is the safe default. +const RPX_BLOCK_DIM: u32 = 128; + +pub(crate) fn rpx_launch_cfg(num_threads: u64) -> LaunchConfig { + debug_assert!( + num_threads <= u32::MAX as u64, + "rpx_launch_cfg: num_threads ({num_threads}) exceeds u32 grid range", + ); + let grid = (num_threads as u32).div_ceil(RPX_BLOCK_DIM); + LaunchConfig { + grid_dim: (grid, 1, 1), + block_dim: (RPX_BLOCK_DIM, 1, 1), + shared_mem_bytes: 0, + } +} + +/// RPX leaf hashing over a base-field column buffer. Twin of +/// [`crate::blake3::leaves_base`], argument for argument. +/// +/// `columns` must hold `num_cols * col_stride` u64s with column `c`'s data at +/// `[c*col_stride .. c*col_stride + num_rows]`. `rows_per_leaf` selects the leaf +/// layout: `1` = one leaf per bit-reversed row (`num_rows` leaves), `2` = one +/// leaf per bit-reversed row pair (`num_rows/2` leaves, the trace-commit +/// layout). Returns `(num_rows / rows_per_leaf) * 32` hash bytes. +pub fn leaves_base( + columns: &[u64], + col_stride: usize, + num_cols: usize, + num_rows: usize, + rows_per_leaf: usize, +) -> Result> { + assert!(num_rows.is_power_of_two()); + assert!(rows_per_leaf == 1 || rows_per_leaf == 2); + assert!( + num_rows >= 2, + "num_rows must be at least 2 for bit-reversed GPU leaf hashing" + ); + assert!( + col_stride >= num_rows, + "col_stride must be >= num_rows to keep per-column reads in-bounds" + ); + let total = num_cols + .checked_mul(col_stride) + .expect("num_cols * col_stride overflows usize"); + assert!(columns.len() >= total); + let be = backend()?; + let stream = be.next_stream(); + let cols_dev = stream.clone_htod(&columns[..total])?; + let mut out_dev = stream.alloc_zeros::((num_rows / rows_per_leaf) * 32)?; + let launch = if rows_per_leaf == 2 { + launch_leaves_base_row_pair + } else { + launch_leaves_base + }; + launch( + stream.as_ref(), + &cols_dev, + col_stride as u64, + num_cols as u64, + num_rows as u64, + &mut out_dev.as_view_mut(), + )?; + let out = stream.clone_dtoh(&out_dev)?; + stream.synchronize()?; + Ok(out) +} + +/// Ext3 variant of [`leaves_base`]: columns arrive as three base slabs per ext3 +/// column, so `columns.len() >= num_cols * 3 * col_stride`. Twin of +/// [`crate::blake3::leaves_ext3`]. +pub fn leaves_ext3( + columns: &[u64], + col_stride: usize, + num_cols: usize, + num_rows: usize, + rows_per_leaf: usize, +) -> Result> { + assert!(num_rows.is_power_of_two()); + assert!(rows_per_leaf == 1 || rows_per_leaf == 2); + assert!( + num_rows >= 2, + "num_rows must be at least 2 for bit-reversed GPU leaf hashing" + ); + assert!( + col_stride >= num_rows, + "col_stride must be >= num_rows to keep per-column reads in-bounds" + ); + let total = num_cols + .checked_mul(3) + .and_then(|v| v.checked_mul(col_stride)) + .expect("num_cols * 3 * col_stride overflows usize"); + assert!(columns.len() >= total); + let be = backend()?; + let stream = be.next_stream(); + let cols_dev = stream.clone_htod(&columns[..total])?; + let mut out_dev = stream.alloc_zeros::((num_rows / rows_per_leaf) * 32)?; + // Row-pair ext3 leaves reuse the comp-poly kernel, as the keccak and BLAKE3 + // paths do: hashing all ext3 columns of rows `2i`, `2i+1` is the same + // traversal whether the columns are called "aux trace" or "parts". + let launch = if rows_per_leaf == 2 { + launch_ext3_row_pair + } else { + launch_leaves_ext3 + }; + launch( + stream.as_ref(), + &cols_dev, + col_stride as u64, + num_cols as u64, + num_rows as u64, + &mut out_dev.as_view_mut(), + )?; + let out = stream.clone_dtoh(&out_dev)?; + stream.synchronize()?; + Ok(out) +} + +pub(crate) fn launch_leaves_base( + stream: &CudaStream, + cols_dev: &CudaSlice, + col_stride: u64, + num_cols: u64, + num_rows: u64, + out_dev: &mut CudaViewMut<'_, u8>, +) -> Result<()> { + // The kernel computes `__brevll(tid) >> (64 - log_num_rows)`, which is UB + // for `log_num_rows == 0` (single-row trees are degenerate anyway). + debug_assert!(num_rows >= 2, "rpx leaf kernel: num_rows must be >= 2"); + let be = backend()?; + let log_num_rows = num_rows.trailing_zeros() as u64; + let cfg = rpx_launch_cfg(num_rows); + unsafe { + stream + .launch_builder(&be.rpx_leaves_base_batched) + .arg(cols_dev) + .arg(&col_stride) + .arg(&num_cols) + .arg(&num_rows) + .arg(&log_num_rows) + .arg(out_dev) + .launch(cfg)?; + } + Ok(()) +} + +pub(crate) fn launch_leaves_base_row_pair( + stream: &CudaStream, + cols_dev: &CudaSlice, + col_stride: u64, + num_cols: u64, + num_rows: u64, + out_dev: &mut CudaViewMut<'_, u8>, +) -> Result<()> { + debug_assert!( + num_rows >= 2, + "rpx row-pair leaf kernel: num_rows must be >= 2" + ); + let be = backend()?; + let log_num_rows = num_rows.trailing_zeros() as u64; + // One thread per leaf (= row pair). + let cfg = rpx_launch_cfg(num_rows >> 1); + unsafe { + stream + .launch_builder(&be.rpx_leaves_base_row_pair_batched) + .arg(cols_dev) + .arg(&col_stride) + .arg(&num_cols) + .arg(&num_rows) + .arg(&log_num_rows) + .arg(out_dev) + .launch(cfg)?; + } + Ok(()) +} + +pub(crate) fn launch_leaves_ext3( + stream: &CudaStream, + cols_dev: &CudaSlice, + col_stride: u64, + num_cols: u64, + num_rows: u64, + out_dev: &mut CudaViewMut<'_, u8>, +) -> Result<()> { + debug_assert!(num_rows >= 2, "rpx leaf kernel: num_rows must be >= 2"); + let be = backend()?; + let log_num_rows = num_rows.trailing_zeros() as u64; + let cfg = rpx_launch_cfg(num_rows); + unsafe { + stream + .launch_builder(&be.rpx_leaves_ext3_batched) + .arg(cols_dev) + .arg(&col_stride) + .arg(&num_cols) + .arg(&num_rows) + .arg(&log_num_rows) + .arg(out_dev) + .launch(cfg)?; + } + Ok(()) +} + +pub(crate) fn launch_ext3_row_pair( + stream: &CudaStream, + cols_dev: &CudaSlice, + col_stride: u64, + num_cols: u64, + num_rows: u64, + out_dev: &mut CudaViewMut<'_, u8>, +) -> Result<()> { + debug_assert!( + num_rows >= 2, + "rpx row-pair leaf kernel: num_rows must be >= 2" + ); + let be = backend()?; + let log_num_rows = num_rows.trailing_zeros() as u64; + let cfg = rpx_launch_cfg(num_rows >> 1); + unsafe { + stream + .launch_builder(&be.rpx_comp_poly_leaves_ext3) + .arg(cols_dev) + .arg(&col_stride) + .arg(&num_cols) + .arg(&num_rows) + .arg(&log_num_rows) + .arg(out_dev) + .launch(cfg)?; + } + Ok(()) +} + +/// Row-major row-pair leaf hashing: leaf `i` hashes the two consecutive +/// bit-reversed rows `reverse_index(2i)`, `reverse_index(2i+1)`, each `m` lanes +/// read contiguously from the row-major `data`. Matches the CPU +/// `commit_bit_reversed(.., 2)`; twin of [`crate::blake3::leaves_base_row_major_row_pair`]. +/// +/// Returns `(num_rows / 2) * 32` hash bytes. +pub fn leaves_base_row_major_row_pair(data: &[u64], m: usize, num_rows: usize) -> Result> { + leaves_row_major_row_pair_inner(data, m, 0, m, num_rows, false) +} + +/// Column-range variant of [`leaves_base_row_major_row_pair`]: each leaf hashes +/// only columns `[col_start, col_end)` of the row pair, while `m` stays the full +/// row stride. Matches the CPU `commit_rows_bit_reversed_subset`. +pub fn leaves_base_row_major_row_pair_range( + data: &[u64], + m: usize, + col_start: usize, + col_end: usize, + num_rows: usize, +) -> Result> { + leaves_row_major_row_pair_inner(data, m, col_start, col_end, num_rows, true) +} + +fn leaves_row_major_row_pair_inner( + data: &[u64], + m: usize, + col_start: usize, + col_end: usize, + num_rows: usize, + ranged: bool, +) -> Result> { + assert!(num_rows.is_power_of_two()); + assert!(num_rows >= 2, "num_rows must be at least 2"); + assert!( + col_start < col_end && col_end <= m, + "column range in bounds" + ); + let total = num_rows + .checked_mul(m) + .expect("num_rows * m overflows usize"); + assert!(data.len() >= total); + + let be = backend()?; + let stream = be.next_stream(); + let data_dev = stream.clone_htod(&data[..total])?; + let mut out_dev = stream.alloc_zeros::((num_rows / 2) * 32)?; + + let m_u64 = m as u64; + let num_rows_u64 = num_rows as u64; + let log_num_rows = num_rows.trailing_zeros() as u64; + if ranged { + launch_leaves_base_row_major_row_pair_range( + stream.as_ref(), + be, + &data_dev, + m_u64, + col_start as u64, + col_end as u64, + num_rows_u64, + log_num_rows, + &mut out_dev.as_view_mut(), + )?; + } else { + launch_leaves_base_row_major_row_pair( + stream.as_ref(), + be, + &data_dev, + m_u64, + num_rows_u64, + log_num_rows, + &mut out_dev.as_view_mut(), + )?; + } + let out = stream.clone_dtoh(&out_dev)?; + stream.synchronize()?; + Ok(out) +} + +/// Row-major ROW-PAIR leaf hashing under RPX: leaf `i` hashes the two +/// consecutive bit-reversed rows `reverse_index(2i)`, `reverse_index(2i+1)` +/// (each `m` lanes, read contiguously from the row-major `buf`), producing +/// `num_rows / 2` leaves. Device-buffer twin of the BLAKE3 launcher the fused +/// LDE pipeline dispatches against; matches the CPU `commit_bit_reversed(.., 2)`. +pub(crate) fn launch_leaves_base_row_major_row_pair( + stream: &CudaStream, + be: &Backend, + buf: &CudaSlice, + m: u64, + num_rows: u64, + log_num_rows: u64, + leaves_out: &mut CudaViewMut<'_, u8>, +) -> Result<()> { + // The kernel derives rows as `__brevll(2*tid + k) >> (64 - log_num_rows)`; + // a 64-bit shift is UB at `log_num_rows == 0`, so require `num_rows >= 2`. + debug_assert!( + num_rows >= 2, + "row-major row-pair rpx requires num_rows >= 2" + ); + let cfg = rpx_launch_cfg(num_rows >> 1); + unsafe { + stream + .launch_builder(&be.rpx_leaves_base_row_major_row_pair) + .arg(buf) + .arg(&m) + .arg(&num_rows) + .arg(&log_num_rows) + .arg(leaves_out) + .launch(cfg)?; + } + Ok(()) +} + +/// Column-range variant of [`launch_leaves_base_row_major_row_pair`]: leaves +/// hash only columns `[col_start, col_end)` of each bit-reversed row pair +/// (`m` stays the full row stride). Matches the CPU +/// `commit_rows_bit_reversed_subset`. +#[allow(clippy::too_many_arguments)] +pub(crate) fn launch_leaves_base_row_major_row_pair_range( + stream: &CudaStream, + be: &Backend, + buf: &CudaSlice, + m: u64, + col_start: u64, + col_end: u64, + num_rows: u64, + log_num_rows: u64, + leaves_out: &mut CudaViewMut<'_, u8>, +) -> Result<()> { + debug_assert!( + num_rows >= 2, + "row-major row-pair rpx requires num_rows >= 2" + ); + debug_assert!( + col_start < col_end && col_end <= m, + "column range in bounds" + ); + let cfg = rpx_launch_cfg(num_rows >> 1); + unsafe { + stream + .launch_builder(&be.rpx_leaves_base_row_major_row_pair_range) + .arg(buf) + .arg(&m) + .arg(&col_start) + .arg(&col_end) + .arg(&num_rows) + .arg(&log_num_rows) + .arg(leaves_out) + .launch(cfg)?; + } + Ok(()) +} + +/// Composition-part leaf hashing under RPX: leaf `i` hashes the ext3 +/// components of every part at the two bit-reversed rows `2i`, `2i+1`, read +/// from per-component slabs with stride `col_stride`. Device-buffer twin of +/// the BLAKE3 launch the comp-poly tree build dispatches against. +#[allow(clippy::too_many_arguments)] +pub(crate) fn launch_comp_poly_leaves_ext3( + stream: &CudaStream, + be: &Backend, + buf: &CudaSlice, + col_stride: u64, + num_parts: u64, + num_rows: u64, + log_num_rows: u64, + leaves_out: &mut CudaViewMut<'_, u8>, +) -> Result<()> { + debug_assert!(num_rows >= 2, "comp-poly rpx leaves require num_rows >= 2"); + let cfg = rpx_launch_cfg(num_rows >> 1); + unsafe { + stream + .launch_builder(&be.rpx_comp_poly_leaves_ext3) + .arg(buf) + .arg(&col_stride) + .arg(&num_parts) + .arg(&num_rows) + .arg(&log_num_rows) + .arg(leaves_out) + .launch(cfg)?; + } + Ok(()) +} + +/// FRI-layer leaf hashing under RPX: leaf `i` hashes the two consecutive ext3 +/// evals `2i`, `2i+1` of an interleaved eval vector (six felts — one block). +/// Device-buffer twin of the BLAKE3 launch the FRI layer commit dispatches +/// against; the host is `AlgebraicPairBackend::hash_data`. +pub(crate) fn launch_fri_leaves_ext3( + stream: &CudaStream, + be: &Backend, + evals: &CudaSlice, + num_leaves: u64, + leaves_out: &mut CudaViewMut<'_, u8>, +) -> Result<()> { + let cfg = rpx_launch_cfg(num_leaves); + unsafe { + stream + .launch_builder(&be.rpx_fri_leaves_ext3) + .arg(evals) + .arg(&num_leaves) + .arg(leaves_out) + .launch(cfg)?; + } + Ok(()) +} + +/// Walk the inner Merkle tree on device under RPX. `nodes_dev` already has the +/// `leaves_len` hashed leaves written into the tail; this fills in the inner +/// nodes bottom-up. Twin of [`crate::blake3::build_inner_tree_levels`], with +/// the same tail cutover: one single-block launch takes over once a level is no +/// wider than the block, where per-level launch overhead dominates the work. +pub(crate) fn build_inner_tree_levels( + stream: &CudaStream, + be: &Backend, + nodes_dev: &mut CudaSlice, + leaves_len: usize, +) -> Result<()> { + const TAIL_MAX_PAIRS: u64 = RPX_BLOCK_DIM as u64; + let mut level_begin: u64 = (leaves_len - 1) as u64; + while level_begin != 0 { + let new_begin = level_begin / 2; + let n_pairs = level_begin - new_begin; + if n_pairs <= TAIL_MAX_PAIRS { + let cfg = LaunchConfig { + grid_dim: (1, 1, 1), + block_dim: (RPX_BLOCK_DIM, 1, 1), + shared_mem_bytes: 0, + }; + unsafe { + stream + .launch_builder(&be.rpx_merkle_tail) + .arg(&mut *nodes_dev) + .arg(&level_begin) + .launch(cfg)?; + } + return Ok(()); + } + let cfg = rpx_launch_cfg(n_pairs); + unsafe { + stream + .launch_builder(&be.rpx_merkle_level) + .arg(&mut *nodes_dev) + .arg(&new_begin) + .arg(&n_pairs) + .launch(cfg)?; + } + level_begin = new_begin; + } + Ok(()) +} + +/// Given `hashed_leaves` of length `leaves_len * 32`, build the full RPX +/// Merkle tree on device and return the `(2*leaves_len - 1) * 32`-byte node +/// buffer in the standard layout: `nodes[0..leaves_len - 1]` are inner nodes +/// (root at index 0) and `nodes[leaves_len - 1..]` are the leaves themselves. +/// +/// Matches the CPU `crypto/crypto/src/merkle_tree/merkle.rs` construction, so +/// the result plugs into `MerkleTree::from_precomputed_nodes` the same way +/// [`crate::blake3::build_merkle_tree_on_device`]'s does. +/// +/// `leaves_len` must be a power of two and >= 2. +pub fn build_merkle_tree_on_device(hashed_leaves: &[u8]) -> Result> { + assert!(hashed_leaves.len().is_multiple_of(32)); + let leaves_len = hashed_leaves.len() / 32; + assert!(leaves_len >= 2, "tree needs at least two leaves"); + assert!( + leaves_len.is_power_of_two(), + "leaves_len must be a power of two" + ); + + let total_nodes = 2 * leaves_len - 1; + let be = backend()?; + let stream = be.next_stream(); + + // SAFETY: every byte is written before it is read — leaves by the H2D + // below, inner nodes by the level walk that follows. + let mut nodes_dev = unsafe { stream.alloc::(total_nodes * 32) }?; + let leaves_offset_bytes = (leaves_len - 1) * 32; + { + let mut slice = + nodes_dev.slice_mut(leaves_offset_bytes..leaves_offset_bytes + hashed_leaves.len()); + stream.memcpy_htod(hashed_leaves, &mut slice)?; + } + + build_inner_tree_levels(stream.as_ref(), be, &mut nodes_dev, leaves_len)?; + + let out = stream.clone_dtoh(&nodes_dev)?; + stream.synchronize()?; + Ok(out) +} + +/// Build the composition Merkle tree under RPX straight from a device-resident +/// slab buffer (`3*m` slabs of `lde_size` u64s, component `k` of part `c` at +/// `(c*3 + k) * lde_size` — the [`crate::lde::GpuLdeExt3`] layout). No host +/// staging and no H2D: the leaf kernel reads `buf` in place on `stream`. +/// +/// Twin of [`crate::blake3::build_comp_poly_tree_from_slabs_dev`]. +pub fn build_comp_poly_tree_from_slabs_dev( + stream: &Arc, + buf: &CudaSlice, + m: usize, + lde_size: usize, +) -> Result { + // Same sticky hook as the keccak and BLAKE3 twins: the comp-tree cliff test + // arms one counter and must reach it under whichever hash the build pins. + #[cfg(feature = "test-faults")] + crate::faults::check_sticky(&crate::faults::FAULT_COMP_TREE_STICKY)?; + assert!(m > 0); + assert!(lde_size.is_power_of_two() && lde_size >= 2); + assert_eq!(buf.len(), 3 * m * lde_size, "slab buffer shape"); + let num_leaves = lde_size / 2; + let tight_total_nodes = 2 * num_leaves - 1; + let be = backend()?; + + // SAFETY: every byte is written before it is read — leaves by the kernel + // below, inner nodes by the level walk after it. + let mut nodes_dev = unsafe { stream.alloc::(tight_total_nodes * 32) }?; + let leaves_offset_bytes = (num_leaves - 1) * 32; + { + let mut leaves_view = + nodes_dev.slice_mut(leaves_offset_bytes..leaves_offset_bytes + num_leaves * 32); + launch_ext3_row_pair( + stream.as_ref(), + buf, + lde_size as u64, + m as u64, + lde_size as u64, + &mut leaves_view, + )?; + } + build_inner_tree_levels(stream.as_ref(), be, &mut nodes_dev, num_leaves)?; + let mut root = [0u8; 32]; + stream.memcpy_dtoh(&nodes_dev.slice(0..32), &mut root)?; + stream.synchronize()?; + Ok(crate::lde::GpuMerkleTree { + nodes: Arc::new(nodes_dev), + leaves_len: num_leaves, + root, + }) +} + +/// Build the composition Merkle tree under RPX from host-side interleaved ext3 +/// parts, keeping the nodes device-resident so openings can gather paths on +/// device. `parts_interleaved` is `num_parts` slices, each `[a0,a1,a2,b0,b1,b2,…]` +/// of length `3*lde_size`. Leaves hash row pairs, so `leaves_len = lde_size / 2`. +/// +/// Twin of [`crate::blake3::build_comp_poly_tree_from_evals_ext3_keep`], and it +/// stages through the same pinned de-interleave buffer for the same reason. +pub fn build_comp_poly_tree_from_evals_ext3_keep( + parts_interleaved: &[&[u64]], +) -> Result { + #[cfg(feature = "test-faults")] + crate::faults::check_sticky(&crate::faults::FAULT_COMP_TREE_STICKY)?; + assert!(!parts_interleaved.is_empty()); + let m = parts_interleaved.len(); + let ext3_elems = parts_interleaved[0].len() / 3; + assert_eq!( + parts_interleaved[0].len(), + 3 * ext3_elems, + "ext3 buffer length must be 3 * lde_size" + ); + for p in parts_interleaved.iter() { + assert_eq!(p.len(), 3 * ext3_elems); + } + let lde_size = ext3_elems; + assert!(lde_size.is_power_of_two() && lde_size >= 2); + + let be = backend()?; + let stream = be.next_stream(); + let staging_slot = be.pinned_staging(); + + // Stage: de-interleave each part into 3 base slabs in pinned memory. + let mb = 3 * m; + let mut staging = staging_slot.lock().unwrap(); + staging.ensure_capacity(mb * lde_size, &be.ctx)?; + let pinned = unsafe { staging.as_mut_slice(mb * lde_size) }; + + pack_ext3_to_pinned_slabs(parts_interleaved, pinned, lde_size); + + // H2D the de-interleaved parts, then release the staging lock: the tree + // build reads the device `buf`, not `pinned`. Synchronize first so the async + // H2D has consumed `pinned` before it can be freed or reused. + let mut buf = stream.alloc_zeros::(mb * lde_size)?; + stream.memcpy_htod(&pinned[..mb * lde_size], &mut buf)?; + stream.synchronize()?; + drop(staging); + + build_comp_poly_tree_from_slabs_dev(&stream, &buf, m, lde_size) +} + +/// Build a FRI-layer Merkle tree on device under RPX from an interleaved ext3 +/// eval vector, returning the full host node buffer so tests can compare it byte +/// for byte against the CPU `AlgebraicPairBackend` tree. Each leaf hashes two +/// consecutive ext3 values; `num_leaves = evals.len() / 6`. Returns +/// `(2*num_leaves - 1) * 32` bytes in standard layout. +/// +/// Twin of [`crate::blake3::build_fri_layer_tree_from_evals_ext3`], and like it +/// a parity harness rather than a production path: production folds and commits +/// through [`crate::fri::FriCommitState::fold_and_commit_layer`], which +/// dispatches to the same two kernels. +pub fn build_fri_layer_tree_from_evals_ext3(evals: &[u64]) -> Result> { + assert!( + evals.len().is_multiple_of(6), + "evals must hold whole pair-leaves" + ); + let num_evals = evals.len() / 3; + let num_leaves = num_evals / 2; + assert!(num_leaves.is_power_of_two() && num_leaves >= 2); + let tight_total_nodes = 2 * num_leaves - 1; + + let be = backend()?; + let stream = be.next_stream(); + + let evals_dev = stream.clone_htod(evals)?; + // SAFETY: leaves are written by the kernel below, inner nodes by the level + // walk after it, before either is read. + let mut nodes_dev = unsafe { stream.alloc::(tight_total_nodes * 32) }?; + + let leaves_offset_bytes = (num_leaves - 1) * 32; + { + let mut leaves_view = + nodes_dev.slice_mut(leaves_offset_bytes..leaves_offset_bytes + num_leaves * 32); + launch_fri_leaves_ext3( + stream.as_ref(), + be, + &evals_dev, + num_leaves as u64, + &mut leaves_view, + )?; + } + + build_inner_tree_levels(stream.as_ref(), be, &mut nodes_dev, num_leaves)?; + + let out = stream.clone_dtoh(&nodes_dev)?; + stream.synchronize()?; + Ok(out) +} + +/// Parity harness: run the device permutation over `states` and return each +/// output state, canonical. +/// +/// Not a production path — the bare device permutation is otherwise +/// unreachable from host code, so without this there would be nothing to check +/// it against the host `Rpx256` (or the host-KAT's oracle tables) with before a +/// whole tree is built. Inputs may be raw `[0, 2^64)` storage. +pub fn permute_probe(states: &[[u64; STATE_FELTS]]) -> Result> { + if states.is_empty() { + return Ok(Vec::new()); + } + let n = states.len(); + let flat: Vec = states.iter().flatten().copied().collect(); + let be = backend()?; + let stream = be.next_stream(); + let states_dev = stream.clone_htod(&flat)?; + let mut out_dev = stream.alloc_zeros::(n * STATE_FELTS)?; + let n_u64 = n as u64; + let cfg = rpx_launch_cfg(n_u64); + unsafe { + stream + .launch_builder(&be.rpx_permute_probe) + .arg(&states_dev) + .arg(&n_u64) + .arg(&mut out_dev) + .launch(cfg)?; + } + let flat_out = stream.clone_dtoh(&out_dev)?; + stream.synchronize()?; + Ok(flat_out + .chunks_exact(STATE_FELTS) + .map(|c| { + let mut s = [0u64; STATE_FELTS]; + s.copy_from_slice(c); + s + }) + .collect()) +} diff --git a/crypto/math-cuda/tests/host_kat/rpx_host_kat.cpp b/crypto/math-cuda/tests/host_kat/rpx_host_kat.cpp index 6a3a6edec..64394d3b1 100644 --- a/crypto/math-cuda/tests/host_kat/rpx_host_kat.cpp +++ b/crypto/math-cuda/tests/host_kat/rpx_host_kat.cpp @@ -22,7 +22,7 @@ // extension against naive polynomial multiplication reduced by // `φ³ = φ + 1` — the same independent algorithms `rpx.rs`'s own tests use. // 3. ★ EXTERNAL: RPX's FB round IS RPO's round with RPO's constants. Seven -// `fb_round` compose to RPO256, and that composition is replayed over +// `fb_round(s, r)` compose to RPO256, and that composition is replayed over // miden-crypto's nineteen `hash_elements` vectors, which nothing in this // tree produced. That pins ARK1/ARK2, the MDS row and orientation, both // S-box chains and the sponge lane convention from outside. @@ -35,12 +35,17 @@ // reaches the output; raw (`≥ p`) and canonical inputs agree; outputs are // canonical. // 6. The cost model, COUNTED rather than asserted from a comment. +// 7. Every leaf kernel, both Merkle compressors and the permutation probe +// replayed thread by thread through the shim against the CPU leaf spec +// and the host parent — the read patterns and the node encoding, with the +// hash over them anchored by the layers above. // // Build and run with `make test-rpx-host-kat`. #include #include #include +#include #include #include "cuda_host_shim.h" @@ -311,13 +316,13 @@ void cubic_extension_matches_naive_polynomial_arithmetic() { // RPO256's permutation composed from the kernel's FB round — rpo.rs:567-583. void rpo_permute(uint64_t s[12]) { - rpx::fb_round<0>(s); - rpx::fb_round<1>(s); - rpx::fb_round<2>(s); - rpx::fb_round<3>(s); - rpx::fb_round<4>(s); - rpx::fb_round<5>(s); - rpx::fb_round<6>(s); + rpx::fb_round(s, 0); + rpx::fb_round(s, 1); + rpx::fb_round(s, 2); + rpx::fb_round(s, 3); + rpx::fb_round(s, 4); + rpx::fb_round(s, 5); + rpx::fb_round(s, 6); for (int i = 0; i < 12; ++i) s[i] = goldilocks::canonical(s[i]); } @@ -376,7 +381,7 @@ void seven_fb_rounds_reproduce_the_miden_rpo_vectors() { for (int d = 0; d < 4; ++d) ok = ok && s[d] == MIDEN_HASH_ELEMENTS[7][d]; check(ok, "permute([0..8 ‖ 0⁴]) must be miden's eight-element vector (compress layout)"); } - printf("★ EXTERNAL: seven fb_round = RPO256 vs miden-crypto hash_elements: %d/19 matched\n", + printf("★ EXTERNAL: seven fb_round(s, r) = RPO256 vs miden-crypto hash_elements: %d/19 matched\n", matched); } @@ -559,13 +564,13 @@ void the_canonicalisation_loop_is_pinned_by_the_witness() { uint64_t s[12]; memcpy(s, w->input, sizeof(s)); - rpx::fb_round<0>(s); - rpx::ext_round<1>(s); - rpx::fb_round<2>(s); - rpx::ext_round<3>(s); - rpx::fb_round<4>(s); - rpx::ext_round<5>(s); - rpx::final_round<6>(s); + rpx::fb_round(s, 0); + rpx::ext_round(s, 1); + rpx::fb_round(s, 2); + rpx::ext_round(s, 3); + rpx::fb_round(s, 4); + rpx::ext_round(s, 5); + rpx::final_round(s, 6); int twins = 0; for (int i = 0; i < 12; ++i) twins += (s[i] >= P) ? 1 : 0; check(twins > 0, "the witness must leave a raw lane >= p before the canonicalisation loop"); @@ -674,9 +679,9 @@ Counted count_ops(F f) { void the_cost_model_is_what_the_header_claims() { uint64_t s[12]; for (int i = 0; i < 12; ++i) s[i] = (uint64_t)i + 1; - const Counted fb = count_ops([&] { rpx::fb_round<0>(s); }); - const Counted ext = count_ops([&] { rpx::ext_round<1>(s); }); - const Counted fin = count_ops([&] { rpx::final_round<6>(s); }); + const Counted fb = count_ops([&] { rpx::fb_round(s, 0); }); + const Counted ext = count_ops([&] { rpx::ext_round(s, 1); }); + const Counted fin = count_ops([&] { rpx::final_round(s, 6); }); const Counted all = count_ops([&] { rpx::permute(s); }); const Counted rpo = count_ops([&] { rpo_permute(s); }); const Counted inv = count_ops([&] { (void)rpx::inv_sbox(s[0]); }); @@ -712,6 +717,284 @@ void the_cost_model_is_what_the_header_claims() { check(rpo.mul == 6384 && rpo.dot3 == 0 && rpo.add == 336, "RPO permutation must be 6384 mul / 336 add"); } +// =========================================================================== +// Layer 7 — the leaf kernels, the Merkle compressors and the probe, replayed +// thread by thread through the shim. +// +// What a leaf hashes is the CPU `leaves_bit_reversed_grouped` sequence — +// bit-reversed rows, each column by column, an ext3 element as its three +// components — and the hash over it is the `sponge_leaf` transcription pinned +// in layer 4. So each kernel is checked for its READ PATTERN and its node +// ENCODING (`digest_to_commitment`: four canonical felts, big-endian), with the +// permutation anchored separately above. Raw `[p, 2^64)` values are fed in, +// since that is what an LDE buffer holds. +// =========================================================================== + +uint64_t reverse_index(uint64_t i, uint32_t log_n) { return __brevll(i) >> (64 - log_n); } + +// The host leaf over `felts`: `sponge_leaf`, then `digest_to_commitment`. +void expected_leaf(const std::vector &felts, uint8_t out[32]) { + uint64_t d[4]; + ref_sponge_leaf(felts.data(), felts.size(), d); + for (int i = 0; i < 4; ++i) { + const uint64_t c = canon(d[i]); + for (int b = 0; b < 8; ++b) out[i * 8 + b] = (uint8_t)(c >> (56 - 8 * b)); + } +} + +std::string hex32(const uint8_t *b) { + std::string s(64, '\0'); + for (int i = 0; i < 32; ++i) snprintf(&s[i * 2], 3, "%02x", (unsigned)b[i]); + return s; +} + +void check_leaves(const std::vector &got, const std::vector> &want, + const char *what) { + if (got.size() != want.size() * 32) { + printf("FAIL %s: leaf count %zu vs %zu\n", what, got.size() / 32, want.size()); + ++failures; + return; + } + for (size_t i = 0; i < want.size(); ++i) { + uint8_t expect[32]; + expected_leaf(want[i], expect); + if (memcmp(got.data() + i * 32, expect, 32) != 0) { + printf("FAIL %s: leaf %zu\n got %s\n want %s\n", what, i, hex32(got.data() + i * 32).c_str(), + hex32(expect).c_str()); + ++failures; + return; + } + } +} + +// The two column-major base kernels: one leaf per bit-reversed row, and one per +// bit-reversed row pair. +void base_leaf_kernels_read_the_specified_felts() { + for (uint32_t log_n : {2u, 4u, 6u}) { + for (uint64_t num_cols : {1ull, 5ull, 8ull, 17ull}) { + const uint64_t n = 1ull << log_n; + std::vector cols(num_cols * n); + uint64_t seed = log_n * 31 + num_cols; + for (size_t i = 0; i < cols.size(); ++i) cols[i] = sample(seed, i); + { + std::vector out(n * 32, 0); + CUDA_HOST_FOR_EACH_THREAD(t, n) { + rpx_leaves_base_batched(cols.data(), n, num_cols, n, log_n, out.data()); + } + std::vector> want(n); + for (uint64_t leaf = 0; leaf < n; ++leaf) { + const uint64_t br = reverse_index(leaf, log_n); + for (uint64_t c = 0; c < num_cols; ++c) want[leaf].push_back(cols[c * n + br]); + } + check_leaves(out, want, "rpx_leaves_base_batched"); + } + { + const uint64_t num_leaves = n / 2; + std::vector out(num_leaves * 32, 0); + CUDA_HOST_FOR_EACH_THREAD(t, num_leaves) { + rpx_leaves_base_row_pair_batched(cols.data(), n, num_cols, n, log_n, out.data()); + } + std::vector> want(num_leaves); + for (uint64_t leaf = 0; leaf < num_leaves; ++leaf) { + for (int k = 0; k < 2; ++k) { + const uint64_t br = reverse_index(2 * leaf + k, log_n); + for (uint64_t c = 0; c < num_cols; ++c) want[leaf].push_back(cols[c * n + br]); + } + } + check_leaves(out, want, "rpx_leaves_base_row_pair_batched"); + } + } + } + printf("base leaf kernels: read pattern + node encoding match the CPU leaf spec\n"); +} + +// The ext3 kernels over the de-interleaved three-slab layout. +void ext3_leaf_kernels_read_the_specified_felts() { + for (uint32_t log_n : {2u, 4u, 6u}) { + for (uint64_t num_cols : {1ull, 3ull, 11ull}) { + const uint64_t n = 1ull << log_n; + std::vector cols(num_cols * 3 * n); + uint64_t seed = log_n * 17 + num_cols; + for (size_t i = 0; i < cols.size(); ++i) cols[i] = sample(seed, i); + { + std::vector out(n * 32, 0); + CUDA_HOST_FOR_EACH_THREAD(t, n) { + rpx_leaves_ext3_batched(cols.data(), n, num_cols, n, log_n, out.data()); + } + std::vector> want(n); + for (uint64_t leaf = 0; leaf < n; ++leaf) { + const uint64_t br = reverse_index(leaf, log_n); + for (uint64_t c = 0; c < num_cols; ++c) { + for (uint64_t k = 0; k < 3; ++k) want[leaf].push_back(cols[(c * 3 + k) * n + br]); + } + } + check_leaves(out, want, "rpx_leaves_ext3_batched"); + } + { + const uint64_t num_leaves = n / 2; + std::vector out(num_leaves * 32, 0); + CUDA_HOST_FOR_EACH_THREAD(t, num_leaves) { + rpx_comp_poly_leaves_ext3(cols.data(), n, num_cols, n, log_n, out.data()); + } + std::vector> want(num_leaves); + for (uint64_t leaf = 0; leaf < num_leaves; ++leaf) { + for (int j = 0; j < 2; ++j) { + const uint64_t br = reverse_index(2 * leaf + j, log_n); + for (uint64_t c = 0; c < num_cols; ++c) { + for (uint64_t k = 0; k < 3; ++k) want[leaf].push_back(cols[(c * 3 + k) * n + br]); + } + } + } + check_leaves(out, want, "rpx_comp_poly_leaves_ext3"); + } + } + } + printf("ext3 + comp-poly leaf kernels: read pattern + node encoding match the CPU leaf spec\n"); +} + +// FRI leaves: two consecutive ext3 values from an interleaved vector, six felts, +// no bit reversal — the Pair backend's `hash_data`. +void fri_leaf_kernel_reads_the_specified_felts() { + for (uint64_t num_leaves : {1ull, 2ull, 8ull, 33ull}) { + std::vector evals(num_leaves * 6); + uint64_t seed = 0xF41; + for (size_t i = 0; i < evals.size(); ++i) evals[i] = sample(seed, i); + std::vector out(num_leaves * 32, 0); + CUDA_HOST_FOR_EACH_THREAD(t, num_leaves) { rpx_fri_leaves_ext3(evals.data(), num_leaves, out.data()); } + std::vector> want(num_leaves); + for (uint64_t leaf = 0; leaf < num_leaves; ++leaf) { + for (int i = 0; i < 6; ++i) want[leaf].push_back(evals[leaf * 6 + i]); + } + check_leaves(out, want, "rpx_fri_leaves_ext3"); + } + printf("FRI leaf kernel: read pattern + node encoding match the Pair backend's leaf\n"); +} + +// The row-major row-pair kernels, plain and column-ranged, every non-empty +// range. +void row_major_leaf_kernels_read_the_specified_felts() { + for (uint32_t log_n : {2u, 4u, 6u}) { + for (uint64_t m : {1ull, 5ull, 13ull}) { + const uint64_t n = 1ull << log_n; + const uint64_t num_leaves = n / 2; + std::vector data(n * m); + uint64_t seed = log_n * 7 + m; + for (size_t i = 0; i < data.size(); ++i) data[i] = sample(seed, i); + { + std::vector out(num_leaves * 32, 0); + CUDA_HOST_FOR_EACH_THREAD(t, num_leaves) { + rpx_leaves_base_row_major_row_pair(data.data(), m, n, log_n, out.data()); + } + std::vector> want(num_leaves); + for (uint64_t leaf = 0; leaf < num_leaves; ++leaf) { + for (int k = 0; k < 2; ++k) { + const uint64_t br = reverse_index(2 * leaf + k, log_n); + for (uint64_t c = 0; c < m; ++c) want[leaf].push_back(data[br * m + c]); + } + } + check_leaves(out, want, "rpx_leaves_base_row_major_row_pair"); + } + for (uint64_t cs = 0; cs < m; ++cs) { + for (uint64_t ce = cs + 1; ce <= m; ++ce) { + std::vector out(num_leaves * 32, 0); + CUDA_HOST_FOR_EACH_THREAD(t, num_leaves) { + rpx_leaves_base_row_major_row_pair_range(data.data(), m, cs, ce, n, log_n, out.data()); + } + std::vector> want(num_leaves); + for (uint64_t leaf = 0; leaf < num_leaves; ++leaf) { + for (int k = 0; k < 2; ++k) { + const uint64_t br = reverse_index(2 * leaf + k, log_n); + for (uint64_t c = cs; c < ce; ++c) want[leaf].push_back(data[br * m + c]); + } + } + check_leaves(out, want, "rpx_leaves_base_row_major_row_pair_range"); + } + } + } + } + printf("row-major leaf kernels: read pattern + node encoding match the CPU leaf spec, all column ranges\n"); +} + +// The host parent over two nodes: decode big-endian, compress, encode. +void expected_parent(const uint8_t *left, const uint8_t *right, uint8_t out[32]) { + uint64_t l[4], r[4], d[4]; + for (int i = 0; i < 4; ++i) { + l[i] = r[i] = 0; + for (int b = 0; b < 8; ++b) { + l[i] = (l[i] << 8) | left[i * 8 + b]; + r[i] = (r[i] << 8) | right[i * 8 + b]; + } + } + rpx::compress(l, r, d); + for (int i = 0; i < 4; ++i) { + for (int b = 0; b < 8; ++b) out[i * 8 + b] = (uint8_t)(canon(d[i]) >> (56 - 8 * b)); + } +} + +// The Merkle level kernel replayed thread by thread up a 16-leaf tree, and the +// tail kernel replayed as a one-thread block (the shim's barrier is a no-op, +// so a single thread walking every pair in order is the tail's sequential +// meaning), both against the host parent over the same node buffer. +void merkle_compressors_match_the_host_parent() { + const uint64_t num_leaves = 16; + const uint64_t total = 2 * num_leaves - 1; + // Nodes must be VALID digests (canonical big-endian felts) for the decode to + // be meaningful, so the leaves are hashes of random felts, not random bytes. + std::vector leaves(num_leaves * 32); + uint64_t seed = 0x3E11; + for (uint64_t i = 0; i < num_leaves; ++i) { + std::vector f = {sample(seed, i), sample(seed, i + 1000)}; + expected_leaf(f, leaves.data() + i * 32); + } + + std::vector want(total * 32, 0); + memcpy(want.data() + (num_leaves - 1) * 32, leaves.data(), leaves.size()); + for (uint64_t parent = num_leaves - 1; parent-- > 0;) { + expected_parent(want.data() + (2 * parent + 1) * 32, want.data() + (2 * parent + 2) * 32, + want.data() + parent * 32); + } + + // Level by level. + std::vector by_level(total * 32, 0); + memcpy(by_level.data() + (num_leaves - 1) * 32, leaves.data(), leaves.size()); + uint64_t level_begin = num_leaves - 1; + while (level_begin != 0) { + const uint64_t new_begin = level_begin / 2; + const uint64_t n_pairs = level_begin - new_begin; + CUDA_HOST_FOR_EACH_THREAD(t, n_pairs) { rpx_merkle_level(by_level.data(), new_begin, n_pairs); } + level_begin = new_begin; + } + check(by_level == want, "rpx_merkle_level must reproduce the host tree"); + + // The tail, in one go. + std::vector by_tail(total * 32, 0); + memcpy(by_tail.data() + (num_leaves - 1) * 32, leaves.data(), leaves.size()); + blockIdx.x = 0; + threadIdx.x = 0; + blockDim.x = 1; + rpx_merkle_tail(by_tail.data(), num_leaves - 1); + check(by_tail == want, "rpx_merkle_tail must reproduce the host tree"); + printf("Merkle compressors: level and tail kernels reproduce the host parent over a 16-leaf tree\n"); +} + +// The permutation probe replayed over the oracle table: pins its indexing. +void permute_probe_matches_the_oracle_table() { + std::vector in(NUM_RPX_PERMUTATION_VECTORS * 12), out(NUM_RPX_PERMUTATION_VECTORS * 12, 0); + for (int n = 0; n < NUM_RPX_PERMUTATION_VECTORS; ++n) { + for (int i = 0; i < 12; ++i) in[n * 12 + i] = RPX_PERMUTATION_VECTORS[n].input[i]; + } + CUDA_HOST_FOR_EACH_THREAD(t, NUM_RPX_PERMUTATION_VECTORS) { + rpx_permute_probe(in.data(), (uint64_t)NUM_RPX_PERMUTATION_VECTORS, out.data()); + } + bool ok = true; + for (int n = 0; n < NUM_RPX_PERMUTATION_VECTORS; ++n) { + for (int i = 0; i < 12; ++i) ok = ok && out[n * 12 + i] == RPX_PERMUTATION_VECTORS[n].output[i]; + } + check(ok, "rpx_permute_probe must reproduce the oracle table, raw"); + printf("permute probe: %d oracle states reproduced through the kernel entry point\n", + NUM_RPX_PERMUTATION_VECTORS); +} + } // namespace int main() { @@ -734,6 +1017,13 @@ int main() { every_input_lane_reaches_the_output(); printf("\n-- layer 6: cost model --\n"); the_cost_model_is_what_the_header_claims(); + printf("\n-- layer 7: leaf kernels, Merkle compressors and the probe, replayed thread by thread --\n"); + base_leaf_kernels_read_the_specified_felts(); + ext3_leaf_kernels_read_the_specified_felts(); + fri_leaf_kernel_reads_the_specified_felts(); + row_major_leaf_kernels_read_the_specified_felts(); + merkle_compressors_match_the_host_parent(); + permute_probe_matches_the_oracle_table(); if (failures != 0) { printf("\n*** %d FAILURE(S) ***\n", failures); return 1; diff --git a/crypto/stark/src/gpu_lde.rs b/crypto/stark/src/gpu_lde.rs index 887a7cc22..35310b409 100644 --- a/crypto/stark/src/gpu_lde.rs +++ b/crypto/stark/src/gpu_lde.rs @@ -1878,9 +1878,10 @@ where math_cuda::DeviceHash::Blake3 => { math_cuda::blake3::build_comp_poly_tree_from_evals_ext3_keep(&raw_parts) } - math_cuda::DeviceHash::Rpo256 - | math_cuda::DeviceHash::Rpx256 - | math_cuda::DeviceHash::Poseidon => unimplemented!( + math_cuda::DeviceHash::Rpx256 => { + math_cuda::rpx::build_comp_poly_tree_from_evals_ext3_keep(&raw_parts) + } + math_cuda::DeviceHash::Rpo256 | math_cuda::DeviceHash::Poseidon => unimplemented!( "{:?} device commit not yet ported (comp-poly tree from ext3 evals)", B::COMMITMENT_HASH ), @@ -1934,9 +1935,13 @@ where handle.m, handle.lde_size, ), - math_cuda::DeviceHash::Rpo256 - | math_cuda::DeviceHash::Rpx256 - | math_cuda::DeviceHash::Poseidon => unimplemented!( + math_cuda::DeviceHash::Rpx256 => math_cuda::rpx::build_comp_poly_tree_from_slabs_dev( + &stream, + handle.buf.as_ref(), + handle.m, + handle.lde_size, + ), + math_cuda::DeviceHash::Rpo256 | math_cuda::DeviceHash::Poseidon => unimplemented!( "{:?} device commit not yet ported (comp-poly tree from resident slabs)", B::COMMITMENT_HASH ), diff --git a/prover/tests/rpx_device_parity.rs b/prover/tests/rpx_device_parity.rs new file mode 100644 index 000000000..8006bd406 --- /dev/null +++ b/prover/tests/rpx_device_parity.rs @@ -0,0 +1,439 @@ +//! The RPX device kernels must produce the host's bytes — the phase-2 gate for +//! lane K, the twin of `crypto/math-cuda/tests/blake3_fused_parity.rs` (Batch +//! backend through the fused LDE+commit pipelines) plus the FRI-layer tree +//! (Pair backend), the comp-poly tree and the bare permutation. +//! +//! Lives in the prover crate rather than `math-cuda` because the host side — +//! `RpxStarkHash`, `AlgebraicBatchBackend`, `AlgebraicPairBackend`, `Rpx256` — +//! lives here; `math-cuda` is a dev-dependency of this crate, not the reverse. +//! +//! Every comparison is byte for byte or lane for lane against the production +//! host path; a tamper arm proves the equalities are not vacuous. Needs a GPU: +//! +//! cargo test -p lambda-vm-prover --release --features cuda --test rpx_device_parity -- --nocapture +#![cfg(feature = "cuda")] + +use crypto::merkle_tree::merkle::MerkleTree; +use crypto::merkle_tree::traits::IsMerkleTreeBackend; +use lambda_vm_prover::lfm::algebraic_commit::{ + AlgebraicBatchBackend, AlgebraicPairBackend, RpxCommit, RpxStarkHash, +}; +use lambda_vm_prover::lfm::hash::LfmHasher; +use lambda_vm_prover::lfm::rpx::Rpx256; +use lambda_vm_prover::tables::types::FE; +use math::fft::two_half_fft::TwoHalfTwiddles; +use math::field::element::FieldElement; +use math::field::extensions_goldilocks::Degree3GoldilocksExtensionField; +use math::field::goldilocks::GoldilocksField; +use math::polynomial::Polynomial; +use stark::prover::{GenericProver, IsStarkProver}; + +/// The RPX prover, named: these tests compare against the CUDA RPX kernels, so +/// the CPU side must say RPX explicitly. +type Prover = GenericProver; + +type Ext3 = Degree3GoldilocksExtensionField; +type Fp3 = FieldElement; +type Fp = FieldElement; + +const P: u64 = 0xFFFF_FFFF_0000_0001; + +/// splitmix64 — deterministic inputs from a seed, with no `rand` dependency +/// (the prover crate carries none for tests). +struct SplitMix(u64); + +impl SplitMix { + fn new(seed: u64) -> Self { + Self(seed) + } + + fn next(&mut self) -> u64 { + self.0 = self.0.wrapping_add(0x9E37_79B9_7F4A_7C15); + let mut z = self.0; + z = (z ^ (z >> 30)).wrapping_mul(0xBF58_476D_1CE4_E5B9); + z = (z ^ (z >> 27)).wrapping_mul(0x94D0_49BB_1331_11EB); + z ^ (z >> 31) + } +} +const COSET_OFFSET: u64 = 7; + +fn coset_weights(n: usize, g: u64) -> Vec { + let inv_n = Fp::from(n as u64).inv().unwrap(); + let g_fp = Fp::from_raw(g); + let mut w = Vec::with_capacity(n); + let mut cur = inv_n; + for _ in 0..n { + w.push(cur); + cur = &cur * &g_fp; + } + w +} + +fn coset_weights_u64(n: usize, g: u64) -> Vec { + coset_weights(n, g).iter().map(|w| *w.value()).collect() +} + +fn canonical(v: u64) -> u64 { + if v >= P { v - P } else { v } +} + +// =========================================================================== +// The bare permutation: device vs the host oracle, lane for lane. +// =========================================================================== + +/// Random states plus the two edge states the host-KAT names; raw (`≥ p`) +/// lanes included, since that is what an LDE buffer holds. +fn probe_states(seed: u64, n: usize) -> Vec<[u64; 12]> { + let mut rng = SplitMix::new(seed); + let mut states = vec![[0u64; 12], [P - 1; 12]]; + for k in 0..n { + let mut s = [0u64; 12]; + for (i, lane) in s.iter_mut().enumerate() { + // Every fifth lane is a raw twin `c + p` of a small canonical `c`. + *lane = if (k + i) % 5 == 0 { + rng.next() % 0xFFFF_FFFF + P + } else { + rng.next() % P + }; + } + states.push(s); + } + states +} + +#[test] +fn rpx_device_permutation_matches_the_host_oracle() { + let states = probe_states(0x0052_5058, 256); + let got = math_cuda::rpx::permute_probe(&states).expect("device permute probe"); + assert_eq!(got.len(), states.len()); + for (n, (input, out)) in states.iter().zip(got.iter()).enumerate() { + let want = Rpx256.permute(core::array::from_fn(|i| FE::from(input[i]))); + for (i, (o, w)) in out.iter().zip(want.iter()).enumerate() { + let w = canonical(*w.value()); + // RAW comparison: the device canonicalises its output, and that + // loop is part of what is pinned (R-952). + assert_eq!(*o, w, "state {n} lane {i}: device {o} vs host {w}"); + } + } +} + +// =========================================================================== +// The Batch backend through the fused row-major LDE + commit pipelines. +// =========================================================================== + +fn cpu_row_major_rpx_root( + columns: &[Vec], + blowup: usize, + weights: &[Fp], + inv_tw: &TwoHalfTwiddles, + fwd_tw: &TwoHalfTwiddles, +) -> [u8; 32] { + let n = columns[0].len(); + let num_cols = columns.len(); + let mut buf: Vec = vec![Fp::from(0u64); n * num_cols]; + for (c, col) in columns.iter().enumerate() { + for (r, &v) in col.iter().enumerate() { + buf[r * num_cols + c] = Fp::from_raw(v); + } + } + Polynomial::::coset_lde_full_expand_row_major::( + &mut buf, num_cols, blowup, weights, inv_tw, fwd_tw, + ) + .expect("CPU row-major LDE"); + let (_, root) = + Prover::::commit_rows_bit_reversed(&buf, num_cols) + .expect("CPU RPX commit"); + root +} + +/// Device fused row-major LDE + RPX leaves + Merkle, root only. +fn gpu_fused_rpx_root(columns: &[Vec], blowup: usize, weights_u64: &[u64]) -> [u8; 32] { + let n = columns[0].len(); + let num_cols = columns.len(); + let mut row_major = vec![0u64; n * num_cols]; + for (c, col) in columns.iter().enumerate() { + for (r, &v) in col.iter().enumerate() { + row_major[r * num_cols + c] = v; + } + } + let (handle, _lde) = math_cuda::lde::coset_lde_row_major_with_merkle_tree_keep( + &row_major, + None, + math_cuda::DeviceHash::Rpx256, + n, + num_cols, + blowup, + weights_u64, + true, + ) + .expect("fused RPX GPU pipeline"); + handle.tree.as_ref().expect("resident merkle tree").root +} + +#[test] +fn rpx_fused_base_root_matches_cpu() { + for log_n in [4usize, 6, 8, 10] { + for blowup in [2usize, 4] { + for num_cols in [1usize, 3, 8, 9] { + let n = 1usize << log_n; + let log_lde = (n * blowup).trailing_zeros() as usize; + let mut rng = SplitMix::new((log_n * 1000 + blowup * 100 + num_cols) as u64); + let columns: Vec> = (0..num_cols) + .map(|_| (0..n).map(|_| rng.next()).collect()) + .collect(); + let weights_u64 = coset_weights_u64(n, COSET_OFFSET); + let weights_fp = coset_weights(n, COSET_OFFSET); + let inv_tw = + TwoHalfTwiddles::::new(log_n, true).expect("inv twiddles"); + let fwd_tw = + TwoHalfTwiddles::::new(log_lde, false).expect("fwd twiddles"); + + let gpu_root = gpu_fused_rpx_root(&columns, blowup, &weights_u64); + let cpu_root = + cpu_row_major_rpx_root(&columns, blowup, &weights_fp, &inv_tw, &fwd_tw); + assert_eq!( + gpu_root, cpu_root, + "RPX fused root mismatch: log_n={log_n} blowup={blowup} num_cols={num_cols}" + ); + } + } + } +} + +fn rand_ext3(rng: &mut SplitMix) -> Fp3 { + Fp3::new([ + Fp::from_raw(rng.next()), + Fp::from_raw(rng.next()), + Fp::from_raw(rng.next()), + ]) +} + +fn cpu_ext3_row_major_rpx_root( + columns: &[Vec], + blowup: usize, + weights: &[Fp], + inv_tw: &TwoHalfTwiddles, + fwd_tw: &TwoHalfTwiddles, +) -> [u8; 32] { + let n = columns[0].len(); + let num_cols = columns.len(); + let mut buf: Vec = vec![Fp3::from(0u64); n * num_cols]; + for (c, col) in columns.iter().enumerate() { + for (r, v) in col.iter().enumerate() { + buf[r * num_cols + c] = *v; + } + } + Polynomial::::coset_lde_full_expand_row_major::( + &mut buf, num_cols, blowup, weights, inv_tw, fwd_tw, + ) + .expect("CPU ext3 row-major LDE"); + let (_, root) = Prover::::commit_rows_bit_reversed(&buf, num_cols) + .expect("CPU ext3 RPX commit"); + root +} + +#[test] +fn rpx_fused_ext3_root_matches_cpu() { + for log_n in [4usize, 6, 8] { + for blowup in [2usize, 4] { + for num_cols in [1usize, 3, 5] { + let n = 1usize << log_n; + let log_lde = (n * blowup).trailing_zeros() as usize; + let mut rng = SplitMix::new((log_n * 1000 + blowup * 100 + num_cols) as u64 + 4242); + let columns: Vec> = (0..num_cols) + .map(|_| (0..n).map(|_| rand_ext3(&mut rng)).collect()) + .collect(); + + // Row-major ext3 = row-major base with 3 * num_cols lanes. + let mut row_major = vec![0u64; n * num_cols * 3]; + for (c, col) in columns.iter().enumerate() { + for (r, v) in col.iter().enumerate() { + for k in 0..3 { + row_major[(r * num_cols + c) * 3 + k] = *v.value()[k].value(); + } + } + } + let weights_u64 = coset_weights_u64(n, COSET_OFFSET); + let weights_fp = coset_weights(n, COSET_OFFSET); + let inv_tw = + TwoHalfTwiddles::::new(log_n, true).expect("inv twiddles"); + let fwd_tw = + TwoHalfTwiddles::::new(log_lde, false).expect("fwd twiddles"); + + let (handle, _lde) = + math_cuda::lde::coset_lde_ext3_row_major_with_merkle_tree_keep( + &row_major, + math_cuda::DeviceHash::Rpx256, + n, + num_cols, + blowup, + &weights_u64, + true, + ) + .expect("fused ext3 RPX GPU pipeline"); + let gpu_root = handle.tree.as_ref().expect("resident merkle tree").root; + let cpu_root = + cpu_ext3_row_major_rpx_root(&columns, blowup, &weights_fp, &inv_tw, &fwd_tw); + assert_eq!( + gpu_root, cpu_root, + "RPX fused ext3 root mismatch: log_n={log_n} blowup={blowup} num_cols={num_cols}" + ); + } + } + } +} + +// =========================================================================== +// The comp-poly tree from interleaved ext3 parts (the `gpu_lde` site), against +// the same row-pair leaf layout committed on the CPU. +// =========================================================================== + +#[test] +fn rpx_comp_poly_tree_root_matches_cpu() { + for (log_lde, m) in [(4usize, 1usize), (6, 2), (10, 3), (12, 4)] { + let lde_size = 1usize << log_lde; + let mut rng = SplitMix::new((log_lde * 10 + m) as u64 + 99); + let parts: Vec> = (0..m) + .map(|_| (0..lde_size).map(|_| rand_ext3(&mut rng)).collect()) + .collect(); + + // Interleaved `[a0,a1,a2,b0,b1,b2,…]` per part for the device. + let parts_u64: Vec> = parts + .iter() + .map(|p| { + p.iter() + .flat_map(|e| e.value().iter().map(|c| *c.value())) + .collect() + }) + .collect(); + let raw_parts: Vec<&[u64]> = parts_u64.iter().map(|p| p.as_slice()).collect(); + let dev_tree = math_cuda::rpx::build_comp_poly_tree_from_evals_ext3_keep(&raw_parts) + .expect("device comp-poly tree"); + assert_eq!(dev_tree.leaves_len, lde_size / 2); + + // Row-major with `m` ext3 columns: row r = [part_0[r], …, part_{m-1}[r]]. + let mut buf: Vec = vec![Fp3::from(0u64); lde_size * m]; + for (c, p) in parts.iter().enumerate() { + for (r, v) in p.iter().enumerate() { + buf[r * m + c] = *v; + } + } + let (_, cpu_root) = Prover::::commit_rows_bit_reversed(&buf, m) + .expect("CPU comp-poly commit"); + assert_eq!( + dev_tree.root, cpu_root, + "RPX comp-poly root mismatch: log_lde={log_lde} m={m}" + ); + } +} + +// =========================================================================== +// The Pair backend: the FRI-layer tree, node for node. +// =========================================================================== + +fn fri_layer_parity(log_num_leaves: u32, seed: u64) { + let num_leaves = 1usize << log_num_leaves; + let mut rng = SplitMix::new(seed); + let evals: Vec = (0..num_leaves * 2).map(|_| rand_ext3(&mut rng)).collect(); + + let mut evals_u64 = Vec::with_capacity(evals.len() * 3); + for e in &evals { + for c in e.value().iter() { + evals_u64.push(*c.value()); + } + } + let leaves: Vec<[Fp3; 2]> = evals.chunks_exact(2).map(|c| [c[0], c[1]]).collect(); + let cpu_tree = MerkleTree::>::build(&leaves).unwrap(); + let cpu_nodes = cpu_tree.nodes(); + + let gpu_bytes = math_cuda::rpx::build_fri_layer_tree_from_evals_ext3(&evals_u64).unwrap(); + assert_eq!(cpu_nodes.len() * 32, gpu_bytes.len(), "node count"); + for (i, expected) in cpu_nodes.iter().enumerate() { + assert_eq!( + &gpu_bytes[i * 32..(i + 1) * 32], + &expected[..], + "node {i} mismatch at log_num_leaves={log_num_leaves}" + ); + } +} + +/// Small trees: every level fits the block, so the tail kernel builds the +/// whole tree in one launch. +#[test] +fn rpx_fri_layer_tree_small() { + for log in 1u32..=6 { + fri_layer_parity(log, 100 + log as u64); + } +} + +/// Deep enough that the per-level kernel runs first and hands over to the tail +/// partway up — the launch path a real commit takes. +#[test] +fn rpx_fri_layer_tree_medium() { + for log in [10u32, 12, 14] { + fri_layer_parity(log, 500 + log as u64); + } +} + +// =========================================================================== +// The column-range leaves (preprocessed tables), leaf for leaf against the +// Batch backend's `hash_data` over the same felts. +// =========================================================================== + +#[test] +fn rpx_row_major_range_leaves_match_cpu() { + let log_n = 6u32; + let n = 1usize << log_n; + let m = 7usize; + let mut rng = SplitMix::new(0xBEEF); + let data: Vec = (0..n * m).map(|_| rng.next()).collect(); + let reverse_index = |i: usize| -> usize { (i as u64).reverse_bits() as usize >> (64 - log_n) }; + + for (cs, ce) in [(0usize, m), (0, 3), (3, m), (2, 5)] { + let gpu = math_cuda::rpx::leaves_base_row_major_row_pair_range(&data, m, cs, ce, n) + .expect("device ranged leaves"); + assert_eq!(gpu.len(), (n / 2) * 32); + for leaf in 0..n / 2 { + let mut felts: Vec = Vec::with_capacity(2 * (ce - cs)); + for k in 0..2 { + let br = reverse_index(2 * leaf + k); + for c in cs..ce { + felts.push(Fp::from_raw(data[br * m + c])); + } + } + let want = + as IsMerkleTreeBackend>::hash_data( + &felts, + ); + assert_eq!( + &gpu[leaf * 32..(leaf + 1) * 32], + &want[..], + "ranged leaf {leaf} mismatch for columns [{cs}, {ce})" + ); + } + } +} + +// =========================================================================== +// Negative control: one corrupted input element must move the device root. +// =========================================================================== + +#[test] +fn rpx_fused_tamper_diverges() { + let n = 1usize << 6; + let num_cols = 3usize; + let mut rng = SplitMix::new(777); + let columns: Vec> = (0..num_cols) + .map(|_| (0..n).map(|_| rng.next()).collect()) + .collect(); + let weights_u64 = coset_weights_u64(n, COSET_OFFSET); + + let honest = gpu_fused_rpx_root(&columns, 2, &weights_u64); + let mut tampered = columns.clone(); + tampered[1][n / 2] ^= 1; + let forged = gpu_fused_rpx_root(&tampered, 2, &weights_u64); + assert_ne!( + honest, forged, + "a corrupted input element must move the root" + ); +}