From 8a569009e784f4ed5671572724b83f7176b8758e Mon Sep 17 00:00:00 2001 From: Arthur Paulino Date: Wed, 26 Aug 2026 14:59:42 +0000 Subject: [PATCH] Add an opt-in first-party CUDA proving backend Accelerate the Goldilocks/BLAKE3 prover with first-party CUDA while keeping the default CPU build independent of nvcc, the CUDA runtime, and NVIDIA hardware. The GPU path keeps DFT/LDE, mixed-height BLAKE3 commitments, lookup construction, quotient evaluation, batched openings, and binary FRI folding device-resident. Preserve the protocol, proof format, proof sizes, transcript order, and CPU verification behavior. Canonicalize lazy Goldilocks representatives at proof serialization so CPU and CUDA proofs are byte-identical. This intentionally changes newly generated proof bytes relative to c72d321; old proofs remain verifiable, and README.md documents the compatibility boundary. Harden the backend after adversarial review: keep pinned traces alive through CUDA handle destruction, reject unsupported Merkle caps and multi-bit GPU FRI, remove the racy generic FRI fold, canonicalize host uploads with grid-stride coverage, place kernel-visible metadata in managed allocations, derive spill headroom from device memory while reserving later-stage scratch, honor device selection, gate dynamic shared memory on device limits, remove the fixed constant-cache ceiling, validate FFI dimensions including one-row quotient steps, and discover nvcc through CUDA_HOME/CUDA_PATH. Generate quotient-coset selectors directly in device memory from compact geometric parameters, using chunked recurrences and in-place batch inversion. This removes the largest remaining happy-path host-to-device upload without increasing peak VRAM. At Ix Vector.extract_append q50 scale, combined inner/outer quotient time falls from 1.73s to 1.45s. Build native cubins for a supported common architecture set, add CUDA compile CI, a threshold-crossing CPU/GPU byte-compatibility smoke test, honest transfer-inclusive microbenchmarks, Plonky3 provenance, and dated benchmark documentation. Organize the Rust backend under src/cuda. The CUDA configuration currently supports Linux x86_64, cap_height=0, and binary FRI. On an RTX PRO 6000 Blackwell, Ix Vector.extract_append recursive q50 reduces the post-execution-and-trace portion of inner and outer proving from 71.87s on CPU to 8.74s with CUDA (8.22x), with identical proof sizes and successful verification. End-to-end proving also includes Aiur execution and trace construction, which remain outside this comparison. Fifty queries is a development benchmark parameter, not a production security recommendation. --- .github/workflows/ci.yml | 30 +- Cargo.lock | 2 + Cargo.toml | 6 + README.md | 29 + build.rs | 205 ++ cuda/README.md | 105 + cuda/kernels.cu | 3310 +++++++++++++++++++++++++ cuda/smoke.sh | 34 + docs/cuda-benchmarks.md | 41 + examples/cuda_blake3_bench.rs | 54 + examples/cuda_dft_bench.rs | 201 ++ examples/proof_compatibility.rs | 83 + src/config.rs | 137 +- src/cuda/mmcs.rs | 335 +++ src/cuda/mod.rs | 2965 ++++++++++++++++++++++ src/cuda/pcs.rs | 1360 ++++++++++ src/lib.rs | 5 + src/lookup.rs | 21 + src/prover.rs | 469 +++- src/test_circuits/baby_bear_config.rs | 10 +- src/types.rs | 649 ++++- src/verifier.rs | 47 +- 22 files changed, 9933 insertions(+), 165 deletions(-) create mode 100644 build.rs create mode 100644 cuda/README.md create mode 100644 cuda/kernels.cu create mode 100755 cuda/smoke.sh create mode 100644 docs/cuda-benchmarks.md create mode 100644 examples/cuda_blake3_bench.rs create mode 100644 examples/cuda_dft_bench.rs create mode 100644 examples/proof_compatibility.rs create mode 100644 src/cuda/mmcs.rs create mode 100644 src/cuda/mod.rs create mode 100644 src/cuda/pcs.rs diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 15cf08b..4f4b749 100644 --- a/.github/workflows/ci.yml +++ b/.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 diff --git a/Cargo.lock b/Cargo.lock index 7d11a6e..e15b793 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -432,6 +432,7 @@ version = "0.1.0" dependencies = [ "bincode", "criterion", + "itertools 0.14.0", "p3-air", "p3-baby-bear", "p3-blake3", @@ -441,6 +442,7 @@ dependencies = [ "p3-field", "p3-fri", "p3-goldilocks", + "p3-interpolation", "p3-keccak", "p3-matrix", "p3-maybe-rayon", diff --git a/Cargo.toml b/Cargo.toml index 7f50d65..a4373f3 100644 --- a/Cargo.toml +++ b/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] diff --git a/README.md b/README.md index 6f1b005..66d2482 100644 --- a/README.md +++ b/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 diff --git a/build.rs b/build.rs new file mode 100644 index 0000000..d3641c7 --- /dev/null +++ b/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 { + 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::>(); + 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 { + 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 { + 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::() + }) + .filter(|architecture| !architecture.is_empty()) + .collect::>(); + architectures.sort_unstable(); + architectures.dedup(); + (!architectures.is_empty()).then(|| architectures.join(",")) +} + +fn cuda_library_directories(nvcc: &std::ffi::OsStr) -> Vec { + 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")] +} diff --git a/cuda/README.md b/cuda/README.md new file mode 100644 index 0000000..8d203f1 --- /dev/null +++ b/cuda/README.md @@ -0,0 +1,105 @@ +# CUDA backend + +This directory contains multi-stark's first-party CUDA implementation of +Goldilocks arithmetic, DFTs, coset low-degree extensions, BLAKE3 Merkle +commitments, lookup construction, quotient evaluation, and FRI proving. It +does not use ICICLE or copy code from it. + +The `cuda` Cargo feature selects `CudaDft` for the production +`GoldilocksBlake3Config`. BabyBear tests remain on their CPU DFT. Without the +feature, `build.rs` exits before looking for nvcc and the normal crate remains +independent of CUDA. + +## Current architecture + +The backend preserves Plonky3's public PCS interfaces while keeping the hot +prover pipeline device-resident: + +1. Rust validates dimensions and builds exact P3-compatible twiddle tables. +2. Trace matrices are uploaded and transformed with fused radix-4/radix-8 + DIF kernels into resident coset LDEs. +3. First-party BLAKE3 kernels commit mixed-height matrices without copying + LDEs back to the host. +4. Lookup traces and quotient LDEs are constructed from resident commitments. +5. Batched openings, reductions, FRI folding, and Merkle authentication remain + resident; only protocol-visible openings and proofs return to Rust. + +Large traces can be retained in pinned host memory and uploaded in bounded +windows when keeping every source matrix resident would exceed device memory. +The spill policy reserves one quarter of reported device memory by default and +uses `cudaMemGetInfo` at runtime instead of assuming a particular GPU size. +This retains the exact CPU storage layout and proof bytes. + +The ordinary `TwoAdicSubgroupDft` implementation remains available for direct +host-matrix users and transfer-inclusive microbenchmarks. + +## Build configuration + +By default, the build emits native cubins for the common architectures among +`sm_80`, `sm_86`, `sm_89`, `sm_90`, `sm_100`, and `sm_120` that the installed +`nvcc` reports as supported. Override the comma-separated list for smaller +deployment artifacts: + +```sh +MULTI_STARK_CUDA_ARCHS=80,90 \ + cargo test --release --locked --features parallel,cuda +``` + +Useful environment variables: + +- `NVCC`: path to nvcc (default: `$CUDA_HOME/bin/nvcc`, then `nvcc` on `PATH`) +- `CUDA_HOME` or `CUDA_PATH`: toolkit root used to find nvcc and `libcudart` +- `MULTI_STARK_CUDA_ARCHS`: numeric compute capabilities, e.g. `80,90` +- `MULTI_STARK_CUDA_DEVICE`: device used by the prover and microbenchmarks + (default: `0`; `CUDA_VISIBLE_DEVICES` remapping is also honored) +- `MULTI_STARK_CUDA_MIN_FREE_BYTES`: device headroom retained by spilling + source traces (default: one quarter of total device memory) + +Only native cubins are emitted. This avoids a newer toolkit's PTX being chosen +by a slightly older driver instead of a compatible native cubin. + +The current native backend supports Linux x86_64, one selected device per +prover, Merkle caps of height 0, and binary FRI folding +(`max_log_arity <= 1`). Select a device with `MULTI_STARK_CUDA_DEVICE` or remap +devices with `CUDA_VISIBLE_DEVICES`. Unsupported cap/arity configurations fail +during configuration construction instead of producing malformed commitments. + +## First GPU run + +From the repository root: + +```sh +./cuda/smoke.sh +``` + +The script prints the git/toolchain/GPU environment, then runs tests in the +failure-localizing order: field arithmetic, DFT, coset LDE, and finally the +entire release suite. It then generates a deterministic wide proof on CPU and +CUDA, crossing the resident-FRI, GPU-MMCS, and parallel-LDE thresholds, and +compares the complete serialized files. Together with the pinned BLAKE3/Merkle +vectors and small-proof digest, this demonstrates production-path protocol +compatibility rather than only transform round trips. + +After correctness passes, run the transfer-inclusive microbenchmark: + +```sh +cargo run --release --locked --features parallel,cuda --example cuda_dft_bench +``` + +It emits CSV to stdout. Shape/iteration controls are documented at the top of +`examples/cuda_dft_bench.rs`. The full proving Criterion benchmark and dated +Ix measurements are documented in [`docs/cuda-benchmarks.md`](../docs/cuda-benchmarks.md). + +## Safety and protocol invariants + +- Goldilocks values cross the ABI as `u64`; P3 declares the type + `repr(transparent)` and permits every `u64` bit pattern. +- CUDA outputs are canonical field representatives. +- Rust checks power-of-two sizes, two-adicity, integer overflow, and buffer + lengths before each synchronous FFI call. +- The CUDA source is licensed under the repository's MIT/Apache-2.0 terms and + depends only on the CUDA runtime/toolkit when enabled. +- CUDA affects prover execution only. Fields, BLAKE3 hashing, transcripts, + proof format, and the CPU verifier are unchanged. Canonical representation + changed newly generated proof bytes relative to pre-CUDA revisions as + documented in the top-level changelog. diff --git a/cuda/kernels.cu b/cuda/kernels.cu new file mode 100644 index 0000000..3c3f13d --- /dev/null +++ b/cuda/kernels.cu @@ -0,0 +1,3310 @@ +// SPDX-License-Identifier: MIT OR Apache-2.0 +// +// First-party CUDA kernels for Goldilocks arithmetic and batched radix-2 +// transforms. The ABI accepts host buffers; device residency is intentionally +// deferred to the later PCS/FRI backend. + +#include + +#include +#include +#include +#include + +namespace { + +constexpr uint64_t GOLDILOCKS_P = 0xffffffff00000001ULL; +constexpr uint64_t GOLDILOCKS_EPSILON = 0x00000000ffffffffULL; +constexpr unsigned int THREADS = 256; +constexpr unsigned int MAX_BLOCKS = 65535; +constexpr uint32_t BLAKE3_CHUNK_START = 1U << 0; +constexpr uint32_t BLAKE3_CHUNK_END = 1U << 1; +constexpr uint32_t BLAKE3_PARENT = 1U << 2; +constexpr uint32_t BLAKE3_ROOT = 1U << 3; +// All work currently uses CUDA's per-thread default stream. Stream-ordered +// allocation preserves the same dependency order while avoiding cudaFree's +// device-wide synchronization on thousands of short-lived PCS buffers. +inline cudaError_t persistent_malloc(void** pointer,size_t bytes){return cudaMalloc(pointer,bytes);} +inline cudaError_t persistent_free(void* pointer){return pointer?cudaFree(pointer):cudaSuccess;} +inline cudaError_t stream_malloc(void** pointer,size_t bytes){return cudaMallocAsync(pointer,bytes,cudaStreamPerThread);} +inline cudaError_t stream_free(void* pointer){return pointer?cudaFreeAsync(pointer,cudaStreamPerThread):cudaSuccess;} +#define cudaMalloc stream_malloc +#define cudaFree stream_free + +struct ConstantCacheEntry{int device=-1;size_t count=0;uint64_t kind=0,key0=0,key1=0;uint64_t* values=nullptr;ConstantCacheEntry* next=nullptr;}; +ConstantCacheEntry* constant_cache=nullptr;volatile int constant_cache_lock=0; +cudaError_t cached_device_constants(int device,const uint64_t* host,size_t count, + uint64_t kind,uint64_t key0,uint64_t key1,const uint64_t** output){ + while(__sync_lock_test_and_set(&constant_cache_lock,1)){} + for(auto* entry=constant_cache;entry;entry=entry->next)if(entry->device==device&& + entry->count==count&&entry->kind==kind&&entry->key0==key0&&entry->key1==key1){ + *output=entry->values;__sync_lock_release(&constant_cache_lock);return cudaSuccess;} + auto* entry=new(std::nothrow) ConstantCacheEntry; + if(!entry){__sync_lock_release(&constant_cache_lock);return cudaErrorMemoryAllocation;} + entry->device=device;entry->count=count;entry->kind=kind;entry->key0=key0;entry->key1=key1; + cudaError_t status=persistent_malloc(reinterpret_cast(&entry->values),count*sizeof(uint64_t)); + if(status==cudaSuccess)status=cudaMemcpy(entry->values,host,count*sizeof(uint64_t),cudaMemcpyHostToDevice); + if(status==cudaSuccess){*output=entry->values;entry->next=constant_cache;constant_cache=entry;} + else{persistent_free(entry->values);delete entry;} + __sync_lock_release(&constant_cache_lock);return status; +} +__device__ __constant__ uint32_t BLAKE3_IV[8] = { + 0x6A09E667U, 0xBB67AE85U, 0x3C6EF372U, 0xA54FF53AU, + 0x510E527FU, 0x9B05688CU, 0x1F83D9ABU, 0x5BE0CD19U, +}; +__device__ __constant__ unsigned int BLAKE3_PERMUTATION[16] = { + 2, 6, 3, 10, 7, 0, 4, 13, 1, 11, 12, 5, 9, 14, 15, 8, +}; + +__device__ __forceinline__ uint64_t canonicalize(uint64_t value) { + return value >= GOLDILOCKS_P ? value - GOLDILOCKS_P : value; +} + +__device__ __forceinline__ uint64_t goldilocks_add(uint64_t left, uint64_t right) { + left = canonicalize(left); + right = canonicalize(right); + // Written this way to avoid relying on overflow behavior in the source + // language. `GOLDILOCKS_P - right` is always representable. + const uint64_t gap = GOLDILOCKS_P - right; + return left >= gap ? left - gap : left + right; +} + +__device__ __forceinline__ uint64_t goldilocks_sub(uint64_t left, uint64_t right) { + left = canonicalize(left); + right = canonicalize(right); + return left >= right ? left - right : GOLDILOCKS_P - (right - left); +} + +__device__ __forceinline__ uint64_t goldilocks_mul(uint64_t left, uint64_t right) { + left = canonicalize(left); + right = canonicalize(right); + + const uint64_t low = left * right; + const uint64_t high = __umul64hi(left, right); + + // Reduce low + high * 2^64 using 2^64 = 2^32 - 1 (mod p). + const uint64_t high_high = high >> 32; + const uint64_t high_low = high & GOLDILOCKS_EPSILON; + uint64_t reduced_low = low - high_high; + if (low < high_high) { + // The wrapped subtraction added 2^64; replace that with +p. + reduced_low -= GOLDILOCKS_EPSILON; + } + const uint64_t reduced_high = high_low * GOLDILOCKS_EPSILON; + return goldilocks_add(reduced_low, reduced_high); +} + +__device__ __forceinline__ uint64_t goldilocks_pow(uint64_t base, uint64_t exponent) { + uint64_t result = 1; + while (exponent != 0) { + if ((exponent & 1U) != 0) { + result = goldilocks_mul(result, base); + } + base = goldilocks_mul(base, base); + exponent >>= 1; + } + return result; +} + +__device__ __forceinline__ uint32_t rotate_right(uint32_t value, + unsigned int count) { + return __funnelshift_r(value, value, count); +} + +__device__ __forceinline__ void blake3_g(uint32_t state[16], unsigned int a, + unsigned int b, unsigned int c, + unsigned int d, uint32_t x, + uint32_t y) { + state[a] = state[a] + state[b] + x; + state[d] = rotate_right(state[d] ^ state[a], 16); + state[c] += state[d]; + state[b] = rotate_right(state[b] ^ state[c], 12); + state[a] = state[a] + state[b] + y; + state[d] = rotate_right(state[d] ^ state[a], 8); + state[c] += state[d]; + state[b] = rotate_right(state[b] ^ state[c], 7); +} + +__device__ __forceinline__ void blake3_compress( + const uint32_t chaining_value[8], const uint32_t block[16], + uint64_t counter, uint32_t block_length, uint32_t flags, + uint32_t output[16]) { + uint32_t state[16]; + uint32_t message[16]; +#pragma unroll + for (unsigned int i = 0; i < 8; ++i) { + state[i] = chaining_value[i]; + state[i + 8] = BLAKE3_IV[i]; + } + state[12] = static_cast(counter); + state[13] = static_cast(counter >> 32); + state[14] = block_length; + state[15] = flags; +#pragma unroll + for (unsigned int i = 0; i < 16; ++i) { + message[i] = block[i]; + } + +#pragma unroll + for (unsigned int round = 0; round < 7; ++round) { + blake3_g(state, 0, 4, 8, 12, message[0], message[1]); + blake3_g(state, 1, 5, 9, 13, message[2], message[3]); + blake3_g(state, 2, 6, 10, 14, message[4], message[5]); + blake3_g(state, 3, 7, 11, 15, message[6], message[7]); + blake3_g(state, 0, 5, 10, 15, message[8], message[9]); + blake3_g(state, 1, 6, 11, 12, message[10], message[11]); + blake3_g(state, 2, 7, 8, 13, message[12], message[13]); + blake3_g(state, 3, 4, 9, 14, message[14], message[15]); + if (round != 6) { + uint32_t permuted[16]; +#pragma unroll + for (unsigned int i = 0; i < 16; ++i) { + permuted[i] = message[BLAKE3_PERMUTATION[i]]; + } +#pragma unroll + for (unsigned int i = 0; i < 16; ++i) { + message[i] = permuted[i]; + } + } + } + +#pragma unroll + for (unsigned int i = 0; i < 8; ++i) { + output[i] = state[i] ^ state[i + 8]; + output[i + 8] = state[i + 8] ^ chaining_value[i]; + } +} + +__device__ __forceinline__ uint32_t load_u32_le(const uint8_t* bytes, + size_t available, + size_t offset) { + uint32_t result = 0; +#pragma unroll + for (unsigned int byte = 0; byte < 4; ++byte) { + if (offset + byte < available) { + result |= static_cast(bytes[offset + byte]) << (8 * byte); + } + } + return result; +} + +__device__ __forceinline__ void blake3_chunk( + const uint8_t* bytes, size_t length, uint64_t chunk_counter, bool root, + uint32_t output[8]) { + uint32_t chaining_value[8]; +#pragma unroll + for (unsigned int i = 0; i < 8; ++i) { + chaining_value[i] = BLAKE3_IV[i]; + } + const size_t blocks = length == 0 ? 1 : (length + 63) / 64; + for (size_t block_index = 0; block_index < blocks; ++block_index) { + const size_t block_offset = block_index * 64; + const size_t block_length = + block_offset < length + ? ((length - block_offset) < 64 ? length - block_offset : 64) + : 0; + uint32_t block[16]; +#pragma unroll + for (unsigned int word = 0; word < 16; ++word) { + block[word] = load_u32_le(bytes + block_offset, block_length, + static_cast(word) * 4); + } + uint32_t flags = block_index == 0 ? BLAKE3_CHUNK_START : 0; + const bool last = block_index + 1 == blocks; + if (last) { + flags |= BLAKE3_CHUNK_END; + if (root) { + flags |= BLAKE3_ROOT; + } + } + uint32_t compressed[16]; + blake3_compress(chaining_value, block, chunk_counter, + static_cast(block_length), flags, compressed); + if (last) { +#pragma unroll + for (unsigned int i = 0; i < 8; ++i) { + output[i] = compressed[i]; + } + } else { +#pragma unroll + for (unsigned int i = 0; i < 8; ++i) { + chaining_value[i] = compressed[i]; + } + } + } +} + +__device__ __forceinline__ void blake3_parent(const uint32_t left[8], + const uint32_t right[8], + bool root, uint32_t output[8]) { + uint32_t block[16]; +#pragma unroll + for (unsigned int i = 0; i < 8; ++i) { + block[i] = left[i]; + block[i + 8] = right[i]; + } + uint32_t compressed[16]; + blake3_compress(BLAKE3_IV, block, 0, 64, + BLAKE3_PARENT | (root ? BLAKE3_ROOT : 0), compressed); +#pragma unroll + for (unsigned int i = 0; i < 8; ++i) { + output[i] = compressed[i]; + } +} + +__device__ __forceinline__ void blake3_hash_digest_pair( + const uint8_t* left, const uint8_t* right, uint8_t* digest) { + uint32_t block[16]; +#pragma unroll + for (unsigned int word = 0; word < 8; ++word) { + block[word] = load_u32_le(left, 32, static_cast(word) * 4); + block[word + 8] = + load_u32_le(right, 32, static_cast(word) * 4); + } + uint32_t compressed[16]; + blake3_compress(BLAKE3_IV, block, 0, 64, + BLAKE3_CHUNK_START | BLAKE3_CHUNK_END | BLAKE3_ROOT, + compressed); +#pragma unroll + for (unsigned int word = 0; word < 8; ++word) { + const uint32_t value = compressed[word]; + digest[word * 4] = static_cast(value); + digest[word * 4 + 1] = static_cast(value >> 8); + digest[word * 4 + 2] = static_cast(value >> 16); + digest[word * 4 + 3] = static_cast(value >> 24); + } +} + +__device__ __forceinline__ size_t reverse_index_bits(size_t index, unsigned int bits) { + if (bits == 0) { + return 0; + } + return static_cast(__brevll(static_cast(index)) >> (64 - bits)); +} + +unsigned int blocks_for(size_t work_items) { + if (work_items == 0) { + return 0; + } + const size_t required = (work_items + THREADS - 1) / THREADS; + return static_cast(required < MAX_BLOCKS ? required : MAX_BLOCKS); +} + +bool is_power_of_two(size_t value) { + return value != 0 && (value & (value - 1)) == 0; +} + +bool product_fits(size_t left, size_t right) { + return left == 0 || right <= SIZE_MAX / left; +} + +unsigned int strict_log2(size_t value) { + unsigned int result = 0; + while (value > 1) { + value >>= 1; + ++result; + } + return result; +} + +class DeviceBuffer { + public: + DeviceBuffer() = default; + DeviceBuffer(const DeviceBuffer&) = delete; + DeviceBuffer& operator=(const DeviceBuffer&) = delete; + + ~DeviceBuffer() { + if (pointer_ != nullptr) { + cudaFree(pointer_); + } + } + + cudaError_t allocate(size_t elements) { + if (!product_fits(elements, sizeof(uint64_t))) { + return cudaErrorInvalidValue; + } + return cudaMalloc(reinterpret_cast(&pointer_), elements * sizeof(uint64_t)); + } + + uint64_t* get() { return pointer_; } + const uint64_t* get() const { return pointer_; } + + private: + uint64_t* pointer_ = nullptr; +}; + +// cudaMemcpy from ordinary Vec storage uses an internal pageable-memory +// staging buffer. Large LDEs are dominated by that staging copy, so pin the +// caller's matrix for the duration of the synchronous operation. Registration +// is deliberately best-effort: constrained hosts can reject pinning without +// making the CUDA backend unusable, in which case cudaMemcpy retains its +// normal pageable fallback. +class HostRegistration { + public: + HostRegistration(void* pointer, size_t bytes) : pointer_(pointer) { + if (pointer != nullptr && bytes != 0) { + registered_ = cudaHostRegister(pointer, bytes, cudaHostRegisterDefault) == + cudaSuccess; + if (!registered_) { + // Do not let an optional registration failure poison the + // thread-local CUDA error observed by the real operation. + cudaGetLastError(); + } + } + } + + HostRegistration(const HostRegistration&) = delete; + HostRegistration& operator=(const HostRegistration&) = delete; + + ~HostRegistration() { + if (registered_) { + cudaHostUnregister(pointer_); + } + } + + private: + void* pointer_ = nullptr; + bool registered_ = false; +}; + +struct ResidentMerkleTree { + uint8_t* rows = nullptr; + uint8_t* digests = nullptr; + size_t row_bytes = 0; + size_t row_count = 0; + + ~ResidentMerkleTree() { + if (digests != nullptr) { + cudaFree(digests); + } + if (rows != nullptr) { + cudaFree(rows); + } + } +}; + +struct ResidentMixedMerkleTree { + uint8_t* digests = nullptr; + size_t row_count = 0; + + ~ResidentMixedMerkleTree() { + if (digests != nullptr) { + cudaFree(digests); + } + } +}; + +struct ResidentLde { + uint64_t* values = nullptr; + // Original row-major evaluations, retained for prover stages (notably + // lookup construction) which consume the witness after its LDE has been + // committed. Null for LDEs synthesized directly on the device. + uint64_t* trace_values = nullptr; + const uint64_t* host_trace_values = nullptr; + bool host_trace_registered = false; + size_t trace_height = 0; + size_t height = 0; + size_t width = 0; + uint8_t* interpolation_scratch = nullptr; + size_t interpolation_scratch_bytes = 0; + + ~ResidentLde() { + if (values != nullptr) { + cudaFree(values); + } + if (trace_values != nullptr) cudaFree(trace_values); + if (host_trace_registered) cudaHostUnregister( + const_cast(host_trace_values)); + if (interpolation_scratch != nullptr) cudaFree(interpolation_scratch); + } +}; + +// Kernels read ResidentLde metadata directly. Keep that small control block in +// CUDA managed memory so device access does not depend on Linux HMM support for +// arbitrary pageable host allocations. +cudaError_t create_resident_lde(ResidentLde** output) { + if (output == nullptr) return cudaErrorInvalidValue; + *output = nullptr; + void* storage = nullptr; + const cudaError_t status = cudaMallocManaged(&storage, sizeof(ResidentLde)); + if (status != cudaSuccess) return status; + *output = new (storage) ResidentLde; + return cudaSuccess; +} + +cudaError_t destroy_resident_lde(ResidentLde* lde) { + if (lde == nullptr) return cudaSuccess; + lde->~ResidentLde(); + return cudaFree(lde); +} + +// Stable C ABI instruction used by Rust's compiled constraint DAG. All +// expressions are base-field-only by this stage. +struct ConstraintNode { + uint64_t value; + uint32_t a; + uint32_t b; + uint32_t op; + uint32_t aux; + uint32_t out; +}; + +struct ConstraintLookup { + uint32_t multiplicity; + uint32_t arg_start; + uint32_t arg_count; + uint32_t emit_after; + uint32_t output; +}; + +struct Ext2 { uint64_t c0; uint64_t c1; }; + +// The resident PCS canonicalizes every committed LDE, and interpolation +// tables are serialized with `as_canonical_u64`. Restrict the faster +// canonical-input arithmetic to this boundary instead of weakening the +// representation guarantees of lookup/quotient code, whose host inputs may +// legitimately use lazy representatives. +__device__ __forceinline__ uint64_t canonical_add(uint64_t a,uint64_t b){ + const uint64_t sum=a+b; + if(sum=GOLDILOCKS_P?sum-GOLDILOCKS_P:sum; +} +__device__ __forceinline__ uint64_t canonical_mul(uint64_t a,uint64_t b){ + const uint64_t low=a*b,high=__umul64hi(a,b),high_high=high>>32; + uint64_t reduced_low=low-high_high; + if(lowvalues[row*lde->width+column]),goldilocks_mul(weight.c1,lde->values[row*lde->width+column])}; + sum=ext2_add(sum,term); + } sums[threadIdx.x]=sum;__syncthreads(); + for(unsigned int stride=blockDim.x/2;stride;stride>>=1){if(threadIdx.x(blockIdx.x)*INTERPOLATION_COLUMNS+threadIdx.x; + const size_t row_begin=static_cast(blockIdx.y)*INTERPOLATION_ROWS; + const size_t row_end=row_begin+INTERPOLATION_ROWSwidth){ + for(size_t row=row_begin+threadIdx.y;rowvalues[row*lde->width+column]; + sum.c0=canonical_add(sum.c0,canonical_mul(weight.c0,value)); + sum.c1=canonical_add(sum.c1,canonical_mul(weight.c1,value)); + } + } + sums[threadIdx.y][threadIdx.x]=sum;__syncthreads(); + for(unsigned int stride=INTERPOLATION_LANES/2;stride;stride>>=1){ + if(threadIdx.ywidth) + partials[static_cast(blockIdx.y)*lde->width+column]=sums[0][threadIdx.x]; +} + +__global__ void interpolate_lde_partials2(Ext2* partials0,Ext2* partials1, + const ResidentLde* lde,const Ext2* inv0,const Ext2* inv1, + const uint64_t* coset,size_t height,uint64_t ext_w){ + __shared__ Ext2 sums0[INTERPOLATION_LANES][INTERPOLATION_COLUMNS]; + __shared__ Ext2 sums1[INTERPOLATION_LANES][INTERPOLATION_COLUMNS]; + __shared__ Ext2 weights0[INTERPOLATION_ROWS]; + __shared__ Ext2 weights1[INTERPOLATION_ROWS]; + const size_t column=static_cast(blockIdx.x)*INTERPOLATION_COLUMNS+threadIdx.x; + const size_t row_begin=static_cast(blockIdx.y)*INTERPOLATION_ROWS; + const size_t row_end=row_begin+INTERPOLATION_ROWSwidth){ + for(size_t row=row_begin+threadIdx.y;rowvalues[row*lde->width+column]; + const Ext2 w0=weights0[row-row_begin],w1=weights1[row-row_begin]; + sum0.c0=canonical_add(sum0.c0,canonical_mul(w0.c0,value)); + sum0.c1=canonical_add(sum0.c1,canonical_mul(w0.c1,value)); + sum1.c0=canonical_add(sum1.c0,canonical_mul(w1.c0,value)); + sum1.c1=canonical_add(sum1.c1,canonical_mul(w1.c1,value)); + } + } + sums0[threadIdx.y][threadIdx.x]=sum0;sums1[threadIdx.y][threadIdx.x]=sum1; + __syncthreads(); + for(unsigned int stride=INTERPOLATION_LANES/2;stride;stride>>=1){ + if(threadIdx.ywidth){const size_t at=static_cast(blockIdx.y)*lde->width+column; + partials0[at]=sums0[0][threadIdx.x];partials1[at]=sums1[0][threadIdx.x];} +} + +__global__ void finish_lde_interpolation(Ext2* output,const Ext2* partials, + size_t partial_rows,size_t width,Ext2 scale,uint64_t ext_w){ + extern __shared__ Ext2 sums[];const size_t column=blockIdx.x;Ext2 sum{0,0}; + for(size_t row=threadIdx.x;row>=1){ + if(threadIdx.x>=1){ + if(threadIdx.y>=1){ + if(threadIdx.x(blockIdx.x)*blockDim.x+threadIdx.x;row(gridDim.x)*blockDim.x){ + Ext2 compressed{0,0};for(size_t column=0;columnwidth;++column){const uint64_t v=lde->values[row*lde->width+column];compressed.c0=goldilocks_add(compressed.c0,goldilocks_mul(alpha_powers[column].c0,v));compressed.c1=goldilocks_add(compressed.c1,goldilocks_mul(alpha_powers[column].c1,v));} + const Ext2 term=ext2_mul(alpha_offset,ext2_mul(ext2_sub(reduced_y,compressed),inv_denoms[row],ext_w),ext_w);output[row]=ext2_add(output[row],term); + } +} + +__global__ void fold_fri_ext2(Ext2* output,const Ext2* input,uint64_t* powers, + size_t height,Ext2 beta,uint64_t ext_w){ + for(size_t row=static_cast(blockIdx.x)*blockDim.x+threadIdx.x;row(gridDim.x)*blockDim.x){ + const Ext2 lo=input[2*row],hi=input[2*row+1]; + const Ext2 even={goldilocks_mul(goldilocks_add(lo.c0,hi.c0),0x7fffffff80000001ULL),goldilocks_mul(goldilocks_add(lo.c1,hi.c1),0x7fffffff80000001ULL)}; + Ext2 odd=ext2_mul(ext2_sub(lo,hi),beta,ext_w);odd.c0=goldilocks_mul(odd.c0,powers[row]);odd.c1=goldilocks_mul(odd.c1,powers[row]);output[row]=ext2_add(even,odd); + } +} +__global__ void init_fri_powers(uint64_t* powers,size_t height,uint64_t g_inv){ + const unsigned int bits=static_cast(__ffsll(height)-1); + for(size_t row=static_cast(blockIdx.x)*blockDim.x+threadIdx.x;row(gridDim.x)*blockDim.x){ + const uint64_t exponent=bits==0?0:__brevll(static_cast(row))>>(64U-bits); + powers[row]=goldilocks_mul(0x7fffffff80000001ULL,goldilocks_pow(g_inv,exponent)); + } +} +__global__ void add_scaled_ext2(Ext2* output,const Ext2* input,size_t count,Ext2 scale,uint64_t ext_w){ + for(size_t i=static_cast(blockIdx.x)*blockDim.x+threadIdx.x;i(gridDim.x)*blockDim.x)output[i]=ext2_add(output[i],ext2_mul(scale,input[i],ext_w)); +} + +__device__ __forceinline__ Ext2 ext2_inverse(Ext2 value,uint64_t ext_w){ + const uint64_t norm=goldilocks_sub(goldilocks_mul(value.c0,value.c0),goldilocks_mul(ext_w,goldilocks_mul(value.c1,value.c1))); + const uint64_t inv=goldilocks_pow(norm,GOLDILOCKS_P-2); + return {goldilocks_mul(value.c0,inv),goldilocks_sub(0,goldilocks_mul(value.c1,inv))}; +} +__global__ void lookup_messages(Ext2* conjugates,uint64_t* norms,const uint64_t* args, + const size_t* arg_offsets,size_t height,size_t num_lookups,size_t args_width, + Ext2 beta,Ext2 gamma,uint64_t ext_w){ + const size_t count=height*num_lookups; + for(size_t index=static_cast(blockIdx.x)*blockDim.x+threadIdx.x;index(gridDim.x)*blockDim.x){ + const size_t row=index/num_lookups,lookup=index%num_lookups;Ext2 fingerprint{0,0}; + for(size_t j=arg_offsets[lookup+1];j>arg_offsets[lookup];--j)fingerprint=ext2_add(ext2_mul(fingerprint,gamma,ext_w),{args[row*args_width+j-1],0}); + const Ext2 message=ext2_add(beta,fingerprint);conjugates[index]={message.c0,goldilocks_sub(0,message.c1)}; + norms[index]=goldilocks_sub(goldilocks_mul(message.c0,message.c0),goldilocks_mul(ext_w,goldilocks_mul(message.c1,message.c1))); + } +} +// Evaluate the lookup-only prefix of a circuit DAG against the original +// row-major witness retained by the stage-1 LDE. This avoids materializing +// and uploading the much wider concrete LookupValues buffer. +__global__ void lookup_messages_graph( + Ext2* conjugates,uint64_t* norms,uint64_t* multiplicities, + const ConstraintNode* nodes,size_t node_count,size_t slot_count, + const ConstraintLookup* lookups,size_t lookup_count,const uint32_t* lookup_args, + const ResidentLde* preprocessed,const uint64_t* main_trace,size_t main_width, + bool main_trace_is_chunk, + Ext2 beta,Ext2 gamma,uint64_t ext_w,size_t height,size_t row_start, + size_t row_count,uint64_t* global_scratch){ + extern __shared__ uint64_t shared_values[]; + const size_t lane=threadIdx.x,tile=blockDim.x; + uint64_t* values=global_scratch?global_scratch+static_cast(blockIdx.x)*slot_count*tile:shared_values; + for(size_t local_row=static_cast(blockIdx.x)*tile+lane;local_row(gridDim.x)*tile){ + const size_t row=row_start+local_row; + const size_t next=(row+1)&(height-1); + size_t next_lookup=0; + for(size_t i=0;i=6?values[n.a*tile+lane]:0; + const uint64_t b=n.op>=6&&n.op!=9?values[n.b*tile+lane]:0;uint64_t v=0; + switch(n.op){ + case 0:v=n.value;break; + case 1:{const bool prep_column=n.aux<2; + const size_t r=prep_column?((n.aux&1U)?next:row): + (main_trace_is_chunk?((n.aux&1U)?local_row+1:local_row):((n.aux&1U)?next:row)); + v=n.aux<2?preprocessed->trace_values[r*preprocessed->width+n.a]:main_trace[r*main_width+n.a];break;} + case 2:v=0;break; + case 3:v=static_cast(row==0);break; + case 4:v=static_cast(row+1==height);break; + case 5:v=static_cast(row+1!=height);break; + case 6:v=goldilocks_add(a,b);break; + case 7:v=goldilocks_sub(a,b);break; + case 8:v=goldilocks_mul(a,b);break; + case 9:v=a==0?0:GOLDILOCKS_P-canonicalize(a);break; + }values[n.out*tile+lane]=v; + while(next_lookup0;--k)fingerprint=ext2_add(ext2_mul(fingerprint,gamma,ext_w),{values[lookup_args[l.arg_start+k-1]*tile+lane],0}); + const Ext2 message=ext2_add(beta,fingerprint);const size_t index=local_row*lookup_count+l.output; + conjugates[index]={message.c0,goldilocks_sub(0,message.c1)}; + norms[index]=goldilocks_sub(goldilocks_mul(message.c0,message.c0),goldilocks_mul(ext_w,goldilocks_mul(message.c1,message.c1))); + multiplicities[index]=values[l.multiplicity*tile+lane]; + } + } + } +} +__global__ void lookup_group_deltas_batched(Ext2* output,const uint64_t* multiplicities, + const Ext2* conjugates,const uint64_t* norm_inverses, + size_t height,size_t num_lookups,size_t group_size,uint64_t ext_w){ + const size_t slots=(num_lookups+group_size-1)/group_size,total=height*slots; + for(size_t index=static_cast(blockIdx.x)*blockDim.x+threadIdx.x;index(gridDim.x)*blockDim.x){ + const size_t row=index/slots,slot=index%slots,begin=slot*group_size,end=(begin+group_size(blockIdx.x)*blockDim.x+threadIdx.x; + products[threadIdx.x]=index>=1){if(threadIdx.x(blockIdx.x)*blockDim.x+threadIdx.x;const uint64_t own=index=stride)prior=products[threadIdx.x-stride];__syncthreads();if(threadIdx.x>=stride)products[threadIdx.x]=goldilocks_mul(products[threadIdx.x],prior);__syncthreads();} + const uint64_t prefix=threadIdx.x?products[threadIdx.x-1]:1; + const uint64_t block_product=products[blockDim.x-1];__syncthreads(); + const uint64_t total_inverse=block_inverses?block_inverses[blockIdx.x]:goldilocks_pow(block_product,GOLDILOCKS_P-2); + products[threadIdx.x]=own;__syncthreads();for(unsigned int stride=1;stride>>(output,input,nullptr,count);return cudaGetLastError();} + if(status==cudaSuccess)status=cudaMalloc(reinterpret_cast(&products),block_count*sizeof(uint64_t)); + if(status==cudaSuccess)status=cudaMalloc(reinterpret_cast(&inverses),block_count*sizeof(uint64_t)); + if(status==cudaSuccess){norm_block_products<<(block_count),THREADS,THREADS*sizeof(uint64_t)>>>(products,input,count);status=cudaGetLastError();} + if(status==cudaSuccess){batch_inverse_norm_blocks<<((block_count+THREADS-1)/THREADS),THREADS,THREADS*sizeof(uint64_t)>>>(inverses,products,nullptr,block_count);status=cudaGetLastError();} + if(status==cudaSuccess){batch_inverse_norm_blocks<<(block_count),THREADS,THREADS*sizeof(uint64_t)>>>(output,input,inverses,count);status=cudaGetLastError();} + cudaFree(inverses);cudaFree(products);return status; +} + +__global__ void initialize_coset_selector_denominators( + uint64_t* denominators, size_t quotient_size, size_t next_step, + uint64_t coset_shift, uint64_t coset_generator, uint64_t trace_last, + uint64_t vanishing_start, uint64_t vanishing_step) { + constexpr size_t ROWS_PER_WORKER = 16; + const size_t worker_count = + (quotient_size + ROWS_PER_WORKER - 1) / ROWS_PER_WORKER; + for (size_t worker = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + worker < worker_count; + worker += static_cast(gridDim.x) * blockDim.x) { + const size_t first_row = worker * ROWS_PER_WORKER; + uint64_t x = goldilocks_mul( + coset_shift, goldilocks_pow(coset_generator, first_row)); + uint64_t vanishing_power = goldilocks_pow( + vanishing_step, first_row & (next_step - 1)); + const size_t last_row = + first_row + ROWS_PER_WORKER < quotient_size + ? first_row + ROWS_PER_WORKER + : quotient_size; + for (size_t row = first_row; row < last_row; ++row) { + const uint64_t vanishing = goldilocks_sub( + goldilocks_mul(vanishing_start, vanishing_power), 1); + denominators[row] = goldilocks_sub(x, 1); + denominators[quotient_size + row] = + goldilocks_sub(x, trace_last); + denominators[2 * quotient_size + row] = + goldilocks_sub(x, trace_last); + denominators[3 * quotient_size + row] = vanishing; + x = goldilocks_mul(x, coset_generator); + vanishing_power = goldilocks_mul(vanishing_power, vanishing_step); + } + } +} + +__global__ void finish_coset_selectors( + uint64_t* selectors, size_t quotient_size) { + for (size_t row = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + row < quotient_size; + row += static_cast(gridDim.x) * blockDim.x) { + const uint64_t vanishing = selectors[3 * quotient_size + row]; + selectors[row] = goldilocks_mul(vanishing, selectors[row]); + selectors[quotient_size + row] = goldilocks_mul( + vanishing, selectors[quotient_size + row]); + } +} + +cudaError_t generate_coset_selectors( + uint64_t* selectors, size_t quotient_size, size_t next_step, + uint64_t coset_shift, uint64_t coset_generator, + uint64_t trace_last, uint64_t vanishing_start, uint64_t vanishing_step) { + constexpr size_t ROWS_PER_WORKER = 16; + const size_t worker_count = + (quotient_size + ROWS_PER_WORKER - 1) / ROWS_PER_WORKER; + initialize_coset_selector_denominators<<>>( + selectors, quotient_size, next_step, coset_shift, coset_generator, + trace_last, vanishing_start, vanishing_step); + cudaError_t status = cudaGetLastError(); + if (status == cudaSuccess) { + status = batch_inverse_norms(selectors, selectors, 2 * quotient_size); + } + if (status == cudaSuccess) { + finish_coset_selectors<<>>( + selectors, quotient_size); + status = cudaGetLastError(); + } + if (status == cudaSuccess) { + status = batch_inverse_norms( + selectors + 3 * quotient_size, + selectors + 3 * quotient_size, + quotient_size); + } + return status; +} + +extern "C" int multi_stark_cuda_generate_coset_selectors( + int device_id, uint64_t* output, size_t quotient_size, size_t next_step, + uint64_t coset_shift, uint64_t coset_generator, uint64_t trace_last, + uint64_t vanishing_start, uint64_t vanishing_step) { + if (output == nullptr || !is_power_of_two(quotient_size) || + !is_power_of_two(next_step) || next_step > quotient_size) { + return static_cast(cudaErrorInvalidValue); + } + cudaError_t status = cudaSetDevice(device_id); + uint64_t* selectors = nullptr; + if (status == cudaSuccess) + status = cudaMalloc(reinterpret_cast(&selectors), + 4 * quotient_size * sizeof(uint64_t)); + if (status == cudaSuccess) + status = generate_coset_selectors( + selectors, quotient_size, next_step, coset_shift, + coset_generator, trace_last, vanishing_start, vanishing_step); + if (status == cudaSuccess) + status = cudaMemcpy(output, selectors, + 4 * quotient_size * sizeof(uint64_t), + cudaMemcpyDeviceToHost); + cudaFree(selectors); + return static_cast(status); +} +__global__ void denominator_norms(uint64_t* norms,const uint64_t* coset, + size_t count,Ext2 point,uint64_t ext_w){ + for(size_t i=static_cast(blockIdx.x)*blockDim.x+threadIdx.x;i(gridDim.x)*blockDim.x){ + const uint64_t real=goldilocks_sub(point.c0,coset[i]); + norms[i]=goldilocks_sub(goldilocks_mul(real,real), + goldilocks_mul(ext_w,goldilocks_mul(point.c1,point.c1))); + } +} +__global__ void finish_inverse_denominators(Ext2* output,const uint64_t* inverses, + const uint64_t* coset,size_t count,Ext2 point){ + for(size_t i=static_cast(blockIdx.x)*blockDim.x+threadIdx.x;i(gridDim.x)*blockDim.x){ + const uint64_t inverse=inverses[i]; + output[i]={goldilocks_mul(goldilocks_sub(point.c0,coset[i]),inverse), + goldilocks_mul(goldilocks_sub(0,point.c1),inverse)}; + } +} +__global__ void lookup_group_deltas(Ext2* output,const uint64_t* multiplicities, + const uint64_t* args,const size_t* arg_offsets,size_t height,size_t num_lookups, + size_t args_width,size_t group_size,Ext2 beta,Ext2 gamma,uint64_t ext_w){ + const size_t slots=(num_lookups+group_size-1)/group_size,total=height*slots; + for(size_t index=static_cast(blockIdx.x)*blockDim.x+threadIdx.x;index(gridDim.x)*blockDim.x){ + const size_t row=index/slots,slot=index%slots,begin=slot*group_size,end=(begin+group_sizearg_offsets[lookup];--j)fingerprint=ext2_add(ext2_mul(fingerprint,gamma,ext_w),{args[row*args_width+j-1],0}); + const Ext2 inverse=ext2_inverse(ext2_add(beta,fingerprint),ext_w);const uint64_t multiplicity=multiplicities[row*num_lookups+lookup]; + delta.c0=goldilocks_add(delta.c0,goldilocks_mul(multiplicity,inverse.c0));delta.c1=goldilocks_add(delta.c1,goldilocks_mul(multiplicity,inverse.c1)); + }output[index]=delta; + } +} +__global__ void reverse_u64(uint64_t* output,const uint64_t* input,size_t count){const size_t i=static_cast(blockIdx.x)*blockDim.x+threadIdx.x;if(i(blockIdx.x)*blockDim.x+threadIdx.x; + const Ext2 own=index=stride)prior=scan_values[threadIdx.x-stride];__syncthreads();if(threadIdx.x>=stride)scan_values[threadIdx.x]=ext2_add(scan_values[threadIdx.x],prior);__syncthreads();} + if(index(blockIdx.x)*blockDim.x+threadIdx.x;if(index(&sums),blocks*sizeof(Ext2)); + if(status==cudaSuccess){scan_ext2_blocks<<(blocks),THREADS,THREADS*sizeof(Ext2)>>>(output,sums,input,count);status=cudaGetLastError();} + if(status==cudaSuccess&&blocks>1)status=cudaMalloc(reinterpret_cast(&offsets),blocks*sizeof(Ext2)); + if(status==cudaSuccess&&blocks>1)status=exclusive_scan_ext2(offsets,sums,blocks); + if(status==cudaSuccess&&blocks>1){add_scan_block_offsets<<(blocks),THREADS>>>(output,offsets,count);status=cudaGetLastError();} + cudaFree(offsets);cudaFree(sums);return status; +} +__global__ void scan_mul_blocks(uint64_t* output,uint64_t* sums,const uint64_t* input,size_t count){ + extern __shared__ uint64_t scan_products[];const size_t index=static_cast(blockIdx.x)*blockDim.x+threadIdx.x; + const uint64_t own=index=stride)prior=scan_products[threadIdx.x-stride];__syncthreads();if(threadIdx.x>=stride)scan_products[threadIdx.x]=goldilocks_mul(scan_products[threadIdx.x],prior);__syncthreads();} + if(index(blockIdx.x)*blockDim.x+threadIdx.x;if(i(&sums),blocks*sizeof(uint64_t)); + if(status==cudaSuccess){scan_mul_blocks<<(blocks),THREADS,THREADS*sizeof(uint64_t)>>>(output,sums,input,count);status=cudaGetLastError();} + if(status==cudaSuccess&&blocks>1)status=cudaMalloc(reinterpret_cast(&offsets),blocks*sizeof(uint64_t));if(status==cudaSuccess&&blocks>1)status=exclusive_scan_mul(offsets,sums,blocks); + if(status==cudaSuccess&&blocks>1){mul_scan_block_offsets<<(blocks),THREADS>>>(output,offsets,count);status=cudaGetLastError();}cudaFree(offsets);cudaFree(sums);return status; +} + +__device__ __forceinline__ size_t reverse_low_bits(size_t value, + unsigned int bits) { + return static_cast(__brevll(static_cast(value)) >> + (64U - bits)); +} + +__global__ void evaluate_constraint_graph( + uint64_t* output, const ConstraintNode* nodes, size_t node_count, + const uint32_t* roots, size_t root_count, const ResidentLde* preprocessed, + const ResidentLde* main, const ResidentLde* stage2, + const uint64_t* publics, const uint64_t* selectors, size_t quotient_size, + size_t next_step) { + extern __shared__ uint64_t values[]; + const size_t lane = threadIdx.x; + const size_t tile = blockDim.x; + const unsigned int log_size = static_cast(__ffsll(quotient_size) - 1); + for (size_t row = static_cast(blockIdx.x) * tile + lane; + row < quotient_size; row += static_cast(gridDim.x) * tile) { + const size_t storage_row = reverse_low_bits(row, log_size); + const size_t next_storage_row = + reverse_low_bits((row + next_step) & (quotient_size - 1), log_size); + for (size_t i = 0; i < node_count; ++i) { + const ConstraintNode node = nodes[i]; + uint64_t result = 0; + const uint64_t left = node.a < i ? values[node.a * tile + lane] : 0; + const uint64_t right = node.b < i ? values[node.b * tile + lane] : 0; + switch (node.op) { + case 0: result = node.value; break; + case 1: { + const ResidentLde* matrix = node.aux < 2 ? preprocessed + : node.aux < 4 ? main : stage2; + const bool next = (node.aux & 1U) != 0; + const size_t r = next ? next_storage_row : storage_row; + result = matrix->values[r * matrix->width + node.a]; + break; + } + case 2: result = publics[node.a]; break; + case 3: result = selectors[row]; break; + case 4: result = selectors[quotient_size + row]; break; + case 5: result = selectors[2 * quotient_size + row]; break; + case 6: result = goldilocks_add(left, right); break; + case 7: result = goldilocks_sub(left, right); break; + case 8: result = goldilocks_mul(left, right); break; + case 9: result = left == 0 ? 0 : GOLDILOCKS_P - canonicalize(left); break; + } + values[node.out * tile + lane] = result; + } + for (size_t root = 0; root < root_count; ++root) { + output[row * root_count + root] = values[roots[root] * tile + lane]; + } + } +} + +__global__ void evaluate_quotient( + uint64_t* output, const ConstraintNode* nodes, size_t node_count, size_t slot_count, + const uint32_t* roots, size_t root_count, const ConstraintLookup* lookups, + size_t lookup_count, const uint32_t* lookup_args, size_t group_size, + const ResidentLde* preprocessed, const ResidentLde* main, + const ResidentLde* stage2, const uint64_t* publics, + const uint64_t* selectors, const uint64_t* alpha, const uint64_t* delta, + uint64_t ext_w, size_t quotient_size, size_t next_step, + uint64_t* global_scratch) { + extern __shared__ uint64_t shared_values[]; + const size_t lane = threadIdx.x, tile = blockDim.x; + uint64_t* values=global_scratch ? global_scratch+static_cast(blockIdx.x)*slot_count*tile : shared_values; + const unsigned int log_size = static_cast(__ffsll(quotient_size) - 1); + for (size_t row = static_cast(blockIdx.x) * tile + lane; + row < quotient_size; row += static_cast(gridDim.x) * tile) { + const size_t sr = reverse_low_bits(row, log_size); + const size_t nr = reverse_low_bits((row + next_step) & (quotient_size - 1), log_size); + for (size_t i = 0; i < node_count; ++i) { + const ConstraintNode n = nodes[i]; + const uint64_t a = n.op>=6 ? values[n.a * tile + lane] : 0; + const uint64_t b = n.op>=6&&n.op!=9 ? values[n.b * tile + lane] : 0; + uint64_t v = 0; + switch (n.op) { + case 0: v = n.value; break; + case 1: { const ResidentLde* m = n.aux < 2 ? preprocessed : n.aux < 4 ? main : stage2; + const size_t r = (n.aux & 1U) ? nr : sr; v = m->values[r * m->width + n.a]; break; } + case 2: v = publics[n.a]; break; + case 3: v = selectors[row]; break; + case 4: v = selectors[quotient_size + row]; break; + case 5: v = selectors[2 * quotient_size + row]; break; + case 6: v = goldilocks_add(a,b); break; + case 7: v = goldilocks_sub(a,b); break; + case 8: v = goldilocks_mul(a,b); break; + case 9: v = a == 0 ? 0 : GOLDILOCKS_P - canonicalize(a); break; + } + values[n.out * tile + lane] = v; + } + const size_t total_constraints = root_count + + (lookup_count == 0 ? 2 : ((lookup_count + group_size - 1) / group_size) * 2); + Ext2 acc{0,0}; size_t ci = 0; + for (; ci < root_count; ++ci) { + const uint64_t c = values[roots[ci] * tile + lane]; + acc.c0 = goldilocks_add(acc.c0, goldilocks_mul(c, alpha[ci])); + acc.c1 = goldilocks_add(acc.c1, goldilocks_mul(c, alpha[total_constraints + ci])); + } + const Ext2 beta{publics[0], publics[1]}, gamma{publics[2], publics[3]}; + const Ext2 injection{goldilocks_mul(selectors[quotient_size + row], delta[0]), + goldilocks_mul(selectors[quotient_size + row], delta[1])}; + const size_t groups = lookup_count == 0 ? 1 : (lookup_count + group_size - 1) / group_size; + for (size_t g = 0; g < groups; ++g) { + Ext2 constraints[8]; size_t count = 0; + if (lookup_count == 0) { + constraints[count++] = ext2_add(ext2_sub( + {stage2->values[nr * stage2->width], stage2->values[nr * stage2->width + 1]}, + {stage2->values[sr * stage2->width], stage2->values[sr * stage2->width + 1]}), injection); + } else { + const size_t begin = g * group_size; + const size_t end = begin + group_size < lookup_count ? begin + group_size : lookup_count; + Ext2 product{1,0}, messages[8]; + for (size_t j=begin;j0;--k) { f=ext2_mul(f,gamma,ext_w); f.c0=goldilocks_add(f.c0,values[lookup_args[l.arg_start+k-1]*tile+lane]); } + messages[j-begin]=ext2_add(f,beta); product=ext2_mul(product,messages[j-begin],ext_w); } + const Ext2 source{stage2->values[sr*stage2->width+2*g],stage2->values[sr*stage2->width+2*g+1]}; + Ext2 target = g+1values[sr*stage2->width+2*g+2],stage2->values[sr*stage2->width+2*g+3]} + : ext2_add({stage2->values[nr*stage2->width],stage2->values[nr*stage2->width+1]},injection); + Ext2 rhs{0,0}; + for(size_t j=begin;j(maximum) < 96 * 1024 + ? static_cast(maximum) : 96 * 1024; + return cudaSuccess; +} + +cudaError_t configure_quotient_shared_memory(int device_id, size_t bytes) { + static volatile int lock = 0; + static size_t configured[64] = {}; + if (device_id < 0 || device_id >= 64) return cudaErrorInvalidDevice; + while (__sync_lock_test_and_set(&lock, 1)) {} + cudaError_t status = cudaSuccess; + if (configured[device_id] < bytes) { + status = cudaFuncSetAttribute( + evaluate_quotient, cudaFuncAttributeMaxDynamicSharedMemorySize, + static_cast(bytes)); + if (status == cudaSuccess) configured[device_id] = bytes; + } + __sync_lock_release(&lock); + return status; +} + +__global__ void canonicalize_goldilocks(uint64_t* values, size_t count) { + for (size_t index = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + index < count; + index += static_cast(gridDim.x) * blockDim.x) { + values[index] = canonicalize(values[index]); + } +} + +__global__ void gather_lde_rows(uint64_t* output, const uint64_t* values, + const uint64_t* rows, size_t row_count, + size_t width) { + const size_t index = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + const size_t elements = row_count * width; + if (index < elements) { + const size_t output_row = index / width; + const size_t column = index - output_row * width; + output[index] = values[rows[output_row] * width + column]; + } +} + +// Convert the raw bit-reversed storage of DFT(quotient evaluations) into +// shifted coefficient slices. This is the device form of +// shifted_quotient_slices: output rows are natural coefficient indices and +// columns are [slice][extension coordinate]. The tail of the blown-up +// allocation is zeroed by the caller before this kernel runs. +__global__ void gather_shifted_quotient_slices( + uint64_t* output, const uint64_t* transformed, const uint64_t* weights, + size_t quotient_height, size_t trace_height, size_t quotient_degree, + size_t extension_degree) { + const size_t width = quotient_degree * extension_degree; + const size_t elements = trace_height * width; + const unsigned int log_height = static_cast(__ffsll(quotient_height) - 1); + const size_t grid_stride = static_cast(blockDim.x) * gridDim.x; + for (size_t index = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + index < elements; index += grid_stride) { + const size_t row = index / width; + const size_t column = index - row * width; + const size_t chunk = column / extension_degree; + const size_t coordinate = column - chunk * extension_degree; + const size_t coefficient = chunk * trace_height + row; + const size_t negated = (quotient_height - coefficient) & (quotient_height - 1); + const size_t source = reverse_low_bits(negated, log_height); + output[index] = goldilocks_mul( + transformed[source * extension_degree + coordinate], weights[chunk]); + } +} + +__global__ void gather_mixed_lde_row(uint64_t* output, + const uint64_t* const* values, + const size_t* widths, + const size_t* rows, + const size_t* offsets, + size_t matrix_count) { + const size_t matrix = blockIdx.x; + if (matrix >= matrix_count) return; + for (size_t column = threadIdx.x; column < widths[matrix]; column += blockDim.x) { + output[offsets[matrix] + column] = + values[matrix][rows[matrix] * widths[matrix] + column]; + } +} + +__global__ void gather_mixed_lde_rows(uint64_t* output, + const uint64_t* const* values, + const size_t* widths, + const size_t* row_shifts, + const size_t* offsets, + const uint64_t* indices, + size_t matrix_count, size_t query_count, + size_t max_height, size_t row_width) { + const size_t task = blockIdx.x; + const size_t matrix = task % matrix_count; + const size_t query = task / matrix_count; + if (query >= query_count) return; + const size_t row = static_cast(indices[query]) >> row_shifts[matrix]; + for (size_t column = threadIdx.x; column < widths[matrix]; column += blockDim.x) { + output[query * row_width + offsets[matrix] + column] = + values[matrix][row * widths[matrix] + column]; + } +} + +__global__ void radix2_dif_stage(uint64_t* values, size_t height, size_t width, + size_t half, const uint64_t* twiddles) { + const size_t total = (height >> 1) * width; + const size_t stride = height / (2 * half); + const size_t grid_stride = static_cast(blockDim.x) * gridDim.x; + for (size_t index = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + index < total; index += grid_stride) { + const size_t butterfly = index / width; + const size_t column = index - butterfly * width; + const size_t offset = butterfly % half; + const size_t group = butterfly / half; + const size_t row_0 = group * (2 * half) + offset; + const size_t row_1 = row_0 + half; + const size_t index_0 = row_0 * width + column; + const size_t index_1 = row_1 * width + column; + const uint64_t left = values[index_0]; + const uint64_t right = values[index_1]; + values[index_0] = goldilocks_add(left, right); + values[index_1] = goldilocks_mul( + goldilocks_sub(left, right), twiddles[offset * stride]); + } +} + +// Fuse two consecutive radix-2 DIF stages. Each thread owns one column of +// one four-row butterfly, so row-major accesses remain coalesced while the +// intermediate values never return to global memory. +__global__ void radix4_dif_stage(uint64_t* values,size_t height,size_t width, + size_t half,const uint64_t* twiddles){ + const size_t quarter=half>>1,total=(height>>2)*width,stride=height/(2*half); + const size_t grid_stride=static_cast(blockDim.x)*gridDim.x; + for(size_t index=static_cast(blockIdx.x)*blockDim.x+threadIdx.x;index>2,total=(height>>3)*width,stride=height/(2*half); + const size_t grid_stride=static_cast(blockDim.x)*gridDim.x; + for(size_t index=static_cast(blockIdx.x)*blockDim.x+threadIdx.x;index>= 1) { + const size_t butterflies = start_half * width; + const size_t stride = height / (2 * half); + for (size_t index = threadIdx.x; index < butterflies; + index += blockDim.x) { + const size_t butterfly = index / width; + const size_t column = index - butterfly * width; + const size_t offset = butterfly % half; + const size_t subgroup = butterfly / half; + const size_t row_0 = subgroup * (2 * half) + offset; + const size_t row_1 = row_0 + half; + const size_t index_0 = row_0 * width + column; + const size_t index_1 = row_1 * width + column; + const uint64_t left = local_values[index_0]; + const uint64_t right = local_values[index_1]; + local_values[index_0] = goldilocks_add(left, right); + local_values[index_1] = goldilocks_mul( + goldilocks_sub(left, right), twiddles[offset * stride]); + } + __syncthreads(); + if (half == 1) { + break; + } + } + + for (size_t index = threadIdx.x; index < elements_per_group; + index += blockDim.x) { + values[global_base + index] = local_values[index]; + } + __syncthreads(); + } +} + +// Wide row-major matrices cannot fit all columns of a row group in shared +// memory. Tile the columns as well, preserving coalesced row-major loads while +// fusing the final DIF stages independently for each column tile. +__global__ void radix2_dif_tail_tiled(uint64_t* values,size_t height,size_t width, + size_t start_half,const uint64_t* twiddles,size_t columns_per_tile){ + extern __shared__ uint64_t local_values[];const size_t rows=2*start_half; + const size_t groups=height/rows,column_tiles=(width+columns_per_tile-1)/columns_per_tile; + const size_t tile_count=groups*column_tiles; + for(size_t tile_index=blockIdx.x;tile_index>=1){const size_t butterflies=start_half*columns,stride=height/(2*half); + for(size_t index=threadIdx.x;index> 1) * width; + const unsigned int blocks = blocks_for(total); + // Tall, wide row-major batches amortize the heavier fused kernel. Width-2 + // FRI codewords retain the shared-memory tail specialized below. + if (width >= 8 && height >= (size_t(1) << 18)) { + size_t half = height >> 1; + while (half >= 4) { + radix8_dif_stage<<> 3) * width), THREADS>>>( + values, height, width, half, twiddles); + const cudaError_t status = cudaGetLastError(); + if (status != cudaSuccess) return status; + half >>= 3; + } + if (half == 2) { + radix4_dif_stage<<> 2) * width), THREADS>>>( + values, height, width, half, twiddles); + return cudaGetLastError(); + } + if (half == 1) { + radix2_dif_stage<<>>(values, height, width, 1, + twiddles); + } + return cudaGetLastError(); + } + if (width > 2) { + for (size_t half = height >> 1;; half >>= 1) { + radix2_dif_stage<<>>(values, height, width, half, + twiddles); + const cudaError_t status = cudaGetLastError(); + if (status != cudaSuccess) return status; + if (half == 1) return cudaSuccess; + } + } + constexpr size_t TAIL_HALF = 128; + for (size_t half = height >> 1; half > TAIL_HALF; half >>= 1) { + radix2_dif_stage<<>>(values, height, width, half, + twiddles); + const cudaError_t status = cudaGetLastError(); + if (status != cudaSuccess) { + return status; + } + } + + const size_t start_half = + height < 2 * TAIL_HALF ? height >> 1 : TAIL_HALF; + const size_t groups = height / (2 * start_half); + const unsigned int tail_blocks = static_cast( + groups < MAX_BLOCKS ? groups : MAX_BLOCKS); + const size_t shared_bytes = 2 * start_half * width * sizeof(uint64_t); + radix2_dif_tail<<>>( + values, height, width, start_half, twiddles); + return cudaGetLastError(); +} + +__global__ void bit_reverse_scale_and_shift(uint64_t* values, size_t height, + size_t width, unsigned int log_height, + uint64_t height_inverse, + const uint64_t* shift_powers) { + const size_t total = height * width; + const size_t grid_stride = static_cast(blockDim.x) * gridDim.x; + for (size_t index = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + index < total; index += grid_stride) { + const size_t row = index / width; + const size_t column = index - row * width; + const size_t reverse_row = reverse_index_bits(row, log_height); + if (row > reverse_row) { + continue; + } + + const size_t reverse_index = reverse_row * width + column; + const uint64_t left = values[index]; + if (row == reverse_row) { + values[index] = goldilocks_mul( + left, goldilocks_mul(height_inverse, shift_powers[row])); + continue; + } + + const uint64_t right = values[reverse_index]; + values[index] = goldilocks_mul( + right, goldilocks_mul(height_inverse, shift_powers[row])); + values[reverse_index] = goldilocks_mul( + left, goldilocks_mul(height_inverse, shift_powers[reverse_row])); + } +} + +__global__ void goldilocks_ops_kernel(uint64_t* sums, uint64_t* differences, + uint64_t* products, uint64_t* inverses, + const uint64_t* left, const uint64_t* right, + size_t len) { + const size_t grid_stride = static_cast(blockDim.x) * gridDim.x; + for (size_t index = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + index < len; index += grid_stride) { + const uint64_t a = left[index]; + const uint64_t b = right[index]; + sums[index] = goldilocks_add(a, b); + differences[index] = goldilocks_sub(a, b); + products[index] = goldilocks_mul(a, b); + const uint64_t canonical_a = canonicalize(a); + inverses[index] = canonical_a == 0 + ? 0 + : goldilocks_pow(canonical_a, GOLDILOCKS_P - 2); + } +} + +// One warp owns one message. Chunk compression is distributed across lanes; +// lane zero then reduces the (at most 32) chunk chaining values into the root. +// This matches the BLAKE3 tree shape while exposing parallelism within wide +// Merkle leaves instead of assigning a long row to one scalar CUDA thread. +__global__ void blake3_hash_rows_kernel(uint8_t* digests, + const uint8_t* messages, + size_t message_bytes, + size_t message_count) { + __shared__ uint32_t chunk_values[THREADS / 32][32][8]; + const unsigned int lane = threadIdx.x & 31U; + const unsigned int warp_in_block = threadIdx.x >> 5; + const size_t warp_index = + static_cast(blockIdx.x) * (blockDim.x / 32) + warp_in_block; + const size_t warp_stride = + static_cast(gridDim.x) * (blockDim.x / 32); + + for (size_t message_index = warp_index; message_index < message_count; + message_index += warp_stride) { + const uint8_t* message = messages + message_index * message_bytes; + const size_t chunks = message_bytes == 0 ? 1 : (message_bytes + 1023) / 1024; + if (lane < chunks) { + const size_t chunk_offset = static_cast(lane) * 1024; + const size_t chunk_length = + chunk_offset < message_bytes + ? ((message_bytes - chunk_offset) < 1024 + ? message_bytes - chunk_offset + : 1024) + : 0; + uint32_t chunk_output[8]; + blake3_chunk(message + chunk_offset, chunk_length, lane, + chunks == 1, chunk_output); +#pragma unroll + for (unsigned int word = 0; word < 8; ++word) { + chunk_values[warp_in_block][lane][word] = chunk_output[word]; + } + } + __syncwarp(); + + if (lane == 0) { + size_t count = chunks; + while (count > 1) { + const bool root_level = count == 2; + size_t next_count = 0; + for (size_t index = 0; index + 1 < count; index += 2) { + uint32_t parent[8]; + blake3_parent(chunk_values[warp_in_block][index], + chunk_values[warp_in_block][index + 1], + root_level, parent); +#pragma unroll + for (unsigned int word = 0; word < 8; ++word) { + chunk_values[warp_in_block][next_count][word] = parent[word]; + } + ++next_count; + } + if ((count & 1U) != 0) { +#pragma unroll + for (unsigned int word = 0; word < 8; ++word) { + chunk_values[warp_in_block][next_count][word] = + chunk_values[warp_in_block][count - 1][word]; + } + ++next_count; + } + count = next_count; + } + + uint8_t* digest = digests + message_index * 32; +#pragma unroll + for (unsigned int word = 0; word < 8; ++word) { + const uint32_t value = chunk_values[warp_in_block][0][word]; + digest[word * 4] = static_cast(value); + digest[word * 4 + 1] = static_cast(value >> 8); + digest[word * 4 + 2] = static_cast(value >> 16); + digest[word * 4 + 3] = static_cast(value >> 24); + } + } + __syncwarp(); + } +} + +__global__ void blake3_hash_digest_pairs_kernel( + uint8_t* digests, const uint8_t* left, const uint8_t* right, + size_t count, bool interleaved) { + const size_t grid_stride = static_cast(blockDim.x) * gridDim.x; + for (size_t index = static_cast(blockIdx.x) * blockDim.x + + threadIdx.x; + index < count; index += grid_stride) { + const uint8_t* left_digest = + left + (interleaved ? 2 * index : index) * 32; + const uint8_t* right_digest = + right + (interleaved ? 2 * index + 1 : index) * 32; + blake3_hash_digest_pair(left_digest, right_digest, + digests + index * 32); + } +} + +cudaError_t launch_blake3_rows(uint8_t* digests, const uint8_t* messages, + size_t message_bytes, size_t message_count) { + constexpr unsigned int WARPS_PER_BLOCK = THREADS / 32; + const size_t required = + (message_count + WARPS_PER_BLOCK - 1) / WARPS_PER_BLOCK; + const unsigned int blocks = static_cast( + required < MAX_BLOCKS ? required : MAX_BLOCKS); + blake3_hash_rows_kernel<<>>( + digests, messages, message_bytes, message_count); + return cudaGetLastError(); +} + +cudaError_t launch_blake3_digest_pairs(uint8_t* digests, + const uint8_t* left, + const uint8_t* right, size_t count, + bool interleaved) { + blake3_hash_digest_pairs_kernel<<>>( + digests, left, right, count, interleaved); + return cudaGetLastError(); +} + +cudaError_t hash_resident_lde_group(uint8_t* digests, + const void* const* handles, + size_t handle_count, size_t height) { + size_t total_width = 0; + for (size_t index = 0; index < handle_count; ++index) { + const ResidentLde* lde = static_cast(handles[index]); + if (lde != nullptr && lde->height == height) { + if (lde->width > SIZE_MAX - total_width) { + return cudaErrorInvalidValue; + } + total_width += lde->width; + } + } + if (total_width == 0 || total_width > (32 * 1024) / sizeof(uint64_t) || + !product_fits(height, total_width)) { + return cudaErrorInvalidValue; + } + + constexpr size_t ROW_STAGING_BYTES = size_t(256) << 20; + const size_t row_bytes = total_width * sizeof(uint64_t); + const size_t rows_per_chunk = + (ROW_STAGING_BYTES / row_bytes) > 0 ? (ROW_STAGING_BYTES / row_bytes) : 1; + DeviceBuffer combined_rows; + cudaError_t status = combined_rows.allocate( + (height < rows_per_chunk ? height : rows_per_chunk) * total_width); + for (size_t row_start = 0; status == cudaSuccess && row_start < height; + row_start += rows_per_chunk) { + const size_t rows = + (height - row_start < rows_per_chunk) ? height - row_start + : rows_per_chunk; + size_t column_offset = 0; + for (size_t index = 0; status == cudaSuccess && index < handle_count; + ++index) { + const ResidentLde* lde = + static_cast(handles[index]); + if (lde == nullptr || lde->height != height) { + continue; + } + status = cudaMemcpy2DAsync( + combined_rows.get() + column_offset, + total_width * sizeof(uint64_t), + lde->values + row_start * lde->width, + lde->width * sizeof(uint64_t), + lde->width * sizeof(uint64_t), rows, cudaMemcpyDeviceToDevice, + cudaStreamPerThread); + column_offset += lde->width; + } + if (status == cudaSuccess) { + status = launch_blake3_rows( + digests + row_start * 32, + reinterpret_cast(combined_rows.get()), row_bytes, + rows); + } + } + return status; +} + +cudaError_t copy_to_device(DeviceBuffer& destination, const uint64_t* source, + size_t elements) { + cudaError_t status = destination.allocate(elements); + if (status != cudaSuccess) { + return status; + } + return cudaMemcpy(destination.get(), source, elements * sizeof(uint64_t), + cudaMemcpyHostToDevice); +} + +cudaError_t copy_to_host(uint64_t* destination, const DeviceBuffer& source, + size_t elements) { + return cudaMemcpy(destination, source.get(), elements * sizeof(uint64_t), + cudaMemcpyDeviceToHost); +} + +} // namespace + +extern "C" int multi_stark_cuda_dft_batch(int device_id, uint64_t* values, + size_t height, size_t width, + const uint64_t* twiddles) { + if (values == nullptr || twiddles == nullptr || !is_power_of_two(height) || + width == 0 || !product_fits(height, width)) { + return static_cast(cudaErrorInvalidValue); + } + cudaError_t status = cudaSetDevice(device_id); + if (status != cudaSuccess) { + return static_cast(status); + } + + const size_t elements = height * width; + HostRegistration registered_values(values, elements * sizeof(uint64_t)); + DeviceBuffer device_values; + DeviceBuffer device_twiddles; + status = copy_to_device(device_values, values, elements); + if (status == cudaSuccess) { + status = copy_to_device(device_twiddles, twiddles, height / 2); + } + if (status == cudaSuccess) { + status = launch_dif(device_values.get(), height, width, device_twiddles.get()); + } + if (status == cudaSuccess) { + status = copy_to_host(values, device_values, elements); + } + return static_cast(status); +} + +extern "C" int multi_stark_cuda_coset_lde_batch( + int device_id, uint64_t* output, const uint64_t* input, size_t height, + size_t width, size_t added_bits, const uint64_t* inverse_twiddles, + const uint64_t* shift_powers, const uint64_t* forward_twiddles, + uint64_t height_inverse) { + if (output == nullptr || input == nullptr || inverse_twiddles == nullptr || + shift_powers == nullptr || forward_twiddles == nullptr || + !is_power_of_two(height) || width == 0 || + added_bits >= sizeof(size_t) * 8 || height > (SIZE_MAX >> added_bits)) { + return static_cast(cudaErrorInvalidValue); + } + const size_t extended_height = height << added_bits; + if (!product_fits(extended_height, width)) { + return static_cast(cudaErrorInvalidValue); + } + + cudaError_t status = cudaSetDevice(device_id); + if (status != cudaSuccess) { + return static_cast(status); + } + + const size_t input_elements = height * width; + const size_t output_elements = extended_height * width; + HostRegistration registered_input( + const_cast(input), input_elements * sizeof(uint64_t)); + HostRegistration registered_output(output, + output_elements * sizeof(uint64_t)); + DeviceBuffer device_values; + DeviceBuffer device_inverse_twiddles; + DeviceBuffer device_shift_powers; + DeviceBuffer device_forward_twiddles; + status = device_values.allocate(output_elements); + if (status == cudaSuccess) { + status = cudaMemset(device_values.get(), 0, output_elements * sizeof(uint64_t)); + } + if (status == cudaSuccess) { + status = cudaMemcpy(device_values.get(), input, + input_elements * sizeof(uint64_t), + cudaMemcpyHostToDevice); + } + if (status == cudaSuccess) { + status = copy_to_device(device_inverse_twiddles, inverse_twiddles, height / 2); + } + if (status == cudaSuccess) { + status = copy_to_device(device_shift_powers, shift_powers, height); + } + if (status == cudaSuccess) { + status = copy_to_device(device_forward_twiddles, forward_twiddles, + extended_height / 2); + } + if (status == cudaSuccess) { + status = launch_dif(device_values.get(), height, width, + device_inverse_twiddles.get()); + } + if (status == cudaSuccess) { + const size_t total = input_elements; + bit_reverse_scale_and_shift<<>>( + device_values.get(), height, width, strict_log2(height), + height_inverse, device_shift_powers.get()); + status = cudaGetLastError(); + } + if (status == cudaSuccess) { + status = launch_dif(device_values.get(), extended_height, width, + device_forward_twiddles.get()); + } + if (status == cudaSuccess) { + status = copy_to_host(output, device_values, output_elements); + } + return static_cast(status); +} + +extern "C" int multi_stark_cuda_coset_lde_create( + int device_id, void** handle, const uint64_t* input, size_t height, + size_t width, size_t added_bits, const uint64_t* inverse_twiddles, + const uint64_t* shift_powers, const uint64_t* forward_twiddles, + uint64_t height_inverse) { + if (handle == nullptr || input == nullptr || (height > 1 && inverse_twiddles == nullptr) || + shift_powers == nullptr || forward_twiddles == nullptr || + !is_power_of_two(height) || width == 0 || + added_bits >= sizeof(size_t) * 8 || height > (SIZE_MAX >> added_bits)) { + return static_cast(cudaErrorInvalidValue); + } + *handle = nullptr; + const size_t extended_height = height << added_bits; + if (!product_fits(extended_height, width)) { + return static_cast(cudaErrorInvalidValue); + } + cudaError_t status = cudaSetDevice(device_id); + if (status != cudaSuccess) { + return static_cast(status); + } + + ResidentLde* lde = nullptr; + status = create_resident_lde(&lde); + if (status != cudaSuccess) return static_cast(status); + lde->height = extended_height; + lde->width = width; + const size_t input_elements = height * width; + const size_t output_elements = extended_height * width; + const size_t input_bytes=input_elements*sizeof(uint64_t); + bool input_registered=false; + status = cudaMalloc(reinterpret_cast(&lde->values), + output_elements * sizeof(uint64_t)); + if (status == cudaSuccess) { + status = cudaMalloc(reinterpret_cast(&lde->trace_values), + input_elements * sizeof(uint64_t)); + } + if (status == cudaSuccess) lde->trace_height = height; + + // Large pageable uploads otherwise serialize through the driver's hidden + // staging pool. Registration is best-effort: constrained hosts retain the + // fully independent pageable path, while recursive proofs can DMA their + // largest trace matrices directly. + if(status==cudaSuccess&&input_bytes>=(size_t(8)<<20)){ + const cudaError_t registration=cudaHostRegister( + const_cast(input),input_bytes,cudaHostRegisterDefault); + if(registration==cudaSuccess)input_registered=true;else cudaGetLastError(); + } + + const uint64_t *device_inverse_twiddles=nullptr,*device_shift_powers=nullptr,*device_forward_twiddles=nullptr; + if (status == cudaSuccess) { + status = cudaMemsetAsync(lde->values, 0, + output_elements * sizeof(uint64_t), + cudaStreamPerThread); + } + if (status == cudaSuccess) { + status = cudaMemcpyAsync(lde->trace_values, input, + input_bytes, + cudaMemcpyHostToDevice, + cudaStreamPerThread); + } + if (status == cudaSuccess) { + status = cudaMemcpyAsync(lde->values, lde->trace_values, + input_elements * sizeof(uint64_t), + cudaMemcpyDeviceToDevice, + cudaStreamPerThread); + } + if (status == cudaSuccess && height > 1) { + status = cached_device_constants(device_id,inverse_twiddles,height/2,1,0,0,&device_inverse_twiddles); + } + if (status == cudaSuccess) { + status = cached_device_constants(device_id,shift_powers,height,3,shift_powers[0],height>1?shift_powers[1]:0,&device_shift_powers); + } + if (status == cudaSuccess) { + status = cached_device_constants(device_id,forward_twiddles,extended_height/2,2,0,0,&device_forward_twiddles); + } + if (status == cudaSuccess) { + status = launch_dif(lde->values, height, width,device_inverse_twiddles); + } + if (status == cudaSuccess) { + bit_reverse_scale_and_shift<<>>( + lde->values, height, width, strict_log2(height), height_inverse, + device_shift_powers); + status = cudaGetLastError(); + } + if (status == cudaSuccess) { + status = launch_dif(lde->values, extended_height, width,device_forward_twiddles); + } + if (status == cudaSuccess) { + canonicalize_goldilocks<<>>( + lde->values, output_elements); + status = cudaGetLastError(); + } + if (status == cudaSuccess) status = cudaStreamSynchronize(cudaStreamPerThread); + if(input_registered){const cudaError_t unregister_status=cudaHostUnregister(const_cast(input)); + if(status==cudaSuccess)status=unregister_status;} + if (status != cudaSuccess) { + destroy_resident_lde(lde); + return static_cast(status); + } + *handle = lde; + return static_cast(cudaSuccess); +} + +extern "C" int multi_stark_cuda_prepare_lde_constants( + int device_id,const uint64_t* inverse_twiddles,size_t inverse_count, + const uint64_t* shift_powers,size_t height,const uint64_t* forward_twiddles, + size_t forward_count) { + if((inverse_count&&!inverse_twiddles)||!shift_powers||!height|| + !forward_twiddles||!forward_count)return static_cast(cudaErrorInvalidValue); + cudaError_t status=cudaSetDevice(device_id);const uint64_t* ignored=nullptr; + if(status==cudaSuccess&&inverse_count) + status=cached_device_constants(device_id,inverse_twiddles,inverse_count,1,0,0,&ignored); + if(status==cudaSuccess) + status=cached_device_constants(device_id,shift_powers,height,3,shift_powers[0],height>1?shift_powers[1]:0,&ignored); + if(status==cudaSuccess) + status=cached_device_constants(device_id,forward_twiddles,forward_count,2,0,0,&ignored); + return static_cast(status); +} + +extern "C" int multi_stark_cuda_lde_create_from_host( + int device_id, void** handle, const uint64_t* input, size_t height, + size_t width) { + if (handle == nullptr || input == nullptr || !is_power_of_two(height) || + width == 0 || !product_fits(height, width)) { + return static_cast(cudaErrorInvalidValue); + } + *handle = nullptr; + cudaError_t status = cudaSetDevice(device_id); + if (status != cudaSuccess) { + return static_cast(status); + } + ResidentLde* lde = nullptr; + status = create_resident_lde(&lde); + if (status != cudaSuccess) return static_cast(status); + lde->height = height; + lde->width = width; + status = cudaMalloc(reinterpret_cast(&lde->values), + height * width * sizeof(uint64_t)); + if (status == cudaSuccess) { + status = cudaMemcpy(lde->values, input, + height * width * sizeof(uint64_t), + cudaMemcpyHostToDevice); + } + if (status == cudaSuccess) { + const size_t count = height * width; + canonicalize_goldilocks<<>>(lde->values, count); + status = cudaGetLastError(); + } + if (status == cudaSuccess) { + status = cudaStreamSynchronize(cudaStreamPerThread); + } + if (status != cudaSuccess) { + destroy_resident_lde(lde); + return static_cast(status); + } + *handle = lde; + return static_cast(cudaSuccess); +} + +extern "C" int multi_stark_cuda_zero_lde_create( + int device_id, void** handle, size_t height, size_t width) { + if (!handle || !is_power_of_two(height) || width == 0 || !product_fits(height,width)) + return static_cast(cudaErrorInvalidValue); + *handle=nullptr;cudaError_t status=cudaSetDevice(device_id); + ResidentLde* lde=nullptr;if(status==cudaSuccess)status=create_resident_lde(&lde); + lde->height=height;lde->width=width; + if(status==cudaSuccess)status=cudaMalloc(reinterpret_cast(&lde->values),height*width*sizeof(uint64_t)); + if(status==cudaSuccess)status=cudaMemset(lde->values,0,height*width*sizeof(uint64_t)); + if(status==cudaSuccess)*handle=lde;else destroy_resident_lde(lde);return static_cast(status); +} + +extern "C" int multi_stark_cuda_lde_copy_to_host(int device_id, + const void* handle, + uint64_t* output) { + if (handle == nullptr || output == nullptr) { + return static_cast(cudaErrorInvalidValue); + } + cudaError_t status = cudaSetDevice(device_id); + if (status != cudaSuccess) { + return static_cast(status); + } + const ResidentLde* lde = static_cast(handle); + HostRegistration registered_output( + output, lde->height * lde->width * sizeof(uint64_t)); + status = cudaMemcpy(output, lde->values, + lde->height * lde->width * sizeof(uint64_t), + cudaMemcpyDeviceToHost); + return static_cast(status); +} + +extern "C" int multi_stark_cuda_lde_copy_row(int device_id, + const void* handle, + size_t row, + uint64_t* output) { + if (handle == nullptr || output == nullptr) { + return static_cast(cudaErrorInvalidValue); + } + cudaError_t status = cudaSetDevice(device_id); + if (status != cudaSuccess) { + return static_cast(status); + } + const ResidentLde* lde = static_cast(handle); + if (row >= lde->height) { + return static_cast(cudaErrorInvalidValue); + } + status = cudaMemcpy(output, lde->values + row * lde->width, + lde->width * sizeof(uint64_t), + cudaMemcpyDeviceToHost); + return static_cast(status); +} + +extern "C" int multi_stark_cuda_lde_copy_rows( + int device_id, const void* handle, const uint64_t* rows, + size_t row_count, uint64_t* output) { + if (handle == nullptr || rows == nullptr || row_count == 0 || + output == nullptr) { + return static_cast(cudaErrorInvalidValue); + } + cudaError_t status = cudaSetDevice(device_id); + if (status != cudaSuccess) { + return static_cast(status); + } + const ResidentLde* lde = static_cast(handle); + if (!product_fits(row_count, lde->width)) { + return static_cast(cudaErrorInvalidValue); + } + for (size_t index = 0; index < row_count; ++index) { + if (rows[index] >= lde->height) { + return static_cast(cudaErrorInvalidValue); + } + } + + DeviceBuffer device_rows; + DeviceBuffer device_output; + status = copy_to_device(device_rows, rows, row_count); + if (status == cudaSuccess) { + status = device_output.allocate(row_count * lde->width); + } + if (status == cudaSuccess) { + gather_lde_rows<<width), THREADS>>>( + device_output.get(), lde->values, device_rows.get(), row_count, + lde->width); + status = cudaGetLastError(); + } + if (status == cudaSuccess) { + status = cudaMemcpy(output, device_output.get(), + row_count * lde->width * sizeof(uint64_t), + cudaMemcpyDeviceToHost); + } + return static_cast(status); +} + +extern "C" int multi_stark_cuda_mixed_lde_open_row( + int device_id, uint64_t* output, const void* const* handles, + size_t handle_count, size_t index) { + if (output == nullptr || handles == nullptr || handle_count == 0) return static_cast(cudaErrorInvalidValue); + cudaError_t status=cudaSetDevice(device_id); size_t max_height=0,total=0; + for(size_t i=0;i(handles[i]); + if(l==nullptr || l->height==0 || (l->height&(l->height-1))!=0 || l->width>SIZE_MAX-total) return static_cast(cudaErrorInvalidValue); + if(l->height>max_height)max_height=l->height; total+=l->width;} + const uint64_t** hv=new(std::nothrow) const uint64_t*[handle_count]; size_t* hw=new(std::nothrow) size_t[3*handle_count]; + if(hv==nullptr||hw==nullptr){delete[] hv;delete[] hw;return static_cast(cudaErrorMemoryAllocation);} size_t* hr=hw+handle_count;size_t* ho=hr+handle_count; + size_t off=0;for(size_t i=0;i(handles[i]);hv[i]=l->values;hw[i]=l->width; + hr[i]=index>>(strict_log2(max_height)-strict_log2(l->height));ho[i]=off;off+=l->width;} + const uint64_t** dv=nullptr;size_t* dm=nullptr;uint64_t* dout=nullptr; + if(status==cudaSuccess)status=cudaMalloc(reinterpret_cast(&dv),handle_count*sizeof(*dv)); + if(status==cudaSuccess)status=cudaMalloc(reinterpret_cast(&dm),3*handle_count*sizeof(size_t)); + if(status==cudaSuccess)status=cudaMalloc(reinterpret_cast(&dout),total*sizeof(uint64_t)); + if(status==cudaSuccess)status=cudaMemcpy(dv,hv,handle_count*sizeof(*dv),cudaMemcpyHostToDevice); + if(status==cudaSuccess)status=cudaMemcpy(dm,hw,3*handle_count*sizeof(size_t),cudaMemcpyHostToDevice); + if(status==cudaSuccess){gather_mixed_lde_row<<(handle_count),THREADS>>>(dout,dv,dm,dm+handle_count,dm+2*handle_count,handle_count);status=cudaGetLastError();} + if(status==cudaSuccess)status=cudaMemcpy(output,dout,total*sizeof(uint64_t),cudaMemcpyDeviceToHost); + cudaFree(dout);cudaFree(dm);cudaFree(dv);delete[] hw;delete[] hv;return static_cast(status); +} + +extern "C" int multi_stark_cuda_mixed_lde_open_rows( + int device_id, uint64_t* output, const void* const* handles, + size_t handle_count, const uint64_t* indices, size_t query_count) { + if (output == nullptr || handles == nullptr || handle_count == 0 || + indices == nullptr || query_count == 0) return static_cast(cudaErrorInvalidValue); + cudaError_t status=cudaSetDevice(device_id);size_t max_height=0,total=0; + const uint64_t** hv=new(std::nothrow) const uint64_t*[handle_count]; + size_t* hm=new(std::nothrow) size_t[3*handle_count]; + if(hv==nullptr||hm==nullptr){delete[] hv;delete[] hm;return static_cast(cudaErrorMemoryAllocation);} + size_t* hw=hm,*hh=hm+handle_count,*ho=hh+handle_count;size_t off=0; + for(size_t i=0;i(handles[i]); + if(l==nullptr||l->height==0||(l->height&(l->height-1))!=0||l->width>SIZE_MAX-total){delete[] hm;delete[] hv;return static_cast(cudaErrorInvalidValue);} + hv[i]=l->values;hw[i]=l->width;hh[i]=l->height;ho[i]=off;off+=l->width;total+=l->width;if(l->height>max_height)max_height=l->height;} + for(size_t i=0;i=max_height){delete[] hm;delete[] hv;return static_cast(cudaErrorInvalidValue);} + const uint64_t** dv=nullptr;size_t* dm=nullptr;uint64_t* di=nullptr;uint64_t* dout=nullptr; + if(status==cudaSuccess)status=cudaMalloc(reinterpret_cast(&dv),handle_count*sizeof(*dv)); + if(status==cudaSuccess)status=cudaMalloc(reinterpret_cast(&dm),3*handle_count*sizeof(size_t)); + if(status==cudaSuccess)status=cudaMalloc(reinterpret_cast(&di),query_count*sizeof(uint64_t)); + if(status==cudaSuccess)status=cudaMalloc(reinterpret_cast(&dout),query_count*total*sizeof(uint64_t)); + if(status==cudaSuccess)status=cudaMemcpy(dv,hv,handle_count*sizeof(*dv),cudaMemcpyHostToDevice); + if(status==cudaSuccess)status=cudaMemcpy(dm,hm,3*handle_count*sizeof(size_t),cudaMemcpyHostToDevice); + if(status==cudaSuccess)status=cudaMemcpy(di,indices,query_count*sizeof(uint64_t),cudaMemcpyHostToDevice); + if(status==cudaSuccess){gather_mixed_lde_rows<<(handle_count*query_count),THREADS>>>(dout,dv,dm,dm+handle_count,dm+2*handle_count,di,handle_count,query_count,max_height,total);status=cudaGetLastError();} + if(status==cudaSuccess)status=cudaMemcpy(output,dout,query_count*total*sizeof(uint64_t),cudaMemcpyDeviceToHost); + cudaFree(dout);cudaFree(di);cudaFree(dm);cudaFree(dv);delete[] hm;delete[] hv;return static_cast(status); +} + +extern "C" int multi_stark_cuda_constraint_graph( + int device_id, uint64_t* output, const void* nodes, size_t node_count, + const uint32_t* roots, size_t root_count, const void* preprocessed_handle, + const void* main_handle, const void* stage2_handle, + const uint64_t* publics, size_t public_count, const uint64_t* selectors, + size_t quotient_size, size_t next_step) { + if (output == nullptr || nodes == nullptr || node_count == 0 || + roots == nullptr || root_count == 0 || main_handle == nullptr || + stage2_handle == nullptr || publics == nullptr || selectors == nullptr || + !is_power_of_two(quotient_size) || !is_power_of_two(next_step) || + next_step > quotient_size) { + return static_cast(cudaErrorInvalidValue); + } + cudaError_t status = cudaSetDevice(device_id); + ConstraintNode* device_nodes = nullptr; + uint32_t* device_roots = nullptr; + uint64_t* device_publics = nullptr; + uint64_t* device_selectors = nullptr; + uint64_t* device_output = nullptr; + auto allocate_copy = [&](void** destination, const void* source, size_t bytes) { + if (status != cudaSuccess) return; + status = cudaMalloc(destination, bytes); + if (status == cudaSuccess) { + status = cudaMemcpy(*destination, source, bytes, cudaMemcpyHostToDevice); + } + }; + allocate_copy(reinterpret_cast(&device_nodes), nodes, + node_count * sizeof(ConstraintNode)); + allocate_copy(reinterpret_cast(&device_roots), roots, + root_count * sizeof(uint32_t)); + allocate_copy(reinterpret_cast(&device_publics), publics, + public_count * sizeof(uint64_t)); + allocate_copy(reinterpret_cast(&device_selectors), selectors, + 3 * quotient_size * sizeof(uint64_t)); + if (status == cudaSuccess) { + status = cudaMalloc(reinterpret_cast(&device_output), + quotient_size * root_count * sizeof(uint64_t)); + } + const size_t shared_budget = 48 * 1024; + size_t tile = shared_budget / (node_count * sizeof(uint64_t)); + if (tile > 32) tile = 32; + if (tile == 0) status = cudaErrorInvalidValue; + if (status == cudaSuccess) { + const size_t block_count = (quotient_size + tile - 1) / tile; + evaluate_constraint_graph<<( + block_count < MAX_BLOCKS ? block_count : MAX_BLOCKS), + static_cast(tile), + node_count * tile * sizeof(uint64_t)>>>( + device_output, device_nodes, node_count, device_roots, root_count, + static_cast(preprocessed_handle), + static_cast(main_handle), + static_cast(stage2_handle), device_publics, + device_selectors, quotient_size, next_step); + status = cudaGetLastError(); + } + if (status == cudaSuccess) { + status = cudaMemcpy(output, device_output, + quotient_size * root_count * sizeof(uint64_t), + cudaMemcpyDeviceToHost); + } + cudaFree(device_output); + cudaFree(device_selectors); + cudaFree(device_publics); + cudaFree(device_roots); + cudaFree(device_nodes); + return static_cast(status); +} + +extern "C" int multi_stark_cuda_quotient_values( + int device_id, uint64_t* output, const void* nodes, size_t node_count, + size_t slot_count, + const uint32_t* roots, size_t root_count, const void* lookups, + size_t lookup_count, const uint32_t* lookup_args, size_t lookup_arg_count, + size_t group_size, const void* preprocessed_handle, const void* main_handle, + const void* stage2_handle, const uint64_t* publics, size_t public_count, + uint64_t coset_shift, uint64_t coset_generator, uint64_t trace_last, + uint64_t vanishing_start, uint64_t vanishing_step, + const uint64_t* alpha, size_t constraint_count, + const uint64_t* delta, uint64_t ext_w, size_t quotient_size, + size_t next_step) { + if (output == nullptr || nodes == nullptr || roots == nullptr || + main_handle == nullptr || stage2_handle == nullptr || publics == nullptr || + alpha == nullptr || delta == nullptr || + node_count == 0 || slot_count == 0 || group_size == 0 || + !is_power_of_two(quotient_size) || !is_power_of_two(next_step) || + next_step > quotient_size || + (lookup_count != 0 && (lookups == nullptr || lookup_args == nullptr))) { + return static_cast(cudaErrorInvalidValue); + } + cudaError_t status = cudaSetDevice(device_id); + auto align8=[](size_t n){return (n+7)&~size_t(7);}; size_t bytes=0; + auto reserve=[&](size_t n){size_t at=bytes;bytes+=align8(n);return at;}; + const size_t on=reserve(node_count*sizeof(ConstraintNode)), oroot=reserve(root_count*sizeof(uint32_t)); + const size_t ol=reserve(lookup_count*sizeof(ConstraintLookup)), oa=reserve(lookup_arg_count*sizeof(uint32_t)); + const size_t op=reserve(public_count*sizeof(uint64_t)), os=reserve(4*quotient_size*sizeof(uint64_t)); + const size_t oalpha=reserve(2*constraint_count*sizeof(uint64_t)), od=reserve(2*sizeof(uint64_t)); + const size_t oo=reserve(2*quotient_size*sizeof(uint64_t)); uint8_t* allocation=nullptr; + if(status==cudaSuccess)status=cudaMalloc(reinterpret_cast(&allocation),bytes); + auto cp=[&](size_t at,const void* src,size_t n){if(status==cudaSuccess&&n!=0)status=cudaMemcpy(allocation+at,src,n,cudaMemcpyHostToDevice);}; + cp(on,nodes,node_count*sizeof(ConstraintNode));cp(oroot,roots,root_count*sizeof(uint32_t)); + cp(ol,lookups,lookup_count*sizeof(ConstraintLookup));cp(oa,lookup_args,lookup_arg_count*sizeof(uint32_t)); + cp(op,publics,public_count*sizeof(uint64_t)); + cp(oalpha,alpha,2*constraint_count*sizeof(uint64_t));cp(od,delta,2*sizeof(uint64_t)); + auto* dn=reinterpret_cast(allocation+on);auto* dr=reinterpret_cast(allocation+oroot); + auto* dl=reinterpret_cast(allocation+ol);auto* da=reinterpret_cast(allocation+oa); + auto* dp=reinterpret_cast(allocation+op);auto* ds=reinterpret_cast(allocation+os); + auto* dal=reinterpret_cast(allocation+oalpha);auto* dd=reinterpret_cast(allocation+od); + auto* dout=reinterpret_cast(allocation+oo); + if(status==cudaSuccess)status=generate_coset_selectors(ds,quotient_size,next_step, + coset_shift,coset_generator,trace_last,vanishing_start,vanishing_step); + size_t budget=0;if(status==cudaSuccess)status=quotient_shared_memory_budget(device_id,&budget); + size_t tile=budget/(slot_count*sizeof(uint64_t));bool global=tile<28; + if(global)tile=128;else if(tile>32)tile=32;size_t blocks=(quotient_size+tile-1)/tile;const size_t block_cap=global?256:1024;if(blocks>block_cap)blocks=block_cap; + uint64_t* scratch=nullptr;if(status==cudaSuccess&&global)status=cudaMalloc(reinterpret_cast(&scratch),blocks*slot_count*tile*sizeof(uint64_t)); + const size_t dynamic_shared=global?0:slot_count*tile*sizeof(uint64_t); + if(status==cudaSuccess&&dynamic_shared>48*1024) + status=configure_quotient_shared_memory(device_id,dynamic_shared); + if(status==cudaSuccess) { + evaluate_quotient<<(blocks),static_cast(tile),dynamic_shared>>>( + dout,dn,node_count,slot_count,dr,root_count,dl,lookup_count,da,group_size, + static_cast(preprocessed_handle),static_cast(main_handle), + static_cast(stage2_handle),dp,ds,dal,dd,ext_w,quotient_size,next_step,scratch); + status=cudaGetLastError(); + } + if(status==cudaSuccess) status=cudaMemcpy(output,dout,2*quotient_size*sizeof(uint64_t),cudaMemcpyDeviceToHost); + cudaFree(scratch); + cudaFree(allocation); + return static_cast(status); +} + +extern "C" int multi_stark_cuda_quotient_lde( + int device_id, void** output_handle, const void* nodes, size_t node_count, + size_t slot_count, const uint32_t* roots, size_t root_count, + const void* lookups, size_t lookup_count, const uint32_t* lookup_args, + size_t lookup_arg_count, size_t group_size, const void* preprocessed_handle, + const void* main_handle, const void* stage2_handle, const uint64_t* publics, + size_t public_count, uint64_t coset_shift, uint64_t coset_generator, + uint64_t trace_last, uint64_t vanishing_start, uint64_t vanishing_step, + const uint64_t* alpha, + size_t constraint_count, const uint64_t* delta, uint64_t ext_w, + size_t quotient_size, size_t next_step, size_t quotient_degree, + size_t log_blowup, const uint64_t* quotient_twiddles, + const uint64_t* lde_twiddles, const uint64_t* slice_weights) { + if (output_handle == nullptr || nodes == nullptr || roots == nullptr || + main_handle == nullptr || stage2_handle == nullptr || publics == nullptr || + alpha == nullptr || delta == nullptr || + quotient_twiddles == nullptr || lde_twiddles == nullptr || + slice_weights == nullptr || node_count == 0 || slot_count == 0 || + group_size == 0 || !is_power_of_two(quotient_size) || + !is_power_of_two(quotient_degree) || quotient_degree > quotient_size || + quotient_size % quotient_degree != 0 || !is_power_of_two(next_step) || + next_step > quotient_size || + log_blowup >= sizeof(size_t) * 8 || + (lookup_count != 0 && (lookups == nullptr || lookup_args == nullptr))) { + return static_cast(cudaErrorInvalidValue); + } + *output_handle = nullptr; + const size_t trace_height = quotient_size / quotient_degree; + if (trace_height > (SIZE_MAX >> log_blowup)) { + return static_cast(cudaErrorInvalidValue); + } + const size_t lde_height = trace_height << log_blowup; + const size_t width = 2 * quotient_degree; + if (!product_fits(lde_height, width)) { + return static_cast(cudaErrorInvalidValue); + } + + cudaError_t status = cudaSetDevice(device_id); + auto align8=[](size_t n){return (n+7)&~size_t(7);}; size_t bytes=0; + auto reserve=[&](size_t n){size_t at=bytes;bytes+=align8(n);return at;}; + const size_t on=reserve(node_count*sizeof(ConstraintNode)), oroot=reserve(root_count*sizeof(uint32_t)); + const size_t ol=reserve(lookup_count*sizeof(ConstraintLookup)), oa=reserve(lookup_arg_count*sizeof(uint32_t)); + const size_t op=reserve(public_count*sizeof(uint64_t)), os=reserve(4*quotient_size*sizeof(uint64_t)); + const size_t oalpha=reserve(2*constraint_count*sizeof(uint64_t)), od=reserve(2*sizeof(uint64_t)); + const size_t oo=reserve(2*quotient_size*sizeof(uint64_t)); uint8_t* allocation=nullptr; + if(status==cudaSuccess)status=cudaMalloc(reinterpret_cast(&allocation),bytes); + auto cp=[&](size_t at,const void* src,size_t n){if(status==cudaSuccess&&n!=0)status=cudaMemcpy(allocation+at,src,n,cudaMemcpyHostToDevice);}; + cp(on,nodes,node_count*sizeof(ConstraintNode));cp(oroot,roots,root_count*sizeof(uint32_t)); + cp(ol,lookups,lookup_count*sizeof(ConstraintLookup));cp(oa,lookup_args,lookup_arg_count*sizeof(uint32_t)); + cp(op,publics,public_count*sizeof(uint64_t)); + cp(oalpha,alpha,2*constraint_count*sizeof(uint64_t));cp(od,delta,2*sizeof(uint64_t)); + auto* dn=reinterpret_cast(allocation+on);auto* dr=reinterpret_cast(allocation+oroot); + auto* dl=reinterpret_cast(allocation+ol);auto* da=reinterpret_cast(allocation+oa); + auto* dp=reinterpret_cast(allocation+op);auto* ds=reinterpret_cast(allocation+os); + auto* dal=reinterpret_cast(allocation+oalpha);auto* dd=reinterpret_cast(allocation+od); + auto* quotient=reinterpret_cast(allocation+oo); + + if(status==cudaSuccess)status=generate_coset_selectors(ds,quotient_size,next_step, + coset_shift,coset_generator,trace_last,vanishing_start,vanishing_step); + + const uint64_t *device_quotient_twiddles=nullptr,*device_lde_twiddles=nullptr,*device_weights=nullptr; + if(status==cudaSuccess)status=cached_device_constants(device_id,quotient_twiddles,quotient_size/2,2,0,0,&device_quotient_twiddles); + if(status==cudaSuccess)status=cached_device_constants(device_id,lde_twiddles,lde_height/2,2,0,0,&device_lde_twiddles); + if(status==cudaSuccess)status=cached_device_constants(device_id,slice_weights,quotient_degree,4,slice_weights[0],quotient_degree>1?slice_weights[1]:0,&device_weights); + + size_t budget=0;if(status==cudaSuccess)status=quotient_shared_memory_budget(device_id,&budget); + size_t tile=budget/(slot_count*sizeof(uint64_t));bool global=tile<28; + if(global)tile=128;else if(tile>32)tile=32;size_t blocks=(quotient_size+tile-1)/tile;const size_t block_cap=global?256:1024;if(blocks>block_cap)blocks=block_cap; + uint64_t* scratch=nullptr;if(status==cudaSuccess&&global)status=cudaMalloc(reinterpret_cast(&scratch),blocks*slot_count*tile*sizeof(uint64_t)); + const size_t dynamic_shared=global?0:slot_count*tile*sizeof(uint64_t); + if(status==cudaSuccess&&dynamic_shared>48*1024) + status=configure_quotient_shared_memory(device_id,dynamic_shared); + if(status==cudaSuccess) { + evaluate_quotient<<(blocks),static_cast(tile),dynamic_shared>>>( + quotient,dn,node_count,slot_count,dr,root_count,dl,lookup_count,da,group_size, + static_cast(preprocessed_handle),static_cast(main_handle), + static_cast(stage2_handle),dp,ds,dal,dd,ext_w,quotient_size,next_step,scratch); + status=cudaGetLastError(); + } + if(status==cudaSuccess)status=launch_dif(quotient,quotient_size,2,device_quotient_twiddles); + + ResidentLde* lde = nullptr; + if(status==cudaSuccess) { + status=create_resident_lde(&lde); + } + if(status==cudaSuccess) { + lde->height=lde_height;lde->width=width; + status=cudaMalloc(reinterpret_cast(&lde->values),lde_height*width*sizeof(uint64_t)); + } + if(status==cudaSuccess)status=cudaMemset(lde->values,0,lde_height*width*sizeof(uint64_t)); + if(status==cudaSuccess) { + gather_shifted_quotient_slices<<>>( + lde->values,quotient,device_weights,quotient_size,trace_height, + quotient_degree,2); + status=cudaGetLastError(); + } + if(status==cudaSuccess)status=launch_dif(lde->values,lde_height,width,device_lde_twiddles); + if(status==cudaSuccess)status=cudaStreamSynchronize(0); + if(status==cudaSuccess) { + *output_handle=lde; + } else if(lde) { + destroy_resident_lde(lde); + } + cudaFree(scratch);cudaFree(allocation); + return static_cast(status); +} + +extern "C" int multi_stark_cuda_lde_interpolate(int device_id,uint64_t* output,const void* handle, + size_t height,const uint64_t* inv_denoms,const uint64_t* coset,const uint64_t* scale,uint64_t ext_w){ + if(!output||!handle||!inv_denoms||!coset||!scale)return static_cast(cudaErrorInvalidValue); + ResidentLde* l=const_cast(static_cast(handle));if(height==0||height>l->height)return static_cast(cudaErrorInvalidValue); + cudaError_t status=cudaSetDevice(device_id);const size_t ib=height*sizeof(Ext2),cb=height*sizeof(uint64_t),ob=l->width*sizeof(Ext2),needed=ib+cb+ob; + if(status==cudaSuccess&&l->interpolation_scratch_bytesinterpolation_scratch)status=cudaFree(l->interpolation_scratch);l->interpolation_scratch=nullptr;if(status==cudaSuccess)status=cudaMalloc(reinterpret_cast(&l->interpolation_scratch),needed);if(status==cudaSuccess)l->interpolation_scratch_bytes=needed;} + uint8_t* mem=l->interpolation_scratch; + if(status==cudaSuccess)status=cudaMemcpy(mem,inv_denoms,ib,cudaMemcpyHostToDevice); + if(status==cudaSuccess)status=cudaMemcpy(mem+ib,coset,cb,cudaMemcpyHostToDevice); + if(status==cudaSuccess){interpolate_lde_columns<<(l->width),THREADS,THREADS*sizeof(Ext2)>>>(reinterpret_cast(mem+ib+cb),l,reinterpret_cast(mem),reinterpret_cast(mem+ib),height,{scale[0],scale[1]},ext_w);status=cudaGetLastError();} + if(status==cudaSuccess)status=cudaMemcpy(output,mem+ib+cb,ob,cudaMemcpyDeviceToHost);return static_cast(status); +} + +extern "C" int multi_stark_cuda_fri_workspace_create(int device_id,void** handle, + const uint64_t* points,const size_t* counts,size_t point_count, + const uint64_t* coset,size_t coset_count,uint64_t ext_w){ + if(!handle||!points||!counts||!point_count||!coset||!coset_count) + return static_cast(cudaErrorInvalidValue);*handle=nullptr; + size_t inv_count=0,max_count=0; + for(size_t i=0;icoset_count||inv_count>SIZE_MAX-counts[i]) + return static_cast(cudaErrorInvalidValue); + inv_count+=counts[i];if(counts[i]>max_count)max_count=counts[i]; + } + cudaError_t status=cudaSetDevice(device_id);auto* w=new(std::nothrow) ResidentFriWorkspace;if(!w)return static_cast(cudaErrorMemoryAllocation); + w->inv_count=inv_count;w->coset_count=coset_count; + if(status==cudaSuccess)status=cudaMalloc(reinterpret_cast(&w->inv_denoms),inv_count*sizeof(Ext2)); + if(status==cudaSuccess)status=cudaMalloc(reinterpret_cast(&w->coset),coset_count*sizeof(uint64_t)); + if(status==cudaSuccess)status=cudaMemcpy(w->coset,coset,coset_count*sizeof(uint64_t),cudaMemcpyHostToDevice); + uint64_t *norms=nullptr,*inverses=nullptr; + if(status==cudaSuccess)status=cudaMalloc(reinterpret_cast(&norms),max_count*sizeof(uint64_t)); + if(status==cudaSuccess)status=cudaMalloc(reinterpret_cast(&inverses),max_count*sizeof(uint64_t)); + size_t offset=0; + for(size_t i=0;status==cudaSuccess&&i>>(norms,w->coset,counts[i],point,ext_w);status=cudaGetLastError(); + if(status==cudaSuccess)status=batch_inverse_norms(inverses,norms,counts[i]); + if(status==cudaSuccess){finish_inverse_denominators<<>>(w->inv_denoms+offset,inverses,w->coset,counts[i],point);status=cudaGetLastError();} + offset+=counts[i]; + } + if(inverses)cudaFree(inverses);if(norms)cudaFree(norms); + if(status!=cudaSuccess){delete w;return static_cast(status);}*handle=w;return static_cast(cudaSuccess); +} + +extern "C" int multi_stark_cuda_fri_interpolate_batch(int device_id,void* handle,uint64_t* output, + size_t output_count,const InterpolationTask* tasks,size_t task_count,uint64_t ext_w){ + if(!handle||!output||!output_count||!tasks||!task_count)return static_cast(cudaErrorInvalidValue); + cudaError_t status=cudaSetDevice(device_id);auto* w=static_cast(handle); + if(status==cudaSuccess&&w->output_capacityoutput)status=cudaFree(w->output);w->output=nullptr;if(status==cudaSuccess)status=cudaMalloc(reinterpret_cast(&w->output),output_count*sizeof(Ext2));if(status==cudaSuccess)w->output_capacity=output_count;} + size_t partial_count=0,block_count=0,finish_count=0; + for(size_t i=0;status==cudaSuccess&&i(t.lde); + if(!l||!t.height||t.height>l->height||t.inv_offset>w->inv_count-t.height||t.output_offset>output_count-l->width){status=cudaErrorInvalidValue;break;} + const size_t partial_rows=(t.height+INTERPOLATION_ROWS-1)/INTERPOLATION_ROWS; + const bool pair=i+1inv_count-t.height&&tasks[i+1].output_offset<=output_count-l->width; + const size_t columns=(l->width+INTERPOLATION_COLUMNS-1)/INTERPOLATION_COLUMNS; + if(partial_rows>SIZE_MAX/l->width||partial_count>SIZE_MAX-partial_rows*l->width*(pair?2:1)|| + columns&&partial_rows>SIZE_MAX/columns||block_count>SIZE_MAX-partial_rows*columns|| + finish_count>SIZE_MAX-l->width*(pair?2:1)){status=cudaErrorInvalidValue;break;} + partial_count+=partial_rows*l->width*(pair?2:1);block_count+=partial_rows*columns; + finish_count+=l->width*(pair?2:1);i+=pair?2:1;} + if(status==cudaSuccess&&w->partial_capacityinterpolation_partials)status=cudaFree(w->interpolation_partials);w->interpolation_partials=nullptr; + if(status==cudaSuccess)status=cudaMalloc(reinterpret_cast(&w->interpolation_partials),partial_count*sizeof(Ext2)); + if(status==cudaSuccess)w->partial_capacity=partial_count;} + InterpolationBlockDesc* host_blocks=status==cudaSuccess?new(std::nothrow) InterpolationBlockDesc[block_count]:nullptr; + InterpolationFinishDesc* host_finishes=status==cudaSuccess?new(std::nothrow) InterpolationFinishDesc[finish_count]:nullptr; + if(status==cudaSuccess&&(!host_blocks||!host_finishes))status=cudaErrorMemoryAllocation; + size_t partial_at=0,block_at=0,finish_at=0; + for(size_t i=0;status==cudaSuccess&&i(t.lde); + const size_t partial_rows=(t.height+INTERPOLATION_ROWS-1)/INTERPOLATION_ROWS; + const size_t one_count=partial_rows*l->width; + const bool pair=i+1inv_count-t.height&&tasks[i+1].output_offset<=output_count-l->width; + Ext2* p0=w->interpolation_partials+partial_at;Ext2* p1=pair?p0+one_count:nullptr; + for(size_t row=0;rowwidth;column+=INTERPOLATION_COLUMNS) + host_blocks[block_at++]={l->values,w->coset,w->inv_denoms+t.inv_offset, + pair?w->inv_denoms+tasks[i+1].inv_offset:nullptr,p0,p1,l->width,column,row, + row+INTERPOLATION_ROWSwidth;++column) + host_finishes[finish_at++]={w->output+t.output_offset+column,p0,partial_rows,l->width,column,{t.scale0,t.scale1}}; + if(pair){const auto& u=tasks[i+1];for(size_t column=0;columnwidth;++column) + host_finishes[finish_at++]={w->output+u.output_offset+column,p1,partial_rows,l->width,column,{u.scale0,u.scale1}};} + partial_at+=one_count*(pair?2:1);i+=pair?2:1;} + InterpolationBlockDesc* device_blocks=nullptr;InterpolationFinishDesc* device_finishes=nullptr; + if(status==cudaSuccess)status=cudaMalloc(reinterpret_cast(&device_blocks),block_count*sizeof(InterpolationBlockDesc)); + if(status==cudaSuccess)status=cudaMalloc(reinterpret_cast(&device_finishes),finish_count*sizeof(InterpolationFinishDesc)); + if(status==cudaSuccess)status=cudaMemcpyAsync(device_blocks,host_blocks,block_count*sizeof(InterpolationBlockDesc),cudaMemcpyHostToDevice,cudaStreamPerThread); + if(status==cudaSuccess)status=cudaMemcpyAsync(device_finishes,host_finishes,finish_count*sizeof(InterpolationFinishDesc),cudaMemcpyHostToDevice,cudaStreamPerThread); + if(status==cudaSuccess)status=cudaStreamSynchronize(cudaStreamPerThread); + delete[] host_finishes;host_finishes=nullptr;delete[] host_blocks;host_blocks=nullptr; + if(status==cudaSuccess){const dim3 block(INTERPOLATION_COLUMNS,INTERPOLATION_LANES); + interpolate_lde_blocks<<(block_count),block>>>(device_blocks,ext_w);status=cudaGetLastError();} + if(status==cudaSuccess){finish_lde_interpolation_blocks<<(finish_count),THREADS,THREADS*sizeof(Ext2)>>>(device_finishes,ext_w);status=cudaGetLastError();} + if(device_finishes)cudaFree(device_finishes);if(device_blocks)cudaFree(device_blocks); + delete[] host_finishes;delete[] host_blocks; + if(status==cudaSuccess)status=cudaMemcpy(output,w->output,output_count*sizeof(Ext2),cudaMemcpyDeviceToHost);return static_cast(status); +} + +extern "C" int multi_stark_cuda_fri_reduce_batch(int device_id,void* handle, + const ReductionTask* tasks,size_t task_count,const uint64_t* alpha,size_t alpha_count,uint64_t ext_w){ + if(!handle||!tasks||!task_count||!alpha||!alpha_count)return static_cast(cudaErrorInvalidValue); + cudaError_t status=cudaSetDevice(device_id);auto* w=static_cast(handle); + if(status==cudaSuccess&&w->alpha_capacityalpha)status=cudaFree(w->alpha);w->alpha=nullptr;if(status==cudaSuccess)status=cudaMalloc(reinterpret_cast(&w->alpha),alpha_count*sizeof(Ext2));if(status==cudaSuccess)w->alpha_capacity=alpha_count;} + if(status==cudaSuccess)status=cudaMemcpy(w->alpha,alpha,alpha_count*sizeof(Ext2),cudaMemcpyHostToDevice); + for(size_t i=0;status==cudaSuccess&&i(t.reduced);auto* l=static_cast(t.lde); + if(!r||!l||t.height!=r->height||t.height>l->height||t.inv_offset+t.height>w->inv_count||l->width>alpha_count){status=cudaErrorInvalidValue;break;} + accumulate_reduced_opening<<>>(r->values,l,w->inv_denoms+t.inv_offset,w->alpha,t.height,{t.y0,t.y1},{t.offset0,t.offset1},ext_w);status=cudaGetLastError();} + return static_cast(status); +} + +extern "C" int multi_stark_cuda_fri_workspace_destroy(int device_id,void* handle){ + cudaError_t status=cudaSetDevice(device_id);if(status==cudaSuccess)delete static_cast(handle);return static_cast(status); +} + +extern "C" int multi_stark_cuda_reduced_to_lde(int device_id,void** output,const void* reduced){ + if(!output||!reduced)return static_cast(cudaErrorInvalidValue);*output=nullptr; + auto* r=static_cast(reduced);cudaError_t status=cudaSetDevice(device_id); + ResidentLde* l=nullptr;if(status==cudaSuccess)status=create_resident_lde(&l);if(status==cudaSuccess){l->height=r->height;l->width=2;} + if(status==cudaSuccess)status=cudaMalloc(reinterpret_cast(&l->values),r->height*sizeof(Ext2)); + if(status==cudaSuccess)status=cudaMemcpy(l->values,r->values,r->height*sizeof(Ext2),cudaMemcpyDeviceToDevice); + if(status!=cudaSuccess){destroy_resident_lde(l);return static_cast(status);}*output=l;return static_cast(cudaSuccess); +} + +extern "C" int multi_stark_cuda_fri_fold_resident(int device_id,void** output,const void* input, + const void* next_reduced,const uint64_t* beta,size_t log_arity,uint64_t beta_power0, + uint64_t beta_power1,uint64_t g_inv,uint64_t ext_w){ + if(!output||!input||!beta||log_arity!=1)return static_cast(cudaErrorInvalidValue);*output=nullptr; + auto* in=static_cast(input);if(in->width!=2||in->height%2)return static_cast(cudaErrorInvalidValue); + const size_t final_height=in->height>>log_arity;auto* next=static_cast(next_reduced); + if(next&&next->height!=final_height)return static_cast(cudaErrorInvalidValue); + cudaError_t status=cudaSetDevice(device_id);Ext2 *a=reinterpret_cast(in->values),*first=nullptr;uint64_t* powers=nullptr; + if(status==cudaSuccess)status=cudaMalloc(reinterpret_cast(&first),(in->height/2)*sizeof(Ext2)); + if(status==cudaSuccess)status=cudaMalloc(reinterpret_cast(&powers),(in->height/2)*sizeof(uint64_t)); + if(status==cudaSuccess){init_fri_powers<<height/2),THREADS>>>(powers,in->height/2,g_inv);status=cudaGetLastError();} + if(status==cudaSuccess){fold_fri_ext2<<>>(first,a,powers,final_height,{beta[0],beta[1]},ext_w);status=cudaGetLastError();} + a=first; + if(status==cudaSuccess&&next){add_scaled_ext2<<>>(a,next->values,final_height,{beta_power0,beta_power1},ext_w);status=cudaGetLastError();} + ResidentLde* result=nullptr;if(status==cudaSuccess)status=create_resident_lde(&result); + if(status==cudaSuccess){result->values=reinterpret_cast(a);result->height=final_height;result->width=2;first=nullptr;*output=result;}else destroy_resident_lde(result); + cudaFree(powers);cudaFree(first);return static_cast(status); +} + +extern "C" int multi_stark_cuda_reduced_create(int device_id,void** handle,size_t height){ + if(!handle||!is_power_of_two(height))return static_cast(cudaErrorInvalidValue);*handle=nullptr;cudaError_t status=cudaSetDevice(device_id); + auto* r=new(std::nothrow) ResidentReducedOpening;if(!r)return static_cast(cudaErrorMemoryAllocation);r->height=height; + if(status==cudaSuccess)status=cudaMalloc(reinterpret_cast(&r->values),height*sizeof(Ext2));if(status==cudaSuccess)status=cudaMemset(r->values,0,height*sizeof(Ext2)); + if(status!=cudaSuccess){delete r;return static_cast(status);}*handle=r;return static_cast(cudaSuccess); +} +extern "C" int multi_stark_cuda_lookup_trace(int device_id,uint64_t* output,uint64_t* total, + const uint64_t* multiplicities,const uint64_t* args,const size_t* arg_offsets, + size_t height,size_t num_lookups,size_t args_width,size_t group_size, + const uint64_t* beta,const uint64_t* gamma,uint64_t ext_w){ + if(!output||!total||!multiplicities||!arg_offsets||!height||!num_lookups||!group_size||!beta||!gamma||(args_width&& !args))return static_cast(cudaErrorInvalidValue); + const size_t slots=(num_lookups+group_size-1)/group_size,count=height*slots;cudaError_t status=cudaSetDevice(device_id); + const size_t message_count=height*num_lookups;uint64_t *dm=nullptr,*da=nullptr,*norms=nullptr,*norm_inverses=nullptr;size_t* offsets=nullptr;Ext2 *conjugates=nullptr,*deltas=nullptr,*scan=nullptr; + if(status==cudaSuccess)status=cudaMalloc(reinterpret_cast(&dm),height*num_lookups*sizeof(uint64_t)); + if(status==cudaSuccess&&args_width)status=cudaMalloc(reinterpret_cast(&da),height*args_width*sizeof(uint64_t)); + if(status==cudaSuccess)status=cudaMalloc(reinterpret_cast(&offsets),(num_lookups+1)*sizeof(size_t)); + if(status==cudaSuccess)status=cudaMalloc(reinterpret_cast(&deltas),count*sizeof(Ext2)); + if(status==cudaSuccess)status=cudaMalloc(reinterpret_cast(&scan),count*sizeof(Ext2)); + if(status==cudaSuccess)status=cudaMalloc(reinterpret_cast(&conjugates),message_count*sizeof(Ext2)); + if(status==cudaSuccess)status=cudaMalloc(reinterpret_cast(&norms),message_count*sizeof(uint64_t)); + if(status==cudaSuccess)status=cudaMalloc(reinterpret_cast(&norm_inverses),message_count*sizeof(uint64_t)); + if(status==cudaSuccess)status=cudaMemcpyAsync(dm,multiplicities,height*num_lookups*sizeof(uint64_t),cudaMemcpyHostToDevice,0); + if(status==cudaSuccess&&args_width)status=cudaMemcpyAsync(da,args,height*args_width*sizeof(uint64_t),cudaMemcpyHostToDevice,0); + if(status==cudaSuccess)status=cudaMemcpyAsync(offsets,arg_offsets,(num_lookups+1)*sizeof(size_t),cudaMemcpyHostToDevice,0); + if(status==cudaSuccess){lookup_messages<<>>(conjugates,norms,da,offsets,height,num_lookups,args_width,{beta[0],beta[1]},{gamma[0],gamma[1]},ext_w);status=cudaGetLastError();} + if(status==cudaSuccess)status=batch_inverse_norms(norm_inverses,norms,message_count); + if(status==cudaSuccess){lookup_group_deltas_batched<<>>(deltas,dm,conjugates,norm_inverses,height,num_lookups,group_size,ext_w);status=cudaGetLastError();} + if(status==cudaSuccess)status=exclusive_scan_ext2(scan,deltas,count); + if(status==cudaSuccess)status=cudaMemcpy(output,scan,count*sizeof(Ext2),cudaMemcpyDeviceToHost); + if(status==cudaSuccess)status=cudaMemcpy(total,scan+count-1,sizeof(Ext2),cudaMemcpyDeviceToHost);if(status==cudaSuccess)status=cudaMemcpy(total+2,deltas+count-1,sizeof(Ext2),cudaMemcpyDeviceToHost); + cudaFree(norm_inverses);cudaFree(norms);cudaFree(conjugates);cudaFree(scan);cudaFree(deltas);cudaFree(offsets);cudaFree(da);cudaFree(dm);return static_cast(status); +} + +extern "C" int multi_stark_cuda_lookup_graph_lde(int device_id,void** output_handle,uint64_t* total, + const void* nodes,size_t node_count,size_t slot_count,const void* lookups,size_t lookup_count, + const uint32_t* lookup_args,size_t lookup_arg_count,const void* preprocessed_handle, + const void* main_handle,size_t group_size,const uint64_t* beta,const uint64_t* gamma, + uint64_t ext_w,size_t added_bits,const uint64_t* inverse_twiddles, + const uint64_t* shift_powers,const uint64_t* forward_twiddles,uint64_t height_inverse){ + auto* main=const_cast(static_cast(main_handle)); + auto* prep=static_cast(preprocessed_handle); + if(!output_handle||!total||!nodes||!node_count||!slot_count||!lookups||!lookup_count|| + (lookup_arg_count&&!lookup_args)||!main|| + (!main->trace_values&&!main->host_trace_values)||!main->trace_height|| + !group_size||!beta||!gamma||!inverse_twiddles||!shift_powers||!forward_twiddles|| + !is_power_of_two(main->trace_height)||added_bits>=sizeof(size_t)*8|| + main->trace_height>(SIZE_MAX>>added_bits)||(prep&&!prep->trace_values)) { + return static_cast(cudaErrorInvalidValue); + } + *output_handle=nullptr;const size_t height=main->trace_height; + const size_t groups=(lookup_count+group_size-1)/group_size,width=2*groups; + const size_t count=height*groups,extended_height=height<(cudaErrorInvalidValue); + cudaError_t status=cudaSetDevice(device_id);uint8_t* metadata=nullptr; + auto align8=[](size_t n){return (n+7)&~size_t(7);};size_t bytes=0; + auto reserve=[&](size_t n){const size_t at=bytes;bytes+=align8(n);return at;}; + const size_t on=reserve(node_count*sizeof(ConstraintNode)); + const size_t ol=reserve(lookup_count*sizeof(ConstraintLookup)); + const size_t oa=reserve(lookup_arg_count*sizeof(uint32_t)); + if(status==cudaSuccess)status=cudaMalloc(reinterpret_cast(&metadata),bytes); + auto cp=[&](size_t at,const void* src,size_t n){if(status==cudaSuccess)status=cudaMemcpy(metadata+at,src,n,cudaMemcpyHostToDevice);}; + cp(on,nodes,node_count*sizeof(ConstraintNode));cp(ol,lookups,lookup_count*sizeof(ConstraintLookup));cp(oa,lookup_args,lookup_arg_count*sizeof(uint32_t)); + auto* dn=reinterpret_cast(metadata+on);auto* dl=reinterpret_cast(metadata+ol);auto* da=reinterpret_cast(metadata+oa); + Ext2 *conjugates=nullptr,*deltas=nullptr;uint64_t *norms=nullptr,*norm_inverses=nullptr,*multiplicities=nullptr,*scratch=nullptr,*trace_chunk=nullptr;ResidentLde* lde=nullptr; + if(status==cudaSuccess)status=cudaMalloc(reinterpret_cast(&conjugates),message_count*sizeof(Ext2)); + if(status==cudaSuccess)status=cudaMalloc(reinterpret_cast(&norms),message_count*sizeof(uint64_t)); + if(status==cudaSuccess)status=cudaMalloc(reinterpret_cast(&norm_inverses),message_count*sizeof(uint64_t)); + if(status==cudaSuccess)status=cudaMalloc(reinterpret_cast(&multiplicities),message_count*sizeof(uint64_t)); + if(status==cudaSuccess)status=cudaMalloc(reinterpret_cast(&deltas),count*sizeof(Ext2)); + if(status==cudaSuccess)status=create_resident_lde(&lde); + if(status==cudaSuccess){lde->height=extended_height;lde->width=width;status=cudaMalloc(reinterpret_cast(&lde->values),extended_height*width*sizeof(uint64_t));} + if(status==cudaSuccess)status=cudaMemset(lde->values,0,extended_height*width*sizeof(uint64_t)); + const size_t budget=48*1024;size_t tile=budget/(slot_count*sizeof(uint64_t));const bool global=tile<32; + if(global)tile=128;else if(tile>32)tile=32;size_t blocks=(chunk_rows+tile-1)/tile;const size_t cap=global?256:1024;if(blocks>cap)blocks=cap; + if(status==cudaSuccess&&global)status=cudaMalloc(reinterpret_cast(&scratch),blocks*slot_count*tile*sizeof(uint64_t)); + if(status==cudaSuccess&&main->host_trace_values)status=cudaMalloc(reinterpret_cast(&trace_chunk),(chunk_rows+1)*main->width*sizeof(uint64_t)); + for(size_t row_start=0;status==cudaSuccess&&row_starttrace_values; + if(main->host_trace_values){ + status=cudaMemcpy(trace_chunk,main->host_trace_values+row_start*main->width,rows*main->width*sizeof(uint64_t),cudaMemcpyHostToDevice); + const size_t next_row=(row_start+rows)&(height-1); + if(status==cudaSuccess)status=cudaMemcpy(trace_chunk+rows*main->width,main->host_trace_values+next_row*main->width,main->width*sizeof(uint64_t),cudaMemcpyHostToDevice); + active_trace=trace_chunk; + } + if(status==cudaSuccess){lookup_messages_graph<<(chunk_blocks),static_cast(tile),global?0:slot_count*tile*sizeof(uint64_t)>>>( + conjugates,norms,multiplicities,dn,node_count,slot_count,dl,lookup_count,da,prep,active_trace,main->width, + main->host_trace_values!=nullptr,{beta[0],beta[1]},{gamma[0],gamma[1]},ext_w,height,row_start,rows,scratch);status=cudaGetLastError();} + if(status==cudaSuccess)status=batch_inverse_norms(norm_inverses,norms,messages); + if(status==cudaSuccess){lookup_group_deltas_batched<<>>( + deltas+row_start*groups,multiplicities,conjugates,norm_inverses,rows,lookup_count,group_size,ext_w);status=cudaGetLastError();} + } + if(status==cudaSuccess)status=exclusive_scan_ext2(reinterpret_cast(lde->values),deltas,count); + if(status==cudaSuccess)status=cudaMemcpy(total,lde->values+2*(count-1),sizeof(Ext2),cudaMemcpyDeviceToHost); + if(status==cudaSuccess)status=cudaMemcpy(total+2,deltas+count-1,sizeof(Ext2),cudaMemcpyDeviceToHost); + const uint64_t *dit=nullptr,*dshift=nullptr,*dft=nullptr; + if(status==cudaSuccess)status=cached_device_constants(device_id,inverse_twiddles,height/2,1,0,0,&dit); + if(status==cudaSuccess)status=cached_device_constants(device_id,shift_powers,height,3,shift_powers[0],height>1?shift_powers[1]:0,&dshift); + if(status==cudaSuccess)status=cached_device_constants(device_id,forward_twiddles,extended_height/2,2,0,0,&dft); + if(status==cudaSuccess)status=launch_dif(lde->values,height,width,dit); + if(status==cudaSuccess){bit_reverse_scale_and_shift<<>>(lde->values,height,width,strict_log2(height),height_inverse,dshift);status=cudaGetLastError();} + if(status==cudaSuccess)status=launch_dif(lde->values,extended_height,width,dft); + if(status==cudaSuccess){canonicalize_goldilocks<<>>(lde->values,extended_height*width);status=cudaGetLastError();} + if(status==cudaSuccess)status=cudaStreamSynchronize(0); + if(status==cudaSuccess)*output_handle=lde;else destroy_resident_lde(lde); + cudaFree(trace_chunk);cudaFree(scratch);cudaFree(deltas);cudaFree(multiplicities);cudaFree(norm_inverses);cudaFree(norms);cudaFree(conjugates);cudaFree(metadata); + return static_cast(status); +} + +extern "C" int multi_stark_cuda_lookup_lde(int device_id,void** output_handle,uint64_t* total, + const uint64_t* multiplicities,const uint64_t* args,const size_t* arg_offsets, + size_t height,size_t num_lookups,size_t args_width,size_t group_size, + const uint64_t* beta,const uint64_t* gamma,uint64_t ext_w,size_t added_bits, + const uint64_t* inverse_twiddles,const uint64_t* shift_powers, + const uint64_t* forward_twiddles,uint64_t height_inverse){ + if(!output_handle||!total||!multiplicities||!arg_offsets||!height||!num_lookups|| + !group_size||!beta||!gamma||!inverse_twiddles||!shift_powers|| + !forward_twiddles||(args_width&&!args)||!is_power_of_two(height)|| + added_bits>=sizeof(size_t)*8||height>(SIZE_MAX>>added_bits)) + return static_cast(cudaErrorInvalidValue); + *output_handle=nullptr;const size_t slots=(num_lookups+group_size-1)/group_size; + const size_t width=2*slots,count=height*slots,extended_height=height<(cudaErrorInvalidValue); + cudaError_t status=cudaSetDevice(device_id);const size_t message_count=height*num_lookups; + uint64_t *dm=nullptr,*da=nullptr,*norms=nullptr,*norm_inverses=nullptr;size_t* offsets=nullptr; + Ext2 *conjugates=nullptr,*deltas=nullptr;ResidentLde* lde=nullptr; + if(status==cudaSuccess)status=cudaMalloc(reinterpret_cast(&dm),height*num_lookups*sizeof(uint64_t)); + if(status==cudaSuccess&&args_width)status=cudaMalloc(reinterpret_cast(&da),height*args_width*sizeof(uint64_t)); + if(status==cudaSuccess)status=cudaMalloc(reinterpret_cast(&offsets),(num_lookups+1)*sizeof(size_t)); + if(status==cudaSuccess)status=cudaMalloc(reinterpret_cast(&deltas),count*sizeof(Ext2)); + if(status==cudaSuccess)status=cudaMalloc(reinterpret_cast(&conjugates),message_count*sizeof(Ext2)); + if(status==cudaSuccess)status=cudaMalloc(reinterpret_cast(&norms),message_count*sizeof(uint64_t)); + if(status==cudaSuccess)status=cudaMalloc(reinterpret_cast(&norm_inverses),message_count*sizeof(uint64_t)); + if(status==cudaSuccess)status=create_resident_lde(&lde); + if(status==cudaSuccess){lde->height=extended_height;lde->width=width;status=cudaMalloc(reinterpret_cast(&lde->values),extended_height*width*sizeof(uint64_t));} + if(status==cudaSuccess)status=cudaMemset(lde->values,0,extended_height*width*sizeof(uint64_t)); + if(status==cudaSuccess)status=cudaMemcpy(dm,multiplicities,height*num_lookups*sizeof(uint64_t),cudaMemcpyHostToDevice); + if(status==cudaSuccess&&args_width)status=cudaMemcpy(da,args,height*args_width*sizeof(uint64_t),cudaMemcpyHostToDevice); + if(status==cudaSuccess)status=cudaMemcpy(offsets,arg_offsets,(num_lookups+1)*sizeof(size_t),cudaMemcpyHostToDevice); + if(status==cudaSuccess){lookup_messages<<>>(conjugates,norms,da,offsets,height,num_lookups,args_width,{beta[0],beta[1]},{gamma[0],gamma[1]},ext_w);status=cudaGetLastError();} + if(status==cudaSuccess)status=batch_inverse_norms(norm_inverses,norms,message_count); + if(status==cudaSuccess){lookup_group_deltas_batched<<>>(deltas,dm,conjugates,norm_inverses,height,num_lookups,group_size,ext_w);status=cudaGetLastError();} + if(status==cudaSuccess)status=exclusive_scan_ext2(reinterpret_cast(lde->values),deltas,count); + if(status==cudaSuccess)status=cudaMemcpy(total,lde->values+2*(count-1),sizeof(Ext2),cudaMemcpyDeviceToHost); + if(status==cudaSuccess)status=cudaMemcpy(total+2,deltas+count-1,sizeof(Ext2),cudaMemcpyDeviceToHost); + const uint64_t *dit=nullptr,*dshift=nullptr,*dft=nullptr; + if(status==cudaSuccess)status=cached_device_constants(device_id,inverse_twiddles,height/2,1,0,0,&dit); + if(status==cudaSuccess)status=cached_device_constants(device_id,shift_powers,height,3,shift_powers[0],height>1?shift_powers[1]:0,&dshift); + if(status==cudaSuccess)status=cached_device_constants(device_id,forward_twiddles,extended_height/2,2,0,0,&dft); + if(status==cudaSuccess)status=launch_dif(lde->values,height,width,dit); + if(status==cudaSuccess){bit_reverse_scale_and_shift<<>>(lde->values,height,width,strict_log2(height),height_inverse,dshift);status=cudaGetLastError();} + if(status==cudaSuccess)status=launch_dif(lde->values,extended_height,width,dft); + if(status==cudaSuccess){canonicalize_goldilocks<<>>(lde->values,extended_height*width);status=cudaGetLastError();} + if(status==cudaSuccess)status=cudaStreamSynchronize(0); + if(status==cudaSuccess)*output_handle=lde;else destroy_resident_lde(lde); + cudaFree(norm_inverses);cudaFree(norms);cudaFree(conjugates);cudaFree(deltas); + cudaFree(offsets);cudaFree(da);cudaFree(dm);return static_cast(status); +} +extern "C" int multi_stark_cuda_reduced_add(int device_id,void* reduced,const void* lde_handle,size_t height, + const uint64_t* inv_denoms,const uint64_t* alpha_powers,const uint64_t* reduced_y,const uint64_t* alpha_offset,uint64_t ext_w){ + if(!reduced||!lde_handle||!inv_denoms||!alpha_powers||!reduced_y||!alpha_offset)return static_cast(cudaErrorInvalidValue); + auto* r=static_cast(reduced);auto* l=static_cast(lde_handle);if(height!=r->height||height>l->height)return static_cast(cudaErrorInvalidValue); + cudaError_t status=cudaSetDevice(device_id);size_t ib=height*sizeof(Ext2),ab=l->width*sizeof(Ext2),needed=ib+ab; + if(status==cudaSuccess&&r->scratch_bytesscratch)status=cudaFree(r->scratch);r->scratch=nullptr;if(status==cudaSuccess)status=cudaMalloc(reinterpret_cast(&r->scratch),needed);if(status==cudaSuccess)r->scratch_bytes=needed;} + uint8_t* mem=r->scratch;if(status==cudaSuccess)status=cudaMemcpy(mem,inv_denoms,ib,cudaMemcpyHostToDevice);if(status==cudaSuccess)status=cudaMemcpy(mem+ib,alpha_powers,ab,cudaMemcpyHostToDevice); + if(status==cudaSuccess){accumulate_reduced_opening<<>>(r->values,l,reinterpret_cast(mem),reinterpret_cast(mem+ib),height,{reduced_y[0],reduced_y[1]},{alpha_offset[0],alpha_offset[1]},ext_w);status=cudaGetLastError();}return static_cast(status); +} +extern "C" int multi_stark_cuda_reduced_copy(int device_id,const void* handle,uint64_t* output){if(!handle||!output)return static_cast(cudaErrorInvalidValue);auto* r=static_cast(handle);cudaError_t status=cudaSetDevice(device_id);if(status==cudaSuccess)status=cudaMemcpy(output,r->values,r->height*sizeof(Ext2),cudaMemcpyDeviceToHost);return static_cast(status);} +extern "C" int multi_stark_cuda_reduced_destroy(int device_id,void* handle){cudaError_t status=cudaSetDevice(device_id);if(status==cudaSuccess)delete static_cast(handle);return static_cast(status);} + +extern "C" int multi_stark_cuda_lde_destroy(int device_id, void* handle) { + if (handle == nullptr) { + return static_cast(cudaSuccess); + } + const cudaError_t status = cudaSetDevice(device_id); + if (status != cudaSuccess) { + return static_cast(status); + } + return static_cast(destroy_resident_lde(static_cast(handle))); +} + +extern "C" int multi_stark_cuda_lde_release_trace(int device_id,void* handle){ + if(!handle)return static_cast(cudaSuccess);cudaError_t status=cudaSetDevice(device_id); + auto* lde=static_cast(handle);if(status==cudaSuccess&&lde->trace_values){status=cudaFree(lde->trace_values);lde->trace_values=nullptr;} + if(status==cudaSuccess&&lde->host_trace_registered){status=cudaHostUnregister(const_cast(lde->host_trace_values));lde->host_trace_registered=false;} + if(status==cudaSuccess){lde->host_trace_values=nullptr;lde->trace_height=0;} + return static_cast(status); +} + +extern "C" int multi_stark_cuda_lde_attach_trace( + int device_id, void* handle, const uint64_t* trace, size_t height, + size_t width) { + if (!handle || !trace || height == 0 || width == 0 || + !product_fits(height, width)) { + return static_cast(cudaErrorInvalidValue); + } + cudaError_t status = cudaSetDevice(device_id); + auto* lde = static_cast(handle); + if (status != cudaSuccess || lde->width != width) { + return status == cudaSuccess ? static_cast(cudaErrorInvalidValue) + : static_cast(status); + } + if (lde->trace_values != nullptr) return static_cast(cudaSuccess); + if (status == cudaSuccess) { + lde->host_trace_values = trace; + lde->trace_height = height; + const size_t bytes=height*width*sizeof(uint64_t); + if(bytes>=(size_t(8)<<20)){ + const cudaError_t registration=cudaHostRegister( + const_cast(trace),bytes,cudaHostRegisterDefault); + if(registration==cudaSuccess)lde->host_trace_registered=true; + else cudaGetLastError(); + } + } + return static_cast(status); +} + +extern "C" int multi_stark_cuda_goldilocks_ops( + int device_id, uint64_t* sums, uint64_t* differences, uint64_t* products, + uint64_t* inverses, const uint64_t* left, const uint64_t* right, + size_t len) { + if (sums == nullptr || differences == nullptr || products == nullptr || + inverses == nullptr || left == nullptr || right == nullptr || len == 0) { + return static_cast(cudaErrorInvalidValue); + } + cudaError_t status = cudaSetDevice(device_id); + if (status != cudaSuccess) { + return static_cast(status); + } + + DeviceBuffer device_left; + DeviceBuffer device_right; + DeviceBuffer device_sums; + DeviceBuffer device_differences; + DeviceBuffer device_products; + DeviceBuffer device_inverses; + status = copy_to_device(device_left, left, len); + if (status == cudaSuccess) { + status = copy_to_device(device_right, right, len); + } + if (status == cudaSuccess) { + status = device_sums.allocate(len); + } + if (status == cudaSuccess) { + status = device_differences.allocate(len); + } + if (status == cudaSuccess) { + status = device_products.allocate(len); + } + if (status == cudaSuccess) { + status = device_inverses.allocate(len); + } + if (status == cudaSuccess) { + goldilocks_ops_kernel<<>>( + device_sums.get(), device_differences.get(), device_products.get(), + device_inverses.get(), device_left.get(), device_right.get(), len); + status = cudaGetLastError(); + } + if (status == cudaSuccess) { + status = copy_to_host(sums, device_sums, len); + } + if (status == cudaSuccess) { + status = copy_to_host(differences, device_differences, len); + } + if (status == cudaSuccess) { + status = copy_to_host(products, device_products, len); + } + if (status == cudaSuccess) { + status = copy_to_host(inverses, device_inverses, len); + } + return static_cast(status); +} + +extern "C" int multi_stark_cuda_blake3_hash_rows( + int device_id, uint8_t* digests, const uint8_t* messages, + size_t message_bytes, size_t message_count) { + if (digests == nullptr || messages == nullptr || message_count == 0 || + message_bytes == 0 || message_bytes > 32 * 1024 || + !product_fits(message_bytes, message_count) || + !product_fits(static_cast(32), message_count)) { + return static_cast(cudaErrorInvalidValue); + } + cudaError_t status = cudaSetDevice(device_id); + if (status != cudaSuccess) { + return static_cast(status); + } + + const size_t input_bytes = message_bytes * message_count; + const size_t output_bytes = 32 * message_count; + DeviceBuffer device_messages; + DeviceBuffer device_digests; + status = device_messages.allocate((input_bytes + 7) / 8); + if (status == cudaSuccess) { + status = device_digests.allocate((output_bytes + 7) / 8); + } + if (status == cudaSuccess) { + status = cudaMemcpy(device_messages.get(), messages, input_bytes, + cudaMemcpyHostToDevice); + } + if (status == cudaSuccess) { + status = launch_blake3_rows( + reinterpret_cast(device_digests.get()), + reinterpret_cast(device_messages.get()), + message_bytes, message_count); + } + if (status == cudaSuccess) { + status = cudaMemcpy(digests, device_digests.get(), output_bytes, + cudaMemcpyDeviceToHost); + } + return static_cast(status); +} + +extern "C" int multi_stark_cuda_blake3_merkle_root( + int device_id, uint8_t* root, const uint8_t* rows, size_t row_bytes, + size_t row_count) { + if (root == nullptr || rows == nullptr || row_bytes == 0 || + row_bytes > 32 * 1024 || !is_power_of_two(row_count) || + !product_fits(row_bytes, row_count)) { + return static_cast(cudaErrorInvalidValue); + } + cudaError_t status = cudaSetDevice(device_id); + if (status != cudaSuccess) { + return static_cast(status); + } + + const size_t input_bytes = row_bytes * row_count; + DeviceBuffer device_rows; + DeviceBuffer layer_a; + DeviceBuffer layer_b; + status = device_rows.allocate((input_bytes + 7) / 8); + if (status == cudaSuccess) { + status = layer_a.allocate(4 * row_count); + } + if (status == cudaSuccess && row_count > 1) { + status = layer_b.allocate(4 * (row_count / 2)); + } + if (status == cudaSuccess) { + status = cudaMemcpy(device_rows.get(), rows, input_bytes, + cudaMemcpyHostToDevice); + } + if (status == cudaSuccess) { + status = launch_blake3_rows( + reinterpret_cast(layer_a.get()), + reinterpret_cast(device_rows.get()), row_bytes, + row_count); + } + + size_t count = row_count; + uint8_t* current = reinterpret_cast(layer_a.get()); + uint8_t* next = reinterpret_cast(layer_b.get()); + while (status == cudaSuccess && count > 1) { + status = launch_blake3_digest_pairs(next, current, current, count / 2, + true); + count /= 2; + uint8_t* swap = current; + current = next; + next = swap; + } + if (status == cudaSuccess) { + status = cudaMemcpy(root, current, 32, cudaMemcpyDeviceToHost); + } + return static_cast(status); +} + +extern "C" int multi_stark_cuda_merkle_create( + int device_id, void** handle, uint8_t* root, const uint8_t* rows, + size_t row_bytes, size_t row_count) { + if (handle == nullptr || root == nullptr || rows == nullptr || + row_bytes == 0 || row_bytes > 32 * 1024 || + !is_power_of_two(row_count) || !product_fits(row_bytes, row_count) || + row_count > SIZE_MAX / 64) { + return static_cast(cudaErrorInvalidValue); + } + *handle = nullptr; + cudaError_t status = cudaSetDevice(device_id); + if (status != cudaSuccess) { + return static_cast(status); + } + + ResidentMerkleTree* tree = new (std::nothrow) ResidentMerkleTree; + if (tree == nullptr) { + return static_cast(cudaErrorMemoryAllocation); + } + tree->row_bytes = row_bytes; + tree->row_count = row_count; + const size_t rows_bytes = row_bytes * row_count; + status = cudaMalloc(reinterpret_cast(&tree->rows), rows_bytes); + if (status == cudaSuccess) { + // A full binary tree has fewer than 2*n digests. Keeping all layers in + // one allocation makes the prover-data handle compact and openings + // independent of host-side allocation lifetimes. + status = cudaMalloc(reinterpret_cast(&tree->digests), + 64 * row_count); + } + if (status == cudaSuccess) { + status = cudaMemcpy(tree->rows, rows, rows_bytes, + cudaMemcpyHostToDevice); + } + if (status == cudaSuccess) { + status = launch_blake3_rows(tree->digests, tree->rows, row_bytes, + row_count); + } + + size_t count = row_count; + size_t offset = 0; + while (status == cudaSuccess && count > 1) { + uint8_t* current = tree->digests + offset; + offset += count * 32; + status = launch_blake3_digest_pairs(tree->digests + offset, current, + current, count / 2, true); + count /= 2; + } + if (status == cudaSuccess) { + status = cudaMemcpy(root, tree->digests + offset, 32, + cudaMemcpyDeviceToHost); + } + if (status != cudaSuccess) { + delete tree; + return static_cast(status); + } + *handle = tree; + return static_cast(cudaSuccess); +} + +__global__ void gather_merkle_siblings(uint8_t* output,const uint8_t* digests, + size_t row_count,size_t index,size_t levels){ + const size_t byte=static_cast(blockIdx.x)*blockDim.x+threadIdx.x; + if(byte>=levels*32)return;const size_t level=byte/32,within=byte%32; + const size_t count=row_count>>level; + const size_t offset=(2*row_count-2*count)*32; + const size_t sibling=(index>>level)^1U; + output[byte]=digests[offset+sibling*32+within]; +} + +__global__ void gather_merkle_siblings_batch(uint8_t* output,const uint8_t* digests, + size_t row_count,const uint64_t* indices,size_t query_count,size_t levels){ + const size_t byte=static_cast(blockIdx.x)*blockDim.x+threadIdx.x; + const size_t path_bytes=levels*32;if(byte>=query_count*path_bytes)return; + const size_t query=byte/path_bytes,path_byte=byte-query*path_bytes; + const size_t level=path_byte/32,within=path_byte%32,count=row_count>>level; + const size_t offset=(2*row_count-2*count)*32; + const size_t sibling=(static_cast(indices[query])>>level)^1U; + output[byte]=digests[offset+sibling*32+within]; +} + +cudaError_t copy_merkle_siblings(uint8_t* output,const uint8_t* digests, + size_t row_count,size_t index){ + const size_t levels=strict_log2(row_count),bytes=levels*32; + if(levels==0)return cudaSuccess; + uint8_t* gathered=nullptr;cudaError_t status=cudaMalloc( + reinterpret_cast(&gathered),bytes); + if(status==cudaSuccess){gather_merkle_siblings<<>>( + gathered,digests,row_count,index,levels);status=cudaGetLastError();} + if(status==cudaSuccess)status=cudaMemcpy(output,gathered,bytes,cudaMemcpyDeviceToHost); + if(gathered)cudaFree(gathered);return status; +} + +cudaError_t copy_merkle_siblings_batch(uint8_t* output,const uint8_t* digests, + size_t row_count,const uint64_t* indices,size_t query_count){ + const size_t levels=strict_log2(row_count),bytes=query_count*levels*32; + uint64_t* device_indices=nullptr;uint8_t* gathered=nullptr; + cudaError_t status=cudaMalloc(reinterpret_cast(&device_indices),query_count*sizeof(uint64_t)); + if(status==cudaSuccess)status=cudaMalloc(reinterpret_cast(&gathered),bytes); + if(status==cudaSuccess)status=cudaMemcpy(device_indices,indices,query_count*sizeof(uint64_t),cudaMemcpyHostToDevice); + if(status==cudaSuccess){gather_merkle_siblings_batch<<>>(gathered,digests,row_count,device_indices,query_count,levels);status=cudaGetLastError();} + if(status==cudaSuccess)status=cudaMemcpy(output,gathered,bytes,cudaMemcpyDeviceToHost); + cudaFree(gathered);cudaFree(device_indices);return status; +} + +extern "C" int multi_stark_cuda_merkle_open( + int device_id, const void* handle, size_t index, uint8_t* row, + uint8_t* siblings) { + if (handle == nullptr || row == nullptr) { + return static_cast(cudaErrorInvalidValue); + } + cudaError_t status = cudaSetDevice(device_id); + if (status != cudaSuccess) { + return static_cast(status); + } + const ResidentMerkleTree* tree = + static_cast(handle); + if (index >= tree->row_count || (tree->row_count > 1 && siblings == nullptr)) { + return static_cast(cudaErrorInvalidValue); + } + status = cudaMemcpy(row, tree->rows + index * tree->row_bytes, + tree->row_bytes, cudaMemcpyDeviceToHost); + if(status==cudaSuccess)status=copy_merkle_siblings( + siblings,tree->digests,tree->row_count,index); + return static_cast(status); +} + +extern "C" int multi_stark_cuda_merkle_destroy(int device_id, void* handle) { + if (handle == nullptr) { + return static_cast(cudaSuccess); + } + const cudaError_t status = cudaSetDevice(device_id); + if (status != cudaSuccess) { + return static_cast(status); + } + delete static_cast(handle); + return static_cast(cudaGetLastError()); +} + +extern "C" int multi_stark_cuda_mixed_merkle_create( + int device_id, void** handle, uint8_t* root, + const uint8_t* const* level_rows, const size_t* level_row_bytes, + const size_t* level_heights, size_t level_count) { + if (handle == nullptr || root == nullptr || level_rows == nullptr || + level_row_bytes == nullptr || level_heights == nullptr || + level_count == 0 || !is_power_of_two(level_heights[0])) { + return static_cast(cudaErrorInvalidValue); + } + for (size_t level = 0; level < level_count; ++level) { + if (level_rows[level] == nullptr || level_row_bytes[level] == 0 || + level_row_bytes[level] > 32 * 1024 || + !is_power_of_two(level_heights[level]) || + (level > 0 && level_heights[level] >= level_heights[level - 1]) || + !product_fits(level_row_bytes[level], level_heights[level])) { + return static_cast(cudaErrorInvalidValue); + } + } + *handle = nullptr; + cudaError_t status = cudaSetDevice(device_id); + if (status != cudaSuccess) { + return static_cast(status); + } + ResidentMixedMerkleTree* tree = + new (std::nothrow) ResidentMixedMerkleTree; + if (tree == nullptr) { + return static_cast(cudaErrorMemoryAllocation); + } + tree->row_count = level_heights[0]; + status = cudaMalloc(reinterpret_cast(&tree->digests), + 64 * tree->row_count); + + DeviceBuffer device_rows; + DeviceBuffer injected_digests; + const size_t first_bytes = level_row_bytes[0] * level_heights[0]; + if (status == cudaSuccess) { + status = device_rows.allocate((first_bytes + 7) / 8); + } + if (status == cudaSuccess) { + status = injected_digests.allocate(4 * tree->row_count); + } + if (status == cudaSuccess) { + status = cudaMemcpy(device_rows.get(), level_rows[0], first_bytes, + cudaMemcpyHostToDevice); + } + if (status == cudaSuccess) { + status = launch_blake3_rows( + tree->digests, + reinterpret_cast(device_rows.get()), + level_row_bytes[0], level_heights[0]); + } + + size_t group = 1; + size_t count = tree->row_count; + size_t offset = 0; + while (status == cudaSuccess && count > 1) { + uint8_t* current = tree->digests + offset; + offset += count * 32; + count /= 2; + uint8_t* next = tree->digests + offset; + status = launch_blake3_digest_pairs(next, current, current, count, + true); + + if (status == cudaSuccess && group < level_count && + level_heights[group] == count) { + const size_t bytes = level_row_bytes[group] * count; + DeviceBuffer injected_rows; + status = injected_rows.allocate((bytes + 7) / 8); + if (status == cudaSuccess) { + status = cudaMemcpy(injected_rows.get(), level_rows[group], + bytes, cudaMemcpyHostToDevice); + } + if (status == cudaSuccess) { + status = launch_blake3_rows( + reinterpret_cast(injected_digests.get()), + reinterpret_cast(injected_rows.get()), + level_row_bytes[group], count); + } + if (status == cudaSuccess) { + status = launch_blake3_digest_pairs( + next, next, + reinterpret_cast(injected_digests.get()), + count, false); + } + ++group; + } + } + if (status == cudaSuccess && group != level_count) { + status = cudaErrorInvalidValue; + } + if (status == cudaSuccess) { + status = cudaMemcpy(root, tree->digests + offset, 32, + cudaMemcpyDeviceToHost); + } + if (status != cudaSuccess) { + delete tree; + return static_cast(status); + } + *handle = tree; + return static_cast(cudaSuccess); +} + +extern "C" int multi_stark_cuda_mixed_merkle_open( + int device_id, const void* handle, size_t index, uint8_t* siblings) { + if (handle == nullptr || siblings == nullptr) { + return static_cast(cudaErrorInvalidValue); + } + cudaError_t status = cudaSetDevice(device_id); + if (status != cudaSuccess) { + return static_cast(status); + } + const ResidentMixedMerkleTree* tree = + static_cast(handle); + if (index >= tree->row_count) { + return static_cast(cudaErrorInvalidValue); + } + status=copy_merkle_siblings(siblings,tree->digests,tree->row_count,index); + return static_cast(status); +} + +extern "C" int multi_stark_cuda_mixed_merkle_open_batch( + int device_id, const void* handle, const uint64_t* indices, + size_t query_count, uint8_t* siblings) { + if(handle==nullptr||indices==nullptr||query_count==0||siblings==nullptr)return static_cast(cudaErrorInvalidValue); + cudaError_t status=cudaSetDevice(device_id);if(status!=cudaSuccess)return static_cast(status); + const ResidentMixedMerkleTree* tree=static_cast(handle); + for(size_t q=0;q=tree->row_count)return static_cast(cudaErrorInvalidValue); + return static_cast(copy_merkle_siblings_batch(siblings,tree->digests,tree->row_count,indices,query_count)); +} + +extern "C" int multi_stark_cuda_mixed_merkle_destroy(int device_id, + void* handle) { + if (handle == nullptr) { + return static_cast(cudaSuccess); + } + const cudaError_t status = cudaSetDevice(device_id); + if (status != cudaSuccess) { + return static_cast(status); + } + delete static_cast(handle); + return static_cast(cudaGetLastError()); +} + +extern "C" int multi_stark_cuda_mixed_merkle_create_from_ldes( + int device_id, void** handle, uint8_t* root, + const void* const* lde_handles, size_t lde_count) { + if (handle == nullptr || root == nullptr || lde_handles == nullptr || + lde_count == 0) { + return static_cast(cudaErrorInvalidValue); + } + size_t max_height = 0; + for (size_t index = 0; index < lde_count; ++index) { + const ResidentLde* lde = + static_cast(lde_handles[index]); + if (lde == nullptr || !is_power_of_two(lde->height) || lde->width == 0) { + return static_cast(cudaErrorInvalidValue); + } + if (lde->height > max_height) { + max_height = lde->height; + } + } + if (max_height > SIZE_MAX / 64) { + return static_cast(cudaErrorInvalidValue); + } + *handle = nullptr; + cudaError_t status = cudaSetDevice(device_id); + if (status != cudaSuccess) { + return static_cast(status); + } + ResidentMixedMerkleTree* tree = + new (std::nothrow) ResidentMixedMerkleTree; + if (tree == nullptr) { + return static_cast(cudaErrorMemoryAllocation); + } + tree->row_count = max_height; + status = cudaMalloc(reinterpret_cast(&tree->digests), + 64 * max_height); + size_t max_injected_height = 0; + for (size_t index = 0; index < lde_count; ++index) { + const ResidentLde* lde = + static_cast(lde_handles[index]); + if (lde->height < max_height && lde->height > max_injected_height) { + max_injected_height = lde->height; + } + } + DeviceBuffer injected_digests; + if (status == cudaSuccess && max_injected_height != 0) { + status = injected_digests.allocate(4 * max_injected_height); + } + if (status == cudaSuccess) { + status = hash_resident_lde_group(tree->digests, lde_handles, + lde_count, max_height); + } + + size_t count = max_height; + size_t offset = 0; + while (status == cudaSuccess && count > 1) { + uint8_t* current = tree->digests + offset; + offset += count * 32; + count /= 2; + uint8_t* next = tree->digests + offset; + status = launch_blake3_digest_pairs(next, current, current, count, + true); + + bool inject = false; + for (size_t index = 0; index < lde_count; ++index) { + const ResidentLde* lde = + static_cast(lde_handles[index]); + inject = inject || lde->height == count; + } + if (status == cudaSuccess && inject) { + status = hash_resident_lde_group( + reinterpret_cast(injected_digests.get()), + lde_handles, lde_count, count); + } + if (status == cudaSuccess && inject) { + status = launch_blake3_digest_pairs( + next, next, + reinterpret_cast(injected_digests.get()), + count, false); + } + } + if (status == cudaSuccess) { + status = cudaMemcpy(root, tree->digests + offset, 32, + cudaMemcpyDeviceToHost); + } + if (status != cudaSuccess) { + delete tree; + return static_cast(status); + } + *handle = tree; + return static_cast(cudaSuccess); +} + +extern "C" int multi_stark_cuda_fri_merkle_create( + int device_id,void** handle,uint8_t* root,const void* codeword_handle,size_t arity){ + if(!codeword_handle||!is_power_of_two(arity))return static_cast(cudaErrorInvalidValue); + const auto* codeword=static_cast(codeword_handle); + if(codeword->width!=2||codeword->height%arity)return static_cast(cudaErrorInvalidValue); + ResidentLde view;view.values=codeword->values;view.height=codeword->height/arity;view.width=2*arity; + const void* views[1]={&view}; + const int status=multi_stark_cuda_mixed_merkle_create_from_ldes(device_id,handle,root,views,1); + view.values=nullptr; + return status; +} + +extern "C" int multi_stark_cuda_mixed_merkle_create_from_host_matrices( + int device_id, void** handle, uint8_t* root, + const uint64_t* const* matrix_values, const size_t* heights, + const size_t* widths, size_t matrix_count) { + if (handle == nullptr || root == nullptr || matrix_values == nullptr || + heights == nullptr || widths == nullptr || matrix_count == 0) { + return static_cast(cudaErrorInvalidValue); + } + cudaError_t status = cudaSetDevice(device_id); + if (status != cudaSuccess) { + return static_cast(status); + } + + void** lde_handles = new (std::nothrow) void*[matrix_count]; + if (lde_handles == nullptr) { + return static_cast(cudaErrorMemoryAllocation); + } + for (size_t index = 0; index < matrix_count; ++index) { + lde_handles[index] = nullptr; + } + size_t created = 0; + for (; status == cudaSuccess && created < matrix_count; ++created) { + if (matrix_values[created] == nullptr || + !is_power_of_two(heights[created]) || widths[created] == 0 || + !product_fits(heights[created], widths[created])) { + status = cudaErrorInvalidValue; + break; + } + ResidentLde* matrix = nullptr; + status = create_resident_lde(&matrix); + if (status != cudaSuccess) break; + matrix->height = heights[created]; + matrix->width = widths[created]; + status = cudaMalloc(reinterpret_cast(&matrix->values), + matrix->height * matrix->width * sizeof(uint64_t)); + if (status == cudaSuccess) { + status = cudaMemcpy(matrix->values, matrix_values[created], + matrix->height * matrix->width * sizeof(uint64_t), + cudaMemcpyHostToDevice); + } + if (status != cudaSuccess) { + destroy_resident_lde(matrix); + break; + } + lde_handles[created] = matrix; + } + if (status == cudaSuccess) { + status = static_cast( + multi_stark_cuda_mixed_merkle_create_from_ldes( + device_id, handle, root, + const_cast(lde_handles), matrix_count)); + } + for (size_t index = 0; index < created; ++index) { + destroy_resident_lde(static_cast(lde_handles[index])); + } + delete[] lde_handles; + return static_cast(status); +} + +extern "C" const char* multi_stark_cuda_error_string(int status) { + return cudaGetErrorString(static_cast(status)); +} + +extern "C" int multi_stark_cuda_memory_info(int device_id, size_t* free_bytes, + size_t* total_bytes) { + if (free_bytes == nullptr || total_bytes == nullptr) { + return static_cast(cudaErrorInvalidValue); + } + cudaError_t status = cudaSetDevice(device_id); + if (status == cudaSuccess) status = cudaMemGetInfo(free_bytes, total_bytes); + return static_cast(status); +} diff --git a/cuda/smoke.sh b/cuda/smoke.sh new file mode 100755 index 0000000..4089dac --- /dev/null +++ b/cuda/smoke.sh @@ -0,0 +1,34 @@ +#!/usr/bin/env bash +# SPDX-License-Identifier: MIT OR Apache-2.0 +set -euo pipefail + +git status --short --branch +git rev-parse HEAD +rustc --version +cargo --version +nvcc --version +nvidia-smi --query-gpu=index,name,uuid,memory.total,driver_version,compute_cap --format=csv + +if [[ -z "${MULTI_STARK_CUDA_ARCHS:-}" ]]; then + MULTI_STARK_CUDA_ARCHS="$(nvidia-smi --query-gpu=compute_cap --format=csv,noheader | sed -n '1p' | tr -d '.[:space:]')" + export MULTI_STARK_CUDA_ARCHS +fi +echo "MULTI_STARK_CUDA_ARCHS=${MULTI_STARK_CUDA_ARCHS}" + +cargo clippy --release --locked --all-targets --features parallel,cuda -- -D warnings +cargo test --release --locked --features parallel,cuda \ + cuda::tests::goldilocks_field_kernels_match_cpu -- --test-threads=1 +cargo test --release --locked --features parallel,cuda \ + cuda::tests::batched_dft_matches_cpu -- --test-threads=1 +cargo test --release --locked --features parallel,cuda \ + cuda::tests::coset_lde_matches_cpu_including_storage_layout -- --test-threads=1 +cargo test --release --locked --features parallel,cuda -- --test-threads=1 + +compat_dir="$(mktemp -d)" +trap 'rm -rf "$compat_dir"' EXIT +cargo run --release --locked --example proof_compatibility -- "$compat_dir/cpu.proof" +echo "c3c8ff942efc36fe508d6dcea3450cf4fa4d92a84f8d16bc651430595f4d244a $compat_dir/cpu.proof" \ + | sha256sum --check --status +cargo run --release --locked --features parallel,cuda \ + --example proof_compatibility -- "$compat_dir/cuda.proof" +cmp "$compat_dir/cpu.proof" "$compat_dir/cuda.proof" diff --git a/docs/cuda-benchmarks.md b/docs/cuda-benchmarks.md new file mode 100644 index 0000000..ed85386 --- /dev/null +++ b/docs/cuda-benchmarks.md @@ -0,0 +1,41 @@ +# CUDA benchmark record + +This file records dated downstream and microbenchmark results without making +hardware-specific numbers part of multi-stark's stable API documentation. + +## 2026-08-26: Ix recursive proving + +- GPU: NVIDIA RTX PRO 6000 Blackwell, 97,887 MiB +- Driver: 595.84 +- Toolkit/compiler: CUDA 13.3 (`nvcc`) +- Workload: Ix `Vector.extract_append`, recursive proving, 50 FRI queries +- CPU combined inner + outer STARK proving: 71.87 s +- CUDA combined inner + outer STARK proving: 8.84 s +- Speedup: 8.13x +- Inner proof: 11,783,606 bytes +- Outer proof: 4,162,775 bytes +- CPU verification: passed + +The 50-query setting was selected for development iteration and is not a +recommended production security parameter. These figures include a downstream +workload and should be re-measured after changes to Ix, multi-stark, the CUDA +toolkit, or the GPU architecture. + +## Reproducing in-repository measurements + +The Criterion benchmark exercises the full proving pipeline: + +```sh +cargo bench --release --locked --features parallel,cuda --bench multi_stark +``` + +Transfer-inclusive DFT/LDE and BLAKE3 comparisons are available separately: + +```sh +cargo run --release --locked --features parallel,cuda --example cuda_dft_bench +cargo run --release --locked --features parallel,cuda --example cuda_blake3_bench +``` + +The CSV DFT benchmark labels shapes below the production CUDA thresholds as +`cpu-fallback`, warms both implementations, uses the same iteration count, and +checks each output against the CPU reference outside the timing window. diff --git a/examples/cuda_blake3_bench.rs b/examples/cuda_blake3_bench.rs new file mode 100644 index 0000000..6625236 --- /dev/null +++ b/examples/cuda_blake3_bench.rs @@ -0,0 +1,54 @@ +#[cfg(not(feature = "cuda"))] +fn main() { + eprintln!("enable --features cuda"); +} + +#[cfg(feature = "cuda")] +fn main() { + use std::hint::black_box; + use std::time::Instant; + + use multi_stark::cuda::blake3_hash_rows; + use p3_blake3::Blake3; + use p3_maybe_rayon::prelude::*; + use p3_symmetric::CryptographicHasher; + + println!("backend,message_bytes,messages,iteration,seconds,megabytes"); + for (message_bytes, message_count) in [(64usize, 1 << 20), (4264, 1 << 15), (7400, 1 << 14)] { + let messages: Vec = (0..message_bytes * message_count) + .map(|index| (index as u64).wrapping_mul(0x9e37_79b9).to_le_bytes()[0]) + .collect(); + let expected: Vec<[u8; 32]> = messages + .par_chunks_exact(message_bytes) + .map(|message| Blake3.hash_iter(message.iter().copied())) + .collect(); + let warm = blake3_hash_rows(0, &messages, message_bytes); + assert_eq!(warm, expected); + + let megabytes = + f64::from(u32::try_from(messages.len()).expect("benchmark input exceeds u32")) + / 1_000_000.0; + for iteration in 0..3 { + let start = Instant::now(); + let cpu: Vec<[u8; 32]> = messages + .par_chunks_exact(message_bytes) + .map(|message| Blake3.hash_iter(message.iter().copied())) + .collect(); + black_box(&cpu); + println!( + "cpu,{message_bytes},{message_count},{iteration},{:.9},{megabytes:.3}", + start.elapsed().as_secs_f64() + ); + + let start = Instant::now(); + let gpu = blake3_hash_rows(0, &messages, message_bytes); + let elapsed = start.elapsed().as_secs_f64(); + black_box(&gpu); + assert_eq!(gpu, expected); + println!( + "cuda,{message_bytes},{message_count},{iteration},{:.9},{megabytes:.3}", + elapsed + ); + } + } +} diff --git a/examples/cuda_dft_bench.rs b/examples/cuda_dft_bench.rs new file mode 100644 index 0000000..5894a6a --- /dev/null +++ b/examples/cuda_dft_bench.rs @@ -0,0 +1,201 @@ +//! Transfer-inclusive CPU/CUDA DFT and coset-LDE comparison. +//! +//! Environment controls (comma-separated where applicable): +//! - `MULTI_STARK_CUDA_BENCH_LOG_HEIGHTS` (default `12,16,19`) +//! - `MULTI_STARK_CUDA_BENCH_WIDTHS` (default `2,8,32`) +//! - `MULTI_STARK_CUDA_BENCH_ITERATIONS` (default `3`) +//! - `MULTI_STARK_CUDA_BENCH_LOG_BLOWUP` (default `1`) +//! - `MULTI_STARK_CUDA_DEVICE` (default `0`) + +#[cfg(feature = "cuda")] +mod enabled { + use std::env; + use std::hint::black_box; + use std::time::Instant; + + use multi_stark::cuda::CudaDft; + use p3_dft::{Radix2DitParallel, TwoAdicSubgroupDft}; + use p3_field::{Field, PrimeCharacteristicRing}; + use p3_goldilocks::Goldilocks; + use p3_matrix::Matrix; + use p3_matrix::dense::RowMajorMatrix; + + pub(super) fn run() { + let log_heights = list("MULTI_STARK_CUDA_BENCH_LOG_HEIGHTS", "12,16,19"); + let widths = list("MULTI_STARK_CUDA_BENCH_WIDTHS", "2,8,32"); + let iterations = scalar("MULTI_STARK_CUDA_BENCH_ITERATIONS", 3); + let log_blowup = scalar("MULTI_STARK_CUDA_BENCH_LOG_BLOWUP", 1); + let device = scalar::("MULTI_STARK_CUDA_DEVICE", 0); + let cpu = Radix2DitParallel::::default(); + let gpu = CudaDft::new(device); + + println!("backend,operation,log_height,width,log_blowup,iteration,seconds,elements"); + for log_height in log_heights { + let height = 1usize << log_height; + for &width in &widths { + let input = deterministic_matrix(height, width); + + let expected_dft = cpu.dft_batch(input.clone()).to_row_major_matrix(); + let warm_dft = gpu.dft_batch(input.clone()).to_row_major_matrix(); + assert_eq!( + warm_dft, expected_dft, + "DFT correctness at 2^{log_height}x{width}" + ); + for iteration in 0..iterations { + let run_input = input.clone(); + let started = Instant::now(); + let output = cpu.dft_batch(run_input).to_row_major_matrix(); + let elapsed = started.elapsed().as_secs_f64(); + assert_eq!(output, expected_dft); + black_box(output); + emit( + "cpu", + "dft", + log_height, + width, + 0, + iteration, + elapsed, + height * width, + ); + + let run_input = input.clone(); + let started = Instant::now(); + let output = gpu.dft_batch(run_input).to_row_major_matrix(); + let elapsed = started.elapsed().as_secs_f64(); + assert_eq!(output, expected_dft); + black_box(output); + emit( + if CudaDft::uses_cuda_dft(height, width) { + "cuda" + } else { + "cpu-fallback" + }, + "dft", + log_height, + width, + 0, + iteration, + elapsed, + height * width, + ); + } + + let expected_lde = cpu + .coset_lde_batch(input.clone(), log_blowup, Goldilocks::GENERATOR) + .to_row_major_matrix(); + let warm_lde = gpu + .coset_lde_batch(input.clone(), log_blowup, Goldilocks::GENERATOR) + .to_row_major_matrix(); + assert_eq!( + warm_lde, expected_lde, + "LDE correctness at 2^{log_height}x{width}" + ); + for iteration in 0..iterations { + let run_input = input.clone(); + let started = Instant::now(); + let output = cpu + .coset_lde_batch(run_input, log_blowup, Goldilocks::GENERATOR) + .to_row_major_matrix(); + let elapsed = started.elapsed().as_secs_f64(); + assert_eq!(output, expected_lde); + black_box(output); + emit( + "cpu", + "coset_lde", + log_height, + width, + log_blowup, + iteration, + elapsed, + (height << log_blowup) * width, + ); + + let run_input = input.clone(); + let started = Instant::now(); + let output = gpu + .coset_lde_batch(run_input, log_blowup, Goldilocks::GENERATOR) + .to_row_major_matrix(); + let elapsed = started.elapsed().as_secs_f64(); + assert_eq!(output, expected_lde); + black_box(output); + emit( + if CudaDft::uses_cuda_coset_lde(height, width, log_blowup) { + "cuda" + } else { + "cpu-fallback" + }, + "coset_lde", + log_height, + width, + log_blowup, + iteration, + elapsed, + (height << log_blowup) * width, + ); + } + } + } + } + + fn deterministic_matrix(height: usize, width: usize) -> RowMajorMatrix { + RowMajorMatrix::new( + (0..height * width) + .map(|index| { + let index = u64::try_from(index).expect("matrix index exceeds u64"); + Goldilocks::from_u64(index.wrapping_mul(0x9e37_79b9_7f4a_7c15) ^ 0xdead_beef) + }) + .collect(), + width, + ) + } + + #[allow(clippy::too_many_arguments)] + fn emit( + backend: &str, + operation: &str, + log_height: usize, + width: usize, + log_blowup: usize, + iteration: usize, + seconds: f64, + elements: usize, + ) { + println!( + "{backend},{operation},{log_height},{width},{log_blowup},{iteration},{seconds:.9},{elements}" + ); + } + + fn list(name: &str, default: &str) -> Vec { + env::var(name) + .unwrap_or_else(|_| default.to_owned()) + .split(',') + .map(|value| { + value + .trim() + .parse() + .unwrap_or_else(|_| panic!("invalid {name}")) + }) + .collect() + } + + fn scalar(name: &str, default: T) -> T + where + T: std::str::FromStr + Copy, + { + env::var(name).map_or(default, |value| { + value.parse().unwrap_or_else(|_| panic!("invalid {name}")) + }) + } +} + +#[cfg(feature = "cuda")] +fn main() { + enabled::run(); +} + +#[cfg(not(feature = "cuda"))] +fn main() { + eprintln!("enable the `cuda` feature to run this benchmark"); + std::process::exit(2); +} diff --git a/examples/proof_compatibility.rs b/examples/proof_compatibility.rs new file mode 100644 index 0000000..eea52cd --- /dev/null +++ b/examples/proof_compatibility.rs @@ -0,0 +1,83 @@ +//! Deterministic CPU/CUDA proof-byte compatibility workload. +//! +//! The default shape deliberately crosses every production CUDA size gate: +//! height > 1024 for resident FRI, more than 2^18 cells for GPU MMCS, and more +//! than 10 million source cells for parallel LDE waves. `cuda/smoke.sh` runs +//! this example once without and once with `cuda`, then compares the complete +//! serialized proof files. + +use std::path::PathBuf; + +use multi_stark::p3_adapter::LookupAir; +use multi_stark::system::{System, SystemWitness}; +use multi_stark::types::{CommitmentParameters, FriParameters, GoldilocksBlake3Config, Val}; +use multi_stark::{ + p3_air::{Air, AirBuilder, BaseAir, WindowAccess}, + p3_field::PrimeCharacteristicRing, + p3_matrix::dense::RowMajorMatrix, +}; + +const WIDTH: usize = 40; +const LOG_HEIGHT: usize = 18; + +struct WidePythagoreanAir; + +impl BaseAir for WidePythagoreanAir { + fn width(&self) -> usize { + WIDTH + } +} + +impl Air for WidePythagoreanAir +where + AB::Var: Copy, +{ + fn eval(&self, builder: &mut AB) { + let main = builder.main(); + let local = main.current_slice(); + builder.assert_eq( + local[0] * local[0] + local[1] * local[1], + local[2] * local[2], + ); + for value in &local[3..] { + builder.assert_zero(*value); + } + } +} + +fn main() { + let output = std::env::args_os() + .nth(1) + .map(PathBuf::from) + .expect("usage: proof_compatibility OUTPUT"); + let config = GoldilocksBlake3Config::new( + CommitmentParameters { + log_blowup: 1, + cap_height: 0, + }, + FriParameters { + log_final_poly_len: 0, + max_log_arity: 1, + num_queries: 2, + commit_proof_of_work_bits: 0, + query_proof_of_work_bits: 0, + }, + ); + let (system, key) = System::new(config, [LookupAir::new(WidePythagoreanAir, vec![])]); + let height = 1 << LOG_HEIGHT; + let mut values = Val::zero_vec(height * WIDTH); + for row in values.chunks_exact_mut(WIDTH) { + row[0] = Val::from_u8(3); + row[1] = Val::from_u8(4); + row[2] = Val::from_u8(5); + } + let proof = system.prove_multiple_claims( + &key, + &[], + SystemWitness::from_stage_1(vec![RowMajorMatrix::new(values, WIDTH)], &system), + ); + system.verify_multiple_claims(&[], &proof).unwrap(); + let bytes = proof.to_bytes().expect("proof serialization failed"); + std::fs::write(&output, &bytes).expect("failed to write proof bytes"); + println!("wrote {} bytes to {}", bytes.len(), output.display()); +} diff --git a/src/config.rs b/src/config.rs index 1ff68c1..bca90b2 100644 --- a/src/config.rs +++ b/src/config.rs @@ -9,7 +9,8 @@ use p3_challenger::{CanObserve, CanSample, FieldChallenger}; use p3_commit::{Pcs, PolynomialSpace}; -use p3_field::{ExtensionField, Field}; +use p3_dft::TwoAdicSubgroupDft; +use p3_field::{ExtensionField, Field, TwoAdicField}; /// The base field of a configuration, as determined by its PCS domain. pub type Val = <<::Pcs as Pcs< @@ -47,6 +48,43 @@ pub type PcsData = <::Pcs as Pcs< ::Challenger, >>::ProverData; +/// Result produced by an accelerated lookup-trace constructor. +pub type AcceleratedLookupTraces = ( + Vec::Challenge>>, + Vec<::Challenge>, +); + +/// Result produced by an accelerated lookup commitment. +pub type AcceleratedLookupCommitment = ( + Com, + PcsData, + Vec<::Challenge>, +); + +/// One circuit's inputs to an optional fused quotient commitment backend. +/// Keeping this protocol-level description free of CUDA types lets the +/// generic CPU prover remain entirely independent of accelerator support. +pub struct QuotientCommitInput<'a, SC: StarkGenericConfig> { + pub circuit: &'a crate::system::Circuit>, + pub lookup_publics: Vec>, + pub trace_domain: Domain, + pub quotient_domain: Domain, + pub preprocessed: Option<(&'a PcsData, usize)>, + pub stage_1: (&'a PcsData, usize), + pub stage_2: (&'a PcsData, usize), + pub constraint_count: usize, +} + +/// One circuit's inputs to an optional fused lookup construction and +/// commitment backend. The committed stage-1 data lets an accelerator reuse +/// the witness already uploaded for the main-trace commitment. +pub struct LookupCommitInput<'a, SC: StarkGenericConfig> { + pub circuit: &'a crate::system::Circuit>, + pub lookup_values: &'a crate::lookup::LookupValues>, + pub preprocessed: Option<(&'a PcsData, usize)>, + pub stage_1: (&'a PcsData, usize), +} + /// Evaluations of committed polynomials over a domain. pub type EvaluationsOnDomain<'a, SC> = <::Pcs as Pcs< ::Challenge, @@ -65,6 +103,15 @@ pub trait StarkGenericConfig { /// The PCS used to commit to trace polynomials. type Pcs: Pcs; + /// The two-adic transform implementation used by prover-side polynomial + /// operations outside the PCS. + /// + /// Configurations should use the same implementation here and inside + /// their PCS. Keeping this transform explicit lets an accelerated backend + /// serve both paths without coupling the generic prover to a concrete CPU + /// DFT or changing the proof protocol. + type Dft: Clone + Default; + /// The field from which random challenges are drawn. Its size bounds the /// Schwartz-Zippel terms of the soundness error, so it must be large /// enough for the target security level (see the soundness argument in @@ -77,6 +124,13 @@ pub trait StarkGenericConfig { /// Returns a reference to the PCS. fn pcs(&self) -> &Self::Pcs; + /// Returns the transform implementation used by quotient polynomial + /// slicing and low-degree extension. + fn dft(&self) -> &Self::Dft + where + Val: TwoAdicField, + Self::Dft: TwoAdicSubgroupDft>; + /// Returns a fresh challenger. /// /// # Transcript contract @@ -120,4 +174,85 @@ pub trait StarkGenericConfig { /// module), and a mismatch produces commitments to the wrong /// evaluations. fn log_blowup(&self) -> usize; + + /// Normalize any backend-dependent field representatives before proof + /// serialization. Most fields have canonical in-memory representations; + /// configurations whose field permits lazy reduction can override this. + fn canonicalize_proof(_proof: &mut crate::prover::Proof) + where + Self: Sized, + { + } + + /// Optional device-resident quotient evaluator. Implementations return + /// `None` to use the portable host evaluator. Keeping this hook on the + /// configuration preserves a CUDA-free PCS and prover for every other + /// field/backend. + #[allow(clippy::too_many_arguments)] + fn accelerated_quotient_values( + &self, + _circuit: &crate::system::Circuit>, + _lookup_publics: &[Val], + _trace_domain: Domain, + _quotient_domain: Domain, + _preprocessed: Option<(&PcsData, usize)>, + _stage_1: (&PcsData, usize), + _stage_2: (&PcsData, usize), + _alpha: Self::Challenge, + _constraint_count: usize, + ) -> Option> + where + Self: Sized, + { + None + } + + /// Optional fused quotient evaluation, coefficient slicing, LDE and PCS + /// commitment. Implementations returning `None` use the portable path. + /// The result must commit to exactly the matrices produced by + /// `shifted_quotient_slices` followed by + /// `lde_from_shifted_coefficients`. + fn accelerated_quotient_commit( + &self, + _inputs: &[QuotientCommitInput<'_, Self>], + _alpha: Self::Challenge, + ) -> Option<(Com, PcsData)> + where + Self: Sized, + { + None + } + + /// Optional accelerator for the logUp message inversion and accumulator + /// scan. The returned matrices retain the protocol's extension-valued + /// row-major layout; the generic prover remains the reference fallback. + fn accelerated_lookup_traces( + &self, + _circuits: &[crate::lookup::LookupValues>], + _group_sizes: &[usize], + _lookup_challenge: Self::Challenge, + _fingerprint_challenge: Self::Challenge, + _accumulator: Self::Challenge, + ) -> Option> + where + Self: Sized, + { + None + } + + /// Optional fused lookup construction, LDE and commitment. This avoids + /// forcing accelerator-native trace storage through host matrices merely + /// to satisfy the portable `Pcs::commit` interface. + fn accelerated_lookup_commit( + &self, + _inputs: &[LookupCommitInput<'_, Self>], + _lookup_challenge: Self::Challenge, + _fingerprint_challenge: Self::Challenge, + _accumulator: Self::Challenge, + ) -> Option> + where + Self: Sized, + { + None + } } diff --git a/src/cuda/mmcs.rs b/src/cuda/mmcs.rs new file mode 100644 index 0000000..df45dbe --- /dev/null +++ b/src/cuda/mmcs.rs @@ -0,0 +1,335 @@ +use p3_commit::{BatchOpening, BatchOpeningRef, Mmcs}; +use p3_goldilocks::Goldilocks; +use p3_matrix::{Dimensions, Matrix, dense::RowMajorMatrix}; +use p3_merkle_tree::{MerkleTreeError, MerkleTreeMmcs}; +use p3_symmetric::MerkleCap; + +use super::{CudaLde, CudaMixedMerkleTree, mixed_lde_open_row, mixed_lde_open_rows}; +use crate::types::Blake3CompressionFunction; +use p3_blake3::Blake3; +use p3_symmetric::SerializingHasher; + +type CpuMmcs = + MerkleTreeMmcs, Blake3CompressionFunction, 2, 32>; +type CpuData = >::ProverData; + +pub enum CudaMmcsData { + Cpu(CpuData), + Cuda { + resident: std::sync::Arc>, + materialize: Option Vec + Send + Sync>>, + // `materialize` can own the last Arc to `resident`; drop it before + // retained host matrices so pinned trace pointers cannot outlive them. + matrices: std::sync::OnceLock>, + tree: CudaMixedMerkleTree, + }, +} + +impl CudaMmcsData { + pub(crate) fn resident(&self, index: usize) -> Option<&CudaLde> { + match self { + Self::Cuda { resident, .. } => resident.get(index), + Self::Cpu(_) => None, + } + } +} + +impl CudaMmcsData> { + pub(crate) fn resident_with_trace(&self, index: usize) -> Option<&CudaLde> { + match self { + Self::Cuda { + resident, matrices, .. + } => { + let lde = resident.get(index)?; + if let Some(retained) = matrices.get() { + // SAFETY: the matrix is retained in this prover data and + // drops after every Arc owner of `resident`. This proving + // path serializes attachment changes with CUDA use. + unsafe { lde.attach_trace(retained.get(index)?) }; + } + Some(lde) + } + Self::Cpu(_) => None, + } + } +} + +#[derive(Clone, Debug)] +pub struct CudaMmcs { + cpu: CpuMmcs, + device_id: i32, +} + +impl CudaMmcs { + pub(crate) fn new(cpu: CpuMmcs) -> Self { + Self { + cpu, + device_id: super::configured_device(), + } + } +} + +pub trait CudaCommitMmcs: Mmcs { + fn cuda_device_id(&self) -> i32; + + fn resident_or_upload>( + &self, + data: &Self::ProverData, + ) -> std::sync::Arc>; + fn commit_cuda_resident( + &self, + ldes: Vec, + ) -> (Self::Commitment, Self::ProverData>); + + fn retain_matrices( + &self, + data: &mut Self::ProverData>, + matrices: Vec>, + ); + + fn commit_cuda_storage( + &self, + ldes: Vec>, + ) -> (Self::Commitment, Self::ProverData>); +} + +pub trait CudaBatchOpenMmcs: Mmcs { + fn open_batches>( + &self, + indices: &[usize], + prover_data: &Self::ProverData, + ) -> Vec> + where + Self: Sized; +} + +impl CudaBatchOpenMmcs for CudaMmcs { + fn open_batches>( + &self, + indices: &[usize], + prover_data: &Self::ProverData, + ) -> Vec> { + match prover_data { + CudaMmcsData::Cpu(_) => indices + .iter() + .map(|&index| self.open_batch(index, prover_data)) + .collect(), + CudaMmcsData::Cuda { resident, tree, .. } => { + let rows = mixed_lde_open_rows(resident, indices); + let paths = tree.open_siblings_batch(indices); + rows.into_iter() + .zip(paths) + .map(|(opened_values, opening_proof)| { + BatchOpening::new(opened_values, opening_proof) + }) + .collect() + } + } + } +} + +impl CudaCommitMmcs for CudaMmcs { + fn cuda_device_id(&self) -> i32 { + self.device_id + } + + fn retain_matrices( + &self, + data: &mut Self::ProverData>, + retained: Vec>, + ) { + if let CudaMmcsData::Cuda { matrices, .. } = data { + matrices + .set(retained) + .expect("fresh CUDA trace matrix cell"); + } + } + + fn resident_or_upload>( + &self, + data: &Self::ProverData, + ) -> std::sync::Arc> { + match data { + CudaMmcsData::Cuda { resident, .. } => std::sync::Arc::clone(resident), + CudaMmcsData::Cpu(cpu) => std::sync::Arc::new( + self.cpu + .get_matrices(cpu) + .into_iter() + .map(|matrix| { + let values = (0..matrix.height()) + .flat_map(|row| matrix.row(row).unwrap()) + .collect(); + CudaLde::from_row_major_matrix( + self.device_id, + &RowMajorMatrix::new(values, matrix.width()), + ) + }) + .collect(), + ), + } + } + fn commit_cuda_resident( + &self, + ldes: Vec, + ) -> ( + Self::Commitment, + Self::ProverData>, + ) { + let ldes = std::sync::Arc::new(ldes); + let tree = CudaMixedMerkleTree::from_ldes(self.device_id, &ldes); + let commitment = MerkleCap::new(vec![tree.root()]); + let materialize_ldes = std::sync::Arc::clone(&ldes); + ( + commitment, + CudaMmcsData::Cuda { + resident: ldes, + materialize: Some(Box::new(move || { + materialize_ldes + .iter() + .map(CudaLde::to_row_major_matrix) + .collect() + })), + matrices: std::sync::OnceLock::new(), + tree, + }, + ) + } + + fn commit_cuda_storage( + &self, + ldes: Vec>, + ) -> ( + Self::Commitment, + Self::ProverData>, + ) { + let resident = ldes + .iter() + .map(|matrix| CudaLde::from_row_major_matrix(self.device_id, matrix)) + .collect(); + let (commitment, mut data) = self.commit_cuda_resident(resident); + if let CudaMmcsData::Cuda { matrices, .. } = &mut data { + let _ = matrices.set(ldes); + } + (commitment, data) + } +} + +impl Mmcs for CudaMmcs { + type ProverData = CudaMmcsData; + type Commitment = MerkleCap; + type Proof = Vec<[u8; 32]>; + type Error = MerkleTreeError; + + fn commit>( + &self, + inputs: Vec, + ) -> (Self::Commitment, Self::ProverData) { + if inputs + .iter() + .map(|matrix| matrix.height().saturating_mul(matrix.width())) + .sum::() + > 1 << 18 + { + let resident: Vec<_> = inputs + .iter() + .map(|matrix| { + let values = (0..matrix.height()) + .flat_map(|row| matrix.row(row).unwrap()) + .collect(); + CudaLde::from_row_major_matrix( + self.device_id, + &RowMajorMatrix::new(values, matrix.width()), + ) + }) + .collect(); + let resident = std::sync::Arc::new(resident); + let tree = CudaMixedMerkleTree::from_ldes(self.device_id, &resident); + let commitment = MerkleCap::new(vec![tree.root()]); + let matrices = std::sync::OnceLock::new(); + matrices.set(inputs).ok().expect("fresh CUDA matrix cell"); + return ( + commitment, + CudaMmcsData::Cuda { + resident, + materialize: None, + matrices, + tree, + }, + ); + } + let (commitment, data) = self.cpu.commit(inputs); + (commitment, CudaMmcsData::Cpu(data)) + } + + fn open_batch>( + &self, + index: usize, + prover_data: &Self::ProverData, + ) -> BatchOpening { + match prover_data { + CudaMmcsData::Cpu(data) => { + let opening = self.cpu.open_batch(index, data); + BatchOpening::new(opening.opened_values, opening.opening_proof) + } + CudaMmcsData::Cuda { resident, tree, .. } => { + let opened_values = mixed_lde_open_row(resident, index); + BatchOpening::new(opened_values, tree.open_siblings(index)) + } + } + } + + fn get_matrices<'a, M: Matrix>( + &self, + prover_data: &'a Self::ProverData, + ) -> Vec<&'a M> { + match prover_data { + CudaMmcsData::Cpu(data) => self.cpu.get_matrices(data), + CudaMmcsData::Cuda { + matrices, + materialize, + .. + } => matrices + .get_or_init(|| { + materialize + .as_ref() + .expect("CUDA matrices have no materializer")() + }) + .iter() + .collect(), + } + } + + fn get_matrix_heights>( + &self, + prover_data: &Self::ProverData, + ) -> Vec { + match prover_data { + CudaMmcsData::Cpu(data) => self.cpu.get_matrix_heights(data), + CudaMmcsData::Cuda { resident, .. } => resident.iter().map(CudaLde::height).collect(), + } + } + + fn get_max_height>(&self, prover_data: &Self::ProverData) -> usize { + match prover_data { + CudaMmcsData::Cpu(data) => self.cpu.get_max_height(data), + CudaMmcsData::Cuda { resident, .. } => { + resident.iter().map(CudaLde::height).max().unwrap() + } + } + } + + fn verify_batch( + &self, + commit: &Self::Commitment, + dimensions: &[Dimensions], + index: usize, + batch_opening: BatchOpeningRef<'_, Goldilocks, Self>, + ) -> Result<(), Self::Error> { + self.cpu.verify_batch( + commit, + dimensions, + index, + BatchOpeningRef::new(batch_opening.opened_values, batch_opening.opening_proof), + ) + } +} diff --git a/src/cuda/mod.rs b/src/cuda/mod.rs new file mode 100644 index 0000000..d7c24a2 --- /dev/null +++ b/src/cuda/mod.rs @@ -0,0 +1,2965 @@ +//! First-party CUDA acceleration for the Goldilocks/BLAKE3 prover pipeline. +//! +//! [`CudaDft`] implements Plonky3's host-oriented [`TwoAdicSubgroupDft`] +//! contract, while the resident types in this module keep LDEs, Merkle trees, +//! lookup traces, quotient evaluations, and FRI codewords on the selected GPU. +//! All public protocol types and serialized proofs remain unchanged. + +pub(crate) mod mmcs; +#[doc(hidden)] +pub mod pcs; + +use core::ffi::{CStr, c_char, c_void}; +use core::mem::{align_of, size_of}; +use core::ptr::NonNull; +use std::collections::BTreeMap; +use std::sync::{Arc, RwLock}; + +use crate::expr::{RowOffset, Source}; +use crate::graph::{ConstraintGraph, Node}; +use p3_dft::{Radix2DitParallel, TwoAdicSubgroupDft}; +use p3_field::{Field, PrimeCharacteristicRing, PrimeField64, TwoAdicField}; +use p3_goldilocks::Goldilocks; +use p3_matrix::Matrix; +use p3_matrix::bitrev::{BitReversalPerm, BitReversedMatrixView}; +use p3_matrix::dense::RowMajorMatrix; +use p3_util::log2_strict_usize; + +const _: () = assert!(size_of::() == size_of::()); +const _: () = assert!(align_of::() == align_of::()); + +type CachedPowers = Arc<[Goldilocks]>; +type SharedPowerCache = Arc>>; + +/// CUDA-backed batched DFT for the Goldilocks field. +/// +/// Clones share the host twiddle cache. This is important because the +/// production configuration gives one clone to the PCS and retains another +/// for quotient transforms. +#[derive(Clone, Debug)] +pub struct CudaDft { + device_id: i32, + cpu: Radix2DitParallel, + twiddles: SharedPowerCache<(usize, bool)>, + shift_powers: SharedPowerCache<(usize, u64)>, +} + +impl Default for CudaDft { + fn default() -> Self { + Self::new(configured_device()) + } +} + +pub(crate) fn configured_device() -> i32 { + std::env::var("MULTI_STARK_CUDA_DEVICE").map_or(0, |value| { + value + .parse() + .expect("MULTI_STARK_CUDA_DEVICE must be a non-negative integer") + }) +} + +impl CudaDft { + /// Selects the zero-based CUDA device used by subsequent transforms. + #[must_use] + pub fn new(device_id: i32) -> Self { + assert!(device_id >= 0, "CUDA device id must be non-negative"); + let _ = device_memory_info(device_id); + Self { + device_id, + cpu: Radix2DitParallel::default(), + twiddles: Arc::default(), + shift_powers: Arc::default(), + } + } + + /// Returns the selected CUDA device id. + #[must_use] + pub const fn device_id(&self) -> i32 { + self.device_id + } + + fn twiddles(&self, log_height: usize, inverse: bool) -> Arc<[Goldilocks]> { + let key = (log_height, inverse); + if let Some(twiddles) = self + .twiddles + .read() + .expect("twiddle cache poisoned") + .get(&key) + { + return Arc::clone(twiddles); + } + + let mut cache = self.twiddles.write().expect("twiddle cache poisoned"); + Arc::clone(cache.entry(key).or_insert_with(|| { + let root = Goldilocks::two_adic_generator(log_height); + let root = if inverse { root.inverse() } else { root }; + root.powers().take((1 << log_height) / 2).collect().into() + })) + } + + fn shift_powers(&self, height: usize, shift: Goldilocks) -> Arc<[Goldilocks]> { + let key = (height, shift.as_canonical_u64()); + if let Some(powers) = self + .shift_powers + .read() + .expect("shift-power cache poisoned") + .get(&key) + { + return Arc::clone(powers); + } + + let mut cache = self + .shift_powers + .write() + .expect("shift-power cache poisoned"); + Arc::clone( + cache + .entry(key) + .or_insert_with(|| shift.powers().take(height).collect().into()), + ) + } + + fn validate_dimensions(height: usize, width: usize) { + assert!( + height.is_power_of_two(), + "DFT height must be a power of two" + ); + let log_height = log2_strict_usize(height); + assert!( + log_height <= Goldilocks::TWO_ADICITY, + "DFT height exceeds Goldilocks two-adicity" + ); + height + .checked_mul(width) + .expect("DFT matrix element count overflows usize"); + } + + #[inline] + fn use_cuda_dft(height: usize, width: usize) -> bool { + height.saturating_mul(width) >= (1 << 15) + } + + #[inline] + fn use_cuda_coset_lde(extended_height: usize, width: usize) -> bool { + extended_height >= (1 << 15) && width <= 2 + } + + /// Reports whether `dft_batch` will execute on CUDA for this shape. + #[must_use] + pub fn uses_cuda_dft(height: usize, width: usize) -> bool { + height.saturating_mul(width) >= (1 << 15) + } + + /// Reports whether `coset_lde_batch` will execute on CUDA for this shape. + #[must_use] + pub fn uses_cuda_coset_lde(height: usize, width: usize, added_bits: usize) -> bool { + let Ok(added_bits) = u32::try_from(added_bits) else { + return false; + }; + let Some(extended_height) = height.checked_shl(added_bits) else { + return false; + }; + extended_height >= (1 << 15) && width <= 2 + } + + /// Computes a coset LDE whose bit-reversed storage remains on the GPU. + #[must_use] + pub(crate) fn coset_lde_batch_resident( + &self, + matrix: &RowMajorMatrix, + added_bits: usize, + shift: Goldilocks, + ) -> CudaLde { + let height = matrix.height(); + let width = matrix.width(); + Self::validate_dimensions(height, width); + assert!(width > 0, "resident CUDA LDE requires at least one column"); + let extended_height = height + .checked_shl(u32::try_from(added_bits).expect("LDE blowup exceeds u32")) + .expect("LDE height overflows usize"); + Self::validate_dimensions(extended_height, width); + + let log_height = log2_strict_usize(height); + let inverse_twiddles = self.twiddles(log_height, true); + let forward_twiddles = self.twiddles(log2_strict_usize(extended_height), false); + let shift_powers = self.shift_powers(height, shift); + let height_inverse = Goldilocks::ONE.div_2exp_u64(log_height as u64); + let mut handle = core::ptr::null_mut(); + // SAFETY: every host buffer has the exact dimensions validated above; + // successful creation transfers the device allocation to `CudaLde`. + let status = unsafe { + multi_stark_cuda_coset_lde_create( + self.device_id, + &mut handle, + matrix.values.as_ptr().cast(), + height, + width, + added_bits, + inverse_twiddles.as_ptr().cast(), + shift_powers.as_ptr().cast(), + forward_twiddles.as_ptr().cast(), + raw_u64(height_inverse), + ) + }; + check_cuda(status, "resident coset LDE"); + CudaLde { + device_id: self.device_id, + handle: NonNull::new(handle).expect("CUDA returned a null LDE handle"), + height: extended_height, + width, + } + } + + pub(crate) fn prepare_coset_lde_constants( + &self, + height: usize, + added_bits: usize, + shift: Goldilocks, + ) { + let extended_height = height + .checked_shl(u32::try_from(added_bits).expect("LDE blowup exceeds u32")) + .expect("LDE height overflows usize"); + Self::validate_dimensions(height, 1); + Self::validate_dimensions(extended_height, 1); + let inverse_twiddles = self.twiddles(log2_strict_usize(height), true); + let shift_powers = self.shift_powers(height, shift); + let forward_twiddles = self.twiddles(log2_strict_usize(extended_height), false); + let status = unsafe { + multi_stark_cuda_prepare_lde_constants( + self.device_id, + inverse_twiddles.as_ptr().cast(), + inverse_twiddles.len(), + shift_powers.as_ptr().cast(), + height, + forward_twiddles.as_ptr().cast(), + forward_twiddles.len(), + ) + }; + check_cuda(status, "prepare resident LDE constants"); + } +} + +/// Bit-reversed coset-LDE storage owned by a CUDA device allocation. +#[doc(hidden)] +pub struct CudaLde { + device_id: i32, + handle: NonNull, + height: usize, + width: usize, +} + +unsafe impl Send for CudaLde {} +unsafe impl Sync for CudaLde {} + +impl CudaLde { + pub(crate) const fn raw_handle(&self) -> *const c_void { + self.handle.as_ptr() + } + + /// # Safety + /// + /// The caller must serialize attachment changes with all CUDA operations + /// that can read the attached trace. + pub(crate) unsafe fn release_trace(&self) { + let status = + unsafe { multi_stark_cuda_lde_release_trace(self.device_id, self.handle.as_ptr()) }; + check_cuda(status, "resident trace release"); + } + + /// # Safety + /// + /// `trace` must outlive this handle or an earlier `release_trace` call, + /// and the caller must serialize attachment changes with CUDA trace use. + pub(crate) unsafe fn attach_trace(&self, trace: &RowMajorMatrix) { + assert_eq!(trace.width(), self.width); + let status = unsafe { + multi_stark_cuda_lde_attach_trace( + self.device_id, + self.handle.as_ptr(), + trace.values.as_ptr().cast(), + trace.height(), + trace.width(), + ) + }; + check_cuda(status, "resident trace attachment"); + } + + pub(crate) fn fri_fold( + &self, + next: Option<&CudaReducedOpening>, + betas: &[[Goldilocks; 2]], + beta_power: [Goldilocks; 2], + g_inv: Goldilocks, + ext_w: Goldilocks, + ) -> Self { + let mut handle = core::ptr::null_mut(); + let status = unsafe { + multi_stark_cuda_fri_fold_resident( + self.device_id, + &mut handle, + self.raw_handle(), + next.map_or(core::ptr::null(), |value| value.handle.as_ptr()), + betas.as_ptr().cast(), + betas.len(), + raw_u64(beta_power[0]), + raw_u64(beta_power[1]), + raw_u64(g_inv), + raw_u64(ext_w), + ) + }; + check_cuda(status, "resident CUDA FRI fold"); + Self { + device_id: self.device_id, + handle: NonNull::new(handle).expect("CUDA returned a null folded codeword"), + height: self.height >> betas.len(), + width: 2, + } + } + /// Uploads an already bit-reversed LDE matrix into resident storage. + #[must_use] + pub(crate) fn from_row_major_matrix( + device_id: i32, + matrix: &RowMajorMatrix, + ) -> Self { + assert!( + matrix.height().is_power_of_two(), + "LDE height is not a power of two" + ); + assert!(matrix.width() > 0, "LDE width is zero"); + let mut handle = core::ptr::null_mut(); + // SAFETY: the matrix is contiguous and remains live for the synchronous upload. + let status = unsafe { + multi_stark_cuda_lde_create_from_host( + device_id, + &mut handle, + matrix.values.as_ptr().cast(), + matrix.height(), + matrix.width(), + ) + }; + check_cuda(status, "resident LDE upload"); + Self { + device_id, + handle: NonNull::new(handle).expect("CUDA returned a null LDE handle"), + height: matrix.height(), + width: matrix.width(), + } + } + + #[must_use] + pub(crate) const fn height(&self) -> usize { + self.height + } + + #[must_use] + pub(crate) const fn width(&self) -> usize { + self.width + } + + /// Copies the bit-reversed committed storage to a host matrix. This is an + /// oracle/debug escape hatch; resident PCS code should retain the handle. + #[must_use] + pub(crate) fn to_row_major_matrix(&self) -> RowMajorMatrix { + let mut values = Goldilocks::zero_vec(self.height * self.width); + // SAFETY: the output allocation matches the immutable resident LDE. + let status = unsafe { + multi_stark_cuda_lde_copy_to_host( + self.device_id, + self.handle.as_ptr(), + values.as_mut_ptr().cast(), + ) + }; + check_cuda(status, "resident LDE host copy"); + RowMajorMatrix::new(values, self.width) + } + + /// Copies one sampled row from the resident bit-reversed LDE. + #[cfg(test)] + #[must_use] + pub(crate) fn row(&self, row: usize) -> Vec { + assert!(row < self.height, "resident LDE row index out of bounds"); + let mut values = Goldilocks::zero_vec(self.width); + // SAFETY: `row` is in bounds and the output has exactly `width` + // elements, matching the contiguous device row. + let status = unsafe { + multi_stark_cuda_lde_copy_row( + self.device_id, + self.handle.as_ptr(), + row, + values.as_mut_ptr().cast(), + ) + }; + check_cuda(status, "resident LDE row copy"); + values + } + + /// Gathers sampled rows with one device kernel and one host transfer. + #[must_use] + pub(crate) fn rows(&self, rows: &[usize]) -> Vec> { + assert!(!rows.is_empty(), "resident LDE row batch is empty"); + assert!( + rows.iter().all(|&row| row < self.height), + "resident LDE row index out of bounds" + ); + let device_rows: Vec = rows + .iter() + .map(|&row| u64::try_from(row).expect("row index exceeds u64")) + .collect(); + let mut values = Goldilocks::zero_vec(rows.len() * self.width); + // SAFETY: all row indices are in bounds and `values` holds one full + // output row for every requested index. + let status = unsafe { + multi_stark_cuda_lde_copy_rows( + self.device_id, + self.handle.as_ptr(), + device_rows.as_ptr(), + device_rows.len(), + values.as_mut_ptr().cast(), + ) + }; + check_cuda(status, "resident LDE row batch copy"); + values + .chunks_exact(self.width) + .map(<[Goldilocks]>::to_vec) + .collect() + } +} + +#[repr(C)] +#[derive(Clone, Copy)] +struct CudaConstraintNode { + value: u64, + a: u32, + b: u32, + op: u32, + aux: u32, + out: u32, +} + +#[repr(C)] +#[derive(Clone, Copy)] +struct CudaConstraintLookup { + multiplicity: u32, + arg_start: u32, + arg_count: u32, + emit_after: u32, + output: u32, +} + +type EncodedLookupGraph = ( + Vec, + usize, + Vec, + Vec, +); + +fn checked_u32(value: usize, what: &str) -> u32 { + u32::try_from(value).unwrap_or_else(|_| panic!("{what} exceeds the CUDA ABI limit")) +} + +fn encode_constraint_nodes(graph: &ConstraintGraph) -> Vec { + graph + .nodes + .iter() + .enumerate() + .map(|(index, node)| { + let mut out = CudaConstraintNode { + value: 0, + a: u32::MAX, + b: u32::MAX, + op: 0, + aux: 0, + out: checked_u32(index, "constraint node index"), + }; + match *node { + Node::Const(value) => out.value = raw_u64(value), + Node::Var(col) => { + out.op = 1; + out.a = col.index; + out.aux = match col.source { + Source::Preprocessed => 0, + Source::Main => 2, + Source::Stage2 => 4, + } + u32::from(matches!(col.offset, RowOffset::Next)); + } + Node::Public(i) => { + out.op = 2; + out.a = i + } + Node::IsFirstRow => out.op = 3, + Node::IsLastRow => out.op = 4, + Node::IsTransition => out.op = 5, + Node::Add(a, b) => { + out.op = 6; + out.a = a.0; + out.b = b.0 + } + Node::Sub(a, b) => { + out.op = 7; + out.a = a.0; + out.b = b.0 + } + Node::Mul(a, b) => { + out.op = 8; + out.a = a.0; + out.b = b.0 + } + Node::Neg(a) => { + out.op = 9; + out.a = a.0 + } + } + out + }) + .collect() +} + +fn encode_quotient_nodes( + graph: &ConstraintGraph, +) -> (Vec, Vec, usize) { + let n = graph.nodes.len(); + let mut last: Vec<_> = (0..n).collect(); + for (i, node) in graph.nodes.iter().enumerate() { + match *node { + Node::Add(a, b) | Node::Sub(a, b) | Node::Mul(a, b) => { + last[a.index()] = i; + last[b.index()] = i + } + Node::Neg(a) => last[a.index()] = i, + _ => {} + } + } + for id in graph.zeros.iter().chain( + graph + .lookups + .iter() + .flat_map(|lookup| core::iter::once(&lookup.multiplicity).chain(lookup.args.iter())), + ) { + last[id.index()] = n; + } + let mut slots = vec![0u32; n]; + let mut free = Vec::new(); + let mut retire = vec![Vec::new(); n + 1]; + let mut count = 0u32; + for i in 0..n { + free.append(&mut retire[i]); + let slot = free.pop().unwrap_or_else(|| { + let s = count; + count += 1; + s + }); + slots[i] = slot; + if last[i] < n { + retire[last[i] + 1].push(slot); + } + } + let mut nodes = encode_constraint_nodes(graph); + for (i, (encoded, node)) in nodes.iter_mut().zip(&graph.nodes).enumerate() { + encoded.out = slots[i]; + match *node { + Node::Add(a, b) | Node::Sub(a, b) | Node::Mul(a, b) => { + encoded.a = slots[a.index()]; + encoded.b = slots[b.index()] + } + Node::Neg(a) => encoded.a = slots[a.index()], + _ => {} + } + } + (nodes, slots, count as usize) +} + +fn encode_lookup_nodes(graph: &ConstraintGraph) -> Option { + let n = graph.lookup_prefix_len; + if n == 0 || graph.lookups.is_empty() { + return None; + } + let mut reachable = vec![false; n]; + let mut pending: Vec = graph + .lookups + .iter() + .flat_map(|lookup| { + core::iter::once(lookup.multiplicity.index()) + .chain(lookup.args.iter().map(|id| id.index())) + }) + .collect(); + while let Some(index) = pending.pop() { + if index >= n || reachable[index] { + continue; + } + reachable[index] = true; + match graph.nodes[index] { + Node::Add(a, b) | Node::Sub(a, b) | Node::Mul(a, b) => { + pending.push(a.index()); + pending.push(b.index()); + } + Node::Neg(a) => pending.push(a.index()), + _ => {} + } + } + let order: Vec<_> = reachable + .iter() + .enumerate() + .filter_map(|(index, &used)| used.then_some(index)) + .collect(); + let mut compact = vec![usize::MAX; n]; + for (position, &index) in order.iter().enumerate() { + compact[index] = position; + } + let node_count = order.len(); + // Lookup witness expressions are base-trace expressions. Public values + // are the logUp challenges sampled only after stage 1, and stage 2 does + // not exist yet, so neither can legally occur in this prefix. + if order.iter().map(|&index| &graph.nodes[index]).any(|node| { + matches!(node, Node::Public(_)) + || matches!(node, Node::Var(col) if col.source == Source::Stage2) + }) { + return None; + } + let mut last = vec![0usize; node_count]; + for (position, &index) in order.iter().enumerate() { + let node = &graph.nodes[index]; + match *node { + Node::Add(a, b) | Node::Sub(a, b) | Node::Mul(a, b) => { + last[compact[a.index()]] = position; + last[compact[b.index()]] = position; + } + Node::Neg(a) => last[compact[a.index()]] = position, + _ => {} + } + } + let emit_after = graph + .lookups + .iter() + .map(|lookup| { + core::iter::once(&lookup.multiplicity) + .chain(lookup.args.iter()) + .map(|id| compact[id.index()]) + .max() + .expect("lookup has a multiplicity") + }) + .collect::>(); + for (lookup, &emit) in graph.lookups.iter().zip(&emit_after) { + for id in core::iter::once(&lookup.multiplicity).chain(lookup.args.iter()) { + last[compact[id.index()]] = last[compact[id.index()]].max(emit); + } + } + let mut slots = vec![0u32; node_count]; + let mut free = Vec::new(); + let mut retire = vec![Vec::new(); node_count + 1]; + let mut slot_count = 0u32; + for i in 0..node_count { + free.append(&mut retire[i]); + let slot = free.pop().unwrap_or_else(|| { + let result = slot_count; + slot_count += 1; + result + }); + slots[i] = slot; + if last[i] < node_count { + retire[last[i] + 1].push(slot); + } + } + let encoded_all = encode_constraint_nodes(graph); + let mut nodes: Vec<_> = order.iter().map(|&index| encoded_all[index]).collect(); + for (position, (&index, encoded)) in order.iter().zip(&mut nodes).enumerate() { + let node = &graph.nodes[index]; + encoded.out = slots[position]; + match *node { + Node::Add(a, b) | Node::Sub(a, b) | Node::Mul(a, b) => { + encoded.a = slots[compact[a.index()]]; + encoded.b = slots[compact[b.index()]]; + } + Node::Neg(a) => encoded.a = slots[compact[a.index()]], + _ => {} + } + } + let mut args = Vec::new(); + let mut lookups = graph + .lookups + .iter() + .zip(emit_after) + .enumerate() + .map(|(output, (lookup, emit_after))| { + let arg_start = checked_u32(args.len(), "lookup argument offset"); + args.extend(lookup.args.iter().map(|id| slots[compact[id.index()]])); + CudaConstraintLookup { + multiplicity: slots[compact[lookup.multiplicity.index()]], + arg_start, + arg_count: checked_u32(lookup.args.len(), "lookup argument count"), + emit_after: checked_u32(emit_after, "lookup emission index"), + output: checked_u32(output, "lookup output index"), + } + }) + .collect::>(); + lookups.sort_unstable_by_key(|lookup| lookup.emit_after); + Some((nodes, slot_count as usize, lookups, args)) +} + +/// Evaluates the compiled base-field constraint roots directly against +/// resident trace LDEs. This is the protocol-independent core used by the +/// CUDA quotient path; selectors and public values remain caller supplied. +#[cfg(test)] +#[allow(clippy::too_many_arguments)] +pub(crate) fn constraint_graph_roots( + graph: &ConstraintGraph, + preprocessed: Option<&CudaLde>, + main: &CudaLde, + stage2: &CudaLde, + publics: &[Goldilocks], + selectors: &[Goldilocks], + quotient_size: usize, + next_step: usize, +) -> RowMajorMatrix { + assert_eq!(selectors.len(), 3 * quotient_size); + assert!(next_step <= quotient_size); + let nodes = encode_constraint_nodes(graph); + let roots: Vec = graph.zeros.iter().map(|root| root.0).collect(); + let mut output = Goldilocks::zero_vec(quotient_size * roots.len()); + let status = unsafe { + multi_stark_cuda_constraint_graph( + main.device_id, + output.as_mut_ptr().cast(), + nodes.as_ptr().cast(), + nodes.len(), + roots.as_ptr().cast(), + roots.len(), + preprocessed.map_or(core::ptr::null(), CudaLde::raw_handle), + main.raw_handle(), + stage2.raw_handle(), + publics.as_ptr().cast(), + publics.len(), + selectors.as_ptr().cast(), + quotient_size, + next_step, + ) + }; + check_cuda(status, "constraint graph evaluation"); + RowMajorMatrix::new(output, roots.len()) +} + +/// Compact description of the four Lagrange-selector columns used on a +/// two-adic quotient coset. CUDA expands these geometric sequences directly +/// into device memory instead of receiving four full host vectors. +#[derive(Clone, Copy)] +pub(crate) struct CudaCosetSelectors { + pub(crate) coset_shift: Goldilocks, + pub(crate) coset_generator: Goldilocks, + pub(crate) trace_last: Goldilocks, + pub(crate) vanishing_start: Goldilocks, + pub(crate) vanishing_step: Goldilocks, +} + +#[allow(clippy::too_many_arguments)] +pub(crate) fn quotient_values_resident( + graph: &ConstraintGraph, + preprocessed: Option<&CudaLde>, + main: &CudaLde, + stage2: &CudaLde, + publics: &[Goldilocks], + selectors: CudaCosetSelectors, + alpha: &[Goldilocks], + delta: &[Goldilocks; 2], + ext_w: Goldilocks, + quotient_size: usize, + next_step: usize, + group_size: usize, +) -> RowMajorMatrix { + assert!((1..=8).contains(&group_size)); + assert!(publics.len() >= 4); + assert!(next_step.is_power_of_two()); + assert!(next_step <= quotient_size); + let (nodes, slots, slot_count) = encode_quotient_nodes(graph); + let roots: Vec<_> = graph.zeros.iter().map(|root| slots[root.index()]).collect(); + let mut args = Vec::new(); + let lookups: Vec<_> = graph + .lookups + .iter() + .map(|lookup| { + let arg_start = checked_u32(args.len(), "lookup argument offset"); + args.extend(lookup.args.iter().map(|arg| slots[arg.index()])); + CudaConstraintLookup { + multiplicity: slots[lookup.multiplicity.index()], + arg_start, + arg_count: checked_u32(lookup.args.len(), "lookup argument count"), + emit_after: 0, + output: 0, + } + }) + .collect(); + let constraint_count = alpha.len() / 2; + let expected_constraints = roots.len() + + if lookups.is_empty() { + 2 + } else { + 2 * lookups.len().div_ceil(group_size) + }; + assert_eq!(alpha.len(), 2 * expected_constraints); + let mut output = Goldilocks::zero_vec(2 * quotient_size); + let status = unsafe { + multi_stark_cuda_quotient_values( + main.device_id, + output.as_mut_ptr().cast(), + nodes.as_ptr().cast(), + nodes.len(), + slot_count, + roots.as_ptr(), + roots.len(), + lookups.as_ptr().cast(), + lookups.len(), + args.as_ptr(), + args.len(), + group_size, + preprocessed.map_or(core::ptr::null(), CudaLde::raw_handle), + main.raw_handle(), + stage2.raw_handle(), + publics.as_ptr().cast(), + publics.len(), + raw_u64(selectors.coset_shift), + raw_u64(selectors.coset_generator), + raw_u64(selectors.trace_last), + raw_u64(selectors.vanishing_start), + raw_u64(selectors.vanishing_step), + alpha.as_ptr().cast(), + constraint_count, + delta.as_ptr().cast(), + raw_u64(ext_w), + quotient_size, + next_step, + ) + }; + check_cuda(status, "resident quotient evaluation"); + RowMajorMatrix::new(output, 2) +} + +/// Evaluate a quotient and carry it through coefficient slicing and the +/// committed low-degree extension without materializing an intermediate host +/// matrix. The returned storage has exactly the bit-reversed layout expected +/// by `Pcs::commit_ldes`. +#[allow(clippy::too_many_arguments)] +pub(crate) fn quotient_lde_resident( + dft: &CudaDft, + graph: &ConstraintGraph, + preprocessed: Option<&CudaLde>, + main: &CudaLde, + stage2: &CudaLde, + publics: &[Goldilocks], + selectors: CudaCosetSelectors, + alpha: &[Goldilocks], + delta: &[Goldilocks; 2], + ext_w: Goldilocks, + quotient_size: usize, + next_step: usize, + group_size: usize, + quotient_degree: usize, + log_blowup: usize, +) -> CudaLde { + assert!((1..=8).contains(&group_size)); + assert!(publics.len() >= 4); + assert!(next_step.is_power_of_two()); + assert!(next_step <= quotient_size); + assert!(quotient_degree.is_power_of_two()); + assert_eq!(quotient_size % quotient_degree, 0); + let (nodes, slots, slot_count) = encode_quotient_nodes(graph); + let roots: Vec<_> = graph.zeros.iter().map(|root| slots[root.index()]).collect(); + let mut args = Vec::new(); + let lookups: Vec<_> = graph + .lookups + .iter() + .map(|lookup| { + let arg_start = checked_u32(args.len(), "lookup argument offset"); + args.extend(lookup.args.iter().map(|arg| slots[arg.index()])); + CudaConstraintLookup { + multiplicity: slots[lookup.multiplicity.index()], + arg_start, + arg_count: checked_u32(lookup.args.len(), "lookup argument count"), + emit_after: 0, + output: 0, + } + }) + .collect(); + let constraint_count = alpha.len() / 2; + let expected_constraints = roots.len() + + if lookups.is_empty() { + 2 + } else { + 2 * lookups.len().div_ceil(group_size) + }; + assert_eq!(alpha.len(), 2 * expected_constraints); + let trace_height = quotient_size / quotient_degree; + let lde_height = trace_height << log_blowup; + let quotient_twiddles = dft.twiddles(log2_strict_usize(quotient_size), false); + let lde_twiddles = dft.twiddles(log2_strict_usize(lde_height), false); + let height_inverse = Goldilocks::ONE.div_2exp_u64(log2_strict_usize(quotient_size) as u64); + let weight_step = Goldilocks::GENERATOR.exp_u64(trace_height as u64).inverse(); + let weights: Vec<_> = weight_step + .powers() + .take(quotient_degree) + .map(|weight| weight * height_inverse) + .collect(); + let mut handle = core::ptr::null_mut(); + let status = unsafe { + multi_stark_cuda_quotient_lde( + dft.device_id, + &mut handle, + nodes.as_ptr().cast(), + nodes.len(), + slot_count, + roots.as_ptr(), + roots.len(), + lookups.as_ptr().cast(), + lookups.len(), + args.as_ptr(), + args.len(), + group_size, + preprocessed.map_or(core::ptr::null(), CudaLde::raw_handle), + main.raw_handle(), + stage2.raw_handle(), + publics.as_ptr().cast(), + publics.len(), + raw_u64(selectors.coset_shift), + raw_u64(selectors.coset_generator), + raw_u64(selectors.trace_last), + raw_u64(selectors.vanishing_start), + raw_u64(selectors.vanishing_step), + alpha.as_ptr().cast(), + constraint_count, + delta.as_ptr().cast(), + raw_u64(ext_w), + quotient_size, + next_step, + quotient_degree, + log_blowup, + quotient_twiddles.as_ptr().cast(), + lde_twiddles.as_ptr().cast(), + weights.as_ptr().cast(), + ) + }; + check_cuda(status, "resident quotient LDE"); + CudaLde { + device_id: dft.device_id, + handle: NonNull::new(handle).expect("CUDA returned a null quotient LDE"), + height: lde_height, + width: 2 * quotient_degree, + } +} + +pub(crate) fn mixed_lde_open_row(ldes: &[CudaLde], index: usize) -> Vec> { + assert!(!ldes.is_empty()); + assert!(index < ldes.iter().map(CudaLde::height).max().unwrap()); + let total: usize = ldes.iter().map(CudaLde::width).sum(); + let handles: Vec<_> = ldes.iter().map(CudaLde::raw_handle).collect(); + let mut flat = Goldilocks::zero_vec(total); + let status = unsafe { + multi_stark_cuda_mixed_lde_open_row( + ldes[0].device_id, + flat.as_mut_ptr().cast(), + handles.as_ptr(), + handles.len(), + index, + ) + }; + check_cuda(status, "resident mixed LDE row opening"); + let mut offset = 0; + ldes.iter() + .map(|lde| { + let row = flat[offset..offset + lde.width()].to_vec(); + offset += lde.width(); + row + }) + .collect() +} + +pub(crate) fn mixed_lde_open_rows( + ldes: &[CudaLde], + indices: &[usize], +) -> Vec>> { + assert!(!ldes.is_empty()); + assert!(!indices.is_empty()); + let max_height = ldes.iter().map(CudaLde::height).max().unwrap(); + assert!(indices.iter().all(|&index| index < max_height)); + let total: usize = ldes.iter().map(CudaLde::width).sum(); + let handles: Vec<_> = ldes.iter().map(CudaLde::raw_handle).collect(); + let device_indices: Vec = indices + .iter() + .map(|&index| u64::try_from(index).expect("row index exceeds u64")) + .collect(); + let mut flat = Goldilocks::zero_vec(indices.len() * total); + let status = unsafe { + multi_stark_cuda_mixed_lde_open_rows( + ldes[0].device_id, + flat.as_mut_ptr().cast(), + handles.as_ptr(), + handles.len(), + device_indices.as_ptr(), + device_indices.len(), + ) + }; + check_cuda(status, "resident mixed LDE row batch opening"); + flat.chunks_exact(total) + .map(|query| { + let mut offset = 0; + ldes.iter() + .map(|lde| { + let row = query[offset..offset + lde.width()].to_vec(); + offset += lde.width(); + row + }) + .collect() + }) + .collect() +} + +#[cfg(test)] +pub(crate) fn lde_interpolate_ext2( + lde: &CudaLde, + height: usize, + inv_denoms: &[[Goldilocks; 2]], + coset: &[Goldilocks], + scale: [Goldilocks; 2], + ext_w: Goldilocks, +) -> Vec<[Goldilocks; 2]> { + assert_eq!(inv_denoms.len(), height); + assert_eq!(coset.len(), height); + let mut output = vec![[Goldilocks::ZERO; 2]; lde.width()]; + let one = [Goldilocks::ONE, Goldilocks::ZERO]; + let status = unsafe { + multi_stark_cuda_lde_interpolate( + lde.device_id, + output.as_mut_ptr().cast(), + lde.raw_handle(), + height, + inv_denoms.as_ptr().cast(), + coset.as_ptr().cast(), + one.as_ptr().cast(), + raw_u64(ext_w), + ) + }; + check_cuda(status, "resident LDE interpolation"); + for value in &mut output { + let [a0, a1] = *value; + *value = [ + a0 * scale[0] + ext_w * (a1 * scale[1]), + a0 * scale[1] + a1 * scale[0], + ]; + } + output +} + +pub(crate) struct CudaReducedOpening { + device_id: i32, + handle: NonNull, + height: usize, +} +impl CudaReducedOpening { + pub(crate) fn new(device_id: i32, height: usize) -> Self { + let mut handle = core::ptr::null_mut(); + let status = unsafe { multi_stark_cuda_reduced_create(device_id, &mut handle, height) }; + check_cuda(status, "reduced opening allocation"); + Self { + device_id, + handle: NonNull::new(handle).unwrap(), + height, + } + } + pub(crate) const fn height(&self) -> usize { + self.height + } + #[cfg(test)] + pub(crate) fn add( + &mut self, + lde: &CudaLde, + inv: &[[Goldilocks; 2]], + alpha: &[[Goldilocks; 2]], + reduced_y: [Goldilocks; 2], + offset: [Goldilocks; 2], + ext_w: Goldilocks, + ) { + assert_eq!(inv.len(), self.height); + assert_eq!(alpha.len(), lde.width()); + let status = unsafe { + multi_stark_cuda_reduced_add( + self.device_id, + self.handle.as_ptr(), + lde.raw_handle(), + self.height, + inv.as_ptr().cast(), + alpha.as_ptr().cast(), + reduced_y.as_ptr().cast(), + offset.as_ptr().cast(), + raw_u64(ext_w), + ) + }; + check_cuda(status, "reduced opening accumulation") + } + #[cfg(test)] + pub(crate) fn to_host(&self) -> Vec<[Goldilocks; 2]> { + let mut out = vec![[Goldilocks::ZERO; 2]; self.height]; + let status = unsafe { + multi_stark_cuda_reduced_copy( + self.device_id, + self.handle.as_ptr(), + out.as_mut_ptr().cast(), + ) + }; + check_cuda(status, "reduced opening copy"); + out + } + pub(crate) fn to_lde(&self) -> CudaLde { + let mut handle = core::ptr::null_mut(); + let status = unsafe { + multi_stark_cuda_reduced_to_lde(self.device_id, &mut handle, self.handle.as_ptr()) + }; + check_cuda(status, "reduced opening to resident FRI codeword"); + CudaLde { + device_id: self.device_id, + handle: NonNull::new(handle).unwrap(), + height: self.height, + width: 2, + } + } +} +impl Drop for CudaReducedOpening { + fn drop(&mut self) { + unsafe { + let _ = multi_stark_cuda_reduced_destroy(self.device_id, self.handle.as_ptr()); + } + } +} + +#[repr(C)] +pub(crate) struct CudaInterpolationTask { + pub(crate) lde: *const c_void, + pub(crate) height: usize, + pub(crate) inv_offset: usize, + pub(crate) output_offset: usize, + pub(crate) scale0: u64, + pub(crate) scale1: u64, +} +#[repr(C)] +pub(crate) struct CudaReductionTask { + pub(crate) reduced: *mut c_void, + pub(crate) lde: *const c_void, + pub(crate) height: usize, + pub(crate) inv_offset: usize, + pub(crate) y0: u64, + pub(crate) y1: u64, + pub(crate) offset0: u64, + pub(crate) offset1: u64, +} + +pub(crate) struct CudaFriWorkspace { + device_id: i32, + handle: NonNull, +} +impl CudaFriWorkspace { + pub(crate) fn new( + device_id: i32, + points: &[[Goldilocks; 2]], + counts: &[usize], + coset: &[Goldilocks], + ext_w: Goldilocks, + ) -> Self { + assert_eq!(points.len(), counts.len()); + let mut handle = core::ptr::null_mut(); + let status = unsafe { + multi_stark_cuda_fri_workspace_create( + device_id, + &mut handle, + points.as_ptr().cast(), + counts.as_ptr(), + points.len(), + coset.as_ptr().cast(), + coset.len(), + raw_u64(ext_w), + ) + }; + check_cuda(status, "FRI workspace allocation"); + Self { + device_id, + handle: NonNull::new(handle).unwrap(), + } + } + pub(crate) fn interpolate( + &mut self, + tasks: &[CudaInterpolationTask], + output_count: usize, + ext_w: Goldilocks, + ) -> Vec<[Goldilocks; 2]> { + let mut output = vec![[Goldilocks::ZERO; 2]; output_count]; + let status = unsafe { + multi_stark_cuda_fri_interpolate_batch( + self.device_id, + self.handle.as_ptr(), + output.as_mut_ptr().cast(), + output_count, + tasks.as_ptr().cast(), + tasks.len(), + raw_u64(ext_w), + ) + }; + check_cuda(status, "batched FRI interpolation"); + output + } + pub(crate) fn reduce( + &mut self, + tasks: &[CudaReductionTask], + alpha: &[[Goldilocks; 2]], + ext_w: Goldilocks, + ) { + let status = unsafe { + multi_stark_cuda_fri_reduce_batch( + self.device_id, + self.handle.as_ptr(), + tasks.as_ptr().cast(), + tasks.len(), + alpha.as_ptr().cast(), + alpha.len(), + raw_u64(ext_w), + ) + }; + check_cuda(status, "batched FRI reduction") + } +} +impl Drop for CudaFriWorkspace { + fn drop(&mut self) { + unsafe { + let _ = multi_stark_cuda_fri_workspace_destroy(self.device_id, self.handle.as_ptr()); + } + } +} + +#[allow(clippy::too_many_arguments)] +pub(crate) fn lookup_lde_resident( + dft: &CudaDft, + multiplicities: &[Goldilocks], + args: &[Goldilocks], + arg_offsets: &[usize], + height: usize, + num_lookups: usize, + group_size: usize, + beta: [Goldilocks; 2], + gamma: [Goldilocks; 2], + ext_w: Goldilocks, + log_blowup: usize, +) -> (CudaLde, [Goldilocks; 2]) { + assert!((1..=8).contains(&group_size)); + assert_eq!(arg_offsets.len(), num_lookups + 1); + assert_eq!(arg_offsets.first(), Some(&0)); + assert!(arg_offsets.windows(2).all(|pair| pair[0] <= pair[1])); + if num_lookups == 0 { + let extended_height = height << log_blowup; + let mut handle = core::ptr::null_mut(); + let status = unsafe { + multi_stark_cuda_zero_lde_create(dft.device_id, &mut handle, extended_height, 2) + }; + check_cuda(status, "zero lookup LDE"); + return ( + CudaLde { + device_id: dft.device_id, + handle: NonNull::new(handle).expect("CUDA returned a null zero LDE"), + height: extended_height, + width: 2, + }, + [Goldilocks::ZERO; 2], + ); + } + let slots = num_lookups.div_ceil(group_size.max(1)); + let args_width = *arg_offsets.last().unwrap(); + assert_eq!(multiplicities.len(), height * num_lookups); + assert_eq!(args.len(), height * args_width); + let extended_height = height << log_blowup; + let inverse_twiddles = dft.twiddles(log2_strict_usize(height), true); + let shift_powers = dft.shift_powers(height, Goldilocks::GENERATOR); + let forward_twiddles = dft.twiddles(log2_strict_usize(extended_height), false); + let height_inverse = Goldilocks::ONE.div_2exp_u64(log2_strict_usize(height) as u64); + let mut tail = [Goldilocks::ZERO; 4]; + let mut handle = core::ptr::null_mut(); + let status = unsafe { + multi_stark_cuda_lookup_lde( + dft.device_id, + &mut handle, + tail.as_mut_ptr().cast(), + multiplicities.as_ptr().cast(), + args.as_ptr().cast(), + arg_offsets.as_ptr(), + height, + num_lookups, + args_width, + group_size.max(1), + beta.as_ptr().cast(), + gamma.as_ptr().cast(), + raw_u64(ext_w), + log_blowup, + inverse_twiddles.as_ptr().cast(), + shift_powers.as_ptr().cast(), + forward_twiddles.as_ptr().cast(), + raw_u64(height_inverse), + ) + }; + check_cuda(status, "resident CUDA lookup LDE"); + ( + CudaLde { + device_id: dft.device_id, + handle: NonNull::new(handle).expect("CUDA returned a null lookup LDE"), + height: extended_height, + width: 2 * slots, + }, + [tail[0] + tail[2], tail[1] + tail[3]], + ) +} + +#[allow(clippy::too_many_arguments)] +pub(crate) fn lookup_graph_lde_resident( + dft: &CudaDft, + graph: &ConstraintGraph, + preprocessed: Option<&CudaLde>, + main: &CudaLde, + height: usize, + group_size: usize, + beta: [Goldilocks; 2], + gamma: [Goldilocks; 2], + ext_w: Goldilocks, + log_blowup: usize, +) -> Option<(CudaLde, [Goldilocks; 2])> { + assert!((1..=8).contains(&group_size)); + let (nodes, slot_count, lookups, args) = encode_lookup_nodes(graph)?; + let num_lookups = lookups.len(); + let groups = num_lookups.div_ceil(group_size.max(1)); + let extended_height = height << log_blowup; + let inverse_twiddles = dft.twiddles(log2_strict_usize(height), true); + let shift_powers = dft.shift_powers(height, Goldilocks::GENERATOR); + let forward_twiddles = dft.twiddles(log2_strict_usize(extended_height), false); + let height_inverse = Goldilocks::ONE.div_2exp_u64(log2_strict_usize(height) as u64); + let mut tail = [Goldilocks::ZERO; 4]; + let mut handle = core::ptr::null_mut(); + let status = unsafe { + multi_stark_cuda_lookup_graph_lde( + dft.device_id, + &mut handle, + tail.as_mut_ptr().cast(), + nodes.as_ptr().cast(), + nodes.len(), + slot_count, + lookups.as_ptr().cast(), + lookups.len(), + args.as_ptr(), + args.len(), + preprocessed.map_or(core::ptr::null(), CudaLde::raw_handle), + main.raw_handle(), + group_size.max(1), + beta.as_ptr().cast(), + gamma.as_ptr().cast(), + raw_u64(ext_w), + log_blowup, + inverse_twiddles.as_ptr().cast(), + shift_powers.as_ptr().cast(), + forward_twiddles.as_ptr().cast(), + raw_u64(height_inverse), + ) + }; + check_cuda(status, "resident CUDA graph lookup LDE"); + Some(( + CudaLde { + device_id: dft.device_id, + handle: NonNull::new(handle).expect("CUDA returned a null graph lookup LDE"), + height: extended_height, + width: 2 * groups, + }, + [tail[0] + tail[2], tail[1] + tail[3]], + )) +} + +impl CudaLde { + pub(crate) fn interpolation_task( + &self, + height: usize, + inv_offset: usize, + output_offset: usize, + scale: [Goldilocks; 2], + ) -> CudaInterpolationTask { + CudaInterpolationTask { + lde: self.raw_handle(), + height, + inv_offset, + output_offset, + scale0: raw_u64(scale[0]), + scale1: raw_u64(scale[1]), + } + } +} +impl CudaReducedOpening { + pub(crate) fn reduction_task( + &mut self, + lde: &CudaLde, + inv_offset: usize, + y: [Goldilocks; 2], + offset: [Goldilocks; 2], + ) -> CudaReductionTask { + CudaReductionTask { + reduced: self.handle.as_ptr(), + lde: lde.raw_handle(), + height: self.height, + inv_offset, + y0: raw_u64(y[0]), + y1: raw_u64(y[1]), + offset0: raw_u64(offset[0]), + offset1: raw_u64(offset[1]), + } + } +} + +impl Drop for CudaLde { + fn drop(&mut self) { + // SAFETY: this wrapper uniquely owns the handle and drops it once. + unsafe { + let _ = multi_stark_cuda_lde_destroy(self.device_id, self.handle.as_ptr()); + } + } +} + +impl TwoAdicSubgroupDft for CudaDft { + type Evaluations = BitReversedMatrixView>; + + fn dft_batch(&self, mut matrix: RowMajorMatrix) -> Self::Evaluations { + let height = matrix.height(); + let width = matrix.width(); + Self::validate_dimensions(height, width); + if height == 1 || width == 0 { + return BitReversalPerm::new_view(matrix); + } + if !Self::use_cuda_dft(height, width) { + return self.cpu.dft_batch(matrix); + } + + let twiddles = self.twiddles(log2_strict_usize(height), false); + // SAFETY: Goldilocks is repr(transparent) over u64 (asserted above), + // every u64 bit pattern is a valid Goldilocks value, all buffers have + // the element counts implied by height/width, and the FFI call is + // synchronous so the borrowed twiddle buffer outlives device use. + let status = unsafe { + multi_stark_cuda_dft_batch( + self.device_id, + matrix.values.as_mut_ptr().cast(), + height, + width, + twiddles.as_ptr().cast(), + ) + }; + check_cuda(status, "batched DFT"); + + // The CUDA DIF kernel writes bit-reversed rows. Wrap that storage so + // callers observe the natural-order evaluations required by the trait. + BitReversalPerm::new_view(matrix) + } + + fn coset_lde_batch( + &self, + matrix: RowMajorMatrix, + added_bits: usize, + shift: Goldilocks, + ) -> Self::Evaluations { + let height = matrix.height(); + let width = matrix.width(); + Self::validate_dimensions(height, width); + let extended_height = height + .checked_shl(u32::try_from(added_bits).expect("LDE blowup exceeds u32")) + .expect("LDE height overflows usize"); + Self::validate_dimensions(extended_height, width); + + if width == 0 { + return BitReversalPerm::new_view(RowMajorMatrix::new(Vec::new(), width)); + } + if height == 1 { + let mut values = Vec::with_capacity(extended_height * width); + for _ in 0..extended_height { + values.extend_from_slice(&matrix.values); + } + return BitReversalPerm::new_view(RowMajorMatrix::new(values, width)); + } + if !Self::use_cuda_coset_lde(extended_height, width) { + return self.cpu.coset_lde_batch(matrix, added_bits, shift); + } + + let log_height = log2_strict_usize(height); + let inverse_twiddles = self.twiddles(log_height, true); + let forward_twiddles = self.twiddles(log2_strict_usize(extended_height), false); + let shift_powers = self.shift_powers(height, shift); + let height_inverse = Goldilocks::ONE.div_2exp_u64(log_height as u64); + let mut output = Goldilocks::zero_vec(extended_height * width); + + // SAFETY: the input/output and cached tables have the exact lengths + // implied by the dimensions, their Goldilocks storage is u64-compatible, + // and the FFI call completes before any borrowed buffer is released. + let status = unsafe { + multi_stark_cuda_coset_lde_batch( + self.device_id, + output.as_mut_ptr().cast(), + matrix.values.as_ptr().cast(), + height, + width, + added_bits, + inverse_twiddles.as_ptr().cast(), + shift_powers.as_ptr().cast(), + forward_twiddles.as_ptr().cast(), + raw_u64(height_inverse), + ) + }; + check_cuda(status, "coset LDE"); + + BitReversalPerm::new_view(RowMajorMatrix::new(output, width)) + } +} + +#[inline] +fn raw_u64(value: Goldilocks) -> u64 { + value.as_canonical_u64() +} + +pub(crate) fn device_memory_info(device_id: i32) -> (usize, usize) { + let mut free_bytes = 0; + let mut total_bytes = 0; + let status = + unsafe { multi_stark_cuda_memory_info(device_id, &mut free_bytes, &mut total_bytes) }; + check_cuda(status, "CUDA device initialization"); + (free_bytes, total_bytes) +} + +fn check_cuda(status: i32, operation: &str) { + if status == 0 { + return; + } + // SAFETY: CUDA returns a pointer to a process-lifetime, NUL-terminated + // static error string for every status code. + let message = unsafe { + let pointer = multi_stark_cuda_error_string(status); + if pointer.is_null() { + "unknown CUDA error".into() + } else { + CStr::from_ptr(pointer).to_string_lossy() + } + }; + panic!("CUDA {operation} failed ({status}): {message}"); +} + +/// Hashes fixed-width byte rows with the first-party CUDA BLAKE3 kernel. +/// +/// This is the byte-compatible leaf/node primitive used by the resident +/// Merkle backend. Messages may span up to 32 BLAKE3 chunks (32 KiB). +#[must_use] +pub fn blake3_hash_rows(device_id: i32, messages: &[u8], message_bytes: usize) -> Vec<[u8; 32]> { + assert!(device_id >= 0, "CUDA device id must be non-negative"); + assert!(message_bytes != 0, "BLAKE3 rows must not be empty"); + assert!(message_bytes <= 32 * 1024, "BLAKE3 rows exceed 32 KiB"); + assert!( + messages.len().is_multiple_of(message_bytes), + "message buffer is not a whole number of rows" + ); + let message_count = messages.len() / message_bytes; + if message_count == 0 { + return Vec::new(); + } + let mut digests = vec![[0u8; 32]; message_count]; + // SAFETY: the input has `message_count` complete fixed-width rows, the + // output has one 32-byte digest per row, and the call is synchronous. + let status = unsafe { + multi_stark_cuda_blake3_hash_rows( + device_id, + digests.as_mut_ptr().cast(), + messages.as_ptr(), + message_bytes, + message_count, + ) + }; + check_cuda(status, "BLAKE3 row hashing"); + digests +} + +/// Builds a binary BLAKE3 Merkle tree for fixed-width rows on the GPU and +/// returns its root. Intermediate digest layers remain device-resident. +#[must_use] +#[cfg(test)] +pub(crate) fn blake3_merkle_root(device_id: i32, rows: &[u8], row_bytes: usize) -> [u8; 32] { + assert!(device_id >= 0, "CUDA device id must be non-negative"); + assert!(row_bytes != 0, "Merkle rows must not be empty"); + assert!( + rows.len().is_multiple_of(row_bytes), + "row buffer is not a whole number of rows" + ); + let row_count = rows.len() / row_bytes; + assert!( + row_count.is_power_of_two(), + "Merkle row count must be a power of two" + ); + let mut root = [0u8; 32]; + // SAFETY: `rows` contains `row_count` fixed-width rows, `root` has the + // required 32 bytes, and the call is synchronous. + let status = unsafe { + multi_stark_cuda_blake3_merkle_root( + device_id, + root.as_mut_ptr(), + rows.as_ptr(), + row_bytes, + row_count, + ) + }; + check_cuda(status, "BLAKE3 Merkle root"); + root +} + +/// Device-resident binary BLAKE3 Merkle tree and its fixed-width source rows. +#[cfg(test)] +pub(crate) struct CudaMerkleTree { + device_id: i32, + handle: NonNull, + root: [u8; 32], + row_bytes: usize, + row_count: usize, +} + +// CUDA allocations are process resources rather than thread-affine handles; +// every operation selects `device_id` before touching the allocation. +#[cfg(test)] +unsafe impl Send for CudaMerkleTree {} + +#[cfg(test)] +impl CudaMerkleTree { + /// Uploads rows once and builds every digest layer on the selected GPU. + #[must_use] + pub(crate) fn new(device_id: i32, rows: &[u8], row_bytes: usize) -> Self { + assert!(device_id >= 0, "CUDA device id must be non-negative"); + assert!(row_bytes != 0, "Merkle rows must not be empty"); + assert!( + rows.len().is_multiple_of(row_bytes), + "row buffer is not a whole number of rows" + ); + let row_count = rows.len() / row_bytes; + assert!( + row_count.is_power_of_two(), + "Merkle row count must be a power of two" + ); + let mut handle = core::ptr::null_mut(); + let mut root = [0u8; 32]; + // SAFETY: input dimensions were checked, the output pointers are valid, + // and successful creation transfers ownership to this RAII wrapper. + let status = unsafe { + multi_stark_cuda_merkle_create( + device_id, + &mut handle, + root.as_mut_ptr(), + rows.as_ptr(), + row_bytes, + row_count, + ) + }; + check_cuda(status, "resident Merkle tree creation"); + Self { + device_id, + handle: NonNull::new(handle).expect("CUDA returned a null Merkle handle"), + root, + row_bytes, + row_count, + } + } + + /// Returns the BLAKE3 root copied out during tree construction. + #[must_use] + pub(crate) const fn root(&self) -> [u8; 32] { + self.root + } + + /// Copies one queried row and its bottom-up authentication siblings. + #[must_use] + pub(crate) fn open(&self, index: usize) -> (Vec, Vec<[u8; 32]>) { + assert!(index < self.row_count, "Merkle opening index out of bounds"); + let mut row = vec![0u8; self.row_bytes]; + let mut siblings = vec![[0u8; 32]; self.row_count.trailing_zeros() as usize]; + // SAFETY: the handle remains owned by `self`; all output allocations + // have the exact sizes implied by the resident tree dimensions. + let status = unsafe { + multi_stark_cuda_merkle_open( + self.device_id, + self.handle.as_ptr(), + index, + row.as_mut_ptr(), + siblings.as_mut_ptr().cast(), + ) + }; + check_cuda(status, "resident Merkle opening"); + (row, siblings) + } +} + +#[cfg(test)] +impl Drop for CudaMerkleTree { + fn drop(&mut self) { + // SAFETY: this wrapper uniquely owns the handle and drops it once. + unsafe { + let _ = multi_stark_cuda_merkle_destroy(self.device_id, self.handle.as_ptr()); + } + } +} + +/// Device-resident digest tree for Plonky3's mixed-height binary MMCS layout. +#[doc(hidden)] +pub struct CudaMixedMerkleTree { + device_id: i32, + handle: NonNull, + root: [u8; 32], + row_count: usize, +} + +unsafe impl Send for CudaMixedMerkleTree {} +unsafe impl Sync for CudaMixedMerkleTree {} + +impl CudaMixedMerkleTree { + /// Builds a mixed-height tree from row groups ordered by descending, + /// distinct power-of-two height. Each group contains already-concatenated + /// rows for all matrices entering the MMCS at that height. + #[cfg(test)] + #[must_use] + pub(crate) fn new(device_id: i32, levels: &[(&[u8], usize, usize)]) -> Self { + assert!(device_id >= 0, "CUDA device id must be non-negative"); + assert!(!levels.is_empty(), "mixed Merkle tree has no row groups"); + for (index, (rows, row_bytes, height)) in levels.iter().copied().enumerate() { + assert!(row_bytes != 0, "mixed Merkle rows must not be empty"); + assert_eq!( + rows.len(), + row_bytes * height, + "mixed Merkle row-group dimensions disagree" + ); + assert!( + height.is_power_of_two(), + "mixed Merkle height is not a power of two" + ); + if index > 0 { + assert!( + height < levels[index - 1].2, + "mixed Merkle heights are not descending" + ); + } + } + + let row_pointers: Vec<*const u8> = + levels.iter().map(|(rows, _, _)| rows.as_ptr()).collect(); + let row_bytes: Vec = levels.iter().map(|(_, width, _)| *width).collect(); + let heights: Vec = levels.iter().map(|(_, _, height)| *height).collect(); + let mut handle = core::ptr::null_mut(); + let mut root = [0u8; 32]; + // SAFETY: every level buffer and dimension was validated above; the + // call copies all host inputs synchronously and returns an owned handle. + let status = unsafe { + multi_stark_cuda_mixed_merkle_create( + device_id, + &mut handle, + root.as_mut_ptr(), + row_pointers.as_ptr(), + row_bytes.as_ptr(), + heights.as_ptr(), + levels.len(), + ) + }; + check_cuda(status, "resident mixed-height Merkle tree creation"); + Self { + device_id, + handle: NonNull::new(handle).expect("CUDA returned a null mixed Merkle handle"), + root, + row_count: heights[0], + } + } + + /// Builds the exact mixed-height tree directly from resident bit-reversed + /// LDE matrices. Matrix rows are concatenated in slice order at each + /// height using device-to-device copies. + #[must_use] + pub(crate) fn from_ldes(device_id: i32, ldes: &[CudaLde]) -> Self { + assert!(device_id >= 0, "CUDA device id must be non-negative"); + assert!(!ldes.is_empty(), "mixed Merkle tree has no resident LDEs"); + assert!( + ldes.iter().all(|lde| lde.device_id == device_id), + "resident LDEs belong to different CUDA devices" + ); + let handles: Vec<*const c_void> = ldes + .iter() + .map(|lde| lde.handle.as_ptr().cast_const()) + .collect(); + let mut handle = core::ptr::null_mut(); + let mut root = [0u8; 32]; + // SAFETY: every handle is live for the synchronous construction call; + // the resulting tree owns only its own digest allocation. + let status = unsafe { + multi_stark_cuda_mixed_merkle_create_from_ldes( + device_id, + &mut handle, + root.as_mut_ptr(), + handles.as_ptr(), + handles.len(), + ) + }; + check_cuda(status, "resident LDE Merkle tree creation"); + Self { + device_id, + handle: NonNull::new(handle).expect("CUDA returned a null mixed Merkle handle"), + root, + row_count: ldes.iter().map(CudaLde::height).max().unwrap(), + } + } + + pub(crate) fn from_fri_codeword(codeword: &CudaLde, arity: usize) -> Self { + assert_eq!(codeword.width, 2); + assert!(arity.is_power_of_two()); + assert_eq!(codeword.height % arity, 0); + let mut handle = core::ptr::null_mut(); + let mut root = [0u8; 32]; + let status = unsafe { + multi_stark_cuda_fri_merkle_create( + codeword.device_id, + &mut handle, + root.as_mut_ptr(), + codeword.raw_handle(), + arity, + ) + }; + check_cuda(status, "resident CUDA FRI Merkle tree creation"); + Self { + device_id: codeword.device_id, + handle: NonNull::new(handle).expect("CUDA returned a null FRI Merkle handle"), + root, + row_count: codeword.height / arity, + } + } + + /// Builds a mixed-height tree by uploading each host matrix directly, + /// without constructing a concatenated host-side row buffer. + #[cfg(test)] + #[must_use] + pub(crate) fn from_host_matrices( + device_id: i32, + matrices: &[RowMajorMatrix], + ) -> Self { + assert!(device_id >= 0, "CUDA device id must be non-negative"); + assert!(!matrices.is_empty(), "mixed Merkle tree has no matrices"); + let values: Vec<*const u64> = matrices + .iter() + .map(|matrix| matrix.values.as_ptr().cast()) + .collect(); + let heights: Vec<_> = matrices.iter().map(Matrix::height).collect(); + let widths: Vec<_> = matrices.iter().map(Matrix::width).collect(); + let mut handle = core::ptr::null_mut(); + let mut root = [0u8; 32]; + // SAFETY: all matrix buffers remain live for the synchronous upload; + // the returned tree owns only its device digest layers. + let status = unsafe { + multi_stark_cuda_mixed_merkle_create_from_host_matrices( + device_id, + &mut handle, + root.as_mut_ptr(), + values.as_ptr(), + heights.as_ptr(), + widths.as_ptr(), + matrices.len(), + ) + }; + check_cuda(status, "host-matrix CUDA Merkle tree creation"); + Self { + device_id, + handle: NonNull::new(handle).expect("CUDA returned a null mixed Merkle handle"), + root, + row_count: heights.into_iter().max().unwrap(), + } + } + + #[must_use] + pub(crate) const fn root(&self) -> [u8; 32] { + self.root + } + + #[must_use] + pub(crate) fn open_siblings(&self, index: usize) -> Vec<[u8; 32]> { + assert!( + index < self.row_count, + "mixed Merkle opening index out of bounds" + ); + let mut siblings = vec![[0u8; 32]; self.row_count.trailing_zeros() as usize]; + // SAFETY: the handle is live and the proof allocation has one sibling + // for every binary layer. + let status = unsafe { + multi_stark_cuda_mixed_merkle_open( + self.device_id, + self.handle.as_ptr(), + index, + siblings.as_mut_ptr().cast(), + ) + }; + check_cuda(status, "resident mixed-height Merkle opening"); + siblings + } + + #[must_use] + pub(crate) fn open_siblings_batch(&self, indices: &[usize]) -> Vec> { + assert!(!indices.is_empty(), "mixed Merkle opening batch is empty"); + assert!( + indices.iter().all(|&index| index < self.row_count), + "mixed Merkle opening index out of bounds" + ); + let levels = self.row_count.trailing_zeros() as usize; + let device_indices: Vec = indices + .iter() + .map(|&index| u64::try_from(index).expect("Merkle index exceeds u64")) + .collect(); + let mut siblings = vec![[0u8; 32]; indices.len() * levels]; + let status = unsafe { + multi_stark_cuda_mixed_merkle_open_batch( + self.device_id, + self.handle.as_ptr(), + device_indices.as_ptr(), + device_indices.len(), + siblings.as_mut_ptr().cast(), + ) + }; + check_cuda(status, "resident mixed-height Merkle opening batch"); + siblings + .chunks_exact(levels) + .map(<[[u8; 32]]>::to_vec) + .collect() + } +} + +impl Drop for CudaMixedMerkleTree { + fn drop(&mut self) { + // SAFETY: this wrapper uniquely owns the handle and drops it once. + unsafe { + let _ = multi_stark_cuda_mixed_merkle_destroy(self.device_id, self.handle.as_ptr()); + } + } +} + +unsafe extern "C" { + fn multi_stark_cuda_memory_info( + device_id: i32, + free_bytes: *mut usize, + total_bytes: *mut usize, + ) -> i32; + + fn multi_stark_cuda_dft_batch( + device_id: i32, + values: *mut u64, + height: usize, + width: usize, + twiddles: *const u64, + ) -> i32; + + fn multi_stark_cuda_coset_lde_batch( + device_id: i32, + output: *mut u64, + input: *const u64, + height: usize, + width: usize, + added_bits: usize, + inverse_twiddles: *const u64, + shift_powers: *const u64, + forward_twiddles: *const u64, + height_inverse: u64, + ) -> i32; + + fn multi_stark_cuda_coset_lde_create( + device_id: i32, + handle: *mut *mut c_void, + input: *const u64, + height: usize, + width: usize, + added_bits: usize, + inverse_twiddles: *const u64, + shift_powers: *const u64, + forward_twiddles: *const u64, + height_inverse: u64, + ) -> i32; + + fn multi_stark_cuda_prepare_lde_constants( + device_id: i32, + inverse_twiddles: *const u64, + inverse_count: usize, + shift_powers: *const u64, + height: usize, + forward_twiddles: *const u64, + forward_count: usize, + ) -> i32; + + fn multi_stark_cuda_lde_create_from_host( + device_id: i32, + handle: *mut *mut c_void, + input: *const u64, + height: usize, + width: usize, + ) -> i32; + fn multi_stark_cuda_zero_lde_create( + device_id: i32, + handle: *mut *mut c_void, + height: usize, + width: usize, + ) -> i32; + + fn multi_stark_cuda_lde_copy_to_host( + device_id: i32, + handle: *const c_void, + output: *mut u64, + ) -> i32; + fn multi_stark_cuda_lde_release_trace(device_id: i32, handle: *mut c_void) -> i32; + fn multi_stark_cuda_lde_attach_trace( + device_id: i32, + handle: *mut c_void, + trace: *const u64, + height: usize, + width: usize, + ) -> i32; + + #[cfg(test)] + fn multi_stark_cuda_lde_copy_row( + device_id: i32, + handle: *const c_void, + row: usize, + output: *mut u64, + ) -> i32; + + fn multi_stark_cuda_lde_copy_rows( + device_id: i32, + handle: *const c_void, + rows: *const u64, + row_count: usize, + output: *mut u64, + ) -> i32; + + fn multi_stark_cuda_lde_destroy(device_id: i32, handle: *mut c_void) -> i32; + + #[cfg(test)] + fn multi_stark_cuda_constraint_graph( + device_id: i32, + output: *mut u64, + nodes: *const c_void, + node_count: usize, + roots: *const u32, + root_count: usize, + preprocessed_handle: *const c_void, + main_handle: *const c_void, + stage2_handle: *const c_void, + publics: *const u64, + public_count: usize, + selectors: *const u64, + quotient_size: usize, + next_step: usize, + ) -> i32; + fn multi_stark_cuda_quotient_values( + device_id: i32, + output: *mut u64, + nodes: *const c_void, + node_count: usize, + slot_count: usize, + roots: *const u32, + root_count: usize, + lookups: *const c_void, + lookup_count: usize, + lookup_args: *const u32, + lookup_arg_count: usize, + group_size: usize, + preprocessed_handle: *const c_void, + main_handle: *const c_void, + stage2_handle: *const c_void, + publics: *const u64, + public_count: usize, + coset_shift: u64, + coset_generator: u64, + trace_last: u64, + vanishing_start: u64, + vanishing_step: u64, + alpha: *const u64, + constraint_count: usize, + delta: *const u64, + ext_w: u64, + quotient_size: usize, + next_step: usize, + ) -> i32; + fn multi_stark_cuda_quotient_lde( + device_id: i32, + output_handle: *mut *mut c_void, + nodes: *const c_void, + node_count: usize, + slot_count: usize, + roots: *const u32, + root_count: usize, + lookups: *const c_void, + lookup_count: usize, + lookup_args: *const u32, + lookup_arg_count: usize, + group_size: usize, + preprocessed_handle: *const c_void, + main_handle: *const c_void, + stage2_handle: *const c_void, + publics: *const u64, + public_count: usize, + coset_shift: u64, + coset_generator: u64, + trace_last: u64, + vanishing_start: u64, + vanishing_step: u64, + alpha: *const u64, + constraint_count: usize, + delta: *const u64, + ext_w: u64, + quotient_size: usize, + next_step: usize, + quotient_degree: usize, + log_blowup: usize, + quotient_twiddles: *const u64, + lde_twiddles: *const u64, + slice_weights: *const u64, + ) -> i32; + fn multi_stark_cuda_mixed_lde_open_row( + device_id: i32, + output: *mut u64, + handles: *const *const c_void, + handle_count: usize, + index: usize, + ) -> i32; + fn multi_stark_cuda_mixed_lde_open_rows( + device_id: i32, + output: *mut u64, + handles: *const *const c_void, + handle_count: usize, + indices: *const u64, + query_count: usize, + ) -> i32; + #[cfg(test)] + fn multi_stark_cuda_lde_interpolate( + device_id: i32, + output: *mut u64, + handle: *const c_void, + height: usize, + inv_denoms: *const u64, + coset: *const u64, + scale: *const u64, + ext_w: u64, + ) -> i32; + fn multi_stark_cuda_fri_workspace_create( + device_id: i32, + handle: *mut *mut c_void, + points: *const u64, + counts: *const usize, + point_count: usize, + coset: *const u64, + coset_count: usize, + ext_w: u64, + ) -> i32; + fn multi_stark_cuda_fri_interpolate_batch( + device_id: i32, + handle: *mut c_void, + output: *mut u64, + output_count: usize, + tasks: *const c_void, + task_count: usize, + ext_w: u64, + ) -> i32; + fn multi_stark_cuda_fri_reduce_batch( + device_id: i32, + handle: *mut c_void, + tasks: *const c_void, + task_count: usize, + alpha: *const u64, + alpha_count: usize, + ext_w: u64, + ) -> i32; + fn multi_stark_cuda_fri_workspace_destroy(device_id: i32, handle: *mut c_void) -> i32; + fn multi_stark_cuda_reduced_to_lde( + device_id: i32, + output: *mut *mut c_void, + reduced: *const c_void, + ) -> i32; + fn multi_stark_cuda_fri_fold_resident( + device_id: i32, + output: *mut *mut c_void, + input: *const c_void, + next_reduced: *const c_void, + beta: *const u64, + log_arity: usize, + beta_power0: u64, + beta_power1: u64, + g_inv: u64, + ext_w: u64, + ) -> i32; + fn multi_stark_cuda_lookup_graph_lde( + device_id: i32, + output_handle: *mut *mut c_void, + total: *mut u64, + nodes: *const c_void, + node_count: usize, + slot_count: usize, + lookups: *const c_void, + lookup_count: usize, + lookup_args: *const u32, + lookup_arg_count: usize, + preprocessed_handle: *const c_void, + main_handle: *const c_void, + group_size: usize, + beta: *const u64, + gamma: *const u64, + ext_w: u64, + added_bits: usize, + inverse_twiddles: *const u64, + shift_powers: *const u64, + forward_twiddles: *const u64, + height_inverse: u64, + ) -> i32; + fn multi_stark_cuda_lookup_lde( + device_id: i32, + output_handle: *mut *mut c_void, + total: *mut u64, + multiplicities: *const u64, + args: *const u64, + arg_offsets: *const usize, + height: usize, + num_lookups: usize, + args_width: usize, + group_size: usize, + beta: *const u64, + gamma: *const u64, + ext_w: u64, + added_bits: usize, + inverse_twiddles: *const u64, + shift_powers: *const u64, + forward_twiddles: *const u64, + height_inverse: u64, + ) -> i32; + fn multi_stark_cuda_reduced_create( + device_id: i32, + handle: *mut *mut c_void, + height: usize, + ) -> i32; + #[cfg(test)] + fn multi_stark_cuda_reduced_add( + device_id: i32, + reduced: *mut c_void, + lde: *const c_void, + height: usize, + inv: *const u64, + alpha: *const u64, + reduced_y: *const u64, + offset: *const u64, + ext_w: u64, + ) -> i32; + #[cfg(test)] + fn multi_stark_cuda_reduced_copy( + device_id: i32, + handle: *const c_void, + output: *mut u64, + ) -> i32; + fn multi_stark_cuda_reduced_destroy(device_id: i32, handle: *mut c_void) -> i32; + + fn multi_stark_cuda_error_string(status: i32) -> *const c_char; + + fn multi_stark_cuda_blake3_hash_rows( + device_id: i32, + digests: *mut u8, + messages: *const u8, + message_bytes: usize, + message_count: usize, + ) -> i32; + + #[cfg(test)] + fn multi_stark_cuda_blake3_merkle_root( + device_id: i32, + root: *mut u8, + rows: *const u8, + row_bytes: usize, + row_count: usize, + ) -> i32; + + #[cfg(test)] + fn multi_stark_cuda_merkle_create( + device_id: i32, + handle: *mut *mut c_void, + root: *mut u8, + rows: *const u8, + row_bytes: usize, + row_count: usize, + ) -> i32; + + #[cfg(test)] + fn multi_stark_cuda_merkle_open( + device_id: i32, + handle: *const c_void, + index: usize, + row: *mut u8, + siblings: *mut u8, + ) -> i32; + + #[cfg(test)] + fn multi_stark_cuda_merkle_destroy(device_id: i32, handle: *mut c_void) -> i32; + + #[cfg(test)] + fn multi_stark_cuda_mixed_merkle_create( + device_id: i32, + handle: *mut *mut c_void, + root: *mut u8, + level_rows: *const *const u8, + level_row_bytes: *const usize, + level_heights: *const usize, + level_count: usize, + ) -> i32; + + fn multi_stark_cuda_mixed_merkle_open( + device_id: i32, + handle: *const c_void, + index: usize, + siblings: *mut u8, + ) -> i32; + fn multi_stark_cuda_mixed_merkle_open_batch( + device_id: i32, + handle: *const c_void, + indices: *const u64, + query_count: usize, + siblings: *mut u8, + ) -> i32; + + fn multi_stark_cuda_mixed_merkle_destroy(device_id: i32, handle: *mut c_void) -> i32; + + fn multi_stark_cuda_mixed_merkle_create_from_ldes( + device_id: i32, + handle: *mut *mut c_void, + root: *mut u8, + lde_handles: *const *const c_void, + lde_count: usize, + ) -> i32; + fn multi_stark_cuda_fri_merkle_create( + device_id: i32, + handle: *mut *mut c_void, + root: *mut u8, + codeword: *const c_void, + arity: usize, + ) -> i32; + + #[cfg(test)] + fn multi_stark_cuda_mixed_merkle_create_from_host_matrices( + device_id: i32, + handle: *mut *mut c_void, + root: *mut u8, + matrix_values: *const *const u64, + heights: *const usize, + widths: *const usize, + matrix_count: usize, + ) -> i32; + + #[cfg(test)] + fn multi_stark_cuda_generate_coset_selectors( + device_id: i32, + output: *mut u64, + quotient_size: usize, + next_step: usize, + coset_shift: u64, + coset_generator: u64, + trace_last: u64, + vanishing_start: u64, + vanishing_step: u64, + ) -> i32; +} + +#[cfg(test)] +mod tests { + use super::*; + use p3_dft::Radix2DitParallel; + use p3_matrix::bitrev::BitReversibleMatrix; + use p3_symmetric::CryptographicHasher; + use rand::{RngExt, SeedableRng, rngs::SmallRng}; + + #[test] + fn device_coset_selectors_match_cpu() { + use p3_commit::PolynomialSpace; + use p3_field::coset::TwoAdicMultiplicativeCoset; + + for (trace_log, quotient_log) in [(4, 5), (8, 11), (12, 16), (20, 24)] { + let trace = + TwoAdicMultiplicativeCoset::::new(Goldilocks::ONE, trace_log).unwrap(); + let quotient = + TwoAdicMultiplicativeCoset::::new(Goldilocks::GENERATOR, quotient_log) + .unwrap(); + let expected = trace.selectors_on_coset(quotient); + let quotient_size = quotient.size(); + let next_step = quotient_size / trace.size(); + let mut actual = vec![Goldilocks::ZERO; 4 * quotient_size]; + let status = unsafe { + multi_stark_cuda_generate_coset_selectors( + 0, + actual.as_mut_ptr().cast(), + quotient_size, + next_step, + raw_u64(quotient.shift()), + raw_u64(quotient.subgroup_generator()), + raw_u64(trace.subgroup_generator().inverse()), + raw_u64(quotient.shift().exp_power_of_2(trace_log)), + raw_u64(Goldilocks::two_adic_generator(quotient_log - trace_log)), + ) + }; + check_cuda(status, "coset selector generation"); + let mut flattened = expected.is_first_row; + flattened.extend(expected.is_last_row); + flattened.extend(expected.is_transition); + flattened.extend(expected.inv_vanishing); + assert_eq!(actual, flattened); + } + } + + #[test] + fn host_uploaded_lde_is_canonicalized() { + const MODULUS: u64 = 0xffff_ffff_0000_0001; + // SAFETY: Goldilocks is repr(transparent) over u64 and accepts every + // bit pattern; this deliberately constructs lazy representatives. + let values = [ + unsafe { core::mem::transmute::(MODULUS + 5) }, + unsafe { core::mem::transmute::(MODULUS + 9) }, + ]; + let lde = CudaLde::from_row_major_matrix(0, &RowMajorMatrix::new(values.to_vec(), 1)); + assert_eq!( + lde.to_row_major_matrix().values, + [Goldilocks::from_u64(5), Goldilocks::from_u64(9)] + ); + } + + #[test] + fn resident_constraint_graph_matches_cpu() { + use crate::expr::{CircuitSpec, ColRef, Expr, RowOffset, Source}; + use crate::graph::{ExtensionParams, compile}; + let x = Expr::Var(ColRef { + source: Source::Main, + offset: RowOffset::Current, + index: 0, + }); + let next = Expr::Var(ColRef { + source: Source::Main, + offset: RowOffset::Next, + index: 0, + }); + let spec = CircuitSpec { + main_width: 1, + preprocessed_width: 0, + stage2_width: 1, + num_publics: 1, + constraints: vec![x.clone() * x + next - Expr::Public(0)], + ext_constraints: vec![], + lookups: vec![], + }; + let graph = compile( + &spec, + &ExtensionParams { + degree: 2, + w: Goldilocks::from_u64(7), + karatsuba: true, + }, + ) + .unwrap(); + let height = 8; + let logical: Vec<_> = (0..height).map(|i| Goldilocks::from_usize(i + 2)).collect(); + let bitrev = RowMajorMatrix::new_col(logical.clone()) + .bit_reverse_rows() + .to_row_major_matrix(); + let main = CudaLde::from_row_major_matrix(0, &bitrev); + let stage2 = CudaLde::from_row_major_matrix( + 0, + &RowMajorMatrix::new_col(vec![Goldilocks::ZERO; height]), + ); + let public = Goldilocks::from_u64(5); + let selectors = vec![Goldilocks::ZERO; 3 * height]; + let got = constraint_graph_roots( + &graph, + None, + &main, + &stage2, + &[public], + &selectors, + height, + 2, + ); + for row in 0..height { + assert_eq!( + got.values[row], + logical[row] * logical[row] + logical[(row + 2) % height] - public + ); + } + } + + #[test] + fn resident_interpolation_matches_cpu() { + use p3_field::coset::TwoAdicMultiplicativeCoset; + use p3_field::{ + BasedVectorSpace, batch_multiplicative_inverse, extension::BinomialExtensionField, + }; + use p3_interpolation::interpolate_coset_with_precomputation; + use p3_util::reverse_slice_index_bits; + type Ext = BinomialExtensionField; + let height = 256; + let storage_height = 512; + let width = 1; + let mut rng = SmallRng::seed_from_u64(991); + let logical = RowMajorMatrix::new( + (0..storage_height * width) + .map(|_| rng.random::()) + .collect(), + width, + ); + let bitrev = logical.clone().bit_reverse_rows().to_row_major_matrix(); + let lde = CudaLde::from_row_major_matrix(0, &bitrev); + let mut coset: Vec<_> = TwoAdicMultiplicativeCoset::new(Goldilocks::GENERATOR, 9) + .unwrap() + .iter() + .collect(); + reverse_slice_index_bits(&mut coset); + let point = Ext::from_basis_coefficients_slice(&[ + Goldilocks::from_u64(5366541241408596314), + Goldilocks::from_u64(11432508775637878761), + ]) + .unwrap(); + let inv = + batch_multiplicative_inverse(&coset.iter().map(|&x| point - x).collect::>()); + let expected = interpolate_coset_with_precomputation( + &bitrev.split_rows(height).0, + Goldilocks::GENERATOR, + point, + &coset[..height], + &inv[..height], + ); + let inv2: Vec<[Goldilocks; 2]> = inv[..height] + .iter() + .map(|x| x.as_basis_coefficients_slice().try_into().unwrap()) + .collect(); + let shift_pow = Goldilocks::GENERATOR.exp_power_of_2(8); + let scale = (point.exp_power_of_2(8) - shift_pow) + * (Goldilocks::from_usize(height) * shift_pow).inverse(); + let got = lde_interpolate_ext2( + &lde, + height, + &inv2, + &coset[..height], + scale.as_basis_coefficients_slice().try_into().unwrap(), + Goldilocks::from_u64(7), + ); + for (got, want) in got.iter().zip(expected) { + assert_eq!(got.as_slice(), want.as_basis_coefficients_slice()); + } + } + + #[test] + fn resident_reduced_opening_matches_cpu() { + use p3_field::coset::TwoAdicMultiplicativeCoset; + use p3_field::{ + BasedVectorSpace, batch_multiplicative_inverse, extension::BinomialExtensionField, + }; + use p3_util::reverse_slice_index_bits; + type Ext = BinomialExtensionField; + let height = 256; + let width = 5; + let logical = RowMajorMatrix::new( + (0..height * width) + .map(|i| Goldilocks::from_usize(i * 17 + 3)) + .collect(), + width, + ); + let bitrev = logical.bit_reverse_rows().to_row_major_matrix(); + let lde = CudaLde::from_row_major_matrix(0, &bitrev); + let mut coset: Vec<_> = TwoAdicMultiplicativeCoset::new(Goldilocks::GENERATOR, 8) + .unwrap() + .iter() + .collect(); + reverse_slice_index_bits(&mut coset); + let point = Ext::from_basis_coefficients_slice(&[ + Goldilocks::from_u64(31), + Goldilocks::from_u64(37), + ]) + .unwrap(); + let inv = + batch_multiplicative_inverse(&coset.iter().map(|&x| point - x).collect::>()); + let inv2: Vec<[Goldilocks; 2]> = inv + .iter() + .map(|x| x.as_basis_coefficients_slice().try_into().unwrap()) + .collect(); + let alpha = Ext::from_basis_coefficients_slice(&[ + Goldilocks::from_u64(41), + Goldilocks::from_u64(43), + ]) + .unwrap(); + let ap: Vec<_> = alpha.powers().take(width).collect(); + let ap2: Vec<[Goldilocks; 2]> = ap + .iter() + .map(|x| x.as_basis_coefficients_slice().try_into().unwrap()) + .collect(); + let ys: Vec<_> = (0..width) + .map(|c| point + Goldilocks::from_usize(c)) + .collect(); + let reduced_y: Ext = ap.iter().zip(&ys).map(|(&a, &y)| a * y).sum(); + let offset = alpha.exp_u64(7); + let mut expected = vec![Ext::ZERO; height]; + for row in 0..height { + let compressed: Ext = (0..width) + .map(|c| ap[c] * bitrev.values[row * width + c]) + .sum(); + expected[row] = offset * (reduced_y - compressed) * inv[row]; + } + let mut got = CudaReducedOpening::new(0, height); + got.add( + &lde, + &inv2, + &ap2, + reduced_y.as_basis_coefficients_slice().try_into().unwrap(), + offset.as_basis_coefficients_slice().try_into().unwrap(), + Goldilocks::from_u64(7), + ); + for (got, want) in got.to_host().iter().zip(expected) { + assert_eq!(got.as_slice(), want.as_basis_coefficients_slice()); + } + } + + #[test] + fn goldilocks_field_kernels_match_cpu() { + let mut rng = SmallRng::seed_from_u64(0xc0da); + let mut left = vec![Goldilocks::ZERO, Goldilocks::ONE]; + let mut right = vec![Goldilocks::ONE, Goldilocks::ZERO]; + left.extend((0..4096).map(|_| rng.random::())); + right.extend((0..4096).map(|_| rng.random::())); + + let mut sums = Goldilocks::zero_vec(left.len()); + let mut differences = Goldilocks::zero_vec(left.len()); + let mut products = Goldilocks::zero_vec(left.len()); + let mut inverses = Goldilocks::zero_vec(left.len()); + // SAFETY: every input/output allocation contains `left.len()` valid + // u64-compatible Goldilocks elements and the call is synchronous. + let status = unsafe { + multi_stark_cuda_goldilocks_ops( + 0, + sums.as_mut_ptr().cast(), + differences.as_mut_ptr().cast(), + products.as_mut_ptr().cast(), + inverses.as_mut_ptr().cast(), + left.as_ptr().cast(), + right.as_ptr().cast(), + left.len(), + ) + }; + check_cuda(status, "Goldilocks arithmetic contract"); + + for i in 0..left.len() { + assert_eq!(sums[i], left[i] + right[i], "sum {i}"); + assert_eq!(differences[i], left[i] - right[i], "difference {i}"); + assert_eq!(products[i], left[i] * right[i], "product {i}"); + let expected_inverse = left[i].try_inverse().unwrap_or(Goldilocks::ZERO); + assert_eq!(inverses[i], expected_inverse, "inverse {i}"); + assert!(sums[i].as_canonical_u64() < Goldilocks::ORDER_U64); + assert!(products[i].as_canonical_u64() < Goldilocks::ORDER_U64); + } + } + + #[test] + fn blake3_rows_match_cpu_across_chunk_boundaries() { + use p3_blake3::Blake3; + + for message_bytes in [1usize, 63, 64, 65, 1023, 1024, 1025, 4264, 7400] { + let message_count = 17; + let messages: Vec = (0..message_bytes * message_count) + .map(|index| (index as u64).wrapping_mul(0x9e37_79b9).to_le_bytes()[0]) + .collect(); + let mut digests = vec![0u8; 32 * message_count]; + // SAFETY: the input contains `message_count` fixed-size messages, + // the output has one 32-byte digest per message, and the call is + // synchronous. + let status = unsafe { + multi_stark_cuda_blake3_hash_rows( + 0, + digests.as_mut_ptr(), + messages.as_ptr(), + message_bytes, + message_count, + ) + }; + check_cuda(status, "BLAKE3 row hashing contract"); + + for (index, message) in messages.chunks_exact(message_bytes).enumerate() { + let expected: [u8; 32] = Blake3.hash_iter(message.iter().copied()); + assert_eq!( + &digests[index * 32..(index + 1) * 32], + &expected, + "message_bytes={message_bytes}, index={index}" + ); + } + } + } + + #[test] + fn blake3_merkle_root_matches_cpu() { + use p3_blake3::Blake3; + + for (row_bytes, row_count) in [(16usize, 1usize), (64, 8), (4264, 1024)] { + let rows: Vec = (0..row_bytes * row_count) + .map(|index| (index as u64).wrapping_mul(0x517c_c1b7).to_le_bytes()[0]) + .collect(); + let mut layer: Vec<[u8; 32]> = rows + .chunks_exact(row_bytes) + .map(|row| Blake3.hash_iter(row.iter().copied())) + .collect(); + while layer.len() > 1 { + layer = layer + .chunks_exact(2) + .map(|children| { + Blake3.hash_iter(children[0].iter().chain(&children[1]).copied()) + }) + .collect(); + } + let expected_root = layer[0]; + assert_eq!(blake3_merkle_root(0, &rows, row_bytes), expected_root); + + let tree = CudaMerkleTree::new(0, &rows, row_bytes); + assert_eq!(tree.root(), expected_root); + let mut index = row_count / 2; + let (opened_row, siblings) = tree.open(index); + assert_eq!(opened_row, rows[index * row_bytes..(index + 1) * row_bytes]); + let mut digest: [u8; 32] = Blake3.hash_iter(opened_row); + for sibling in siblings { + digest = if index & 1 == 0 { + Blake3.hash_iter(digest.iter().chain(&sibling).copied()) + } else { + Blake3.hash_iter(sibling.iter().chain(&digest).copied()) + }; + index >>= 1; + } + assert_eq!(digest, expected_root); + } + } + + #[test] + fn mixed_height_merkle_matches_cpu_layout() { + use p3_blake3::Blake3; + + let level_8: Vec = (0usize..8 * 16) + .map(|index| index.to_le_bytes()[0]) + .collect(); + let level_4: Vec = (0usize..4 * 24) + .map(|index| index.wrapping_mul(17).to_le_bytes()[0]) + .collect(); + let level_1: Vec = (0usize..8) + .map(|index| index.wrapping_mul(29).to_le_bytes()[0]) + .collect(); + let levels = [ + (level_8.as_slice(), 16usize, 8usize), + (level_4.as_slice(), 24, 4), + (level_1.as_slice(), 8, 1), + ]; + + let mut layer: Vec<[u8; 32]> = level_8 + .chunks_exact(16) + .map(|row| Blake3.hash_iter(row.iter().copied())) + .collect(); + let mut digest_layers = vec![layer.clone()]; + while layer.len() > 1 { + layer = layer + .chunks_exact(2) + .map(|children| Blake3.hash_iter(children.iter().flatten().copied())) + .collect(); + let injected = match layer.len() { + 4 => Some((level_4.as_slice(), 24)), + 1 => Some((level_1.as_slice(), 8)), + _ => None, + }; + if let Some((rows, row_bytes)) = injected { + let row_digests = rows + .chunks_exact(row_bytes) + .map(|row| Blake3.hash_iter(row.iter().copied())); + layer = layer + .into_iter() + .zip(row_digests) + .map(|(digest, row_digest)| { + Blake3.hash_iter(digest.iter().chain(&row_digest).copied()) + }) + .collect(); + } + digest_layers.push(layer.clone()); + } + + let tree = CudaMixedMerkleTree::new(0, &levels); + assert_eq!(tree.root(), layer[0]); + let mut index = 5usize; + for (actual, expected_layer) in tree.open_siblings(index).into_iter().zip(&digest_layers) { + assert_eq!(actual, expected_layer[index ^ 1]); + index >>= 1; + } + } + + #[test] + fn batched_dft_matches_cpu() { + let mut rng = SmallRng::seed_from_u64(0xd17); + let cpu = Radix2DitParallel::::default(); + let gpu = CudaDft::default(); + for log_height in [0usize, 1, 2, 5, 8, 12, 14] { + for width in [1usize, 2, 7, 31] { + let height = 1 << log_height; + let matrix = + RowMajorMatrix::new((0..height * width).map(|_| rng.random()).collect(), width); + let expected = cpu.dft_batch(matrix.clone()).to_row_major_matrix(); + let actual = gpu.dft_batch(matrix).to_row_major_matrix(); + assert_eq!(actual, expected, "height=2^{log_height}, width={width}"); + } + } + } + + #[test] + fn coset_lde_matches_cpu_including_storage_layout() { + let mut rng = SmallRng::seed_from_u64(0x1de); + let cpu = Radix2DitParallel::::default(); + let gpu = CudaDft::default(); + for log_height in [0usize, 1, 2, 5, 8, 12, 14] { + for added_bits in [0usize, 1, 2, 3] { + for width in [1usize, 2, 7] { + let height = 1 << log_height; + let matrix = RowMajorMatrix::new( + (0..height * width).map(|_| rng.random()).collect(), + width, + ); + let expected = cpu + .coset_lde_batch(matrix.clone(), added_bits, Goldilocks::GENERATOR) + .bit_reverse_rows() + .to_row_major_matrix(); + let actual = gpu + .coset_lde_batch(matrix, added_bits, Goldilocks::GENERATOR) + .bit_reverse_rows() + .to_row_major_matrix(); + assert_eq!( + actual, expected, + "height=2^{log_height}, blowup=2^{added_bits}, width={width}" + ); + } + } + } + } + + #[test] + fn resident_coset_lde_matches_cpu_storage() { + let mut rng = SmallRng::seed_from_u64(0x51de); + let cpu = Radix2DitParallel::::default(); + let gpu = CudaDft::default(); + for (log_height, width, added_bits) in [ + (10usize, 925usize, 1usize), + (12, 17, 1), + (12, 2, 1), + (14, 2, 2), + (16, 1, 1), + ] { + let height = 1 << log_height; + let matrix = + RowMajorMatrix::new((0..height * width).map(|_| rng.random()).collect(), width); + let expected = cpu + .coset_lde_batch(matrix.clone(), added_bits, Goldilocks::GENERATOR) + .bit_reverse_rows() + .to_row_major_matrix(); + let resident = gpu.coset_lde_batch_resident(&matrix, added_bits, Goldilocks::GENERATOR); + assert_eq!(resident.height(), expected.height()); + assert_eq!(resident.width(), expected.width()); + assert_eq!(resident.to_row_major_matrix(), expected); + } + } + + #[test] + fn resident_ldes_feed_mixed_merkle_without_host_round_trip() { + let mut rng = SmallRng::seed_from_u64(0x1de_cafe); + let gpu = CudaDft::default(); + let inputs: Vec> = [(4usize, 2usize), (4, 1), (2, 3)] + .into_iter() + .map(|(height, width)| { + RowMajorMatrix::new((0..height * width).map(|_| rng.random()).collect(), width) + }) + .collect(); + let ldes: Vec<_> = inputs + .iter() + .map(|matrix| gpu.coset_lde_batch_resident(matrix, 1, Goldilocks::GENERATOR)) + .collect(); + let host_ldes: Vec<_> = ldes.iter().map(CudaLde::to_row_major_matrix).collect(); + + let mut packed_levels = Vec::new(); + for height in [8usize, 4] { + let matrices: Vec<_> = host_ldes + .iter() + .filter(|matrix| matrix.height() == height) + .collect(); + let total_width: usize = matrices.iter().map(|matrix| matrix.width()).sum(); + let mut bytes = Vec::with_capacity(height * total_width * 8); + for row in 0..height { + for matrix in &matrices { + bytes.extend( + matrix + .row(row) + .unwrap() + .into_iter() + .flat_map(|value| value.as_canonical_u64().to_le_bytes()), + ); + } + } + packed_levels.push((bytes, total_width * 8, height)); + } + let level_refs: Vec<_> = packed_levels + .iter() + .map(|(bytes, row_bytes, height)| (bytes.as_slice(), *row_bytes, *height)) + .collect(); + let host_tree = CudaMixedMerkleTree::new(0, &level_refs); + let direct_host_tree = CudaMixedMerkleTree::from_host_matrices(0, &host_ldes); + let resident_tree = CudaMixedMerkleTree::from_ldes(0, &ldes); + assert_eq!(resident_tree.root(), host_tree.root()); + assert_eq!(direct_host_tree.root(), host_tree.root()); + assert_eq!(resident_tree.open_siblings(5), host_tree.open_siblings(5)); + assert_eq!( + direct_host_tree.open_siblings(5), + host_tree.open_siblings(5) + ); + for (resident, host) in ldes.iter().zip(&host_ldes) { + let rows = [0, resident.height() / 2, resident.height() - 1]; + for row in rows { + assert_eq!( + resident.row(row), + host.row(row).unwrap().into_iter().collect::>() + ); + } + assert_eq!( + resident.rows(&rows), + rows.map(|row| host.row(row).unwrap().into_iter().collect::>()) + ); + } + } + + unsafe extern "C" { + fn multi_stark_cuda_goldilocks_ops( + device_id: i32, + sums: *mut u64, + differences: *mut u64, + products: *mut u64, + inverses: *mut u64, + left: *const u64, + right: *const u64, + len: usize, + ) -> i32; + + } +} +#[test] +fn ext2_schoolbook_matches_p3() { + use p3_field::BasedVectorSpace; + use p3_field::extension::BinomialExtensionField; + type E = BinomialExtensionField; + let a = [ + Goldilocks::from_u64(16806794362951923611), + Goldilocks::from_u64(17407403654981854636), + ]; + let b = [ + Goldilocks::from_u64(15157121085234511159), + Goldilocks::from_u64(9277427398096025811), + ]; + let mut got = a; + let p0 = got[0] * b[0]; + let p1 = got[1] * b[1]; + got = [ + p0 + Goldilocks::from_u64(7) * p1, + got[0] * b[1] + got[1] * b[0], + ]; + let expected = E::from_basis_coefficients_slice(&a).unwrap() + * E::from_basis_coefficients_slice(&b).unwrap(); + assert_eq!(got.as_slice(), expected.as_basis_coefficients_slice()); +} diff --git a/src/cuda/pcs.rs b/src/cuda/pcs.rs new file mode 100644 index 0000000..99f3c90 --- /dev/null +++ b/src/cuda/pcs.rs @@ -0,0 +1,1360 @@ +//! The FRI PCS protocol over two-adic fields. +//! +//! The following implements a slight variant of the usual FRI protocol. As usual we start +//! with a polynomial `F(x)` of degree `n` given as evaluations over the coset `gH` with `|H| = 2^n`. +//! +//! Now consider the polynomial `G(x) = F(gx)`. Note that `G(x)` has the same degree as `F(x)` and +//! the evaluations of `F(x)` over `gH` are identical to the evaluations of `G(x)` over `H`. +//! +//! Hence we can reinterpret our vector of evaluations as evaluations of `G(x)` over `H` and apply +//! the standard FRI protocol to this evaluation vector. This makes it easier to apply FRI to a collection +//! of polynomials defined over different cosets as we don't need to keep track of the coset shifts. We +//! can just assume that every polynomial is defined over the subgroup of the relevant size. +//! +//! If we changed our domain construction (e.g., using multiple cosets), we would need to carefully reconsider these assumptions. +//! +//! The CPU PCS control flow in this module is derived from Plonky3's +//! `two_adic_pcs.rs` at revision e9d75614dd6816f9b5dbb4413c69be63536efd64 +//! (MIT/Apache-2.0). CUDA-resident commitments, openings, and FRI are maintained +//! here so their transcript order can be reviewed directly against that source. + +use core::fmt::Debug; +use core::iter; +use core::marker::PhantomData; +use core::mem::size_of; +use std::borrow::Cow; +use std::vec; +use std::vec::Vec; + +use itertools::{Itertools, izip}; +use p3_challenger::{CanObserve, FieldChallenger, GrindingChallenger}; +use p3_commit::{ + BatchOpening, BuildPeriodicLdeTableFast, ExtensionMmcs, Mmcs, OpenedValues, Pcs, + PeriodicLdeTable, +}; +use p3_dft::{Radix2DFTSmallBatch, TwoAdicSubgroupDft}; +use p3_field::coset::TwoAdicMultiplicativeCoset; +use p3_field::{ + BasedVectorSpace, ExtensionField, PackedFieldExtension, PrimeCharacteristicRing, PrimeField64, + TwoAdicField, batch_multiplicative_inverse, dot_product, +}; +use p3_interpolation::interpolate_coset_with_precomputation; +use p3_matrix::Matrix; +use p3_matrix::bitrev::{BitReversedMatrixView, BitReversibleMatrix}; +use p3_matrix::dense::{RowMajorMatrix, RowMajorMatrixCow}; +use p3_maybe_rayon::prelude::*; +use p3_util::linear_map::LinearMap; +use p3_util::{log2_strict_usize, reverse_bits_len, reverse_slice_index_bits}; +use tracing::{debug_span, instrument}; + +use super::mmcs::{CudaBatchOpenMmcs, CudaCommitMmcs}; +use super::{CudaFriWorkspace, CudaLde, CudaMixedMerkleTree, CudaReducedOpening}; +use p3_goldilocks::Goldilocks; +use p3_symmetric::MerkleCap; + +pub trait CudaPcsDft: TwoAdicSubgroupDft { + fn prepare_coset_lde_constants(&self, height: usize, added_bits: usize, shift: T); + + fn coset_lde_batch_resident( + &self, + matrix: &RowMajorMatrix, + added_bits: usize, + shift: T, + ) -> CudaLde; +} + +use p3_fri::{ + CommitPhaseProofStep, FriFoldingStrategy, FriParameters, FriProof, QueryProof, + build_periodic_lde_table_two_adic, compute_log_arity_for_round, prover, + verifier::{self, FriError}, +}; + +struct CudaFriRound { + codeword: CudaLde, + tree: CudaMixedMerkleTree, + arity: usize, +} + +trait CudaFriMmcs: Mmcs { + fn commit_cuda_fri( + &self, + codeword: CudaLde, + log_arity: usize, + ) -> (Self::Commitment, CudaFriRound); + fn open_cuda_fri_batch( + &self, + round: &CudaFriRound, + rows: &[usize], + ) -> Vec<(Vec, Self::Proof)>; +} + +impl CudaFriMmcs + for ExtensionMmcs +where + Challenge: ExtensionField, +{ + fn commit_cuda_fri( + &self, + codeword: CudaLde, + log_arity: usize, + ) -> (Self::Commitment, CudaFriRound) { + let arity = 1 << log_arity; + let tree = CudaMixedMerkleTree::from_fri_codeword(&codeword, arity); + let commitment = MerkleCap::new(vec![tree.root()]); + ( + commitment, + CudaFriRound { + codeword, + tree, + arity, + }, + ) + } + + fn open_cuda_fri_batch( + &self, + round: &CudaFriRound, + rows: &[usize], + ) -> Vec<(Vec, Self::Proof)> { + let codeword_rows = rows + .iter() + .flat_map(|&row| (0..round.arity).map(move |column| row * round.arity + column)) + .collect_vec(); + let opened = round.codeword.rows(&codeword_rows); + let paths = round.tree.open_siblings_batch(rows); + opened + .chunks_exact(round.arity) + .zip(paths) + .map(|(query_rows, opening_proof)| { + let values = query_rows + .iter() + .map(|pair| { + Challenge::from_basis_coefficients_slice(pair) + .expect("quadratic extension row") + }) + .collect(); + (values, opening_proof) + }) + .collect() + } +} + +/// A polynomial commitment scheme using FRI to generate opening proofs. +/// +/// We commit to a polynomial `f` via its evaluation vectors over a coset +/// `gH` where `|H| >= 2 * deg(f)`. A value `f(z)` is opened by using a FRI +/// proof to show that the evaluations of `(f(x) - f(z))/(x - z)` over +/// `gH` are low degree. +#[derive(Clone, Debug)] +pub struct CudaTwoAdicFriPcs { + pub(crate) dft: Dft, + pub(crate) mmcs: InputMmcs, + pub(crate) fri: FriParameters, + _phantom: PhantomData, +} + +struct CudaFriFolding(PhantomData<(InputProof, InputError)>); +type CudaFriFoldingForMmcs = CudaFriFolding>, >::Error>; +impl FriFoldingStrategy + for CudaFriFolding +where + Val: TwoAdicField + PrimeField64, + Challenge: ExtensionField, +{ + type InputProof = InputProof; + type InputError = InputError; + fn extra_query_index_bits(&self) -> usize { + 0 + } + fn fold_row( + &self, + index: usize, + log_height: usize, + log_arity: usize, + beta: Challenge, + evals: impl Iterator, + ) -> Challenge { + TwoAdicFriFolding::(PhantomData) + .fold_row(index, log_height, log_arity, beta, evals) + } + fn fold_matrix>( + &self, + beta: Challenge, + log_arity: usize, + m: M, + ) -> Vec { + TwoAdicFriFolding::(PhantomData).fold_matrix(beta, log_arity, m) + } +} + +fn prove_fri_cuda_resident( + params: &FriParameters, + mut inputs: Vec, + challenger: &mut Challenger, + log_global_max_height: usize, + prover_data_with_opening_points: &[ProverDataWithOpeningPoints< + '_, + Challenge, + InputMmcs::ProverData>, + >], + input_mmcs: &InputMmcs, + ext_w: Goldilocks, +) -> FriProof>> +where + Val: TwoAdicField + PrimeField64, + Challenge: ExtensionField, + InputMmcs: Mmcs + CudaBatchOpenMmcs, + FriMmcs: CudaFriMmcs, + Challenger: FieldChallenger + GrindingChallenger + CanObserve, +{ + assert!(!inputs.is_empty()); + assert!( + inputs + .windows(2) + .all(|pair| pair[0].height() > pair[1].height()) + ); + let to_pair = |value: Challenge| { + let coefficients = value.as_basis_coefficients_slice(); + [ + Goldilocks::from_u64(coefficients[0].as_canonical_u64()), + Goldilocks::from_u64(coefficients[1].as_canonical_u64()), + ] + }; + let from_pair = |pair: [Goldilocks; 2]| { + Challenge::from_basis_coefficients_slice(&[ + Val::from_u64(pair[0].as_canonical_u64()), + Val::from_u64(pair[1].as_canonical_u64()), + ]) + .expect("quadratic extension element") + }; + + let mut codeword = inputs.remove(0).to_lde(); + let mut commits = Vec::new(); + let mut rounds = Vec::new(); + let mut log_arities = Vec::new(); + let mut commit_pow_witnesses = Vec::new(); + let final_height = params.blowup() * params.final_poly_len(); + + while codeword.height() > final_height { + let log_height = log2_strict_usize(codeword.height()); + let next_log_height = inputs + .first() + .map(|input| log2_strict_usize(input.height())); + let log_arity = compute_log_arity_for_round( + log_height, + next_log_height, + params.log_blowup + params.log_final_poly_len, + params.max_log_arity, + ); + let arity = 1 << log_arity; + let (commitment, round) = params.mmcs.commit_cuda_fri(codeword, log_arity); + challenger.observe(commitment.clone()); + commits.push(commitment); + commit_pow_witnesses.push(challenger.grind(params.commit_proof_of_work_bits)); + let beta: Challenge = challenger.sample_algebra_element(); + let mut beta_step = beta; + let betas = (0..log_arity) + .map(|_| { + let result = to_pair(beta_step); + beta_step = beta_step.square(); + result + }) + .collect_vec(); + let folded_height = round.codeword.height() >> log_arity; + let add_next = inputs + .first() + .is_some_and(|input| input.height() == folded_height); + let beta_power = to_pair(beta.exp_power_of_2(log_arity)); + let g_inv = Goldilocks::from_u64( + Val::two_adic_generator(log_height) + .inverse() + .as_canonical_u64(), + ); + codeword = round.codeword.fri_fold( + add_next.then(|| &inputs[0]), + &betas, + beta_power, + g_inv, + ext_w, + ); + if add_next { + inputs.remove(0); + } + rounds.push(round); + log_arities.push(log_arity); + debug_assert_eq!(arity, 1 << log_arity); + } + assert!(inputs.is_empty()); + + let mut final_values = codeword + .to_row_major_matrix() + .values + .chunks_exact(2) + .take(params.final_poly_len()) + .map(|pair| from_pair([pair[0], pair[1]])) + .collect_vec(); + reverse_slice_index_bits(&mut final_values); + let final_poly = Radix2DFTSmallBatch::default().idft_algebra(final_values); + challenger.observe_algebra_slice(&final_poly); + + for &log_arity in &log_arities { + challenger.observe(Val::from_usize(log_arity)); + } + let query_pow_witness = challenger.grind(params.query_proof_of_work_bits); + let query_indices = iter::repeat_with(|| { + challenger.sample_bits(log2_strict_usize(rounds[0].codeword.height())) + }) + .take(params.num_queries) + .collect_vec(); + let input_batches = prover_data_with_opening_points + .iter() + .map(|(data, _)| { + let log_height = log2_strict_usize(input_mmcs.get_max_height(data)); + let indices = query_indices + .iter() + .map(|&index| index >> (log_global_max_height - log_height)) + .collect_vec(); + input_mmcs.open_batches(&indices, data) + }) + .collect_vec(); + let mut current_indices = query_indices; + let commit_batches = rounds + .iter() + .zip(&log_arities) + .map(|(round, &log_arity)| { + let arity = 1 << log_arity; + let positions = current_indices + .iter() + .map(|&index| index % arity) + .collect_vec(); + let group_indices = current_indices + .iter() + .map(|&index| index >> log_arity) + .collect_vec(); + let openings = params.mmcs.open_cuda_fri_batch(round, &group_indices); + current_indices = group_indices; + positions + .into_iter() + .zip(openings) + .map(|(index_in_group, (opened, opening_proof))| { + let sibling_values = opened + .into_iter() + .enumerate() + .filter_map(|(column, value)| (column != index_in_group).then_some(value)) + .collect(); + CommitPhaseProofStep { + log_arity: u8::try_from(log_arity).expect("FRI arity exceeds u8"), + sibling_values, + opening_proof, + } + }) + .collect_vec() + }) + .collect_vec(); + let mut input_iters = input_batches.into_iter().map(Vec::into_iter).collect_vec(); + let mut commit_iters = commit_batches.into_iter().map(Vec::into_iter).collect_vec(); + let query_proofs = (0..params.num_queries) + .map(|_| QueryProof { + input_proof: input_iters + .iter_mut() + .map(|openings| openings.next().expect("complete input opening batch")) + .collect(), + commit_phase_openings: commit_iters + .iter_mut() + .map(|openings| openings.next().expect("complete FRI opening batch")) + .collect(), + }) + .collect(); + + FriProof { + commit_phase_commits: commits, + commit_pow_witnesses, + query_proofs, + final_poly, + query_pow_witness, + } +} + +impl CudaTwoAdicFriPcs { + pub const fn new(dft: Dft, mmcs: InputMmcs, fri: FriParameters) -> Self { + Self { + dft, + mmcs, + fri, + _phantom: PhantomData, + } + } +} + +/// The Prover Data associated to a commitment to a collection of matrices +/// and a list of points to open each matrix at. +pub type ProverDataWithOpeningPoints<'a, EF, ProverData> = ( + // The matrices and auxiliary prover data + &'a ProverData, + // for each matrix, + Vec< + // points to open + Vec, + >, +); + +/// A joint commitment to a collection of matrices and their opening at +/// a collection of points. +pub type CommitmentWithOpeningPoints = ( + Commitment, + // For each matrix in the commitment: + Vec<( + // The domain of the matrix + Domain, + // A vector of (point, claimed_evaluation) pairs + Vec<(Challenge, Vec)>, + )>, +); + +pub struct TwoAdicFriFolding(pub PhantomData<(InputProof, InputError)>); + +pub type TwoAdicFriFoldingForMmcs = + TwoAdicFriFolding>, >::Error>; + +impl> + FriFoldingStrategy for TwoAdicFriFolding +{ + type InputProof = InputProof; + type InputError = InputError; + + fn extra_query_index_bits(&self) -> usize { + 0 + } + + fn fold_row( + &self, + index: usize, + log_height: usize, + log_arity: usize, + beta: EF, + evals: impl Iterator, + ) -> EF { + let arity = 1 << log_arity; + let evals: Vec<_> = evals.collect(); + assert_eq!(evals.len(), arity, "Expected {} evaluations", arity); + + // Compute the evaluation points in the subgroup + let subgroup_start = F::two_adic_generator(log_height + log_arity) + .exp_u64(reverse_bits_len(index, log_height) as u64); + let mut xs: Vec = F::two_adic_generator(log_arity) + .shifted_powers(subgroup_start) + .take(arity) + .collect(); + reverse_slice_index_bits(&mut xs); + + // Lagrange interpolation at beta + lagrange_interpolate_at(&xs, &evals, beta) + } + + #[instrument(skip_all)] + fn fold_matrix>(&self, beta: EF, log_arity: usize, m: M) -> Vec { + if log_arity == 1 { + // Optimized path for arity 2 + // We use the fact that + // p_e(x^2) = (p(x) + p(-x)) / 2 + // p_o(x^2) = (p(x) - p(-x)) / (2 x) + // that is, + // p_e(g^(2i)) = (p(g^i) + p(g^(n/2 + i))) / 2 + // p_o(g^(2i)) = (p(g^i) - p(g^(n/2 + i))) / (2 g^i) + // so + // result(g^(2i)) = p_e(g^(2i)) + beta p_o(g^(2i)) + // + // As p_e, p_o will be in the extension field we want to find ways to avoid extension multiplications. + // We should only need a single one (namely multiplication by beta). + let g_inv = F::two_adic_generator(log2_strict_usize(m.height()) + 1).inverse(); + + // As beta is in the extension field, we want to avoid multiplying by it + // for as long as possible. Here we precompute the powers `g_inv^i / 2` in the base field. + let mut halve_inv_powers = g_inv.shifted_powers(F::ONE.halve()).collect_n(m.height()); + reverse_slice_index_bits(&mut halve_inv_powers); + + m.par_rows() + .zip(halve_inv_powers) + .map(|(mut row, halve_inv_power)| { + let (lo, hi) = row.next_tuple().unwrap(); + (lo + hi).halve() + (lo - hi) * beta * halve_inv_power + }) + .collect() + } else { + // Decompose arity-2^k fold into k sequential arity-2 folds. + // This way, an arity-2^k fold with a single challenge beta is equivalent to + // k arity-2 folds with challenges beta, beta^2, beta^4, ..., beta^{2^{k-1}}. + // + // For arity 4 with evaluation points {s, -s, si, -si}: + // Step 1 (beta): fold pairs → g(s^2), g(-s^2) where g = f_e + beta*f_o + // Step 2 (beta^2): fold pair → g(beta^2) = f(beta) + + let mut data = m.to_row_major_matrix().values; + + let initial_height = data.len() / 2; + let g_inv = F::two_adic_generator(log2_strict_usize(initial_height) + 1).inverse(); + let mut halve_inv_powers = g_inv + .shifted_powers(F::ONE.halve()) + .collect_n(initial_height); + reverse_slice_index_bits(&mut halve_inv_powers); + + let two = F::ONE + F::ONE; + let mut current_beta = beta; + let mut next_data = EF::zero_vec(initial_height); + + for step in 0..log_arity { + let current_len = data.len(); + let height = current_len / 2; + // Since j << 1 is always >= j, we never overwrite data we haven't read yet. + if step > 0 { + for j in 0..height { + halve_inv_powers[j] = two * halve_inv_powers[j << 1].square(); + } + } + next_data[..height] + .par_iter_mut() + .zip(data.par_chunks_exact(2)) + .zip(&halve_inv_powers[..height]) + .for_each(|((out, chunk), &halve_inv_power)| { + // chunk is guaranteed to be size 2 by par_chunks_exact + let lo = chunk[0]; + let hi = chunk[1]; + + *out = (lo + hi).halve() + (lo - hi) * current_beta * halve_inv_power; + }); + current_beta = current_beta.square(); + + // Swap buffers conceptually (just truncate data and copy back, or ping-pong). + data.truncate(height); + data.copy_from_slice(&next_data[..height]); + } + + data + } + } +} + +/// Lagrange interpolation: given points (xs[i], ys[i]), evaluate at z. +/// +/// Uses the barycentric formula for efficiency when xs are roots of unity. +fn lagrange_interpolate_at>( + xs: &[F], + ys: &[EF], + z: EF, +) -> EF { + debug_assert_eq!(xs.len(), ys.len()); + let n = xs.len(); + + if n == 0 { + return EF::ZERO; + } + + // If z equals one of the interpolation points, return early. + for i in 0..n { + if (z - xs[i]).is_zero() { + return ys[i]; + } + } + + let log_n = log2_strict_usize(n); + + // All xs lie in a coset of the 2^log_n roots of unity. + let coset_power = xs[0].exp_power_of_2(log_n); + let weight_scale = (F::from_usize(n) * coset_power).inverse(); + + // Compute (z - x_i)^{-1} as a batch inversion + let diffs: Vec<_> = xs.iter().map(|&x| z - x).collect(); + let diff_invs = batch_multiplicative_inverse(&diffs); + + // Compute L(z) = prod_i (z - x_i) + let l_z = diffs.iter().copied().product::(); + + // Barycentric formula: sum_i (w_i * y_i / (z - x_i)) + // where w_i = 1 / prod_{j != i} (x_i - x_j) = x_i * weight_scale. + let mut result = EF::ZERO; + for ((&x, &y), &diff_inv) in xs.iter().zip(ys).zip(diff_invs.iter()) { + let weight = x * weight_scale; + result += y * weight * diff_inv; + } + result * l_z +} + +impl Pcs + for CudaTwoAdicFriPcs +where + Val: TwoAdicField + PrimeField64, + Dft: TwoAdicSubgroupDft + CudaPcsDft + Sync, + InputMmcs: Mmcs + CudaCommitMmcs + CudaBatchOpenMmcs, + FriMmcs: Mmcs + CudaFriMmcs, + Challenge: ExtensionField, + Challenger: + FieldChallenger + CanObserve + GrindingChallenger, +{ + type Domain = TwoAdicMultiplicativeCoset; + type Commitment = InputMmcs::Commitment; + type ProverData = InputMmcs::ProverData>; + type EvaluationsOnDomain<'a> = BitReversedMatrixView>; + type Proof = FriProof>>; + type Error = FriError; + const ZK: bool = false; + + /// Get the unique subgroup `H` of size `|H| = degree`. + /// + /// # Panics: + /// This function will panic if `degree` is not a power of 2 or `degree > (1 << Val::TWO_ADICITY)`. + fn natural_domain_for_degree(&self, degree: usize) -> Self::Domain { + TwoAdicMultiplicativeCoset::new(Val::ONE, log2_strict_usize(degree)).unwrap() + } + + /// Commit to a collection of evaluation matrices. + /// + /// Each element of `evaluations` contains a coset `shift * H` and a matrix `mat` with `mat.height() = |H|`. + /// Interpreting each column of `mat` as the evaluations of a polynomial `p_i(x)` over `shift * H`, + /// this computes the evaluations of `p_i` over `gK` where `g` is the chosen generator of the multiplicative group + /// of `Val` and `K` is the unique subgroup of order `|H| << self.fri.log_blowup`. + /// + /// This then outputs a Merkle commitment to these evaluations. + fn commit( + &self, + evaluations: impl IntoIterator)>, + ) -> (Self::Commitment, Self::ProverData) { + let evaluations: Vec<_> = evaluations.into_iter().collect(); + let source_cells = evaluations + .iter() + .map(|(_, matrix)| matrix.height() * matrix.width()) + .sum::(); + let dft = &self.dft; + let log_blowup = self.fri.log_blowup; + if source_cells >= 10_000_000 { + for (domain, evals) in &evaluations { + dft.prepare_coset_lde_constants( + evals.height(), + log_blowup, + Val::GENERATOR / domain.shift(), + ); + } + } + // Each Rayon worker uses its own per-thread CUDA stream. A bounded + // admission window keeps pageable upload staging within the driver's + // reliable concurrency range while retaining substantial overlap. + const CUDA_LDE_WAVE: usize = 16; + let mut ldes = Vec::with_capacity(evaluations.len()); + for wave in evaluations.chunks(CUDA_LDE_WAVE) { + let transform = + |(domain, evals): &(TwoAdicMultiplicativeCoset, RowMajorMatrix)| { + assert_eq!(domain.size(), evals.height()); + let shift = Val::GENERATOR / domain.shift(); + dft.coset_lde_batch_resident(evals, log_blowup, shift) + }; + // Ix's inner proof contains hundreds of tiny matrices. Entering + // CUDA concurrently for that low-volume batch costs more than it + // saves and can exhaust driver-side pageable-copy workers. + let wave_ldes: Vec<_> = if source_cells < 10_000_000 { + wave.iter().map(transform).collect() + } else { + wave.par_iter().map(transform).collect() + }; + ldes.extend(wave_ldes); + } + // Preserve enough device headroom for lookup, quotient, and FRI + // allocations. Spill the largest retained traces first, deriving the + // policy from this device rather than a 96-GiB development machine. + let (free_bytes, total_bytes) = crate::cuda::device_memory_info(self.mmcs.cuda_device_id()); + let minimum_free = std::env::var("MULTI_STARK_CUDA_MIN_FREE_BYTES") + .ok() + .and_then(|value| value.parse().ok()) + .unwrap_or(total_bytes / 4); + // A free-memory snapshot immediately after commitment construction is + // not enough: lookup construction temporarily needs buffers + // proportional to the original trace. Proactively spill large traces + // once this commitment itself is a meaningful fraction of VRAM. The + // ratios retain the policy that was validated on a 96-GiB device while + // scaling it to smaller cards. + let source_bytes = source_cells.saturating_mul(size_of::()); + let proactive_spill = source_bytes >= total_bytes / 12; + let large_matrix = total_bytes / 120; + let mut projected_free = free_bytes; + let mut by_size = evaluations + .iter() + .enumerate() + .map(|(index, (_, matrix))| { + ( + index, + matrix + .height() + .saturating_mul(matrix.width()) + .saturating_mul(size_of::()), + ) + }) + .collect_vec(); + by_size.sort_unstable_by_key(|&(_, bytes)| core::cmp::Reverse(bytes)); + let mut spilled = false; + for (index, bytes) in by_size { + let needs_headroom = projected_free < minimum_free; + let crowds_later_stages = proactive_spill && bytes >= large_matrix; + if !needs_headroom && !crowds_later_stages { + break; + } + // SAFETY: construction and commitment are synchronous; no CUDA + // operation can observe this trace while it is released. + unsafe { ldes[index].release_trace() }; + projected_free = projected_free.saturating_add(bytes); + spilled = true; + } + // Commit to the bit-reversed LDEs. + let (commitment, mut data) = self.mmcs.commit_cuda_resident(ldes); + if spilled { + self.mmcs.retain_matrices( + &mut data, + evaluations.into_iter().map(|(_, matrix)| matrix).collect(), + ); + } + (commitment, data) + } + + fn get_quotient_ldes( + &self, + evaluations: impl IntoIterator)>, + _num_chunks: usize, + ) -> Vec> { + evaluations + .into_iter() + .map(|(domain, evals)| { + assert_eq!(domain.size(), evals.height()); + // coset_lde_batch converts from evaluations over `xH` to evaluations over `shift * x * K`. + // Hence, letting `shift = g/x` the output will be evaluations over `gK` as desired. + // When `x = g`, we could just use the standard LDE but currently this doesn't seem + // to give a meaningful performance boost. + let shift = Val::GENERATOR / domain.shift(); + // Compute the LDE with blowup factor fri.log_blowup. + // We bit reverse as this is required by our implementation of the FRI protocol. + self.dft + .coset_lde_batch(evals, self.fri.log_blowup, shift) + .bit_reverse_rows() + .to_row_major_matrix() + }) + .collect() + } + + fn commit_ldes(&self, ldes: Vec>) -> (Self::Commitment, Self::ProverData) { + self.mmcs.commit_cuda_storage(ldes) + } + + /// Given the evaluations on a domain `gH`, return the evaluations on a different domain `g'K`. + /// + /// Arguments: + /// - `prover_data`: The prover data containing all committed evaluation matrices. + /// - `idx`: The index of the matrix containing the evaluations we want. These evaluations + /// are assumed to be over the coset `gH` where `g = Val::GENERATOR`. + /// - `domain`: The domain `g'K` on which to get evaluations on. + /// + /// When `g' = g` (i.e. `Val::GENERATOR`) and `K` is a subgroup of `H`, this is a simple + /// truncation of the bit-reversed LDE. Otherwise, we recover the polynomial coefficients + /// from the committed LDE and re-evaluate on the requested domain. + fn get_evaluations_on_domain<'a>( + &self, + prover_data: &'a Self::ProverData, + idx: usize, + domain: Self::Domain, + ) -> Self::EvaluationsOnDomain<'a> { + let lde = self.mmcs.get_matrices(prover_data)[idx]; + if domain.shift() == Val::GENERATOR && lde.height() >= domain.size() { + return lde.split_rows(domain.size()).0.as_cow().bit_reverse_rows(); + } + + // The committed LDE contains bit-reversed evaluations over `gH`. + // Un-bit-reverse, coset iDFT to recover coefficients, truncate to + // the original polynomial degree, then coset DFT onto the target domain. + let poly_height = lde.height() >> self.fri.log_blowup; + let lde_mat = lde.as_view().bit_reverse_rows().to_row_major_matrix(); + let mut coeffs = self.dft.coset_idft_batch(lde_mat, Val::GENERATOR); + let width = coeffs.width(); + coeffs.values.truncate(poly_height * width); + coeffs.values.resize(domain.size() * width, Val::ZERO); + let result = self + .dft + .coset_dft_batch(coeffs, domain.shift()) + .to_row_major_matrix(); + let result_width = result.width(); + + RowMajorMatrixCow::new(Cow::Owned(result.values), result_width).bit_reverse_rows() + } + + /// Open a batch of matrices at a collection of points. + /// + /// Returns the opened values along with a proof. + /// + /// This function assumes that all matrices correspond to evaluations over the + /// coset `gH` where `g = Val::GENERATOR` and `H` is a subgroup of appropriate size depending on the + /// matrix. + fn open( + &self, + // For each multi-matrix commitment, + commitment_data_with_opening_points: Vec<( + // The matrices and auxiliary prover data + &Self::ProverData, + // for each matrix, + Vec< + // points to open + Vec, + >, + )>, + challenger: &mut Challenger, + ) -> (OpenedValues, Self::Proof) { + /* + + A quick rundown of the optimizations in this function: + We are trying to compute sum_i alpha^i * (p(X) - y)/(X - z), + for each z an opening point, y = p(z). Each p(X) is given as evaluations in bit-reversed order + in the columns of the matrices. y is computed by barycentric interpolation. + X and p(X) are in the base field; alpha, y and z are in the extension. + The primary goal is to minimize extension multiplications. + + - Instead of computing all alpha^i, we just compute alpha^i for i up to the largest width + of a matrix, then multiply by an "alpha offset" when accumulating. + a^0 x0 + a^1 x1 + a^2 x2 + a^3 x3 + ... + = ( a^0 x0 + a^1 x1 ) + a^2 ( a^0 x2 + a^1 x3 ) + ... + (see `alpha_pows`, `alpha_pow_offset`, `num_reduced`) + + - For each unique point z, we precompute 1/(X-z) for the largest subgroup opened at this point. + Since we compute it in bit-reversed order, smaller subgroups can simply truncate the vector. + (see `inv_denoms`) + + - Then, for each matrix (with columns p_i) and opening point z, we want: + for each row (corresponding to subgroup element X): + reduced[X] += alpha_offset * sum_i [ alpha^i * inv_denom[X] * (p_i[X] - y[i]) ] + + We can factor out inv_denom, and expand what's left: + reduced[X] += alpha_offset * inv_denom[X] * sum_i [ alpha^i * p_i[X] - alpha^i * y[i] ] + + And separate the sum: + reduced[X] += alpha_offset * inv_denom[X] * [ sum_i [ alpha^i * p_i[X] ] - sum_i [ alpha^i * y[i] ] ] + + And now the last sum doesn't depend on X, so we can precompute that for the matrix, too. + So the hot loop (that depends on both X and i) is just: + sum_i [ alpha^i * p_i[X] ] + + with alpha^i an extension, p_i[X] a base + + */ + + // Keep CUDA commitments resident through barycentric interpolation + // and construction of the reduced FRI codewords. Only the opened + // values and final extension codewords cross back to the host. + let resident_rounds = debug_span!("cuda prepare resident rounds").in_scope(|| { + commitment_data_with_opening_points + .iter() + .map(|(data, points)| (self.mmcs.resident_or_upload(data), points)) + .collect_vec() + }); + let resident_max_height = resident_rounds + .iter() + .flat_map(|(ldes, _)| ldes.iter()) + .map(CudaLde::height) + .max() + .unwrap_or(0); + let final_fri_height = self.fri.blowup() * self.fri.final_poly_len(); + if resident_max_height > 1024 && resident_max_height > final_fri_height { + let _resident_guard = debug_span!("cuda resident fri").entered(); + let rounds = resident_rounds; + let device_id = self.mmcs.cuda_device_id(); + assert_eq!(>::DIMENSION, 2); + let to_gold = |v: Val| Goldilocks::from_u64(v.as_canonical_u64()); + let to_pair = |v: Challenge| { + let c = >::as_basis_coefficients_slice(&v); + [to_gold(c[0]), to_gold(c[1])] + }; + let from_pair = |v: [Goldilocks; 2]| { + Challenge::from_basis_coefficients_slice(&[ + Val::from_u64(v[0].as_canonical_u64()), + Val::from_u64(v[1].as_canonical_u64()), + ]) + .unwrap() + }; + let global_max_height = rounds + .iter() + .flat_map(|(ldes, _)| ldes.iter().map(CudaLde::height)) + .max() + .unwrap(); + let global_max_width = rounds + .iter() + .flat_map(|(ldes, _)| ldes.iter().map(CudaLde::width)) + .max() + .unwrap(); + let log_global_max_height = log2_strict_usize(global_max_height); + let coset_domain = + TwoAdicMultiplicativeCoset::new(Val::GENERATOR, log_global_max_height).unwrap(); + let mut coset: Vec = coset_domain.iter().collect(); + reverse_slice_index_bits(&mut coset); + let mut max_log: LinearMap = LinearMap::new(); + for (ldes, points) in &rounds { + for (lde, ps) in ldes.iter().zip(points.iter()) { + for &z in ps { + if let Some(h) = max_log.get_mut(&z) { + *h = (*h).max(log2_strict_usize(lde.height())) + } else { + max_log.insert(z, log2_strict_usize(lde.height())); + } + } + } + } + let ext_x = Challenge::from_basis_coefficients_slice(&[Val::ZERO, Val::ONE]).unwrap(); + let ext_w = to_pair(ext_x * ext_x)[0]; + assert_eq!(ext_w, Goldilocks::from_u64(7)); + let coset_gold: Vec<_> = coset.iter().copied().map(to_gold).collect(); + let mut inv_offsets = LinearMap::new(); + let mut inverse_points = Vec::new(); + let mut inverse_counts = Vec::new(); + let mut inverse_count = 0usize; + for (point, lh) in max_log { + inv_offsets.insert(point, inverse_count); + inverse_points.push(to_pair(point)); + let count = 1usize << lh; + inverse_counts.push(count); + inverse_count += count; + } + let mut workspace = CudaFriWorkspace::new( + device_id, + &inverse_points, + &inverse_counts, + &coset_gold, + ext_w, + ); + let mut interpolation_tasks = Vec::new(); + let mut output_count = 0usize; + let layouts = rounds + .iter() + .map(|(ldes, points)| { + ldes.iter() + .zip(points.iter()) + .map(|(lde, ps)| { + let h = lde.height() >> self.fri.log_blowup; + let lh = log2_strict_usize(h); + ps.iter() + .map(|&point| { + let offset = output_count; + output_count += lde.width(); + let shift_pow = Val::GENERATOR.exp_power_of_2(lh); + let scale = (point.exp_power_of_2(lh) - shift_pow) + * (Val::from_usize(h) * shift_pow).inverse(); + interpolation_tasks.push(lde.interpolation_task( + h, + *inv_offsets.get(&point).unwrap(), + offset, + to_pair(scale), + )); + (offset, lde.width()) + }) + .collect_vec() + }) + .collect_vec() + }) + .collect_vec(); + let interpolated = debug_span!("cuda interpolate openings") + .in_scope(|| workspace.interpolate(&interpolation_tasks, output_count, ext_w)); + let all_opened_values = layouts + .into_iter() + .map(|round| { + round + .into_iter() + .map(|matrix| { + matrix + .into_iter() + .map(|(offset, width)| { + interpolated[offset..offset + width] + .iter() + .copied() + .map(from_pair) + .collect_vec() + }) + .collect_vec() + }) + .collect_vec() + }) + .collect_vec(); + for round in &all_opened_values { + for matrix in round { + for values in matrix { + challenger.observe_algebra_slice(values); + } + } + } + let alpha: Challenge = challenger.sample_algebra_element(); + let alpha_powers: Vec<_> = alpha.powers().take(global_max_width).collect(); + let alpha_pairs: Vec<_> = alpha_powers.iter().copied().map(to_pair).collect(); + let mut num_reduced = [0usize; 32]; + let mut reduced: [Option; 32] = core::array::from_fn(|_| None); + let mut reduction_tasks = Vec::new(); + for ((ldes, points), openings_round) in rounds.iter().zip(all_opened_values.iter()) { + for ((lde, ps), openings) in + ldes.iter().zip(points.iter()).zip(openings_round.iter()) + { + let lh = log2_strict_usize(lde.height()); + let target = reduced[lh] + .get_or_insert_with(|| CudaReducedOpening::new(device_id, lde.height())); + for (&point, ys) in ps.iter().zip(openings.iter()) { + let reduced_y = + dot_product(alpha_powers.iter().copied(), ys.iter().copied()); + let offset = alpha.exp_u64(num_reduced[lh] as u64); + reduction_tasks.push(target.reduction_task( + lde, + *inv_offsets.get(&point).unwrap(), + to_pair(reduced_y), + to_pair(offset), + )); + num_reduced[lh] += lde.width(); + } + } + } + debug_span!("cuda reduce openings") + .in_scope(|| workspace.reduce(&reduction_tasks, &alpha_pairs, ext_w)); + let fri_input = reduced.into_iter().rev().flatten().collect_vec(); + let fri_proof = debug_span!("cuda prove fri").in_scope(|| { + prove_fri_cuda_resident( + &self.fri, + fri_input, + challenger, + log_global_max_height, + &commitment_data_with_opening_points, + &self.mmcs, + ext_w, + ) + }); + return (all_opened_values, fri_proof); + } + + // Contained in each `Self::ProverData` is a list of matrices which have been committed to. + // We extract those matrices to be able to refer to them directly. + let mats_and_points = commitment_data_with_opening_points + .iter() + .map(|(data, points)| { + let mats = self + .mmcs + .get_matrices(data) + .into_iter() + .map(|m| m.as_view()) + .collect_vec(); + debug_assert_eq!( + mats.len(), + points.len(), + "each matrix should have a corresponding set of evaluation points" + ); + (mats, points) + }) + .collect_vec(); + + // Find the maximum height and the maximum width of matrices in the batch. + // These do not need to correspond to the same matrix. + let (global_max_height, global_max_width) = mats_and_points + .iter() + .flat_map(|(mats, _)| mats.iter().map(|m| (m.height(), m.width()))) + .reduce(|(hmax, wmax), (h, w)| (hmax.max(h), wmax.max(w))) + .expect("No Matrices Supplied?"); + let log_global_max_height = log2_strict_usize(global_max_height); + + // Get all values of the coset `gH` for the largest necessary subgroup `H`. + // We also bit reverse which means that coset has the nice property that + // `coset[..2^i]` contains the values of `gK` for `|K| = 2^i`. + let coset = { + let coset = + TwoAdicMultiplicativeCoset::new(Val::GENERATOR, log_global_max_height).unwrap(); + let mut coset_points = coset.iter().collect(); + reverse_slice_index_bits(&mut coset_points); + coset_points + }; + + // For each unique opening point z, we will find the largest degree bound + // for that point, and precompute 1/(z - X) for the largest subgroup (in bitrev order). + let inv_denoms = compute_inverse_denominators(&mats_and_points, &coset); + + // Evaluate coset representations and write openings to the challenger + let all_opened_values = mats_and_points + .iter() + .map(|(mats, points)| { + // For each collection of matrices + izip!(mats.iter(), points.iter()) + .map(|(mat, points_for_mat)| { + // TODO: This assumes that every input matrix has a blowup of at least self.fri.log_blowup. + // If the blow_up factor is smaller than self.fri.log_blowup, this will lead to errors. + // If it is bigger, we shouldn't get any errors but it will be slightly slower. + // Ideally, polynomials could be passed in with their blow_up factors known. + + // The point of this correction is that each column of the matrix corresponds to a low degree polynomial. + // Hence we can save time by restricting the height of the matrix to be the minimal height which + // uniquely identifies the polynomial. + let h = mat.height() >> self.fri.log_blowup; + + // `subgroup` and `mat` are both in bit-reversed order, so we can truncate. + let (low_coset, _) = mat.split_rows(h); + let coset_h = &coset[..h]; + + points_for_mat + .iter() + .map(|&point| { + let _guard = + debug_span!("evaluate matrix", dims = %mat.dimensions()) + .entered(); + + // Use Barycentric interpolation to evaluate each column of the matrix at the given point. + let ys = debug_span!( + "compute opened values with Lagrange interpolation" + ) + .in_scope(|| { + // Get the relevant inverse denominators for this point and use these to + // interpolate to get the evaluation of each polynomial in the matrix + // at the desired point. + let inv_denoms = &inv_denoms.get(&point).unwrap()[..h]; + interpolate_coset_with_precomputation( + &low_coset, + Val::GENERATOR, + point, + coset_h, + inv_denoms, + ) + }); + + challenger.observe_algebra_slice(&ys); + ys + }) + .collect_vec() + }) + .collect_vec() + }) + .collect_vec(); + + // Batch combination challenge + + // Soundness Error: + // See the discussion in the doc comment of [`prove_fri`]. Essentially, the soundness error + // for this sample is tightly tied to the soundness error of the FRI protocol. + // Roughly speaking, at a minimum is it k/|EF| where `k` is the sum of, for each function, the number of + // points it needs to be opened at. This comes from the fact that we are taking a large linear combination + // of `(f(zeta) - f(x))/(zeta - x)` for each function `f` and all of `f`'s opening points. + // In our setup, k is two times the trace width plus the number of quotient polynomials. + let alpha: Challenge = challenger.sample_algebra_element(); + + // We precompute powers of alpha as we need the same powers for each matrix. + // We compute both a vector of unpacked powers and a vector of packed powers. + // TODO: It should be possible to refactor this to only use the packed powers but + // this is not a bottleneck so is not a priority. + let packed_alpha_powers = + Challenge::ExtensionPacking::packed_ext_powers_capped(alpha, global_max_width) + .collect_vec(); + let alpha_powers = + Challenge::ExtensionPacking::to_ext_iter(packed_alpha_powers.iter().copied()) + .collect_vec(); + + // Now that we have sent the openings to the verifier, it remains to prove + // that those openings are correct. + + // Given a low degree polynomial `f(x)` with claimed evaluation `f(zeta)`, we can check + // that `f(zeta)` is correct by doing a low degree test on `(f(zeta) - f(x))/(zeta - x)`. + // We will use `alpha` to batch together both different claimed openings `zeta` and + // different polynomials `f` whose evaluation vectors have the same height. + + // TODO: If we allow different polynomials to have different blow_up factors + // we may need to revisit this and to ensure it is safe to batch them together. + + // num_reduced records the number of (function, opening point) pairs for each `log_height`. + // TODO: This should really be `[0; Val::TWO_ADICITY]` but that runs into issues with generics. + let mut num_reduced = [0; 32]; + + // For each `log_height` from 2^1 -> 2^32, reduced_openings will contain either `None` + // if there are no matrices of that height, or `Some(vec)` where `vec` is equal to + // a weighted sum of `(f(zeta) - f(x))/(zeta - x)` over all `f`'s of that height and + // for each `f`, all opening points `zeta`. The sum is weighted by powers of the challenge alpha. + let mut reduced_openings: [_; 32] = core::array::from_fn(|_| None); + + for ((mats, points), openings_for_round) in + mats_and_points.iter().zip(all_opened_values.iter()) + { + for (mat, points_for_mat, openings_for_mat) in + izip!(mats.iter(), points.iter(), openings_for_round.iter()) + { + let _guard = + debug_span!("reduce matrix quotient", dims = %mat.dimensions()).entered(); + + let log_height = log2_strict_usize(mat.height()); + + // If this is our first matrix at this height, initialise reduced_openings to zero. + // Otherwise, get a mutable reference to it. + let reduced_opening_for_log_height = reduced_openings[log_height] + .get_or_insert_with(|| vec![Challenge::ZERO; mat.height()]); + debug_assert_eq!(reduced_opening_for_log_height.len(), mat.height()); + + // Treating our matrix M as the evaluations of functions f_0, f_1, ... + // Compute the evaluations of `Mred(x) = f_0(x) + alpha*f_1(x) + ...` + let mat_compressed = debug_span!("compress mat").in_scope(|| { + // This will be reused for all points z which M is opened at so we collect into a vector. + mat.rowwise_packed_dot_product::(&packed_alpha_powers) + .collect::>() + }); + + for (&point, openings) in points_for_mat.iter().zip(openings_for_mat) { + // If we have multiple matrices at the same height, we need to scale alpha to combine them. + // This means that reduced_openings will contain: + // Mred_0(x) + alpha^{M_0.width()}Mred_1(x) + alpha^{M_0.width() + M_1.width()}Mred_2(x) + ... + // Where M_0, M_1, ... are the matrices of the same height. + let alpha_pow_offset = alpha.exp_u64(num_reduced[log_height] as u64); + + // As we have all the openings `f_i(z)`, we can combine them using `alpha` + // in an identical way to before to compute `Mred(z)`. + let reduced_openings: Challenge = + dot_product(alpha_powers.iter().copied(), openings.iter().copied()); + + mat_compressed + .par_iter() + .zip(reduced_opening_for_log_height.par_iter_mut()) + // inv_denoms contains `1/(z - x)` for `x` in a coset `gK`. + // If `|K| =/= mat.height()` we actually want a subset of this + // corresponding to the evaluations over `gH` for `|H| = mat.height()`. + // As inv_denoms is bit reversed, the evaluations over `gH` are exactly + // the evaluations over `gK` at the indices `0..mat.height()`. + // So zip will truncate to the desired smaller length. + .zip(inv_denoms.get(&point).unwrap().par_iter()) + // Map the function `Mred(x) -> (Mred(z) - Mred(x))/(z - x)` + // across the evaluation vector of `Mred(x)`. Adjust by alpha_pow_offset + // as needed. + .for_each(|((&reduced_row, ro), &inv_denom)| { + *ro += alpha_pow_offset * (reduced_openings - reduced_row) * inv_denom; + }); + num_reduced[log_height] += mat.width(); + } + } + } + + // It remains to prove that all evaluation vectors in reduced_openings correspond to + // low degree functions. + let fri_input = reduced_openings.into_iter().rev().flatten().collect_vec(); + + let folding: CudaFriFoldingForMmcs = CudaFriFolding(PhantomData); + + // Produce the FRI proof. + let fri_proof = prover::prove_fri( + &folding, + &self.fri, + fri_input, + challenger, + log_global_max_height, + &commitment_data_with_opening_points, + &self.mmcs, + ); + + (all_opened_values, fri_proof) + } + + fn verify( + &self, + // For each commitment: + commitments_with_opening_points: Vec< + CommitmentWithOpeningPoints, + >, + proof: &Self::Proof, + challenger: &mut Challenger, + ) -> Result<(), Self::Error> { + // Write all evaluations to challenger. + // Need to ensure to do this in the same order as the prover. + for (_, round) in &commitments_with_opening_points { + for (_, mat) in round { + for (_, point) in mat { + challenger.observe_algebra_slice(point); + } + } + } + + let folding: TwoAdicFriFoldingForMmcs = TwoAdicFriFolding(PhantomData); + + verifier::verify_fri( + &folding, + &self.fri, + proof, + challenger, + &commitments_with_opening_points, + &self.mmcs, + )?; + + Ok(()) + } +} + +impl BuildPeriodicLdeTableFast + for CudaTwoAdicFriPcs +where + Val: TwoAdicField, + Dft: TwoAdicSubgroupDft + CudaPcsDft + Default, +{ + type PeriodicDomain = TwoAdicMultiplicativeCoset; + + fn maybe_build_periodic_lde_table_fast( + &self, + periodic_cols: &[Vec>], + trace_domain: Self::PeriodicDomain, + quotient_domain: Self::PeriodicDomain, + ) -> Option>> + where + p3_commit::Val: Clone, + { + let periodic_cols_val: &[Vec] = unsafe { core::mem::transmute(periodic_cols) }; + let table = build_periodic_lde_table_two_adic::( + periodic_cols_val, + &trace_domain, + "ient_domain, + ); + Some(table) + } +} + +/// Compute vectors of inverse denominators for each unique opening point. +/// +/// Arguments: +/// - `mats_and_points` is a list of matrices and for each matrix a list of points. We assume that +/// the total number of distinct points is very small as several methods contained herein are `O(n^2)` +/// in the number of points. +/// - `coset` is the set of points `gH` where `H` a two-adic subgroup such that `|H|` is greater +/// than or equal to the largest height of any matrix in `mats_and_points`. The values +/// in `coset` must be in bit-reversed order. +/// +/// For each point `z`, let `M` be the matrix of largest height which opens at `z`. +/// let `H_z` be the unique subgroup of order `M.height()`. Compute the vector of +/// `1/(z - x)` for `x` in `gH_z`. +/// +/// Return a LinearMap which allows us to recover the computed vectors for each `z`. +#[instrument(skip_all)] +fn compute_inverse_denominators, M: Matrix>( + mats_and_points: &[(Vec, &Vec>)], + coset: &[F], +) -> LinearMap> { + // For each `z`, find the maximal height of any matrix which we need to + // open at `z`. + let mut max_log_height_for_point: LinearMap = LinearMap::new(); + for (mats, points) in mats_and_points { + for (mat, points_for_mat) in izip!(mats, *points) { + let log_height = log2_strict_usize(mat.height()); + for &z in points_for_mat { + if let Some(lh) = max_log_height_for_point.get_mut(&z) { + *lh = core::cmp::max(*lh, log_height); + } else { + max_log_height_for_point.insert(z, log_height); + } + } + } + } + + // Compute the inverse denominators for each point `z`. + max_log_height_for_point + .into_iter() + .map(|(z, log_height)| { + ( + z, + batch_multiplicative_inverse( + // As coset is stored in bit-reversed order, + // we can just take the first `2^log_height` elements. + &coset[..(1 << log_height)] + .iter() + .map(|&x| z - x) + .collect_vec(), + ), + ) + }) + .collect() +} diff --git a/src/lib.rs b/src/lib.rs index 0df70f6..a3afd0a 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,4 +1,9 @@ pub mod config; +#[cfg(feature = "cuda")] +pub mod cuda; +#[cfg(feature = "cuda")] +#[doc(hidden)] +pub use cuda::pcs as cuda_pcs; pub mod eval; pub mod expr; pub mod graph; diff --git a/src/lookup.rs b/src/lookup.rs index cf94ad7..642fb62 100644 --- a/src/lookup.rs +++ b/src/lookup.rs @@ -702,6 +702,27 @@ impl LookupValues { } } +#[cfg(feature = "cuda")] +impl LookupValues { + pub(crate) fn cuda_parts( + &self, + ) -> ( + usize, + usize, + &[p3_goldilocks::Goldilocks], + &[p3_goldilocks::Goldilocks], + &[usize], + ) { + ( + self.height, + self.num_lookups, + &self.multiplicities, + &self.args, + &self.arg_offsets, + ) + } +} + /// Incremental, allocation-free constructor for [`LookupValues`]. /// /// Rows start zeroed — multiplicity zero and zero arguments in every slot — diff --git a/src/prover.rs b/src/prover.rs index 146a239..57a574c 100644 --- a/src/prover.rs +++ b/src/prover.rs @@ -190,7 +190,7 @@ use bincode::error::{DecodeError, EncodeError}; use bincode::serde::{decode_from_slice, encode_to_vec}; use p3_challenger::{CanObserve, FieldChallenger}; use p3_commit::{LagrangeSelectors, OpenedValuesForRound, Pcs, PolynomialSpace}; -use p3_dft::{Radix2DitParallel, TwoAdicSubgroupDft}; +use p3_dft::TwoAdicSubgroupDft; use p3_field::{ Algebra, BasedVectorSpace, Field, PackedValue, PrimeCharacteristicRing, TwoAdicField, }; @@ -200,7 +200,7 @@ use p3_util::{log2_strict_usize, reverse_bits_len}; use serde::{Deserialize, Serialize}; /// Polynomial commitments included in the proof. -#[derive(Serialize, Deserialize)] +#[derive(Clone, Serialize, Deserialize)] pub struct Commitments { /// Commitment to the stage 1 (main) execution traces. pub stage_1_trace: Com, @@ -238,6 +238,22 @@ pub struct Proof { pub stage_2_opened_values: OpenedValuesForRound, } +impl Clone for Proof { + fn clone(&self) -> Self { + Self { + active: self.active.clone(), + commitments: self.commitments.clone(), + intermediate_accumulators: self.intermediate_accumulators.clone(), + log_degrees: self.log_degrees.clone(), + opening_proof: self.opening_proof.clone(), + quotient_opened_values: self.quotient_opened_values.clone(), + preprocessed_opened_values: self.preprocessed_opened_values.clone(), + stage_1_opened_values: self.stage_1_opened_values.clone(), + stage_2_opened_values: self.stage_2_opened_values.clone(), + } + } +} + impl Proof { fn serde_config() -> Configuration { standard().with_little_endian().with_fixed_int_encoding() @@ -245,7 +261,9 @@ impl Proof { #[inline] pub fn to_bytes(&self) -> Result, EncodeError> { - encode_to_vec(self, Self::serde_config()) + let mut proof = self.clone(); + SC::canonicalize_proof(&mut proof); + encode_to_vec(&proof, Self::serde_config()) } #[inline] @@ -261,7 +279,8 @@ where // Two-adicity is needed to slice the quotient into coefficient slices // and rebuild their committed LDE from those coefficients; every // FRI-based config is two-adic anyway. - Val: TwoAdicField + Ord, + Val: TwoAdicField, + SC::Dft: TwoAdicSubgroupDft>, { /// Generates a STARK proof for the system with a single claim. /// @@ -402,28 +421,69 @@ where .iter() .map(|&ci| self.circuits[ci].lookup_group_size) .collect(); - let (stage_2_traces, intermediate_accumulators) = LookupValues::stage_2_traces( - &active_lookups, - &group_sizes, + let lookup_inputs: Vec<_> = active_indices + .iter() + .enumerate() + .zip(&active_lookups) + .map( + |((position, &circuit_index), lookup_values)| crate::config::LookupCommitInput { + circuit: &self.circuits[circuit_index], + lookup_values, + preprocessed: key + .preprocessed_data + .as_ref() + .zip(self.preprocessed_indices[circuit_index]), + stage_1: (&stage_1_trace_data, position), + }, + ) + .collect(); + let accelerated_lookup_commit = self.config.accelerated_lookup_commit( + &lookup_inputs, lookup_argument_challenge, - &fingerprint_challenge, + fingerprint_challenge, acc, ); - // The lookup witness can be as large as the traces themselves; free it - // now instead of holding it through the commit/quotient/FRI stages. - drop(active_lookups); - drop(_g); + let (stage_2_trace_commit, stage_2_trace_data, intermediate_accumulators) = + if let Some(result) = accelerated_lookup_commit { + drop(active_lookups); + drop(_g); + result + } else { + let (stage_2_traces, intermediate_accumulators) = self + .config + .accelerated_lookup_traces( + &active_lookups, + &group_sizes, + lookup_argument_challenge, + fingerprint_challenge, + acc, + ) + .unwrap_or_else(|| { + LookupValues::stage_2_traces( + &active_lookups, + &group_sizes, + lookup_argument_challenge, + &fingerprint_challenge, + acc, + ) + }); + // The lookup witness can be as large as the traces themselves; free it + // now instead of holding it through the commit/quotient/FRI stages. + drop(active_lookups); + drop(_g); - // Cost: "Stage 2 commit" — LDE + Merkle for flattened extension traces. - // FFT work: Σ w2_i · D · (B+1) · n_i · log₂(n_i). - let _g = tracing::info_span!("stark/stage2_commit").entered(); - let evaluations = stage_2_traces.into_iter().map(|trace| { - let degree = trace.height(); - let trace_domain = pcs.natural_domain_for_degree(degree); - (trace_domain, trace.flatten_to_base()) - }); - let (stage_2_trace_commit, stage_2_trace_data) = pcs.commit(evaluations); - drop(_g); + // Cost: "Stage 2 commit" — LDE + Merkle for flattened extension traces. + // FFT work: Σ w2_i · D · (B+1) · n_i · log₂(n_i). + let _g = tracing::info_span!("stark/stage2_commit").entered(); + let evaluations = stage_2_traces.into_iter().map(|trace| { + let degree = trace.height(); + let trace_domain = pcs.natural_domain_for_degree(degree); + (trace_domain, trace.flatten_to_base()) + }); + let (commitment, data) = pcs.commit(evaluations); + drop(_g); + (commitment, data, intermediate_accumulators) + }; challenger.observe(stage_2_trace_commit.clone()); // Observe the intermediate accumulators. They enter the constraints as @@ -443,93 +503,135 @@ where let _g = tracing::info_span!("stark/quotient").entered(); debug_assert_eq!(intermediate_accumulators.len(), active_indices.len()); debug_assert_eq!(log_degrees.len(), active_indices.len()); - let dft = Radix2DitParallel::>::default(); + let dft = self.config.dft(); let log_blowup = self.config.log_blowup(); - let quotient_ldes: Vec<_> = active_indices + let mut quotient_inputs = Vec::with_capacity(active_indices.len()); + for (pos, ((&ci, log_degree), next_acc)) in active_indices .iter() .zip(log_degrees.iter()) .zip(intermediate_accumulators.iter()) .enumerate() - .map(|(pos, ((&ci, log_degree), next_acc))| { - let circuit = &self.circuits[ci]; - let quotient_degree = circuit.quotient_degree(); - let log_quotient_degree = log2_strict_usize(quotient_degree); - let trace_domain = pcs.natural_domain_for_degree(1 << log_degree); - let quotient_domain = - trace_domain.create_disjoint_domain(1 << (log_degree + log_quotient_degree)); - let preprocessed_trace_on_quotient_domain = key + { + let circuit = &self.circuits[ci]; + let quotient_degree = circuit.quotient_degree(); + let log_quotient_degree = log2_strict_usize(quotient_degree); + let trace_domain = pcs.natural_domain_for_degree(1 << log_degree); + let quotient_domain = + trace_domain.create_disjoint_domain(1 << (log_degree + log_quotient_degree)); + // The four lookup publics (β, γ, acc, next_acc) as flat base + // coordinates, in the layout the synthesized constraints read. + let mut lookup_publics: Vec> = Vec::new(); + for ef in [ + lookup_argument_challenge, + fingerprint_challenge, + acc, + *next_acc, + ] { + lookup_publics.extend_from_slice(ef.as_basis_coefficients_slice()); + } + + quotient_inputs.push(crate::config::QuotientCommitInput { + circuit, + lookup_publics, + trace_domain, + quotient_domain, + preprocessed: key .preprocessed_data .as_ref() - .zip(self.preprocessed_indices[ci]) - .map(|(preprocessed_trace_data, preprocessed_idx)| { - pcs.get_evaluations_on_domain( - preprocessed_trace_data, - preprocessed_idx, - quotient_domain, - ) - }); - let stage_1_trace_on_quotient_domain = - pcs.get_evaluations_on_domain(&stage_1_trace_data, pos, quotient_domain); - let stage_2_trace_on_quotient_domain = - pcs.get_evaluations_on_domain(&stage_2_trace_data, pos, quotient_domain); - - // The four lookup publics (β, γ, acc, next_acc) as flat base - // coordinates, in the layout the synthesized constraints read. - let mut lookup_publics: Vec> = Vec::new(); - for ef in [ - lookup_argument_challenge, - fingerprint_challenge, - acc, - *next_acc, - ] { - lookup_publics.extend_from_slice(ef.as_basis_coefficients_slice()); - } + .zip(self.preprocessed_indices[ci]), + stage_1: (&stage_1_trace_data, pos), + stage_2: (&stage_2_trace_data, pos), + constraint_count: circuit.constraint_count(), + }); + acc = *next_acc; + } - // compute the quotient values which are elements of the extension field and flatten it to the base field - let quotient_values = quotient_values::( - circuit, - &lookup_publics, - trace_domain, - quotient_domain, - &preprocessed_trace_on_quotient_domain, - &stage_1_trace_on_quotient_domain, - &stage_2_trace_on_quotient_domain, - constraint_challenge, - circuit.constraint_count(), - ); - let quotient_flat = - RowMajorMatrix::new_col(quotient_values).flatten_to_base::>(); - // The quotient has degree greater than the trace polynomials, - // so for FRI to work it must be split into `quotient_degree` - // sub-polynomials of trace degree. Slice its COEFFICIENTS: - // `Q(X) = Σᵢ X^{i·n}·cᵢ(X)` with each `cᵢ` of degree < n, and - // commit all slices as ONE `q·D`-column matrix on the trace - // domain — instead of one matrix per slice on the split - // cosets — so the opening phase pays its per-matrix costs - // once per circuit rather than once per slice. The committed - // LDE is built straight from these coefficients: evaluating - // them onto the trace domain only for `Pcs::commit` to - // inverse-DFT that evaluation right back would waste two - // size-n transforms per column. The slicing itself is fused - // with the iDFT's scaling passes into one parallel gather - // (see `shifted_quotient_slices`). - acc = *next_acc; - let sliced = shifted_quotient_slices( - &dft, - quotient_flat, - quotient_domain.first_point(), - quotient_degree, - ); - lde_from_shifted_coefficients(&dft, sliced, log_blowup) - }) - .collect(); - // `commit_ldes` skips the randomization a hiding PCS applies inside - // `commit`; this prover targets non-hiding configurations only. + // `commit_ldes` and the fused accelerator both bypass the + // randomization a hiding PCS applies inside `commit`. assert!( !>::ZK, "committing the quotient from coefficients bypasses hiding-PCS randomization" ); - let (quotient_commit, quotient_data) = pcs.commit_ldes(quotient_ldes); + let (quotient_commit, quotient_data) = if let Some(committed) = self + .config + .accelerated_quotient_commit("ient_inputs, constraint_challenge) + { + committed + } else { + let quotient_ldes: Vec<_> = quotient_inputs + .into_iter() + .map(|input| { + let circuit = input.circuit; + let quotient_degree = circuit.quotient_degree(); + + // compute the quotient values which are elements of the extension field and flatten it to the base field + let quotient_values = self + .config + .accelerated_quotient_values( + circuit, + &input.lookup_publics, + input.trace_domain, + input.quotient_domain, + input.preprocessed, + input.stage_1, + input.stage_2, + constraint_challenge, + input.constraint_count, + ) + .unwrap_or_else(|| { + let preprocessed_trace_on_quotient_domain = + input.preprocessed.map(|(data, idx)| { + pcs.get_evaluations_on_domain(data, idx, input.quotient_domain) + }); + let stage_1_trace_on_quotient_domain = pcs.get_evaluations_on_domain( + input.stage_1.0, + input.stage_1.1, + input.quotient_domain, + ); + let stage_2_trace_on_quotient_domain = pcs.get_evaluations_on_domain( + input.stage_2.0, + input.stage_2.1, + input.quotient_domain, + ); + quotient_values::( + circuit, + &input.lookup_publics, + input.trace_domain, + input.quotient_domain, + &preprocessed_trace_on_quotient_domain, + &stage_1_trace_on_quotient_domain, + &stage_2_trace_on_quotient_domain, + constraint_challenge, + input.constraint_count, + ) + }); + let quotient_flat = + RowMajorMatrix::new_col(quotient_values).flatten_to_base::>(); + // The quotient has degree greater than the trace polynomials, + // so for FRI to work it must be split into `quotient_degree` + // sub-polynomials of trace degree. Slice its COEFFICIENTS: + // `Q(X) = Σᵢ X^{i·n}·cᵢ(X)` with each `cᵢ` of degree < n, and + // commit all slices as ONE `q·D`-column matrix on the trace + // domain — instead of one matrix per slice on the split + // cosets — so the opening phase pays its per-matrix costs + // once per circuit rather than once per slice. The committed + // LDE is built straight from these coefficients: evaluating + // them onto the trace domain only for `Pcs::commit` to + // inverse-DFT that evaluation right back would waste two + // size-n transforms per column. The slicing itself is fused + // with the iDFT's scaling passes into one parallel gather + // (see `shifted_quotient_slices`). + let sliced = shifted_quotient_slices( + dft, + quotient_flat, + input.quotient_domain.first_point(), + quotient_degree, + ); + lde_from_shifted_coefficients(dft, sliced, log_blowup) + }) + .collect(); + pcs.commit_ldes(quotient_ldes) + }; challenger.observe(quotient_commit.clone()); drop(_g); @@ -634,12 +736,16 @@ where /// natural trace domain guarantees it). What survives is a single parallel /// gather off the DFT storage with ONE constant weight per slice: /// `wₖ = N⁻¹ · GENERATOR^{−k·n}`. -fn shifted_quotient_slices( - dft: &Radix2DitParallel, +fn shifted_quotient_slices( + dft: &Dft, quotient_evals: RowMajorMatrix, domain_shift: F, quotient_degree: usize, -) -> RowMajorMatrix { +) -> RowMajorMatrix +where + F: TwoAdicField, + Dft: TwoAdicSubgroupDft, +{ assert_eq!( domain_shift, F::GENERATOR, @@ -651,9 +757,13 @@ fn shifted_quotient_slices( debug_assert_eq!(big_height % quotient_degree, 0); let n = big_height / quotient_degree; let width = quotient_degree * ext_degree; - // Raw storage of the forward DFT: natural index `k` lives at row - // `rev(k)`, and the unwrap out of the bit-reversed view is copy-free. - let storage = dft.dft_batch(quotient_evals).bit_reverse_rows(); + // Bit-reversed storage of the forward DFT: natural index `k` lives at + // row `rev(k)`. This conversion is copy-free for Radix2DitParallel's + // native output and materializes the same layout for other backends. + let storage = dft + .dft_batch(quotient_evals) + .bit_reverse_rows() + .to_row_major_matrix(); let n_inv = F::ONE.div_2exp_u64(log_big_height as u64); let weight_step = F::GENERATOR.exp_u64(n as u64).inverse(); let weights: Vec = weight_step @@ -712,14 +822,20 @@ fn shifted_quotient_slices( /// `coset_dft_batch`, no reassembly copies, and `Radix2DitParallel`'s /// native output order is already the bit-reversed storage order, so the /// final unwrap is copy-free. -fn lde_from_shifted_coefficients( - dft: &Radix2DitParallel, +fn lde_from_shifted_coefficients( + dft: &Dft, mut coefficients: RowMajorMatrix, log_blowup: usize, -) -> RowMajorMatrix { +) -> RowMajorMatrix +where + F: TwoAdicField, + Dft: TwoAdicSubgroupDft, +{ let height = coefficients.height(); coefficients.pad_to_height(height << log_blowup, F::ZERO); - dft.dft_batch(coefficients).bit_reverse_rows() + dft.dft_batch(coefficients) + .bit_reverse_rows() + .to_row_major_matrix() } /// Reference form of [`lde_from_shifted_coefficients`] taking PLAIN @@ -727,11 +843,15 @@ fn lde_from_shifted_coefficients( /// pinning tests need it; the prover gets the shift for free inside /// [`shifted_quotient_slices`]. #[cfg(test)] -fn lde_from_coefficients( - dft: &Radix2DitParallel, +fn lde_from_coefficients( + dft: &Dft, mut coefficients: RowMajorMatrix, log_blowup: usize, -) -> RowMajorMatrix { +) -> RowMajorMatrix +where + F: TwoAdicField, + Dft: TwoAdicSubgroupDft, +{ scale_rows_by_powers(&mut coefficients, F::GENERATOR); lde_from_shifted_coefficients(dft, coefficients, log_blowup) } @@ -772,7 +892,7 @@ fn quotient_values( ) -> Vec where SC: StarkGenericConfig, - Val: TwoAdicField + Ord, + Val: TwoAdicField, { let quotient_size = quotient_domain.size(); let main_width = circuit.main_width; @@ -889,7 +1009,7 @@ fn quotient_values_inner( ) -> impl Iterator where SC: StarkGenericConfig, - Val: TwoAdicField + Ord, + Val: TwoAdicField, { let i_range = i_start..i_start + PackedVal::::WIDTH; let is_first_row = *PackedVal::::from_slice(&sels.is_first_row[i_range.clone()]); @@ -972,8 +1092,48 @@ where mod tests { use super::*; use crate::types::Val; + use p3_dft::{NaiveDft, Radix2DitParallel}; use rand::{RngExt, SeedableRng, rngs::SmallRng}; + /// Host model of the CUDA radix-2 DIF kernel. It intentionally mirrors + /// `cuda/kernels.cu` stage/index/twiddle ordering and returns raw + /// bit-reversed storage. + fn cuda_dif_model(mut matrix: RowMajorMatrix, inverse: bool) -> RowMajorMatrix { + let height = matrix.height(); + if height <= 1 { + return matrix; + } + let width = matrix.width(); + let log_height = log2_strict_usize(height); + let root = Val::two_adic_generator(log_height); + let root = if inverse { root.inverse() } else { root }; + let twiddles: Vec<_> = root.powers().take(height / 2).collect(); + + let mut half = height / 2; + loop { + let stride = height / (2 * half); + for butterfly in 0..height / 2 { + let offset = butterfly % half; + let group = butterfly / half; + let row_0 = group * (2 * half) + offset; + let row_1 = row_0 + half; + for column in 0..width { + let index_0 = row_0 * width + column; + let index_1 = row_1 * width + column; + let left = matrix.values[index_0]; + let right = matrix.values[index_1]; + matrix.values[index_0] = left + right; + matrix.values[index_1] = (left - right) * twiddles[offset * stride]; + } + } + if half == 1 { + break; + } + half >>= 1; + } + matrix + } + /// `lde_from_coefficients` must reproduce, value for value, the matrix /// `TwoAdicFriPcs::commit` stores for the same polynomials given as /// trace-domain evaluations: `coset_lde_batch` with the generator shift @@ -1046,4 +1206,89 @@ mod tests { } } } + + /// The quotient path must depend only on `TwoAdicSubgroupDft` semantics, + /// not on `Radix2DitParallel`'s bit-reversed storage representation. This + /// is the contract a CUDA adapter will implement. + #[test] + fn quotient_transforms_are_dft_backend_independent() { + let mut rng = SmallRng::seed_from_u64(2); + let radix = Radix2DitParallel::::default(); + let naive = NaiveDft; + + for log_n in [0usize, 1, 3] { + for quotient_degree in [1usize, 2, 4] { + for ext_degree in [1usize, 2] { + let big_height = (1 << log_n) * quotient_degree; + let evals = RowMajorMatrix::new( + (0..big_height * ext_degree).map(|_| rng.random()).collect(), + ext_degree, + ); + let radix_slices = shifted_quotient_slices( + &radix, + evals.clone(), + Val::GENERATOR, + quotient_degree, + ); + let naive_slices = + shifted_quotient_slices(&naive, evals, Val::GENERATOR, quotient_degree); + assert_eq!(naive_slices, radix_slices); + + for log_blowup in [1usize, 2] { + let radix_lde = + lde_from_shifted_coefficients(&radix, radix_slices.clone(), log_blowup); + let naive_lde = + lde_from_shifted_coefficients(&naive, naive_slices.clone(), log_blowup); + assert_eq!(naive_lde, radix_lde); + } + } + } + } + } + + /// Pins the first-party CUDA kernel's DIF ordering and fused coset-LDE + /// pipeline even on machines without a CUDA toolkit or GPU. + #[test] + fn cuda_dif_and_coset_lde_model_match_cpu() { + let mut rng = SmallRng::seed_from_u64(3); + let radix = Radix2DitParallel::::default(); + for log_height in [0usize, 1, 2, 5, 8] { + for width in [1usize, 2, 7] { + let height = 1 << log_height; + let matrix = + RowMajorMatrix::new((0..height * width).map(|_| rng.random()).collect(), width); + let expected_dft = radix + .dft_batch(matrix.clone()) + .bit_reverse_rows() + .to_row_major_matrix(); + assert_eq!(cuda_dif_model(matrix.clone(), false), expected_dft); + + for added_bits in [0usize, 1, 2, 3] { + // CUDA pipeline: inverse DIF (bit-reversed coefficients), + // bit-reverse + normalize + shift, zero-pad, forward DIF. + let mut coefficients = cuda_dif_model(matrix.clone(), true) + .bit_reverse_rows() + .to_row_major_matrix(); + let height_inverse = Val::ONE.div_2exp_u64(log_height as u64); + let mut shift = Val::ONE; + for row in 0..height { + for value in &mut coefficients.values[row * width..(row + 1) * width] { + *value *= height_inverse * shift; + } + shift *= Val::GENERATOR; + } + coefficients.pad_to_height(height << added_bits, Val::ZERO); + let actual = cuda_dif_model(coefficients, false); + let expected = radix + .coset_lde_batch(matrix.clone(), added_bits, Val::GENERATOR) + .bit_reverse_rows() + .to_row_major_matrix(); + assert_eq!( + actual, expected, + "height=2^{log_height}, blowup=2^{added_bits}, width={width}" + ); + } + } + } + } } diff --git a/src/test_circuits/baby_bear_config.rs b/src/test_circuits/baby_bear_config.rs index ca35637..5a8f63f 100644 --- a/src/test_circuits/baby_bear_config.rs +++ b/src/test_circuits/baby_bear_config.rs @@ -39,6 +39,7 @@ type Pcs = TwoAdicFriPcs; struct BabyBearPoseidon2Config { pcs: Pcs, + dft: Dft, perm: Perm, /// Field elements observed into every fresh challenger: a domain tag /// plus a digest of the protocol parameters (see the transcript contract @@ -66,7 +67,8 @@ impl BabyBearPoseidon2Config { query_proof_of_work_bits: fri_parameters.query_proof_of_work_bits, mmcs: challenge_mmcs, }; - let pcs = Pcs::new(Dft::default(), val_mmcs, inner_parameters); + let dft = Dft::default(); + let pcs = Pcs::new(dft.clone(), val_mmcs, inner_parameters); let mut challenger_seed: Vec = b"multi-stark/v0" .iter() .map(|&byte| Val::from_u8(byte)) @@ -87,6 +89,7 @@ impl BabyBearPoseidon2Config { let max_quotient_degree = 1 << commitment_parameters.log_blowup; Self { pcs, + dft, perm, challenger_seed, max_log_degree, @@ -98,6 +101,7 @@ impl BabyBearPoseidon2Config { impl StarkGenericConfig for BabyBearPoseidon2Config { type Pcs = Pcs; + type Dft = Dft; type Challenge = Challenge; type Challenger = Challenger; @@ -105,6 +109,10 @@ impl StarkGenericConfig for BabyBearPoseidon2Config { &self.pcs } + fn dft(&self) -> &Dft { + &self.dft + } + fn initialise_challenger(&self) -> Challenger { let mut challenger = Challenger::new(self.perm.clone()); for value in &self.challenger_seed { diff --git a/src/types.rs b/src/types.rs index ee42b65..3128af0 100644 --- a/src/types.rs +++ b/src/types.rs @@ -13,14 +13,27 @@ use p3_challenger::{ }; use p3_commit::{ExtensionMmcs, Pcs as PcsTrait}; use p3_dft::Radix2DitParallel; +use p3_field::BasedVectorSpace; use p3_field::{ - ExtensionField, Field, PrimeCharacteristicRing, TwoAdicField, extension::BinomialExtensionField, + ExtensionField, Field, PrimeCharacteristicRing, PrimeField64, TwoAdicField, + extension::BinomialExtensionField, }; -use p3_fri::{FriParameters as InnerFriParameters, TwoAdicFriPcs}; +use p3_fri::FriParameters as InnerFriParameters; +#[cfg(not(feature = "cuda"))] +use p3_fri::TwoAdicFriPcs; use p3_goldilocks::Goldilocks; +#[cfg(feature = "cuda")] +use p3_matrix::dense::RowMajorMatrix; +#[cfg(feature = "cuda")] +use p3_maybe_rayon::prelude::*; use p3_merkle_tree::MerkleTreeMmcs; use p3_symmetric::{CompressionFunctionFromHasher, SerializingHasher}; +#[cfg(feature = "cuda")] +use crate::cuda::CudaDft; +#[cfg(feature = "cuda")] +use crate::cuda::pcs::CudaPcsDft; + pub type Val = Goldilocks; pub type PackedVal = ::Packing; pub type ExtVal = BinomialExtensionField; @@ -79,10 +92,16 @@ impl GrindingChallenger for DeterministicPow { self.0.grind(bits) } } -pub type Mmcs = - MerkleTreeMmcs, Blake3CompressionFunction, 2, 32>; +type CpuMmcs = MerkleTreeMmcs, Blake3CompressionFunction, 2, 32>; +#[cfg(not(feature = "cuda"))] +pub type Mmcs = CpuMmcs; +#[cfg(feature = "cuda")] +pub type Mmcs = crate::cuda::mmcs::CudaMmcs; pub type ExtMmcs = ExtensionMmcs; -pub type Pcs = TwoAdicFriPcs; +#[cfg(not(feature = "cuda"))] +pub type Pcs = TwoAdicFriPcs; +#[cfg(feature = "cuda")] +pub type Pcs = crate::cuda::pcs::CudaTwoAdicFriPcs; pub type Commitment = >::Commitment; pub type Domain = >::Domain; @@ -91,10 +110,30 @@ pub type EvaluationsOnDomain<'a> = >::Evalua pub type PcsError = >::Error; pub type PcsProof = >::Proof; +#[cfg(feature = "cuda")] +fn cuda_coset_selectors( + trace_domain: Domain, + quotient_domain: Domain, +) -> crate::cuda::CudaCosetSelectors { + let rate_bits = quotient_domain.log_size() - trace_domain.log_size(); + crate::cuda::CudaCosetSelectors { + coset_shift: quotient_domain.shift(), + coset_generator: quotient_domain.subgroup_generator(), + trace_last: trace_domain.subgroup_generator().inverse(), + vanishing_start: quotient_domain + .shift() + .exp_power_of_2(trace_domain.log_size()), + vanishing_step: Val::two_adic_generator(rate_bits), + } +} + /// The reference [`StarkGenericConfig`] implementation. pub struct GoldilocksBlake3Config { /// The PCS used to commit polynomials and prove opening proofs. pcs: Pcs, + /// The same transform implementation used inside `pcs`, exposed for + /// prover-side quotient transforms which live outside the PCS API. + dft: Dft, /// Seed for fresh challengers: a domain-separation tag followed by a /// digest of all protocol parameters. challenger_seed: Vec, @@ -109,7 +148,18 @@ pub struct GoldilocksBlake3Config { impl GoldilocksBlake3Config { pub fn new(commitment_parameters: CommitmentParameters, fri_parameters: FriParameters) -> Self { - let pcs = new_pcs(commitment_parameters, fri_parameters); + #[cfg(feature = "cuda")] + { + assert_eq!( + commitment_parameters.cap_height, 0, + "the CUDA backend currently supports only cap_height = 0" + ); + assert!( + fri_parameters.max_log_arity <= 1, + "the CUDA backend currently supports only binary FRI folds" + ); + } + let (pcs, dft) = new_pcs(commitment_parameters, fri_parameters); // Seed the challenger with a protocol tag for domain separation, // followed by every protocol parameter. Binding the parameters into // the seed means transcripts produced under different parameters @@ -132,6 +182,7 @@ impl GoldilocksBlake3Config { let max_quotient_degree = 1 << commitment_parameters.log_blowup; Self { pcs, + dft, challenger_seed, max_log_degree, max_quotient_degree, @@ -142,6 +193,7 @@ impl GoldilocksBlake3Config { impl StarkGenericConfig for GoldilocksBlake3Config { type Pcs = Pcs; + type Dft = Dft; type Challenge = ExtVal; type Challenger = Challenger; @@ -149,6 +201,10 @@ impl StarkGenericConfig for GoldilocksBlake3Config { &self.pcs } + fn dft(&self) -> &Dft { + &self.dft + } + fn initialise_challenger(&self) -> Challenger { Challenger::from_hasher(self.challenger_seed.clone(), Blake3) } @@ -164,6 +220,263 @@ impl StarkGenericConfig for GoldilocksBlake3Config { fn log_blowup(&self) -> usize { self.log_blowup } + + fn canonicalize_proof(proof: &mut crate::prover::Proof) { + fn canonical_base(value: &mut Val) { + *value = Val::from_u64(value.as_canonical_u64()); + } + + fn canonical_ext(value: &mut ExtVal) { + let coefficients: &[Val] = value.as_basis_coefficients_slice(); + *value = ExtVal::from_basis_coefficients_slice(&[ + Val::from_u64(coefficients[0].as_canonical_u64()), + Val::from_u64(coefficients[1].as_canonical_u64()), + ]) + .expect("quadratic extension element"); + } + + fn canonical_opened(round: &mut p3_commit::OpenedValuesForRound) { + for matrix in round { + for point in matrix { + point.iter_mut().for_each(canonical_ext); + } + } + } + + proof + .intermediate_accumulators + .iter_mut() + .for_each(canonical_ext); + canonical_opened(&mut proof.quotient_opened_values); + if let Some(round) = &mut proof.preprocessed_opened_values { + canonical_opened(round); + } + canonical_opened(&mut proof.stage_1_opened_values); + canonical_opened(&mut proof.stage_2_opened_values); + + let fri = &mut proof.opening_proof; + fri.commit_pow_witnesses.iter_mut().for_each(canonical_base); + canonical_base(&mut fri.query_pow_witness); + fri.final_poly.iter_mut().for_each(canonical_ext); + for query in &mut fri.query_proofs { + for opening in &mut query.input_proof { + for row in &mut opening.opened_values { + row.iter_mut().for_each(canonical_base); + } + } + for step in &mut query.commit_phase_openings { + step.sibling_values.iter_mut().for_each(canonical_ext); + } + } + } + + #[cfg(feature = "cuda")] + fn accelerated_quotient_values( + &self, + circuit: &crate::system::Circuit, + lookup_publics: &[Val], + trace_domain: crate::config::Domain, + quotient_domain: crate::config::Domain, + preprocessed: Option<(&crate::config::PcsData, usize)>, + stage_1: (&crate::config::PcsData, usize), + stage_2: (&crate::config::PcsData, usize), + alpha: ExtVal, + constraint_count: usize, + ) -> Option> { + let main = stage_1.0.resident(stage_1.1)?; + let s2 = stage_2.0.resident(stage_2.1)?; + let prep = match preprocessed { + Some((data, index)) => Some(data.resident(index)?), + None => None, + }; + let qsize = quotient_domain.size(); + let selectors = cuda_coset_selectors(trace_domain, quotient_domain); + let mut powers = Vec::with_capacity(constraint_count); + let mut power = ExtVal::ONE; + for _ in 0..constraint_count { + powers.push(power); + power *= alpha; + } + powers.reverse(); + let mut alpha_flat = Vec::with_capacity(2 * constraint_count); + for coordinate in 0..2 { + alpha_flat.extend(powers.iter().map(|x| { + >::as_basis_coefficients_slice(x)[coordinate] + })); + } + let n = Val::from_usize(trace_domain.size()); + let g = Val::two_adic_generator(trace_domain.size().ilog2() as usize); + let norm = (n * g).inverse(); + let delta = [ + (lookup_publics[6] - lookup_publics[4]) * norm, + (lookup_publics[7] - lookup_publics[5]) * norm, + ]; + let next_step = 1usize << (quotient_domain.size().ilog2() - trace_domain.size().ilog2()); + let flat = crate::cuda::quotient_values_resident( + &circuit.graph, + prep, + main, + s2, + lookup_publics, + selectors, + &alpha_flat, + &delta, + crate::system::extension_params::().w, + qsize, + next_step, + circuit.lookup_group_size, + ); + Some( + flat.values + .chunks_exact(2) + .map(|coords| { + ExtVal::from_basis_coefficients_slice(coords) + .expect("CUDA quotient has two coordinates") + }) + .collect(), + ) + } + + #[cfg(feature = "cuda")] + fn accelerated_quotient_commit( + &self, + inputs: &[crate::config::QuotientCommitInput<'_, Self>], + alpha: ExtVal, + ) -> Option<(crate::config::Com, crate::config::PcsData)> { + use crate::cuda::mmcs::CudaCommitMmcs; + let ldes: Option> = inputs + .par_iter() + .map(|input| { + let main = input.stage_1.0.resident(input.stage_1.1)?; + let stage2 = input.stage_2.0.resident(input.stage_2.1)?; + let preprocessed = match input.preprocessed { + Some((data, index)) => Some(data.resident(index)?), + None => None, + }; + let quotient_size = input.quotient_domain.size(); + let quotient_degree = input.circuit.quotient_degree(); + let selectors = cuda_coset_selectors(input.trace_domain, input.quotient_domain); + + let mut powers = Vec::with_capacity(input.constraint_count); + let mut power = ExtVal::ONE; + for _ in 0..input.constraint_count { + powers.push(power); + power *= alpha; + } + powers.reverse(); + let mut alpha_flat = Vec::with_capacity(2 * input.constraint_count); + for coordinate in 0..2 { + alpha_flat.extend(powers.iter().map(|value| { + >::as_basis_coefficients_slice(value) + [coordinate] + })); + } + let trace_size = input.trace_domain.size(); + let n = Val::from_usize(trace_size); + let generator = Val::two_adic_generator(trace_size.ilog2() as usize); + let normalization = (n * generator).inverse(); + let delta = [ + (input.lookup_publics[6] - input.lookup_publics[4]) * normalization, + (input.lookup_publics[7] - input.lookup_publics[5]) * normalization, + ]; + let next_step = quotient_size / trace_size; + Some(crate::cuda::quotient_lde_resident( + &self.pcs.dft, + &input.circuit.graph, + preprocessed, + main, + stage2, + &input.lookup_publics, + selectors, + &alpha_flat, + &delta, + crate::system::extension_params::().w, + quotient_size, + next_step, + input.circuit.lookup_group_size, + quotient_degree, + self.log_blowup, + )) + }) + .collect(); + let ldes = ldes?; + Some(self.pcs.mmcs.commit_cuda_resident(ldes)) + } + + #[cfg(feature = "cuda")] + fn accelerated_lookup_commit( + &self, + inputs: &[crate::config::LookupCommitInput<'_, Self>], + lookup_challenge: ExtVal, + fingerprint_challenge: ExtVal, + mut accumulator: ExtVal, + ) -> Option<( + crate::config::Com, + crate::config::PcsData, + Vec, + )> { + use crate::cuda::mmcs::CudaCommitMmcs; + let pair = |value: ExtVal| { + let coordinates = value.as_basis_coefficients_slice(); + [coordinates[0], coordinates[1]] + }; + let beta = pair(lookup_challenge); + let gamma = pair(fingerprint_challenge); + let extension_generator = ExtVal::from_basis_coefficients_slice(&[Val::ZERO, Val::ONE])?; + let ext_w = pair(extension_generator * extension_generator)[0]; + let evaluate = |input: &crate::config::LookupCommitInput<'_, Self>| { + let main = input.stage_1.0.resident_with_trace(input.stage_1.1)?; + let preprocessed = match input.preprocessed { + Some((data, index)) => Some(data.resident(index)?), + None => None, + }; + let (height, num_lookups, multiplicities, args, arg_offsets) = + input.lookup_values.cuda_parts(); + let result = if num_lookups == 0 { + Some(crate::cuda::lookup_lde_resident( + &self.pcs.dft, + multiplicities, + args, + arg_offsets, + height, + num_lookups, + input.circuit.lookup_group_size.max(1), + beta, + gamma, + ext_w, + self.log_blowup, + )) + } else { + crate::cuda::lookup_graph_lde_resident( + &self.pcs.dft, + &input.circuit.graph, + preprocessed, + main, + height, + input.circuit.lookup_group_size.max(1), + beta, + gamma, + ext_w, + self.log_blowup, + ) + }; + // SAFETY: evaluation is synchronous and complete, while the + // retained matrix remains owned by the prover data. + unsafe { main.release_trace() }; + result + }; + let results: Option> = inputs.iter().map(evaluate).collect(); + let results = results?; + let mut ldes = Vec::with_capacity(results.len()); + let mut intermediates = Vec::with_capacity(inputs.len()); + for (lde, total) in results { + accumulator += ExtVal::from_basis_coefficients_slice(&total)?; + intermediates.push(accumulator); + ldes.push(lde); + } + let (commitment, data) = self.pcs.mmcs.commit_cuda_resident(ldes); + Some((commitment, data, intermediates)) + } } /// Parameters of the polynomial commitment: Reed-Solomon rate and Merkle @@ -196,17 +509,45 @@ pub struct FriParameters { pub query_proof_of_work_bits: usize, } -type Blake3CompressionFunction = CompressionFunctionFromHasher; +pub(crate) type Blake3CompressionFunction = CompressionFunctionFromHasher; + +#[cfg(feature = "cuda")] +impl CudaPcsDft for CudaDft { + fn prepare_coset_lde_constants(&self, height: usize, added_bits: usize, shift: Val) { + self.prepare_coset_lde_constants(height, added_bits, shift); + } + + fn coset_lde_batch_resident( + &self, + matrix: &RowMajorMatrix, + added_bits: usize, + shift: Val, + ) -> crate::cuda::CudaLde { + self.coset_lde_batch_resident(matrix, added_bits, shift) + } +} + type Dft = Radix2DitParallel; +#[cfg(not(feature = "cuda"))] +type PcsDft = Dft; +#[cfg(feature = "cuda")] +type PcsDft = CudaDft; fn new_mmcs(cap_height: usize) -> Mmcs { let byte_hash = Blake3; let field_hash = SerializingHasher::new(byte_hash); let compress = Blake3CompressionFunction::new(byte_hash); - Mmcs::new(field_hash, compress, cap_height) + let cpu = CpuMmcs::new(field_hash, compress, cap_height); + #[cfg(not(feature = "cuda"))] + return cpu; + #[cfg(feature = "cuda")] + crate::cuda::mmcs::CudaMmcs::new(cpu) } -fn new_pcs(commitment_parameters: CommitmentParameters, fri_parameters: FriParameters) -> Pcs { +fn new_pcs( + commitment_parameters: CommitmentParameters, + fri_parameters: FriParameters, +) -> (Pcs, Dft) { let val_mmcs = new_mmcs(commitment_parameters.cap_height); let mmcs = ExtensionMmcs::new(val_mmcs.clone()); let inner_parameters = InnerFriParameters { @@ -219,17 +560,57 @@ fn new_pcs(commitment_parameters: CommitmentParameters, fri_parameters: FriParam mmcs, }; let dft = Dft::default(); - Pcs::new(dft, val_mmcs, inner_parameters) + let pcs = Pcs::new(PcsDft::default(), val_mmcs, inner_parameters); + (pcs, dft) } #[cfg(test)] mod pcs_ref_gen { use super::*; use p3_commit::Mmcs as _; - use p3_field::PrimeCharacteristicRing; + use p3_field::{ + BasedVectorSpace, PrimeCharacteristicRing, PrimeField64, batch_multiplicative_inverse, + }; use p3_matrix::dense::RowMajorMatrix; use p3_symmetric::{CryptographicHasher, PseudoCompressionFunction}; + #[cfg(feature = "cuda")] + fn test_fri_parameters(max_log_arity: usize) -> FriParameters { + FriParameters { + log_final_poly_len: 0, + max_log_arity, + num_queries: 1, + commit_proof_of_work_bits: 0, + query_proof_of_work_bits: 0, + } + } + + #[cfg(feature = "cuda")] + #[test] + #[should_panic(expected = "cap_height = 0")] + fn cuda_rejects_nonzero_merkle_caps() { + let _ = GoldilocksBlake3Config::new( + CommitmentParameters { + log_blowup: 1, + cap_height: 1, + }, + test_fri_parameters(1), + ); + } + + #[cfg(feature = "cuda")] + #[test] + #[should_panic(expected = "binary FRI folds")] + fn cuda_rejects_multi_bit_fri_folds() { + let _ = GoldilocksBlake3Config::new( + CommitmentParameters { + log_blowup: 1, + cap_height: 0, + }, + test_fri_parameters(2), + ); + } + fn limbs(d: [u8; 32]) -> [u64; 4] { core::array::from_fn(|i| u64::from_le_bytes(d[i * 8..i * 8 + 8].try_into().unwrap())) } @@ -241,20 +622,126 @@ mod pcs_ref_gen { o } - /// Generates the Blake3 reference values asserted by `Ix/MultiStark/Tests.lean` - /// (`pcs_hash_test`, `pcs_merkle_test`). Run with `--nocapture` and copy. + /// Pins canonical Goldilocks base/extension arithmetic and Montgomery + /// batch inversion at values chosen around the Solinas modulus. These are + /// compact host-side oracles for the first CUDA field kernels. + #[test] + fn goldilocks_arithmetic_contract() { + let canonical = |value: Val| value.as_canonical_u64(); + for (a, b, sum, difference, product) in [ + ( + 1_311_768_467_463_790_320, + 1_147_797_409_030_816_545, + 2_459_565_876_494_606_865, + 163_971_058_432_973_775, + 14_965_091_924_900_821_934, + ), + ( + 18_446_744_069_414_584_319, + 4_294_967_296, + 4_294_967_294, + 18_446_744_065_119_617_023, + 18_446_744_060_824_649_729, + ), + ( + 4_294_967_295, + 4_294_967_297, + 8_589_934_592, + 18_446_744_069_414_584_319, + 4_294_967_294, + ), + ] { + let a = Val::from_u64(a); + let b = Val::from_u64(b); + assert_eq!(canonical(a + b), sum); + assert_eq!(canonical(a - b), difference); + assert_eq!(canonical(a * b), product); + } + + let values = [2, 3, 4_294_967_296, 1_311_768_467_463_790_320].map(Val::from_u64); + let inverses = batch_multiplicative_inverse(&values); + assert_eq!( + inverses.into_iter().map(canonical).collect::>(), + vec![ + 9_223_372_034_707_292_161, + 12_297_829_379_609_722_881, + 18_446_744_065_119_617_026, + 14_736_413_637_906_284_881, + ] + ); + + let left = ExtVal::from_basis_coefficients_fn(|i| [Val::from_u8(3), Val::from_u8(5)][i]); + let right = ExtVal::from_basis_coefficients_fn(|i| [Val::from_u8(11), Val::from_u8(13)][i]); + let product = left * right; + assert_eq!( + product + .as_basis_coefficients_slice() + .iter() + .copied() + .map(canonical) + .collect::>(), + vec![488, 94] + ); + } + + /// Pins the Blake3 and Merkle values shared with `Ix/MultiStark/Tests.lean` + /// (`pcs_hash_test`, `pcs_merkle_test`). A future GPU implementation must + /// match these byte-for-byte, including field serialization and tree + /// layout. #[test] - fn gen_pcs_refs() { + fn blake3_and_merkle_contract() { let f = Val::from_u32; let fh = SerializingHasher::new(Blake3); - for n in [3u32, 17, 22, 20] { + for (n, expected) in [ + ( + 3u32, + [ + 4163513704854067712, + 9384471110237386207, + 13671380075168847140, + 1533933974187331481, + ], + ), + ( + 17, + [ + 8431665677194841246, + 4495111673672851816, + 7709594803249897978, + 12683511314940902790, + ], + ), + ( + 22, + [ + 14017803411919507972, + 9236340131056405306, + 11356520758956579629, + 2008168271701183309, + ], + ), + ( + 20, + [ + 8822819174011220231, + 9835070768970864367, + 9646176123001837413, + 1210344881395534089, + ], + ), + ] { let row: Vec = (1..=n).map(f).collect(); - println!("LEAF{} {:?}", n, limbs(fh.hash_iter(row))); + assert_eq!(limbs(fh.hash_iter(row)), expected, "leaf width {n}"); } let comp = Blake3CompressionFunction::new(Blake3); - println!( - "COMPRESS {:?}", - limbs(comp.compress([dig([1, 2, 3, 4]), dig([5, 6, 7, 8])])) + assert_eq!( + limbs(comp.compress([dig([1, 2, 3, 4]), dig([5, 6, 7, 8])])), + [ + 16432952784711837466, + 12565756115161032165, + 6915939387221618258, + 11123773279136987111, + ], ); // Merkle tree: matrices of heights 8/4/2 and widths 2/3/1, opened at index 5. @@ -268,23 +755,115 @@ mod pcs_ref_gen { let mut m2 = vec![f(0); 2]; m2[1] = f(202); // row 1 = [202] let mmcs = new_mmcs(0); + let (commit, pd) = mmcs.commit(vec![ + RowMajorMatrix::new(m0.clone(), 2), + RowMajorMatrix::new(m1.clone(), 3), + RowMajorMatrix::new(m2.clone(), 1), + ]); + let bo = mmcs.open_batch(5, &pd); + assert_eq!( + bo.opened_values, + vec![ + vec![f(11), f(12)], + vec![f(107), f(108), f(109)], + vec![f(202)] + ] + ); + assert_eq!( + bo.opening_proof + .iter() + .copied() + .map(limbs) + .collect::>(), + vec![ + [ + 824163284354560741, + 10184227291309369989, + 7314170388788081421, + 2210258918235055872, + ], + [ + 16321412416894375658, + 13817763133082311448, + 4555362725758189505, + 13946835461337436585, + ], + [ + 11035117895010660519, + 10627114985553641692, + 18209541265052796223, + 11062544859664569990, + ], + ] + ); + assert_eq!( + commit.roots(), + &[[ + 45, 230, 248, 40, 61, 21, 136, 65, 180, 102, 50, 238, 76, 222, 102, 39, 123, 114, + 106, 220, 182, 223, 92, 68, 228, 55, 152, 7, 80, 209, 237, 16, + ]] + ); + + let mmcs = new_mmcs(2); let (commit, pd) = mmcs.commit(vec![ RowMajorMatrix::new(m0, 2), RowMajorMatrix::new(m1, 3), RowMajorMatrix::new(m2, 1), ]); let bo = mmcs.open_batch(5, &pd); - println!("OPENED {:?}", bo.opened_values); - for (i, s) in bo.opening_proof.iter().enumerate() { - println!("SIB{} {:?}", i, limbs(*s)); - } - println!("COMMIT {:?}", commit); + assert_eq!( + commit + .roots() + .iter() + .copied() + .map(limbs) + .collect::>(), + vec![ + [ + 16321412416894375658, + 13817763133082311448, + 4555362725758189505, + 13946835461337436585, + ], + [ + 16321412416894375658, + 13817763133082311448, + 4555362725758189505, + 13946835461337436585, + ], + [ + 2755952710066137292, + 16563663342057344133, + 5946896676730904047, + 10390238708790769607, + ], + [ + 16321412416894375658, + 13817763133082311448, + 4555362725758189505, + 13946835461337436585, + ], + ] + ); + assert_eq!( + bo.opening_proof + .iter() + .copied() + .map(limbs) + .collect::>(), + vec![[ + 824163284354560741, + 10184227291309369989, + 7314170388788081421, + 2210258918235055872, + ]] + ); } - /// Regenerates the Blake3-challenger reference values for `sample_bits_test` - /// and `pcs_challenger4_test`. + /// Pins the Blake3-challenger reference values used by `sample_bits_test` + /// and `pcs_challenger4_test` in Ix. #[test] - fn gen_challenger_refs() { + fn blake3_challenger_contract() { use p3_challenger::{CanObserve, CanSampleBits, FieldChallenger}; use p3_field::{BasedVectorSpace, PrimeField64}; let g = Val::from_u64; @@ -295,26 +874,20 @@ mod pcs_ref_gen { // sample_bits_test: observe 0x0102030405060708, sample_bits(20). let mut ch = Challenger::from_hasher(vec![], Blake3); ch.observe(g(0x0102030405060708)); - println!( - "SAMPLE_BITS {}", - CanSampleBits::::sample_bits(&mut ch, 20) - ); + assert_eq!(CanSampleBits::::sample_bits(&mut ch, 20), 1019203); // pcs_challenger4_test: the α_pcs/α_fri/β/index continuation. let mut ch = Challenger::from_hasher(vec![], Blake3); ch.observe(g(0x0102030405060708)); ch.observe(g(0x1122334455667788)); let apcs: ExtVal = ch.sample_algebra_element(); let afri: ExtVal = ch.sample_algebra_element(); - println!("APCS {:?}", el(apcs)); - println!("AFRI {:?}", el(afri)); + assert_eq!(el(apcs), (17795849114622667264, 4116843485681689527)); + assert_eq!(el(afri), (11768399386651893439, 10948618071942561750)); ch.observe(g(0x00000000deadbeef)); let beta: ExtVal = ch.sample_algebra_element(); - println!("BETA {:?}", el(beta)); + assert_eq!(el(beta), (12096272534537655203, 11431251745744402868)); ch.observe(g(0x0a0b0c0d01020304)); ch.observe(g(0x0000000000000002)); - println!( - "SAMPLE_BITS2 {}", - CanSampleBits::::sample_bits(&mut ch, 20) - ); + assert_eq!(CanSampleBits::::sample_bits(&mut ch, 20), 458922); } } diff --git a/src/verifier.rs b/src/verifier.rs index 40398a7..51a7cff 100644 --- a/src/verifier.rs +++ b/src/verifier.rs @@ -715,7 +715,9 @@ mod tests { types::{CommitmentParameters, ExtVal, FriParameters, GoldilocksBlake3Config, Val}, }; use p3_air::{Air, AirBuilder, BaseAir, WindowAccess}; + use p3_blake3::Blake3; use p3_matrix::dense::RowMajorMatrix; + use p3_symmetric::CryptographicHasher; enum CS { Pythagorean, @@ -801,7 +803,7 @@ mod tests { } #[test] - fn multi_stark_prove_verify_serialize() { + fn multi_stark_proof_bytes_contract() { let (system, key) = system(); let f = Val::from_u32; // 2^4 = 16 rows — small enough for fast CI @@ -811,17 +813,42 @@ mod tests { pythagorean_trace.extend(pythagorean_trace.clone()); complex_trace.extend(complex_trace.clone()); } - let witness = SystemWitness::from_stage_1( - vec![ - RowMajorMatrix::new(pythagorean_trace, 3), - RowMajorMatrix::new(complex_trace, 6), - ], - &system, - ); + let traces = vec![ + RowMajorMatrix::new(pythagorean_trace, 3), + RowMajorMatrix::new(complex_trace, 6), + ]; let no_claims = &[]; - let proof = system.prove_multiple_claims(&key, no_claims, witness); - // Serialization round-trip + let proof = system.prove_multiple_claims( + &key, + no_claims, + SystemWitness::from_stage_1(traces.clone(), &system), + ); + let second_proof = system.prove_multiple_claims( + &key, + no_claims, + SystemWitness::from_stage_1(traces, &system), + ); + + // Proving is deterministic at these zero-PoW parameters, so this + // serialization is a protocol-compatibility contract for alternate + // backends rather than merely a round-trip smoke test. let proof_bytes = proof.to_bytes().expect("Failed to serialize proof"); + let second_proof_bytes = second_proof.to_bytes().expect("Failed to serialize proof"); + assert_eq!(proof_bytes, second_proof_bytes); + assert_eq!( + proof_bytes.len(), + 77_637, + "proof encoding changed; update only with an intentional protocol review" + ); + assert_eq!( + Blake3.hash_slice(&proof_bytes), + [ + 132, 122, 135, 163, 73, 111, 225, 81, 221, 201, 107, 28, 30, 21, 49, 58, 253, 13, + 161, 49, 19, 184, 213, 239, 107, 152, 43, 42, 67, 255, 151, 193, + ], + "proof bytes changed; update only with an intentional protocol review" + ); + let proof2 = Proof::from_bytes(&proof_bytes).expect("Failed to deserialize proof"); system.verify_multiple_claims(no_claims, &proof2).unwrap(); }