Add opt-in CUDA acceleration for Aiur proofs

.github/workflows/ci.yml

@@ -101,9 +101,9 @@ jobs:
       - name: Check Rustfmt code style
         uses: actions-rust-lang/rustfmt@v1
       - name: Check clippy warnings
-        run: cargo clippy --workspace --all-targets --all-features
+        run: cargo clippy --release --workspace --all-targets --features ix-ffi/parallel,ix-ffi/net,ix-ffi/test-ffi -- -D warnings
       - name: Check *everything* compiles
-        run: cargo check --all-targets --all-features --workspace
+        run: cargo check --release --workspace --all-targets --features ix-ffi/parallel,ix-ffi/net,ix-ffi/test-ffi
       - name: Tests
         run: cargo nextest run --release --profile ci --workspace
       - name: Get Rust version
@@ -114,6 +114,34 @@ jobs:
         with:
           rust-version: ${{ env.RUST_VERSION }}
 
+  # Compile and link the opt-in backend without requiring a GPU. Runtime and
+  # proof-byte equivalence are covered by multi-stark's NVIDIA smoke suite.
+  cuda-compile:
+    runs-on: warp-ubuntu-latest-x64-8x
+    container: nvidia/cuda:12.8.1-devel-ubuntu24.04
+    env:
+      # This job validates feature plumbing and the Rust-staticlib-to-Lean
+      # link without paying for a GPU runner. Multi-stark CI compiles the
+      # complete supported cubin set; one representative architecture is
+      # sufficient for this downstream integration gate.
+      MULTI_STARK_CUDA_ARCHS: "80"
+    steps:
+      - uses: actions/checkout@v7
+      - name: Install toolchain bootstrap dependencies
+        run: |
+          apt-get update
+          apt-get install --yes --no-install-recommends \
+            build-essential ca-certificates curl git libclang-dev
+      - uses: ./.github/actions/setup-rust-toolchain
+      - uses: leanprover/lean-action@v1
+        with:
+          auto-config: false
+          use-github-cache: false
+      - name: Compile CUDA Rust targets
+        run: cargo check --release --locked --workspace --all-targets --features ix-ffi/parallel,ix-ffi/cuda
+      - name: Link a CUDA-enabled Lean executable
+        run: IX_CUDA=1 lake build bench-typecheck
+
   # zkVM host build + execute gate: do the Zisk/SP1 hosts (and their guest
   # ELFs, via each workspace's build.rs) still compile AND run? rust-test
   # doesn't build these workspaces (special toolchains), and bench-main.yml's

Benchmarks/Aiur.lean

@@ -88,7 +88,7 @@ def main : IO Unit := do
     let system ← benchStep "build AiurSystem"
         (Aiur.AiurSystem.build compiled.bytecode commitmentParameters) friParameters
     let funIdx := compiled.getFuncIdx `main |>.get!
-    let (claim, proof, _) ← benchStep "prove fib 10"
+    let (claim, proof, _) ← benchStepE "prove fib 10"
         (Aiur.AiurSystem.prove system funIdx #[10]) default (oneShot := true)
     let _ ← benchStepE "verify fib 10"
         (Aiur.AiurSystem.verify system claim) proof

Benchmarks/RecursionDebug.lean

@@ -92,8 +92,9 @@ def proveConst (ixePath constName : String) (skipDeps : Bool)
   let (claim, proof) ←
     if skipDeps then
       let witness := IxVM.ClaimHarness.buildVerifyConst ixonEnv addr
-      let (claim, proof, _) := aiurSystem.proveIxVM funIdx witness.input witness.inputIOBuffer
-      pure (claim, proof)
+      match aiurSystem.proveIxVM funIdx witness.input witness.inputIOBuffer with
+      | .error e => IO.eprintln s!"proveIxVM failed: {e}"; return none
+      | .ok (claim, proof, _) => pure (claim, proof)
     else do
       let envHandle ← match Aiur.EnvHandle.fromIxe ixePath with
         | .error e => IO.eprintln s!"EnvHandle.fromIxe: {e}"; return none

Benchmarks/Typecheck.lean

@@ -537,13 +537,10 @@ def runTypecheckCmd (p : Cli.Parsed) : IO UInt32 := do
       let (proveRes, proveSec) ← timed fun _ =>
         if skipDeps then
           let witness := IxVM.ClaimHarness.buildVerifyConst ixonEnv addr
-          let (claim, proof, ioBuf) :=
-            if useInterp then
-              aiurSystem.prove funIdx witness.input witness.inputIOBuffer
-            else
-              aiurSystem.proveIxVM funIdx witness.input witness.inputIOBuffer
-          (.ok (claim, proof, ioBuf) :
-            Except String (Array Aiur.G × Aiur.Proof × Aiur.IOBuffer))
+          if useInterp then
+            aiurSystem.prove funIdx witness.input witness.inputIOBuffer
+          else
+            aiurSystem.proveIxVM funIdx witness.input witness.inputIOBuffer
         else
           match aiurSystem.proveAddrWithEnv funIdx envHandle addr.hash useInterp with
           | .error e => .error e
@@ -623,9 +620,14 @@ def runTypecheckCmd (p : Cli.Parsed) : IO UInt32 := do
             IO.println s!"  [{i + 1}/{ordered.size}] proving the verifier over {r.name} …"
             (← IO.getStdout).flush
             TracingTexray.resetPeakTreeRss
-            let ((rvClaim, rvProof), rvProveSec) ← timed fun _ =>
+            let (rvProveRes, rvProveSec) ← timed fun _ =>
               vSystem.proveMultiStark vIdx pubInput proofBytes vkBytes
                 claimBytes useInterp
+            let (rvClaim, rvProof) ← match rvProveRes with
+              | .ok result => pure result
+              | .error e =>
+                IO.eprintln s!"  outer prove {r.name} failed: {e}"
+                continue
             let rvPeak ← TracingTexray.peakTreeRssBytes
             let rvProofBytes := Aiur.Proof.toBytes rvProof
             let (rvVerifyRes, rvVerifySec) ← timed fun _ =>

Cargo.toml

@@ -56,7 +56,7 @@ libc = "0.2"
 log = "0.4"
 memmap2 = "0.9"
 mimalloc = { version = "0.1", default-features = false }
-multi-stark = { git = "https://github.com/argumentcomputer/multi-stark.git", rev = "c72d32130eb084327b73849e39c92734a8829011" }
+multi-stark = { git = "https://github.com/argumentcomputer/multi-stark.git", rev = "e43201522a95775d1495b4467fbab69fbf7a1ae8" }
 nom = "7.1.3"
 num-bigint = "0.4.6"
 quickcheck = "1.0.3"

Ix/Aiur/Protocol.lean

@@ -109,24 +109,24 @@ private opaque prove' : @& AiurSystem →
   @& Bytecode.FunIdx → @& Array G →
   (ioData : @& Array (G × Array G)) →
   (ioMap : @& Array ((G × Array G) × IOKeyInfo)) →
-    ProveResult
+    Except String ProveResult
 
 /-- Executes the bytecode function `funIdx` with the given `args` and `ioBuffer`,
 then generates a proof of the computation. Returns the claim
 (`#[functionChannel, funIdx] ++ args ++ output`), the `Proof`, and the
 updated `IOBuffer`. -/
 def prove (system : @& AiurSystem)
   (funIdx : @& Bytecode.FunIdx) (args : @& Array G) (ioBuffer : IOBuffer) :
-    Array G × Proof × IOBuffer :=
-  let r := prove' system funIdx args ioBuffer.data.toArray ioBuffer.map.toArray
-  (r.claim, r.proof, .ofArrays r.ioData r.ioMap)
+    Except String (Array G × Proof × IOBuffer) :=
+  (prove' system funIdx args ioBuffer.data.toArray ioBuffer.map.toArray).map
+    fun r => (r.claim, r.proof, .ofArrays r.ioData r.ioMap)
 
 @[extern "rs_aiur_system_prove_ixvm"]
 private opaque proveIxVM' : @& AiurSystem →
   @& Bytecode.FunIdx → @& Array G →
   (ioData : @& Array (G × Array G)) →
   (ioMap : @& Array ((G × Array G) × IOKeyInfo)) →
-    ProveResult
+    Except String ProveResult
 
 /-- IxVM-native prove: same shape as `prove`, but routes execution
     through the codegen'd Rust kernel (`execute_generated`) instead
@@ -135,9 +135,9 @@ private opaque proveIxVM' : @& AiurSystem →
     `system.toplevel` is the IxVM kernel's bytecode. -/
 def proveIxVM (system : @& AiurSystem)
   (funIdx : @& Bytecode.FunIdx) (args : @& Array G) (ioBuffer : IOBuffer) :
-    Array G × Proof × IOBuffer :=
-  let r := proveIxVM' system funIdx args ioBuffer.data.toArray ioBuffer.map.toArray
-  (r.claim, r.proof, .ofArrays r.ioData r.ioMap)
+    Except String (Array G × Proof × IOBuffer) :=
+  (proveIxVM' system funIdx args ioBuffer.data.toArray ioBuffer.map.toArray).map
+    fun r => (r.claim, r.proof, .ofArrays r.ioData r.ioMap)
 
 /-- Prove the MultiStark recursive verifier over raw proof/vk/claims
     byte blobs. The IO advice buffer is built natively in Rust (see
@@ -152,7 +152,7 @@ def proveIxVM (system : @& AiurSystem)
 opaque proveMultiStark (system : @& AiurSystem)
   (funIdx : @& Bytecode.FunIdx) (pubInput : @& Array G)
   (proofBytes vkBytes claimBytes : @& ByteArray) (useBytecode : Bool := false) :
-    Array G × Proof
+    Except String (Array G × Proof)
 
 @[extern "rs_aiur_system_prove_addr_with_env"]
 private opaque proveAddrWithEnv' : @& AiurSystem →

Ix/Cli/ProveCmd.lean

@@ -88,9 +88,11 @@ def proveOne (aiurSystem : Aiur.AiurSystem)
         return 1
       | .ok (_claimBytes, proof, _outIO) => pure proof
     | .leanW witness, _ =>
-      let (_aiurClaim, proof, _outIO) :=
-        aiurSystem.proveIxVM funIdx witness.input witness.inputIOBuffer
-      pure proof
+      match aiurSystem.proveIxVM funIdx witness.input witness.inputIOBuffer with
+      | .error e =>
+        IO.eprintln s!"{label}: proveIxVM error: {e}"
+        return 1
+      | .ok (_aiurClaim, proof, _outIO) => pure proof
     | _, none =>
       IO.eprintln s!"{label}: internal: addr/shard target with no envHandle"
       return 1

README.md

@@ -179,6 +179,27 @@ https://bencher.dev/console/projects/ix/plots. `ix bench` runs the same
 cells locally, and `!benchmark` runs them on a PR — see
 [docs/benchmarking.md](docs/benchmarking.md).
 
+### CUDA-accelerated Aiur proving
+
+Aiur can use multi-stark's first-party CUDA backend. CUDA is opt-in: ordinary
+Cargo and Lake builds remain CPU-only and do not require a CUDA toolkit or
+runtime. First generate the benchmark environment, then set `IX_CUDA=1` (also
+accepts `true` or `yes`) for the proving command:
+
+```sh
+lake exe ix compile Benchmarks/Compile/CompileInitStd.lean --out InitStd.ixe
+IX_CUDA=1 lake exe bench-typecheck --ixe InitStd.ixe \
+  --consts Vector.extract_append --recursive
+```
+
+Rust consumers can instead enable the `cuda` feature on `aiur` or `ix-ffi`.
+The backend requires an NVIDIA GPU and a CUDA toolkit with `nvcc`; build and
+architecture controls are documented in the multi-stark repository. It keeps
+the Goldilocks/BLAKE3 protocol and proof format unchanged, and GPU proofs remain
+verifiable by the CPU implementation. Dated hardware measurements belong in
+[BENCHMARKS.md](BENCHMARKS.md) and [docs/benchmarking.md](docs/benchmarking.md),
+not this stable overview.
+
 ## Usage
 
 ### Prerequisites

Tests/Aiur/Common.lean

@@ -123,16 +123,18 @@ def AiurTestEnv.runTestCase (env : AiurTestEnv) (testCase : AiurTestCase) : Test
       else .done
     if !testCase.withProof then execTest ++ interpTest
     else
-      let (claim, proof, ioBuffer) := env.aiurSystem.prove
-        funIdx testCase.input testCase.inputIOBuffer
-      let claimTest := test s!"Claim matches for {label}"
-        (claim == Aiur.buildClaim funIdx testCase.input testCase.expectedOutput)
-      let ioTest := test s!"IOBuffer matches for {label}"
-        (ioBuffer == testCase.expectedIOBuffer)
-      let proof := .ofBytes proof.toBytes
-      let pvTest := withExceptOk s!"Prove/verify works for {label}"
-        (env.aiurSystem.verify claim proof) fun _ => .done
-      execTest ++ interpTest ++ claimTest ++ ioTest ++ pvTest
+      match env.aiurSystem.prove funIdx testCase.input testCase.inputIOBuffer with
+      | .error _ =>
+        execTest ++ interpTest ++ test s!"Prove succeeds for {label}" false
+      | .ok (claim, proof, ioBuffer) =>
+        let claimTest := test s!"Claim matches for {label}"
+          (claim == Aiur.buildClaim funIdx testCase.input testCase.expectedOutput)
+        let ioTest := test s!"IOBuffer matches for {label}"
+          (ioBuffer == testCase.expectedIOBuffer)
+        let proof := .ofBytes proof.toBytes
+        let pvTest := withExceptOk s!"Prove/verify works for {label}"
+          (env.aiurSystem.verify claim proof) fun _ => .done
+        execTest ++ interpTest ++ claimTest ++ ioTest ++ pvTest
 
 def mkAiurTests (toplevelFn : Except Aiur.Global Aiur.Source.Toplevel)
     (cases : List AiurTestCase) : TestSeq :=

Tests/MultiStark.lean

@@ -153,7 +153,9 @@ def endToEndSuite : IO UInt32 := do
 
   -- ── prove factorial(5) = 120 (`G` is a reserved DSL token, spell it qualified)
   let input := #[Aiur.G.ofNat 5]
-  let (claim, proof, _) := facSystem.prove facIdx input default
+  let (claim, proof, _) ← match facSystem.prove facIdx input default with
+    | .ok result => pure result
+    | .error e => IO.eprintln s!"factorial prove failed: {e}"; return 1
   let expectedClaim := buildClaim facIdx input #[Aiur.G.ofNat 120]
   let proofBytes := proof.toBytes
 

crates/aiur/Cargo.toml

@@ -18,6 +18,7 @@ tracing-texray = { workspace = true }
 [features]
 default = []
 parallel = ["multi-stark/parallel"]
+cuda = ["multi-stark/cuda"]
 
 [lints]
 workspace = true

crates/ffi/Cargo.toml

@@ -48,6 +48,7 @@ serde = { version = "1.0.219", features = ["derive"], optional = true }
 [features]
 default = []
 parallel = ["aiur/parallel"]
+cuda = ["aiur/cuda"]
 test-ffi = []
 net = ["bytes", "tokio", "iroh", "iroh-base", "n0-error", "getrandom", "bincode", "serde"]