refactor(stark): unify & clean up the commitment layer

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];
@@ -181,6 +181,51 @@ extern "C" __global__ void keccak256_leaves_base_batched(
     finalize_keccak256(st, rate_pos, hashed_leaves_out + tid * 32);
 }
 
+// ---------------------------------------------------------------------------
+// Goldilocks BASE-FIELD row-pair leaf hashing.
+//
+// Leaf `leaf_idx` hashes TWO consecutive bit-reversed rows
+//   br_0 = reverse_index(2*leaf_idx),  br_1 = reverse_index(2*leaf_idx + 1)
+// each written column-by-column in canonical BE (same per-row byte layout as
+// `keccak256_leaves_base_batched`), in (br_0 row: col 0..K-1) then (br_1 row:
+// col 0..K-1) order. `num_leaves = num_rows / 2`; writes 32 bytes to
+// `hashed_leaves_out[leaf_idx * 32 ..]`. Matches the CPU
+// `keccak_leaves_row_pair_bit_reversed` (rows_per_leaf = 2) — the base-field
+// analog of `keccak_comp_poly_leaves_ext3`.
+// ---------------------------------------------------------------------------
+extern "C" __global__ void keccak256_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);
+
+    uint64_t st[25];
+    #pragma unroll
+    for (int i = 0; i < 25; ++i) st[i] = 0;
+
+    uint32_t rate_pos = 0;
+    // First row (br_0): col 0..K-1.
+    for (uint64_t c = 0; c < num_cols; ++c) {
+        uint64_t v = columns_base_ptr[c * col_stride + br_0];
+        absorb_lane(st, rate_pos, bswap64(goldilocks::canonical(v)));
+    }
+    // Second row (br_1): col 0..K-1.
+    for (uint64_t c = 0; c < num_cols; ++c) {
+        uint64_t v = columns_base_ptr[c * col_stride + br_1];
+        absorb_lane(st, rate_pos, bswap64(goldilocks::canonical(v)));
+    }
+
+    finalize_keccak256(st, rate_pos, hashed_leaves_out + tid * 32);
+}
+
 // ---------------------------------------------------------------------------
 // Goldilocks EXT3 leaf hashing (3 base-field components per ext3 element).
 //
@@ -349,35 +394,49 @@ extern "C" __global__ void keccak_merkle_level(
 }
 
 // ---------------------------------------------------------------------------
-// Row-major base leaf hashing.
+// Row-major ROW-PAIR leaf hashing.
+//
+// Row-major analog of `keccak256_leaves_base_row_pair_batched` (which reads a
+// column-major slab): each leaf hashes TWO consecutive bit-reversed rows.
+// Leaf `tid` hashes row `reverse_index(2*tid)` followed by row
+// `reverse_index(2*tid + 1)`, each as `m` canonical big-endian lanes read from
+// the contiguous row-major buffer (`data + br * m`). `num_leaves = num_rows/2`;
+// writes 32 bytes to `hashed_leaves_out[tid*32 ..]`.
 //
-// Input layout: data[row * m + col] for `num_rows` rows and `m` columns.
-// For leaf `tid`, reads the bit-reversed row `br(tid)` — a contiguous slice
-// of `m` elements starting at data[br * m]. Coalesced when multiple threads
-// in the same warp process consecutive `tid` values (they read non-overlapping
-// rows, each a contiguous block of m u64s in order).
+// `m` is the row stride in u64s: base trace = num columns; ext3 trace = 3 *
+// num columns (an ext3 element's components c0,c1,c2 are consecutive, matching
+// the CPU `write_bytes_be`). Byte layout therefore equals the CPU
+// `commit_bit_reversed(.., ROWS_PER_LEAF=2)` and the verifier's
+// `verify_opening_pair` (queried row ‖ its symmetric counterpart, one leaf).
 // ---------------------------------------------------------------------------
-extern "C" __global__ void keccak256_leaves_base_row_major(
+extern "C" __global__ void keccak256_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;
-    if (tid >= num_rows) return;
-    uint64_t br = __brevll(tid) >> (64 - log_num_rows);
-    const uint64_t *row = data + br * m;
+    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;
 
     uint64_t st[25];
     #pragma unroll
     for (int i = 0; i < 25; ++i) st[i] = 0;
 
     uint32_t rate_pos = 0;
+    // First row (br_0): cols 0..m-1.
     for (uint64_t c = 0; c < m; ++c) {
-        uint64_t canon = goldilocks::canonical(row[c]);
-        uint64_t lane  = bswap64(canon);
-        absorb_lane(st, rate_pos, lane);
+        absorb_lane(st, rate_pos, bswap64(goldilocks::canonical(row_0[c])));
+    }
+    // Second row (br_1): cols 0..m-1.
+    for (uint64_t c = 0; c < m; ++c) {
+        absorb_lane(st, rate_pos, bswap64(goldilocks::canonical(row_1[c])));
     }
     finalize_keccak256(st, rate_pos, hashed_leaves_out + tid * 32);
 }

crypto/math-cuda/src/device.rs

@@ -147,8 +147,9 @@ pub struct Backend {
     pub matrix_transpose_strided: CudaFunction,
 
     // keccak.ptx
-    pub keccak256_leaves_base_row_major: CudaFunction,
+    pub keccak256_leaves_base_row_major_row_pair: CudaFunction,
     pub keccak256_leaves_base_batched: CudaFunction,
+    pub keccak256_leaves_base_row_pair_batched: CudaFunction,
     pub keccak256_leaves_ext3_batched: CudaFunction,
     pub keccak_comp_poly_leaves_ext3: CudaFunction,
     pub keccak_fri_leaves_ext3: CudaFunction,
@@ -247,9 +248,11 @@ impl Backend {
             ntt_dit_level_row_major: ntt.load_function("ntt_dit_level_row_major")?,
             pointwise_mul_row_major: ntt.load_function("pointwise_mul_row_major")?,
             matrix_transpose_strided: ntt.load_function("matrix_transpose_strided")?,
-            keccak256_leaves_base_row_major: keccak
-                .load_function("keccak256_leaves_base_row_major")?,
+            keccak256_leaves_base_row_major_row_pair: keccak
+                .load_function("keccak256_leaves_base_row_major_row_pair")?,
             keccak256_leaves_base_batched: keccak.load_function("keccak256_leaves_base_batched")?,
+            keccak256_leaves_base_row_pair_batched: keccak
+                .load_function("keccak256_leaves_base_row_pair_batched")?,
             keccak256_leaves_ext3_batched: keccak.load_function("keccak256_leaves_ext3_batched")?,
             keccak_comp_poly_leaves_ext3: keccak.load_function("keccak_comp_poly_leaves_ext3")?,
             keccak_fri_leaves_ext3: keccak.load_function("keccak_fri_leaves_ext3")?,