feat: integrate GPU prover (#987)

### Description

### Performence
- fibonacci $\text{fib}(2^{20})$: runs `12M` steps (warmup + rerun)
- machine(node-19): `AMD EPYC 7702 64-Core`, `1 x RTX 3090`
```
# CPU
RUST_LOG=debug cargo run --release --package ceno_zkvm --bin e2e -- --platform=ceno --hints=20 --public-io=3130 --profiling=2   examples/target/riscv32im-ceno-zkvm-elf/release/examples/fibonacci 2>&1 | tee "fib_perf_cpu.log"
```
```
# GPU
make enable-gpu
RUST_LOG=debug cargo run --release --features gpu --package ceno_zkvm --bin e2e -- --platform=ceno --hints=20 --public-io=3130 --profiling=2 --field=goldilocks  examples/target/riscv32im-ceno-zkvm-elf/release/examples/keccak_syscall 2>&1 | tee "fib_perf_gpu.log"
make disable-gpu
```
- **phase1: basefold commit & open**

| metric | CPU time | GPU time | speedup |
|------|----------|----------|--------|
| ZKVM_create_proof | 13.5s | 8.0s | 1.68x |
| batch commit to traces | 6.07s | 1.15s | 5.27x |
| transfer MLEs to host | 0 | 550ms | N/A |
| transfer pk to device | 0 | 4.53ms | N/A |
| main_proofs (no GPU yet) | 6.04s | 6.15s | 0.98x |
| pcs_opening | 1.42s | 168ms | 8.45x |

- **(WIP) phase2: main & tower prove**

| metric | CPU time | GPU time | speedup |
|------|----------|----------|--------|
| ZKVM_create_proof | 13.6s | 6.72s | 2.02x |
| batch commit to traces | 6.07s | 1.15s | 5.27x |
| transfer MLEs to host | 0 | 550ms (next: eliminate) | N/A |
| transfer pk to device | 0 | 4.53ms | N/A |
| main & tower prove | 6.04s | 4.85s | N/A |
| -- main (todo: update)  | - | - | - |
| -- tower  | 4.12s | 2.72s (next: optimize) | 1.51x |
| pcs_opening | 1.42s | 168ms | 8.45x |

### Current Status
- [x] Add GPU Prover Hal
- [x] Manually checked that `batch_commit` and `batch_open` produce
consistent results between CPU and GPU provers
    - `cargo test --release test_single_add_instance_e2e -- --nocapture`
- [x] API & IO, to pass ceno verifier
- [x] full GPU pipeline for basefold commit -> open
- [x] tower: witness, prove
- [x] main: layer witness, prove

### Future improvements
- code quality
- performance issues
- optimize GPU memory usage

---------

Co-authored-by: sm.wu <[email protected]>

Author: Velaciela

Cargo.lock

@@ -64,6 +64,8 @@ tracing-forest = { version = "0.1.6" }
 tracing-subscriber = { version = "0.3", features = ["env-filter"] }
 uint = "0.8"
 
+ceno_gpu = { path = "utils/cuda_hal", package = "cuda_hal" }
+
 [profile.dev]
 lto = "thin"
 # We are running our tests with optimizations turned on to make them faster.
@@ -76,3 +78,17 @@ opt-level = 3
 
 [profile.release]
 lto = "thin"
+
+#[patch."ssh://[email protected]/scroll-tech/ceno-gpu.git"]
+#ceno_gpu = { path = "../ceno-gpu/cuda_hal", package = "cuda_hal" }
+
+#[patch."https://github.com/scroll-tech/gkr-backend"]
+#ff_ext = { path = "../gkr-backend/crates/ff_ext", package = "ff_ext" }
+#mpcs = { path = "../gkr-backend/crates/mpcs", package = "mpcs" }
+#multilinear_extensions = { path = "../gkr-backend/crates/multilinear_extensions", package = "multilinear_extensions" }
+#p3 = { path = "../gkr-backend/crates/p3", package = "p3" }
+#poseidon = { path = "../gkr-backend/crates/poseidon", package = "poseidon" }
+#sumcheck = { path = "../gkr-backend/crates/sumcheck", package = "sumcheck" }
+#transcript = { path = "../gkr-backend/crates/transcript", package = "transcript" }
+#whir = { path = "../gkr-backend/crates/whir", package = "whir" }
+#witness = { path = "../gkr-backend/crates/witness", package = "witness" }

Cargo.toml

@@ -0,0 +1,24 @@
+# Makefile for conditional GPU builds
+
+.PHONY: enable-gpu disable-gpu clean help
+
+help:
+	@echo "Available targets:"
+	@echo "  enable-gpu   - Switch to GPU mode (uses remote implementation, requires private repo access)"
+	@echo "  disable-gpu  - Switch to CPU mode (uses local placeholder, default state)"
+	@echo "  clean        - Clean build artifacts and reset to CPU mode"
+	@echo ""
+	@echo "Normal usage:"
+	@echo "  cargo build                    # CPU build (default, no private repo fetch)"
+	@echo "  make enable-gpu && cargo build # GPU build (requires private repo access)"
+
+enable-gpu:
+	@./build-scripts/conditional-patch.sh enable-gpu
+
+disable-gpu:
+	@./build-scripts/conditional-patch.sh disable-gpu
+
+clean:
+	@cargo clean
+	@./build-scripts/conditional-patch.sh disable-gpu
+	@echo "Cleaned build artifacts and reset to CPU mode"

Makefile

@@ -0,0 +1,49 @@
+#!/bin/bash
+
+# Script to switch between local placeholder and remote GPU implementation
+# Usage: ./build-scripts/conditional-patch.sh [enable-gpu|disable-gpu]
+
+WORKSPACE_CARGO="Cargo.toml"
+
+# Workspace dependency declarations
+LOCAL_DEP='ceno_gpu = { path = "utils/cuda_hal", package = "cuda_hal" }'
+REMOTE_DEP='ceno_gpu = { git = "ssh://[email protected]/scroll-tech/ceno-gpu.git", package = "cuda_hal", branch = "dev/integrate-into-ceno-as-dep" }'
+
+if [ "$1" = "enable-gpu" ]; then
+    echo "Switching to GPU mode (using remote implementation)..."
+    
+    # Replace local path with remote git in workspace dependencies
+    if [[ "$OSTYPE" == "darwin"* ]]; then
+        # macOS sed
+        sed -i '' "s|${LOCAL_DEP}|${REMOTE_DEP}|g" "$WORKSPACE_CARGO"
+    else
+        # Linux sed
+        sed -i "s|${LOCAL_DEP}|${REMOTE_DEP}|g" "$WORKSPACE_CARGO"
+    fi
+    
+    echo "✅ Switched to remote GPU implementation"
+    echo "Now you can run: cargo build -p ceno_zkvm -F gpu"
+    
+elif [ "$1" = "disable-gpu" ]; then
+    echo "Switching to CPU mode (using local placeholder)..."
+    
+    # Replace remote git with local path in workspace dependencies
+    if [[ "$OSTYPE" == "darwin"* ]]; then
+        # macOS sed
+        sed -i '' "s|${REMOTE_DEP}|${LOCAL_DEP}|g" "$WORKSPACE_CARGO"
+    else
+        # Linux sed
+        sed -i "s|${REMOTE_DEP}|${LOCAL_DEP}|g" "$WORKSPACE_CARGO"
+    fi
+    
+    echo "✅ Switched to local placeholder implementation"
+    echo "Now you can run: cargo build -p ceno_zkvm --no-default-features"
+    
+else
+    echo "Usage: $0 [enable-gpu|disable-gpu]"
+    echo "  enable-gpu   - Switch to remote GPU implementation (requires private repo access)"
+    echo "  disable-gpu  - Switch to local placeholder (default, no private repo access)"
+    exit 1
+fi
+
+echo "Done."

build-scripts/conditional-patch.sh

@@ -6,20 +6,20 @@ use ceno_host::{CenoStdin, memory_from_file};
 use ceno_zkvm::{
     e2e::*,
     scheme::{
-        constants::MAX_NUM_VARIABLES, mock_prover::LkMultiplicityKey, verifier::ZKVMVerifier,
+        constants::MAX_NUM_VARIABLES, create_backend, create_prover,
+        mock_prover::LkMultiplicityKey, verifier::ZKVMVerifier,
     },
 };
 use clap::Args;
 use ff_ext::{BabyBearExt4, ExtensionField, GoldilocksExt2};
-use gkr_iop::cpu::{CpuBackend, CpuProver};
+
 use mpcs::{
     Basefold, BasefoldRSParams, PolynomialCommitmentScheme, SecurityLevel, Whir, WhirDefaultSpec,
 };
 use serde::Serialize;
 use std::{
     fs::File,
     path::{Path, PathBuf},
-    rc::Rc,
 };
 
 /// Ceno options
@@ -373,12 +373,9 @@ fn run_elf_inner<
         platform.hints.len()
     );
 
-    // TODO support GPU backend + prover
-    let backend: Rc<_> =
-        CpuBackend::<E, PCS>::new(options.max_num_variables, options.security_level).into();
-
+    let backend = create_backend(options.max_num_variables, options.security_level);
     Ok(run_e2e_with_checkpoint::<E, PCS, _, _>(
-        CpuProver::new(backend.clone()),
+        create_prover(backend.clone()),
         program,
         platform,
         &hints,

ceno_cli/src/commands/common_args/ceno.rs

@@ -13,12 +13,15 @@ version.workspace = true
 base64 = "0.22"
 ceno-examples = { path = "../examples-builder" }
 ceno_emul = { path = "../ceno_emul" }
+ceno_gpu = { workspace = true, optional = true }
 ceno_host = { path = "../ceno_host" }
+cudarc = { version = "0.13.0", features = ["driver", "cuda-version-from-build-system"], optional = true }
 either.workspace = true
 ff_ext.workspace = true
 gkr_iop = { path = "../gkr_iop" }
 mpcs.workspace = true
 multilinear_extensions.workspace = true
+once_cell = "1.21.3"
 p3.workspace = true
 rayon.workspace = true
 serde.workspace = true
@@ -64,6 +67,7 @@ default = ["forbid_overflow", "nightly-features", "u16limb_circuit", "parallel"]
 flamegraph = ["pprof2/flamegraph", "pprof2/criterion"]
 forbid_overflow = []
 goldilocks = ["forbid_overflow", "nightly-features"]
+gpu = ["gkr_iop/gpu", "dep:ceno_gpu"]
 jemalloc = ["dep:tikv-jemallocator", "dep:tikv-jemalloc-ctl"]
 jemalloc-prof = ["jemalloc", "tikv-jemallocator?/profiling"]
 nightly-features = [

ceno_zkvm/Cargo.toml

@@ -36,7 +36,9 @@ fn keccak_f_fn(c: &mut Criterion) {
 
                         #[allow(clippy::unit_arg)]
                         run_bitwise_keccakf::<GoldilocksExt2, BasefoldDefault<GoldilocksExt2>>(
-                            setup_bitwise_keccak_gkr_circuit().expect("setup circuit error"),
+                            setup_bitwise_keccak_gkr_circuit()
+                                .expect("setup circuit error")
+                                .1,
                             black_box(states),
                             false,
                             false,

ceno_zkvm/benches/bitwise_keccakf.rs

@@ -1,17 +1,19 @@
-use std::{rc::Rc, time::Duration};
+use std::time::Duration;
 
 use ceno_emul::{Platform, Program};
 use ceno_host::CenoStdin;
 use ceno_zkvm::{
     self,
     e2e::{Checkpoint, Preset, run_e2e_with_checkpoint, setup_platform},
-    scheme::verifier::ZKVMVerifier,
+    scheme::{create_backend, create_prover},
 };
 mod alloc;
 use criterion::*;
 
 use ff_ext::BabyBearExt4;
-use gkr_iop::cpu::{CpuBackend, CpuProver};
+use gkr_iop::cpu::default_backend_config;
+
+use ceno_zkvm::scheme::verifier::ZKVMVerifier;
 use mpcs::BasefoldDefault;
 use transcript::BasicTranscript;
 
@@ -40,14 +42,16 @@ fn setup() -> (Program, Platform) {
 
 fn fibonacci_prove(c: &mut Criterion) {
     let (program, platform) = setup();
-    let backend: Rc<_> = CpuBackend::<E, Pcs>::default().into();
+    let (max_num_variables, security_level) = default_backend_config();
+    let backend = create_backend::<E, Pcs>(max_num_variables, security_level);
+
     for max_steps in [1usize << 20, 1usize << 21, 1usize << 22] {
         // retrive 1 << 20th fibonacci element >> max_steps
         let mut hints = CenoStdin::default();
         let _ = hints.write(&20);
         // estimate proof size data first
         let result = run_e2e_with_checkpoint::<E, Pcs, _, _>(
-            CpuProver::new(backend.clone()),
+            create_prover(backend.clone()),
             program.clone(),
             platform.clone(),
             &Vec::from(&hints),
@@ -84,7 +88,7 @@ fn fibonacci_prove(c: &mut Criterion) {
                     let mut time = Duration::new(0, 0);
                     for _ in 0..iters {
                         let result = run_e2e_with_checkpoint::<E, Pcs, _, _>(
-                            CpuProver::new(backend.clone()),
+                            create_prover(backend.clone()),
                             program.clone(),
                             platform.clone(),
                             &Vec::from(&hints),

ceno_zkvm/benches/fibonacci.rs

@@ -3,12 +3,14 @@ use ceno_host::CenoStdin;
 use ceno_zkvm::{
     self,
     e2e::{Checkpoint, Preset, run_e2e_with_checkpoint, setup_platform},
+    scheme::{create_backend, create_prover},
 };
-use std::{fs, path::PathBuf, rc::Rc, time::Duration};
+use std::{fs, path::PathBuf, time::Duration};
 mod alloc;
 use criterion::*;
+
 use ff_ext::BabyBearExt4;
-use gkr_iop::cpu::{CpuBackend, CpuProver};
+use gkr_iop::cpu::default_backend_config;
 use mpcs::BasefoldDefault;
 
 criterion_group! {
@@ -38,7 +40,8 @@ fn setup() -> (Program, Platform) {
 
 fn fibonacci_witness(c: &mut Criterion) {
     let (program, platform) = setup();
-    let backend: Rc<_> = CpuBackend::<E, Pcs>::default().into();
+    let (max_num_variables, security_level) = default_backend_config();
+    let backend = create_backend::<E, Pcs>(max_num_variables, security_level);
 
     let max_steps = usize::MAX;
     let mut group = c.benchmark_group(format!("fib_wit_max_steps_{}", max_steps));
@@ -59,7 +62,7 @@ fn fibonacci_witness(c: &mut Criterion) {
                 let mut time = Duration::new(0, 0);
                 for _ in 0..iters {
                     let result = run_e2e_with_checkpoint::<E, Pcs, _, _>(
-                        CpuProver::new(backend.clone()),
+                        create_prover(backend.clone()),
                         program.clone(),
                         platform.clone(),
                         &Vec::from(&hints),

ceno_zkvm/benches/fibonacci_witness.rs

@@ -1,15 +1,16 @@
-use std::{rc::Rc, time::Duration};
+use std::time::Duration;
 
 use ceno_emul::{Platform, Program};
 use ceno_host::CenoStdin;
 use ceno_zkvm::{
     self,
     e2e::{Checkpoint, Preset, run_e2e_with_checkpoint, setup_platform},
+    scheme::{create_backend, create_prover},
 };
 mod alloc;
 use criterion::*;
 use ff_ext::BabyBearExt4;
-use gkr_iop::cpu::{CpuBackend, CpuProver};
+use gkr_iop::cpu::default_backend_config;
 use mpcs::BasefoldDefault;
 
 criterion_group! {
@@ -36,7 +37,9 @@ fn setup() -> (Program, Platform) {
 
 fn is_prime_1(c: &mut Criterion) {
     let (program, platform) = setup();
-    let backend: Rc<_> = CpuBackend::<E, Pcs>::default().into();
+
+    let (max_num_variables, security_level) = default_backend_config();
+    let backend = create_backend::<E, Pcs>(max_num_variables, security_level);
 
     for n in [100u32, 10000u32, 50000u32] {
         let max_steps = usize::MAX;
@@ -56,7 +59,7 @@ fn is_prime_1(c: &mut Criterion) {
 
                     for _ in 0..iters {
                         let result = run_e2e_with_checkpoint::<E, Pcs, _, _>(
-                            CpuProver::new(backend.clone()),
+                            create_prover(backend.clone()),
                             program.clone(),
                             platform.clone(),
                             &hints,