feat: Multi-GPU E2E proving pipeline

.github/workflows/ci.yml

@@ -60,6 +60,11 @@ jobs:
       - name: Check codegen'd IxVM kernel is up to date
         if: github.event_name != 'push'
         run: lake exe ix codegen --check
+      # A stale generated trace module is otherwise caught only when the
+      # generated GPU provider registers at proving time.
+      - name: Check generated trace writers, CUDA units and manifest are up to date
+        if: github.event_name != 'push'
+        run: lake exe ix codegen --trace-bundle --check
       - name: Check Lean versions match for Ix and compiler bench
         if: github.event_name != 'push'
         run: diff lean-toolchain Benchmarks/Compile/lean-toolchain

Cargo.toml

@@ -53,7 +53,7 @@ log = "0.4"
 memmap2 = "0.9"
 # Bundle v3.5.1 to fix large-arena failures (microsoft/mimalloc#1309).
 mimalloc = { git = "https://github.com/argumentcomputer/mimalloc_rust.git", rev = "616cc1f130d9ea2a1111a46b4e98c140e078f882", default-features = false }
-multi-stark = { git = "https://github.com/argumentcomputer/multi-stark.git", rev = "a155ee47687e3a9b7e94f158eeff9d710a8262d5" }
+multi-stark = { git = "https://github.com/argumentcomputer/multi-stark.git", rev = "acfc370a2e74affcb7e8ee218fbbaf80b9da1eef" }
 nom = "7.1.3"
 num-bigint = "0.4.6"
 quickcheck = "1.0.3"

Ix/Aggr/FFI.lean

@@ -97,6 +97,33 @@ opaque aggregateExpected (envHandle : @& EnvHandle)
     Except String AggregateExpected
 
 
+
+/-- The Stage 2 plan cut for `ix prove --lanes`: the tree is split from the
+root while any frontier node holds more than `subtreeSize` leaves (a lone
+leaf may then stand as a subtree of its own). Rows in post-order:
+`#[0, slot, ids…]` for a subtree (its root slot, as `ix aggregate --subtree`
+numbers slots, then the manifest ids of the shards under it, sorted) and
+`#[1, slot, left, right]` for each join above the frontier, the root
+last. -/
+@[extern "rs_aiur_aggregate_subtree_plan"]
+opaque aggregateSubtreePlan (envHandle : @& EnvHandle) (manifestPath : @& String)
+  (structuralAbove subtreeSize : @& Nat) : Except String (Array (Array Nat))
+
+/-- `ix prove --lanes`: the whole multi-device run in this process, one
+resident prover per GPU built from these two systems (`AiurSystem.on_device`
+in Rust; nothing is copied across the FFI), every claim and join scheduled
+over them as their inputs become ready, the root wrapped, and every claim
+proof verified natively at the end. `maxRamBytes` is one worker's host
+budget (0 detects); `execJobs` the claim executions ahead of each prover.
+Oversized environment claims are bisected and checkpointed; completed
+proofs are reused. A nonempty `outManifest` receives the final partition.
+Returns the verified root address; progress is printed on stderr. -/
+@[extern "rs_aiur_prove_lanes"]
+opaque proveLanes (ixvmSystem aggrSystem : @& AiurSystem)
+  (envHandle : @& EnvHandle) (manifestPath : @& String)
+  (verifyIdx aggrIdx lanes maxRamBytes execJobs structuralAbove : @& Nat)
+  (cacheFriBytes : @& ByteArray) (outManifest : @& String) : Except String String
+
 end AiurSystem
 end Aiur
 end

Ix/Aiur/Codegen/Rust.lean

@@ -35,6 +35,7 @@ inductive MatchPat where
   | litU64 (n : Nat)
   | wildcard
   | tuple (elems : Array MatchPat)
+  | array (elems : Array MatchPat)
   | constructor (path : Array String) (args : Array MatchPat)
   deriving Inhabited, BEq
 
@@ -73,6 +74,7 @@ inductive RustExpr where
 inductive RustStmt where
   | letStmt (isMut : Bool) (name : String) (ty : Option RustType) (val : RustExpr)
   | letPattern (pat : MatchPat) (ty : Option RustType) (val : RustExpr)
+  | assign (target : RustExpr) (val : RustExpr)
   | exprStmt (e : RustExpr)
   | returnStmt (e : RustExpr)
   | ifStmt (cond : RustExpr) (thenStmts : Array RustStmt)
@@ -143,6 +145,7 @@ partial def MatchPat.toStr : MatchPat → String
   | .litU64 n => s!"{n}u64"
   | .wildcard => "_"
   | .tuple ps => tupleText (ps.map MatchPat.toStr)
+  | .array ps => "[" ++ commaSep (ps.map MatchPat.toStr) ++ "]"
   | .constructor path ps => "::".intercalate path.toList ++
       "(" ++ commaSep (ps.map MatchPat.toStr) ++ ")"
 
@@ -167,11 +170,11 @@ partial def RustExpr.toStr : RustExpr → String
     let callee := match f with
       | .field .. => "(" ++ f.toStr ++ ")"
       | _ => f.postfixBase
-    callee ++ "(" ++ commaSep (xs.map (·.toStr)) ++ ")"
+    callee ++ "(" ++ commaSep (xs.map (·.toStrBare)) ++ ")"
   | .methodCall e n xs => e.postfixBase ++ "." ++ n ++
-      "(" ++ commaSep (xs.map (·.toStr)) ++ ")"
+      "(" ++ commaSep (xs.map (·.toStrBare)) ++ ")"
   | .constGeneric f xs => f.postfixBase ++ "::<" ++ commaSep (xs.map (·.toStr)) ++ ">"
-  | .index a i => a.postfixBase ++ "[" ++ i.toStr ++ "]"
+  | .index a i => a.postfixBase ++ "[" ++ i.toStrBare ++ "]"
   | .field e n => e.postfixBase ++ "." ++ n
   | .binop op a b => "(" ++ a.toStr ++ " " ++ op.toStr ++ " " ++ b.toStr ++ ")"
   -- Parenthesize prefix/cast expressions so postfix composition preserves precedence.
@@ -182,12 +185,12 @@ partial def RustExpr.toStr : RustExpr → String
   | .tryExpr e => e.postfixBase ++ "?"
   | .range a b => (a.map (·.toStr) |>.getD "") ++ ".." ++
       (b.map (·.toStr) |>.getD "")
-  | .macroCall n xs => n ++ "!(" ++ commaSep (xs.map (·.toStr)) ++ ")"
-  | .arrayLit xs => "[" ++ commaSep (xs.map (·.toStr)) ++ "]"
-  | .arrayRepeat e n => "[" ++ e.toStr ++ "; " ++ n.toStr ++ "]"
-  | .tuple xs => tupleText (xs.map (·.toStr))
+  | .macroCall n xs => n ++ "!(" ++ commaSep (xs.map (·.toStrBare)) ++ ")"
+  | .arrayLit xs => "[" ++ commaSep (xs.map (·.toStrBare)) ++ "]"
+  | .arrayRepeat e n => "[" ++ e.toStrBare ++ "; " ++ n.toStr ++ "]"
+  | .tuple xs => tupleText (xs.map (·.toStrBare))
   | .structLit p fs => "::".intercalate p.toList ++ " { " ++
-      commaSep (fs.map fun (n, e) => n ++ ": " ++ e.toStr) ++ " }"
+      commaSep (fs.map fun (n, e) => n ++ ": " ++ e.toStrBare) ++ " }"
   | .block b => b.toStr
   | .unsafeBlock b => "unsafe " ++ b.toStr
   | .labeledBlock label b => "'" ++ label ++ ": " ++ b.toStr
@@ -202,13 +205,24 @@ partial def RustExpr.postfixBase (e : RustExpr) : String :=
   | .closure .. | .ifExpr .. | .matchExpr .. | .range .. => "(" ++ e.toStr ++ ")"
   | _ => e.toStr
 
+/-- An expression in a position the grammar already delimits (an argument,
+    an element, an index, an initializer), where its own parentheses would
+    be redundant and trip the lints. -/
+partial def RustExpr.toStrBare : RustExpr → String
+  | .binop op a b => a.toStr ++ " " ++ op.toStr ++ " " ++ b.toStr
+  | .deref e => "*" ++ e.toStr
+  | .ref e => "&" ++ e.toStr
+  | .not e => "!" ++ e.toStr
+  | .cast e t => e.toStr ++ " as " ++ t.toStr
+  | e => e.toStr
+
 /-- Generated source stays compact: whitespace separates tokens, but nested
     syntax never expands into additional lines. Each top-level item gets one
     line. String contents are escaped before rendering, not whitespace-stripped. -/
 partial def RustBlock.toStr (b : RustBlock) : String :=
   "{ " ++ stmtsToStr b.stmts ++
     (if b.stmts.isEmpty then "" else " ") ++
-    (b.tail.map RustExpr.toStr |>.getD "") ++ " }"
+    (b.tail.map RustExpr.toStrBare |>.getD "") ++ " }"
 
 partial def renderMatch (scrut : RustExpr) (arms : Array MatchArm) : String :=
   "match " ++ scrut.toStr ++ " { " ++
@@ -219,19 +233,21 @@ partial def renderMatch (scrut : RustExpr) (arms : Array MatchArm) : String :=
 
 partial def RustStmt.toStr : RustStmt → String
   | .letStmt isMut n ty e => "let " ++ (if isMut then "mut " else "") ++ n ++
-      (ty.map (fun t => ": " ++ t.toStr) |>.getD "") ++ " = " ++ e.toStr ++ ";"
+      (ty.map (fun t => ": " ++ t.toStr) |>.getD "") ++ " = " ++ e.toStrBare ++ ";"
   | .letPattern p ty e => "let " ++ p.toStr ++
-      (ty.map (fun t => ": " ++ t.toStr) |>.getD "") ++ " = " ++ e.toStr ++ ";"
-  | .exprStmt e => e.toStr ++ ";"
-  | .returnStmt e => "return " ++ e.toStr ++ ";"
+      (ty.map (fun t => ": " ++ t.toStr) |>.getD "") ++ " = " ++ e.toStrBare ++ ";"
+  | .assign t e => t.toStr ++ " = " ++ e.toStrBare ++ ";"
+  | .exprStmt e => e.toStrBare ++ ";"
+  | .returnStmt (.tuple #[]) => "return;"
+  | .returnStmt e => "return " ++ e.toStrBare ++ ";"
   | .ifStmt c t e => "if " ++ c.toStr ++ " " ++
       RustBlock.toStr { stmts := t } ++
       (e.map (fun ss => " else " ++ RustBlock.toStr { stmts := ss }) |>.getD "")
   | .matchStmt e arms => renderMatch e arms
   | .block ss => RustBlock.toStr { stmts := ss }
   | .forStmt p e ss => "for " ++ p.toStr ++ " in " ++ e.toStr ++
       " " ++ RustBlock.toStr { stmts := ss }
-  | .breakWith label e => "break '" ++ label ++ " " ++ e.toStr ++ ";"
+  | .breakWith label e => "break '" ++ label ++ " " ++ e.toStrBare ++ ";"
 
 partial def stmtsToStr (ss : Array RustStmt) : String :=
   " ".intercalate (ss.toList.map RustStmt.toStr)
@@ -243,10 +259,12 @@ def RustItem.toStr : RustItem → String
     let vis := match f.visibility with | .internal => "" | .crate => "pub(crate) "
     let generics := if f.constParams.isEmpty then "" else
       "<" ++ commaSep (f.constParams.map fun (n, t) => "const " ++ n ++ ": " ++ t.toStr) ++ ">"
+    -- A unit return type is left implicit, as the lints prefer.
+    let returns := if f.returnTy == .tuple #[] then "" else "-> " ++ f.returnTy.toStr ++ " "
     vis ++ "fn " ++ f.name ++ generics ++ "(" ++
       commaSep (f.params.map fun (n, t) => n ++ ": " ++ t.toStr) ++
       (if f.params.isEmpty then "" else ",") ++
-      ") -> " ++ f.returnTy.toStr ++ " " ++ f.body.toStr ++ "\n"
+      ") " ++ returns ++ f.body.toStr ++ "\n"
 
 end Aiur.Codegen
 end