diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 6ae6220a6..267f0b62e 100644 --- a/.github/workflows/ci.yml +++ b/.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 diff --git a/Benchmarks/Aiur.lean b/Benchmarks/Aiur.lean index 931d855a6..9542806f9 100644 --- a/Benchmarks/Aiur.lean +++ b/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 diff --git a/Benchmarks/RecursionDebug.lean b/Benchmarks/RecursionDebug.lean index 1b5bff892..41167de24 100644 --- a/Benchmarks/RecursionDebug.lean +++ b/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 diff --git a/Benchmarks/Typecheck.lean b/Benchmarks/Typecheck.lean index 4bdd9e125..177905eae 100644 --- a/Benchmarks/Typecheck.lean +++ b/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 _ => diff --git a/Cargo.lock b/Cargo.lock index eebdffbbc..db881f753 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2084,9 +2084,10 @@ dependencies = [ [[package]] name = "multi-stark" version = "0.1.0" -source = "git+https://github.com/argumentcomputer/multi-stark.git?rev=c72d32130eb084327b73849e39c92734a8829011#c72d32130eb084327b73849e39c92734a8829011" +source = "git+https://github.com/argumentcomputer/multi-stark.git?rev=e43201522a95775d1495b4467fbab69fbf7a1ae8#e43201522a95775d1495b4467fbab69fbf7a1ae8" dependencies = [ "bincode", + "itertools 0.14.0", "p3-air", "p3-blake3", "p3-challenger", @@ -2095,6 +2096,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 f535dd771..1791ae265 100644 --- a/Cargo.toml +++ b/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" diff --git a/Ix/Aiur/Protocol.lean b/Ix/Aiur/Protocol.lean index 57f954a8c..412ad59ea 100644 --- a/Ix/Aiur/Protocol.lean +++ b/Ix/Aiur/Protocol.lean @@ -109,7 +109,7 @@ 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 @@ -117,16 +117,16 @@ then generates a proof of the computation. Returns the claim 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 → diff --git a/Ix/Cli/ProveCmd.lean b/Ix/Cli/ProveCmd.lean index a35609d9a..f5802c09e 100644 --- a/Ix/Cli/ProveCmd.lean +++ b/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 diff --git a/README.md b/README.md index 7dffe9ceb..5cd891617 100644 --- a/README.md +++ b/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 diff --git a/Tests/Aiur/Common.lean b/Tests/Aiur/Common.lean index f90e64f66..f16bbc739 100644 --- a/Tests/Aiur/Common.lean +++ b/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 := diff --git a/Tests/MultiStark.lean b/Tests/MultiStark.lean index a1b09e149..26d62327d 100644 --- a/Tests/MultiStark.lean +++ b/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 diff --git a/crates/aiur/Cargo.toml b/crates/aiur/Cargo.toml index 9061574a0..cf9e4321e 100644 --- a/crates/aiur/Cargo.toml +++ b/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 diff --git a/crates/ffi/Cargo.toml b/crates/ffi/Cargo.toml index 1e02a9a6d..7a0bc7896 100644 --- a/crates/ffi/Cargo.toml +++ b/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"] diff --git a/crates/ffi/src/aiur/protocol.rs b/crates/ffi/src/aiur/protocol.rs index acf20c60e..9adb0cd3c 100644 --- a/crates/ffi/src/aiur/protocol.rs +++ b/crates/ffi/src/aiur/protocol.rs @@ -225,15 +225,17 @@ extern "C" fn rs_aiur_system_prove( args: LeanArray>, io_data_arr: LeanArray>, io_map_arr: LeanArray>, -) -> LeanAiurProveResult { - let fun_idx = lean_unbox_nat_as_usize(fun_idx.inner()); - let args = args.map(|x| lean_unbox_g(&x)); - let mut io_buffer = decode_io_buffer(&io_data_arr, &io_map_arr); +) -> LeanExcept { + ffi_catch_unwind("AiurSystem.prove", || { + let fun_idx = lean_unbox_nat_as_usize(fun_idx.inner()); + let args = args.map(|x| lean_unbox_g(&x)); + let mut io_buffer = decode_io_buffer(&io_data_arr, &io_map_arr); - let (claim, proof) = - aiur_system_obj.get().prove(fun_idx, &args, &mut io_buffer); + let (claim, proof) = + aiur_system_obj.get().prove(fun_idx, &args, &mut io_buffer); - build_prove_result(&claim, proof, &io_buffer) + build_prove_result(&claim, proof, &io_buffer).into() + }) } // ============================================================================= @@ -850,44 +852,47 @@ extern "C" fn rs_aiur_system_prove_addr_with_env( addr_bytes: LeanByteArray>, use_bytecode: bool, ) -> LeanExcept { - let fun_idx = lean_unbox_nat_as_usize(fun_idx.inner()); - let addr = match decode_addr(&addr_bytes) { - Ok(a) => a, - Err(e) => return LeanExcept::error_string(&e), - }; - let env = &env_handle.get().env; - - let (claim, input, mut io_buffer) = - match ixvm_codegen::aiur_ixvm_witness::build_claim_check_witness(env, &addr) - { - Ok(t) => t, - Err(e) => { - return LeanExcept::error_string(&format!("witness build: {e}")); - }, + ffi_catch_unwind_except("AiurSystem.proveAddrWithEnv", || { + let fun_idx = lean_unbox_nat_as_usize(fun_idx.inner()); + let addr = match decode_addr(&addr_bytes) { + Ok(a) => a, + Err(e) => return LeanExcept::error_string(&e), + }; + let env = &env_handle.get().env; + + let (claim, input, mut io_buffer) = + match ixvm_codegen::aiur_ixvm_witness::build_claim_check_witness( + env, &addr, + ) { + Ok(t) => t, + Err(e) => { + return LeanExcept::error_string(&format!("witness build: {e}")); + }, + }; + + // `use_bytecode` selects the generic Aiur bytecode interpreter over the + // codegen'd IxVM kernel (same toggle as + // `rs_aiur_toplevel_check_addr_with_env`). + let (_aiur_claim_arr, proof) = if use_bytecode { + aiur_system_obj.get().prove_ixvm( + fun_idx, + &input, + &mut io_buffer, + |toplevel, fun_idx, input, io_buffer| { + toplevel.execute(fun_idx, input, io_buffer) + }, + ) + } else { + aiur_system_obj.get().prove_ixvm( + fun_idx, + &input, + &mut io_buffer, + ixvm_codegen::aiur_ixvm_runner::execute_ixvm, + ) }; - // `use_bytecode` selects the generic Aiur bytecode interpreter over the - // codegen'd IxVM kernel (same toggle as - // `rs_aiur_toplevel_check_addr_with_env`). - let (_aiur_claim_arr, proof) = if use_bytecode { - aiur_system_obj.get().prove_ixvm( - fun_idx, - &input, - &mut io_buffer, - |toplevel, fun_idx, input, io_buffer| { - toplevel.execute(fun_idx, input, io_buffer) - }, - ) - } else { - aiur_system_obj.get().prove_ixvm( - fun_idx, - &input, - &mut io_buffer, - ixvm_codegen::aiur_ixvm_runner::execute_ixvm, - ) - }; - - LeanExcept::ok(build_prove_env_result(&claim, proof, &io_buffer)) + LeanExcept::ok(build_prove_env_result(&claim, proof, &io_buffer)) + }) } /// `AiurSystem.shardProveWithEnv`: per-shard prove against a @@ -903,31 +908,33 @@ extern "C" fn rs_aiur_system_shard_prove_with_env( >, owned_blob: LeanByteArray>, ) -> LeanExcept { - let fun_idx = lean_unbox_nat_as_usize(fun_idx.inner()); - let owned = match decode_owned_blob(&owned_blob) { - Ok(v) => v, - Err(e) => return LeanExcept::error_string(&e), - }; - let env = &env_handle.get().env; - - let (claim, input, mut io_buffer) = - match ixvm_codegen::aiur_ixvm_witness::build_shard_check_env_witness( - env, &owned, - ) { - Ok(t) => t, - Err(e) => { - return LeanExcept::error_string(&format!("witness build: {e}")); - }, + ffi_catch_unwind_except("AiurSystem.shardProveWithEnv", || { + let fun_idx = lean_unbox_nat_as_usize(fun_idx.inner()); + let owned = match decode_owned_blob(&owned_blob) { + Ok(v) => v, + Err(e) => return LeanExcept::error_string(&e), }; + let env = &env_handle.get().env; + + let (claim, input, mut io_buffer) = + match ixvm_codegen::aiur_ixvm_witness::build_shard_check_env_witness( + env, &owned, + ) { + Ok(t) => t, + Err(e) => { + return LeanExcept::error_string(&format!("witness build: {e}")); + }, + }; + + let (_aiur_claim_arr, proof) = aiur_system_obj.get().prove_ixvm( + fun_idx, + &input, + &mut io_buffer, + ixvm_codegen::aiur_ixvm_runner::execute_ixvm, + ); - let (_aiur_claim_arr, proof) = aiur_system_obj.get().prove_ixvm( - fun_idx, - &input, - &mut io_buffer, - ixvm_codegen::aiur_ixvm_runner::execute_ixvm, - ); - - LeanExcept::ok(build_prove_env_result(&claim, proof, &io_buffer)) + LeanExcept::ok(build_prove_env_result(&claim, proof, &io_buffer)) + }) } /// `AiurSystem.proveIxVM`: IxVM-native prove path. Same return shape @@ -942,19 +949,21 @@ extern "C" fn rs_aiur_system_prove_ixvm( args: LeanArray>, io_data_arr: LeanArray>, io_map_arr: LeanArray>, -) -> LeanAiurProveResult { - let fun_idx = lean_unbox_nat_as_usize(fun_idx.inner()); - let args = args.map(|x| lean_unbox_g(&x)); - let mut io_buffer = decode_io_buffer(&io_data_arr, &io_map_arr); +) -> LeanExcept { + ffi_catch_unwind("AiurSystem.proveIxVM", || { + let fun_idx = lean_unbox_nat_as_usize(fun_idx.inner()); + let args = args.map(|x| lean_unbox_g(&x)); + let mut io_buffer = decode_io_buffer(&io_data_arr, &io_map_arr); - let (claim, proof) = aiur_system_obj.get().prove_ixvm( - fun_idx, - &args, - &mut io_buffer, - ixvm_codegen::aiur_ixvm_runner::execute_ixvm, - ); + let (claim, proof) = aiur_system_obj.get().prove_ixvm( + fun_idx, + &args, + &mut io_buffer, + ixvm_codegen::aiur_ixvm_runner::execute_ixvm, + ); - build_prove_result(&claim, proof, &io_buffer) + build_prove_result(&claim, proof, &io_buffer).into() + }) } /// `Bytecode.Toplevel.executeMultiStark`: run the MultiStark recursive @@ -1022,38 +1031,76 @@ extern "C" fn rs_aiur_multi_stark_prove( vk_bytes: LeanByteArray>, claims_bytes: LeanByteArray>, use_bytecode: bool, -) -> LeanOwned { - let fun_idx = lean_unbox_nat_as_usize(fun_idx.inner()); - let mut io_buffer = ixvm_codegen::aiur_multi_stark_runner::verifier_io_buffer( - proof_bytes.as_bytes(), - vk_bytes.as_bytes(), - claims_bytes.as_bytes(), - ); - let args = pub_input.map(|x| lean_unbox_g(&x)); - - let system = aiur_system_obj.get(); - let (claim, proof) = if use_bytecode { - system.prove(fun_idx, &args, &mut io_buffer) - } else { - system.prove_ixvm( - fun_idx, - &args, - &mut io_buffer, - ixvm_codegen::aiur_multi_stark_runner::execute_multi_stark, - ) - }; +) -> LeanExcept { + ffi_catch_unwind("AiurSystem.proveMultiStark", || { + let fun_idx = lean_unbox_nat_as_usize(fun_idx.inner()); + let mut io_buffer = + ixvm_codegen::aiur_multi_stark_runner::verifier_io_buffer( + proof_bytes.as_bytes(), + vk_bytes.as_bytes(), + claims_bytes.as_bytes(), + ); + let args = pub_input.map(|x| lean_unbox_g(&x)); + + let system = aiur_system_obj.get(); + let (claim, proof) = if use_bytecode { + system.prove(fun_idx, &args, &mut io_buffer) + } else { + system.prove_ixvm( + fun_idx, + &args, + &mut io_buffer, + ixvm_codegen::aiur_multi_stark_runner::execute_multi_stark, + ) + }; - let lean_proof: LeanOwned = - LeanExternal::alloc(&AIUR_PROOF_CLASS, proof).into(); - // Array G × Proof - let result = LeanProd::new(build_g_array(&claim), lean_proof); - result.into() + let lean_proof: LeanOwned = + LeanExternal::alloc(&AIUR_PROOF_CLASS, proof).into(); + // Array G × Proof + LeanProd::new(build_g_array(&claim), lean_proof).into() + }) } // ============================================================================= // Helpers // ============================================================================= +/// Prevent Rust panics (including CUDA runtime failures) from unwinding across +/// the Lean C ABI. Lean receives the failure as an ordinary `Except.error`. +fn ffi_catch_unwind( + context: &str, + f: impl FnOnce() -> LeanOwned, +) -> LeanExcept { + match std::panic::catch_unwind(std::panic::AssertUnwindSafe(f)) { + Ok(value) => LeanExcept::ok(value), + Err(payload) => { + let message = payload + .downcast_ref::<&str>() + .copied() + .or_else(|| payload.downcast_ref::().map(String::as_str)) + .unwrap_or("unknown Rust panic"); + LeanExcept::error_string(&format!("{context}: {message}")) + }, + } +} + +fn ffi_catch_unwind_except( + context: &str, + f: impl FnOnce() -> LeanExcept, +) -> LeanExcept { + match std::panic::catch_unwind(std::panic::AssertUnwindSafe(f)) { + Ok(result) => result, + Err(payload) => { + let message = payload + .downcast_ref::<&str>() + .copied() + .or_else(|| payload.downcast_ref::().map(String::as_str)) + .unwrap_or("unknown Rust panic"); + LeanExcept::error_string(&format!("{context}: {message}")) + }, + } +} + /// Build a Lean `Array G` from a slice of field elements. fn build_g_array(values: &[G]) -> LeanArray { let arr = LeanArray::alloc(values.len()); diff --git a/flake.nix b/flake.nix index ce74d895f..546e2f6df 100644 --- a/flake.nix +++ b/flake.nix @@ -103,14 +103,14 @@ pkgs.libiconv ]; }; - # Build dependencies once with every feature enabled so the `net` - # stack (tokio/iroh) is compiled and cached here, then shared by the - # package builds and the all-features clippy check instead of being - # rebuilt per consumer. + # Build dependencies once with every host feature enabled so the + # `net` stack (tokio/iroh) is compiled and cached here, then shared + # by the package builds and clippy. CUDA remains opt-in and is + # compiled separately in CI with the CUDA toolkit available. cargoArtifacts = craneLib.buildDepsOnly ( craneArgs // { - cargoExtraArgs = "--locked --all-features"; + cargoExtraArgs = "--locked --features parallel,test-ffi,net"; } ); @@ -151,15 +151,21 @@ # Lake package lake2nix = pkgs.callPackage lean4-nix.lake { inherit lean; }; - # Restrict the Lake build inputs to Lean-relevant files so edits to - # unrelated files (flake.nix, CI, docs) don't invalidate the whole - # Lean build. The Rust side gets the same via cleanCargoSource. + # Restrict the Lake build inputs to files traced by the Lean and + # Rust archive targets. Cargo itself is still handled by Crane, but + # Lake needs the Rust sources and manifests to calculate the archive + # dependency trace before copying the prebuilt static library. leanSrc = pkgs.lib.fileset.toSource { root = ./.; fileset = pkgs.lib.fileset.unions [ ./lakefile.lean ./lake-manifest.json ./lean-toolchain + ./Cargo.toml + ./Cargo.lock + (pkgs.lib.fileset.fileFilter + (f: f.hasExt "rs" || f.hasExt "toml") + ./crates) (pkgs.lib.fileset.fileFilter (f: f.hasExt "lean") ./.) ]; }; @@ -297,7 +303,7 @@ craneArgs // { inherit cargoArtifacts; - cargoExtraArgs = "--locked --all-features"; + cargoExtraArgs = "--locked --features parallel,test-ffi,net"; cargoClippyExtraArgs = "--all-targets -- -D warnings"; } ); diff --git a/lakefile.lean b/lakefile.lean index 409361693..3f585ebb5 100644 --- a/lakefile.lean +++ b/lakefile.lean @@ -41,50 +41,83 @@ The Rust static libraries use `target` + `moreLinkObjs` instead of `extern_lib` - `ix` uses `ix_rs_net` (`parallel,net`) for networking support (iroh). - `IxTests` uses `ix_rs_test` (`parallel,test-ffi`) for test-only FFI code. -- Everything else inherits `ix_rs` (`parallel` only) from the `Ix` `lean_lib`. +- Everything else inherits `ix_rs` (`parallel`, plus opt-in `cuda`) from the + `Ix` `lean_lib`. The `ix_rs_test` and `ix_rs_net` targets fetch `ix_rs` first to guarantee ordering -before overwriting the lib, since they write to the same lib path. The second cargo build is incremental — only the feature-affected crates recompile. +before Cargo overwrites its release archive, then snapshot distinct Lake artifacts. +The second Cargo build is incremental — only feature-affected crates recompile. `extern_lib` only runs at link time, so `lake build` on a `lean_lib` alone wouldn't trigger the Cargo build. With `target` + `moreLinkObjs`, the Rust static lib is built during module compilation on the default `Ix` lib, allowing Lake to conditional compile the Rust lib per build target. -/ section FFI -/-- Build args for `cargo build --release` with feature flags from env vars. +/-- Build args for `cargo build --release` with opt-in feature overrides. Cargo output is visible with `lake -v build`. -/ def cargoArgs (testFfi : Bool := false) (net : Bool := false) : IO (Array String) := do - -- IX_NO_PAR=1 disables parallel + -- IX_NO_PAR=1 disables parallel; IX_CUDA=1/true/yes enables CUDA. let ixNoPar ← IO.getEnv "IX_NO_PAR" + let ixCuda ← IO.getEnv "IX_CUDA" let mut features : Array String := #[] if ixNoPar != some "1" then features := features.push "parallel" + if ixCuda == some "1" || ixCuda == some "true" || ixCuda == some "yes" then + features := features.push "cuda" if net && !System.Platform.isOSX then features := features.push "net" if testFfi then features := features.push "test-ffi" + IO.println s!"Ix Rust features: {if features.isEmpty then "none" else ",".intercalate features.toList}" let buildArgs := #["build", "--release", "-p", "ix-ffi"] if features.isEmpty then return buildArgs else return buildArgs ++ #["--features", ",".intercalate features.toList] +/-- Build and snapshot one feature selection of the Rust static library. +The copied output has a Lake trace containing both Rust sources and Cargo +arguments, so changing `IX_CUDA` cannot silently reuse a differently-featured +archive from a previous invocation. -/ +def buildRustStatic (pkg : Package) (args : Array String) (tag : String) : + SpawnM (Job FilePath) := do + let sources ← inputDir (pkg.dir / "crates") true fun path => + path.extension == some "rs" || path.fileName == "Cargo.toml" + let manifests := Job.collectArray #[ + ← inputTextFile (pkg.dir / "Cargo.toml"), + ← inputTextFile (pkg.dir / "Cargo.lock") + ] + let deps := sources.zipWith (fun sourceFiles manifestFiles => + (sourceFiles, manifestFiles)) manifests + let output := pkg.buildDir / "lib" / s!"libix_ffi_{tag}.a" + buildFileAfterDep output deps (fun _ => do + proc { cmd := "cargo", args, cwd := pkg.dir } (quiet := true) + let built := pkg.dir / "target" / "release" / nameToStaticLib "ix_ffi" + copyFile built output + ) (extraDepTrace := pure <| .ofHash (pureHash args) s!"cargo args: {args}") + /-- Build the Rust static lib with default features (`parallel`). -/ target ix_rs pkg : FilePath := do - let args ← cargoArgs - proc { cmd := "cargo", args, cwd := pkg.dir } (quiet := true) - inputBinFile $ pkg.dir / "target" / "release" / nameToStaticLib "ix_ffi" + buildRustStatic pkg (← cargoArgs) "default" /-- Rebuild the Rust static lib with `test-ffi`. Only triggered by `lake test` (via `moreLinkObjs` on `IxTests`). Fetches `ix_rs` first to guarantee ordering before overwriting the lib. -/ target ix_rs_test pkg : FilePath := do - let _ ← ix_rs.fetch - let args ← cargoArgs (testFfi := true) - proc { cmd := "cargo", args, cwd := pkg.dir } (quiet := true) - inputBinFile $ pkg.dir / "target" / "release" / nameToStaticLib "ix_ffi" + let base ← ix_rs.fetch + base.mapM fun _ => do + let args ← cargoArgs (testFfi := true) + proc { cmd := "cargo", args, cwd := pkg.dir } (quiet := true) + let built := pkg.dir / "target" / "release" / nameToStaticLib "ix_ffi" + let output := pkg.buildDir / "lib" / "libix_ffi_test.a" + copyFile built output + return output /-- Build the Rust static lib with `net` for the `ix` CLI. Fetches `ix_rs` first to guarantee ordering before overwriting the lib. -/ target ix_rs_net pkg : FilePath := do - let _ ← ix_rs.fetch - let args ← cargoArgs (net := true) - proc { cmd := "cargo", args, cwd := pkg.dir } (quiet := true) - inputBinFile $ pkg.dir / "target" / "release" / nameToStaticLib "ix_ffi" + let base ← ix_rs.fetch + base.mapM fun _ => do + let args ← cargoArgs (net := true) + proc { cmd := "cargo", args, cwd := pkg.dir } (quiet := true) + let built := pkg.dir / "target" / "release" / nameToStaticLib "ix_ffi" + let output := pkg.buildDir / "lib" / "libix_ffi_net.a" + copyFile built output + return output /-- The `ix-ffi-dyn` cdylib: Ix's own raw `@[extern]` symbols (currently the `toLEBytes` operations) as a small standalone shared library. Consumed by