From e51b0a27dd5dbdf2ed78191a41a71b0876f9399a Mon Sep 17 00:00:00 2001 From: MauroFab Date: Fri, 26 Jun 2026 16:24:26 -0300 Subject: [PATCH] fix(cuda): align review cleanup with row-pair commits --- crypto/math-cuda/kernels/keccak.cu | 4 +-- crypto/math-cuda/src/lde.rs | 45 ++++++++++++------------ crypto/math-cuda/src/merkle.rs | 18 +++++++++- crypto/math-cuda/tests/comp_poly_tree.rs | 4 +-- 4 files changed, 44 insertions(+), 27 deletions(-) diff --git a/crypto/math-cuda/kernels/keccak.cu b/crypto/math-cuda/kernels/keccak.cu index 272b847dc..32dd75d47 100644 --- a/crypto/math-cuda/kernels/keccak.cu +++ b/crypto/math-cuda/kernels/keccak.cu @@ -159,8 +159,8 @@ extern "C" __global__ void keccak256_leaves_base_batched( uint64_t tid = (uint64_t)blockIdx.x * blockDim.x + threadIdx.x; if (tid >= num_rows) return; - // Bit-reverse the row index so we read columns at `br` but write the - // hashed leaf at `tid` — matching the CPU `commit_columns_bit_reversed`. + // Bit-reverse the row index so we read columns at `br` but write the hashed + // leaf at `tid` — matching the CPU per-row `commit_bit_reversed(.., 1)`. uint64_t br = __brevll(tid) >> (64 - log_num_rows); uint64_t st[25]; diff --git a/crypto/math-cuda/src/lde.rs b/crypto/math-cuda/src/lde.rs index 8dc05ba01..c3f59fd90 100644 --- a/crypto/math-cuda/src/lde.rs +++ b/crypto/math-cuda/src/lde.rs @@ -37,11 +37,11 @@ fn assert_u32_domain(n: usize, what: &str) { /// Output shape requested from the fused LDE + Keccak entry points. #[derive(Copy, Clone, PartialEq, Eq)] enum KeccakCommit { - /// Only the `lde_size` keccak-256 leaves; no inner-tree build. Caller - /// receives `lde_size * 32` bytes. + /// Only the keccak-256 leaves; no inner-tree build. Caller receives + /// `num_leaves * 32` bytes. LeavesOnly, /// Full Merkle tree: leaves at the tail + inner nodes built on-device. - /// Caller receives `(2*lde_size - 1) * 32` bytes. + /// Caller receives `(2*num_leaves - 1) * 32` bytes. FullTree, } @@ -595,12 +595,12 @@ pub fn coset_lde_batch_base_into( Ok(()) } -/// Fused LDE + Keccak-256 leaf hashing. Caller receives the `lde_size * 32` -/// bytes of leaf hashes in `hashed_leaves_out` (one 32-byte digest per output -/// row, in natural row order; leaves are computed reading columns at -/// bit-reversed rows, matching `commit_columns_bit_reversed` on the CPU -/// side). Thin wrapper over `coset_lde_batch_base_into_with_merkle_tree_inner` -/// with `LeavesOnly` — no inner-tree build, no device handle. +/// Fused LDE + row-pair Keccak-256 leaf hashing. Caller receives +/// `(lde_size / 2) * 32` bytes of leaf hashes in `hashed_leaves_out` (one +/// 32-byte digest per bit-reversed row pair, in natural leaf order, matching +/// `commit_bit_reversed(.., 2)` on the CPU side). Thin wrapper over +/// `coset_lde_batch_base_into_with_merkle_tree_inner` with `LeavesOnly` — no +/// inner-tree build, no device handle. pub fn coset_lde_batch_base_into_with_leaf_hash( columns: &[&[u64]], blowup_factor: usize, @@ -616,14 +616,15 @@ pub fn coset_lde_batch_base_into_with_leaf_hash( hashed_leaves_out, KeccakCommit::LeavesOnly, false, - 1, + 2, ) .map(|_| ()) } /// Like `coset_lde_batch_base_into_with_leaf_hash`, but also builds the full -/// Merkle tree on device and returns the `2*lde_size - 1` node buffer back -/// to the caller in `merkle_nodes_out` (byte length `(2*lde_size - 1) * 32`). +/// row-pair Merkle tree on device and returns the `2*(lde_size/2) - 1` node +/// buffer back to the caller in `merkle_nodes_out` (byte length +/// `(2*(lde_size/2) - 1) * 32`). /// /// The leaf hashes are never exposed to the caller — they stay on device and /// feed straight into the pair-hash tree kernel, avoiding the @@ -644,7 +645,7 @@ pub fn coset_lde_batch_base_into_with_merkle_tree( merkle_nodes_out, KeccakCommit::FullTree, false, - 1, + 2, ) .map(|_| ()) } @@ -799,8 +800,8 @@ fn coset_lde_batch_base_into_with_merkle_tree_inner( )?; // Allocate the device output buffer. In `LeavesOnly` mode this is just - // `lde_size * 32` bytes (the leaves themselves); in `FullTree` mode it's - // `(2*lde_size - 1) * 32` bytes (leaves in the tail + inner nodes filled + // `num_leaves * 32` bytes (the leaves themselves); in `FullTree` mode it's + // `(2*num_leaves - 1) * 32` bytes (leaves in the tail + inner nodes filled // below). `alloc` (not `alloc_zeros`) is safe because every byte is // written before any reader sees it: the keccak kernel fills the // leaves slab, the inner-tree pass (when present) fills the head. @@ -857,8 +858,8 @@ fn coset_lde_batch_base_into_with_merkle_tree_inner( } } -/// Ext3 variant of `coset_lde_batch_base_into_with_leaf_hash`: fused -/// LDE + Keccak-256 leaf hashing over ext3 columns. Thin wrapper over +/// Ext3 variant of `coset_lde_batch_base_into_with_leaf_hash`: fused LDE + +/// row-pair Keccak-256 leaf hashing over ext3 columns. Thin wrapper over /// `coset_lde_batch_ext3_into_with_merkle_tree_inner` with `LeavesOnly`. pub fn coset_lde_batch_ext3_into_with_leaf_hash( columns: &[&[u64]], @@ -877,14 +878,14 @@ pub fn coset_lde_batch_ext3_into_with_leaf_hash( hashed_leaves_out, KeccakCommit::LeavesOnly, false, - 1, + 2, ) .map(|_| ()) } /// Ext3 variant of the fused `coset_lde_batch_base_into_with_merkle_tree`. /// LDE + leaf hashing + inner-tree build, all on device; D2Hs only the LDE -/// evaluations and the full `2*lde_size - 1` node buffer. +/// evaluations and the full `2*(lde_size/2) - 1` row-pair node buffer. pub fn coset_lde_batch_ext3_into_with_merkle_tree( columns: &[&[u64]], n: usize, @@ -902,7 +903,7 @@ pub fn coset_lde_batch_ext3_into_with_merkle_tree( merkle_nodes_out, KeccakCommit::FullTree, false, - 1, + 2, ) .map(|_| ()) } @@ -1055,8 +1056,8 @@ fn coset_lde_batch_ext3_into_with_merkle_tree_inner( mb_u32, )?; - // Allocate device output buffer (LeavesOnly → lde_size*32; FullTree → - // (2*lde_size - 1)*32). Leaf kernel writes to the leaves slab; the + // Allocate device output buffer (LeavesOnly -> num_leaves*32; FullTree -> + // (2*num_leaves - 1)*32). Leaf kernel writes to the leaves slab; the // inner-tree pass (when present) fills the head. let mut nodes_dev = unsafe { stream.alloc::(nodes_dev_bytes) }?; let leaves_offset_bytes = commit.leaves_offset_bytes(num_leaves); diff --git a/crypto/math-cuda/src/merkle.rs b/crypto/math-cuda/src/merkle.rs index ef84c6da4..27f38ce0a 100644 --- a/crypto/math-cuda/src/merkle.rs +++ b/crypto/math-cuda/src/merkle.rs @@ -3,7 +3,7 @@ //! Matches `FieldElementVectorBackend::hash_data` in //! `crypto/crypto/src/merkle_tree/backends/field_element_vector.rs`, combined //! with the `reverse_index` row read pattern used in -//! `commit_columns_bit_reversed` at `crypto/stark/src/prover.rs`. +//! `commit_bit_reversed` at `crypto/stark/src/commitment.rs`. //! //! Caller supplies base-field column slabs already laid out as //! `[col * col_stride + row]` (the same layout `coset_lde_batch_base_into` @@ -38,6 +38,14 @@ pub fn keccak_leaves_base( ) -> Result> { assert!(num_rows.is_power_of_two()); assert!(rows_per_leaf == 1 || rows_per_leaf == 2); + assert!( + num_rows >= rows_per_leaf, + "num_rows must be at least rows_per_leaf" + ); + 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" @@ -80,6 +88,14 @@ pub fn keccak_leaves_ext3( ) -> Result> { assert!(num_rows.is_power_of_two()); assert!(rows_per_leaf == 1 || rows_per_leaf == 2); + assert!( + num_rows >= rows_per_leaf, + "num_rows must be at least rows_per_leaf" + ); + 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" diff --git a/crypto/math-cuda/tests/comp_poly_tree.rs b/crypto/math-cuda/tests/comp_poly_tree.rs index 29e33b6fe..51b826dd1 100644 --- a/crypto/math-cuda/tests/comp_poly_tree.rs +++ b/crypto/math-cuda/tests/comp_poly_tree.rs @@ -1,6 +1,6 @@ //! Parity: GPU fused `evaluate_poly_coset_batch_ext3_into_with_merkle_tree` //! (LDE + row-pair Keccak leaves + Merkle inner tree) against the same CPU -//! pipeline produced by `commit_composition_polynomial`. +//! row-pair commitment layout used by `commit_bit_reversed(.., 2)`. use math::field::element::FieldElement; use math::field::extensions_goldilocks::Degree3GoldilocksExtensionField; @@ -83,7 +83,7 @@ fn cpu_hash_pair(left: &[u8; 32], right: &[u8; 32]) -> [u8; 32] { out } -/// CPU: `commit_composition_polynomial`-style tree root over num_rows/2 leaves. +/// CPU: `commit_bit_reversed(.., 2)`-style tree root over num_rows/2 leaves. fn cpu_tree_nodes(parts: &[Vec]) -> Vec<[u8; 32]> { let num_rows = parts[0].len(); let num_parts = parts.len();