Add an opt-in first-party CUDA proving backend

.github/workflows/ci.yml

@@ -10,6 +10,9 @@ concurrency:
   group: ${{ github.workflow }}-${{ github.head_ref || github.run_id }}
   cancel-in-progress: true
 
+permissions:
+  contents: read
+
 jobs:
   linux-test:
     # Change to warp-ubuntu-latest-x64-16x for a more powerful runner (GitHub App must be enabled for this repo)
@@ -39,10 +42,12 @@ jobs:
       - uses: Swatinem/rust-cache@v2
       - name: Check Rustfmt code style
         run: cargo fmt --all --check
-      - name: Check *everything* compiles
-        run: cargo check --all-targets --all-features --workspace
+      # Runtime correctness is validated on NVIDIA runners via cuda/smoke.sh. Keep the
+      # ordinary CI path independent of the CUDA toolkit.
+      - name: Check CPU targets compile
+        run: cargo check --release --all-targets --features parallel --workspace
       - name: Check clippy lints
-        run: cargo clippy --workspace --all-targets --all-features -- -D warnings
+        run: cargo clippy --release --workspace --all-targets --features parallel -- -D warnings
       - name: Doctests
         run: cargo test --doc --workspace
       - name: Get Rust version
@@ -53,3 +58,22 @@ jobs:
         uses: EmbarkStudios/cargo-deny-action@v2
         with:
           rust-version: ${{ env.RUST_VERSION }}
+
+  cuda-compile:
+    # A toolkit container is sufficient for compile coverage; no GPU runner
+    # is needed. Warp's larger x64 runner keeps the full cubin-set build fast
+    # and provides the same native-codegen compatibility domain as Ix CI.
+    runs-on: warp-ubuntu-latest-x64-8x
+    container: nvidia/cuda:12.8.1-devel-ubuntu24.04
+    steps:
+      - uses: actions/checkout@v6
+      - name: Install Rust bootstrap dependencies
+        run: |
+          apt-get update
+          apt-get install --yes --no-install-recommends build-essential ca-certificates curl
+      - uses: actions-rust-lang/setup-rust-toolchain@v1
+        with:
+          rustflags: -Ctarget-cpu=native -Dwarnings
+          cache-key: warp-x64
+      - name: Compile and lint CUDA targets
+        run: cargo clippy --release --locked --all-targets --features parallel,cuda -- -D warnings

Cargo.toml

@@ -13,6 +13,7 @@ rust-version = "1.91"
 
 [dependencies]
 tracing = "0.1"
+itertools = { version = "0.14", optional = true }
 serde = { version = "1", features = ["derive"] }
 bincode = { version = "2", features = ["serde"] }
 p3-air = { git = "https://github.com/Plonky3/Plonky3", rev = "e9d75614dd6816f9b5dbb4413c69be63536efd64" }
@@ -23,6 +24,7 @@ p3-dft = { git = "https://github.com/Plonky3/Plonky3", rev = "e9d75614dd6816f9b5
 p3-field = { git = "https://github.com/Plonky3/Plonky3", rev = "e9d75614dd6816f9b5dbb4413c69be63536efd64" }
 p3-fri = { git = "https://github.com/Plonky3/Plonky3", rev = "e9d75614dd6816f9b5dbb4413c69be63536efd64" }
 p3-keccak = { git = "https://github.com/Plonky3/Plonky3", rev = "e9d75614dd6816f9b5dbb4413c69be63536efd64" }
+p3-interpolation = { git = "https://github.com/Plonky3/Plonky3", rev = "e9d75614dd6816f9b5dbb4413c69be63536efd64", optional = true }
 p3-blake3 = { git = "https://github.com/Plonky3/Plonky3", rev = "e9d75614dd6816f9b5dbb4413c69be63536efd64" }
 p3-matrix = { git = "https://github.com/Plonky3/Plonky3", rev = "e9d75614dd6816f9b5dbb4413c69be63536efd64" }
 p3-maybe-rayon = { git = "https://github.com/Plonky3/Plonky3", rev = "e9d75614dd6816f9b5dbb4413c69be63536efd64" }
@@ -43,6 +45,10 @@ harness = false
 
 [features]
 parallel = ["p3-maybe-rayon/parallel"]
+# Use the first-party CUDA Goldilocks DFT/LDE backend. Enabling this feature
+# requires a CUDA toolkit at build time and an NVIDIA GPU at runtime; the
+# default CPU build never invokes nvcc or links the CUDA runtime.
+cuda = ["dep:itertools", "dep:p3-interpolation"]
 
 # Similar to `release`, but preserves debug info
 [profile.dev-ci]

README.md

@@ -17,6 +17,8 @@ lookup arguments for shared state.
   a batteries-included Goldilocks/Blake3 instantiation is provided
 - **Serialization** — `Proof::to_bytes` / `Proof::from_bytes` via bincode
 - **Parallel proving** — opt-in via the `parallel` feature flag
+- **CUDA transforms** — opt-in first-party Goldilocks DFT/LDE backend via the
+  `cuda` feature; normal builds remain independent of CUDA
 
 ## Reference configuration
 
@@ -75,6 +77,33 @@ using a U32 addition circuit with lookups and a preprocessed byte table.
 Use `--features parallel` for representative numbers. Native SIMD instructions
 are enabled by default via `.cargo/config.toml`.
 
+## CUDA acceleration
+
+The optional `cuda` feature routes the production Goldilocks configuration's
+PCS and quotient transforms through first-party CUDA kernels. CUDA and CPU
+proofs generated by the same revision are byte-identical and use the same CPU
+verifier. Canonical field serialization changes newly generated CPU proof bytes
+relative to revisions before this backend; previously generated proofs remain
+verifiable, but byte-addressed proof caches must treat this as a compatibility
+boundary.
+The default build does not invoke `nvcc` or link the CUDA runtime. Enable the
+backend explicitly:
+
+```sh
+cargo test --release --features parallel,cuda
+./cuda/smoke.sh
+cargo run --release --locked --features parallel,cuda --example cuda_dft_bench
+```
+
+The CUDA build requires an NVIDIA GPU and CUDA toolkit. It keeps trace LDEs,
+mixed-height BLAKE3 Merkle commitments, lookup construction, quotient
+evaluation, batched openings, and FRI folding resident on the device. Dated
+hardware and downstream Ix measurements live in
+[docs/cuda-benchmarks.md](docs/cuda-benchmarks.md).
+
+See [cuda/README.md](cuda/README.md) for architecture, build settings, the NVIDIA
+correctness harness, platform limitations, and benchmark commands.
+
 ## License
 
 MIT or Apache-2.0

build.rs

@@ -0,0 +1,205 @@
+//! CUDA build isolation.
+//!
+//! This build script is deliberately a no-op unless Cargo enables the
+//! `cuda` feature. Normal CPU builds therefore need neither nvcc nor CUDA
+//! headers/libraries.
+
+use std::env;
+use std::path::{Path, PathBuf};
+use std::process::Command;
+
+fn main() {
+    println!("cargo:rerun-if-changed=build.rs");
+    println!("cargo:rerun-if-changed=cuda/kernels.cu");
+    println!("cargo:rerun-if-env-changed=NVCC");
+    println!("cargo:rerun-if-env-changed=CUDA_HOME");
+    println!("cargo:rerun-if-env-changed=CUDA_PATH");
+    println!("cargo:rerun-if-env-changed=MULTI_STARK_CUDA_ARCHS");
+
+    if env::var_os("CARGO_FEATURE_CUDA").is_none() {
+        return;
+    }
+
+    assert_eq!(
+        env::var("CARGO_CFG_TARGET_OS").as_deref(),
+        Ok("linux"),
+        "the first-party CUDA backend currently supports Linux targets only"
+    );
+    assert_eq!(
+        env::var("CARGO_CFG_TARGET_ARCH").as_deref(),
+        Ok("x86_64"),
+        "the first-party CUDA backend currently supports x86_64 targets only"
+    );
+
+    let nvcc = nvcc_path();
+    let out_dir = PathBuf::from(env::var_os("OUT_DIR").expect("Cargo did not set OUT_DIR"));
+    let library = out_dir.join("libmulti_stark_cuda.a");
+    let architectures = cuda_architectures(&nvcc);
+
+    let mut command = Command::new(&nvcc);
+    command
+        .arg("--lib")
+        .arg("--std=c++17")
+        .arg("--cudart=static")
+        .arg("--default-stream=per-thread")
+        .arg("-O3")
+        .arg("-lineinfo")
+        .arg("--compiler-options=-fPIC")
+        .arg("-o")
+        .arg(&library)
+        .arg("cuda/kernels.cu");
+
+    for architecture in &architectures {
+        command.arg(format!(
+            "-gencode=arch=compute_{architecture},code=sm_{architecture}"
+        ));
+    }
+    // Emit native cubins only. A toolkit can produce PTX newer than the
+    // installed driver understands even when both support the GPU's native
+    // ISA (for example nvcc 13.3 with a CUDA-13.2-capable production
+    // driver). In that case the runtime may select the incompatible PTX and
+    // reject an otherwise usable exact-architecture cubin with error 222.
+
+    let status = command.status().unwrap_or_else(|error| {
+        panic!(
+            "failed to execute {:?}: {error}; install the CUDA toolkit or set NVCC",
+            nvcc
+        )
+    });
+    assert!(status.success(), "nvcc failed with status {status}");
+
+    println!("cargo:rustc-link-search=native={}", out_dir.display());
+    println!("cargo:rustc-link-lib=static=multi_stark_cuda");
+    // Bundle the CUDA runtime into Rust staticlib consumers as well as normal
+    // Cargo binaries. Dynamic native dependencies are not propagated through
+    // a Rust staticlib to a foreign final linker (for example, Lean/Lake),
+    // which otherwise leaves the CUDA registration and runtime symbols
+    // unresolved. This remains entirely behind the `cuda` feature.
+    println!("cargo:rustc-link-lib=static=cudart_static");
+    println!("cargo:rustc-link-lib=dylib=dl");
+    println!("cargo:rustc-link-lib=dylib=rt");
+    println!("cargo:rustc-link-lib=dylib=pthread");
+    println!("cargo:rustc-link-lib=dylib=stdc++");
+
+    for directory in cuda_library_directories(&nvcc) {
+        if directory.is_dir() {
+            println!("cargo:rustc-link-search=native={}", directory.display());
+        }
+    }
+}
+
+fn nvcc_path() -> std::ffi::OsString {
+    if let Some(nvcc) = env::var_os("NVCC") {
+        return nvcc;
+    }
+    if let Some(root) = env::var_os("CUDA_HOME").or_else(|| env::var_os("CUDA_PATH")) {
+        let candidate = PathBuf::from(root).join("bin/nvcc");
+        if candidate.is_file() {
+            return candidate.into_os_string();
+        }
+    }
+    "nvcc".into()
+}
+
+fn cuda_architectures(nvcc: &std::ffi::OsStr) -> Vec<String> {
+    let supported = supported_architectures(nvcc);
+    let configured = env::var("MULTI_STARK_CUDA_ARCHS").ok().map_or_else(
+        || {
+            let preferred = ["80", "86", "89", "90", "100", "120"];
+            let selected = preferred
+                .into_iter()
+                .filter(|architecture| supported.iter().any(|item| item == architecture))
+                .collect::<Vec<_>>();
+            let defaults = if selected.is_empty() {
+                detect_nvidia_architectures().unwrap_or_else(|| "80".to_owned())
+            } else {
+                selected.join(",")
+            };
+            (defaults, false)
+        },
+        |value| (value, true),
+    );
+    let (configured, explicitly_configured) = configured;
+    let architectures: Vec<_> = configured
+        .split(',')
+        .map(str::trim)
+        .filter(|architecture| !architecture.is_empty())
+        .map(|architecture| {
+            assert!(
+                (2..=3).contains(&architecture.len())
+                    && architecture.chars().all(|character| character.is_ascii_digit()),
+                "invalid CUDA architecture {architecture:?}; expected comma-separated numbers such as 80,90"
+            );
+            assert!(
+                !explicitly_configured
+                    || supported.is_empty()
+                    || supported.iter().any(|item| item == architecture),
+                "nvcc does not support CUDA architecture sm_{architecture}"
+            );
+            architecture.to_owned()
+        })
+        .collect();
+    assert!(
+        !architectures.is_empty(),
+        "MULTI_STARK_CUDA_ARCHS must contain at least one architecture"
+    );
+    architectures
+}
+
+fn supported_architectures(nvcc: &std::ffi::OsStr) -> Vec<String> {
+    let Ok(output) = Command::new(nvcc).arg("--list-gpu-code").output() else {
+        return Vec::new();
+    };
+    if !output.status.success() {
+        return Vec::new();
+    }
+    String::from_utf8_lossy(&output.stdout)
+        .split_whitespace()
+        .filter_map(|code| code.strip_prefix("sm_"))
+        .filter(|code| code.chars().all(|character| character.is_ascii_digit()))
+        .map(str::to_owned)
+        .collect()
+}
+
+/// Detect the architecture of the installed GPU when building on the target
+/// machine. This avoids requiring users to translate `nvidia-smi`'s `12.0`
+/// compute capability into nvcc's `sm_120` spelling. Cross builds and hosts
+/// without a visible GPU retain the portable sm_80 default and can still set
+/// `MULTI_STARK_CUDA_ARCHS` explicitly.
+fn detect_nvidia_architectures() -> Option<String> {
+    let output = Command::new("nvidia-smi")
+        .args(["--query-gpu=compute_cap", "--format=csv,noheader"])
+        .output()
+        .ok()?;
+    if !output.status.success() {
+        return None;
+    }
+    let capabilities = String::from_utf8(output.stdout).ok()?;
+    let mut architectures = capabilities
+        .lines()
+        .map(|capability| {
+            capability
+                .chars()
+                .filter(|character| character.is_ascii_digit())
+                .collect::<String>()
+        })
+        .filter(|architecture| !architecture.is_empty())
+        .collect::<Vec<_>>();
+    architectures.sort_unstable();
+    architectures.dedup();
+    (!architectures.is_empty()).then(|| architectures.join(","))
+}
+
+fn cuda_library_directories(nvcc: &std::ffi::OsStr) -> Vec<PathBuf> {
+    let root = env::var_os("CUDA_HOME")
+        .or_else(|| env::var_os("CUDA_PATH"))
+        .map(PathBuf::from)
+        .or_else(|| {
+            let nvcc = Path::new(nvcc);
+            nvcc.is_absolute()
+                .then(|| nvcc.parent()?.parent().map(Path::to_path_buf))
+                .flatten()
+        })
+        .unwrap_or_else(|| PathBuf::from("/usr/local/cuda"));
+    vec![root.join("lib64"), root.join("targets/x86_64-linux/lib")]
+}