From b78d0d5f601309c74c22c36cc444b9dbd2f718cd Mon Sep 17 00:00:00 2001 From: MauroFab Date: Mon, 7 Sep 2026 16:37:35 -0300 Subject: [PATCH 1/4] feat(math-cuda): RPX256 leaf and tree kernels, launch module and device dispatch Lane K phase 2 of the per-table GPU redo. kernels/rpx.cu gains the device surface, kernel for kernel the twin of blake3.cu:338-620: seven leaf kernels (column-major base per row and per row pair, ext3 per row, comp-poly ext3 row pairs, FRI ext3 pairs, row-major row pairs plain and column-ranged), the Merkle level and tail compressors, and a permutation probe. A node is four canonical felts as big-endian bytes (digest_to_commitment), so device trees are the host's bytes; a leaf absorbs exactly the felt sequence the host leaf hashes (leaves_bit_reversed_grouped), which is the sequence felts_from_bytes rebuilds, so hash_bytes == hash_data holds on device by construction; a parent decodes its children as commitment_to_digest and runs compress. src/rpx.rs is the launch module, launcher for launcher the twin of blake3.rs at block dim 128 (keccak's register-file setting, to be measured with -Xptxas -v). build.rs compiles rpx.cu, lib.rs exports the module, device.rs loads the cubin and its ten functions, and the nine loud DeviceHash::Rpx256 arms in lde.rs, fri.rs and stark's gpu_lde.rs become launches; Rpo256 and Poseidon keep their loud arms. The host-KAT gains a seventh layer that replays every leaf kernel, both compressors and the probe thread by thread through the shim against the CPU leaf spec and the host parent, so the read patterns and the node encoding are pinned with no GPU. prover/tests/rpx_device_parity.rs (cuda-gated) is the device gate: the bare permutation against Rpx256, the fused base and ext3 roots against RpxStarkHash, the comp-poly tree root, the FRI-layer tree node for node against AlgebraicPairBackend, column-ranged leaves against hash_data, and a tamper control. --- crypto/math-cuda/build.rs | 4 + crypto/math-cuda/kernels/rpx.cu | 319 +++++++- crypto/math-cuda/src/device.rs | 32 + crypto/math-cuda/src/fri.rs | 17 +- crypto/math-cuda/src/lde.rs | 59 +- crypto/math-cuda/src/lib.rs | 10 +- crypto/math-cuda/src/rpx.rs | 766 +++++++++++++++++- .../math-cuda/tests/host_kat/rpx_host_kat.cpp | 290 +++++++ crypto/stark/src/gpu_lde.rs | 17 +- prover/tests/rpx_device_parity.rs | 439 ++++++++++ 10 files changed, 1915 insertions(+), 38 deletions(-) create mode 100644 prover/tests/rpx_device_parity.rs 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..e3520535f 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, @@ -449,3 +450,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..a2a5fa759 100644 --- a/crypto/math-cuda/src/rpx.rs +++ b/crypto/math-cuda/src/rpx.rs @@ -1,21 +1,751 @@ -//! 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 { + 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 { + 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..c47bf93a2 100644 --- a/crypto/math-cuda/tests/host_kat/rpx_host_kat.cpp +++ b/crypto/math-cuda/tests/host_kat/rpx_host_kat.cpp @@ -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" @@ -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..69237b7ba --- /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(0x5250_58, 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" + ); +} From 1b22e663cd5f6295eb6859474aa8a94813a30386 Mon Sep 17 00:00:00 2001 From: MauroFab Date: Mon, 7 Sep 2026 17:30:08 -0300 Subject: [PATCH 2/4] =?UTF-8?q?fix(math-cuda):=20give=20the=20RPX=20permut?= =?UTF-8?q?ation=20a=20compilable=20shape=20=E2=80=94=20one=20non-inlined?= =?UTF-8?q?=20body,=20rolled=20loops?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The first cubin build of rpx.cu ran 41 minutes and emitted 56 MB of PTX (1.33 M lines, 10 entries, zero functions): every device function was force-inlined and every loop fully unrolled, so one permutation was ~49k straight-line lines — the 72-step inverse S-box chain unrolled over twelve lanes three times — and each leaf kernel carried one copy per permute call site, seven in the comp-poly kernel. permute is now a real device function (RPX_NOINLINE, __noinline__ under nvcc; the host shim has no such attribute and the host KAT's inlining is irrelevant), the rounds take a runtime index inside a rolled round loop, the lane loops, the squaring chains, the extension triples and the constant adds are rolled with #pragma unroll 1, and the MDS output-lane loop is rolled over a doubled circulant row so the entry needs no modulo. The MDS inner loop stays unrolled (24 MACs). Nothing arithmetic moves: the host KAT passes unchanged, layer 6 counts the same 2736/144/300 ops, and the canonicalisation witness still pins the loop. Code-size proxy (x86-64 -O2 through the shim, instructions): the whole file drops from 5393 to 1581, with the ten kernels unchanged at 29–137 each (they only call now) and permute's inlined body a few hundred; PTX is expected in the low thousands of lines for the whole file. --- crypto/math-cuda/kernels/rpx.cu | 102 +++++++++++------- .../math-cuda/tests/host_kat/rpx_host_kat.cpp | 38 +++---- 2 files changed, 85 insertions(+), 55 deletions(-) diff --git a/crypto/math-cuda/kernels/rpx.cu b/crypto/math-cuda/kernels/rpx.cu index e3520535f..96553060f 100644 --- a/crypto/math-cuda/kernels/rpx.cu +++ b/crypto/math-cuda/kernels/rpx.cu @@ -20,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. @@ -64,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 { @@ -142,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` @@ -210,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]; } @@ -244,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; } @@ -311,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; @@ -352,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]); } 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 c47bf93a2..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. @@ -316,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]); } @@ -381,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); } @@ -564,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"); @@ -679,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]); }); From 856091bf4b2e62684bf7b295ec8d340730a2b633 Mon Sep 17 00:00:00 2001 From: MauroFab Date: Mon, 7 Sep 2026 17:38:17 -0300 Subject: [PATCH 3/4] test(math-cuda): mirror the sticky comp-tree fault hook into the RPX builders #961 gave both BLAKE3 comp-poly tree builders the test-faults-gated sticky hook the keccak twins carry, because the fault suite's comp-tree cliff test arms one counter and must reach it under whichever hash the build pins. The two RPX builders lacked it, so the suite would have gone silent again the day the pinned hash is RPX. Same four lines, first statement of build_comp_poly_tree_from_slabs_dev and build_comp_poly_tree_from_evals_ext3_keep. rpx.cu is untouched. --- crypto/math-cuda/src/rpx.rs | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/crypto/math-cuda/src/rpx.rs b/crypto/math-cuda/src/rpx.rs index a2a5fa759..ba8be5dab 100644 --- a/crypto/math-cuda/src/rpx.rs +++ b/crypto/math-cuda/src/rpx.rs @@ -582,6 +582,10 @@ pub fn build_comp_poly_tree_from_slabs_dev( 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"); @@ -626,6 +630,8 @@ pub fn build_comp_poly_tree_from_slabs_dev( 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; From f8d41476cfb572eaf99a38c821e65c4ad4a43a78 Mon Sep 17 00:00:00 2001 From: MauroFab Date: Mon, 7 Sep 2026 18:23:13 -0300 Subject: [PATCH 4/4] test(prover): regroup the RPX parity probe seed so the cuda-arm clippy accepts it make lint's cuda pass compiles the cuda-gated rpx_device_parity test and denies unusual_byte_groupings: probe_states(0x5250_58, 256) had digit groups of four and two. 0x0052_5058 is the same seed in groups of four. No other hex literal in the lane's Rust files mixes group sizes. --- prover/tests/rpx_device_parity.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/prover/tests/rpx_device_parity.rs b/prover/tests/rpx_device_parity.rs index 69237b7ba..8006bd406 100644 --- a/prover/tests/rpx_device_parity.rs +++ b/prover/tests/rpx_device_parity.rs @@ -103,7 +103,7 @@ fn probe_states(seed: u64, n: usize) -> Vec<[u64; 12]> { #[test] fn rpx_device_permutation_matches_the_host_oracle() { - let states = probe_states(0x5250_58, 256); + 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() {