feat: migrate mod builder

Cargo.toml

@@ -253,6 +253,9 @@ num-traits = { version = "0.2.19", default-features = false }
 ff = { version = "0.13.1", default-features = false }
 sha2 = { version = "0.10", default-features = false }
 
+# specific to CUDA and GPU
+cuda-runtime-sys = "0.3.0-alpha.1"
+
 # For local development. Add to your `.cargo/config.toml`
 # [patch."https://github.com/Plonky3/Plonky3.git"]
 # p3-field = { path = "../Plonky3/field" }

crates/circuits/mod-builder/Cargo.toml

@@ -21,9 +21,16 @@ rand.workspace = true
 num-bigint.workspace = true
 num-traits.workspace = true
 tracing.workspace = true
-
 itertools.workspace = true
 
+# CUDA dependencies (only included when cuda feature is enabled)
+openvm-cuda-backend = { workspace = true, optional = true }
+openvm-cuda-common = { workspace = true, optional = true }
+cuda-runtime-sys = { workspace = true, optional = true }
+
+[build-dependencies]
+openvm-cuda-builder.workspace = true
+
 [dev-dependencies]
 openvm-circuit-primitives = { workspace = true }
 openvm-pairing-guest = { workspace = true, features = ["halo2curves"] }
@@ -32,6 +39,12 @@ openvm-circuit = { workspace = true, features = ["test-utils"] }
 
 [features]
 default = []
+cuda = [
+    "dep:openvm-cuda-backend",
+    "dep:openvm-cuda-common",
+    "dep:cuda-runtime-sys",
+    "openvm-circuit-primitives/cuda",
+]
 parallel = ["openvm-stark-backend/parallel"]
 test-utils = [
     "dep:halo2curves-axiom",

crates/circuits/mod-builder/build.rs

@@ -0,0 +1,31 @@
+#[cfg(feature = "cuda")]
+use openvm_cuda_builder::{cuda_available, CudaBuilder};
+
+#[cfg(feature = "cuda")]
+fn main() {
+    if !cuda_available() {
+        return; // Skip CUDA compilation
+    }
+
+    let builder = CudaBuilder::new()
+        .include_from_dep("DEP_CUDA_COMMON_INCLUDE") // Import headers
+        .include("../primitives/cuda/include")
+        .include("cuda/include")
+        .include("../../../extensions/rv32-adapters/cuda/include")
+        .include("../poseidon2-air/cuda/include")
+        .include("../../../tracegen-gpu/cuda/src") // TODO[arayi]: change once system is migrated
+        .watch("src/cuda_abi.rs")
+        .watch("cuda")
+        .watch("../primitives/cuda/include")
+        .watch("../../../extensions/rv32-adapters/cuda/include")
+        .watch("../poseidon2-air/cuda/include")
+        .watch("../../../tracegen-gpu/cuda/src") // TODO[arayi]: change once system is migrated
+        .library_name("tracegen_mod_builder")
+        .file("cuda/src/field_expression.cu");
+
+    builder.emit_link_directives();
+    builder.build();
+}
+
+#[cfg(not(feature = "cuda"))]
+fn main() {}

crates/circuits/mod-builder/cuda/include/records.cuh

@@ -1,6 +1,6 @@
 #pragma once
 
-#include "../../src/extensions/rv32_adapters/vec_heap.cuh"
+#include "vec_heap.cuh"
 
 struct FieldExprCoreRecord {
     uint8_t opcode;

crates/circuits/mod-builder/cuda/include/rv32_vec_heap_router.cuh

@@ -1,7 +1,7 @@
 #pragma once
 
-#include "../../src/extensions/rv32_adapters/vec_heap.cuh"
-#include "mod_builder/meta.cuh"
+#include "vec_heap.cuh"
+#include "meta.cuh"
 #include "trace_access.h"
 
 struct Rv32VecHeapConfig {

tracegen-gpu/cuda/src/mod_builder/field_expression.cu → crates/circuits/mod-builder/cuda/src/field_expression.cu

@@ -1,11 +1,11 @@
 #include "fp.h"
 #include "launcher.cuh"
-#include "mod_builder/bigint_ops.cuh"
-#include "mod_builder/expr_codec.cuh"
-#include "mod_builder/meta.cuh"
-#include "mod_builder/overflow_ops.cuh"
-#include "mod_builder/records.cuh"
-#include "mod_builder/rv32_vec_heap_router.cuh"
+#include "bigint_ops.cuh"
+#include "expr_codec.cuh"
+#include "meta.cuh"
+#include "overflow_ops.cuh"
+#include "records.cuh"
+#include "rv32_vec_heap_router.cuh"
 #include "trace_access.h"
 #include <cstdint>
 

tracegen-gpu/src/mod_builder/field_expression/chip.rs → crates/circuits/mod-builder/src/cuda/chip.rs

@@ -7,29 +7,26 @@ use cuda_runtime_sys::{cudaDeviceSetLimit, cudaLimit};
 use num_bigint::BigUint;
 use num_traits::{FromBytes, One};
 use openvm_circuit::utils::next_power_of_two_or_zero;
-use openvm_mod_circuit_builder::{
-    utils::biguint_to_limbs_vec, FieldExpressionCoreAir, SymbolicExpr,
+use openvm_circuit_primitives::{
+    bitwise_op_lookup::cuda::BitwiseOperationLookupChipGPU,
+    var_range::cuda::VariableRangeCheckerChipGPU,
 };
-use openvm_stark_backend::p3_air::BaseAir;
-use stark_backend_gpu::{
-    base::DeviceMatrix,
-    cuda::{
-        copy::{MemCopyD2H, MemCopyH2D},
-        d_buffer::DeviceBuffer,
-    },
-    prelude::F,
+use openvm_cuda_backend::{base::DeviceMatrix, types::F};
+use openvm_cuda_common::{
+    copy::{MemCopyD2H, MemCopyH2D},
+    d_buffer::DeviceBuffer,
 };
+use openvm_stark_backend::p3_air::BaseAir;
 
-use super::{
-    constants::*,
-    cuda::field_expression::tracegen,
-    types::{ExprMeta, ExprNode, FieldExprMeta, FieldExpressionChipGPU},
-};
 use crate::{
-    mod_builder::expr_op::ExprOp,
-    primitives::{
-        bitwise_op_lookup::BitwiseOperationLookupChipGPU, var_range::VariableRangeCheckerChipGPU,
+    cuda::{
+        constants::{ExprType, LIMB_BITS, MAX_LIMBS},
+        expr_op::ExprOp,
     },
+    cuda_abi::field_expression::tracegen,
+    utils::biguint_to_limbs_vec,
+    ExprMeta, ExprNode, FieldExprMeta, FieldExpressionChipGPU, FieldExpressionCoreAir,
+    SymbolicExpr,
 };
 
 impl FieldExpressionChipGPU {
@@ -336,7 +333,6 @@ impl FieldExpressionChipGPU {
     }
 
     fn calculate_ast_depth(expr: &SymbolicExpr) -> u32 {
-        use openvm_mod_circuit_builder::SymbolicExpr;
         match expr {
             SymbolicExpr::Input(_) | SymbolicExpr::Var(_) | SymbolicExpr::Const(_, _, _) => 1,
             SymbolicExpr::Add(left, right)

tracegen-gpu/src/mod_builder/expr_op.rs → crates/circuits/mod-builder/src/cuda/expr_op.rs

@@ -1,4 +1,4 @@
-use crate::mod_builder::field_expression::ExprNode;
+use crate::ExprNode;
 
 #[repr(C)]
 #[derive(Debug, Clone, Copy)]

crates/circuits/mod-builder/src/cuda/mod.rs

@@ -0,0 +1,6 @@
+mod chip;
+mod constants;
+mod expr_op;
+mod types;
+
+pub use types::*;

tracegen-gpu/src/mod_builder/field_expression/types.rs → crates/circuits/mod-builder/src/cuda/types.rs

@@ -1,12 +1,14 @@
 use std::sync::Arc;
 
-use stark_backend_gpu::cuda::d_buffer::DeviceBuffer;
+use openvm_circuit_primitives::{
+    bitwise_op_lookup::cuda::BitwiseOperationLookupChipGPU,
+    var_range::cuda::VariableRangeCheckerChipGPU,
+};
+use openvm_cuda_common::d_buffer::DeviceBuffer;
 
 use crate::{
-    mod_builder::{expr_op::ExprOp, field_expression::constants::LIMB_BITS},
-    primitives::{
-        bitwise_op_lookup::BitwiseOperationLookupChipGPU, var_range::VariableRangeCheckerChipGPU,
-    },
+    cuda::{constants::LIMB_BITS, expr_op::ExprOp},
+    FieldExpressionCoreAir,
 };
 
 #[repr(C)]
@@ -75,7 +77,7 @@ unsafe impl Send for FieldExprMeta {}
 unsafe impl Sync for FieldExprMeta {}
 
 pub struct FieldExpressionChipGPU {
-    pub air: openvm_mod_circuit_builder::FieldExpressionCoreAir,
+    pub air: FieldExpressionCoreAir,
     pub records: Arc<DeviceBuffer<u8>>,
     pub num_records: usize,
     pub record_stride: usize,

tracegen-gpu/src/mod_builder/field_expression/cuda.rs → crates/circuits/mod-builder/src/cuda_abi.rs

@@ -1,9 +1,10 @@
 #![allow(clippy::missing_safety_doc)]
 #![allow(clippy::too_many_arguments)]
 
-use stark_backend_gpu::cuda::{d_buffer::DeviceBuffer, error::CudaError};
+use openvm_cuda_common::{d_buffer::DeviceBuffer, error::CudaError};
 
 pub mod field_expression {
+
     use super::*;
     extern "C" {
         fn _field_expression_tracegen(

crates/circuits/mod-builder/src/lib.rs

@@ -14,3 +14,11 @@ pub mod utils;
 
 #[cfg(any(test, feature = "test-utils"))]
 pub mod test_utils;
+
+#[cfg(feature = "cuda")]
+mod cuda;
+#[cfg(feature = "cuda")]
+mod cuda_abi;
+
+#[cfg(feature = "cuda")]
+pub use cuda::*;