From 1bae8294f5d29517279df640f79fd25a036c4428 Mon Sep 17 00:00:00 2001 From: samuelburnham <45365069+samuelburnham@users.noreply.github.com> Date: Thu, 10 Sep 2026 03:02:19 +0000 Subject: [PATCH 1/4] NUMA lanes: single-process Stage 1, pinned Stage 2 scheduler, lane pipelines crates/ffi/src/numa.rs: topology from sysfs and sched_getaffinity, per-thread pinning (sched_setaffinity + set_mempolicy MPOL_BIND), one Rayon pool per domain, cgroup memory limit, resident bytes per node. Knobs IX_NUMA, IX_NUMA_POLICY, IX_NUMA_THREADS, IX_NUMA_PACK. Stage 1: `ix prove --lookahead` runs one execute-next-while-proving pipeline per NUMA domain inside a single process (leaves split by measured peak, IX_PROVE_LANES), sharing the environment and proving systems. Stage 2: every join is placed on a domain under a per-lane budget (90% of the node, capped by --max-ram and the cgroup limit), at most two per node, solo tail unpinned. Dependency-free joins run on per-lane queues that prepare the next slot's execution record while the current one proves; with packing on a lane runs two such queues when both fit. Static RAM weights per join shape: direct/mixed 180 GiB (a direct join is ~200 GiB resident after main's function groups), lifts/structural 195 GiB + 1.25 MiB per subject with a 390 GiB floor above 65,536 subjects. Parallel proof import; every slot logs its node's resident peak; `ix aggregate --texray`. IX_AGGREGATE_SHARDS=a-b,c aggregates one subtree of a manifest from an existing run's leaf proofs (experiments; leaf claims do not depend on the manifest size). Rebased on main's function groups (#619/#620): the verify command's backend construction compiles with the IxVM and ix_aggr groupings. --- Cargo.lock | 1 + Ix/Aiur/Protocol.lean | 12 + Ix/Cli/AggregateCmd.lean | 56 +- Ix/Cli/ProveCmd.lean | 72 +- Ix/Cli/ShardProofIndex.lean | 49 +- Ix/Cli/VerifyCmd.lean | 55 +- Tests/Aggr.lean | 3 +- Tests/AggrSemantics.lean | 31 + Tests/Main.lean | 2 + Tests/ShardPipeline.lean | 127 ++ crates/aiur/src/synthesis.rs | 5 + crates/ffi/Cargo.toml | 1 + crates/ffi/src/aiur/aggregate.rs | 1563 +++++++++++++++-- .../ffi/src/aiur/aggregate/shard_pipeline.rs | 1082 ++++++++++++ crates/ffi/src/lib.rs | 1 + crates/ffi/src/numa.rs | 369 ++++ 16 files changed, 3216 insertions(+), 213 deletions(-) create mode 100644 Tests/ShardPipeline.lean create mode 100644 crates/ffi/src/aiur/aggregate/shard_pipeline.rs create mode 100644 crates/ffi/src/numa.rs diff --git a/Cargo.lock b/Cargo.lock index c8c03e78f..dff5501b5 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1823,6 +1823,7 @@ dependencies = [ "ixon", "ixvm-codegen", "lean-ffi", + "libc", "memmap2", "mimalloc", "multi-stark", diff --git a/Ix/Aiur/Protocol.lean b/Ix/Aiur/Protocol.lean index cbb7ae891..9fa21e7ce 100644 --- a/Ix/Aiur/Protocol.lean +++ b/Ix/Aiur/Protocol.lean @@ -311,6 +311,18 @@ opaque proofToAdviceBytes : @& AiurSystem → end AiurSystem +/-- Native manifest-leaf proving with private split healing. `blocks` and +`owned` use the counted address-list format; `ids` are newline-separated. +`maxRam` is a positive byte count limiting the predicted proving peak of +each executed query record. Oversized records are dropped and split. +Only original leaf proofs are printed. An empty index path disables indexing. +Every split is rejoined canonically and checked against the original claim. -/ +@[extern "rs_aiur_shard_pipeline"] +opaque shardPipeline : @& AiurSystem → @& AiurSystem → @& EnvHandle → + @& ByteArray → @& ByteArray → @& String → @& Nat → @& Nat → + @& Nat → @& String → @& String → @& String → + Bool → Bool → Bool → IO (Except String String) + /-- Write a `.ixes` manifest for an EXPLICIT partition — the block lists a run actually produced (splits included) rather than a planner's output. `shardsBlob`: per shard, a 4-byte LE block count followed by diff --git a/Ix/Cli/AggregateCmd.lean b/Ix/Cli/AggregateCmd.lean index 1b19531e5..e3bdbcdda 100644 --- a/Ix/Cli/AggregateCmd.lean +++ b/Ix/Cli/AggregateCmd.lean @@ -16,6 +16,7 @@ -/ module import Std.Sync +import Ix.TracingTexray public import Cli public import Ix.Aggr public import Ix.Cli.CheckCmd @@ -101,8 +102,7 @@ def aggregateLiftRamBytes : Nat := 195 * aggregateGiB benchmark/test union is migrated in M1-e/M1-f. -/ def aggregateWrapRamBytes : Nat := aggregateLiftRamBytes -/-- Structural joins are dominated by the same two recursive-proof checks as -lifts. Keep the conservative lift reserve until the real E2E calibration. -/ +/-- Base reserve for self-pairs, before the flat or structural subject term. -/ def aggregateStructuralJoinRamBytes : Nat := aggregateLiftRamBytes /-- Native verification and serialization of a raw shard proof in direct mode @@ -110,10 +110,11 @@ is charged to its consuming pair. -/ def aggregateRawShardRamBytes : Nat := 4 * aggregateGiB /-- Measured upper envelope for an `IxVM + IxVM` pair (shapes 2/6). -/ -def aggregateDirectJoinRamBytes : Nat := 390 * aggregateGiB +def aggregateDirectJoinRamBytes : Nat := 180 * aggregateGiB -/-- Measured upper envelope for a mixed recursive/IxVM pair (shapes 3/4/7/8). -/ -def aggregateMixedJoinRamBytes : Nat := 340 * aggregateGiB +/-- Mixed recursive/IxVM pairs (shapes 3/4/7/8) predicted up to 385 GiB in +the 2026-09-09 Mathlib run, exceeding the previous 340 GiB reserve. -/ +def aggregateMixedJoinRamBytes : Nat := 180 * aggregateGiB /-- Flat joins add canonical subject-tree work to the recursive-proof base. One MiB per subject is a deliberately conservative placeholder: at Init's @@ -121,9 +122,19 @@ One MiB per subject is a deliberately conservative placeholder: at Init's default structural threshold caps this term near 4 GiB in production. -/ def aggregateFlatJoinRamPerSubjectBytes : Nat := 1024 * 1024 -/-- Per-shape RAM weight used by the Lean admission gate. Shape 5 retains the -flat subject-count reserve; shape 9 is the O(1)-subject structural arm. The -direct/mixed values are conservative round-ups of the §3.4 measurements. -/ +/-- Mathlib structural self-joins need a subject-dependent reserve despite +their O(1) subject-root fold: assumption/path checks and child verification +can grow. Use a subject term and a doubled base above 64k subjects to cover +trace-size steps (380.5 GiB predicted at 91068 subjects). See +`exp/design/numa-slot-pinning.md` §12. Keep these constants in sync with +Rust's `STRUCTURAL_RAM_PER_SUBJECT` and `STRUCTURAL_LARGE_SUBJECTS`. -/ +def aggregateStructuralJoinRamPerSubjectBytes : Nat := 5 * 1024 * 1024 / 4 + +def aggregateStructuralLargeSubjects : Nat := 64 * 1024 + +/-- Per-shape RAM weight used by the Lean admission gate. Both self-pair +shapes (5/9) reserve subject-dependent work. The direct/mixed values are +conservative round-ups of the §3.4 measurements. -/ def aggregateShapeRamBytes (shape subjectCount : Nat) : Nat := match shape with | 0 | 1 => aggregateWrapRamBytes @@ -131,7 +142,12 @@ def aggregateShapeRamBytes (shape subjectCount : Nat) : Nat := | 3 | 4 | 7 | 8 => aggregateMixedJoinRamBytes | 5 => aggregateStructuralJoinRamBytes + subjectCount * aggregateFlatJoinRamPerSubjectBytes - | 9 => aggregateStructuralJoinRamBytes + | 9 => + let weight := aggregateStructuralJoinRamBytes + + subjectCount * aggregateStructuralJoinRamPerSubjectBytes + if subjectCount > aggregateStructuralLargeSubjects then + max weight (2 * aggregateStructuralJoinRamBytes) + else weight | _ => aggregateDirectJoinRamBytes /-- Calibration-pending per-slot RAM weight used by the Lean admission gate. @@ -143,7 +159,7 @@ def aggregateSlotRamBytes (item : ScheduledFold) : Nat := | .leaf _ => if item.kind == .ixvm then aggregateRawShardRamBytes else aggregateWrapRamBytes | .join _ _ => - if item.structural then aggregateStructuralJoinRamBytes + if item.structural then aggregateShapeRamBytes 9 item.subjectCount else aggregateStructuralJoinRamBytes + item.subjectCount * aggregateFlatJoinRamPerSubjectBytes @@ -972,6 +988,9 @@ private def runAggregateCmdNativeWith | .ok backend => pure backend let verifyIdx := ixvmBackend.compiled.getFuncIdx `verify_claim |>.get! let aggrIdx := aggrBackend.compiled.getFuncIdx `ix_aggr |>.get! + -- Streamed `[texray]` per-span lines (execute / witness / STARK) for every + -- Stage 2 slot, as `ix prove --texray` does for shards. + if p.hasFlag "texray" then TracingTexray.init {} let nativeResult ← IO.lazyPure fun _ => ixvmBackend.system.aggregateStage2 aggrBackend.system envHandle manifestPath proofHexes verifyIdx aggrIdx jobs ramBudgetBytes @@ -1017,6 +1036,9 @@ private def runAggregateCmdLeanReferenceWith let structuralAbove := ((p.flag? "structural-above").map (·.as! Nat)).getD defaultStructuralAbove let directJoins := p.hasFlag "direct-joins" + -- Same streamed `[texray]` per-span lines as `ix prove --texray`: the + -- execute / witness / STARK split of every Stage 2 slot. + if p.hasFlag "texray" then TracingTexray.init {} let plan ← match schedulePlan view.aggregationTree.foldPlan shardCounts structuralAbove directJoins with | .error e => IO.eprintln e; return 1 @@ -1035,13 +1057,20 @@ private def runAggregateCmdLeanReferenceWith let budgetSource := if maxRamGb?.isSome then "--max-ram" else "92% MemTotal" IO.println s!"[aggregate] scheduler: jobs={jobsLabel}, RAM budget \ {formatAggregateGiB ramBudgetBytes} GiB ({budgetSource}); \ - wrap/self reserve {formatAggregateGiB aggregateWrapRamBytes} GiB, \ + wrap/self base {formatAggregateGiB aggregateWrapRamBytes} GiB, \ direct {formatAggregateGiB aggregateDirectJoinRamBytes} GiB, mixed \ - {formatAggregateGiB aggregateMixedJoinRamBytes} GiB, flat +1 MiB/subject" + {formatAggregateGiB aggregateMixedJoinRamBytes} GiB, flat self +1 MiB/subject, \ + structural self +1.25 MiB/subject (minimum \ + {formatAggregateGiB (2 * aggregateStructuralJoinRamBytes)} GiB above \ + {aggregateStructuralLargeSubjects} subjects)" if p.hasFlag "plan-only" then return 0 let proofHexes := (p.variableArgsAs! String).toList - if proofHexes.length != view.shards.size then + -- `IX_AGGREGATE_SHARDS` (experiment knob, see `shard_selection` in + -- aggregate.rs) aggregates a subtree; the native side then requires one + -- proof per selected shard and ignores the rest. + let partialRun := (← IO.getEnv "IX_AGGREGATE_SHARDS").isSome + if !partialRun && proofHexes.length != view.shards.size then IO.eprintln s!"aggregate requires exactly {view.shards.size} shard proofs; got {proofHexes.length}" return 1 @@ -1210,6 +1239,7 @@ def aggregateCmd : Cli.Cmd := `[Cli| "max-ram" : Nat; "Aggregate in-flight RAM budget in GiB (default: 92% of MemTotal). An estimated-oversized slot runs alone." "structural-above" : Nat; "Use structural joins when a node contains more than N subject leaves (default 4096; 0 means every join)." "direct-joins"; "Keep IxVM leaves raw until their first pair instead of wrapping first (non-default; substantially higher RAM)." + "texray"; "Stream per-phase `[texray]` timing/RSS lines (execute, witness, STARK stages) for every slot to stderr, as `ix prove --texray` does." ARGS: ...proofs : String; "Persisted shard-proof wrapper addresses, in any order (one per nonempty shard, except --plan-only or replay with aggregate children)." diff --git a/Ix/Cli/ProveCmd.lean b/Ix/Cli/ProveCmd.lean index 6e5e52b9b..59d8f6e37 100644 --- a/Ix/Cli/ProveCmd.lean +++ b/Ix/Cli/ProveCmd.lean @@ -237,6 +237,58 @@ def reportPartition (proven : Array (Array Address × Nat)) (planned : Nat) : IO ({proven.size - planned} from splits) — re-shard with this partition to \ skip the splits next run" +private def runShardPipeline (p : Cli.Parsed) (ixe manifest : String) + (aiurSystem : Aiur.AiurSystem) (compiled : Aiur.CompiledToplevel) + (maxRamBytes : Nat) (indexDir? : Option System.FilePath) : IO UInt32 := do + let (ixonEnv, shards) ← match ← Ix.Cli.CheckCmd.loadEnvAndShards manifest ixe with + | .error e => IO.eprintln e; return 1 + | .ok value => pure value + let selected ← match (p.flag? "shard").map (·.as! Nat), (p.flag? "shards").map (·.as! String) with + | some k, none => pure #[k] + | none, some s => match Ix.Cli.CheckCmd.parseShardSelection s with + | .error e => IO.eprintln e; return 1 + | .ok ids => pure ids + | none, none => pure (Array.range shards.size) + | some _, some _ => IO.eprintln "use only one of --shard and --shards"; return 1 + if let some k := selected.find? (· ≥ shards.size) then + IO.eprintln s!"shard {k} out of range ({shards.size} shards)" + return 1 + let envHandle ← match Aiur.EnvHandle.fromIxe ixe with + | .error e => IO.eprintln e; return 1 + | .ok handle => pure handle + let some verifyIdx := compiled.getFuncIdx `verify_claim + | IO.eprintln "verify_claim entrypoint missing"; return 1 + let recursion ← match ShardProofIndex.buildRecursionBackend aiurSystem verifyIdx with + | .error e => IO.eprintln e; return 1 + | .ok backend => pure backend + let owned := Ix.Cli.CheckCmd.ownedConstsPer ixonEnv shards + let storePath ← StoreIO.toIO Store.storeDir + let plans ← StoreIO.toIO (Store.cacheDir "shard-splits") + -- `idmeasuredPeakBytes` per selected leaf: the native pipeline balances + -- its NUMA lanes by measured prover peak when the manifest carries one + -- (`ix shard refine`), and falls back to block counts otherwise. + let peaks : Array Nat ← do + match Ix.Cli.CheckCmd.parseIxesManifest (← IO.FS.readBinFile manifest) with + | .ok view => pure (selected.map fun k => + match view.shardIds.findIdx? (· == k) with + | some i => (view.measuredPeakBytes[i]?).getD 0 + | none => (view.measuredPeakBytes[k]?).getD 0) + | .error _ => pure (selected.map fun _ => 0) + let ids := String.intercalate "\n" + ((selected.zip peaks).toList.map fun (k, peak) => s!"{k}\t{peak}") + match ← Aiur.shardPipeline aiurSystem recursion.system envHandle + (Ix.Cli.CheckCmd.addrListsBlob (selected.map (shards[·]!))) + (Ix.Cli.CheckCmd.addrListsBlob (selected.map (owned[·]!))) ids + verifyIdx recursion.aggrIdx maxRamBytes storePath.toString + (indexDir?.map (·.toString) |>.getD "") plans.toString + (p.hasFlag "lookahead") (p.hasFlag "skip-proven") (p.hasFlag "keep-going") with + | .error e => IO.eprintln s!"[prove] {e}"; return 1 + | .ok summary => + IO.eprintln s!"[prove] {summary}; original manifest claims preserved" + if let some out := (p.flag? "out-ixes").map (·.as! String) then + IO.FS.writeBinFile out (← IO.FS.readBinFile manifest) + return 0 + def runProveCmd (p : Cli.Parsed) : IO UInt32 := do -- Streamed `[texray] : ── RAM Δ/peak` lines on stderr as -- each `aiur/` / `stark/` span closes: the per-phase wall + RSS @@ -245,10 +297,19 @@ def runProveCmd (p : Cli.Parsed) : IO UInt32 := do let keepGoing := p.hasFlag "keep-going" -- Same units as `ix shard --max-ram`: the per-shard prover budget the -- partition was sized against, re-checked here against each shard's - -- measured peak. 0 = detect (85% of `MemAvailable`, the check batch's - -- gate policy — see `shardProveWithEnv`). + -- predicted peak. The legacy path detects a budget when this is 0 + -- (see `shardProveWithEnv`); split healing requires an explicit budget. let maxRamBytes := ((p.flag? "max-ram").map (·.as! Nat)).getD 0 * gibBytes + let healSplits := p.hasFlag "heal-splits" || p.hasFlag "lookahead" + if healSplits then + if !(p.hasFlag "ixe" && p.hasFlag "ixes") || p.hasFlag "exec-only" || + p.hasFlag "claim" || !(p.variableArgsAs! String).isEmpty then + IO.eprintln "--heal-splits/--lookahead require --ixe and --ixes, without --exec-only, --claim or names" + return 1 + if maxRamBytes == 0 then + IO.eprintln "--heal-splits/--lookahead require explicit positive --max-ram (GiB)" + return 1 let execOnly := p.hasFlag "exec-only" let outIxes := (p.flag? "out-ixes").map (·.as! String) -- The shard-proof index: written for every proof this run persists, @@ -277,6 +338,9 @@ def runProveCmd (p : Cli.Parsed) : IO UInt32 := do | .error e => IO.eprintln s!"compilation failed: {e}"; return 1 | .ok c => pure c let aiurSystem := Aiur.AiurSystem.build compiled.bytecode commitmentParameters friParameters + if healSplits then + return ← runShardPipeline p (p.flag! "ixe" |>.as! String) + (p.flag! "ixes" |>.as! String) aiurSystem compiled maxRamBytes indexDir? let runOne := proveOne aiurSystem compiled match ixePath, (p.flag? "ixes").map (·.as! String), (p.flag? "shard").map (·.as! Nat) with | some ixe, some manifest, some k => @@ -373,9 +437,11 @@ def proveCmd : Cli.Cmd := `[Cli| "shards" : String; "With --ixes and no --shard: prove only these leaves — `K`, `a-b`, or a comma list of those — in one process (one env load); every other leaf is carried over unchanged by --out-ixes." "out-ixes" : String; "Write the partition this run actually proved — splits included — as a `.ixes` manifest to this path: the manifest `ix verify --ixes` checks these proofs against, and the one the next run of this env should start from. Skipped if any shard failed." "exec-only"; "Execute each shard and measure its projected prover peak, splitting over-budget shards as usual, but never start a STARK. The cheap way to audit a partition's split behavior at scale." + "heal-splits"; "With --ixe/--ixes: split oversized shards privately and prove flat joins back to each original CheckEnv claim. Requires explicit positive --max-ram; --out-ixes keeps the original manifest." + "lookahead"; "Enable split healing and execute at most one next shard while the current proof runs. Check each executed record's predicted proving peak against --max-ram before proving it." "skip-proven"; "With --ixes: before executing a leaf, look its claim up in the shard-proof index (`~/.ix/cache/shard-proofs/`); a recorded proof that decodes, bundles exactly that claim and verifies natively is reused — its address printed, nothing executed — instead of proving again. How a partially proved partition resumes after a refinement." "no-index"; "Neither read nor write the shard-proof index (every persisted proof is normally recorded there under its claim digest)." - "max-ram" : Nat; "Per-shard prover-RAM budget, GiB — normally the same value the partition was sized with (`ix shard --max-ram`). Each shard is executed, its projected prover peak measured on the resulting record, and the proof attempted only if it fits; an over-budget shard is cut into the part count the peak model projects will fit, and each part re-gated, instead of being taken into the FFT phases that would exhaust the box. Omit to detect: 85% of the machine's available RAM." + "max-ram" : Nat; "Per-shard prover-RAM budget, GiB — normally the same value the partition was sized with (`ix shard --max-ram`). Each shard is executed, its projected prover peak measured on the resulting record, and the proof attempted only if it fits; an over-budget shard is cut into the part count the peak model projects will fit, and each part re-gated, instead of being taken into the FFT phases that would exhaust the box. Required with --heal-splits/--lookahead; otherwise omit to detect 85% of the machine's available RAM." ARGS: ...names : String; "Fully-qualified Lean.Name(s) to prove. With none, iterate every named constant in the env (sorted)." diff --git a/Ix/Cli/ShardProofIndex.lean b/Ix/Cli/ShardProofIndex.lean index 500ac2aa9..551e861e7 100644 --- a/Ix/Cli/ShardProofIndex.lean +++ b/Ix/Cli/ShardProofIndex.lean @@ -11,6 +11,8 @@ module public import Ix.Address public import Ix.Aiur.Compiler public import Ix.Aiur.Protocol +public import Ix.Aggr +public import Ix.MultiStark public import Ix.Claim public import Ix.Ixon public import Ix.Store @@ -20,6 +22,42 @@ public section namespace Ix.Cli.ShardProofIndex +structure RecursionBackend where + system : Aiur.AiurSystem + aggrIdx : Aiur.Bytecode.FunIdx + allowed : ByteArray + +def buildRecursionBackend (ixvm : Aiur.AiurSystem) (verifyIdx : Nat) + (parameters : MultiStark.RecursionParameters := MultiStark.defaultRecursionParameters) : + Except String RecursionBackend := do + let top ← Aggr.ixAggr.mapError (fun e => s!"recursion toplevel: {e}") + let compiled ← top.compileWithGroups Aggr.functionGroups + |>.mapError (fun e => s!"recursion compilation: {e}") + let some aggrIdx := compiled.getFuncIdx `ix_aggr + | throw "recursion entrypoint missing" + let system := MultiStark.buildRecursionSystem compiled.bytecode parameters + pure { system, aggrIdx, allowed := Aggr.allowedBlob ixvm.vkBytes verifyIdx system.vkBytes aggrIdx } + +/-- Authenticate the backend from the proof. Healed CheckEnv leaves bind the +same claim bytes through the configured recursion system and allowed keys. -/ +def verifyProof (ixvm : Aiur.AiurSystem) (verifyIdx : Nat) + (claim : Ix.Claim) (proof : Aiur.Proof) + (recursion? : Option RecursionBackend := none) + (parameters : MultiStark.RecursionParameters := MultiStark.defaultRecursionParameters) : + Except String Unit := do + let bytes := Ix.Claim.ser claim + let input := IxVM.ClaimHarness.packedDigestKey (Address.blake3 bytes) + match ixvm.verify (Aiur.buildClaim verifyIdx input #[]) proof with + | .ok () => pure () + | .error rawError => + let .checkEnv _ _ := claim | throw s!"IxVM verification failed: {rawError}" + let backend ← match recursion? with + | some backend => pure backend + | none => buildRecursionBackend ixvm verifyIdx parameters + backend.system.verify + (Aiur.buildClaim backend.aggrIdx (Aggr.pubInput backend.allowed bytes) #[]) proof + |>.mapError (fun e => s!"neither IxVM nor healed CheckEnv verification succeeded: {e}") + /-- The index directory: the global `~/.ix/cache/shard-proofs`, or a hermetic root for tests. -/ def indexDir (cacheRoot? : Option System.FilePath := none) : IO System.FilePath := do @@ -48,12 +86,13 @@ def writeAddress (dir : System.FilePath) (digest addr : Address) : IO Unit := do /-- Native verification of a persisted shard-proof wrapper against the claim the caller expects: the wrapper must decode, bundle exactly `expected`, - and its proof must verify under `verify_claim`'s public input for that - claim (the same check `ix verify --shard K` performs). -/ + and its proof must verify as either an IxVM or healed CheckEnv proof. -/ def verifyWrapper (aiurSystem : Aiur.AiurSystem) (compiled : Aiur.CompiledToplevel) (expected : Ix.Claim) (proofAddr : Address) : IO (Except String Unit) := do try let bytes ← StoreIO.toIO (Store.read proofAddr) + if Address.blake3 bytes != proofAddr then + return .error s!"wrapper {proofAddr} has a different content digest" match Ixon.Proof.de bytes with | .error e => pure (.error s!"wrapper {proofAddr} does not decode: {e}") | .ok wrapper => @@ -64,11 +103,7 @@ def verifyWrapper (aiurSystem : Aiur.AiurSystem) (compiled : Aiur.CompiledToplev | .ok proof => let some funIdx := compiled.getFuncIdx `verify_claim | return .error "`verify_claim` entrypoint missing from compiled toplevel" - let input := IxVM.ClaimHarness.packedDigestKey - (Address.blake3 (Ix.Claim.ser wrapper.claim)) - match aiurSystem.verify (Aiur.buildClaim funIdx input #[]) proof with - | .ok () => pure (.ok ()) - | .error e => pure (.error s!"proof {proofAddr} does not verify: {e}") + pure (verifyProof aiurSystem funIdx wrapper.claim proof) catch e => pure (.error s!"{e}") /-- The address of a verified proof of `expected`, if the index has one. diff --git a/Ix/Cli/VerifyCmd.lean b/Ix/Cli/VerifyCmd.lean index 2f03bb565..121cdc1fa 100644 --- a/Ix/Cli/VerifyCmd.lean +++ b/Ix/Cli/VerifyCmd.lean @@ -51,8 +51,12 @@ private def friParameters : Aiur.FriParameters := /-- Verify one persisted `Ixon.Proof` wrapper (by store address) against its bundled claim, using an already-built Aiur backend. -/ def verifyOneProof (aiurSystem : Aiur.AiurSystem) (compiled : Aiur.CompiledToplevel) - (proofAddr : Address) : IO UInt32 := do + (proofAddr : Address) + (recursionParameters : MultiStark.RecursionParameters := MultiStark.defaultRecursionParameters) : IO UInt32 := do let bytes ← StoreIO.toIO (Store.read proofAddr) + if Address.blake3 bytes != proofAddr then + IO.eprintln s!"error: proof {proofAddr} has a different content digest" + return 1 let wrapper ← IO.ofExcept (Ixon.Proof.de bytes) let proof ← match Aiur.Proof.ofBytesChecked wrapper.proof with | .ok proof => pure proof @@ -67,9 +71,7 @@ def verifyOneProof (aiurSystem : Aiur.AiurSystem) (compiled : Aiur.CompiledTople | none => IO.eprintln "error: `verify_claim` entrypoint missing from compiled toplevel" return 1 - let input : Array Aiur.G := IxVM.ClaimHarness.packedDigestKey claimDigest - let aiurClaim := Aiur.buildClaim funIdx input #[] - match aiurSystem.verify aiurClaim proof with + match ShardProofIndex.verifyProof aiurSystem funIdx wrapper.claim proof none recursionParameters with | .ok () => IO.println s!"ok: proof {proofAddr} verifies claim {claimDigest}" return 0 @@ -87,10 +89,7 @@ def buildBackend : IO (Except String (Aiur.AiurSystem × Aiur.CompiledToplevel)) | .ok compiled => return .ok (Aiur.AiurSystem.build compiled.bytecode commitmentParameters friParameters, compiled) -structure AggregateBackend where - system : Aiur.AiurSystem - aggrIdx : Aiur.Bytecode.FunIdx - allowed : ByteArray +abbrev AggregateBackend := ShardProofIndex.RecursionBackend structure ExpectedAggregate where claim : Ix.Claim @@ -152,30 +151,11 @@ aggregate root: the IxVM vk and the single-entrypoint recursion vk. -/ private def buildAggregateBackend (recursionParameters : MultiStark.RecursionParameters) : IO (Except String AggregateBackend) := do - let ixvmCompiled ← match IxVM.ixVM with - | .error e => return .error s!"IxVM toplevel merging failed: {e}" - | .ok top => match top.compileWithGroups IxVM.functionGroups with - | .error e => return .error s!"IxVM compilation failed: {e}" - | .ok compiled => pure compiled - let aggrCompiled ← match Aggr.ixAggr with - | .error e => return .error s!"recursion toplevel merging failed: {e}" - | .ok top => match top.compileWithGroups Aggr.functionGroups with - | .error e => return .error s!"recursion compilation failed: {e}" - | .ok compiled => pure compiled - let verifyIdx := ixvmCompiled.getFuncIdx `verify_claim |>.get! - let aggrIdx := aggrCompiled.getFuncIdx `ix_aggr |>.get! - let ixvmSystem := Aiur.AiurSystem.build ixvmCompiled.bytecode - commitmentParameters friParameters - let aggrSystem := MultiStark.buildRecursionSystem aggrCompiled.bytecode - recursionParameters - let ixvmVk := ixvmSystem.vkBytes - let aggrVk := aggrSystem.vkBytes - let allowed := Aggr.allowedBlob ixvmVk verifyIdx aggrVk aggrIdx - return .ok { - system := aggrSystem - aggrIdx - allowed - } + let (ixvmSystem, compiled) ← match ← buildBackend with + | .error e => return .error e + | .ok backend => pure backend + let verifyIdx := compiled.getFuncIdx `verify_claim |>.get! + return ShardProofIndex.buildRecursionBackend ixvmSystem verifyIdx recursionParameters private def shardStatement (env : Ixon.Env) (owned : Array Address) : Except String Aggr.CheckEnvTrees := do @@ -259,7 +239,8 @@ private def verifyAggregateProof (backend : AggregateBackend) - no `--shard` + proofs: composed verdict — coverage, every proof bound to a shard, and every shard covered by a valid proof. -/ def verifyShardComposition (ixePath manifestPath : String) (shardK? : Option Nat) - (proofs : List String) (record : Bool := false) : IO UInt32 := do + (proofs : List String) (record : Bool := false) + (recursionParameters : MultiStark.RecursionParameters := MultiStark.defaultRecursionParameters) : IO UInt32 := do let (ixonEnv, shards) ← match (← Ix.Cli.CheckCmd.loadEnvAndShards manifestPath ixePath) with | .error e => IO.eprintln e; return 1 | .ok r => pure r @@ -295,7 +276,7 @@ def verifyShardComposition (ixePath manifestPath : String) (shardK? : Option Nat if d != expected then IO.eprintln s!"[verify] FAIL: proof {proofAddr} (claim {d}) is not shard {k} (claim {expected})" rc := 1 - else if (← verifyOneProof aiurSystem compiled proofAddr) != 0 then rc := 1 + else if (← verifyOneProof aiurSystem compiled proofAddr recursionParameters) != 0 then rc := 1 else recordProof d proofAddr return rc | none => @@ -315,7 +296,7 @@ def verifyShardComposition (ixePath manifestPath : String) (shardK? : Option Nat match digestToShard.get? d with | none => IO.eprintln s!"[verify] FAIL: proof {proofAddr} (claim {d}) matches no shard"; rc := 1 | some k => - if (← verifyOneProof aiurSystem compiled proofAddr) != 0 then rc := 1 + if (← verifyOneProof aiurSystem compiled proofAddr recursionParameters) != 0 then rc := 1 else covered := covered.insert k recordProof d proofAddr @@ -397,7 +378,7 @@ def runVerifyCmdWith (recursionParameters : MultiStark.RecursionParameters) match (p.flag? "ixe").map (·.as! String), (p.flag? "ixes").map (·.as! String) with | some ixe, some manifest => verifyShardComposition ixe manifest ((p.flag? "shard").map (·.as! Nat)) proofs - (p.hasFlag "record") + (p.hasFlag "record") recursionParameters | _, _ => if proofs.isEmpty then p.printError "error: must specify ... (or --ixe + --ixes for a shard partition)" @@ -408,7 +389,7 @@ def runVerifyCmdWith (recursionParameters : MultiStark.RecursionParameters) let mut rc : UInt32 := 0 for hex in proofs do let proofAddr ← addrOfHex! "proof" hex - if (← verifyOneProof aiurSystem compiled proofAddr) != 0 then rc := 1 + if (← verifyOneProof aiurSystem compiled proofAddr recursionParameters) != 0 then rc := 1 return rc def runVerifyCmd (p : Cli.Parsed) : IO UInt32 := diff --git a/Tests/Aggr.lean b/Tests/Aggr.lean index eceb1146f..41850e3e5 100644 --- a/Tests/Aggr.lean +++ b/Tests/Aggr.lean @@ -200,7 +200,8 @@ def smokeSuite : IO UInt32 := do Ix.Cli.AggregateCmd.aggregateStructuralJoinRamBytes + 3 * Ix.Cli.AggregateCmd.aggregateFlatJoinRamPerSubjectBytes, Ix.Cli.AggregateCmd.aggregateWrapRamBytes, - Ix.Cli.AggregateCmd.aggregateStructuralJoinRamBytes] && + Ix.Cli.AggregateCmd.aggregateStructuralJoinRamBytes + + 4 * Ix.Cli.AggregateCmd.aggregateStructuralJoinRamPerSubjectBytes] && Ix.Cli.AggregateCmd.aggregateSlotRamWeights direct == #[ Ix.Cli.AggregateCmd.aggregateRawShardRamBytes, Ix.Cli.AggregateCmd.aggregateRawShardRamBytes, diff --git a/Tests/AggrSemantics.lean b/Tests/AggrSemantics.lean index f042f936e..c615a8658 100644 --- a/Tests/AggrSemantics.lean +++ b/Tests/AggrSemantics.lean @@ -599,6 +599,29 @@ def semanticSuite : IO UInt32 := do { op := .join 0 1, subjectCount := 7, structural := false } == Ix.Cli.AggregateCmd.aggregateStructuralJoinRamBytes + 7 * Ix.Cli.AggregateCmd.aggregateFlatJoinRamPerSubjectBytes + let structuralWeight := Ix.Cli.AggregateCmd.aggregateShapeRamBytes 9 + let nodeBudget := 453 * Ix.Cli.AggregateCmd.aggregateGiB + -- Mathlib's low joins can still share a node. The pair live on node 0 at + -- the OOM (slots 140/355) must not, even with process-wide headroom. + let structuralPackingBySize : Bool := + structuralWeight 9480 + structuralWeight 10271 <= nodeBudget && + structuralWeight 23993 + structuralWeight 24805 <= nodeBudget && + structuralWeight 187668 <= nodeBudget && + structuralWeight 187668 + structuralWeight 13023 > nodeBudget && + structuralWeight 91620 + structuralWeight 96048 > nodeBudget + let structuralWeightsCoverMeasuredPeaks : Bool := + (#[(5371, 196), (11972, 203), (19751, 208), (55496, 212), + (91068, 381), (91620, 257), (96048, 249), (126527, 381), + (187668, 384), (314195, 455)] : Array (Nat × Nat)).all + fun (subjects, peakGiB) => + peakGiB * Ix.Cli.AggregateCmd.aggregateGiB <= structuralWeight subjects + let structuralWeightStep := + structuralWeight 65536 == 275 * Ix.Cli.AggregateCmd.aggregateGiB && + structuralWeight 65537 == 390 * Ix.Cli.AggregateCmd.aggregateGiB + let structuralFallbackWeight := + Ix.Cli.AggregateCmd.aggregateSlotRamBytes + { op := .join 0 1, subjectCount := 187668, structural := true } == + structuralWeight 187668 let memTotalParsing := Ix.Cli.AggregateCmd.aggregateMemTotalBytes "MemTotal: 1024 kB\nMemFree: 512 kB\n" == some (1024 * 1024) let invalidScheduleRejected : Bool := match @@ -784,6 +807,14 @@ def semanticSuite : IO UInt32 := do oversizedRunsAlone, test "flat self-pair RAM reserve is affine in subject leaves" flatWeightAffine, + test "structural RAM weights preserve small packing and reject the OOM pair" + structuralPackingBySize, + test "structural RAM weights cover the measured Mathlib peak envelope" + structuralWeightsCoverMeasuredPeaks, + test "large structural joins reserve the trace-size step" + structuralWeightStep, + test "hand-built structural schedules use the subject-dependent weight" + structuralFallbackWeight, test "aggregate scheduler parses MemTotal for its default budget" memTotalParsing, test "invalid non-post-order schedules are rejected" invalidScheduleRejected, diff --git a/Tests/Main.lean b/Tests/Main.lean index fb6d83d8b..e49579da0 100644 --- a/Tests/Main.lean +++ b/Tests/Main.lean @@ -59,6 +59,7 @@ import Tests.AggrActivation import Tests.Cli import Tests.Ix.Ixes import Tests.ShardMap +import Tests.ShardPipeline import Tests.Ix.EnvBody import Tests.Ix.Lean4Lean import Tests.Ix.MetaEnv @@ -208,6 +209,7 @@ def primaryRunners : List (String × IO UInt32) := [ /-- Ignored test runners - expensive, deferred IO actions run only when explicitly requested -/ def ignoredRunners (env : Lean.Environment) : List (String × IO UInt32) := [ + ("shard-pipeline", Tests.ShardPipeline.suite), ("ixvm", do let kernelChecks ← kernelChecks env -- the kernel CheckEnv smokes . diff --git a/Tests/ShardPipeline.lean b/Tests/ShardPipeline.lean new file mode 100644 index 000000000..62aa58991 --- /dev/null +++ b/Tests/ShardPipeline.lean @@ -0,0 +1,127 @@ +module + +import Ix.Cli.CheckCmd +import Ix.Cli.ShardProofIndex + +/-! +Opt-in integration test using the real IxVM and recursion systems. Tiny FRI +parameters keep the fixture smaller; this still constructs recursive proofs +and must run separately from measurements. All writes stay in a temporary +directory. Production parameters are untouched. +-/ + +namespace Tests.ShardPipeline + +private def ensure (ok : Bool) (message : String) : IO Unit := do + unless ok do throw (IO.userError message) + +private def fixture : Ixon.Env × Array Address := Id.run do + let a : Ixon.Constant := ⟨.axio ⟨false, 0, .sort 0⟩, #[], #[], #[.succ .zero]⟩ + let d : Ixon.Constant := ⟨.axio ⟨false, 0, .sort 0⟩, #[], #[], #[.succ (.succ .zero)]⟩ + let aa := Address.blake3 (Ixon.serConstant a) + let da := Address.blake3 (Ixon.serConstant d) + let b : Ixon.Constant := ⟨.axio ⟨false, 0, .ref 0 #[]⟩, #[], #[aa], #[]⟩ + let c : Ixon.Constant := ⟨.axio ⟨false, 0, .ref 0 #[]⟩, #[], #[da], #[]⟩ + let ba := Address.blake3 (Ixon.serConstant b) + let ca := Address.blake3 (Ixon.serConstant c) + let env := ({} : Ixon.Env).storeConst aa a |>.storeConst ba b + |>.storeConst ca c |>.storeConst da d + return (env, #[aa, ba, ca, da]) + +private def objectPath (dir : System.FilePath) (address : Address) : System.FilePath := + let s := (toString address).toSlice + dir / (s.take 2).toString / (s.drop 2 |>.take 2).toString / + (s.drop 4 |>.take 2).toString / (s.drop 6).toString + +private def indexed (dir : System.FilePath) (claim : Ix.Claim) : IO Address := do + let some address ← Ix.Cli.ShardProofIndex.readAddress dir (Address.blake3 (Ix.Claim.ser claim)) + | throw (IO.userError "missing shard proof index entry") + return address + +private def smoke : IO Unit := do + let (env, addresses) := fixture + let a := addresses[0]! + let b := addresses[1]! + let c := addresses[2]! + let d := addresses[3]! + let claimOf (owned : Array Address) := + IxVM.ClaimHarness.shardCheckEnvClaimTrees env owned |>.map (·.1) + let original ← IO.ofExcept (claimOf #[a, b, c]) + let ab ← IO.ofExcept (claimOf #[a, b]) + let ca ← IO.ofExcept (claimOf #[a]) + let top ← IO.ofExcept (IxVM.ixVM.mapError (fun e => s!"{e}")) + let compiled ← IO.ofExcept (top.compile.mapError (fun e => s!"{e}")) + let verifyIdx := compiled.getFuncIdx `verify_claim |>.get! + let cp : Aiur.CommitmentParameters := { logBlowup := 1, capHeight := 0 } + let fp : Aiur.FriParameters := + { logFinalPolyLen := 0, maxLogArity := 1, numQueries := 4, + commitProofOfWorkBits := 0, queryProofOfWorkBits := 0 } + let ixvm := Aiur.AiurSystem.build compiled.bytecode cp fp + let recursion ← IO.ofExcept (Ix.Cli.ShardProofIndex.buildRecursionBackend ixvm verifyIdx + { commitment := cp, fri := fp }) + let handle ← IO.ofExcept (Aiur.EnvHandle.fromBytes (← IO.ofExcept (Ixon.serEnv env))) + let dir ← IO.FS.createTempDir + IO.println s!"shard-pipeline smoke artifacts: {dir}" + let storeDir := dir / "store" + let indexDir := dir / "index" + let planDir := dir / "splits" + IO.FS.createDirAll planDir + let maxRam := 128 * 1024 * 1024 * 1024 + -- Recreate two validated journals from an interrupted run. This forces + -- nested healing without a production-only fault-injection switch. + let journal (claim : Ix.Claim) (parts : Array (Array Address)) : IO Unit := do + let key := "ix-shard-splits-v2".toUTF8 ++ recursion.allowed ++ + maxRam.toUInt64.toLEBytes ++ Ix.Claim.ser claim + let path := planDir / s!"{Address.blake3 key}.json" + IO.FS.writeFile path (Lean.toJson (parts.map (fun p => p.map toString))).compress + journal original #[#[a, b], #[c]] + journal ab #[#[a], #[b]] + let blocks := Ix.Cli.CheckCmd.addrListsBlob #[#[a, b, c], #[d]] + let run (lookahead : Bool) : IO String := do + IO.ofExcept (← Aiur.shardPipeline ixvm recursion.system handle blocks blocks "0\n1" + verifyIdx recursion.aggrIdx maxRam storeDir.toString indexDir.toString + planDir.toString lookahead true false) + let verifyOriginal : IO Address := do + let address ← indexed indexDir original + let bytes ← IO.FS.readBinFile (objectPath storeDir address) + ensure (Address.blake3 bytes == address) "stored wrapper hash changed" + let wrapper ← IO.ofExcept (Ixon.Proof.de bytes) + ensure (Ix.Claim.ser wrapper.claim == Ix.Claim.ser original) "healing changed the original claim bytes" + let proof ← IO.ofExcept (Aiur.Proof.ofBytesChecked wrapper.proof) + IO.ofExcept (Ix.Cli.ShardProofIndex.verifyProof ixvm verifyIdx original proof (some recursion)) + let rawClaim := Aiur.buildClaim verifyIdx + (IxVM.ClaimHarness.packedDigestKey (Address.blake3 (Ix.Claim.ser original))) #[] + ensure (!((ixvm.verify rawClaim proof).toBool)) "split root should be a recursion proof" + return address + let first ← run true + IO.println first + -- The summary is one segment per in-process NUMA lane (`a | b | …`), so + -- counts are summed over segments: a lane holding a single shard has + -- nothing to overlap, and cached proofs are reported per lane. + let countOf (summary label : String) : Nat := + (summary.splitOn label).dropLast.foldl (fun acc piece => + acc + (((piece.trimRight.splitOn " ").getLast?.bind String.toNat?).getD 0)) 0 + ensure (countOf first " preparation overlap(s)" > 0) "lookahead did not overlap any preparation" + let firstAddress ← verifyOriginal + let resumed ← run false + ensure (countOf resumed " cached proof(s)" == 2) "resume did not reuse both original proofs" + ensure ((← verifyOriginal) == firstAddress) "resume replaced an already verified original proof" + -- Remove completed parents and corrupt one persisted child. The journal + -- must resume through verified siblings and replace the rejected child. + IO.FS.removeFile (indexDir / toString (Address.blake3 (Ix.Claim.ser original))) + IO.FS.removeFile (indexDir / toString (Address.blake3 (Ix.Claim.ser ab))) + let childAddress ← indexed indexDir ca + IO.FS.writeBinFile (objectPath storeDir childAddress) "corrupt child".toUTF8 + IO.println (← run false) + let _ ← verifyOriginal + IO.println "shard-pipeline: nested healing, unchanged claim, overlap, verified resume and corrupt-child recovery passed" + +public def suite : IO UInt32 := do + try + smoke + return 0 + catch e => + IO.eprintln s!"shard-pipeline: {e}" + return 1 + +end Tests.ShardPipeline diff --git a/crates/aiur/src/synthesis.rs b/crates/aiur/src/synthesis.rs index ae7e2ac89..28d5435b5 100644 --- a/crates/aiur/src/synthesis.rs +++ b/crates/aiur/src/synthesis.rs @@ -504,6 +504,11 @@ impl AiurSystem { } } + /// The bytecode layout used by native executors to size their query record. + pub fn toplevel(&self) -> &Toplevel { + &self.toplevel + } + /// `prove_ixvm`, but the record's projected prover peak has to fit /// `max_bytes` before any proving starts (`None` skips the check), /// and `exec_only` stops after execution + measurement — the split diff --git a/crates/ffi/Cargo.toml b/crates/ffi/Cargo.toml index 7a0bc7896..38c44fc6a 100644 --- a/crates/ffi/Cargo.toml +++ b/crates/ffi/Cargo.toml @@ -27,6 +27,7 @@ multi-stark = { workspace = true } mimalloc = { workspace = true } num-bigint = { workspace = true } rayon = { workspace = true } +libc = { workspace = true } rustc-hash = { workspace = true } serde_json = { workspace = true } sha2 = { workspace = true } diff --git a/crates/ffi/src/aiur/aggregate.rs b/crates/ffi/src/aiur/aggregate.rs index fc76e93cc..dccec396f 100644 --- a/crates/ffi/src/aiur/aggregate.rs +++ b/crates/ffi/src/aiur/aggregate.rs @@ -14,6 +14,8 @@ #![allow(clippy::too_many_arguments)] +mod shard_pipeline; + use std::{ cmp::Ordering, fs, @@ -24,7 +26,9 @@ use std::{ }; use aiur::{ - G, function_channel, + G, + execute::IOBuffer, + function_channel, synthesis::{AiurProof, AiurSystem, GatedProve}, }; use ix_common::address::Address; @@ -62,9 +66,19 @@ const GIB: usize = 1024 * 1024 * 1024; const WRAP_RAM_BYTES: usize = 195 * GIB; const STRUCTURAL_RAM_BYTES: usize = 195 * GIB; const RAW_SHARD_RAM_BYTES: usize = 4 * GIB; -const DIRECT_RAM_BYTES: usize = 390 * GIB; -const MIXED_RAM_BYTES: usize = 340 * GIB; +const DIRECT_RAM_BYTES: usize = 180 * GIB; +// Mixed joins measured 378–385 GiB projected peak on the 2026-09-09 Mathlib +// Stage 2 against the previous 340 GiB reserve; 390 matches a direct pair. +const MIXED_RAM_BYTES: usize = 180 * GIB; const FLAT_RAM_PER_SUBJECT: usize = 1024 * 1024; +// Structural subject roots are O(1), but assumption/path work and child +// verification can grow. Reserve a subject term plus a doubled base above +// 64k subjects: Mathlib peaks jumped to 380.5 GiB at 91,068 subjects, whereas +// a similar-sized join used 256.5 GiB. The term alone misses that trace-size +// step. A flat 195 GiB reserve caused a packed-node OOM at 187,668 subjects. +// Calibration and limitations: exp/design/numa-slot-pinning.md §12. +const STRUCTURAL_RAM_PER_SUBJECT: usize = 5 * MIB / 4; +const STRUCTURAL_LARGE_SUBJECTS: usize = 64 * 1024; fn format_gib(bytes: usize) -> String { let tenths = bytes.saturating_mul(10) / GIB; @@ -338,6 +352,44 @@ struct PreparedRun { env_root: Address, env_count: usize, expected_shards: ShardSet, + /// `IX_AGGREGATE_SHARDS` restricted the run to a subset of the manifest's + /// shards (experiments only): the root is the aggregate of that subtree and + /// may retain assumptions on constants owned by unselected shards. + partial: bool, +} + +/// Experiment knob: `IX_AGGREGATE_SHARDS=K|a-b|a,b-c,…` aggregates only the +/// named shard ids (the manifest tree pruned to that subtree) from the +/// corresponding subset of the supplied Stage 1 proofs. Shard claims depend +/// only on a shard's own constants and frontier, so the leaf proofs of a full +/// run stay valid; the root is then a partial, assumption-carrying aggregate. +fn shard_selection() -> Result>, String> { + let Ok(spec) = std::env::var("IX_AGGREGATE_SHARDS") else { + return Ok(None); + }; + let mut ids = FxHashSet::default(); + for piece in spec.split(',').map(str::trim).filter(|p| !p.is_empty()) { + let bad = |error: std::num::ParseIntError| { + format!("IX_AGGREGATE_SHARDS: malformed `{piece}`: {error}") + }; + match piece.split_once('-') { + Some((a, b)) => { + let a: u32 = a.trim().parse().map_err(bad)?; + let b: u32 = b.trim().parse().map_err(bad)?; + if b < a { + return Err(format!("IX_AGGREGATE_SHARDS: empty range `{piece}`")); + } + ids.extend(a..=b); + }, + None => { + ids.insert(piece.parse().map_err(bad)?); + }, + } + } + if ids.is_empty() { + return Err("IX_AGGREGATE_SHARDS: empty selection".into()); + } + Ok(Some(ids)) } #[derive(Clone, Copy, Debug)] @@ -380,14 +432,13 @@ struct Slot { struct ProveContext<'a> { specs: &'a [SlotSpec], prepared: &'a [PreparedShard], - proofs: Option<&'a [Arc]>, + proofs: Option<&'a [Arc]>, owner_by_address: &'a FxHashMap, ixvm_system: &'a AiurSystem, aggr_system: &'a AiurSystem, ixvm_vk: &'a [u8], aggr_vk: &'a [u8], allowed: &'a [u8], - verify_idx: usize, aggr_idx: usize, store_dir: &'a Path, cache_dir: Option<&'a Path>, @@ -506,14 +557,38 @@ fn prepare_run( ); } + let selection = shard_selection()?; let retained_old: Vec = owned .iter() .enumerate() + .filter(|(index, _)| { + selection + .as_ref() + .is_none_or(|ids| ids.contains(&manifest.shards[*index].id)) + }) .filter_map(|(index, addresses)| (!addresses.is_empty()).then_some(index)) .collect(); if retained_old.is_empty() { return Err("manifest has no shard owning an environment constant".into()); } + let partial = + selection.is_some() && retained_old.len() < manifest.shards.len(); + // A partial run's root covers only the selected shards' constants. + let root_addresses: Vec
= if partial { + let mut selected: Vec
= + retained_old.iter().flat_map(|old| owned[*old].iter().cloned()).collect(); + selected.par_sort_unstable(); + eprintln!( + "[aggregate] IX_AGGREGATE_SHARDS: partial aggregate over {} of {} shards ({} of {} constants); the root may retain assumptions", + retained_old.len(), + manifest.shards.len(), + selected.len(), + all_addresses.len() + ); + selected + } else { + all_addresses.clone() + }; let retained_ids: FxHashSet = retained_old.iter().map(|index| manifest.shards[*index].id).collect(); let source_tree = manifest.tree.clone().unwrap_or_else(|| { @@ -532,7 +607,11 @@ fn prepare_run( .iter() .cloned() .zip(owners_old) - .map(|(address, old)| (address, old_to_retained[&old])) + .map(|(address, old)| { + // Constants of unselected shards keep a sentinel owner that no retained + // index equals and no subject tree contains: they stay assumptions. + (address, old_to_retained.get(&old).copied().unwrap_or(usize::MAX)) + }) .collect(); let shard_inputs: Vec<(usize, u32, Vec
, Vec
)> = @@ -580,7 +659,7 @@ fn prepare_run( // resolve any errors in that same order for deterministic diagnostics. let shards = shard_results.into_iter().collect::, _>>()?; - let env_root = merkle_root_canonical_sorted(&all_addresses) + let env_root = merkle_root_canonical_sorted(&root_addresses) .ok_or("cannot aggregate an empty environment")?; let expected_shards = ShardSet( (0..retained_old.len().div_ceil(64)) @@ -596,8 +675,9 @@ fn prepare_run( owner_by_address, tree, env_root, - env_count: all_addresses.len(), + env_count: root_addresses.len(), expected_shards, + partial, }) } @@ -751,7 +831,16 @@ fn shape_ram_bytes(shape: u8, subject_count: usize) -> usize { 3 | 4 | 7 | 8 => MIXED_RAM_BYTES, 5 => STRUCTURAL_RAM_BYTES .saturating_add(subject_count.saturating_mul(FLAT_RAM_PER_SUBJECT)), - 9 => STRUCTURAL_RAM_BYTES, + 9 => { + let weight = STRUCTURAL_RAM_BYTES.saturating_add( + subject_count.saturating_mul(STRUCTURAL_RAM_PER_SUBJECT), + ); + if subject_count > STRUCTURAL_LARGE_SUBJECTS { + weight.max(2 * STRUCTURAL_RAM_BYTES) + } else { + weight + } + }, // Shapes 2/6 are direct pairs; unknown shapes retain the conservative // direct-pair fallback used by the Lean reference scheduler. _ => DIRECT_RAM_BYTES, @@ -859,7 +948,7 @@ fn validate_root_statement( prepared.env_root.hex() )); } - if root.assumptions.is_some() { + if root.assumptions.is_some() && !prepared.partial { return Err("aggregate root retains undischarged assumptions".into()); } Ok(()) @@ -966,12 +1055,43 @@ fn read_store(root: &Path, address: &Address) -> Result, String> { fn write_store(root: &Path, bytes: &[u8]) -> Result { let address = Address::hash(bytes); let path = store_path(root, &address); - let parent = path.parent().ok_or("store path has no parent")?; + write_atomic(&path, bytes)?; + Ok(address) +} + +fn write_atomic(path: &Path, bytes: &[u8]) -> Result<(), String> { + use std::{ + io::Write, + sync::atomic::{AtomicU64, Ordering}, + time::{SystemTime, UNIX_EPOCH}, + }; + static NEXT: AtomicU64 = AtomicU64::new(0); + let parent = path.parent().ok_or("output path has no parent")?; + let nonce = + SystemTime::now().duration_since(UNIX_EPOCH).unwrap_or_default().as_nanos(); + let tmp = path.with_extension(format!( + "{}.{nonce}.{}.tmp", + std::process::id(), + NEXT.fetch_add(1, Ordering::Relaxed) + )); fs::create_dir_all(parent) .map_err(|error| format!("create {}: {error}", parent.display()))?; - fs::write(&path, bytes) - .map_err(|error| format!("write {}: {error}", path.display()))?; - Ok(address) + // A failed create must never remove another writer's temporary file. + let mut file = fs::OpenOptions::new() + .write(true) + .create_new(true) + .open(&tmp) + .map_err(|error| format!("create {}: {error}", tmp.display()))?; + let result = (|| -> std::io::Result<()> { + file.write_all(bytes)?; + file.sync_all()?; + fs::rename(&tmp, path)?; + fs::File::open(parent)?.sync_all() + })(); + if result.is_err() { + let _ = fs::remove_file(&tmp); + } + result.map_err(|error| format!("write {}: {error}", path.display())) } fn decode_wrapper(bytes: &[u8]) -> Result { @@ -987,10 +1107,11 @@ fn load_input_proofs( proof_hexes: &str, store_dir: &Path, prepared: &[PreparedShard], + partial: bool, ) -> Result>, String> { let values: Vec<&str> = proof_hexes.lines().filter(|line| !line.is_empty()).collect(); - if values.len() != prepared.len() { + if !partial && values.len() != prepared.len() { return Err(format!( "aggregate requires exactly {} shard proofs; got {}", prepared.len(), @@ -1005,27 +1126,39 @@ fn load_input_proofs( if by_digest.len() != prepared.len() { return Err("two reconstructed shard claims have the same digest".into()); } + // Read, hash and decode every wrapper in parallel (this was 40+ s serial + // for 246 Mathlib proofs); claim matching and duplicate checks stay + // sequential below so the error semantics are unchanged. + let decoded: Vec<(Address, IxonProof)> = values + .par_iter() + .map(|value| -> Result<(Address, IxonProof), String> { + let address = Address::from_hex(value).ok_or_else(|| { + format!("shard proof is not a 64-character address: {value}") + })?; + let bytes = read_store(store_dir, &address)?; + if Address::hash(&bytes) != address { + return Err(format!( + "shard proof store object {} has the wrong digest", + address.hex() + )); + } + let wrapper = decode_wrapper(&bytes).map_err(|error| { + format!("decode shard proof {}: {error}", address.hex()) + })?; + Ok((address, wrapper)) + }) + .collect::, _>>()?; let mut proofs: Vec>> = vec![None; prepared.len()]; - for value in values { - let address = Address::from_hex(value).ok_or_else(|| { - format!("shard proof is not a 64-character address: {value}") - })?; - let bytes = read_store(store_dir, &address)?; - if Address::hash(&bytes) != address { - return Err(format!( - "shard proof store object {} has the wrong digest", - address.hex() - )); - } - let wrapper = decode_wrapper(&bytes).map_err(|error| { - format!("decode shard proof {}: {error}", address.hex()) - })?; + for (address, wrapper) in decoded { let mut claim_bytes = Vec::new(); wrapper.claim.put(&mut claim_bytes); let digest = Address::hash(&claim_bytes); - let shard = by_digest.get(&digest).copied().ok_or_else(|| { - format!("proof {} matches no manifest shard", address.hex()) - })?; + let Some(shard) = by_digest.get(&digest).copied() else { + if partial { + continue; // a proof for an unselected shard + } + return Err(format!("proof {} matches no manifest shard", address.hex())); + }; if wrapper.claim != prepared[shard].statement.claim { return Err(format!( "proof {} hit a claim-digest collision for shard {}", @@ -1052,6 +1185,102 @@ fn load_input_proofs( .collect() } +/// Authenticate a shard certificate under one of the two supported systems. +/// The backend is established by verification, never by an untrusted tag. +fn verify_shard_proof( + ixvm: &AiurSystem, + aggr: &AiurSystem, + verify_idx: usize, + aggr_idx: usize, + allowed: &[u8], + claim: &Claim, + proof: &AiurProof, +) -> Result<(ChildKind, Vec), String> { + let mut bytes = Vec::new(); + claim.put(&mut bytes); + let inner = inner_claim(verify_idx, &bytes); + if ixvm.verify(&inner, proof).is_ok() { + return Ok((ChildKind::Ixvm, inner)); + } + if !matches!(claim, Claim::CheckEnv { .. }) { + return Err("proof does not verify under the IxVM system".into()); + } + let outer = aggregate_outer_claim(aggr_idx, allowed, &bytes); + aggr.verify(&outer, proof).map_err(|error| { + format!("proof verifies under neither IxVM nor ix_aggr: {error:?}") + })?; + Ok((ChildKind::Aggr, outer)) +} + +fn import_shard_proof( + ixvm: &AiurSystem, + aggr: &AiurSystem, + verify_idx: usize, + aggr_idx: usize, + allowed: &[u8], + statement: Arc, + wrapper: &IxonProof, + proof_address: Option
, +) -> Result, String> { + if wrapper.claim != statement.claim { + return Err("shard proof bundles a different CheckEnv claim".into()); + } + let proof = AiurProof::from_bytes(&wrapper.proof) + .map_err(|error| format!("shard proof does not decode: {error}"))?; + let (kind, outer_claim) = verify_shard_proof( + ixvm, + aggr, + verify_idx, + aggr_idx, + allowed, + &statement.claim, + &proof, + )?; + let claims_bytes = serialize_claims(&[&outer_claim]); + Ok(Arc::new(Slot { + kind, + statement, + outer_claim, + proof, + proof_address, + claims_bytes, + })) +} + +/// Statement roots and slot indices do not depend on the input proof kind. +/// Imported healed leaves are already complete; only proof shapes/reserves +/// above them change. Cache keys continue to bind the same output statements. +fn bind_imported_specs( + specs: &mut [SlotSpec], + inputs: &[Arc], + aggr_vk: &[u8], + cache_fri_bytes: &[u8], +) { + for index in 0..specs.len() { + match specs[index].op { + PlanOp::Leaf(shard) if inputs[shard].kind == ChildKind::Aggr => { + let spec = &mut specs[index]; + spec.kind = ChildKind::Aggr; + spec.shape = None; + spec.outer_claim = inputs[shard].outer_claim.clone(); + spec.cache_key = cache_key(aggr_vk, cache_fri_bytes, &spec.outer_claim); + spec.ram_bytes = RAW_SHARD_RAM_BYTES; + }, + PlanOp::Join(left, right) => { + let shape = if specs[index].structural { + structural_shape_code(specs[left].kind, specs[right].kind) + } else { + shape_code(specs[left].kind, Some(specs[right].kind)) + }; + specs[index].shape = Some(shape); + specs[index].ram_bytes = + shape_ram_bytes(shape, specs[index].subject_count); + }, + PlanOp::Leaf(_) => {}, + } + } +} + fn cache_address(cache_dir: &Path, key: &Address) -> Option
{ let path = cache_dir.join(key.hex()); let raw = fs::read_to_string(path).ok()?; @@ -1285,25 +1514,13 @@ fn assumption_count(statement: &Statement) -> usize { statement.assumptions.as_ref().map_or(0, |tree| tree.leaves.len()) } -fn prove_aggregate( +/// Advice construction shared by Stage 2 and budgeted local shard healing. +fn aggregate_io( ctx: ProveContext<'_>, spec: &SlotSpec, left: &Slot, right: Option<&Slot>, - slot_index: usize, -) -> Result<(AiurProof, Option
), String> { - let replaying = ctx.reprove_slot == Some(slot_index); - if !replaying { - if let Some((proof, address)) = load_cached(ctx, slot_index, spec) { - return Ok((proof, Some(address))); - } - } else { - eprintln!( - "[aggregate] replay slot {slot_index}: bypassing its cache entry" - ); - } - let started = Instant::now(); - +) -> Result<(IOBuffer, String), String> { let left_system = match left.kind { ChildKind::Ixvm => ctx.ixvm_system, ChildKind::Aggr => ctx.aggr_system, @@ -1376,7 +1593,7 @@ fn prove_aggregate( let right_claims = right.map_or(empty.as_slice(), |slot| slot.claims_bytes.as_slice()); let shape = spec.shape.ok_or("aggregate proof slot has no shape")?; - let mut io = aggr_io_buffer(&AggrAdvice { + let io = aggr_io_buffer(&AggrAdvice { shape, proof_advice: [&left_advice, &right_advice], ixvm_vk: ctx.ixvm_vk, @@ -1388,6 +1605,97 @@ fn prove_aggregate( trees: &trees, paths: &paths, }); + let sizes = format!( + "proof advice {}+{} MiB, {} trees/{} MiB, {} paths/{} MiB, preimages {} MiB", + format_mib(left_advice.len()), + format_mib(right_advice.len()), + tree_storage.len(), + format_mib(tree_storage.iter().map(|t| t.bytes.len()).sum()), + path_storage.len(), + format_mib(path_storage.iter().map(|(_, p)| p.len()).sum()), + format_mib(preimage_storage.iter().map(|(_, p)| p.len()).sum()), + ); + Ok((io, sizes)) +} + +/// Extra RAM for one lookahead execution record (measured 32.6 GiB for a +/// Mathlib direct join), charged to both the process and any bound node. +const LOOKAHEAD_RAM_BYTES: usize = 40 * GIB; + +/// The overlappable front half of a slot: advice construction plus the +/// `ix_aggr` execution into a query record. Runs on a lane's prep thread +/// while that lane proves its current slot; `None` when the slot's proof is +/// already cached (nothing to prepare). +fn prepare_aggregate<'a>( + ctx: ProveContext<'a>, + spec: &SlotSpec, + left: &Slot, + right: Option<&Slot>, + slot_index: usize, +) -> Result>, String> { + if ctx.reprove_slot != Some(slot_index) + && load_cached(ctx, slot_index, spec).is_some() + { + return Ok(None); + } + let (io, _advice_sizes) = aggregate_io(ctx, spec, left, right)?; + let mut public_input = packed_digest(ctx.allowed); + public_input.extend(packed_digest(&spec.statement.claim_bytes)); + shard_pipeline::Execution::new( + ctx.aggr_system, + ctx.aggr_idx, + public_input, + io, + execute_ix_aggr, + ) + .map(Some) +} + +/// The back half: prove from a prepared record and persist. Mirrors the +/// tail of [`prove_aggregate`]. +fn prove_prepared( + ctx: ProveContext<'_>, + spec: &SlotSpec, + execution: shard_pipeline::Execution<'_>, + slot_index: usize, +) -> Result<(AiurProof, Option
), String> { + let started = Instant::now(); + let peak = execution.peak(); + let (outer_claim, proof) = execution.prove(); + let proved_at = Instant::now(); + if outer_claim != spec.outer_claim { + return Err("aggregate prover returned an unexpected outer claim".into()); + } + let address = persist_cached(ctx, slot_index, spec, &proof); + eprintln!( + "[aggregate] slot {slot_index}: prepared ahead; prove {:.1}s, persist {:.1}s, peak {} GiB", + (proved_at - started).as_secs_f64(), + proved_at.elapsed().as_secs_f64(), + format_gib(peak), + ); + Ok((proof, address)) +} + +fn prove_aggregate( + ctx: ProveContext<'_>, + spec: &SlotSpec, + left: &Slot, + right: Option<&Slot>, + slot_index: usize, +) -> Result<(AiurProof, Option
), String> { + let replaying = ctx.reprove_slot == Some(slot_index); + if !replaying { + if let Some((proof, address)) = load_cached(ctx, slot_index, spec) { + return Ok((proof, Some(address))); + } + } else { + eprintln!( + "[aggregate] replay slot {slot_index}: bypassing its cache entry" + ); + } + let started = Instant::now(); + + let (mut io, advice_sizes) = aggregate_io(ctx, spec, left, right)?; let mut public_input = packed_digest(ctx.allowed); public_input.extend(packed_digest(&spec.statement.claim_bytes)); let proving_started = Instant::now(); @@ -1411,28 +1719,24 @@ fn prove_aggregate( return Err("aggregate prover returned an unexpected outer claim".into()); } let address = persist_cached(ctx, slot_index, spec, &proof); + // Per-slot phase timings for every slot (not only replays): the numbers a + // Stage 2 throughput model needs — how much of a slot is advice/execute + // (overlappable) vs prove vs persistence — plus the record's peak. + eprintln!( + "[aggregate] slot {slot_index}: advice {:.1}s, execute+prove {:.1}s, persist {:.1}s, peak {} GiB", + (proving_started - started).as_secs_f64(), + (proved_at - proving_started).as_secs_f64(), + proved_at.elapsed().as_secs_f64(), + format_gib(peak), + ); if replaying { - let tree_bytes: usize = - tree_storage.iter().map(|tree| tree.bytes.len()).sum(); - let path_bytes: usize = - path_storage.iter().map(|(_, path)| path.len()).sum(); - let preimage_bytes: usize = - preimage_storage.iter().map(|(_, bytes)| bytes.len()).sum(); - let right_assumptions = - right.map_or(0, |slot| assumption_count(&slot.statement)); eprintln!( - "[aggregate] replay slot {slot_index}: shape {shape}, {} subjects, assumptions {}/{}/{}, proof advice {}+{} MiB, {} trees/{} MiB, {} paths/{} MiB, preimages {} MiB, query-record peak {} GiB ({} bytes)", + "[aggregate] replay slot {slot_index}: shape {}, {} subjects, assumptions {}/{}/{}, {advice_sizes}, query-record peak {} GiB ({} bytes)", + spec.shape.expect("aggregate slot has a shape"), spec.subject_count, assumption_count(&left.statement), - right_assumptions, + right.map_or(0, |slot| assumption_count(&slot.statement)), assumption_count(&spec.statement), - format_mib(left_advice.len()), - format_mib(right_advice.len()), - tree_storage.len(), - format_mib(tree_bytes), - path_storage.len(), - format_mib(path_bytes), - format_mib(preimage_bytes), format_gib(peak), peak, ); @@ -1451,56 +1755,34 @@ fn prove_slot( ctx: ProveContext<'_>, slot_index: usize, children: &[Arc], + prepared: Option>, ) -> Result, String> { let spec = ctx.specs.get(slot_index).ok_or("missing aggregate slot spec")?; + let mut prepared = prepared; match spec.op { PlanOp::Leaf(shard) => { - let prepared = &ctx.prepared[shard]; - let wrapper = + let prepared_shard = &ctx.prepared[shard]; + let raw = ctx.proofs.and_then(|proofs| proofs.get(shard)).ok_or_else(|| { format!( "shard {} proof was not loaded for replay", - prepared.original_id + prepared_shard.original_id ) })?; - let proof = AiurProof::from_bytes(&wrapper.proof).map_err(|error| { - format!("shard {} proof does not decode: {error}", prepared.original_id) - })?; - let inner = inner_claim(ctx.verify_idx, &prepared.statement.claim_bytes); - ctx.ixvm_system.verify(&inner, &proof).map_err(|error| { - format!( - "shard {} proof fails native verification: {error:?}", - prepared.original_id - ) - })?; - let inner_claims = serialize_claims(&[&inner]); - if spec.kind == ChildKind::Ixvm { - if spec.outer_claim != inner { - return Err("direct shard slot has an unexpected outer claim".into()); + if spec.shape.is_none() { + if spec.outer_claim != raw.outer_claim || spec.kind != raw.kind { + return Err("imported shard slot has an unexpected identity".into()); } - return Ok(Arc::new(Slot { - kind: ChildKind::Ixvm, - statement: spec.statement.clone(), - outer_claim: inner, - proof, - proof_address: None, - claims_bytes: inner_claims, - })); + return Ok(raw.clone()); } eprintln!( "[aggregate] wrapping shard {} into slot {slot_index}", - prepared.original_id + prepared_shard.original_id ); - let raw = Slot { - kind: ChildKind::Ixvm, - statement: spec.statement.clone(), - outer_claim: inner, - proof, - proof_address: None, - claims_bytes: inner_claims, + let (proof, proof_address) = match prepared.take() { + Some(execution) => prove_prepared(ctx, spec, execution, slot_index)?, + None => prove_aggregate(ctx, spec, raw, None, slot_index)?, }; - let (proof, proof_address) = - prove_aggregate(ctx, spec, &raw, None, slot_index)?; Ok(Arc::new(Slot { kind: ChildKind::Aggr, statement: spec.statement.clone(), @@ -1520,8 +1802,10 @@ fn prove_slot( eprintln!( "[aggregate] {mode}-joining slots {left_index}, {right_index} into {slot_index}" ); - let (proof, proof_address) = - prove_aggregate(ctx, spec, left, Some(right), slot_index)?; + let (proof, proof_address) = match prepared.take() { + Some(execution) => prove_prepared(ctx, spec, execution, slot_index)?, + None => prove_aggregate(ctx, spec, left, Some(right), slot_index)?, + }; Ok(Arc::new(Slot { kind: ChildKind::Aggr, statement: spec.statement.clone(), @@ -1555,12 +1839,17 @@ fn plan_replay( "--reprove-slot {target} selects a raw IxVM leaf, not a Stage 2 proof" )); } + if spec.shape.is_none() { + return Err(format!( + "--reprove-slot {target} selects an imported healed leaf; it has no Stage 2 execution to replay" + )); + } let children = match spec.op { PlanOp::Leaf(_) => Vec::new(), PlanOp::Join(left, right) => vec![left, right], }; let needs_input_proofs = children.is_empty() - || children.iter().any(|index| specs[*index].kind == ChildKind::Ixvm); + || children.iter().any(|index| specs[*index].shape.is_none()); Ok(ReplayPlan { children, needs_input_proofs }) } @@ -1573,8 +1862,10 @@ fn load_replay_child( .specs .get(child_index) .ok_or("replay target has a missing child slot")?; - if spec.kind == ChildKind::Ixvm { - return prove_slot(ctx, child_index, &[]); + if spec.kind == ChildKind::Ixvm + || (matches!(spec.op, PlanOp::Leaf(_)) && spec.shape.is_none()) + { + return prove_slot(ctx, child_index, &[], None); } let (proof, proof_address) = load_cached(ctx, child_index, spec).ok_or_else(|| { format!( @@ -1607,7 +1898,7 @@ fn run_replay( .map(|child| load_replay_child(ctx, target, *child)) .collect::>()?; let children_loaded_at = Instant::now(); - let slot = prove_slot(ctx, target, &children)?; + let slot = prove_slot(ctx, target, &children, None)?; ctx.aggr_system.verify(&slot.outer_claim, &slot.proof).map_err(|error| { format!("replayed slot {target} proof failed verification: {error:?}") })?; @@ -1639,26 +1930,495 @@ fn dependencies_complete(spec: &SlotSpec, completed: &[bool]) -> bool { } } -fn run_scheduler( - ctx: ProveContext<'_>, +/// One NUMA domain used as a scheduling lane: a pinned rayon pool plus its +/// own RAM reservation (see `crate::numa`). Slots proved on a lane run with +/// their threads and first-touch memory confined to that domain. +struct NumaLane { + domain: crate::numa::Domain, + pool: Arc, + budget: usize, + reserved: usize, + active: usize, + /// Highest resident memory observed on this lane's node at any slot + /// completion (observability, not accounting). + peak_resident: usize, +} + +/// Resident bytes on `node` right now (0 when unreadable). +fn resident_on(node: u32) -> usize { + crate::numa::resident_by_node() + .into_iter() + .find(|(n, _)| *n == node) + .map_or(0, |(_, bytes)| bytes) +} + +/// Run `work` while a sampler thread reads this process's resident memory on +/// `node` once a second; returns the work's result and the peak seen. +/// Observability only (a `/proc/self/numa_maps` parse per sample). +fn with_resident_peak( + node: u32, + work: impl FnOnce() -> R + Send, +) -> (R, usize) { + let stop = std::sync::atomic::AtomicBool::new(false); + thread::scope(|scope| { + let sampler = scope.spawn(|| { + let mut peak = 0usize; + while !stop.load(std::sync::atomic::Ordering::Relaxed) { + peak = peak.max(resident_on(node)); + thread::sleep(std::time::Duration::from_secs(1)); + } + peak.max(resident_on(node)) + }); + let result = work(); + stop.store(true, std::sync::atomic::Ordering::Relaxed); + let peak = sampler.join().unwrap_or(0); + (result, peak) + }) +} + +/// Pick the lane for a slot of `weight` bytes: an idle lane with the most free +/// RAM, else (when packing is allowed) the least-loaded lane that fits, at +/// most two slots per lane. `None` when no lane can take it now — the caller +/// then waits, or runs the slot unpinned if nothing else is in flight (the +/// over-budget-runs-alone rule). +fn choose_numa_lane( + lanes: &[NumaLane], + weight: usize, + pack: bool, +) -> Option { + let mut best: Option<((usize, usize), usize)> = None; + for (index, lane) in lanes.iter().enumerate() { + if weight > lane.budget.saturating_sub(lane.reserved) { + continue; + } + if lane.active > 0 && (!pack || lane.active >= 2) { + continue; + } + let free = lane.budget - lane.reserved; + let key = (lane.active, usize::MAX - free); + if best.is_none_or(|(k, _)| key < k) { + best = Some((key, index)); + } + } + best.map(|(_, index)| index) +} + +fn numa_lanes(budget: usize) -> Result, String> { + let numa = crate::numa::detect(); + if !numa.enabled() { + eprintln!( + "[aggregate] numa: disabled (single domain, IX_NUMA=off, or unsupported)" + ); + return Ok(Vec::new()); + } + let mut lanes = Vec::with_capacity(numa.domains.len()); + for domain in &numa.domains { + let pool = crate::numa::pool(numa, domain)?; + let lane_budget = (domain.mem_bytes / 10 * 9).min(budget); + lanes.push(NumaLane { + domain: domain.clone(), + pool, + budget: lane_budget, + reserved: 0, + active: 0, + peak_resident: 0, + }); + } + let described: Vec = lanes + .iter() + .map(|lane| { + format!( + "node {} ({} cpus, {} GiB)", + lane.domain.node, + lane.domain.cpus.len(), + format_gib(lane.budget) + ) + }) + .collect(); + eprintln!( + "[aggregate] numa: {} lanes: {}; policy={:?} pack={} threads={}", + lanes.len(), + described.join(", "), + numa.policy, + numa.pack, + numa.threads.map_or("cpuset".to_string(), |n| n.to_string()), + ); + Ok(lanes) +} + +#[derive(Debug)] +struct PipelineQueue { + /// None uses the ordinary Rayon pool and inherits the process affinity. + lane: Option, + slots: Vec, + /// Largest proving reservation in this queue, excluding lookahead. + weight: usize, + lookahead: bool, +} + +/// Partition independent jobs without exceeding the job, process or node +/// limits. Reserve proving capacity first, then enable at most one prepared +/// record per queue from the remaining RAM. This keeps lookahead from +/// reducing the number of concurrent proofs. Unassigned jobs stay on the +/// ordinary scheduler, including its existing over-budget-runs-alone path. +fn plan_pipelines( + batch: &[(usize, usize)], + lane_budgets: &[usize], + max_jobs: usize, + budget: usize, + per_lane: usize, +) -> Vec { + let mut ordered = batch.to_vec(); + ordered.sort_by_key(|&(index, weight)| (std::cmp::Reverse(weight), index)); + let mut pipelines: Vec = Vec::new(); + let mut load = Vec::::new(); + let mut reserved = 0usize; + // Proving reservation (plus enabled lookahead records) per NUMA lane; up + // to `per_lane` queues share a lane's pool when both fit its budget. + let mut lane_load = vec![0usize; lane_budgets.len()]; + let lane_count = |pipelines: &[PipelineQueue], k: usize| { + pipelines.iter().filter(|p| p.lane == Some(k)).count() + }; + for (index, weight) in ordered { + if weight == 0 { + continue; + } + if pipelines.len() < max_jobs && weight <= budget - reserved { + let placement = if lane_budgets.is_empty() { + Some(None) + } else { + (0..lane_budgets.len()) + .filter(|&k| { + lane_count(&pipelines, k) < per_lane.max(1) + && weight <= lane_budgets[k].saturating_sub(lane_load[k]) + }) + // An idle lane first, then the emptiest; ties to the lowest node. + .min_by_key(|&k| (lane_count(&pipelines, k), lane_load[k], k)) + .map(Some) + }; + if let Some(lane) = placement { + if let Some(k) = lane { + lane_load[k] += weight; + } + pipelines.push(PipelineQueue { + lane, + slots: vec![index], + weight, + lookahead: false, + }); + load.push(weight); + reserved += weight; + continue; + } + } + if let Some(k) = (0..pipelines.len()) + .filter(|&k| weight <= pipelines[k].weight) + .min_by_key(|&k| (load[k], k)) + { + pipelines[k].slots.push(index); + load[k] = load[k].saturating_add(weight); + } + } + for pipeline in &mut pipelines { + pipeline.slots.sort_unstable(); + let lane_room = pipeline + .lane + .map_or(usize::MAX, |k| lane_budgets[k].saturating_sub(lane_load[k])); + if pipeline.slots.len() > 1 + && LOOKAHEAD_RAM_BYTES <= budget - reserved + && LOOKAHEAD_RAM_BYTES <= lane_room + { + pipeline.lookahead = true; + reserved += LOOKAHEAD_RAM_BYTES; + if let Some(k) = pipeline.lane { + lane_load[k] += LOOKAHEAD_RAM_BYTES; + } + } + } + pipelines +} + +/// Prove the initial independent jobs on bounded queues, overlapping the +/// next preparation where the plan has reserved room for its record. NUMA +/// placement is optional; dependent joins use the ordinary scheduler after +/// this batch. Returns completed slots and peaks keyed by NUMA lane index. +fn run_pipelines<'a>( + ctx: ProveContext<'a>, + lanes: &[NumaLane], + slots: &[Option>], + pipelines: &[PipelineQueue], +) -> Result<(Vec<(usize, Arc)>, Vec<(usize, usize)>), String> { + let numa = crate::numa::detect(); + eprintln!( + "[aggregate] pipelines: {} independent slots over {} workers ({}); prepare-next overlap on {}/{} workers", + pipelines.iter().map(|p| p.slots.len()).sum::(), + pipelines.len(), + pipelines + .iter() + .enumerate() + .map(|(i, p)| { + let placement = p.lane.map_or_else( + || format!("unpinned worker {i}"), + |k| format!("node {}", lanes[k].domain.node), + ); + format!( + "{placement}: {} slots, {} GiB reserved", + p.slots.len(), + format_gib( + p.weight + if p.lookahead { LOOKAHEAD_RAM_BYTES } else { 0 } + ) + ) + }) + .collect::>() + .join(", "), + pipelines.iter().filter(|p| p.lookahead).count(), + pipelines.len(), + ); + let children_of = |index: usize| -> Vec> { + match ctx.specs[index].op { + PlanOp::Leaf(_) => Vec::new(), + PlanOp::Join(left, right) => vec![ + slots[left].as_ref().expect("completed left slot").clone(), + slots[right].as_ref().expect("completed right slot").clone(), + ], + } + }; + let prepare = |index: usize, + children: &[Arc]| + -> Result>, String> { + let spec = &ctx.specs[index]; + match spec.op { + PlanOp::Join(..) => { + prepare_aggregate(ctx, spec, &children[0], Some(&children[1]), index) + }, + PlanOp::Leaf(shard) => { + let raw = ctx + .proofs + .and_then(|proofs| proofs.get(shard)) + .ok_or("shard proof was not loaded")?; + prepare_aggregate(ctx, spec, raw, None, index) + }, + } + }; + let results: Vec)>, usize), String>> = + thread::scope(|scope| { + let handles: Vec<_> = pipelines + .iter() + .enumerate() + .map(|(worker, pipeline)| { + let children_of = &children_of; + let prepare = &prepare; + let lane = pipeline.lane.map(|k| &lanes[k]); + scope.spawn(move || { + let run = || -> Result<(Vec<(usize, Arc)>, usize), String> { + let queue = &pipeline.slots; + let mut done = Vec::with_capacity(queue.len()); + let mut peak_resident = 0usize; + let mut prepared: Option<( + usize, + Result>, String>, + )> = None; + for (position, &index) in queue.iter().enumerate() { + let started = Instant::now(); + let children = children_of(index); + // This slot's record: prepared during the previous prove, or + // built now for the first slot of the queue. + let execution = match prepared.take() { + Some((prepared_index, result)) if prepared_index == index => { + result? + }, + _ => prepare(index, &children)?, + }; + let next = if pipeline.lookahead { + queue.get(position + 1).copied() + } else { + None + }; + let (proved, node_peak, next_prepared) = thread::scope(|inner| { + let producer = next.map(|next_index| { + let next_children = children_of(next_index); + inner.spawn(move || { + if let Some(lane) = lane { + crate::numa::pin_current_thread(&lane.domain, numa.policy); + } + (next_index, prepare(next_index, &next_children)) + }) + }); + let prove = || prove_slot(ctx, index, &children, execution); + let (proved, node_peak) = match lane { + Some(lane) => with_resident_peak(lane.domain.node, prove), + None => (prove(), 0), + }; + let next_prepared = producer.map(|handle| { + handle.join().unwrap_or_else(|payload| { + ( + next.expect("producer exists only with a next slot"), + Err(format!( + "preparation panicked: {}", + panic_text(&payload) + )), + ) + }) + }); + (proved, node_peak, next_prepared) + }); + prepared = next_prepared; + let slot = proved?; + let resident = node_peak; + peak_resident = peak_resident.max(resident); + let placement = lane.map_or_else( + || format!("unpinned worker {worker}"), + |lane| format!( + "node {} (node peak resident {} GiB)", + lane.domain.node, format_gib(resident) + ), + ); + eprintln!( + "[aggregate] slot {index}: completed in {:.1}s on {placement} ({} GiB weight); pipeline {}/{}", + started.elapsed().as_secs_f64(), + format_gib(ctx.specs[index].ram_bytes), + position + 1, + queue.len(), + ); + done.push((index, slot)); + } + Ok((done, peak_resident)) + }; + match lane { + Some(lane) => { + crate::numa::pin_current_thread(&lane.domain, numa.policy); + lane.pool.install(run) + }, + None => run(), + } + }) + }) + .collect(); + handles + .into_iter() + .map(|h| { + h.join().unwrap_or_else(|payload| { + Err(format!( + "aggregate pipeline panicked: {}", + panic_text(&payload) + )) + }) + }) + .collect() + }); + let mut completed = + Vec::with_capacity(pipelines.iter().map(|p| p.slots.len()).sum()); + let mut peaks = Vec::new(); + for (pipeline, result) in pipelines.iter().zip(results) { + let (done, peak) = result?; + completed.extend(done); + if let Some(k) = pipeline.lane { + peaks.push((k, peak)); + } + } + Ok((completed, peaks)) +} + +fn run_scheduler<'a>( + ctx: ProveContext<'a>, jobs: usize, budget: usize, ) -> Result>, String> { if budget == 0 { return Err("aggregate scheduler RAM budget must be positive".into()); } - let max_jobs = if jobs == 0 { ctx.specs.len().max(1) } else { jobs.max(1) }; + // Adapt to the cgroup this process was launched in: never admit more than + // 92 % of its memory limit, whatever `--max-ram` or the default said. + let budget = match crate::numa::cgroup_memory_max() { + Some(limit) if limit / 100 * 92 < budget => { + let clamped = limit / 100 * 92; + eprintln!( + "[aggregate] RAM budget {} GiB exceeds 92 % of the cgroup limit {} GiB; clamping to {} GiB", + format_gib(budget), + format_gib(limit), + format_gib(clamped), + ); + clamped + }, + _ => budget, + }; + let numa = crate::numa::detect(); + let mut lanes = numa_lanes(budget)?; + // Prepare-next overlap is independent of placement. The existing switch + // also controls unpinned workers on single-node or unsupported hosts. + let lookahead = !matches!( + std::env::var("IX_NUMA_LOOKAHEAD").as_deref().map(str::trim), + Ok("0" | "off" | "false") + ); + let max_jobs = if jobs == 0 { + if lanes.is_empty() { + ctx.specs.len().max(1) + } else { + lanes.len() * if numa.pack { 2 } else { 1 } + } + } else { + jobs.max(1) + }; + let n = ctx.specs.len(); + let mut slots: Vec>> = vec![None; n]; + let mut completed = vec![false; n]; + let mut completed_count = 0usize; + + // Imported raw leaves (no shape) complete without proving: do them inline. + for index in 0..n { + if ctx.specs[index].shape.is_none() + && matches!(ctx.specs[index].op, PlanOp::Leaf(_)) + { + let slot = prove_slot(ctx, index, &[], None) + .map_err(|error| format!("slot {index}: {error}"))?; + slots[index] = Some(slot); + completed[index] = true; + completed_count += 1; + } + } + + // Phase A: everything ready now has no dependency on another proof. + if lookahead { + let batch: Vec<(usize, usize)> = (0..n) + .filter(|&i| { + !completed[i] + && ctx.specs[i].shape.is_some() + && dependencies_complete(&ctx.specs[i], &completed) + }) + .map(|i| (i, ctx.specs[i].ram_bytes)) + .collect(); + let lane_budgets: Vec = lanes.iter().map(|l| l.budget).collect(); + let pipelines = plan_pipelines( + &batch, + &lane_budgets, + max_jobs, + budget, + if numa.pack { 2 } else { 1 }, + ); + // Without room or enough work to overlap, retain ordinary dynamic + // admission and packing instead of introducing a batch barrier. + if pipelines.iter().any(|p| p.lookahead) { + let (done, peaks) = run_pipelines(ctx, &lanes, &slots, &pipelines)?; + for (index, slot) in done { + slots[index] = Some(slot); + completed[index] = true; + completed_count += 1; + } + for (k, peak) in peaks { + lanes[k].peak_resident = lanes[k].peak_resident.max(peak); + } + } + } + let (sender, receiver) = mpsc::channel(); thread::scope(|scope| -> Result>, String> { - let mut slots: Vec>> = vec![None; ctx.specs.len()]; - let mut completed = vec![false; ctx.specs.len()]; - let mut in_flight = vec![false; ctx.specs.len()]; - let mut completed_count = 0usize; + let mut in_flight = vec![false; n]; + let mut admitted_at: Vec> = vec![None; n]; let mut active = 0usize; let mut reserved = 0usize; let mut failures: Vec<(usize, String)> = Vec::new(); - while completed_count < ctx.specs.len() { + while completed_count < n { if failures.is_empty() && active < max_jobs { let mut ready: Vec = ctx .specs @@ -1677,15 +2437,30 @@ fn run_scheduler( .cmp(&ctx.specs[*left].ram_bytes) .then_with(|| left.cmp(right)) }); + let ready_count = ready.len(); for index in ready { if active >= max_jobs { break; } let weight = ctx.specs[index].ram_bytes; - let fits = reserved.saturating_add(weight) <= budget; + let fits = weight <= budget.saturating_sub(reserved); if !fits && active != 0 { continue; } + // A slot that is the only runnable work with nothing else live + // (the dependency tail near the root) is faster unpinned: it can + // use every core and all memory channels (measured ~12 % over one + // domain), and there is no neighbour to isolate it from. + let solo_tail = active == 0 && ready_count == 1; + let lane = if solo_tail { + None + } else { + choose_numa_lane(&lanes, weight, numa.pack) + }; + if !lanes.is_empty() && lane.is_none() && active != 0 { + // Fits the global budget but no domain can hold it yet. + continue; + } let children = match ctx.specs[index].op { PlanOp::Leaf(_) => Vec::new(), PlanOp::Join(left, right) => vec![ @@ -1694,30 +2469,74 @@ fn run_scheduler( ], }; in_flight[index] = true; + admitted_at[index] = Some(Instant::now()); active += 1; reserved = reserved.saturating_add(weight); let over = if weight > budget { "; over-budget slot runs alone" } else { "" }; + let placement = match lane { + Some(k) => { + lanes[k].reserved = lanes[k].reserved.saturating_add(weight); + lanes[k].active += 1; + format!( + " on node {} (node reserved {}/{} GiB, node active {})", + lanes[k].domain.node, + format_gib(lanes[k].reserved), + format_gib(lanes[k].budget), + lanes[k].active, + ) + }, + None if !lanes.is_empty() && solo_tail => { + " unpinned (solo tail)".to_string() + }, + None if !lanes.is_empty() => " unpinned".to_string(), + None => String::new(), + }; eprintln!( - "[aggregate] slot {index}: admitted {} GiB; reserved {}/{} GiB; active {active}/{max_jobs}{over}", + "[aggregate] slot {index}: admitted {} GiB{placement}; reserved {}/{} GiB; active {active}/{max_jobs}{over}", format_gib(weight), format_gib(reserved), format_gib(budget), ); let sender = sender.clone(); + let pinned = + lane.map(|k| (lanes[k].pool.clone(), lanes[k].domain.clone())); + let unpin = lane.is_none() && !lanes.is_empty(); scope.spawn(move || { - let result = - std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { - prove_slot(ctx, index, &children) - })) - .unwrap_or_else(|payload| { - Err(format!( - "Rust proof worker panicked: {}", - panic_text(&payload) - )) + let node = pinned.as_ref().map(|(_, domain)| domain.node); + let (result, node_peak) = + with_resident_peak(node.unwrap_or(0), || { + std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + match &pinned { + Some((pool, domain)) => { + crate::numa::pin_current_thread(domain, numa.policy); + pool.install(|| prove_slot(ctx, index, &children, None)) + }, + None => { + if unpin { + crate::numa::unpin_current_thread(numa); + } + prove_slot(ctx, index, &children, None) + }, + } + })) + .unwrap_or_else(|payload| { + Err(format!( + "Rust proof worker panicked: {}", + panic_text(&payload) + )) + }) }); - let _ = sender.send((index, weight, result)); + let node_peak = if node.is_some() { node_peak } else { 0 }; + let _ = sender.send((index, weight, lane, node_peak, result)); }); + if unpin { + // No other slot is active. Wait for this one before admitting + // neighbours: a slot too large for one node uses all nodes and + // holds no per-node reservation, even if the global budget has + // room left. The next completion is necessarily this slot's. + break; + } } } @@ -1729,9 +2548,10 @@ fn run_scheduler( break; } - let (index, weight, result) = receiver.recv().map_err(|error| { - format!("aggregate scheduler channel closed: {error}") - })?; + let (index, weight, lane, node_peak, result) = + receiver.recv().map_err(|error| { + format!("aggregate scheduler channel closed: {error}") + })?; if !in_flight.get(index).copied().unwrap_or(false) { failures.push((index, "duplicate or unknown scheduler result".into())); continue; @@ -1739,24 +2559,54 @@ fn run_scheduler( in_flight[index] = false; active -= 1; reserved = reserved.saturating_sub(weight); + if let Some(k) = lane { + lanes[k].reserved = lanes[k].reserved.saturating_sub(weight); + lanes[k].active = lanes[k].active.saturating_sub(1); + } + let elapsed = admitted_at[index] + .take() + .map_or(0.0, |started| started.elapsed().as_secs_f64()); + let where_ = lane.map_or(String::new(), |k| { + lanes[k].peak_resident = lanes[k].peak_resident.max(node_peak); + format!( + " on node {} (node peak resident {} GiB)", + lanes[k].domain.node, + format_gib(node_peak) + ) + }); match result { Ok(slot) => { + eprintln!( + "[aggregate] slot {index}: completed in {elapsed:.1}s{where_} ({} GiB weight); active {}/{max_jobs}", + format_gib(weight), + active, + ); slots[index] = Some(slot); completed[index] = true; completed_count += 1; }, - Err(error) => failures.push((index, error)), + Err(error) => { + eprintln!( + "[aggregate] slot {index}: FAILED after {elapsed:.1}s{where_}" + ); + failures.push((index, error)) + }, } } while active > 0 { - let (index, weight, result) = receiver.recv().map_err(|error| { - format!("aggregate scheduler drain failed: {error}") - })?; + let (index, weight, lane, _node_peak, result) = + receiver.recv().map_err(|error| { + format!("aggregate scheduler drain failed: {error}") + })?; if in_flight.get(index).copied().unwrap_or(false) { in_flight[index] = false; active -= 1; reserved = reserved.saturating_sub(weight); + if let Some(k) = lane { + lanes[k].reserved = lanes[k].reserved.saturating_sub(weight); + lanes[k].active = lanes[k].active.saturating_sub(1); + } } match result { Ok(slot) => { @@ -1775,6 +2625,21 @@ fn run_scheduler( error }); } + if !lanes.is_empty() { + eprintln!( + "[aggregate] lane peaks (resident at slot completions): {}", + lanes + .iter() + .map(|lane| format!( + "node {}: {} GiB of {} GiB", + lane.domain.node, + format_gib(lane.peak_resident), + format_gib(lane.domain.mem_bytes) + )) + .collect::>() + .join(", ") + ); + } slots .into_iter() .enumerate() @@ -1794,16 +2659,21 @@ fn print_plan( specs.iter().filter(|spec| matches!(spec.op, PlanOp::Leaf(_))).count(); let wraps = specs .iter() - .filter(|spec| { - matches!(spec.op, PlanOp::Leaf(_)) && spec.kind == ChildKind::Aggr - }) + .filter(|spec| matches!(spec.op, PlanOp::Leaf(_)) && spec.shape.is_some()) .count(); let structural = specs.iter().filter(|spec| spec.structural).count(); - let policy = if wraps == leaves { - format!("{wraps} wraps") - } else { - format!("{} direct IxVM leaves", leaves - wraps) - }; + let imported = specs + .iter() + .filter(|s| { + matches!(s.op, PlanOp::Leaf(_)) + && s.kind == ChildKind::Aggr + && s.shape.is_none() + }) + .count(); + let policy = format!( + "{wraps} wraps, {imported} imported healed leaves, {} direct IxVM leaves", + leaves - wraps - imported + ); eprintln!( "[aggregate] plan: {policy} + {} binary joins ({structural} structural; threshold > {threshold} subject leaves)", specs.len() - leaves @@ -1811,8 +2681,13 @@ fn print_plan( for (index, spec) in specs.iter().enumerate() { match spec.op { PlanOp::Leaf(shard) => { - let mode = - if spec.kind == ChildKind::Ixvm { "raw shard" } else { "wrap shard" }; + let mode = if spec.kind == ChildKind::Ixvm { + "raw shard" + } else if spec.shape.is_none() { + "healed shard" + } else { + "wrap shard" + }; eprintln!( " slot {index}: {mode} {} ({} subjects)", prepared[shard].original_id, spec.subject_count @@ -1859,7 +2734,7 @@ fn run(config: RunConfig<'_>) -> Result { .map_err(|error| format!("ixAggr VK serialization failed: {error}"))?; let allowed = allowed_blob(&ixvm_vk, config.verify_idx, &aggr_vk, config.aggr_idx); - let specs = build_specs( + let mut specs = build_specs( &prepared, config.verify_idx, config.aggr_idx, @@ -1869,13 +2744,13 @@ fn run(config: RunConfig<'_>) -> Result { &allowed, config.cache_fri_bytes, )?; - let replay_plan = config + let mut replay_plan = config .reprove_slot .map(|target| plan_replay(&specs, target)) .transpose()?; let specs_at = Instant::now(); - print_plan(&specs, &prepared.shards, config.structural_above); if config.plan_only { + print_plan(&specs, &prepared.shards, config.structural_above); eprintln!( "[aggregate] Rust plan startup: manifest {:.3}s, env/claims {:.3}s, plan/statements {:.3}s; total {:.3}s", (parsed_at - started).as_secs_f64(), @@ -1908,10 +2783,43 @@ fn run(config: RunConfig<'_>) -> Result { if !config.write_outputs { eprintln!("[aggregate] output writes disabled (--no-write)"); } - let needs_input_proofs = - replay_plan.as_ref().is_none_or(|plan| plan.needs_input_proofs); + let needs_input_proofs = replay_plan.as_ref().is_none_or(|plan| { + plan.needs_input_proofs + || (!config.proof_hexes.trim().is_empty() + && plan + .children + .iter() + .any(|index| matches!(specs[*index].op, PlanOp::Leaf(_)))) + }); let proofs = if needs_input_proofs { - Some(load_input_proofs(config.proof_hexes, &store_dir, &prepared.shards)?) + let wrappers = load_input_proofs( + config.proof_hexes, + &store_dir, + &prepared.shards, + prepared.partial, + )?; + let inputs = wrappers + .into_par_iter() + .enumerate() + .map(|(index, wrapper)| { + import_shard_proof( + config.ixvm_system, + config.aggr_system, + config.verify_idx, + config.aggr_idx, + &allowed, + prepared.shards[index].statement.clone(), + &wrapper, + None, + ) + }) + .collect::, _>>()?; + bind_imported_specs(&mut specs, &inputs, &aggr_vk, config.cache_fri_bytes); + replay_plan = config + .reprove_slot + .map(|target| plan_replay(&specs, target)) + .transpose()?; + Some(inputs) } else { let supplied = config.proof_hexes.lines().filter(|line| !line.is_empty()).count(); @@ -1921,6 +2829,7 @@ fn run(config: RunConfig<'_>) -> Result { None }; let proofs_at = Instant::now(); + print_plan(&specs, &prepared.shards, config.structural_above); eprintln!( "[aggregate] Rust startup: manifest {:.3}s, env/claims {:.3}s, plan/statements {:.3}s, proofs {:.3}s; total {:.3}s", (parsed_at - started).as_secs_f64(), @@ -1939,7 +2848,6 @@ fn run(config: RunConfig<'_>) -> Result { ixvm_vk: &ixvm_vk, aggr_vk: &aggr_vk, allowed: &allowed, - verify_idx: config.verify_idx, aggr_idx: config.aggr_idx, store_dir: &store_dir, cache_dir, @@ -1958,8 +2866,13 @@ fn run(config: RunConfig<'_>) -> Result { config.jobs.to_string() }; eprintln!( - "[aggregate] scheduler: jobs={jobs_label}, RAM budget {} GiB; wrap/self 195.0 GiB, direct 390.0 GiB, mixed 340.0 GiB, flat +1 MiB/subject", - format_gib(config.ram_budget_bytes) + "[aggregate] scheduler: jobs={jobs_label}, RAM budget {} GiB; wrap/self base {} GiB, direct {} GiB, mixed {} GiB, flat self +1 MiB/subject, structural self +1.25 MiB/subject (minimum {} GiB above {} subjects)", + format_gib(config.ram_budget_bytes), + format_gib(STRUCTURAL_RAM_BYTES), + format_gib(DIRECT_RAM_BYTES), + format_gib(MIXED_RAM_BYTES), + format_gib(2 * STRUCTURAL_RAM_BYTES), + STRUCTURAL_LARGE_SUBJECTS, ); let slots = run_scheduler(context, config.jobs, config.ram_budget_bytes)?; let root = slots.last().ok_or("aggregate plan produced no root slot")?; @@ -2083,6 +2996,250 @@ extern "C" fn rs_aiur_stage2_aggregate( #[cfg(test)] mod tests { use super::*; + + #[test] + fn structural_ram_covers_measured_mathlib_peaks() { + // Subject counts and query-record peaks (rounded up to GiB) from the + // 2026-09-09 Mathlib run, including the slot behind the packed-node OOM. + for (subjects, peak_gib) in [ + (5_371, 196), + (11_972, 203), + (19_751, 208), + (55_496, 212), + (91_068, 381), + (91_620, 257), + (96_048, 249), + (126_527, 381), + (187_668, 384), + (314_195, 455), + ] { + assert!(shape_ram_bytes(9, subjects) >= peak_gib * GIB); + } + // The new term belongs to structural self-pairs. Keep flat, direct, + // mixed and wrap reservations distinct. + for (shape, gib) in [(0, 195), (2, 180), (5, 199), (8, 180), (9, 200)] { + assert_eq!(shape_ram_bytes(shape, 4096), gib * GIB); + } + assert_eq!(shape_ram_bytes(9, 65_536), 275 * GIB); + assert_eq!(shape_ram_bytes(9, 65_537), 390 * GIB); + } + + /// The pre-function-groups direct-pair reservation the pipeline + /// arithmetic below was written against (one queue per 453 GiB lane). + const TEST_DIRECT: usize = 390 * GIB; + + fn direct_batch() -> Vec<(usize, usize)> { + (0..6).map(|i| (i, TEST_DIRECT)).collect() + } + + #[test] + fn pipelines_pack_two_queues_per_lane_when_both_fit() { + // Six 180 GiB direct joins over three 453 GiB lanes: two queues per + // lane, each with its 40 GiB lookahead record (2 x 220 <= 453; the + // process budget must hold all six, 6 x 220 = 1320 GiB). + let batch: Vec<(usize, usize)> = + (0..12).map(|i| (i, DIRECT_RAM_BYTES)).collect(); + let packed = plan_pipelines(&batch, &[453 * GIB; 3], 6, 1400 * GIB, 2); + assert_eq!(packed.len(), 6); + assert!(packed.iter().all(|p| p.lookahead && p.slots.len() == 2)); + for k in 0..3 { + assert_eq!(packed.iter().filter(|p| p.lane == Some(k)).count(), 2); + } + assert_eq!(pipeline_reservation(&packed), 6 * 220 * GIB); + // One queue per lane keeps the previous placement. + let single = plan_pipelines(&batch, &[453 * GIB; 3], 6, 1400 * GIB, 1); + assert_eq!(single.len(), 3); + assert!(single.iter().all(|p| p.lookahead && p.slots.len() == 4)); + // A second queue that would not fit beside the first stays off the lane. + let tight = plan_pipelines(&batch, &[300 * GIB; 3], 6, 1400 * GIB, 2); + assert_eq!(tight.len(), 3); + // With 1300 GiB the sixth queue has no room for its record. + let capped = plan_pipelines(&batch, &[453 * GIB; 3], 6, 1300 * GIB, 2); + assert_eq!(capped.len(), 6); + assert_eq!(capped.iter().filter(|p| p.lookahead).count(), 5); + } + + fn pipeline_reservation(pipelines: &[PipelineQueue]) -> usize { + pipelines + .iter() + .map(|p| p.weight + if p.lookahead { LOOKAHEAD_RAM_BYTES } else { 0 }) + .sum() + } + + #[test] + fn pipelines_obey_jobs_one_on_multiple_nodes() { + let pipelines = + plan_pipelines(&direct_batch(), &[453 * GIB; 3], 1, 1300 * GIB, 1); + assert_eq!(pipelines.len(), 1); + assert_eq!(pipelines[0].slots, (0..6).collect::>()); + assert!(pipelines[0].lookahead); + assert_eq!(pipeline_reservation(&pipelines), 430 * GIB); + } + + #[test] + fn pipelines_obey_combined_budget_including_lookahead() { + for (budget_gib, overlaps) in [(800, 0), (820, 1), (860, 2)] { + let pipelines = plan_pipelines( + &direct_batch(), + &[453 * GIB; 3], + 6, + budget_gib * GIB, + 1, + ); + // Keep two concurrent provers even when neither can prepare ahead. + assert_eq!(pipelines.len(), 2); + assert_eq!(pipelines.iter().filter(|p| p.lookahead).count(), overlaps); + assert!(pipeline_reservation(&pipelines) <= budget_gib * GIB); + let mut assigned: Vec<_> = + pipelines.iter().flat_map(|p| p.slots.iter().copied()).collect(); + assigned.sort_unstable(); + assert_eq!(assigned, (0..6).collect::>()); + } + } + + #[test] + fn pipelines_overlap_without_numa() { + let pipelines = plan_pipelines(&direct_batch(), &[], 3, 1300 * GIB, 1); + assert_eq!(pipelines.len(), 3); + assert!(pipelines.iter().all(|p| p.lane.is_none() && p.lookahead)); + assert!(pipelines.iter().all(|p| p.slots.len() == 2)); + assert_eq!(pipeline_reservation(&pipelines), 1290 * GIB); + + let single = plan_pipelines(&direct_batch(), &[], 1, 430 * GIB, 1); + assert_eq!(single.len(), 1); + assert!(single[0].lane.is_none() && single[0].lookahead); + } + + #[test] + fn pipelines_require_local_room_for_the_next_record() { + let pipelines = plan_pipelines( + &direct_batch(), + &[390 * GIB, 453 * GIB], + 2, + 1000 * GIB, + 1, + ); + assert_eq!(pipelines.len(), 2); + let tight = pipelines.iter().find(|p| p.lane == Some(0)).unwrap(); + let roomy = pipelines.iter().find(|p| p.lane == Some(1)).unwrap(); + assert!(!tight.lookahead); + assert!(roomy.lookahead); + assert_eq!(pipeline_reservation(&pipelines), 820 * GIB); + } + + #[test] + fn pipelines_handle_mixed_weights_and_leave_oversized_jobs() { + let batch = vec![ + (0, 500 * GIB), + (1, TEST_DIRECT), + (2, TEST_DIRECT), + (3, STRUCTURAL_RAM_BYTES), + ]; + let pipelines = + plan_pipelines(&batch, &[453 * GIB, 220 * GIB], 3, 700 * GIB, 1); + assert_eq!(pipelines.len(), 2); + assert_eq!(pipelines[0].slots, vec![1, 2]); + assert_eq!(pipelines[1].slots, vec![3]); + assert!(pipelines[0].lookahead); + assert!(!pipelines[1].lookahead); + assert_eq!(pipeline_reservation(&pipelines), 625 * GIB); + } + + #[test] + fn pipelines_skip_overlap_without_spare_ram_or_another_job() { + let full = plan_pipelines(&direct_batch(), &[], 1, TEST_DIRECT, 1); + assert_eq!(full.len(), 1); + assert!(!full[0].lookahead); + let single = plan_pipelines(&[(0, TEST_DIRECT)], &[], 1, 1300 * GIB, 1); + assert_eq!(single.len(), 1); + assert!(!single[0].lookahead); + assert!(plan_pipelines(&direct_batch(), &[], 3, 0, 1).is_empty()); + assert!(plan_pipelines(&direct_batch(), &[], 0, 1300 * GIB, 1).is_empty()); + } + + fn lane( + node: u32, + budget: usize, + reserved: usize, + active: usize, + ) -> NumaLane { + NumaLane { + domain: crate::numa::Domain { node, cpus: vec![0], mem_bytes: budget }, + pool: Arc::new( + rayon::ThreadPoolBuilder::new().num_threads(1).build().unwrap(), + ), + budget, + reserved, + active, + peak_resident: 0, + } + } + + #[test] + fn numa_lane_prefers_idle_then_most_free() { + let lanes = + vec![lane(0, 400, 200, 1), lane(1, 400, 0, 0), lane(2, 400, 100, 0)]; + // idle lanes 1 and 2 beat the busy lane 0; lane 1 has more free RAM. + assert_eq!(choose_numa_lane(&lanes, 195, true), Some(1)); + // a 390 slot fits only lane 1. + assert_eq!(choose_numa_lane(&lanes, 390, true), Some(1)); + } + + #[test] + fn numa_lane_packs_at_most_two_and_only_when_allowed() { + let lanes = vec![lane(0, 400, 195, 1), lane(1, 400, 390, 2)]; + assert_eq!(choose_numa_lane(&lanes, 195, true), Some(0)); + assert_eq!(choose_numa_lane(&lanes, 195, false), None); + // lane 1 already holds two slots: never a third even when packing. + let lanes = vec![lane(1, 900, 390, 2)]; + assert_eq!(choose_numa_lane(&lanes, 195, true), None); + } + + #[test] + fn numa_lane_none_when_nothing_fits() { + let lanes = vec![lane(0, 400, 0, 0)]; + assert_eq!(choose_numa_lane(&lanes, 401, true), None); + } + + #[test] + fn structural_packing_preserves_small_pairs_and_rejects_oom_pair() { + let budget = 453 * GIB; + for (left, right, fits) in [ + (9_480, 10_271, true), + (23_993, 24_805, true), + (91_620, 96_048, false), + // Mathlib slots 140/355 were live together when node 0 OOMed. + (187_668, 13_023, false), + ] { + let left = shape_ram_bytes(9, left); + let right = shape_ram_bytes(9, right); + // Either join fits individually, but packing depends on their sum. + assert!(left <= budget && right <= budget); + for (reserved, weight) in [(left, right), (right, left)] { + let lanes = vec![lane(0, budget, reserved, 1)]; + assert_eq!(choose_numa_lane(&lanes, weight, true), fits.then_some(0)); + } + } + } + + #[test] + fn large_structural_join_needs_unpinned_fallback() { + // Mathlib slot 266: larger than a node reservation, smaller than the + // process budget. The scheduler must wait and run it alone unpinned. + let weight = shape_ram_bytes(9, 314_195); + assert!(weight < 1300 * GIB); + let lanes = vec![lane(0, 453 * GIB, 0, 0), lane(1, 453 * GIB, 0, 0)]; + assert_eq!(choose_numa_lane(&lanes, weight, true), None); + } + + #[test] + fn structural_reservation_overflow_cannot_enable_packing() { + let weight = shape_ram_bytes(9, usize::MAX); + assert_eq!(weight, usize::MAX); + let lanes = vec![lane(0, usize::MAX, GIB, 1)]; + assert_eq!(choose_numa_lane(&lanes, weight, true), None); + } + use ix_kernel::shard::ShardInfo; use ixon::{Axiom, Expr}; @@ -2267,4 +3424,106 @@ mod tests { assert_eq!(paths[0].0, dependency); assert_eq!(paths[0].1.first(), Some(&1)); } + + /// Small real proofs for testing host transport and backend authentication. + /// Circuit semantics are exercised separately by the real IxVM smoke test. + pub(super) fn transport_system(input_size: usize) -> AiurSystem { + use aiur::bytecode::{ + Block, Circuit, Ctrl, Function, FunctionLayout, Toplevel, + }; + use multi_stark::types::{CommitmentParameters, FriParameters}; + let layout = + FunctionLayout { input_size, selectors: 1, auxiliaries: 1, lookups: 1 }; + AiurSystem::build( + Toplevel { + functions: vec![Function { + body: Block { ops: vec![], ctrl: Ctrl::Return(0, vec![]) }, + layout: layout.clone(), + entry: true, + constrained: true, + }], + memory_sizes: vec![], + // The singleton partition the Lean compiler emits by default. + circuits: vec![Circuit { members: vec![0], layout }], + }, + CommitmentParameters { log_blowup: 1, cap_height: 0 }, + FriParameters { + log_final_poly_len: 0, + max_log_arity: 1, + num_queries: 4, + commit_proof_of_work_bits: 0, + query_proof_of_work_bits: 0, + }, + ) + } + + #[test] + fn stage2_authenticates_healed_leaves_and_updates_mixed_shapes() { + let env = ixon::Env::new(); + let a = store_axiom(&env, Expr::sort(0), vec![]); + let b = store_axiom(&env, Expr::reference(0, vec![]), vec![a.clone()]); + let manifest = ShardManifest { + num_shards: 2, + shards: vec![shard(0, a), shard(1, b)], + total_cross_ingress: 0, + tree: None, + }; + let prepared = prepare_run(&env, &manifest).unwrap(); + let ixvm = transport_system(8); + let aggr = transport_system(16); + let ixvm_vk = aiur::vk_codec::aiur_system_to_bytes(&ixvm).unwrap(); + let aggr_vk = aiur::vk_codec::aiur_system_to_bytes(&aggr).unwrap(); + let allowed = allowed_blob(&ixvm_vk, 0, &aggr_vk, 0); + let mut specs = + build_specs(&prepared, 0, 0, 4096, true, &aggr_vk, &allowed, &[0; 40]) + .unwrap(); + let root_claim = specs[2].statement.claim_bytes.clone(); + let make_wrapper = |index: usize, healed| { + let statement = &prepared.shards[index].statement; + let outer = if healed { + aggregate_outer_claim(0, &allowed, &statement.claim_bytes) + } else { + inner_claim(0, &statement.claim_bytes) + }; + let mut io = + IOBuffer { data: FxHashMap::default(), map: FxHashMap::default() }; + let system = if healed { &aggr } else { &ixvm }; + let (_, proof) = system.prove(0, &outer[2..], &mut io); + IxonProof::new(statement.claim.clone(), proof.to_bytes().unwrap()) + }; + let healed = make_wrapper(0, true); + let raw = make_wrapper(1, false); + let import = |index: usize, wrapper: &IxonProof, identity: &[u8]| { + import_shard_proof( + &ixvm, + &aggr, + 0, + 0, + identity, + prepared.shards[index].statement.clone(), + wrapper, + None, + ) + }; + let inputs = vec![ + import(0, &healed, &allowed).unwrap(), + import(1, &raw, &allowed).unwrap(), + ]; + assert_eq!(inputs[0].kind, ChildKind::Aggr); + assert_eq!(inputs[1].kind, ChildKind::Ixvm); + bind_imported_specs(&mut specs, &inputs, &aggr_vk, &[0; 40]); + assert!(specs[0].shape.is_none()); + assert_eq!(specs[2].shape, Some(4)); + assert_eq!(specs[2].statement.claim_bytes, root_claim); + assert!( + plan_replay(&specs, 0).unwrap_err().contains("imported healed leaf") + ); + assert!(plan_replay(&specs, 2).unwrap().needs_input_proofs); + assert!(import(1, &healed, &allowed).is_err()); + let mut other_identity = allowed.clone(); + other_identity[0] ^= 1; + assert!(import(0, &healed, &other_identity).is_err()); + let malformed = IxonProof::new(healed.claim.clone(), vec![0xff; 3]); + assert!(import(0, &malformed, &allowed).is_err()); + } } diff --git a/crates/ffi/src/aiur/aggregate/shard_pipeline.rs b/crates/ffi/src/aiur/aggregate/shard_pipeline.rs new file mode 100644 index 000000000..c1b4fe0aa --- /dev/null +++ b/crates/ffi/src/aiur/aggregate/shard_pipeline.rs @@ -0,0 +1,1082 @@ +//! One proving consumer and at most one next execution in flight. Splits +//! are private balanced subtrees of flat joins: the published claim stays +//! byte-for-byte equal to the original manifest leaf. + +use super::*; +use crate::aiur::protocol::decode_addr_lists; +use aiur::{ + bytecode::{FunIdx, Toplevel}, + execute::{ExecError, QueryRecord}, +}; +use ixvm_codegen::{ + aiur_ixvm_runner::execute_ixvm, + aiur_ixvm_witness::build_shard_check_env_witness, +}; +use lean_ffi::object::LeanIOResult; +use std::collections::VecDeque; + +pub(super) type Executor = fn( + &Toplevel, + FunIdx, + Vec, + &mut IOBuffer, +) -> Result<(QueryRecord, Vec), ExecError>; + +// Scheduling state, not a second prover API. Reuse the native executor's +// record/output and move the record into the existing prove_from_execution. +pub(super) struct Execution<'a> { + system: &'a AiurSystem, + fun_idx: usize, + input: Vec, + io: IOBuffer, + record: QueryRecord, + output: Vec, +} + +impl<'a> Execution<'a> { + pub(super) fn new( + system: &'a AiurSystem, + fun_idx: usize, + input: Vec, + mut io: IOBuffer, + executor: Executor, + ) -> Result { + let _span = tracing::info_span!("aiur/execute_ixvm").entered(); + let (record, output) = + executor(system.toplevel(), fun_idx, input.clone(), &mut io) + .map_err(|e| format!("native execution failed: {e:?}"))?; + Ok(Self { system, fun_idx, input, io, record, output }) + } + + pub(super) fn peak(&self) -> usize { + self.system.peak_prove_bytes(&self.record).peak + } + + pub(super) fn prove(self) -> (Vec, AiurProof) { + let _span = tracing::info_span!("aiur/prove_prepared").entered(); + self.system.prove_from_execution( + self.fun_idx, + &self.input, + &self.io, + self.record, + &self.output, + ) + } +} + +#[derive(Clone)] +enum Job { + Shard { blocks: Vec
, owned: Vec
}, + Aggregate { left: Address, right: Option
}, +} + +#[derive(Clone)] +struct Work { + job: Job, + statement: Arc, + original: usize, + publish: bool, +} + +impl Work { + fn digest(&self) -> Address { + Address::hash(&self.statement.claim_bytes) + } + + fn ready(&self, slots: &FxHashMap>) -> bool { + match &self.job { + Job::Shard { .. } => true, + Job::Aggregate { left, right } => { + slots.contains_key(left) + && right.as_ref().is_none_or(|r| slots.contains_key(r)) + }, + } + } +} + +#[allow(clippy::large_enum_variant)] +enum Prepared<'a> { + Execute(Execution<'a>), + Reused(Arc), + Split(Vec>), + WrapChildren, +} + +struct Pipeline<'a> { + ctx: ProveContext<'a>, + env: &'a ixon::Env, + verify_idx: usize, + max_ram: usize, + index: Option<&'a Path>, + plans: &'a Path, + skip_proven: bool, + lookahead: bool, + keep_going: bool, +} + +fn shard_statement( + env: &ixon::Env, + owned: &[Address], +) -> Result, String> { + let mut sorted = owned.to_vec(); + sorted.sort_unstable(); + let (_, frontier) = ixon::shard_claim::shard_check_env_claim(env, owned) + .ok_or("shard owns no constants")?; + Ok(Statement::new( + SubjectTree::canonical(sorted, ShardSet(Vec::new()))?, + CanonicalTree::from_sorted(frontier)?, + )) +} + +/// Local children share an original manifest owner. Discharge by actual +/// subject membership, including at intermediate joins, not that owner id. +fn flat_join( + left: &Arc, + right: &Arc, +) -> Result, String> { + let subjects = SubjectTree::flat(&left.subjects, &right.subjects)?; + let leaves = &subjects + .canonical_tree() + .ok_or("healing requires canonical subjects")? + .leaves; + let remaining = merge_optional_sets( + left.assumptions.as_deref(), + right.assumptions.as_deref(), + ) + .into_iter() + .filter(|a| leaves.binary_search(a).is_err()) + .collect(); + Ok(Statement::new(subjects, CanonicalTree::from_sorted(remaining)?)) +} + +fn validate_parts( + blocks: &[Address], + parts: &[Vec
], +) -> Result<(), String> { + if parts.len() < 2 || parts.iter().any(Vec::is_empty) { + return Err("split must have at least two nonempty parts".into()); + } + let mut actual: Vec<_> = parts.iter().flatten().cloned().collect(); + let mut expected = blocks.to_vec(); + actual.sort_unstable(); + expected.sort_unstable(); + if actual != expected || !actual.windows(2).all(|w| w[0] < w[1]) { + return Err( + "split is not an exact disjoint cover of the original blocks".into(), + ); + } + Ok(()) +} + +fn cut(blocks: &[Address], count: usize) -> Result>, String> { + if blocks.len() < 2 { + return Err("indivisible block exceeds the RAM budget".into()); + } + let count = count.clamp(2, blocks.len()); + Ok( + (0..count) + .map(|i| { + blocks[i * blocks.len() / count..(i + 1) * blocks.len() / count] + .to_vec() + }) + .collect(), + ) +} + +impl<'a> Pipeline<'a> { + fn plan_path(&self, work: &Work) -> PathBuf { + let mut key = b"ix-shard-splits-v2".to_vec(); + key.extend_from_slice(self.ctx.allowed); + key.extend_from_slice(&self.max_ram.to_le_bytes()); + key.extend_from_slice(&work.statement.claim_bytes); + self.plans.join(format!("{}.json", Address::hash(&key).hex())) + } + + fn cached(&self, work: &Work) -> Option> { + if !self.skip_proven { + return None; + } + let address = cache_address(self.index?, &work.digest())?; + let result = (|| { + let bytes = read_store(self.ctx.store_dir, &address)?; + if Address::hash(&bytes) != address { + return Err("store object hash mismatch".into()); + } + let wrapper = decode_wrapper(&bytes)?; + import_shard_proof( + self.ctx.ixvm_system, + self.ctx.aggr_system, + self.verify_idx, + self.ctx.aggr_idx, + self.ctx.allowed, + work.statement.clone(), + &wrapper, + Some(address), + ) + })(); + match result { + // A wrap must actually change the backend or a budget fallback loops. + Ok(slot) + if matches!(work.job, Job::Aggregate { right: None, .. }) + && slot.kind != ChildKind::Aggr => + { + None + }, + Ok(slot) => Some(slot), + Err(e) => { + eprintln!( + "[shard-pipeline] ignored index entry {}: {e}", + work.digest().hex() + ); + None + }, + } + } + + fn over_budget( + &self, + work: &Work, + slots: &FxHashMap>, + parts: usize, + ) -> Result, String> { + match &work.job { + Job::Shard { blocks, .. } => Ok(Prepared::Split(cut(blocks, parts)?)), + Job::Aggregate { left, right: Some(right) } + if slots[left].kind == ChildKind::Ixvm + || slots[right].kind == ChildKind::Ixvm => + { + Ok(Prepared::WrapChildren) + }, + Job::Aggregate { .. } => Err(format!( + "flat healing proof cannot fit --max-ram for claim {}; increase the budget", + work.digest().hex() + )), + } + } + + fn prepare( + &self, + work: &Work, + slots: &FxHashMap>, + ) -> Result, String> { + if let Some(slot) = self.cached(work) { + return Ok(Prepared::Reused(slot)); + } + if !work.ready(slots) { + return Err("healing job has an unavailable child proof".into()); + } + if let Job::Shard { blocks, .. } = &work.job + && self.skip_proven + && let Ok(bytes) = fs::read(self.plan_path(work)) + { + let hint = (|| { + let hexes: Vec> = + serde_json::from_slice(&bytes).map_err(|e| e.to_string())?; + let parts = hexes + .into_iter() + .map(|p| { + p.into_iter() + .map(|s| { + Address::from_hex(&s) + .ok_or_else(|| "invalid block address".to_string()) + }) + .collect() + }) + .collect::>, _>>()?; + validate_parts(blocks, &parts)?; + Ok::<_, String>(parts) + })(); + match hint { + Ok(parts) => { + eprintln!("[shard-pipeline] restored split {}", work.digest().hex()); + return Ok(Prepared::Split(parts)); + }, + Err(e) => eprintln!("[shard-pipeline] ignored split journal: {e}"), + } + } + let execution = match &work.job { + Job::Shard { owned, .. } => { + let (claim, input, io) = + build_shard_check_env_witness(self.env, owned)?; + if claim != work.statement.claim { + return Err("prepared shard claim changed".into()); + } + Execution::new( + self.ctx.ixvm_system, + self.verify_idx, + input, + io, + execute_ixvm, + )? + }, + Job::Aggregate { left, right } => { + let left = &slots[left]; + let right = right.as_ref().map(|r| slots[r].as_ref()); + let outer = aggregate_outer_claim( + self.ctx.aggr_idx, + self.ctx.allowed, + &work.statement.claim_bytes, + ); + let spec = SlotSpec { + op: PlanOp::Leaf(0), + statement: work.statement.clone(), + subject_count: work.statement.subjects.count, + structural: false, + kind: ChildKind::Aggr, + shape: Some(shape_code(left.kind, right.map(|r| r.kind))), + outer_claim: outer.clone(), + cache_key: work.digest(), + ram_bytes: 0, + }; + let (io, _) = aggregate_io(self.ctx, &spec, left, right)?; + Execution::new( + self.ctx.aggr_system, + self.ctx.aggr_idx, + outer[2..].to_vec(), + io, + execute_ix_aggr, + )? + }, + }; + let peak = execution.peak(); + eprintln!( + "[shard-pipeline] shard {} claim {}: projected prove {} GiB, budget {} GiB", + work.original, + work.digest().hex(), + format_gib(peak), + format_gib(self.max_ram) + ); + // Gate proving on the executed record's predicted peak. An oversized + // record is dropped before its smaller parts are executed and checked. + if peak > self.max_ram { + let parts = + execution.system.suggested_split_parts(&execution.record, self.max_ram); + drop(execution); + return self.over_budget(work, slots, parts); + } + Ok(Prepared::Execute(execution)) + } + + fn split( + &self, + work: &Work, + parts: &[Vec
], + ) -> Result, String> { + let Job::Shard { blocks, owned } = &work.job else { + return Err("only shards can split".into()); + }; + validate_parts(blocks, parts)?; + let mut owners = FxHashMap::default(); + for (i, part) in parts.iter().enumerate() { + for block in part { + owners.insert(block.clone(), i); + } + } + let mut owned_parts = vec![Vec::new(); parts.len()]; + for address in owned { + let constant = + self.env.try_get_const(address).ok_or("missing owned constant")??; + let block = projection_block(address, &constant); + let owner = owners.get(&block).ok_or("split omitted an owned block")?; + owned_parts[*owner].push(address.clone()); + } + let children = parts + .iter() + .zip(owned_parts) + .map(|(blocks, owned)| { + Ok(Work { + statement: shard_statement(self.env, &owned)?, + job: Job::Shard { blocks: blocks.clone(), owned }, + original: work.original, + publish: false, + }) + }) + .collect::, String>>()?; + let mut jobs = Vec::new(); + let root = append_balanced(&children, &mut jobs)?; + if root.statement.claim_bytes != work.statement.claim_bytes { + return Err("split healing would change the original claim bytes".into()); + } + jobs.last_mut().ok_or("split produced no jobs")?.publish = work.publish; + let hexes: Vec> = + parts.iter().map(|p| p.iter().map(Address::hex).collect()).collect(); + write_atomic( + &self.plan_path(work), + &serde_json::to_vec(&hexes).map_err(|e| e.to_string())?, + )?; + Ok(jobs) + } + + fn persist( + &self, + work: &Work, + outer: Vec, + proof: AiurProof, + ) -> Result, String> { + let kind = match work.job { + Job::Shard { .. } => ChildKind::Ixvm, + Job::Aggregate { .. } => ChildKind::Aggr, + }; + let expected = match kind { + ChildKind::Ixvm => { + inner_claim(self.verify_idx, &work.statement.claim_bytes) + }, + ChildKind::Aggr => aggregate_outer_claim( + self.ctx.aggr_idx, + self.ctx.allowed, + &work.statement.claim_bytes, + ), + }; + if outer != expected { + return Err("prover returned a different public claim".into()); + } + let system = if kind == ChildKind::Ixvm { + self.ctx.ixvm_system + } else { + self.ctx.aggr_system + }; + system + .verify(&outer, &proof) + .map_err(|e| format!("new shard proof failed verification: {e:?}"))?; + let address = persist_wrapper(self.ctx.store_dir, &work.statement, &proof)?; + if let Some(index) = self.index { + write_atomic( + &index.join(work.digest().hex()), + format!("{}\n", address.hex()).as_bytes(), + )?; + } + Ok(Arc::new(Slot { + kind, + statement: work.statement.clone(), + claims_bytes: serialize_claims(&[&outer]), + outer_claim: outer, + proof, + proof_address: Some(address), + })) + } + + fn run(&self, mut queue: VecDeque) -> Result { + // Ingress uses Rayon too. A small separate pool prevents it from taking + // all witness/FFT workers while a STARK occupies the main pool. + let pool = rayon::ThreadPoolBuilder::new() + .num_threads(2) + .build() + .map_err(|e| e.to_string())?; + let mut slots = FxHashMap::default(); + let mut pending = None; + let mut completed = 0; + let mut failures = Vec::new(); + let mut overlaps = 0; + let mut reused = 0; + loop { + let (work, prepared) = if let Some(pending) = pending.take() { + pending + } else if let Some(work) = queue.pop_front() { + let prepared = pool.install(|| self.prepare(&work, &slots)); + (work, prepared) + } else { + break; + }; + let result = (|| -> Result>, String> { + match prepared? { + Prepared::Reused(slot) => { + reused += 1; + Ok(Some(slot)) + }, + Prepared::Split(parts) => { + eprintln!( + "[shard-pipeline] shard {}: splitting into {} parts", + work.original, + parts.len() + ); + for child in self.split(&work, &parts)?.into_iter().rev() { + queue.push_front(child); + } + Ok(None) + }, + Prepared::WrapChildren => { + let Job::Aggregate { left, right: Some(right) } = &work.job else { + unreachable!() + }; + queue.push_front(work.clone()); + for key in [right, left] { + if slots[key].kind == ChildKind::Ixvm { + queue.push_front(Work { + job: Job::Aggregate { left: key.clone(), right: None }, + statement: slots[key].statement.clone(), + original: work.original, + publish: false, + }); + } + } + eprintln!( + "[shard-pipeline] wrapping raw children to reduce healing RAM" + ); + Ok(None) + }, + Prepared::Execute(execution) => { + let prefetch = self.lookahead + && queue.front().is_some_and(|next| next.ready(&slots)); + let (outer, proof) = if prefetch { + let next = queue.pop_front().expect("prefetch has a next job"); + overlaps += 1; + eprintln!( + "[shard-pipeline] overlap: proving shard {}, preparing shard {}", + work.original, next.original + ); + let (proved, prepared) = thread::scope(|scope| { + let producer = + scope.spawn(|| pool.install(|| self.prepare(&next, &slots))); + let proved = execution.prove(); + let prepared = producer + .join() + .map_err(|p| { + format!("preparation panicked: {}", panic_text(&p)) + }) + .and_then(|result| result); + (proved, prepared) + }); + pending = Some((next, prepared)); + proved + } else { + execution.prove() + }; + Ok(Some(self.persist(&work, outer, proof)?)) + }, + } + })(); + match result { + Ok(Some(slot)) => { + if let Job::Aggregate { left, right } = &work.job { + slots.remove(left); + if let Some(right) = right { + slots.remove(right); + } + } + if work.publish { + println!( + "claim {}\n{}", + work.digest().hex(), + slot + .proof_address + .as_ref() + .ok_or("published proof has no address")? + .hex() + ); + completed += 1; + } else { + slots.insert(work.digest(), slot); + } + }, + Ok(None) => {}, + Err(error) => { + let error = format!("shard {}: {error}", work.original); + eprintln!("[shard-pipeline] {error}"); + if !self.keep_going { + return Err(error); + } + failures.push(error); + queue.retain(|w| w.original != work.original); + if pending.as_ref().is_some_and(|(w, _)| w.original == work.original) + { + pending = None; + } + slots.clear(); + }, + } + } + let summary = format!( + "{completed} original shard(s) proven; {reused} cached proof(s); {overlaps} preparation overlap(s)" + ); + if failures.is_empty() { + Ok(summary) + } else { + Err(format!( + "{summary}; {} failure(s): {}", + failures.len(), + failures.join("; ") + )) + } + } +} + +fn append_balanced( + children: &[Work], + jobs: &mut Vec, +) -> Result { + if children.is_empty() { + return Err("cannot join an empty partition".into()); + } + if children.len() == 1 { + jobs.push(children[0].clone()); + return Ok(children[0].clone()); + } + let (left, right) = children.split_at(children.len() / 2); + let left = append_balanced(left, jobs)?; + let right = append_balanced(right, jobs)?; + let work = Work { + job: Job::Aggregate { left: left.digest(), right: Some(right.digest()) }, + statement: flat_join(&left.statement, &right.statement)?, + original: left.original, + publish: false, + }; + jobs.push(work.clone()); + Ok(work) +} + +/// How many NUMA lanes to prove on: `IX_PROVE_LANES=N` caps it (`0`/`off` = +/// one lane, i.e. the single pipeline above); default is one lane per domain +/// visible to this process (`crate::numa`), and never more lanes than +/// original shards. +fn numa_lane_count(shards: usize, max_ram: usize) -> usize { + let numa = crate::numa::detect(); + let domains = numa.domains.len(); + let requested = match std::env::var("IX_PROVE_LANES") { + Ok(v) if matches!(v.trim(), "off" | "0") => 1, + Ok(v) => v.trim().parse::().unwrap_or(domains).max(1), + Err(_) => domains, + }; + if !numa.enabled() { + return 1; + } + let mut lanes = requested.min(domains).min(shards.max(1)); + // Each lane needs its proving budget plus ~15 % for the overlapped + // execution and the shared environment; never run more lanes than the + // cgroup this process lives in can hold. + if let Some(limit) = crate::numa::cgroup_memory_max() { + let per_lane = max_ram / 100 * 115; + let fit = (limit / per_lane.max(1)).max(1); + if fit < lanes { + eprintln!( + "[shard-pipeline] cgroup memory limit {} GiB holds {fit} lane(s) at --max-ram {} GiB (+15 %); reducing from {lanes}", + format_gib(limit), + format_gib(max_ram), + ); + lanes = fit; + } + } + lanes +} + +/// Longest-processing-time assignment of original shards to `lanes` queues. +/// Weighted by the manifest's measured prover peak when every selected shard +/// has one (prove time tracks peak: ~0.5 s/GiB), else by block count (over +/// ~80 shards per lane that proxy balances to within a few percent). Each +/// lane's queue keeps manifest order so lookahead prefetch stays predictable. +fn split_lanes( + queue: VecDeque, + lanes: usize, + peaks: &FxHashMap, +) -> Vec> { + let mut indexed: Vec<(usize, Work)> = queue.into_iter().enumerate().collect(); + let all_measured = !indexed.is_empty() + && indexed.iter().all(|(_, w)| peaks.contains_key(&w.original)); + let weight = |w: &Work| match &w.job { + Job::Shard { blocks, .. } if all_measured => peaks[&w.original].max(1), + Job::Shard { blocks, .. } => blocks.len().max(1), + Job::Aggregate { .. } => 1, + }; + indexed.sort_by_key(|(i, w)| (std::cmp::Reverse(weight(w)), *i)); + let mut load = vec![0usize; lanes]; + let mut assigned: Vec> = vec![Vec::new(); lanes]; + for (i, w) in indexed { + let k = (0..lanes).min_by_key(|&k| (load[k], k)).expect("lanes > 0"); + load[k] += weight(&w); + assigned[k].push((i, w)); + } + assigned + .into_iter() + .map(|mut lane| { + lane.sort_by_key(|(i, _)| *i); + lane.into_iter().map(|(_, w)| w).collect() + }) + .collect() +} + +impl Pipeline<'_> { + /// One pipeline per NUMA domain inside this process: the environment and + /// both systems are shared, each lane's proving threads and first-touch + /// memory are confined to its domain (`crate::numa`), and each lane runs + /// the single-lane pipeline over its own queue (lookahead, split healing, + /// publication unchanged). Stdout lines (`claim …` + address) come from all + /// lanes; each is one atomic `println!`. + fn run_lanes( + &self, + queue: VecDeque, + lanes: usize, + peaks: &FxHashMap, + ) -> Result { + let numa = crate::numa::detect(); + let balanced_by = if !queue.is_empty() + && queue.iter().all(|w| peaks.contains_key(&w.original)) + { + "measured peak" + } else { + "block count" + }; + let queues = split_lanes(queue, lanes, peaks); + let described: Vec = queues + .iter() + .zip(&numa.domains) + .map(|(q, d)| format!("node {} ({} shards)", d.node, q.len())) + .collect(); + eprintln!( + "[shard-pipeline] numa: {lanes} lanes: {}; policy={:?}; balanced by {balanced_by}", + described.join(", "), + numa.policy + ); + let pools = numa + .domains + .iter() + .take(lanes) + .map(|d| crate::numa::pool(numa, d)) + .collect::, _>>()?; + let results: Vec<(u32, Result)> = thread::scope(|scope| { + let handles: Vec<_> = queues + .into_iter() + .zip(pools.iter()) + .zip(numa.domains.iter()) + .map(|((queue, pool), domain)| { + scope.spawn(move || { + crate::numa::pin_current_thread(domain, numa.policy); + let result = + std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + pool.install(|| self.run(queue)) + })) + .unwrap_or_else(|p| { + Err(format!("lane panicked: {}", panic_text(&p))) + }); + (domain.node, result) + }) + }) + .collect(); + handles + .into_iter() + .map(|h| { + h.join().unwrap_or_else(|p| { + (u32::MAX, Err(format!("lane thread panicked: {}", panic_text(&p)))) + }) + }) + .collect() + }); + let mut summaries = Vec::new(); + let mut failures = Vec::new(); + for (node, result) in results { + match result { + Ok(summary) => summaries.push(format!("node {node}: {summary}")), + Err(error) => failures.push(format!("node {node}: {error}")), + } + } + let summary = summaries.join(" | "); + if failures.is_empty() { + Ok(summary) + } else { + Err(format!("{summary} | {}", failures.join(" | "))) + } + } +} + +/// Explicit directories keep tests hermetic and make the native call's +/// persistence boundary identical to the CLI's existing store/index policy. +#[unsafe(no_mangle)] +extern "C" fn rs_aiur_shard_pipeline( + ixvm_system: LeanExternal>, + aggr_system: LeanExternal>, + env_handle: LeanExternal>, + blocks_blob: LeanByteArray>, + owned_blob: LeanByteArray>, + ids: LeanString>, + verify_idx: LeanNat>, + aggr_idx: LeanNat>, + max_ram: LeanNat>, + store: LeanString>, + index: LeanString>, + plans: LeanString>, + lookahead: bool, + skip_proven: bool, + keep_going: bool, +) -> LeanIOResult { + let result = (|| -> Result { + let blocks = decode_addr_lists(blocks_blob.as_bytes())?; + let owned = decode_addr_lists(owned_blob.as_bytes())?; + // Each line is `id` or `idmeasuredPeakBytes` (0 = unmeasured). + let mut peaks: FxHashMap = FxHashMap::default(); + let ids = ids + .as_str() + .lines() + .map(|line| -> Result { + let mut fields = line.split('\t'); + let id = fields + .next() + .unwrap_or("") + .trim() + .parse::() + .map_err(|e| e.to_string())?; + if let Some(peak) = fields.next() + && let Ok(peak) = peak.trim().parse::() + && peak > 0 + { + peaks.insert(id, peak); + } + Ok(id) + }) + .collect::, _>>()?; + if blocks.is_empty() + || blocks.len() != owned.len() + || blocks.len() != ids.len() + { + return Err( + "shard pipeline requires equally sized, nonempty block/owned/id lists" + .into(), + ); + } + let max_ram = lean_unbox_nat_as_usize(max_ram.inner()); + if max_ram == 0 { + return Err("shard pipeline requires positive --max-ram".into()); + } + let ixvm = ixvm_system.get(); + let aggr = aggr_system.get(); + let env = &env_handle.get().env; + let verify_idx = lean_unbox_nat_as_usize(verify_idx.inner()); + let aggr_idx = lean_unbox_nat_as_usize(aggr_idx.inner()); + let ixvm_vk = aiur::vk_codec::aiur_system_to_bytes(ixvm)?; + let aggr_vk = aiur::vk_codec::aiur_system_to_bytes(aggr)?; + let allowed = allowed_blob(&ixvm_vk, verify_idx, &aggr_vk, aggr_idx); + let mut queue = VecDeque::new(); + let mut seen = FxHashSet::default(); + let mut seen_blocks = FxHashSet::default(); + for ((blocks, owned), id) in blocks.into_iter().zip(owned).zip(ids) { + if !seen.insert(id) { + return Err("duplicate original shard id".into()); + } + if blocks.is_empty() { + return Err("empty original shard".into()); + } + let block_set: FxHashSet<_> = blocks.iter().collect(); + for block in &blocks { + if !seen_blocks.insert(block.clone()) { + return Err("selected shard block lists overlap".into()); + } + } + for address in &owned { + let constant = env + .try_get_const(address) + .ok_or("owned constant is missing from the environment")??; + if !block_set.contains(&projection_block(address, &constant)) { + return Err( + "owned constant belongs to a block outside its shard".into(), + ); + } + } + queue.push_back(Work { + statement: shard_statement(env, &owned)?, + job: Job::Shard { blocks, owned }, + original: id, + publish: true, + }); + } + let lanes = numa_lane_count(queue.len(), max_ram); + let pipeline = Pipeline { + ctx: ProveContext { + specs: &[], + prepared: &[], + proofs: None, + owner_by_address: &FxHashMap::default(), + ixvm_system: ixvm, + aggr_system: aggr, + ixvm_vk: &ixvm_vk, + aggr_vk: &aggr_vk, + allowed: &allowed, + aggr_idx, + store_dir: Path::new(store.as_str()), + cache_dir: None, + reprove_slot: None, + write_outputs: true, + }, + env, + verify_idx, + max_ram, + index: (!index.as_str().is_empty()).then(|| Path::new(index.as_str())), + plans: Path::new(plans.as_str()), + skip_proven, + lookahead, + keep_going, + }; + if lanes <= 1 { + return pipeline.run(queue); + } + pipeline.run_lanes(queue, lanes, &peaks) + })(); + LeanIOResult::ok(match result { + Ok(summary) => LeanExcept::ok(LeanString::new(&summary)), + Err(error) => LeanExcept::error_string(&error), + }) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn shard_work(id: usize, blocks: usize) -> Work { + Work { + job: Job::Shard { + blocks: (0..blocks) + .map(|i| { + Address::hash(&((id as u64) * 1000 + i as u64).to_le_bytes()) + }) + .collect(), + owned: Vec::new(), + }, + statement: statement(&[Address::hash(&(id as u64).to_le_bytes())], &[]), + original: id, + publish: true, + } + } + + #[test] + fn split_lanes_balances_by_block_count_and_keeps_order() { + let queue: VecDeque = + [(0, 5), (1, 1), (2, 4), (3, 1), (4, 3), (5, 2)] + .into_iter() + .map(|(id, b)| shard_work(id, b)) + .collect(); + let lanes = split_lanes(queue, 3, &FxHashMap::default()); + assert_eq!(lanes.len(), 3); + let loads: Vec = lanes + .iter() + .map(|l| { + l.iter() + .map(|w| match &w.job { + Job::Shard { blocks, .. } => blocks.len(), + Job::Aggregate { .. } => 1, + }) + .sum() + }) + .collect(); + assert!( + loads.iter().max().unwrap() - loads.iter().min().unwrap() <= 1, + "{loads:?}" + ); + for lane in &lanes { + let ids: Vec = lane.iter().map(|w| w.original).collect(); + let mut sorted = ids.clone(); + sorted.sort_unstable(); + assert_eq!(ids, sorted); + } + assert_eq!(lanes.iter().map(VecDeque::len).sum::(), 6); + } + + #[test] + fn split_lanes_prefers_measured_peaks_when_all_present() { + // Block counts say shard 0 is heaviest; measured peaks say shard 1 is. + let queue: VecDeque = [(0, 9), (1, 1), (2, 1)] + .into_iter() + .map(|(id, b)| shard_work(id, b)) + .collect(); + let peaks: FxHashMap = + [(0, 100), (1, 500), (2, 100)].into_iter().collect(); + let lanes = split_lanes(queue, 2, &peaks); + // LPT by peak: shard 1 (500) alone on one lane, 0 and 2 (100+100) together. + let alone: Vec<&VecDeque> = + lanes.iter().filter(|l| l.len() == 1).collect(); + assert_eq!(alone.len(), 1); + assert_eq!(alone[0][0].original, 1); + } + use std::sync::atomic::{AtomicUsize, Ordering}; + + fn address(n: usize) -> Address { + Address::hash(&n.to_le_bytes()) + } + + fn statement( + subjects: &[Address], + assumptions: &[Address], + ) -> Arc { + let mut subjects = subjects.to_vec(); + let mut assumptions = assumptions.to_vec(); + subjects.sort_unstable(); + assumptions.sort_unstable(); + Statement::new( + SubjectTree::canonical(subjects, ShardSet(vec![])).unwrap(), + CanonicalTree::from_sorted(assumptions).unwrap(), + ) + } + + #[test] + fn nested_flat_healing_keeps_external_and_unproven_sibling_assumptions() { + // Exercise a canonical original larger than Stage 2's structural cutoff. + let all: Vec<_> = (0..5000).map(address).collect(); + let outside = address(5000); + let a = statement( + &all[..2000], + &[all[2000].clone(), all[4000].clone(), outside.clone()], + ); + let b = statement(&all[2000..4000], &[all[0].clone()]); + let c = statement(&all[4000..], std::slice::from_ref(&outside)); + let ab = flat_join(&a, &b).unwrap(); + assert_eq!(ab.assumptions.as_ref().unwrap().leaves.len(), 2); + assert!(ab.assumptions.as_ref().unwrap().leaves.contains(&all[4000])); + let healed = flat_join(&ab, &c).unwrap(); + let original = statement(&all, &[outside]); + assert_eq!(healed.claim_bytes, original.claim_bytes); + assert_eq!( + Address::hash(&healed.claim_bytes), + Address::hash(&original.claim_bytes) + ); + assert!(healed.subjects.canonical_tree().is_some()); + assert!(flat_join(&a, &a).is_err()); + } + + #[test] + fn split_hints_reject_missing_repeated_foreign_and_empty_blocks() { + let blocks: Vec<_> = (0..5).map(address).collect(); + let parts = cut(&blocks, 3).unwrap(); + validate_parts(&blocks, &parts).unwrap(); + let mut bad = parts.clone(); + bad[0].clear(); + assert!(validate_parts(&blocks, &bad).is_err()); + let mut bad = parts.clone(); + bad[0][0] = address(100); + assert!(validate_parts(&blocks, &bad).is_err()); + let mut bad = parts.clone(); + bad[0].push(blocks[0].clone()); + assert!(validate_parts(&blocks, &bad).is_err()); + assert!(validate_parts(&blocks, &parts[1..]).is_err()); + assert!(cut(&blocks[..1], 2).is_err()); + assert!(cut(&[], 2).is_err()); + } + + #[test] + fn execution_moves_between_threads_without_reexecution_or_io_copy() { + static CALLS: AtomicUsize = AtomicUsize::new(0); + fn execute( + top: &Toplevel, + idx: FunIdx, + args: Vec, + io: &mut IOBuffer, + ) -> Result<(QueryRecord, Vec), ExecError> { + CALLS.fetch_add(1, Ordering::SeqCst); + top.execute(idx, args, io) + } + let system = super::super::tests::transport_system(1); + let mut io = + IOBuffer { data: FxHashMap::default(), map: FxHashMap::default() }; + io.data.insert(G::ZERO, Vec::with_capacity(1024)); + let arena = io.data[&G::ZERO].as_ptr(); + let execution = thread::scope(|scope| { + scope + .spawn(|| { + Execution::new(&system, 0, vec![G::ONE], io, execute).unwrap() + }) + .join() + .unwrap() + }); + assert_eq!(execution.io.data[&G::ZERO].as_ptr(), arena); + let (claim, proof) = execution.prove(); + system.verify(&claim, &proof).unwrap(); + assert_eq!(CALLS.load(Ordering::SeqCst), 1); + let io = IOBuffer { data: FxHashMap::default(), map: FxHashMap::default() }; + drop(Execution::new(&system, 0, vec![G::ONE], io, execute).unwrap()); + assert_eq!(CALLS.load(Ordering::SeqCst), 2); + } +} diff --git a/crates/ffi/src/lib.rs b/crates/ffi/src/lib.rs index 93157799b..05f4f1e4c 100644 --- a/crates/ffi/src/lib.rs +++ b/crates/ffi/src/lib.rs @@ -20,6 +20,7 @@ pub mod lean_env; not(all(target_os = "macos", target_arch = "aarch64")) ))] pub mod lean_iroh; +pub mod numa; pub mod texray; pub mod unsigned; diff --git a/crates/ffi/src/numa.rs b/crates/ffi/src/numa.rs new file mode 100644 index 000000000..5e512f2d7 --- /dev/null +++ b/crates/ffi/src/numa.rs @@ -0,0 +1,369 @@ +//! NUMA domains as scheduling lanes. +//! +//! A single STARK prove cannot use more than one SNC/NUMA domain's worth of +//! cores (measured: a full-box prove is 1.08–1.18x one 64-thread lane), while +//! three proves each pinned to their own domain scale 3.00x. Concurrent slots +//! that share one process and one rayon pool interleave on the same cores and +//! first-touch their buffers on whatever node the faulting thread ran, which +//! measured ~0.55x per slot (~1.7x aggregate at three slots). +//! +//! This module gives the scheduler one rayon pool per domain whose workers are +//! pinned with `sched_setaffinity(2)` and `set_mempolicy(2)`. Running a slot's +//! whole body inside `pool.install` puts every `par_iter` the prover reaches on +//! that pool, and pinning the calling thread first makes the slot's serial +//! allocations local too. Memory policy and CPU affinity are per task on +//! Linux, so both are applied on every thread that works for the slot. +//! +//! Configuration (environment, read once): +//! - `IX_NUMA=off` disables pinning (default `auto`: pin when ≥2 domains are +//! visible to this process' cpuset). +//! - `IX_NUMA_POLICY=bind|preferred` (default `bind`): `bind` fails/OOMs rather +//! than spilling to another node when a domain is full; `preferred` spills +//! (measured ~10 % slower when it does). +//! - `IX_NUMA_THREADS=N` workers per domain pool (default: the domain's cpuset). +//! - `IX_NUMA_PACK=0|1` (default 1): allow two slots on one domain when it has +//! the RAM and no domain is idle. +//! +//! Non-Linux targets and single-domain hosts get an empty topology and the +//! scheduler behaves exactly as before. + +use std::sync::{Arc, OnceLock}; + +/// One NUMA domain visible to this process. +#[derive(Clone, Debug)] +pub struct Domain { + pub node: u32, + pub cpus: Vec, + pub mem_bytes: usize, +} + +/// Memory placement policy applied to pinned threads. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum Policy { + Bind, + Preferred, +} + +/// Detected topology and the env-derived knobs. +#[derive(Clone, Debug)] +pub struct Topology { + pub domains: Vec, + pub policy: Policy, + pub threads: Option, + pub pack: bool, +} + +impl Topology { + pub fn enabled(&self) -> bool { + self.domains.len() >= 2 + } +} + +fn env_flag(name: &str, default: bool) -> bool { + match std::env::var(name) { + Ok(v) => !matches!( + v.trim().to_ascii_lowercase().as_str(), + "0" | "off" | "false" | "no" + ), + Err(_) => default, + } +} + +/// Parse a sysfs cpulist such as `0-31,96-127`. +fn parse_cpulist(text: &str) -> Vec { + let mut cpus = Vec::new(); + for part in text.trim().split(',') { + let part = part.trim(); + if part.is_empty() { + continue; + } + if let Some((lo, hi)) = part.split_once('-') { + if let (Ok(lo), Ok(hi)) = (lo.parse::(), hi.parse::()) { + cpus.extend(lo..=hi); + } + } else if let Ok(cpu) = part.parse::() { + cpus.push(cpu); + } + } + cpus +} + +/// Parse `Node N MemTotal: X kB` out of a node meminfo file. +fn parse_mem_total(text: &str) -> Option { + text.lines().find(|l| l.contains("MemTotal")).and_then(|l| { + l.split_whitespace() + .filter_map(|w| w.parse::().ok()) + .nth(1) + .map(|kb| kb * 1024) + }) +} + +#[cfg(target_os = "linux")] +fn allowed_cpus() -> Option> { + // SAFETY: cpu_set_t is plain data; sched_getaffinity fills it for the + // calling thread. Its size is the documented argument. + unsafe { + let mut set: libc::cpu_set_t = std::mem::zeroed(); + let rc = libc::sched_getaffinity(0, size_of::(), &mut set); + if rc != 0 { + return None; + } + let mut cpus = Vec::new(); + for cpu in 0..libc::CPU_SETSIZE as usize { + if libc::CPU_ISSET(cpu, &set) { + cpus.push(cpu); + } + } + Some(cpus) + } +} + +#[cfg(not(target_os = "linux"))] +fn allowed_cpus() -> Option> { + None +} + +fn detect_uncached() -> Topology { + let policy = match std::env::var("IX_NUMA_POLICY") + .map(|v| v.trim().to_ascii_lowercase()) + .as_deref() + { + Ok("preferred") => Policy::Preferred, + _ => Policy::Bind, + }; + let threads = std::env::var("IX_NUMA_THREADS") + .ok() + .and_then(|v| v.trim().parse::().ok()) + .filter(|&n| n > 0); + let pack = env_flag("IX_NUMA_PACK", true); + let mut topology = Topology { domains: Vec::new(), policy, threads, pack }; + if !env_flag("IX_NUMA", true) || !cfg!(target_os = "linux") { + return topology; + } + let Some(allowed) = allowed_cpus() else { + return topology; + }; + let Ok(entries) = std::fs::read_dir("/sys/devices/system/node") else { + return topology; + }; + let mut domains = Vec::new(); + for entry in entries.flatten() { + let name = entry.file_name(); + let name = name.to_string_lossy(); + let Some(node) = + name.strip_prefix("node").and_then(|n| n.parse::().ok()) + else { + continue; + }; + let base = entry.path(); + let Ok(cpulist) = std::fs::read_to_string(base.join("cpulist")) else { + continue; + }; + let cpus: Vec = parse_cpulist(&cpulist) + .into_iter() + .filter(|cpu| allowed.contains(cpu)) + .collect(); + if cpus.is_empty() { + continue; + } + let mem_bytes = std::fs::read_to_string(base.join("meminfo")) + .ok() + .and_then(|t| parse_mem_total(&t)) + .unwrap_or(0); + if mem_bytes == 0 { + continue; + } + domains.push(Domain { node, cpus, mem_bytes }); + } + domains.sort_by_key(|d| d.node); + if domains.len() >= 2 { + topology.domains = domains; + } + topology +} + +/// The tightest cgroup-v2 memory limit that applies to this process +/// (`memory.max` of its cgroup and every ancestor), in bytes; `None` when no +/// limit is set or the hierarchy is unreadable. Lets a stage adapt to the +/// scope it was launched in instead of trusting an inherited environment. +pub fn cgroup_memory_max() -> Option { + let cgroup = std::fs::read_to_string("/proc/self/cgroup").ok()?; + let path = + cgroup.lines().find_map(|l| l.strip_prefix("0::"))?.trim().to_string(); + let mut tightest: Option = None; + let mut dir = std::path::PathBuf::from(format!("/sys/fs/cgroup{path}")); + loop { + if let Ok(text) = std::fs::read_to_string(dir.join("memory.max")) + && let Ok(limit) = text.trim().parse::() + { + tightest = Some(tightest.map_or(limit, |t| t.min(limit))); + } + if dir == std::path::Path::new("/sys/fs/cgroup") { + break; + } + if !dir.pop() { + break; + } + } + tightest +} + +/// This process's resident memory per NUMA node, in bytes, from +/// `/proc/self/numa_maps` (`N=` at that mapping's +/// `kernelpagesize_kB`). Observability only — a few milliseconds for a few +/// hundred mappings. Empty when the file is unreadable (non-Linux, no NUMA). +pub fn resident_by_node() -> Vec<(u32, usize)> { + let Ok(text) = std::fs::read_to_string("/proc/self/numa_maps") else { + return Vec::new(); + }; + let mut totals: std::collections::BTreeMap = + std::collections::BTreeMap::new(); + for line in text.lines() { + let mut page_kb = 4usize; + let mut counts: Vec<(u32, usize)> = Vec::new(); + for field in line.split_whitespace().skip(2) { + if let Some(kb) = field.strip_prefix("kernelpagesize_kB=") { + page_kb = kb.parse().unwrap_or(4); + } else if let Some(rest) = field.strip_prefix('N') + && let Some((node, pages)) = rest.split_once('=') + && let (Ok(node), Ok(pages)) = + (node.parse::(), pages.parse::()) + { + counts.push((node, pages)); + } + } + for (node, pages) in counts { + *totals.entry(node).or_insert(0) += pages * page_kb * 1024; + } + } + totals.into_iter().collect() +} + +/// The process-wide topology (detected once; env read once). +pub fn detect() -> &'static Topology { + static TOPOLOGY: OnceLock = OnceLock::new(); + TOPOLOGY.get_or_init(detect_uncached) +} + +/// Pin the calling thread to `domain`'s CPUs and memory node. +pub fn pin_current_thread(domain: &Domain, policy: Policy) { + pin_current_thread_to(&domain.cpus, domain.node, policy); +} + +/// Pin the calling thread to an explicit CPU list (a whole domain or one of +/// its halves) and to memory node `node`. +#[cfg(target_os = "linux")] +pub fn pin_current_thread_to(cpus: &[usize], node: u32, policy: Policy) { + // SAFETY: cpu_set_t is plain data manipulated through libc's CPU_* helpers; + // set_mempolicy takes a node bitmask with its bit length. Both syscalls + // only affect the calling thread. + unsafe { + let mut set: libc::cpu_set_t = std::mem::zeroed(); + libc::CPU_ZERO(&mut set); + for &cpu in cpus { + if cpu < libc::CPU_SETSIZE as usize { + libc::CPU_SET(cpu, &mut set); + } + } + let _ = libc::sched_setaffinity(0, size_of::(), &set); + let mode = match policy { + Policy::Bind => libc::MPOL_BIND, + Policy::Preferred => libc::MPOL_PREFERRED, + }; + let words = (node as usize / 64) + 1; + let mut mask = vec![0u64; words]; + mask[node as usize / 64] |= 1u64 << (node % 64); + let _ = libc::syscall( + libc::SYS_set_mempolicy, + libc::c_long::from(mode), + mask.as_ptr(), + (words * 64) as libc::c_ulong, + ); + } +} + +#[cfg(not(target_os = "linux"))] +pub fn pin_current_thread_to(_cpus: &[usize], _node: u32, _policy: Policy) {} + +/// Undo `pin_current_thread` on the calling thread: all allowed CPUs, default +/// memory policy. +#[cfg(target_os = "linux")] +pub fn unpin_current_thread(topology: &Topology) { + // SAFETY: as in `pin_current_thread`; MPOL_DEFAULT takes no mask. + unsafe { + let mut set: libc::cpu_set_t = std::mem::zeroed(); + libc::CPU_ZERO(&mut set); + for domain in &topology.domains { + for &cpu in &domain.cpus { + if cpu < libc::CPU_SETSIZE as usize { + libc::CPU_SET(cpu, &mut set); + } + } + } + let _ = libc::sched_setaffinity(0, size_of::(), &set); + let _ = libc::syscall( + libc::SYS_set_mempolicy, + libc::c_long::from(libc::MPOL_DEFAULT), + std::ptr::null::(), + 0_u64, + ); + } +} + +#[cfg(not(target_os = "linux"))] +pub fn unpin_current_thread(_topology: &Topology) {} + +/// A rayon pool whose workers are pinned to one domain. +pub fn pool( + topology: &Topology, + domain: &Domain, +) -> Result, String> { + let threads = topology.threads.unwrap_or(domain.cpus.len()).max(1); + let node = domain.node; + let pin_cpus = domain.cpus.clone(); + let policy = topology.policy; + rayon::ThreadPoolBuilder::new() + .num_threads(threads) + .thread_name(move |i| format!("numa{node}-{i}")) + .start_handler(move |_| pin_current_thread_to(&pin_cpus, node, policy)) + .build() + .map(Arc::new) + .map_err(|e| format!("numa pool for node {node}: {e}")) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn cpulist_ranges_and_singletons() { + assert_eq!(parse_cpulist("0-3,8,10-11\n"), vec![0, 1, 2, 3, 8, 10, 11]); + assert_eq!(parse_cpulist(""), Vec::::new()); + } + + #[test] + fn meminfo_total() { + let text = "Node 2 MemTotal: 516001 kB\nNode 2 MemFree: 1 kB\n"; + assert_eq!(parse_mem_total(text), Some(516001 * 1024)); + assert_eq!(parse_mem_total("nothing"), None); + } + + #[test] + fn resident_by_node_reads_this_process() { + // On Linux the current process has some resident memory on some node. + let resident = resident_by_node(); + if cfg!(target_os = "linux") { + assert!(resident.iter().map(|(_, b)| *b).sum::() > 0); + } + } + + #[test] + fn detect_never_panics_and_is_consistent() { + let t = detect(); + assert!(t.domains.len() != 1); + for d in &t.domains { + assert!(!d.cpus.is_empty()); + assert!(d.mem_bytes > 0); + } + } +} From f8f845b187cdde28cca5708191cb526f28092d63 Mon Sep 17 00:00:00 2001 From: samuelburnham <45365069+samuelburnham@users.noreply.github.com> Date: Thu, 10 Sep 2026 03:02:19 +0000 Subject: [PATCH 2/4] aiur: peak model sums a grouped circuit's member rows After function groups (#619) CircuitType::Function { idx } enumerates circuits, but raw_of still read record.function_queries[idx], charging circuit i with function i's rows. Projections came out 0.83-1.72x (median 1.30x) off on Mathlib shards while measured peaks were unchanged, splitting shard 217 (481 GiB projected, 284 real) at prove time. Sum the members' query counts instead. Shard 0: 378 -> 250 GiB projected vs 248 measured. --- crates/aiur/src/synthesis.rs | 28 ++++++++++++++++++++-------- 1 file changed, 20 insertions(+), 8 deletions(-) diff --git a/crates/aiur/src/synthesis.rs b/crates/aiur/src/synthesis.rs index 28d5435b5..cbd2e4877 100644 --- a/crates/aiur/src/synthesis.rs +++ b/crates/aiur/src/synthesis.rs @@ -112,14 +112,23 @@ pub struct CircuitShape { /// gadgets keep their fixed heights: they are the same size in every /// shard and are most of the peak model's floor, which dividing cannot /// shrink. -fn raw_of( - record: &QueryRecord, +/// Unique queried rows per circuit, before padding. A function circuit's +/// trace concatenates its members' queried rows (see `trace.rs`), so its +/// height is the sum over the circuit's member functions — `idx` is a +/// CIRCUIT index, which only coincides with the function index while every +/// circuit is a singleton. +fn raw_of<'a>( + record: &'a QueryRecord, + circuits: &'a [crate::bytecode::Circuit], parts: usize, -) -> impl Fn(usize, &CircuitType) -> usize + '_ { +) -> impl Fn(usize, &CircuitType) -> usize + 'a { move |_, ct| match ct { - CircuitType::Function { idx } => { - record.function_queries[*idx].len().div_ceil(parts) - }, + CircuitType::Function { idx } => circuits[*idx] + .members + .iter() + .map(|&member| record.function_queries[member].len()) + .sum::() + .div_ceil(parts), CircuitType::Memory { width } => { record.memory_queries.get(width).map_or(0, |m| m.len().div_ceil(parts)) }, @@ -280,7 +289,7 @@ impl AiurSystem { /// which per-fft models blur. pub fn peak_prove_bytes(&self, record: &QueryRecord) -> PeakProveBytes { self.peak_prove_bytes_by( - raw_of(record, 1), + raw_of(record, &self.toplevel.circuits, 1), crate::execute::record_retained_bytes(record), ) } @@ -371,7 +380,10 @@ impl AiurSystem { // count; stop rather than search forever. while parts < (1 << 20) { let peak = self - .peak_prove_bytes_by(raw_of(record, parts), record_bytes / parts) + .peak_prove_bytes_by( + raw_of(record, &self.toplevel.circuits, parts), + record_bytes / parts, + ) .peak; if peak <= max_bytes { break; From eacdfe24b98d273f4095a6c3416fc5167133b548 Mon Sep 17 00:00:00 2001 From: samuelburnham <45365069+samuelburnham@users.noreply.github.com> Date: Thu, 10 Sep 2026 03:02:19 +0000 Subject: [PATCH 3/4] ix verify --ixes: native parallel composed verdict The composed verdict (no --shard, proofs supplied) reconstructed every shard claim in Lean, one shard at a time on one core (~15 s per Mathlib shard, 58 min for the 246-leaf manifest before it was killed). Route it through the Stage 2 import instead: a verify_only mode of the native aggregate entry reconstructs all claims in Rust, binds each proof to its shard by claim digest, verifies every proof in parallel (IxVM or healed ix_aggr) and requires exactly one valid proof per shard. Mathlib, 246 proofs: 1.3 s claims + 0.3 s verification inside a 54 s process (environment load). --record index writes are kept; --shard K is unchanged. --- Ix/Aiur/Protocol.lean | 7 +++-- Ix/Cli/AggregateCmd.lean | 2 +- Ix/Cli/VerifyCmd.lean | 45 +++++++++++++++++--------------- Tests/AggrSemantics.lean | 2 +- crates/ffi/src/aiur/aggregate.rs | 29 +++++++++++++++++--- 5 files changed, 57 insertions(+), 28 deletions(-) diff --git a/Ix/Aiur/Protocol.lean b/Ix/Aiur/Protocol.lean index 9fa21e7ce..e4a1f0234 100644 --- a/Ix/Aiur/Protocol.lean +++ b/Ix/Aiur/Protocol.lean @@ -211,13 +211,16 @@ proving, and persistence. `proofHexes` is one store address per line; `reproveSlotCode` is zero for a full run and `slot + 1` for a targeted replay; the latter loads and verifies only the target's immediate cached children. When `writeOutputs` is false, proofs are hashed but neither the store nor cache -is changed. Returns the root or replayed proof address. -/ +is changed. Returns the root or replayed proof address. With `verifyOnly`, the +run stops after the parallel proof import — every shard claim reconstructed +natively and every supplied proof bound to its shard and verified, exactly one +per shard — returning the empty string (`ix verify --ixes` composed verdict). -/ @[extern "rs_aiur_stage2_aggregate"] opaque aggregateStage2 (ixvmSystem aggrSystem : @& AiurSystem) (envHandle : @& EnvHandle) (manifestPath proofHexes : @& String) (verifyIdx aggrIdx jobs ramBudgetBytes structuralAbove reproveSlotCode : @& Nat) (directJoins planOnly : Bool) (cacheFriBytes : @& ByteArray) - (useCache writeOutputs : Bool) : + (useCache writeOutputs verifyOnly : Bool) : Except String String /-- Reconstruct and audit the manifest-relative aggregate root entirely in diff --git a/Ix/Cli/AggregateCmd.lean b/Ix/Cli/AggregateCmd.lean index e3bdbcdda..88767e15f 100644 --- a/Ix/Cli/AggregateCmd.lean +++ b/Ix/Cli/AggregateCmd.lean @@ -997,7 +997,7 @@ private def runAggregateCmdNativeWith structuralAbove reproveSlotCode (p.hasFlag "direct-joins") (p.hasFlag "plan-only") recursionParameters.cacheFriBytes (!(p.hasFlag "no-cache")) - (!(p.hasFlag "no-write")) + (!(p.hasFlag "no-write")) false match nativeResult with | .error e => IO.eprintln s!"aggregate failed: {e}"; return 1 | .ok _ => return 0 diff --git a/Ix/Cli/VerifyCmd.lean b/Ix/Cli/VerifyCmd.lean index 121cdc1fa..077053c02 100644 --- a/Ix/Cli/VerifyCmd.lean +++ b/Ix/Cli/VerifyCmd.lean @@ -282,31 +282,34 @@ def verifyShardComposition (ixePath manifestPath : String) (shardK? : Option Nat | none => if !(← Ix.Cli.CheckCmd.shardsCover ixonEnv shards) then return 1 if proofs.isEmpty then return 0 - let mut digestToShard : Std.HashMap Address Nat := {} - for k in [0:shards.size] do - let some d ← digestOf k | return 1 - digestToShard := digestToShard.insert d k - let (aiurSystem, compiled) ← match (← buildBackend) with + -- Composed verdict through the native Stage 2 import: Rust reconstructs + -- every shard claim from the manifest (the Lean reconstruction above is + -- one shard at a time on one core — ~15 s per Mathlib shard, an hour for + -- the manifest) and verifies all proofs in parallel, IxVM or healed + -- `ix_aggr` alike, requiring exactly one valid proof per shard. + let (ixvmSystem, compiled) ← match (← buildBackend) with | .error e => IO.eprintln e; return 1 | .ok b => pure b - let mut covered : Std.HashSet Nat := {} - let mut rc : UInt32 := 0 + let verifyIdx := compiled.getFuncIdx `verify_claim |>.get! + let backend ← match ShardProofIndex.buildRecursionBackend ixvmSystem verifyIdx + recursionParameters with + | .error e => IO.eprintln s!"recursion backend: {e}"; return 1 + | .ok backend => pure backend + let envHandle ← match Aiur.EnvHandle.fromIxe ixePath with + | .error e => IO.eprintln s!"EnvHandle.fromIxe {ixePath}: {e}"; return 1 + | .ok handle => pure handle + let proofHexes := String.intercalate "\n" proofs + let verdict ← IO.lazyPure fun _ => + ixvmSystem.aggregateStage2 backend.system envHandle manifestPath proofHexes + verifyIdx backend.aggrIdx 0 0 Ix.Cli.AggregateCmd.defaultStructuralAbove 0 + true false recursionParameters.cacheFriBytes false false true + match verdict with + | .error e => IO.eprintln s!"[verify] FAIL: {e}"; return 1 + | .ok _ => pure () for hex in proofs do let (proofAddr, d) ← claimDigestOfProof hex - match digestToShard.get? d with - | none => IO.eprintln s!"[verify] FAIL: proof {proofAddr} (claim {d}) matches no shard"; rc := 1 - | some k => - if (← verifyOneProof aiurSystem compiled proofAddr recursionParameters) != 0 then rc := 1 - else - covered := covered.insert k - recordProof d proofAddr - let missing := (List.range shards.size).filter (fun k => !covered.contains k) - if !missing.isEmpty then - IO.eprintln s!"[verify] FAIL: shards lacking a valid proof: {missing}" - rc := 1 - if rc == 0 then - IO.println s!"[verify] OK: composed verdict — all {shards.size} shards proven + disjoint cover" - return rc + recordProof d proofAddr + return 0 /-- Verify with an explicit aggregate-recursion configuration. Ordinary IxVM proof verification remains pinned to its independent canonical parameters. -/ diff --git a/Tests/AggrSemantics.lean b/Tests/AggrSemantics.lean index c615a8658..0e09c8e3f 100644 --- a/Tests/AggrSemantics.lean +++ b/Tests/AggrSemantics.lean @@ -431,7 +431,7 @@ def semanticSuite : IO UInt32 := do let planWorks := (ixvmSystem.aggregateStage2 selfSystem handle ixesPath.toString "" verifyIdx fakeAggrIdx 1 (16 * 1024 * 1024 * 1024) 4096 0 false true - childRecursionParameters.cacheFriBytes false true).isOk + childRecursionParameters.cacheFriBytes false true false).isOk let expectedMatches := match Aiur.AiurSystem.aggregateExpected handle ixesPath.toString 4096 with | .error _ => false diff --git a/crates/ffi/src/aiur/aggregate.rs b/crates/ffi/src/aiur/aggregate.rs index dccec396f..8157c2a8f 100644 --- a/crates/ffi/src/aiur/aggregate.rs +++ b/crates/ffi/src/aiur/aggregate.rs @@ -464,6 +464,11 @@ struct RunConfig<'a> { cache_fri_bytes: &'a [u8], use_cache: bool, write_outputs: bool, + /// `ix verify --ixes `: stop after the parallel proof import — + /// every shard claim reconstructed natively, every supplied proof bound + /// to its shard and verified (IxVM or healed `ix_aggr`), exactly one per + /// shard — and report that composed verdict instead of proving. + verify_only: bool, } fn projection_block(addr: &Address, constant: &Constant) -> Address { @@ -2777,10 +2782,10 @@ fn run(config: RunConfig<'_>) -> Result { dir.display() )); } - } else { + } else if !config.verify_only { eprintln!("[aggregate] cache disabled (--no-cache)"); } - if !config.write_outputs { + if !config.write_outputs && !config.verify_only { eprintln!("[aggregate] output writes disabled (--no-write)"); } let needs_input_proofs = replay_plan.as_ref().is_none_or(|plan| { @@ -2829,6 +2834,22 @@ fn run(config: RunConfig<'_>) -> Result { None }; let proofs_at = Instant::now(); + if config.verify_only { + let shards = prepared.shards.len(); + let imported = proofs.as_ref().map_or(0, Vec::len); + if imported != shards { + return Err(format!( + "verified {imported} shard proofs but the manifest has {shards} shards" + )); + } + eprintln!( + "[verify] OK: composed verdict — all {shards} shards proven + disjoint cover ({} proofs verified natively in {:.1}s; claims {:.1}s)", + imported, + (proofs_at - specs_at).as_secs_f64(), + (prepared_at - parsed_at).as_secs_f64(), + ); + return Ok(String::new()); + } print_plan(&specs, &prepared.shards, config.structural_above); eprintln!( "[aggregate] Rust startup: manifest {:.3}s, env/claims {:.3}s, plan/statements {:.3}s, proofs {:.3}s; total {:.3}s", @@ -2960,6 +2981,7 @@ extern "C" fn rs_aiur_stage2_aggregate( cache_fri_bytes: LeanByteArray>, use_cache: bool, write_outputs: bool, + verify_only: bool, ) -> LeanExcept { let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { let reprove_slot = @@ -2981,6 +3003,7 @@ extern "C" fn rs_aiur_stage2_aggregate( cache_fri_bytes: cache_fri_bytes.as_bytes(), use_cache, write_outputs, + verify_only, }) })); match result { @@ -3438,7 +3461,7 @@ mod tests { Toplevel { functions: vec![Function { body: Block { ops: vec![], ctrl: Ctrl::Return(0, vec![]) }, - layout: layout.clone(), + layout, entry: true, constrained: true, }], From fa67a34d3554e499dfc8a106d9149203b4eff3d9 Mon Sep 17 00:00:00 2001 From: samuelburnham <45365069+samuelburnham@users.noreply.github.com> Date: Thu, 10 Sep 2026 03:02:19 +0000 Subject: [PATCH 4/4] Proving docs and the NUMA pinning design note docs/shard-pipeline.md: proving a whole environment on one NUMA box (shard -> refine -> Stage 1 -> Stage 2 under one cgroup slice, THP always), a self-contained end-to-end reproduce section (build, per-boot box setup, slice, every command, what to look for in the logs, resume semantics, subtree experiments), and the measured Mathlib budget: 4 h 36 min total on the rebased binary (Stage 1 2 h 44 min, Stage 2 1 h 52 min), 6 h 20 min before main's function groups, 28 h in production. docs/numa-slot-pinning.md: the pinning design and the calibration behind the slot weights. --- docs/numa-slot-pinning.md | 269 +++++++++++++++++++++++++ docs/shard-pipeline.md | 409 ++++++++++++++++++++++++++++++++++++++ 2 files changed, 678 insertions(+) create mode 100644 docs/numa-slot-pinning.md create mode 100644 docs/shard-pipeline.md diff --git a/docs/numa-slot-pinning.md b/docs/numa-slot-pinning.md new file mode 100644 index 000000000..16b8afe14 --- /dev/null +++ b/docs/numa-slot-pinning.md @@ -0,0 +1,269 @@ +# Design: per-slot NUMA pinning for `ix aggregate` (and lanes for `ix prove`) + +Status: IMPLEMENTED 2026-09-08 (uncommitted; see §10 for measured results). Author: claude (session with S. Burnham). Targets +`crates/ffi/src/aiur/aggregate.rs::run_scheduler` + a new `crates/ffi/src/numa.rs`. +Coordinate with the agent working in `aggregate.rs`/`shard_pipeline.rs`: this design +does not touch `prove_aggregate`, `aggregate_io`, replay, or the shard pipeline. + +## 1. Why (measured) + +| fact | measurement | +|---|---| +| one prove cannot use more than one domain | E1: full box 1.08–1.18x over one 64-thread lane | +| isolated lanes scale perfectly | E1(c): three membind lanes = 3.00x, zero per-lane slowdown | +| in-process slots interfere | InitStd Stage 2: `ix_aggr` join 85 s solo vs 133–159 s with 3 live; direct join ~310 s with 3 live; ⇒ ~0.55x per slot, **~1.7x aggregate at `--jobs 3`** | +| cross-domain memory is not the main cost | interleave vs membind ≈ 10 % (m2/m5/m12) | +| so the loss is shared rayon pool + unplaced memory | three slots' `par_iter` work interleaves on 192 threads; first-touch places each slot's buffers wherever the faulting thread ran | + +Goal: make each admitted slot behave like an E1 lane — its own 64 threads, its own +domain's memory — inside the single `ix aggregate` process (no env reload, no IPC, +no extra binaries). Expected: Stage 2 ≈ 3.0x at three slots instead of ≈ 1.7x. + +## 2. Mechanism + +Linux gives everything needed per thread: `sched_setaffinity(2)` (CPU set) and +`set_mempolicy(2)` (`MPOL_BIND`/`MPOL_PREFERRED` to a node mask). Both are +per-task, inherited by nothing — so they must be applied on every thread that +allocates or computes for the slot. Rayon lets us do exactly that: a +`ThreadPool` built with `.start_handler(|i| pin(domain))` runs the pin on each +worker at spawn, and `pool.install(f)` runs `f` on a worker of that pool, so every +`par_iter` reached from inside `f` — the whole multi-stark prover — executes on +that pool. mimalloc's huge objects are fresh `mmap`s freed on drop (memcfg +series), so their pages are first-touched by pool threads and land on the bound +node; per-thread small heaps are local by construction. THP faults honor the +policy. + +``` +run_scheduler + ├─ topology = numa::detect() // sysfs; honors current cpuset; 1 node ⇒ no-op + ├─ pools[k] = numa::pool(domain k) // lazy, 64 threads, start_handler pins + └─ admit slot → choose domain k (RAM-aware, §4) + scope.spawn(move || { + numa::pin_current_thread(k); // advice/child-proof decode/persist allocate locally + pools[k].install(|| prove_slot(ctx, index, &children)) // prover on domain k + }) +``` + +## 3. Module `crates/ffi/src/numa.rs` (~150 lines, no new deps; `libc` is already a workspace dep) + +```rust +pub struct Domain { pub node: u32, pub cpus: Vec, pub mem_bytes: usize } +pub struct Topology { pub domains: Vec } // empty ⇒ pinning disabled + +pub fn detect() -> Topology + // /sys/devices/system/node/node*/cpulist ("0-31,96-127") + // /sys/devices/system/node/node*/meminfo ("Node 0 MemTotal: N kB") + // intersect each cpulist with sched_getaffinity(0) so `numactl`-launched or + // cgroup-limited processes only see their allowed CPUs; drop empty domains; + // if <2 domains remain ⇒ Topology{domains: []}. + // Env: IX_NUMA=off disables; IX_NUMA=auto (default). + +pub fn pin_current_thread(d: &Domain, policy: Policy) // Policy::Bind | Policy::Preferred + // sched_setaffinity(0, cpu_set of d.cpus); set_mempolicy(MPOL_BIND|MPOL_PREFERRED, &[1< rayon::ThreadPool + // ThreadPoolBuilder::new().num_threads(threads).thread_name(|i| format!("numa{}-{i}", d.node)) + // .start_handler(move |_| pin_current_thread(&d, policy)).build() + +pub fn unpin_current_thread() // affinity = all allowed; MPOL_DEFAULT +``` + +Policy default `Bind` (max throughput; the RAM gate uses estimated weights). `Preferred` +(env `IX_NUMA_POLICY=preferred`) spills to other nodes instead of failing when a +domain is full — the safe mode for untrusted weights. + +Threads per pool: the domain's full cpuset (64 = 32 cores × 2 HT). m7/m8 vs m12 +showed 32 vs 64 within noise under THP; 64 is the default, `IX_NUMA_THREADS` +overrides. + +## 4. Scheduler changes (`run_scheduler`, ~60 lines) + +Today: one global `reserved`/`budget` (`--max-ram`, default 92 % MemTotal), +heaviest-ready-first, `--jobs` cap, "over-budget slot runs alone". + +New state when topology is non-empty: + +``` +domain_budget[k] = min(domain.mem_bytes × 0.90, budget) // ≈ 460 GiB here +domain_reserved[k], domain_active[k] +``` + +Admission for a ready slot of weight `w` (still heaviest first): + +1. candidates = domains with `domain_reserved[k] + w ≤ domain_budget[k]` and global `reserved + w ≤ budget`. +2. prefer an **idle** domain (`domain_active[k] == 0`); among those the most free RAM. +3. else, if `IX_NUMA_PACK != 0` (default 1), the domain with the fewest active slots that fits — two 195 GiB slots may share a domain (PR #598 Init report: 2 proves on one 64-thread box = 1.53x, i.e. still +50 % throughput; a 390 GiB direct join never shares). +4. none fits and nothing active ⇒ the existing over-budget rule: run alone, **unpinned with `MPOL_PREFERRED`-interleave semantics** (`numa::unpin_current_thread`, global pool) so a giant slot can use all memory. +5. `max_jobs` default (when `--jobs 0`) becomes `domains × (pack ? 2 : 1)` instead of "all ready slots". + +Slot completion releases both the global and the domain reservation. Placement is +logged: `[aggregate] slot N: admitted W GiB on node K; node reserved r/b GiB; active a/m`. + +Nothing else changes: `prove_slot`, cache, persistence, replay, `--plan-only`, +the Lean-side plan simulation (`runAggregateDag`) and the FFI signature are +untouched. Configuration is env-only for the first cut (`IX_NUMA`, +`IX_NUMA_POLICY`, `IX_NUMA_PACK`, `IX_NUMA_THREADS`); a `--numa` CLI flag can be +added once the FFI/Lean side is quiet. + +## 5. Memory-placement details that decide whether this reaches 3x + +- **Everything the slot allocates must happen on a pinned thread.** The scope + thread is pinned before `install`, so `aggregate_io` (advice + child-proof + decode), the public-input packing, and `persist_cached` are local. Inside + `install`, all multi-stark `par_iter`s run on pool workers (pinned). +- **Children's proofs live on the child's domain** (`Arc` holds the proof + bytes). A parent reads 8–25 MB remotely once; negligible. +- **mimalloc reuse across slots**: freed small-object pages sit in per-thread + heaps (local); huge buffers are unmapped on free, so a later slot on another + domain never inherits remote pages. Verified by the fault counts: with THP the + process faults ~0.7 M pages per bench, i.e. buffers really are fresh each phase. +- **Global rayon pool remains** for the deterministic prepare/plan work (unpinned); + it is idle during proving. +- **Verify with** `numastat -p ` during a run (each live slot's memory should + sit on one node) and `AnonHugePages` ≈ RSS in `smaps_rollup`. + +## 6. Failure modes + +- `MPOL_BIND` + domain full ⇒ the kernel OOM-kills the process (there is no swap), + taking all live slots with it; the aggregate cache resumes completed slots. The + original 195 GiB structural weight underestimated upper Mathlib joins and + caused this failure on 2026-09-09; see §12 for the corrected model. + `IX_NUMA_POLICY=preferred` trades ~10 % for graceful spill. +- A panic inside `install` propagates to the scope thread and is caught by the + existing `catch_unwind`; the pool stays usable. +- 1-node hosts (CI, laptops): topology empty ⇒ byte-identical behavior to today. +- Running under `numactl --cpunodebind=k`: `sched_getaffinity` intersection leaves + one domain ⇒ pinning disabled, behaves as today. + +## 7. Expected numbers + +Solo per-slot rates (from pass A's root slots and the 0.55x ratio): `ix_aggr` join +≈ 85 s, wrap ≈ 85–90 s, direct join ≈ 170 s (only the 3-way 310 s is measured; solo +is inferred). + +| | today (`--jobs 3`, 1.7x) | pinned (3 idle domains) | pinned + pack (2×195 per domain) | +|---|---:|---:|---:| +| InitStd Stage 2 (pass A shape) | 1244 s measured | ≈ 700 s | ≈ 650 s | +| Mathlib Stage 2 (123 direct + 122 aggr joins) | ≈ 5.5 h | ≈ 2.9 h | ≈ 2.6 h | + +Acceptance test: rerun InitStd pass A (`--direct-joins --jobs 3 --max-ram 1300`, +cache cleared or `--no-cache`) — target ≤ 750 s, per-slot times within 15 % of +solo, `numastat -p` showing one node per live slot. + +## 8. Follow-up: Stage 1 lanes in one process + +`ix prove --ixes --shards …` currently needs three `numactl` processes (three env +loads, three manifests to reconcile). With `numa.rs`, the shard pipeline in +`shard_pipeline.rs` can own one lane per domain (each lane = a pinned pool + +pinned consumer thread + its lookahead executor), giving `ix prove --lanes 3` with +one env load and one `--out-ixes`. Same module, same pin helpers; do it after the +other agent's pipeline lands. + +## 9. Alternatives considered + +- **Per-domain worker processes** (controller spawns `numactl`'d children, IPC over + a pipe): equally 3.0x, better crash isolation, and the natural step to + multi-box; but each worker reloads the 3.3 GB env or must be long-lived with a + slot protocol — more code and a second binary path. Not needed for one box. +- **`numactl` on the whole aggregate process**: caps it at one domain. No. +- **libnuma**: unnecessary; two syscalls and two sysfs files cover it. + +## 10. Implementation and measured results (2026-09-08 20:40–22:00 UTC) + +Implemented as designed: `crates/ffi/src/numa.rs` (topology, pin helpers, pools; +env knobs `IX_NUMA`, `IX_NUMA_POLICY`, `IX_NUMA_THREADS`, `IX_NUMA_PACK`), +`run_scheduler` in `aggregate.rs` (`NumaLane`, `choose_numa_lane`, `numa_lanes`, +per-lane budgets, pin + `install`, per-slot completion log lines, and one rule +beyond the design: a slot that is the only runnable work with nothing live runs +**unpinned**, because a solo pinned slot is ~12 % slower than a solo unpinned +one). `libc` added to the ffi crate. fmt/clippy(-D warnings)/unit tests pass. + +Controlled InitStd Stage 2 (15 leaves, direct joins, `--max-ram 1300`, +`--no-cache`, same binary): + +| run | wall | vs unpinned | +|---|---:|---:| +| `IX_NUMA=off --jobs 3` | 1265.6 s | 1.00x | +| pinned `--jobs 6` (pack) | 1033.6 s | **1.22x** | + +Per slot: direct join 289–332 s unpinned (2 others live) → **209–223 s pinned, +flat under load** (= solo speed, 1.45x). `ix_aggr` join: 136–176 s unpinned with +others → 149 s pinned when packed two per node, 95–103 s pinned alone vs 83–88 s +unpinned alone. Placement verified with `numastat -p` (293/291/281 GB on nodes +0/1/2, one direct join each) and per-thread `Cpus_allowed_list`. + +Why 1.22x and not 1.7x: the interference removed was 1.45x per slot (the solo +direct join is ~210 s, not the 170 s §7 inferred), and on a 15-leaf tree the +serial dependency tail (last ~5 slots) is a third of the wall. On Mathlib +(246 leaves) the parallel phase dominates: **Stage 2 ≈ 3.6 h pinned vs ≈ 5.3 h +unpinned**. + +## 11. Half-pools for packed slots (2026-09-08 22:30–23:07 UTC) + +`IX_NUMA_SPLIT` (default on): each domain's physical cores are split into two +halves (HT siblings kept together, via sysfs `thread_siblings_list`); a lane owns +two 32-thread half-pools besides its full pool. Light slots that pack take a free +half; heavy or lone slots keep the full pool. Wrap-first pinned InitStd Stage 2: + +| | wall | packed slot time | +|---|---:|---:| +| shared 64-thread pool | 977.6 s | 96–347 s (median 160) | +| core-disjoint half-pools | **938.0 s** | 104–159 s (median 136) | + +With isolation + packing, wrap-first (six 195 GiB slots live) beats direct joins +(three 390 GiB slots live) by 8 % (938 vs 1018–1036 s), reversing the unpinned +ordering. The user chose to keep direct joins as the Stage 2 mode (wrap-first recorded for information); half-pools then apply to the upper `ix_aggr` joins. + +## 12. Subject-aware structural reservations (2026-09-09) + +The first Mathlib attempt packed slots 140 and 355 onto node 0, reserving +195 GiB each against a 453.3 GiB node budget (~504 GiB physical). Slot 140 +has 187,668 subjects; slot 355 has 13,023. The subsequent run with packing +disabled completed slot 140 with a query-record prediction of 383.4 GiB and +a sampled node resident peak of 362.0 GiB. Global slice headroom could not +help allocations bound to that node. + +The shape-9 weight is now **195 GiB + 1.25 MiB × subject count**, with a +**390 GiB minimum above 65,536 subjects**, shared by the Rust scheduler and +Lean reference scheduler. The subject term covers growth through the +187,668- and 314,195-subject joins. The minimum covers a separate jump: +slot 337 predicted 380.5 GiB at 91,068 subjects, compared with slot 62's +256.5 GiB at 91,620 subjects. A linear term alone underestimated slot 337 +and slot 265. This model covers the observed peaks; it is not a least-squares +fit. The structural subject-root fold is constant work, but assumption/path +checks and recursive verification of growing children are not. Trace padding +also makes peaks grow in steps. + +Representative shape-9 measurements from +[`stage2-mathlib.log`](../logs/stage2-mathlib.log), matching plan subject counts +with completed proof `peak` entries: + +| slot | subjects | query-record predicted peak (GiB) | new reservation (GiB) | +|---|---:|---:|---:| +| 95 | 11,972 | 202.3 | 209.6 | +| 61 | 55,496 | 211.2 | 262.7 | +| 337 | 91,068 | 380.5 | 390.0 | +| 62 | 91,620 | 256.5 | 390.0 | +| 139 | 96,048 | 248.3 | 390.0 | +| 265 | 126,527 | 380.6 | 390.0 | +| 140 | 187,668 | 383.4 | 424.1 | +| 266 | 314,195 | 454.9 | 578.5 | + +These `peak` entries come from the executed query record's proving-memory +model, not an RSS sampler. Node resident samples during packing include both +slots and must not be attributed to a single proof. Cache hits are excluded. + +The incident pair now reserves 635.0 GiB and cannot share a ~453 GiB node +budget. Small pairs still pack: 23,993 + 24,805 subjects reserve 449.6 GiB +together. Larger-than-node slots use the ordinary unpinned pool; admission +waits for that slot to finish before starting any neighbours, even when its +weight leaves room in the global budget. A concurrent calibration also raised +the mixed-pair reserve from 340 to 390 GiB (Mathlib predicted up to 385 GiB); +the Lean reference and scheduler diagnostics agree with that reserve. Proof +construction, claims and cache keys are unchanged. + +This remains a calibrated scheduling estimate, not a hard RSS bound. The +recorded envelope covers this run through the 314,195-subject join; higher +joins and other workloads require checking against their completed records. +The running benchmark and its `IX_NUMA_PACK=0` setting were not changed. diff --git a/docs/shard-pipeline.md b/docs/shard-pipeline.md new file mode 100644 index 000000000..ba62ff97d --- /dev/null +++ b/docs/shard-pipeline.md @@ -0,0 +1,409 @@ +# Shard proving with lookahead and split healing + +Opt in with `ix prove --ixe ENV.ixe --ixes PLAN.ixes --lookahead +--max-ram PROOF_GIB`. The budget must be an explicit, positive GiB value. +`--heal-splits` enables the same claim-preserving split behavior with serial +execution. Commands without these flags retain the +existing partition-refinement behavior. + +The native scheduler executes at most one next job while the current proof +runs, whenever that job's child proofs are ready. The existing prover consumes +the current record directly. A separate two-thread Rayon pool handles +preparation; at most one next execution or completed record waits for the prover. + +`--max-ram` gates each completed query record's predicted proving peak, including +recursive healing proofs. Execution produces the record, then the existing +peak model determines whether it can be proved within the budget. The same +check applies to serial and lookahead preparation. This budget describes a +single proof's predicted peak; process RSS, free RAM and total machine RAM +do not control execution or overlap. Process placement and memory allocation +across concurrent provers remain the launcher's responsibility. + +An oversized shard is dropped and re-executed in smaller parts, cut only at +block boundaries to preserve mutual recursion groups. Each part is gated again. +All local joins use canonical subject trees, including intermediates larger +than Stage 2's structural cutoff. A join discharges only assumptions in its +actual subject set. Before scheduling a split, the final folded claim bytes +must equal the original leaf's bytes. Recursive proofs are also verified before +publication. If a direct or mixed join does not fit, raw children are wrapped +and a self join is attempted. An indivisible block, wrap, or final flat join +that still cannot fit stops that shard with an error. + +Only original shard proof addresses appear on stdout, each preceded by its +claim digest. Progress, predicted peaks and split decisions go to stderr. +`--shard` and `--shards` select original leaves; `--keep-going` continues to other originals +after a failure. `--out-ixes` copies the original manifest after a successful +run. Child peaks never overwrite original manifest peaks. + +Every completed part and intermediate proof is stored and indexed by its claim +digest. `--skip-proven` verifies index hits under the current IxVM or recursion +backend before using them. Missing objects, wrong hashes, wrong claims and +invalid proofs are misses. Split plans live under `~/.ix/cache/shard-splits/`; +their versioned identity includes both verifier identities, the proving budget +and the original claim. Plans are validated as exact, disjoint block partitions +when restored. Store objects, native index entries and split journals are +published atomically. `--no-index` disables proof-index reads and writes; +proof objects and split plans are still persisted. + +`ix verify`, shard-aware verification, the refinement proof guard, and native +Stage 2 accept healed leaves. Stage 2 authenticates each imported proof's +backend and selects the appropriate raw, mixed or self join while preserving +the existing manifest statement fold. `--plan-only` has no authenticated input +proofs and shows the plan for raw inputs. A replay of an imported healed leaf +has no Stage 2 execution to repeat; replay its consuming join instead, supplying +the shard proofs when an immediate child is an imported leaf. + +No circuit, witness-generation rule, proving parameter, verifying-key encoding +or claim format changes. The only synthesis addition is an immutable bytecode +accessor; execution ownership stays private to the scheduler. Aggregate advice +construction and shard-proof verification reuse existing backends. + +## Validation after benchmarks finish + +The latest changes and new tests have not been built or run: implementation +was completed without Cargo or Lake commands after the request to protect +concurrent benchmark measurements. An earlier intermediate Rust check and CLI +module build passed; those do not validate the final revision. + +Run the Rust aggregate tests and the existing aggregate/manifest suites first. +The additional `IxTests --ignored shard-pipeline` runner uses real IxVM and +recursive proofs with test-only FRI parameters. It exercises nested split +journals, exact original claims with an external frontier, lookahead, verified +resume and recovery from a corrupt child. All artifacts stay in a reported +temporary directory. This runner can use substantial RAM and should run +separately from benchmarks. Source unit tests also cover thread handoff, +canonical joins above the structural cutoff, invalid split partitions, and +authenticated mixed Stage 2 inputs. + +Use `CFLAGS=-std=gnu17` with build commands on the current benchmark host, as +described in the committed handoff. Production shard-scale RSS calibration +and throughput measurements remain necessary before enabling lookahead for a +full proving campaign. + +--- + +# Proving a whole environment: shard → refine → prove → aggregate + +The recommended workflow for a full-environment proof (Mathlib-scale), as +measured on a 3-domain r8i.metal-48xl (2026-09-08/09), using the pipeline +above for Stage 1. The one rule this part exists to state: **run +`ix shard refine` before proving a whole environment.** Everything else +follows from it. + +## 1. Compile and shard + +```sh +ix compile Benchmarks/Compile/CompileMathlib.lean --out mathlib.ixe # ~1 min +ix shard mathlib.ixe --max-ram 400 --out mathlib-400.ixes # ~1 min +``` + +`--max-ram` in the static strategy is a *seed*, not a bound: it sets the +shard count (`233 × (400 / budget)^1.2` for Mathlib) and balances a block-shape +score across it. Real prover peaks are only known after execution, and the +static model's error tail runs to ~1.7x on a few leaves. Seed at 400 GiB for a +430 GiB prove budget: measured 236 → 246 leaves with 8 splits; seeding at 430 +gave 217 → 237 leaves with 19 splits and leaves within 4 GiB of the gate. + +## 2. Refine — always, before proving + +```sh +ix shard refine mathlib.ixe --ixes mathlib-400.ixes --max-ram 400 \ + --out mathlib-refined.ixes --report mathlib-refined.json # ~14 min +``` + +Refine executes every leaf in one parallel batch (admission gated on the +execute-peak estimates, ~11–23 GiB each, so all 236 leaves ran at once at +1.26 TB), records each leaf's measured projected prover peak in the manifest, +and cuts every leaf whose peak exceeds `--max-ram` into a balanced subtree of +parts (part 0 keeps the leaf's id; new parts get new ids at the end). The +output manifest is a refinement of the source: untouched leaves keep their +records and tree positions, and it is the manifest every later step binds to +(`ix prove --ixes`, `ix aggregate --ixes`, `ix verify --aggregate --ixes`). + +Why refine rather than letting the prover split at run time: + +| over-budget leaf handled by | extra cost | +|---|---:| +| refine (before proving) | one more leaf pair and one more Stage 2 join: ~215 s of one domain in the parallel phase ≈ 70 s of wall | +| prove-time split + heal (`ix prove --lookahead`, see `shard-pipeline.md`) | the parts are proved and joined back serially on that lane: measured 589 s vs 160 s for the unsplit leaf (≈ 300 s extra when the direct heal join fits the budget) | + +Refine's fixed cost (~14 min for Mathlib, paid even if nothing splits) breaks +even at 2–3 split leaves per run; Mathlib had 8 (3.4 %). Refine also yields +the measured peaks the lane assignment needs (three lanes balanced to within +0.1 % of peak-sum stayed in lockstep for the whole run) and guarantees a +margin under the prove-time gate (refine at 400, prove at 430: every leaf +≥ 30 GiB under the gate, so no lane splits mid-run and the three lanes' output +manifests stay identical). Prove-time healing remains the right *insurance* +for the rare leaf the model still misses; it should not be the plan. + +## 3. Stage 1: one process, one lane per NUMA domain, THP always + +Set `transparent_hugepage/enabled` to `always` first (1.5–1.7x on the prover: +its per-phase buffers are fresh mappings, and 4K first-touch faults serialize on +the mm lock). Then a single `ix prove`: + +```sh +systemd-run --user --scope --slice=ix-pipeline.slice -- \ + ix prove --ixe mathlib.ixe --ixes mathlib-refined.ixes \ + --lookahead --max-ram 430 --skip-proven --keep-going --texray \ + --out-ixes mathlib-proved.ixes +``` + +With `--lookahead`, the pipeline detects the NUMA topology and runs one lane +per domain inside the process (`IX_PROVE_LANES=N|off` overrides): the +environment and both systems are loaded once, the selected leaves are split +across lanes (longest-first by block count, manifest order within a lane), and +each lane runs the pipeline above on a thread pinned to its domain's cores and +memory node (`crate::numa`), so lanes never share a core or a memory channel. +Every lane keeps lookahead, split healing and publication unchanged, and there +is one `--out-ixes`. A single prove cannot use more than one domain (full-box +prove: 1.08–1.18x one lane); three isolated lanes scale 3.00x. `--lookahead` +hides the execute phase (~24 % of a leaf) under the previous proof: measured +116 s per ~300 GiB leaf per lane vs 132 s without. + +Resume after any failure by rerunning the same command (`--skip-proven`). +Because the lanes share one process, a lane that outgrows its domain takes the +process down (its memory policy is `MPOL_BIND`); the restart loses at most one +in-flight leaf per lane. Expect `--out-ixes` to be byte-identical to the input: +prove-time splits are healed privately and never change the manifest. + +### Memory bounding: one slice for the whole pipeline + +Do not pass per-stage caps around. Size one cgroup slice from the machine and +run every stage as a child scope of it: + +```sh +systemctl --user set-property ix-pipeline.slice MemoryMax=$((MemTotalGiB-24))G MemorySwapMax=0 +systemd-run --user --scope --slice=ix-pipeline.slice -- ix prove … # Stage 1 +systemd-run --user --scope --slice=ix-pipeline.slice -- ix aggregate … # Stage 2 +``` + +The binaries adapt to whatever cgroup they run in: the prove pipeline reads the +tightest `memory.max` over its cgroup ancestors and caps its lane count at +`limit / (max_ram × 1.15)`; the aggregate scheduler clamps its admission budget +to 92 % of the limit. Within a process, per-lane bounding comes from the NUMA +memory policy, so the machine-wide cap is the only number that has to be right. +Section 5 below runs all of the above end to end and is safe to rerun. + +## 4. Collect, verify, aggregate + +Proof addresses come from the shard-proof index +(`~/.ix/cache/shard-proofs/`; claims via `ix shard claims`). +Verify them with `ix verify --ixe --ixes `: the composed verdict +runs through the native Stage 2 import (all shard claims reconstructed in +Rust, every proof bound to its shard and verified in parallel, exactly one +valid proof per shard), about a minute for Mathlib, almost all of it +environment load. It is a convenience gate: Stage 2 re-verifies every leaf +before building a join, and the final `ix verify --aggregate` covers the +environment. Then: + +```sh +systemd-run --user --scope --slice=ix-pipeline.slice -- \ + ix aggregate --ixe mathlib.ixe --ixes mathlib-refined.ixes \ + --direct-joins --jobs 0 --max-ram 1300 +ix verify --aggregate --ixe mathlib.ixe --ixes mathlib-refined.ixes +``` + +The aggregate scheduler pins each slot to a NUMA domain (`crate::numa`; +`IX_NUMA=off` disables). Direct and mixed joins reserve 180 GiB: since the +IxVM function groups (main #619) halved what a leaf proof opens, a direct join +verifying two leaves peaks near 200 GiB resident (it was ≈ 400 GiB at the old +390 GiB reservation). A structural self-join reserves 195 GiB + 1.25 MiB per +subject, with a 390 GiB minimum above 65,536 subjects to cover trace-size +jumps; that subject term is the join's own claim work and did not shrink. Up to two joins can share a +node when their combined weights fit its budget (90% of node total, capped +by the process budget). This retains packing for small joins and separates +large upper joins; the old flat 195 GiB weight caused a Mathlib node OOM. See the +[calibration](numa-slot-pinning.md#12-subject-aware-structural-reservations-2026-09-09) +for measured peaks and the model's limits. + +The solo dependency tail runs unpinned. A join whose weight exceeds every +node's budget also runs unpinned, with no other slots active until it finishes, +even when the process budget has spare room. `--jobs 0` derives the cap from +the topology. A join's identity depends only on its children's claims, so the +aggregate cache resumes any interrupted run and is shared between direct-join +and wrap-first modes. + +The initial independent joins also overlap preparation of the next proof, +including on single-node hosts and with `IX_NUMA=off`. Each worker holds at +most one next execution record. With packing on, a NUMA lane runs up to two +such queues when both fit its budget (2 × (180 + 40) GiB ≤ 453 GiB): Mathlib's +105 direct joins then run six-wide at ≈ 155 s each instead of three-wide at +≈ 94 s, about 20 % more throughput, at the cost of the shared-pool +straggler tail (a few slots at 2–6x). `--max-ram` must hold all six queues +(1350 GiB on this box); `IX_NUMA_PACK=0` returns to one queue per lane. The batch obeys `--jobs`, the process RAM +budget (including its cgroup cap), and any NUMA node budgets. Proving capacity +is assigned first; a worker enables overlap only when its queue has another +job and there is room for the additional 40 GiB record allowance. Without +room for overlap, the ordinary scheduler handles the jobs. Dependent joins +continue to use dynamic admission and packing. `IX_NUMA_LOOKAHEAD=0` disables +overlap on both pinned and unpinned workers. + +## 5. Reproduce end to end + +Everything below was measured on an r8i.metal-48xl (Xeon 6975P-C, 96 cores / +192 threads, 1511 GiB, three sub-NUMA domains of ~504 GiB, Ubuntu 26.04). No +helper scripts are needed; each step is one command, and every step can be +rerun after a failure. + +### Build + +```sh +sudo apt install clang # bindgen needs libclang's headers +export CFLAGS=-std=gnu17 # gcc ≥ 14 defaults to C23; Lean's sysroot lacks __isoc23_strtol +lake build ix +``` + +### Box setup, once per boot + +```sh +echo always | sudo tee /sys/kernel/mm/transparent_hugepage/enabled +echo defer+madvise | sudo tee /sys/kernel/mm/transparent_hugepage/defrag +sudo sysctl -w kernel.numa_balancing=0 +numactl --hardware # expect the domains the lanes will use +``` + +THP `always` is the largest single lever (prover 1.5–1.7x: 87 M → 0.66 M +first-touch faults per leaf). `numa_balancing=0` keeps the kernel from +migrating pages the lanes have deliberately bound. No CPU governor or clock +changes are needed. Then one cgroup slice for the whole pipeline, sized from +the machine (24 GiB left to the OS), swap off: + +```sh +systemctl --user set-property ix-pipeline.slice \ + MemoryMax=$(( $(awk '/MemTotal/{print $2}' /proc/meminfo) / 1048576 - 24 ))G MemorySwapMax=0 +``` + +Every `ix` invocation below runs as `systemd-run --user --scope +--slice=ix-pipeline.slice -- ix …`; the binaries read that limit (lane count, +admission budget). A lane that outgrows its NUMA node is killed by the kernel +under `MPOL_BIND`, not by the slice; the slice bounds the process as a whole. + +### Environment, shards, refine (sections 1–2) + +```sh +ix compile Benchmarks/Compile/CompileMathlib.lean --out mathlib.ixe +ix shard mathlib.ixe --max-ram 400 --out mathlib-400.ixes +ix shard refine mathlib.ixe --ixes mathlib-400.ixes --max-ram 400 \ + --out mathlib-refined.ixes --report mathlib-refined.json +``` + +The manifest to use from here on is `mathlib-refined.ixes` (246 leaves for +Mathlib as of 2026-09-09). A manifest's `measured_peak_bytes` come from the +prover peak model; a manifest measured with a binary that mis-projected (see +the note under the budget table) should be re-run through `refine`. + +### Stage 1 + +```sh +systemd-run --user --scope --slice=ix-pipeline.slice -- \ + ix prove --ixe mathlib.ixe --ixes mathlib-refined.ixes \ + --lookahead --max-ram 430 --skip-proven --keep-going \ + --out-ixes mathlib-proved.ixes +``` + +What to look for in the first minutes: + +- `[shard-pipeline] numa: 3 lanes: node 0 (82 shards), … balanced by measured peak` + — one lane per domain. `IX_PROVE_LANES=1` forces a single lane. +- `[shard-pipeline] shard K claim …: projected prove P GiB, budget 430.0 GiB` + — Mathlib leaves project 250–400 GiB; a projection above 430 means a + prove-time split (`splitting into N parts`), which is healed privately and + costs ~5 min of that lane. +- ~118 s per leaf per lane once the overlap is running; 246 leaves ⇒ + ≈ 2 h 40 min. Node usage stays under ~420 GiB. + +Rerun the same command to resume: `--skip-proven` finds each leaf's proof in +`~/.ix/cache/shard-proofs/` and re-verifies it under the current +verifying key before skipping. Proofs made by a binary with a different +verifying key (any change to the IxVM circuits, e.g. the function groupings) +are not reused; start such a run on fresh `~/.ix/cache/{shard-proofs,aggregate}` +directories (move the old ones aside). + +### Collect and verify the leaf proofs + +```sh +ix shard claims mathlib.ixe --ixes mathlib-refined.ixes > claims.txt # "id digest blocks consts" per leaf +while read -r id digest rest; do + [[ "$id" =~ ^[0-9]+$ ]] && cat ~/.ix/cache/shard-proofs/$digest && echo +done < claims.txt | sed '/^$/d' > proofs.txt +wc -l proofs.txt # must equal the leaf count +ix verify --ixe mathlib.ixe --ixes mathlib-refined.ixes $(cat proofs.txt) +``` + +The composed verdict (`[verify] OK: composed verdict — all 246 shards proven + +disjoint cover`) runs through the native Stage 2 import: about a minute, +almost all of it environment load. Stage 2 repeats the same check on import. + +### Stage 2 + +```sh +systemd-run --user --scope --slice=ix-pipeline.slice -- \ + ix aggregate --ixe mathlib.ixe --ixes mathlib-refined.ixes \ + --direct-joins --jobs 0 --max-ram 1350 $(cat proofs.txt) +``` + +`--max-ram 1350` is what lets two prepare-ahead queues run on every node +(6 × (180 + 40) GiB); the scheduler clamps it to 92 % of the slice anyway. +First lines to check: + +- `[aggregate] plan: 0 wraps, 0 imported healed leaves, 246 direct IxVM leaves + 245 binary joins` +- `[aggregate] numa: 3 lanes … policy=Bind pack=true` +- `[aggregate] pipelines: 105 independent slots over 6 workers (node 0: 18 slots, 220.0 GiB reserved, …); prepare-next overlap on 6/6 workers` +- direct joins complete at ≈ 155 s each packed two per node (≈ 94 s alone); + the direct phase takes ≈ 50 min; upper joins 45–150 s; the last three + levels run nearly serially (~12 min). Root after ≈ 1 h 50 min: + `[aggregate] root proof:
`. + +Rerun the same command to resume: every finished join is in +`~/.ix/cache/aggregate`, keyed by its children's claims, so a restart replays +the cache and continues. `IX_NUMA_PACK=0` returns to one queue per lane +(three-wide, ≈ 300 GiB headroom per node instead of ≈ 95 GiB); +`IX_NUMA_LOOKAHEAD=0` disables prepare-ahead; `IX_NUMA=off` disables pinning. + +### Root + +```sh +ix verify --aggregate --ixe mathlib.ixe --ixes mathlib-refined.ixes +``` + +Expected: `[verify] all 679499 included constants certified well-typed; 0 +undischarged assumptions`, in a few seconds. The root wrapper is +`~/.ix/store////` for address `abc…rest` (≈ 4.9 MB). + +### Reproducing a measurement on a subtree + +Stage 2 experiments do not need the full tree. With a completed Stage 1, +`IX_AGGREGATE_SHARDS=0-31 ix aggregate … $(cat proofs.txt)` aggregates only +that subtree of the manifest (the leaf claims do not depend on the manifest +size; the root then keeps assumptions on the other shards). Use a fresh +`~/.ix/cache/aggregate` per configuration (move the previous one aside) so +nothing is shared between the A and B runs; join times, per-slot peaks and +`lane peaks` are on stderr, and `memory.current` of the slice's cgroup is the +process peak. Shards 0–31 of the Mathlib manifest (32 leaves, 31 joins) run +in 16–36 minutes depending on the configuration and were the basis for the +direct-vs-wrap-first, function-groups and two-queues comparisons in the +budget table. + +## Measured budget (Mathlib, 679,499 constants, 246 leaves) + +| step | wall | +|---|---:| +| compile + static shard | ~2 min | +| refine | ~14 min | +| claims + lane split | ~3 min | +| Stage 1 (3 lanes, lookahead) | 2 h 44 min measured (2 h 38 min the run before) | +| collect + verify | ~1 min (native composed verdict) | +| Stage 2 (pinned direct joins, two queues per lane) | 1 h 52 min measured (3 h 25 min before main's function groups) | +| root verification | seconds | + +Total proving on the rebased binary: **4 h 36 min** (run `mathlib2`, 2026-09-10; +root 4.9 MB, leaf proofs 11 MB each). The previous run on this box took +6 h 20 min; production on one 64-vCPU box (PR #598) took 28 h for the two +proving stages. + +The prover peak model (`AiurSystem::peak_prove_bytes`) gates each leaf against +`--max-ram`; after function groups it must sum a circuit's member rows (fixed +on this branch) — a manifest measured with the broken projection carries +inflated `measured_peak_bytes` and should be re-measured with `ix shard +refine`.