Skip to content

Commit 0c86bbf

Browse files
kunxian-xiakunxian-xia
authored
Add CUDA tracegen (#1379)
## Problem Stacked on #1378. Recursion v2 had CPU-side trace generation and proof wiring, but the CUDA path was not usable end to end for inner recursion proving. That left GPU proving unable to exercise the verifier, PCS, transcript, proof-shape, and public-value traces in the same path used by aggregation. ## Design Rationale The main design is to keep recursion proving backend selection behind the existing generic engine/backend type and enable the CUDA specialization under the `cuda` feature. CUDA tracegen receives explicit, typed inputs for proof shape, public values, PCS, transcript, and verifier data instead of duplicating host-side layout assumptions. For the OpenVM CUDA prover path, `PcsJaggedAssistQAir` is kept within degree 4 by materializing the high-degree factor in the trace. This avoids the logup zerocheck round0 5-coset CUDA limit while preserving the AIR relation through explicit constraints. ## Change Highlights - `ceno_recursion_v2`: add CUDA tracegen ABI, build wiring, fixtures, and test harness. - `ceno_recursion_v2`: wire CUDA inner tracegen/proving into `e2e_aggregate`. - `ceno_recursion_v2`: add CUDA layouts for proof shape, public values, transcript, verifier, PCS, and main recursion traces. - `ceno_recursion_v2`: reduce `PcsJaggedAssistQAir` max constraint degree for CUDA proving compatibility. ## Benchmark / Performance Impact This PR enables GPU recursion proving. There is no direct `master` baseline for the CUDA recursion path because the path was not functional before this PR. ### Operation | Operation | master (s) | this PR (s) | Improve (master -> this PR) | |-----------|------------|-------------|-----------------------------| | recursion aggregation prove, fibonacci | N/A | 0.235 | enables CUDA path | | recursion aggregation verify, fibonacci | N/A | 0.037 | enables CUDA path | | recursion aggregation prove, keccak_syscall | N/A | 0.185 | enables CUDA path | | recursion aggregation verify, keccak_syscall | N/A | 0.033 | enables CUDA path | ### Layer | Layer | master (s) | this PR (s) | Improve (master -> this PR) | |-------|------------|-------------|-----------------------------| | CUDA inner recursion proving | N/A | 0.185-0.235 | new working path | Benchmark command(s): ```sh RUST_LOG=openvm_cuda_backend::logup_zerocheck=debug,ceno_recursion_v2::continuation::prover::inner=info RUST_MIN_STACK=67108864 cargo run --release --features cuda --bin e2e_aggregate -- /home/zbx/blockchain/ceno/examples/target/riscv32im-ceno-zkvm-elf/release/examples/fibonacci --platform ceno --max-cycle-per-shard 16000 --hints 10 --public-io 4191 RUST_LOG=openvm_cuda_backend::logup_zerocheck=debug,ceno_recursion_v2::continuation::prover::inner=info RUST_MIN_STACK=67108864 cargo run --release --features cuda --bin e2e_aggregate -- /home/zbx/blockchain/ceno/examples/target/riscv32im-ceno-zkvm-elf/release/examples/keccak_syscall --platform ceno --max-cycle-per-shard 16000 ``` Environment (CPU/GPU, core count, rust toolchain, commit hash): - local CUDA dev machine - commit: `25d8ec89` raw data: - master: CUDA recursion proving path not available - this PR: - fibonacci: proved recursion aggregation in 0.235s, verified in 0.037s; final exit is the known internal aggregation tree-reduction gap for 4 leaf proofs - keccak_syscall: proved recursion aggregation in 0.185s, verified in 0.033s; exit code 0 ## Testing ```sh cargo fmt --check git diff --check -- src/pcs/mod.rs src/bin/e2e_aggregate.rs RUST_MIN_STACK=67108864 cargo run --release --features cuda --bin e2e_aggregate -- /home/zbx/blockchain/ceno/examples/target/riscv32im-ceno-zkvm-elf/release/examples/keccak_syscall --platform ceno --max-cycle-per-shard 16000 RUST_MIN_STACK=67108864 cargo run --release --features cuda --bin e2e_aggregate -- /home/zbx/blockchain/ceno/examples/target/riscv32im-ceno-zkvm-elf/release/examples/fibonacci --platform ceno --max-cycle-per-shard 16000 --hints 10 --public-io 4191 ``` ## Risks and Rollout The main risk is CUDA/host layout drift. The PR mitigates this with typed ABI structs, generated fixtures, and CUDA tracegen tests. Rollout is feature-gated behind `cuda`, so CPU proving remains the fallback path. ## Follow-ups (optional) - Implement internal aggregation tree reduction for multiple leaf proofs. - Remove temporary CUDA diagnostic prints from the local OpenVM checkout or rebuild clean CUDA objects before collecting final benchmark logs. ## Copilot Reviewer Directive (keep this section) When Copilot reviews this PR, apply `.github/copilot-instructions.md` strictly.
1 parent 107577b commit 0c86bbf

28 files changed

Lines changed: 1604 additions & 218 deletions

File tree

‎Cargo.lock‎

Lines changed: 3 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎ceno_recursion_v2/Cargo.lock‎

Lines changed: 1 addition & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎ceno_recursion_v2/Cargo.toml‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -56,10 +56,15 @@ verify-stark = { git = "https://github.com/openvm-org/openvm.git", package = "op
5656
whir = { git = "https://github.com/scroll-tech/gkr-backend.git", package = "whir", branch = "feat/bump-p3" }
5757
witness = { git = "https://github.com/scroll-tech/gkr-backend.git", package = "witness", branch = "feat/bump-p3" }
5858

59+
[build-dependencies]
60+
openvm-cuda-builder = { git = "https://github.com/openvm-org/stark-backend.git", branch = "develop-v2", optional = true }
61+
5962
[features]
6063
cuda = [
64+
"dep:openvm-cuda-builder",
6165
"dep:openvm-cuda-backend",
6266
"dep:openvm-cuda-common",
67+
"openvm-circuit-primitives/cuda",
6368
]
6469
default = ["parallel"]
6570
parallel = [

‎ceno_recursion_v2/build.rs‎

Lines changed: 44 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,44 @@
1+
#[cfg(feature = "cuda")]
2+
use {
3+
openvm_cuda_builder::{CudaBuilder, cuda_available},
4+
std::{env, process::Command},
5+
};
6+
7+
fn main() {
8+
#[cfg(feature = "cuda")]
9+
{
10+
if !cuda_available() {
11+
println!(
12+
"cargo:warning=CUDA toolkit is not available; skipping ceno_recursion_v2 CUDA kernels"
13+
);
14+
return;
15+
}
16+
17+
if env::var_os("CUDA_ARCH").is_none()
18+
&& Command::new("nvidia-smi")
19+
.args(["--query-gpu=compute_cap", "--format=csv,noheader"])
20+
.output()
21+
.map(|output| !output.status.success() || output.stdout.is_empty())
22+
.unwrap_or(true)
23+
{
24+
println!(
25+
"cargo:warning=nvidia-smi did not report a CUDA device; defaulting CUDA_ARCH=80 for kernel compilation"
26+
);
27+
// The build script is single-threaded here; this only controls CudaBuilder arch detection.
28+
unsafe { env::set_var("CUDA_ARCH", "80") };
29+
}
30+
31+
let common = CudaBuilder::new()
32+
.include_from_dep("DEP_CUDA_COMMON_INCLUDE")
33+
.include("cuda/include")
34+
.flag("-Xcompiler=-Wno-maybe-uninitialized");
35+
36+
common.emit_link_directives();
37+
38+
common
39+
.clone()
40+
.library_name("cuda-ceno-recursion-v2")
41+
.files_from_glob("cuda/src/**/*.cu")
42+
.build();
43+
}
44+
}

‎ceno_recursion_v2/cuda/include/primitives/trace_access.h‎

Lines changed: 67 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,67 @@
1+
#pragma once
2+
3+
#include "fp.h"
4+
#include <cstddef>
5+
6+
struct RowSlice {
7+
Fp *ptr;
8+
size_t stride;
9+
10+
__device__ RowSlice(Fp *ptr, size_t stride) : ptr(ptr), stride(stride) {}
11+
12+
__device__ __forceinline__ Fp &operator[](size_t column_index) const {
13+
return ptr[column_index * stride];
14+
}
15+
16+
__device__ static RowSlice null() { return RowSlice(nullptr, 0); }
17+
18+
__device__ bool is_valid() const { return ptr != nullptr; }
19+
20+
template <typename T>
21+
__device__ __forceinline__ void write(size_t column_index, T value) const {
22+
ptr[column_index * stride] = value;
23+
}
24+
25+
template <typename T>
26+
__device__ __forceinline__ void write_array(size_t column_index, size_t length, const T *values)
27+
const {
28+
#pragma unroll
29+
for (size_t i = 0; i < length; i++) {
30+
ptr[(column_index + i) * stride] = values[i];
31+
}
32+
}
33+
34+
template <typename T>
35+
__device__ __forceinline__ void write_bits(size_t column_index, const T value) const {
36+
#pragma unroll
37+
for (size_t i = 0; i < sizeof(T) * 8; i++) {
38+
ptr[(column_index + i) * stride] = (value >> i) & 1;
39+
}
40+
}
41+
42+
__device__ __forceinline__ void fill_zero(size_t column_index_from, size_t length) const {
43+
#pragma unroll
44+
for (size_t i = 0, c = column_index_from; i < length; i++, c++) {
45+
ptr[c * stride] = 0;
46+
}
47+
}
48+
49+
__device__ __forceinline__ RowSlice slice_from(size_t column_index) const {
50+
return RowSlice(ptr + column_index * stride, stride);
51+
}
52+
53+
__device__ __forceinline__ RowSlice shift_row(size_t n) const {
54+
return RowSlice(ptr + n, stride);
55+
}
56+
};
57+
58+
#define COL_INDEX(STRUCT, FIELD) (offsetof(STRUCT<uint8_t>, FIELD))
59+
#define COL_ARRAY_LEN(STRUCT, FIELD) (sizeof(static_cast<STRUCT<uint8_t> *>(nullptr)->FIELD))
60+
#define COL_WRITE_VALUE(ROW, STRUCT, FIELD, VALUE) (ROW).write(COL_INDEX(STRUCT, FIELD), VALUE)
61+
#define COL_WRITE_ARRAY(ROW, STRUCT, FIELD, VALUES) \
62+
(ROW).write_array(COL_INDEX(STRUCT, FIELD), COL_ARRAY_LEN(STRUCT, FIELD), VALUES)
63+
#define COL_WRITE_BITS(ROW, STRUCT, FIELD, VALUE) (ROW).write_bits(COL_INDEX(STRUCT, FIELD), VALUE)
64+
#define COL_FILL_ZERO(ROW, STRUCT, FIELD) \
65+
(ROW).fill_zero( \
66+
COL_INDEX(STRUCT, FIELD), sizeof(static_cast<STRUCT<uint8_t> *>(nullptr)->FIELD) \
67+
)

‎ceno_recursion_v2/cuda/include/ptr_array.h‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,16 @@
1+
#pragma once
2+
3+
#include <cstddef>
4+
5+
template <typename T, size_t N> struct PtrArray {
6+
T *arr[N];
7+
8+
__device__ __host__ PtrArray(T **raw_ptr) {
9+
for (size_t i = 0; i < N; i++) {
10+
arr[i] = raw_ptr[i];
11+
}
12+
}
13+
14+
__device__ __host__ T *operator[](size_t i) { return arr[i]; }
15+
__device__ __host__ const T *operator[](size_t i) const { return arr[i]; }
16+
};

‎ceno_recursion_v2/cuda/include/switch_macro.h‎

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,33 @@
1+
#pragma once
2+
3+
#define UNWRAP(...) __VA_ARGS__
4+
5+
#define CASE_BLOCK(N, name, body) \
6+
case N: { \
7+
constexpr size_t name = N; \
8+
UNWRAP body break; \
9+
}
10+
11+
#define APPLY_CASES_1(name, body, c1) CASE_BLOCK(c1, name, body)
12+
#define APPLY_CASES_2(name, body, c1, c2) APPLY_CASES_1(name, body, c1) CASE_BLOCK(c2, name, body)
13+
#define APPLY_CASES_3(name, body, c1, c2, c3) \
14+
APPLY_CASES_2(name, body, c1, c2) CASE_BLOCK(c3, name, body)
15+
#define APPLY_CASES_4(name, body, c1, c2, c3, c4) \
16+
APPLY_CASES_3(name, body, c1, c2, c3) CASE_BLOCK(c4, name, body)
17+
#define APPLY_CASES_5(name, body, c1, c2, c3, c4, c5) \
18+
APPLY_CASES_4(name, body, c1, c2, c3, c4) CASE_BLOCK(c5, name, body)
19+
#define APPLY_CASES_6(name, body, c1, c2, c3, c4, c5, c6) \
20+
APPLY_CASES_5(name, body, c1, c2, c3, c4, c5) CASE_BLOCK(c6, name, body)
21+
#define APPLY_CASES_7(name, body, c1, c2, c3, c4, c5, c6, c7) \
22+
APPLY_CASES_6(name, body, c1, c2, c3, c4, c5, c6) CASE_BLOCK(c7, name, body)
23+
#define APPLY_CASES_8(name, body, c1, c2, c3, c4, c5, c6, c7, c8) \
24+
APPLY_CASES_7(name, body, c1, c2, c3, c4, c5, c6, c7) CASE_BLOCK(c8, name, body)
25+
26+
#define GET_MACRO(_1, _2, _3, _4, _5, _6, _7, _8, NAME, ...) NAME
27+
28+
#define SWITCH_BLOCK(value, name, body, ...) \
29+
switch (value) { \
30+
GET_MACRO(__VA_ARGS__, APPLY_CASES_8, APPLY_CASES_7, APPLY_CASES_6, APPLY_CASES_5, APPLY_CASES_4, APPLY_CASES_3, APPLY_CASES_2, APPLY_CASES_1)( \
31+
name, body, __VA_ARGS__ \
32+
) default : return cudaErrorInvalidConfiguration; \
33+
}

‎ceno_recursion_v2/cuda/include/types.h‎

Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,37 @@
1+
#pragma once
2+
3+
#include "fp.h"
4+
5+
#include <cstddef>
6+
#include <cstdint>
7+
8+
constexpr size_t D_EF = 4;
9+
constexpr size_t DIGEST_SIZE = 8;
10+
typedef Fp Digest[DIGEST_SIZE];
11+
12+
typedef struct {
13+
size_t air_idx;
14+
uint8_t log_height;
15+
} TraceHeight;
16+
17+
typedef struct {
18+
size_t cached_idx;
19+
size_t total_interactions;
20+
size_t num_air_id_lookups;
21+
} TraceMetadata;
22+
23+
typedef struct {
24+
size_t air_idx;
25+
size_t air_num_pvs;
26+
size_t num_airs;
27+
size_t pv_idx;
28+
Fp value;
29+
} PublicValueData;
30+
31+
typedef struct {
32+
size_t num_cached;
33+
size_t num_interactions_per_row;
34+
size_t total_width;
35+
bool has_preprocessed;
36+
bool need_rot;
37+
} AirData;

‎ceno_recursion_v2/cuda/src/proof_shape/pvs.cu‎

Lines changed: 83 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,83 @@
1+
#include "fp.h"
2+
#include "launcher.cuh"
3+
#include "primitives/trace_access.h"
4+
#include "ptr_array.h"
5+
#include "switch_macro.h"
6+
#include "types.h"
7+
8+
#include <cassert>
9+
#include <cstddef>
10+
#include <cstdint>
11+
12+
template <typename T> struct PublicValuesCols {
13+
T is_valid;
14+
T proof_idx;
15+
T air_idx;
16+
T pv_idx;
17+
T is_first_in_proof;
18+
T is_first_in_air;
19+
T tidx;
20+
T value;
21+
};
22+
23+
template <size_t NUM_PROOFS>
24+
__global__ void public_values_tracegen(
25+
Fp *trace,
26+
size_t height,
27+
PtrArray<PublicValueData, NUM_PROOFS> pvs_data,
28+
PtrArray<size_t, NUM_PROOFS> pvs_tidx,
29+
size_t num_pvs
30+
) {
31+
uint32_t idx = blockIdx.x * blockDim.x + threadIdx.x;
32+
RowSlice row(trace + idx, height);
33+
if (idx < NUM_PROOFS * num_pvs) {
34+
size_t proof_idx = idx / num_pvs;
35+
size_t record_idx = idx % num_pvs;
36+
PublicValueData pv_data = pvs_data[proof_idx][record_idx];
37+
size_t starting_tidx = pvs_tidx[proof_idx][pv_data.air_idx];
38+
39+
COL_WRITE_VALUE(row, PublicValuesCols, is_valid, Fp::one());
40+
COL_WRITE_VALUE(row, PublicValuesCols, proof_idx, proof_idx);
41+
COL_WRITE_VALUE(row, PublicValuesCols, air_idx, pv_data.air_idx);
42+
COL_WRITE_VALUE(row, PublicValuesCols, pv_idx, pv_data.pv_idx);
43+
COL_WRITE_VALUE(row, PublicValuesCols, is_first_in_proof, record_idx == 0);
44+
COL_WRITE_VALUE(row, PublicValuesCols, is_first_in_air, pv_data.pv_idx == 0);
45+
COL_WRITE_VALUE(row, PublicValuesCols, tidx, starting_tidx + pv_data.pv_idx);
46+
COL_WRITE_VALUE(row, PublicValuesCols, value, pv_data.value);
47+
} else {
48+
row.fill_zero(0, sizeof(PublicValuesCols<uint8_t>));
49+
}
50+
}
51+
52+
extern "C" int _public_values_recursion_tracegen(
53+
Fp *d_trace,
54+
size_t height,
55+
PublicValueData **d_pvs_data,
56+
size_t **d_pvs_tidx,
57+
size_t num_proofs,
58+
size_t num_pvs,
59+
cudaStream_t stream
60+
) {
61+
assert((height & (height - 1)) == 0);
62+
auto [grid, block] = kernel_launch_params(height);
63+
SWITCH_BLOCK(
64+
num_proofs,
65+
NUM_PROOFS,
66+
(public_values_tracegen<NUM_PROOFS><<<grid, block, 0, stream>>>(
67+
d_trace,
68+
height,
69+
PtrArray<PublicValueData, NUM_PROOFS>(d_pvs_data),
70+
PtrArray<size_t, NUM_PROOFS>(d_pvs_tidx),
71+
num_pvs
72+
);),
73+
1,
74+
2,
75+
3,
76+
4,
77+
5,
78+
6,
79+
7,
80+
8
81+
)
82+
return CHECK_KERNEL();
83+
}

‎ceno_recursion_v2/scripts/run_cuda_tracegen_test.sh‎

Lines changed: 59 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,59 @@
1+
#!/usr/bin/env bash
2+
# Run recursion v2 CPU-vs-GPU trace generation comparison tests.
3+
set -euo pipefail
4+
5+
REPO_ROOT="$(cd "$(dirname "$0")/../.." && pwd)"
6+
FIXTURE_DIR="${1:-${CENO_RECURSION_V2_FIXTURE_DIR:-$REPO_ROOT/ceno_recursion_v2/src/imported}}"
7+
PROOF_FILE="$FIXTURE_DIR/proof.bin"
8+
VK_FILE="$FIXTURE_DIR/vk.bin"
9+
FORCE_REGEN="${CENO_RECURSION_V2_FORCE_REGEN:-0}"
10+
MAX_CYCLE_PER_SHARD="${CENO_RECURSION_V2_MAX_CYCLE_PER_SHARD:-16000}"
11+
HINTS="${CENO_RECURSION_V2_HINTS:-10}"
12+
PUBLIC_IO="${CENO_RECURSION_V2_PUBLIC_IO:-4191}"
13+
PCS="${CENO_RECURSION_V2_PCS:-jagged}"
14+
15+
export CENO_RECURSION_V2_FIXTURE_DIR="$FIXTURE_DIR"
16+
17+
if [[ "$FORCE_REGEN" != "1" && -f "$PROOF_FILE" && -f "$VK_FILE" ]]; then
18+
echo "[fixtures] found:"
19+
ls -lh "$PROOF_FILE" "$VK_FILE"
20+
else
21+
echo "[fixtures] generating base-layer proofs (fibonacci, --pcs=$PCS, --max-cycle-per-shard=$MAX_CYCLE_PER_SHARD) ..."
22+
23+
ELF="$REPO_ROOT/examples/target/riscv32im-ceno-zkvm-elf/release/examples/fibonacci"
24+
if [[ ! -f "$ELF" ]]; then
25+
echo "[fixtures] building example ELFs..."
26+
cargo build --release --manifest-path "$REPO_ROOT/examples/Cargo.toml" --examples
27+
fi
28+
29+
mkdir -p "$FIXTURE_DIR"
30+
cargo run --release --package ceno_zkvm --bin e2e -- \
31+
--platform=ceno \
32+
--pcs="$PCS" \
33+
--max-cycle-per-shard="$MAX_CYCLE_PER_SHARD" \
34+
--hints="$HINTS" \
35+
--public-io="$PUBLIC_IO" \
36+
"$ELF" \
37+
"$PROOF_FILE" \
38+
"$VK_FILE"
39+
40+
echo "[fixtures] generated:"
41+
ls -lh "$PROOF_FILE" "$VK_FILE"
42+
fi
43+
44+
cd "$REPO_ROOT"
45+
TESTS=(
46+
system::cuda_tracegen_tests::test_cuda_tracegen_compare_single_fixture_proof
47+
system::cuda_tracegen_tests::test_cuda_tracegen_compare_multi_fixture_proofs
48+
system::cuda_tracegen_tests::test_cuda_tracegen_required_heights_match_cpu
49+
)
50+
51+
for test_name in "${TESTS[@]}"; do
52+
echo ""
53+
echo "[test] running $test_name ..."
54+
RUST_MIN_STACK="${RUST_MIN_STACK:-33554432}" \
55+
cargo test --release --manifest-path ceno_recursion_v2/Cargo.toml \
56+
--features cuda \
57+
"$test_name" \
58+
-- --exact --nocapture
59+
done

0 commit comments

Comments
 (0)