fix(cuda): address PR 691 AI review cleanups

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];

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::<u8>(nodes_dev_bytes) }?;
     let leaves_offset_bytes = commit.leaves_offset_bytes(num_leaves);

crypto/math-cuda/src/merkle.rs

@@ -3,7 +3,7 @@
 //! Matches `FieldElementVectorBackend<F, Keccak256, 32>::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<Vec<u8>> {
     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<Vec<u8>> {
     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"

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<Fp3>]) -> Vec<[u8; 32]> {
     let num_rows = parts[0].len();
     let num_parts = parts.len();