chore(cuda): BabyBear unified + CUDA_DEBUG + opener in natural order (#2149)

Co-authored-by: Jonathan Wang <[email protected]>

Author: gaxiom

Cargo.lock

@@ -1,5 +1,5 @@
 [workspace.package]
-version = "1.4.1-rc.0"
+version = "1.4.1-rc.1"
 edition = "2021"
 rust-version = "1.86.0"
 authors = ["OpenVM Authors"]
@@ -113,11 +113,11 @@ lto = "thin"
 
 [workspace.dependencies]
 # Stark Backend
-openvm-stark-backend = { git = "https://github.com/openvm-org/stark-backend.git", tag = "v1.2.1-rc.2", default-features = false }
-openvm-stark-sdk = { git = "https://github.com/openvm-org/stark-backend.git", tag = "v1.2.1-rc.2", default-features = false }
-openvm-cuda-backend = { git = "https://github.com/openvm-org/stark-backend.git", tag = "v1.2.1-rc.2", default-features = false }
-openvm-cuda-builder = { git = "https://github.com/openvm-org/stark-backend.git", tag = "v1.2.1-rc.2", default-features = false }
-openvm-cuda-common = { git = "https://github.com/openvm-org/stark-backend.git", tag = "v1.2.1-rc.2", default-features = false }
+openvm-stark-backend = { git = "https://github.com/openvm-org/stark-backend.git", tag = "v1.2.1-rc.3", default-features = false }
+openvm-stark-sdk = { git = "https://github.com/openvm-org/stark-backend.git", tag = "v1.2.1-rc.3", default-features = false }
+openvm-cuda-backend = { git = "https://github.com/openvm-org/stark-backend.git", tag = "v1.2.1-rc.3", default-features = false }
+openvm-cuda-builder = { git = "https://github.com/openvm-org/stark-backend.git", tag = "v1.2.1-rc.3", default-features = false }
+openvm-cuda-common = { git = "https://github.com/openvm-org/stark-backend.git", tag = "v1.2.1-rc.3", default-features = false }
 
 # OpenVM
 openvm-sdk = { path = "crates/sdk", default-features = false }

Cargo.toml

@@ -371,5 +371,5 @@ extern "C" int _field_expression_tracegen(
         workspace_per_thread
     );
 
-    return cudaGetLastError();
+    return CHECK_KERNEL();
 }

crates/circuits/mod-builder/cuda/src/field_expression.cu

@@ -40,5 +40,5 @@ extern "C" int _poseidon2_dummy_tracegen(Fp *output, Fp *inputs, uint32_t sbox_r
     default:
         return cudaErrorInvalidConfiguration;
     }
-    return cudaGetLastError();
+    return CHECK_KERNEL();
 }

crates/circuits/poseidon2-air/cuda/src/dummy.cu

@@ -48,7 +48,7 @@ struct Encoder {
     __device__ uint32_t width() const { return k; }
 
     __device__ void write_flag_pt(RowSlice pt, uint32_t idx) const {
-#ifdef DEBUG
+#ifdef CUDA_DEBUG
         assert(idx < num_flags);
         assert(this->k > 0);
 #endif

crates/circuits/primitives/cuda/include/primitives/encoder.cuh

@@ -1,6 +1,7 @@
 #pragma once
 
 #include "launcher.cuh"
+#include "utils.cuh"
 #include "trace_access.h"
 
 /**
@@ -64,8 +65,8 @@ struct VariableRangeChecker {
         const size_t limbs_len
     ) {
         size_t range_max_bits = max_bits();
-#ifdef DEBUG
-        assert(limbs_len >= div_ceil(bits, range_max_bits));
+#ifdef CUDA_DEBUG
+        assert(limbs_len >= d_div_ceil(bits, range_max_bits));
 #endif
         uint32_t mask = (1 << range_max_bits) - 1;
         size_t bits_remaining = bits;
@@ -77,7 +78,7 @@ struct VariableRangeChecker {
             x >>= range_max_bits;
             bits_remaining -= min(bits_remaining, range_max_bits);
         }
-#ifdef DEBUG
+#ifdef CUDA_DEBUG
         assert(bits_remaining == 0 && x == 0);
 #endif
     }
@@ -105,9 +106,6 @@ template <uint32_t N> struct RangeTupleChecker {
     }
 
     __device__ void add_count(RowSlice values) {
-#ifdef DEBUG
-        assert(values.length == N);
-#endif
         uint32_t idx = 0;
         for (int i = 0; i < N; i++) {
             idx = idx * sizes[i] + values[i].asUInt32();

crates/circuits/primitives/cuda/include/primitives/histogram.cuh

@@ -24,5 +24,5 @@ extern "C" int _bitwise_op_lookup_tracegen(
     uint32_t num_rows = 1 << (2 * num_bits);
     auto [grid, block] = kernel_launch_params(num_rows);
     bitwise_op_lookup_tracegen<<<grid, block>>>(d_count, d_cpu_count, d_trace, num_rows);
-    return cudaGetLastError();
+    return CHECK_KERNEL();
 }

crates/circuits/primitives/cuda/src/bitwise_op_lookup.cu

@@ -53,5 +53,5 @@ extern "C" int _bitwise_dummy_tracegen(
 ) {
     auto [grid, block] = kernel_launch_params(records.len());
     bitwise_dummy_tracegen<<<grid, block>>>(d_trace, records, bitwise_count, bitwise_num_bits);
-    return cudaGetLastError();
+    return CHECK_KERNEL();
 }

crates/circuits/primitives/cuda/src/dummy/bitwise_op_lookup.cu

@@ -34,5 +34,5 @@ extern "C" int _encoder_tracegen(
 
     cukernel_encoder_tracegen<<<grid, block>>>(trace, num_flags, max_degree, reserve_invalid, k);
 
-    return cudaGetLastError();
+    return CHECK_KERNEL();
 }

crates/circuits/primitives/cuda/src/dummy/encoder.cu

@@ -1,4 +1,5 @@
 #include "fp.h"
+#include "launcher.cuh"
 
 __global__ void cukernel_fibair_tracegen(Fp *output, uint32_t a, uint32_t b, uint32_t n) {
     if (blockIdx.x != 0 || threadIdx.x != 0)
@@ -19,5 +20,5 @@ extern "C" int _fibair_tracegen(Fp *output, uint32_t a, uint32_t b, uint32_t n)
     dim3 grid(1);
     dim3 block(1);
     cukernel_fibair_tracegen<<<grid, block>>>(output, a, b, n);
-    return cudaGetLastError();
+    return CHECK_KERNEL();
 }