fix: limit mulh kernel threads (#2337)

shuklaayush · web-flow · commit d7eab708f434 · 2026-01-01T22:06:57.000-08:00
- limit threads launched by mulh cuda kernel - use `inline constexpr` for compile time constants instead of `static const` in cuda kernels these are unrelated changes picked from #2338
diff --git a/crates/circuits/primitives/cuda/include/primitives/constants.h b/crates/circuits/primitives/cuda/include/primitives/constants.h
@@ -3,91 +3,91 @@
 #include <cstddef>
 
 namespace riscv {
-static const size_t RV32_REGISTER_NUM_LIMBS = 4;
-static const size_t RV32_CELL_BITS = 8;
-static const size_t RV_J_TYPE_IMM_BITS = 21;
+inline constexpr size_t RV32_REGISTER_NUM_LIMBS = 4;
+inline constexpr size_t RV32_CELL_BITS = 8;
+inline constexpr size_t RV_J_TYPE_IMM_BITS = 21;
 
-static const size_t RV32_IMM_AS = 0;
+inline constexpr size_t RV32_IMM_AS = 0;
 } // namespace riscv
 
 namespace program {
-static const size_t PC_BITS = 30;
-static const size_t DEFAULT_PC_STEP = 4;
+inline constexpr size_t PC_BITS = 30;
+inline constexpr size_t DEFAULT_PC_STEP = 4;
 } // namespace program
 
 namespace native {
-static const size_t AS_IMMEDIATE = 0;
-static const size_t AS_NATIVE = 4;
-static const size_t EXT_DEG = 4;
-static const size_t BETA = 11;
+inline constexpr size_t AS_IMMEDIATE = 0;
+inline constexpr size_t AS_NATIVE = 4;
+inline constexpr size_t EXT_DEG = 4;
+inline constexpr size_t BETA = 11;
 } // namespace native
 
 namespace poseidon2 {
-static const size_t CHUNK = 8;
+inline constexpr size_t CHUNK = 8;
 } // namespace poseidon2
 
 namespace p3_keccak_air {
-static const size_t NUM_ROUNDS = 24;
-static const size_t BITS_PER_LIMB = 16;
-static const size_t U64_LIMBS = 64 / BITS_PER_LIMB;
-static const size_t RATE_BITS = 1088;
-static const size_t RATE_LIMBS = RATE_BITS / BITS_PER_LIMB;
+inline constexpr size_t NUM_ROUNDS = 24;
+inline constexpr size_t BITS_PER_LIMB = 16;
+inline constexpr size_t U64_LIMBS = 64 / BITS_PER_LIMB;
+inline constexpr size_t RATE_BITS = 1088;
+inline constexpr size_t RATE_LIMBS = RATE_BITS / BITS_PER_LIMB;
 } // namespace p3_keccak_air
 
 namespace keccak256 {
 /// Total number of sponge bytes: number of rate bytes + number of capacity bytes.
-static const size_t KECCAK_WIDTH_BYTES = 200;
+inline constexpr size_t KECCAK_WIDTH_BYTES = 200;
 /// Total number of 16-bit limbs in the sponge.
-static const size_t KECCAK_WIDTH_U16S = KECCAK_WIDTH_BYTES / 2;
+inline constexpr size_t KECCAK_WIDTH_U16S = KECCAK_WIDTH_BYTES / 2;
 /// Number of rate bytes.
-static const size_t KECCAK_RATE_BYTES = 136;
+inline constexpr size_t KECCAK_RATE_BYTES = 136;
 /// Number of 16-bit rate limbs.
-static const size_t KECCAK_RATE_U16S = KECCAK_RATE_BYTES / 2;
+inline constexpr size_t KECCAK_RATE_U16S = KECCAK_RATE_BYTES / 2;
 /// Number of absorb rounds, equal to rate in u64s.
-static const size_t NUM_ABSORB_ROUNDS = KECCAK_RATE_BYTES / 8;
+inline constexpr size_t NUM_ABSORB_ROUNDS = KECCAK_RATE_BYTES / 8;
 /// Number of capacity bytes.
-static const size_t KECCAK_CAPACITY_BYTES = 64;
+inline constexpr size_t KECCAK_CAPACITY_BYTES = 64;
 /// Number of 16-bit capacity limbs.
-static const size_t KECCAK_CAPACITY_U16S = KECCAK_CAPACITY_BYTES / 2;
+inline constexpr size_t KECCAK_CAPACITY_U16S = KECCAK_CAPACITY_BYTES / 2;
 /// Number of output digest bytes used during the squeezing phase.
-static const size_t KECCAK_DIGEST_BYTES = 32;
+inline constexpr size_t KECCAK_DIGEST_BYTES = 32;
 /// Number of 64-bit digest limbs.
-static const size_t KECCAK_DIGEST_U64S = KECCAK_DIGEST_BYTES / 8;
+inline constexpr size_t KECCAK_DIGEST_U64S = KECCAK_DIGEST_BYTES / 8;
 
 // ==== Constants for register/memory adapter ====
 /// Register reads to get dst, src, len
-static const size_t KECCAK_REGISTER_READS = 3;
+inline constexpr size_t KECCAK_REGISTER_READS = 3;
 /// Number of cells to read/write in a single memory access
-static const size_t KECCAK_WORD_SIZE = 4;
+inline constexpr size_t KECCAK_WORD_SIZE = 4;
 /// Memory reads for absorb per row
-static const size_t KECCAK_ABSORB_READS = KECCAK_RATE_BYTES / KECCAK_WORD_SIZE;
+inline constexpr size_t KECCAK_ABSORB_READS = KECCAK_RATE_BYTES / KECCAK_WORD_SIZE;
 /// Memory writes for digest per row
-static const size_t KECCAK_DIGEST_WRITES = KECCAK_DIGEST_BYTES / KECCAK_WORD_SIZE;
+inline constexpr size_t KECCAK_DIGEST_WRITES = KECCAK_DIGEST_BYTES / KECCAK_WORD_SIZE;
 /// keccakf parameters
-static const size_t KECCAK_ROUND = 24;
-static const size_t KECCAK_STATE_SIZE = 25;
-static const size_t KECCAK_Q_SIZE = 192;
+inline constexpr size_t KECCAK_ROUND = 24;
+inline constexpr size_t KECCAK_STATE_SIZE = 25;
+inline constexpr size_t KECCAK_Q_SIZE = 192;
 /// From memory config
-static const size_t KECCAK_POINTER_MAX_BITS = 29;
+inline constexpr size_t KECCAK_POINTER_MAX_BITS = 29;
 } // namespace keccak256
 
 namespace mod_builder {
-static const size_t MAX_LIMBS = 97;
+inline constexpr size_t MAX_LIMBS = 97;
 } // namespace mod_builder
 
 namespace sha256 {
-static const size_t SHA256_BLOCK_BITS = 512;
-static const size_t SHA256_BLOCK_U8S = 64;
-static const size_t SHA256_BLOCK_WORDS = 16;
-static const size_t SHA256_WORD_U8S = 4;
-static const size_t SHA256_WORD_BITS = 32;
-static const size_t SHA256_WORD_U16S = 2;
-static const size_t SHA256_HASH_WORDS = 8;
-static const size_t SHA256_NUM_READ_ROWS = 4;
-static const size_t SHA256_ROWS_PER_BLOCK = 17;
-static const size_t SHA256_ROUNDS_PER_ROW = 4;
-static const size_t SHA256_ROW_VAR_CNT = 5;
-static const size_t SHA256_REGISTER_READS = 3;
-static const size_t SHA256_READ_SIZE = 16;
-static const size_t SHA256_WRITE_SIZE = 32;
-} // namespace sha256
+inline constexpr size_t SHA256_BLOCK_BITS = 512;
+inline constexpr size_t SHA256_BLOCK_U8S = 64;
+inline constexpr size_t SHA256_BLOCK_WORDS = 16;
+inline constexpr size_t SHA256_WORD_U8S = 4;
+inline constexpr size_t SHA256_WORD_BITS = 32;
+inline constexpr size_t SHA256_WORD_U16S = 2;
+inline constexpr size_t SHA256_HASH_WORDS = 8;
+inline constexpr size_t SHA256_NUM_READ_ROWS = 4;
+inline constexpr size_t SHA256_ROWS_PER_BLOCK = 17;
+inline constexpr size_t SHA256_ROUNDS_PER_ROW = 4;
+inline constexpr size_t SHA256_ROW_VAR_CNT = 5;
+inline constexpr size_t SHA256_REGISTER_READS = 3;
+inline constexpr size_t SHA256_READ_SIZE = 16;
+inline constexpr size_t SHA256_WRITE_SIZE = 32;
+} // namespace sha256
diff --git a/crates/circuits/primitives/cuda/include/primitives/less_than.cuh b/crates/circuits/primitives/cuda/include/primitives/less_than.cuh
@@ -3,7 +3,7 @@
 #include "histogram.cuh"
 #include "fp_array.cuh"
 
-static const size_t AUX_LEN = 2;
+inline constexpr size_t AUX_LEN = 2;
 
 template <typename T, size_t AUX_LEN = AUX_LEN> struct LessThanAuxCols {
     T lower_decomp[AUX_LEN];
@@ -19,14 +19,14 @@ template <typename T, size_t NUM, size_t AUX_LEN = AUX_LEN> struct LessThanArray
 namespace AssertLessThan {
 /**
  * @brief Generates columns needed to constrain that x < y
- * 
+ *
  * @section Trace Context Parameters
  * @param rc Range checker histogram reference
  * @param max_bits Maximum number of bits the respresntation of x and y can be
  * @param x First value to compare (must be strictly less than y)
  * @param y Second value to compare
  * @param lower_decomp_len Number of columns needed to constrain x < y
- * 
+ *
  * @section Mutable Column Parameters
  * @param lower_decomp Columns used to constrain x < y
  */
@@ -45,14 +45,14 @@ __device__ __forceinline__ void generate_subrow(
 namespace IsLessThan {
 /**
  * @brief Generates columns needed to constrain that out_flag == (x < y)
- * 
+ *
  * @section Trace Context Parameters
  * @param rc Range checker histogram reference
  * @param max_bits Maximum number of bits the respresntation of x and y can be
  * @param x First value to compare
  * @param y Second value to compare
  * @param lower_decomp_len Number of columns needed to constrain out_flag == (x < y)
- * 
+ *
  * @section Mutable Column Parameters
  * @param lower_decomp Columns used to constrain out_flag == (x < y)
  * @param out_flag Boolean value equal to x < y
@@ -77,15 +77,15 @@ namespace IsLessThanArray {
 /**
  * @brief Generates columns needed to constrain that out_flag == (x < y),
  *        where x and y are represented by array_len limbs.
- * 
+ *
  * @section Trace Context Parameters
  * @param rc Range checker histogram reference
  * @param max_bits Maximum number of bits each limb of x and y can be
  * @param x First value to compare
  * @param y Second value to compare
  * @param array_len Number of limbs to represent x and y
  * @param aux_len Number of additional columns needed to constrain outflag == (x < y)
- * 
+ *
  * @section Mutable Column Parameters
  * @param diff_marker Array that marks the most significant limb difference in x and y
  * @param diff_inv Field inverse of the first differing y[i] - x[i], or 0
diff --git a/crates/vm/cuda/src/system/boundary.cu b/crates/vm/cuda/src/system/boundary.cu
@@ -5,8 +5,8 @@
 #include "primitives/trace_access.h"
 #include <cassert>
 
-static const size_t PERSISTENT_CHUNK = 8;
-static const size_t VOLATILE_CHUNK = 1;
+inline constexpr size_t PERSISTENT_CHUNK = 8;
+inline constexpr size_t VOLATILE_CHUNK = 1;
 
 template <size_t CHUNK> struct BoundaryRecord {
     uint32_t address_space;
@@ -24,8 +24,8 @@ template <typename T> struct PersistentBoundaryCols {
     T timestamp;
 };
 
-static const size_t ADDR_ELTS = 2;
-static const size_t NUM_AS_LIMBS = 1;
+inline constexpr size_t ADDR_ELTS = 2;
+inline constexpr size_t NUM_AS_LIMBS = 1;
 
 template <typename T> struct VolatileBoundaryCols {
     T address_space_limbs[NUM_AS_LIMBS];
@@ -61,7 +61,7 @@ __global__ void cukernel_persistent_boundary_tracegen(
             // TODO better address space handling
             FpArray<8> init_values;
             if (initial_mem[record.address_space - 1]) {
-                init_values = 
+                init_values =
                     record.address_space == 4
                         ? FpArray<8>::from_raw_array(
                             reinterpret_cast<uint32_t const *>(
diff --git a/extensions/native/circuit/cuda/include/native/fri.cuh b/extensions/native/circuit/cuda/include/native/fri.cuh
@@ -103,7 +103,7 @@ template <typename T> struct Instruction2Cols {
 };
 
 // Size constants for the three column types
-static const size_t WORKLOAD_SIZE = sizeof(WorkloadCols<uint8_t>);
-static const size_t INSN1_SIZE = sizeof(Instruction1Cols<uint8_t>);
-static const size_t INSN2_SIZE = sizeof(Instruction2Cols<uint8_t>);
-static const size_t OVERALL_SIZE = std::max({WORKLOAD_SIZE, INSN1_SIZE, INSN2_SIZE});
+inline constexpr size_t WORKLOAD_SIZE = sizeof(WorkloadCols<uint8_t>);
+inline constexpr size_t INSN1_SIZE = sizeof(Instruction1Cols<uint8_t>);
+inline constexpr size_t INSN2_SIZE = sizeof(Instruction2Cols<uint8_t>);
+inline constexpr size_t OVERALL_SIZE = std::max({WORKLOAD_SIZE, INSN1_SIZE, INSN2_SIZE});
diff --git a/extensions/native/circuit/cuda/src/poseidon2.cu b/extensions/native/circuit/cuda/src/poseidon2.cu
@@ -7,13 +7,13 @@
 
 using namespace poseidon2;
 
-static const size_t WIDTH = 16;
-static const size_t SBOX_DEGREE = Poseidon2DefaultParams::SBOX_DEGREE;
-static const size_t HALF_FULL_ROUNDS = Poseidon2DefaultParams::HALF_FULL_ROUNDS;
-static const size_t PARTIAL_ROUNDS = Poseidon2DefaultParams::PARTIAL_ROUNDS;
+inline constexpr size_t WIDTH = 16;
+inline constexpr size_t SBOX_DEGREE = Poseidon2DefaultParams::SBOX_DEGREE;
+inline constexpr size_t HALF_FULL_ROUNDS = Poseidon2DefaultParams::HALF_FULL_ROUNDS;
+inline constexpr size_t PARTIAL_ROUNDS = Poseidon2DefaultParams::PARTIAL_ROUNDS;
 
-static const uint32_t NUM_INITIAL_READS = 6;
-// static const uint32_t NUM_SIMPLE_ACCESSES = 7;
+inline constexpr uint32_t NUM_INITIAL_READS = 6;
+// inline constexpr uint32_t NUM_SIMPLE_ACCESSES = 7;
 
 template <typename T, size_t SBOX_REGISTERS> struct NativePoseidon2Cols {
     Poseidon2SubCols<T, WIDTH, SBOX_DEGREE, SBOX_REGISTERS, HALF_FULL_ROUNDS, PARTIAL_ROUNDS> inner;
diff --git a/extensions/rv32im/circuit/cuda/src/mulh.cu b/extensions/rv32im/circuit/cuda/src/mulh.cu
@@ -203,7 +203,7 @@ extern "C" int _mulh_tracegen(
     assert(height >= d_records.len());
     assert(width == sizeof(MulHCols<uint8_t>));
 
-    auto [grid, block] = kernel_launch_params(height);
+    auto [grid, block] = kernel_launch_params(height, 512);
 
     mulh_tracegen<<<grid, block>>>(
         d_trace,
