chore: migrate CUDA primitive chips + test (#1974)

Author: stephenh-axiom-xyz

Cargo.lock

@@ -114,6 +114,8 @@ lto = "thin"
 # Stark Backend
 openvm-stark-backend = { git = "https://github.com/openvm-org/stark-backend.git", tag = "v1.2.0-rc.1", default-features = false }
 openvm-stark-sdk = { git = "https://github.com/openvm-org/stark-backend.git", tag = "v1.2.0-rc.1", default-features = false }
+# TODO[stephenh]: Replace this once open sourced
+stark-backend-gpu = { git = "https://github.com/axiom-crypto/axiom-gpu.git", default-features = false }
 
 # OpenVM
 openvm-sdk = { path = "crates/sdk", default-features = false }

Cargo.toml

@@ -17,6 +17,7 @@ zkhash = { workspace = true }
 
 openvm-stark-backend = { workspace = true }
 openvm-stark-sdk = { workspace = true }
+stark-backend-gpu = { workspace = true, optional = true }
 
 rand.workspace = true
 lazy_static.workspace = true
@@ -26,5 +27,6 @@ derivative.workspace = true
 p3-symmetric = { workspace = true }
 
 [features]
-default = ["parallel"]
+default = ["parallel", "cuda"]
 parallel = ["openvm-stark-backend/parallel"]
+cuda = ["dep:stark-backend-gpu"]

crates/circuits/poseidon2-air/Cargo.toml

@@ -28,7 +28,7 @@ __global__ void cukernel_poseidon2_tracegen(Fp *output, Fp *inputs, uint32_t n)
         PARTIAL_ROUNDS>(row, state);
 }
 
-extern "C" int _poseidon2_tracegen(Fp *output, Fp *inputs, uint32_t sbox_regs, uint32_t n) {
+extern "C" int _poseidon2_dummy_tracegen(Fp *output, Fp *inputs, uint32_t sbox_regs, uint32_t n) {
 
     auto [grid, block] = kernel_launch_params(n);
     switch (sbox_regs) {

tracegen-gpu/cuda/src/dummy/poseidon2.cu …circuits/poseidon2-air/cuda/src/dummy.cutracegen-gpu/cuda/src/dummy/poseidon2.cu renamed to crates/circuits/poseidon2-air/cuda/src/dummy.cu tracegen-gpu/cuda/src/dummy/poseidon2.cu renamed to crates/circuits/poseidon2-air/cuda/src/dummy.cu

@@ -0,0 +1,45 @@
+#![allow(clippy::missing_safety_doc)]
+
+use stark_backend_gpu::{
+    cuda::{d_buffer::DeviceBuffer, error::CudaError},
+    prelude::F,
+};
+
+pub mod poseidon2 {
+    /// Poseidon2 tracegen on GPU (parallelized over rows)
+    ///
+    /// # Arguments
+    ///
+    /// * `d_output` - DeviceBuffer for the output (column major)
+    /// * `d_inputs` - DeviceBuffer for the inputs (column major)
+    /// * `sbox_regs` - Number of sbox registers (0 or 1)
+    /// * `n` - Number of rows
+    ///
+    /// Currently only supports same constants as  
+    /// https://github.com/openvm-org/openvm/blob/08bbf79368b07437271aeacb25fb8857980ca863/crates/circuits/poseidon2-air/src/lib.rs
+    /// so:
+    /// * `WIDTH` - 16
+    /// * `SBOX_DEGREE` - 7
+    /// * `HALF_FULL_ROUNDS` - 4
+    /// * `PARTIAL_ROUNDS` - 13
+    use super::*;
+
+    extern "C" {
+        fn _poseidon2_dummy_tracegen(output: *mut F, inputs: *mut F, sbox_regs: u32, n: u32)
+            -> i32;
+    }
+
+    pub unsafe fn dummy_tracegen(
+        d_output: &DeviceBuffer<F>,
+        d_inputs: &DeviceBuffer<F>,
+        sbox_regs: u32,
+        n: u32,
+    ) -> Result<(), CudaError> {
+        CudaError::from_result(_poseidon2_dummy_tracegen(
+            d_output.as_mut_ptr(),
+            d_inputs.as_mut_ptr(),
+            sbox_regs,
+            n,
+        ))
+    }
+}

crates/circuits/poseidon2-air/src/cuda_abi.rs

@@ -28,6 +28,9 @@ pub use babybear::*;
 pub use config::*;
 pub use permute::*;
 
+#[cfg(all(feature = "cuda", test))]
+mod cuda_abi;
+
 #[cfg(test)]
 mod tests;
 

crates/circuits/poseidon2-air/src/lib.rs

@@ -19,6 +19,13 @@ use rand::{rngs::StdRng, Rng, RngCore};
 use super::{Poseidon2Config, Poseidon2Constants, Poseidon2SubChip};
 use crate::BABY_BEAR_POSEIDON2_HALF_FULL_ROUNDS;
 
+#[cfg(feature = "cuda")]
+use {
+    crate::cuda_abi::poseidon2,
+    openvm_stark_backend::p3_field::PrimeField32,
+    stark_backend_gpu::{base::DeviceMatrix, cuda::copy::MemCopyH2D as _, types::F},
+};
+
 fn run_poseidon2_subchip_test(subchip: Arc<Poseidon2SubChip<BabyBear, 0>>, rng: &mut StdRng) {
     // random state and trace generation
     let num_rows = 1 << 4;
@@ -93,3 +100,48 @@ fn test_poseidon2_random_constants() {
     let poseidon2_subchip = Arc::new(Poseidon2SubChip::<BabyBear, 0>::new(constants));
     run_poseidon2_subchip_test(poseidon2_subchip, &mut rng);
 }
+
+#[cfg(feature = "cuda")]
+#[test]
+fn test_cuda_tracegen_poseidon2() {
+    const WIDTH: usize = 16; // constant for BabyBear
+    const N: usize = 16;
+    const SBOX_REGS: usize = 1;
+    const HALF_FULL_ROUNDS: usize = 4; // Constant for BabyBear
+    const PARTIAL_ROUNDS: usize = 13; // Constant for BabyBear
+
+    // Generate random states and prepare GPU inputs
+    let mut rng = create_seeded_rng();
+    let cpu_inputs: Vec<[F; WIDTH]> = (0..N)
+        .map(|_| std::array::from_fn(|_| F::from_canonical_u32(rng.gen_range(0..F::ORDER_U32))))
+        .collect();
+
+    // Flatten inputs in row-major order for GPU (same layout as cpu_inputs)
+    let inputs_dev = cpu_inputs
+        .iter()
+        .flat_map(|r| r.iter().copied())
+        .collect::<Vec<_>>()
+        .to_device()
+        .unwrap();
+
+    // Launch GPU tracegen
+    let num_cols = 1
+        + WIDTH
+        + HALF_FULL_ROUNDS * (WIDTH * SBOX_REGS + WIDTH)
+        + PARTIAL_ROUNDS * (SBOX_REGS + 1)
+        + HALF_FULL_ROUNDS * (WIDTH * SBOX_REGS + WIDTH);
+
+    let gpu_mat = DeviceMatrix::<F>::with_capacity(N, num_cols);
+
+    unsafe {
+        poseidon2::dummy_tracegen(gpu_mat.buffer(), &inputs_dev, SBOX_REGS as u32, N as u32)
+            .expect("GPU tracegen failed");
+    }
+
+    // Run CPU tracegen and compare results
+    let config = Poseidon2Config::<BabyBear>::default();
+    let chip: Poseidon2SubChip<_, SBOX_REGS> = Poseidon2SubChip::new(config.constants);
+    let _cpu_trace = Arc::new(chip.generate_trace(cpu_inputs));
+    // TODO[stephenh]: Uncomment this when we decide where to put it
+    // assert_eq_cpu_and_gpu_matrix(cpu_trace, &gpu_mat);
+}

crates/circuits/poseidon2-air/src/tests.rs

@@ -11,6 +11,7 @@ license.workspace = true
 [dependencies]
 openvm-stark-backend = { workspace = true }
 openvm-circuit-primitives-derive = { workspace = true }
+stark-backend-gpu = { workspace = true, optional = true }
 
 rand.workspace = true
 derive-new.workspace = true
@@ -24,5 +25,6 @@ test-case.workspace = true
 openvm-stark-sdk = { workspace = true }
 
 [features]
-default = ["parallel"]
+default = ["parallel", "cuda"]
 parallel = ["openvm-stark-backend/parallel"]
+cuda = ["dep:stark-backend-gpu"]