diff --git a/Ix/Aiur/Protocol.lean b/Ix/Aiur/Protocol.lean index 43ce405ec..6b9dc4130 100644 --- a/Ix/Aiur/Protocol.lean +++ b/Ix/Aiur/Protocol.lean @@ -173,17 +173,56 @@ def proveAddrWithEnv (system : @& AiurSystem) (proveAddrWithEnv' system funIdx envHandle addrBytes useBytecode).map fun r => (r.claimBytes, r.proof, .ofArrays r.ioData r.ioMap) +/-- Result of a per-shard prove: the claim's wire bytes, the proof, the + projected prover RAM peak of the record that produced it + (`AiurSystem::peak_prove_bytes`), and the part count the peak model + projects will fit the budget + (`AiurSystem::suggested_split_parts`). + + `proof` is `none` exactly when the peak exceeded the budget — a + RESULT rather than an error, since the caller's answer is to cut + the shard into `suggestedParts` parts and prove those. The count is + computed Rust-side because only there does the executed record + still exist to read per-circuit heights from; it is optimistic + (parts re-execute dependencies shared across the cut), so each part + must still be gated on its own record. `suggestedParts` is 1 + whenever the prove ran. The claim bytes are filled either way (the + claim is known before proving starts). + + The final IO buffer is not returned — it is the shard's whole + ingested byte scope and no caller reads it. -/ +structure ShardProveResult where + claimBytes : ByteArray + proof : Option Proof + peakBytes : Nat + suggestedParts : Nat + @[extern "rs_aiur_system_shard_prove_with_env"] private opaque shardProveWithEnv' : @& AiurSystem → - @& Bytecode.FunIdx → @& EnvHandle → @& ByteArray → - Except String ProveEnvResult - -/-- Per-shard prove against a Rust-owned `EnvHandle`. -/ + @& Bytecode.FunIdx → @& EnvHandle → @& ByteArray → @& Nat → Bool → + Except String ShardProveResult + +/-- Per-shard prove against a Rust-owned `EnvHandle`: ONE execution, + whose record is proven from directly. + + `maxRamBytes` is a per-shard prover-RAM budget checked against that + record's projected peak before the witness phase begins; `0` means + detect (85% of `MemAvailable`, the policy the check batch's RAM gate + uses), and an unreadable `/proc/meminfo` disables the check rather + than guessing. Over budget, the record is dropped and `proof` is + `none` — learning that here costs one execution instead of an OOM + part-way through an FFT. The peak comes back either way, so a prove + run yields the same split/merge signal a check run does. + + `execOnly` stops after execution + measurement (`proof` is `none` + either way; `suggestedParts` is 1 exactly when the peak fits): the + split loop runs on executions alone, never starting a STARK. -/ def shardProveWithEnv (system : @& AiurSystem) - (funIdx : @& Bytecode.FunIdx) (envHandle : @& EnvHandle) (ownedBlob : ByteArray) : - Except String (ByteArray × Proof × IOBuffer) := - (shardProveWithEnv' system funIdx envHandle ownedBlob).map - fun r => (r.claimBytes, r.proof, .ofArrays r.ioData r.ioMap) + (funIdx : @& Bytecode.FunIdx) (envHandle : @& EnvHandle) + (ownedBlob : ByteArray) (maxRamBytes : Nat := 0) + (execOnly : Bool := false) : + Except String ShardProveResult := + shardProveWithEnv' system funIdx envHandle ownedBlob maxRamBytes execOnly @[extern "rs_aiur_system_verify"] opaque verify : @& AiurSystem → @@ -197,36 +236,67 @@ opaque proofToAdviceBytes : @& AiurSystem → end AiurSystem +/-- 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 + that many 32-byte block addresses; every env block must appear in + exactly one shard. `peaksBlob`: one 8-byte LE measured prover peak + per shard in order, recorded on the manifest for schedulers. Own + sizes, foreign blocks, cross-ingress and assumption roots are + recomputed from the env's static profile; prints the manifest + summary to stderr. -/ +@[extern "rs_shard_manifest_from_partition"] +opaque shardManifestFromPartition : @& EnvHandle → + @& ByteArray → @& ByteArray → @& String → IO Unit + namespace Bytecode.Toplevel +/-- One shard's result from `shardCheckBatchWithEnv`. -/ +structure ShardResult where + error : String + peakBytes : Nat + /-- 1 when `peakBytes` fits the batch's `maxRamBytes` (or no budget + was given); otherwise the part count the peak model projects will + fit (`AiurSystem::suggested_split_parts`, measured on the record + in-task). -/ + suggestedParts : Nat + deriving Inhabited + @[extern "rs_aiur_toplevel_shard_check_batch"] private opaque shardCheckBatchWithEnv' : @& Bytecode.Toplevel → @& Bytecode.FunIdx → @& EnvHandle → @& ByteArray → Bool → @& Nat → - @& CommitmentParameters → @& FriParameters → - Except String (Array (String × Nat)) + @& CommitmentParameters → @& FriParameters → @& Nat → + Except String (Array ShardResult) /-- Check EVERY shard of a partition in one call: rayon over the shard list with true work-stealing (no chunk barriers), each shard through the exact single-shard machinery over its own private record and witness io. `shardsBlob` encodes, per shard, a 4-byte LE owned-constant count followed by that many 32-byte addresses. - Returns one `(error, peakBytes)` pair per shard in shard order: - empty error = clean, and `peakBytes` is the analytic prover RAM - peak ([`AiurSystem::peak_prove_bytes`] Rust-side) of the shard's - executed record — the split/merge input (0 on failure). + Returns one `ShardResult` per shard in shard order: empty error = + clean, and `peakBytes` is the analytic prover RAM peak + ([`AiurSystem::peak_prove_bytes`] Rust-side) of the shard's executed + record — the split/merge input (0 on failure). `jobs = 0` uses rayon's default pool width (all cores): peak RSS is bounded by the Rust-side RAM gate (a byte-weighted admission semaphore over estimated per-shard execution RSS vs available system RAM), not by thread count — pass `jobs` only to narrow - CPU use. -/ + CPU use. + + `maxRamBytes > 0` is a per-shard prover-RAM budget: each result's + `suggestedParts` is 1 when its peak fits and the model's projected + part count otherwise, so a caller can cut over-budget shards and + re-batch the parts — the wave loop that audits a partition's split + behavior on executions alone. -/ def shardCheckBatchWithEnv (toplevel : @& Bytecode.Toplevel) (funIdx : @& Bytecode.FunIdx) (envHandle : @& EnvHandle) (shardsBlob : ByteArray) (useBytecode : Bool := false) (jobs : Nat := 0) (commitmentParameters : CommitmentParameters := defaultCommitmentParameters) (friParameters : FriParameters := defaultFriParameters) - : Except String (Array (String × Nat)) := + (maxRamBytes : Nat := 0) + : Except String (Array ShardResult) := shardCheckBatchWithEnv' toplevel funIdx envHandle shardsBlob useBytecode - jobs commitmentParameters friParameters + jobs commitmentParameters friParameters maxRamBytes end Bytecode.Toplevel diff --git a/Ix/Aiur/Semantics/BytecodeFfi.lean b/Ix/Aiur/Semantics/BytecodeFfi.lean index e699b165a..19fe47f19 100644 --- a/Ix/Aiur/Semantics/BytecodeFfi.lean +++ b/Ix/Aiur/Semantics/BytecodeFfi.lean @@ -207,7 +207,8 @@ private opaque checkAddrsWithEnv' : @& Bytecode.Toplevel → through the exact single-claim machinery over task-private data (own witness io, own query record), nothing shared between tasks but the read-only toplevel and env. Returns the FAILURES as - `(addrHex, error)` pairs; empty means all passed. Per-claim + `(batch index as a decimal string, error)` pairs, resolving back to + the caller's label order; empty means all passed. Per-claim outputs/records are not returned — use `checkAddrWithEnv` for a single claim's full result. -/ def checkAddrsWithEnv (toplevel : @& Bytecode.Toplevel) diff --git a/Ix/Aiur/Statistics.lean b/Ix/Aiur/Statistics.lean index 9ff9200c6..aa91de166 100644 --- a/Ix/Aiur/Statistics.lean +++ b/Ix/Aiur/Statistics.lean @@ -53,19 +53,6 @@ structure ExecutionStats where totalCacheHits : Nat deriving Inhabited -/-- Coarse PROVER RAM projection for a record with these circuit - heights: the committed trace bytes over the LDE domain, - `Σ nextPowerOfTwo(height) · committedWidth · 8 · 2^logBlowup`, with empty - circuits contributing nothing (the prover deactivates them). A - surrogate for relative shard sizing — how many shards must split - (over a prover budget) or could merge (far under it) — not an - exact allocator-level peak. -/ -def ExecutionStats.projectedProverBytes (stats : ExecutionStats) - (logBlowup : Nat := defaultCommitmentParameters.logBlowup) : Nat := - stats.circuits.foldl (init := 0) fun acc c => - if c.height == 0 then acc - else acc + Nat.nextPowerOfTwo c.height * c.width * 8 * (2 ^ logBlowup) - /-- Continuous transform-size surrogate: `0` at `0` (an empty circuit is deactivated by the prover), else `x·log2(max(x, 2))` — the clamp keeps a height-1 transform nonzero. -/ diff --git a/Ix/Cli/BenchCmd.lean b/Ix/Cli/BenchCmd.lean index f43c9a932..ecc699f74 100644 --- a/Ix/Cli/BenchCmd.lean +++ b/Ix/Cli/BenchCmd.lean @@ -276,16 +276,18 @@ def backendSpecs : List BackendSpec := [ ("fri-verifier-verify-time", "0.10", "_")] }, -- aiur-sharded-env: whole-env Aiur execution — the sharded feeder pipeline -- end-to-end at env scale, one row per env. Shards the `.ixe` for the - -- runner's RAM (`ix shard --max-ram 100`: naive sizing → ~3.5 GB - -- execution RSS per shard on a 128 GB runner), then one gated - -- full-width rayon batch (`ix check --ixe --ixes`) over the whole - -- manifest — the byte-weighted RamGate, not a thread cap, bounds peak - -- RSS, so the same entry is correct on any runner class. ISLB only - -- for now (~10 min/run); add "FLT" / "Mathlib" to `envs` for the - -- env-scale tiers (~30-45 min each on the 32x runner) when their - -- per-push cost is warranted. `shards` is deterministic per - -- (env bytes, budget) and only drops on a real compression win → - -- upper-only pin. + -- runner's RAM (`ix shard --max-ram 100`: seed sizing), then one gated + -- full-width rayon batch with the split audit (`ix check --ram-budget + -- 100`): a seed the spread pushes over the budget is cut in place and + -- its parts re-measured, so the row describes the LEAF partition the + -- run validated. The byte-weighted RamGate, not a thread cap, bounds + -- peak RSS, so the same entry is correct on any runner class. The + -- measured window (`check-time`) is the wave-0 batch call; split + -- waves are audit extras outside it. ISLB only for now (~10 min/run); + -- add "FLT" / "Mathlib" to `envs` for the env-scale tiers when their + -- per-push cost is warranted. `shards` counts leaves — deterministic + -- per (env static block profile, budget), rising only when the seed under-counts → + -- upper-only pin, re-pin on a justified seed or split change. { name := "aiur-sharded-env", defaultMode := "execute", inputs := .perEnv, envs := some ["ISLB"], testbeds := [("execute", "aiur-sharded-env-check-x64-32x")], @@ -747,23 +749,29 @@ is not a benchmark run" if exit != 0 && exit != exitRejected then IO.eprintln s!"[bench] per-constant checks failed (exit {exit})" | "aiur-sharded-env" => - -- Whole-env sharded Aiur execution: shard the env for the runner's - -- RAM (naive `--max-ram 100` sizing → ~3.5 GB execution RSS per - -- shard), then ONE gated full-width rayon batch over the manifest — - -- the RamGate bounds peak RSS, so no `--jobs` is passed. The check - -- writes the env-keyed row itself (`--json`): check-time, - -- throughput, peak-rss, constants, shards. The shard step is - -- deterministic setup, not part of the measured window. + -- Whole-env sharded Aiur execution: seed the env for the runner's RAM + -- from its static block-shape score (`--max-ram 100`), then ONE gated + -- full-width rayon batch over the manifest — the RamGate bounds peak + -- RSS, so no `--jobs` is passed. The check writes the env-keyed row + -- itself (`--json`): check-time, throughput, peak-rss, constants, + -- shards. The shard step is deterministic setup, not part of the + -- measured window. let ixe ← ensureIxe repo info ((p.flag? "ixe").map (·.as! String)) let ix ← resolveBin repo "ix" let manifest := s!"{env}-exec.ixes" + -- One budget for both stages: the shard step seeds for it, and the + -- check's split audit (`--ram-budget`) cuts any seed the spread + -- pushes over it in place — the row's `shards` counts the LEAF + -- partition the run actually validated. 100 GiB = the 128 GB + -- runner class the testbed pins. + let budgetGib := "100" let exit ← runGuarded watchdog ceilingGb ix - #["shard", ixe, "--max-ram", "100", "--out", manifest] + #["shard", ixe, "--max-ram", budgetGib, "--out", manifest] if exit != 0 then IO.eprintln s!"[bench] ix shard failed (exit {exit})" return 1 let exit ← runGuarded watchdog ceilingGb ix - #["check", "--ixe", ixe, "--ixes", manifest, + #["check", "--ixe", ixe, "--ixes", manifest, "--ram-budget", budgetGib, "--json", out, "--json-name", info.name] if exit != 0 && exit != exitRejected then IO.eprintln s!"[bench] whole-env aiur check failed (exit {exit})" diff --git a/Ix/Cli/CheckCmd.lean b/Ix/Cli/CheckCmd.lean index c9a3caba1..092f0520c 100644 --- a/Ix/Cli/CheckCmd.lean +++ b/Ix/Cli/CheckCmd.lean @@ -34,6 +34,7 @@ public import Ix.Meta public import Ix.Store public import Ix.Cli.NameResolve public import Ix.Benchmark.Results +public import Ix.KernelCheck public import Ix.TracingTexray public section @@ -197,7 +198,8 @@ def runBatchCheck (ixePath : String) (names : List String) (jobs : Nat) The `envHandle?` is `none` only for `.leanW` targets (`--interp` fallback); the addr/shard arms require it. -/ def runCompiled (compiled : Aiur.CompiledToplevel) (printStats : Bool) - (statsOut : Option String) (useBytecode : Bool) + (statsOut : Option String) + (useBytecode : Bool) (envHandle? : Option Aiur.EnvHandle) (target : Target) (label : String) : IO UInt32 := do IO.println s!"Typechecking {label}" @@ -467,20 +469,54 @@ private def blockAddrOf (addr : Address) (c : Ixon.Constant) : Address := | .dPrj prj => prj.block | _ => addr -/-- Owned constants of a shard: every env constant whose check-schedule block - is in `blocks`. - - Constants whose bytes do not parse are skipped and therefore owned by - NOBODY. That is safe only because `shardsCover` fails the run when any - exist, so this is never reached with one present; without that gate a - silent skip here means a constant no shard ever checks. -/ -def ownedConstsForBlocks (ixonEnv : Ixon.Env) (blocks : Array Address) : Array Address := Id.run do - let blockSet : Std.HashSet Address := blocks.foldl (·.insert ·) {} - let mut o : Array Address := #[] - for (addr, lc) in ixonEnv.consts do - let some c := lc.get? | continue - if blockSet.contains (blockAddrOf addr c) then o := o.push addr - return o +/-- Each block address mapped to the index of the list owning it. -/ +def blockIndexOf (lists : Array (Array Address)) : + Std.HashMap Address Nat := + (lists.mapIdx fun k l => (k, l)).foldl (init := {}) fun m (k, l) => + l.foldl (fun m blk => m.insert blk k) m + +/-- Owned constants per entry, in ONE env pass: `result[k]` is every env + constant whose check-schedule block is in `lists[k]`, in + env-iteration order (identical to a per-entry filter, so claim + digests are unchanged). Per-entry filtering rescans all consts each + call — at env scale (241 shards × 688k consts) that is ~30 min of + setup, vs seconds here. + + Constants whose bytes do not parse are owned by NOBODY. That is safe + only because `shardsCover` fails the run when any exist; without + that gate a silent skip here means a constant no shard ever + checks. -/ +def ownedConstsPer (ixonEnv : Ixon.Env) (lists : Array (Array Address)) : + Array (Array Address) := + let blockTo := blockIndexOf lists + ixonEnv.consts.fold (init := Array.replicate lists.size #[]) + fun owned addr lc => + match lc.get? with + | none => owned + | some c => + match blockTo.get? (blockAddrOf addr c) with + | some k => owned.modify k (·.push addr) + | none => owned + +/-- Owned constants of one shard: `ownedConstsPer` over a singleton. -/ +def ownedConstsForBlocks (ixonEnv : Ixon.Env) (blocks : Array Address) : + Array Address := + (ownedConstsPer ixonEnv #[blocks])[0]! + +/-- Partition a shard's already-known `owned` constants among `parts` + (block lists) by check-schedule block: one pass over the owned + consts, none over the env — the split-time companion of + `ownedConstsPer`, whose full env pass runs once per run. -/ +def partitionOwned (ixonEnv : Ixon.Env) (owned : Array Address) + (parts : Array (Array Address)) : Array (Array Address) := + let blockTo := blockIndexOf parts + owned.foldl (init := Array.replicate parts.size #[]) fun acc a => + match (ixonEnv.consts.get? a).bind (·.get?) with + | none => acc + | some c => + match blockTo.get? (blockAddrOf a c) with + | some k => acc.modify k (·.push a) + | none => acc /-- The `CheckEnv` claim digest a shard's proof commits to — reconstructed deterministically from the env + the shard's owned blocks. Matches the @@ -516,8 +552,8 @@ def runShardOwned (ixonEnv : Ixon.Env) (blocks : Array Address) (shardK : Nat) the pre-built envHandle so all shards in an all-shards run share one env parse. -/ def runShardOwnedNative (envHandle : Aiur.EnvHandle) (compiled : Aiur.CompiledToplevel) - (printStats : Bool) (statsOut : Option String) (useBytecode : Bool) - (shapes : Array Aiur.CircuitShape) + (printStats : Bool) (statsOut : Option String) + (useBytecode : Bool) (ixonEnv : Ixon.Env) (blocks : Array Address) (shardK : Nat) : IO UInt32 := do let owned := ownedConstsForBlocks ixonEnv blocks IO.println s!"[shard] shard {shardK}: {blocks.size} owned blocks → \ @@ -534,13 +570,6 @@ def runShardOwnedNative (envHandle : Aiur.EnvHandle) (compiled : Aiur.CompiledTo IO.eprintln s!"{label}: IxVM-native shard check error: {e}" return 1 | .ok (_output, _ioBuffer, queryCounts) => - -- Prover RAM projection from this shard's executed heights: the - -- input to split/merge decisions against a prover budget. - let stats := Aiur.computeStats compiled queryCounts shapes - let bytes := stats.projectedProverBytes - let gib := Float.ofNat bytes / 1073741824.0 - IO.println s!"[shard {shardK}] projected prover RAM: \ - {gib} GiB (padded committed traces × blowup)" if printStats then emitStats compiled queryCounts statsOut pure 0 @@ -557,7 +586,8 @@ def runShardCheckManifest (manifestPath ixePath : String) (shardK : Nat) once for this one call. -/ def runShardCheckManifestNative (manifestPath ixePath : String) (shardK : Nat) (compiled : Aiur.CompiledToplevel) (printStats : Bool) - (statsOut : Option String) (useBytecode : Bool) : IO UInt32 := do + (statsOut : Option String) + (useBytecode : Bool) : IO UInt32 := do match (← loadEnvAndShards manifestPath ixePath) with | .error e => IO.eprintln e; return 1 | .ok (ixonEnv, shards) => match shards[shardK]? with @@ -566,9 +596,8 @@ def runShardCheckManifestNative (manifestPath ixePath : String) (shardK : Nat) let envHandle ← match Aiur.EnvHandle.fromIxe ixePath with | .error e => IO.eprintln s!"EnvHandle.fromIxe {ixePath}: {e}"; return 1 | .ok h => pure h - let shapes := Aiur.circuitShapes compiled.bytecode - Aiur.defaultCommitmentParameters Aiur.defaultFriParameters - runShardOwnedNative envHandle compiled printStats statsOut useBytecode shapes ixonEnv blocks shardK + runShardOwnedNative envHandle compiled printStats statsOut + useBytecode ixonEnv blocks shardK /-- Coverage check over already-loaded env + shards: every constant's check-schedule block is owned by **exactly one** shard. That is the whole @@ -618,6 +647,47 @@ def shardsCover (ixonEnv : Ixon.Env) (shards : Array (Array Address)) : IO Bool IO.println s!"[shards] OK: partition covers all {ixonEnv.consts.size} consts, disjoint" pure ok +/-- Cut `blocks` into contiguous, nonempty, roughly equal-count parts + (`p` clamped to `[2, blocks.size]` — a rejection always needs at + least two parts, and the block list caps how many a cut can + produce). Rows are what the prover-RAM model responds to, and equal + block counts are the measured-best proxy for equal rows + (equal-vspan cuts lost on parts, depth and executions on every env + tried). Shared by the prove split recursion and the check wave + loop. -/ +def cutBlocks (blocks : Array Address) (p : Nat) : Array (Array Address) := + let p := min (max p 2) blocks.size + (Array.range p).map fun i => + blocks.extract (blocks.size * i / p) (blocks.size * (i + 1) / p) + +/-- Measured-peaks blob for the manifest emit: one 8-byte LE value per + shard, in shard order. -/ +def peaksBlob (peaks : Array Nat) : ByteArray := + peaks.foldl (init := ByteArray.empty) fun blob pk => + blob ++ pk.toUInt64.toLEBytes + +/-- Wire encoding shared by the shard batch FFI and the manifest emit: + per entry, a 4-byte LE count followed by the 32-byte addresses. -/ +def addrListsBlob (lists : Array (Array Address)) : ByteArray := + lists.foldl (init := ByteArray.empty) fun blob l => + l.foldl (fun b a => b ++ a.hash) (blob ++ l.size.toUInt32.toLEBytes) + +/-- Emit the partition a run actually validated as a `.ixes` manifest, + measured peaks included — the shared tail of `ix prove --out-ixes` + and `ix check --ram-budget --out-ixes`. Skipped with a note when any + shard failed: a corrected manifest describes a fully-validated + partition. -/ +def emitCorrectedManifest (tag : String) (envHandle : Aiur.EnvHandle) + (out : String) (partition : Array (Array Address × Nat)) + (failures : Nat) : IO Unit := do + if failures == 0 then + Aiur.shardManifestFromPartition envHandle + (addrListsBlob (partition.map (·.1))) + (peaksBlob (partition.map (·.2))) out + IO.println s!"[{tag}] corrected manifest → {out}" + else + IO.eprintln s!"--out-ixes {out} skipped: {failures} failure(s)" + /-- Whole-partition check as ONE Rust rayon batch: work-stealing across shards, no chunk barriers (measured 2.3–2.5x faster than the per-shard Lean-task scheduler it replaced, whose chunk-of-N full @@ -631,7 +701,9 @@ def shardsCover (ixonEnv : Ixon.Env) (shards : Array (Array Address)) : IO Bool soundness contract as `runShardCheckAll`. -/ def runShardBatchNative (manifestPath ixePath : String) (jobs? : Option Nat) (compiled : Aiur.CompiledToplevel) (useBytecode : Bool) - (json? : Option (String × String) := none) : IO UInt32 := do + (json? : Option (String × String)) + (maxRamBytes : Nat) (outIxes : Option String) : + IO UInt32 := do -- The row's peak-rss needs the process-tree RSS sampler running -- (`peakTreeRssBytes` reports 0 otherwise); started before the env -- load so the peak covers the whole run, like `check-rs`. @@ -645,71 +717,111 @@ def runShardBatchNative (manifestPath ixePath : String) (jobs? : Option Nat) | .ok h => pure h let funIdx := compiled.getFuncIdx `verify_claim |>.get! let jobs := jobs?.getD 0 - -- Assign every const to its shard in ONE env pass. Calling - -- `ownedConstsForBlocks` per shard rescans all consts each time — - -- at env scale (241 shards × 688k consts) that is ~30 min of setup. - -- Each shard's array keeps env-iteration order, identical to what - -- the per-shard filter produces, so claim digests are unchanged. - let mut blockToShard : Std.HashMap Address Nat := {} - for (blocks, k) in shards.mapIdx (fun k b => (b, k)) do - for blk in blocks do blockToShard := blockToShard.insert blk k - let mut ownedPerShard : Array (Array Address) := Array.replicate shards.size #[] - for (addr, lc) in ixonEnv.consts do - let some c := lc.get? | continue - match blockToShard.get? (blockAddrOf addr c) with - | some k => ownedPerShard := ownedPerShard.modify k (·.push addr) - | none => pure () - let mut blob := ByteArray.empty - for owned in ownedPerShard do - let n := owned.size.toUInt32 - blob := blob.push n.toUInt8 - blob := blob.push (n >>> 8).toUInt8 - blob := blob.push (n >>> 16).toUInt8 - blob := blob.push (n >>> 24).toUInt8 - for a in owned do blob := blob ++ a.hash - IO.println s!"Typechecking {shards.size} shard(s) in one rayon \ - batch, {jobs} thread(s) (0 = all)" - (← IO.getStdout).flush - let totalConsts := ownedPerShard.foldl (· + ·.size) 0 - let start ← IO.monoMsNow - match compiled.bytecode.shardCheckBatchWithEnv funIdx envHandle blob - useBytecode jobs Aiur.defaultCommitmentParameters - Aiur.defaultFriParameters with - | .error e => IO.eprintln s!"shard batch: {e}"; return 1 - | .ok results => - let elapsedMs := (← IO.monoMsNow) - start - let mut failures : Nat := 0 - for k in [:results.size] do - let (err, peak) := results[k]! - if err.isEmpty then - let gib := Float.ofNat peak / 1073741824.0 - IO.println s!"[shard {k}] ok, projected prover peak {gib} GiB" - else - IO.eprintln s!"[shard {k}] FAILED: {err}" - failures := failures + 1 - -- One env-keyed results row for `ix bench run --backend aiur-sharded-env`: - -- the measured window is the batch FFI call (env load and blob - -- setup are excluded, matching what the benchmark tracks — the - -- execution engine, not the loader). - if let some (path, key) := json? then - let secs := elapsedMs.toFloat / 1000.0 - let tput := if elapsedMs > 0 - then totalConsts.toFloat * 1000.0 / elapsedMs.toFloat else 0.0 - let peakRss ← TracingTexray.peakTreeRssBytes - let status := if failures == 0 then "ok" else "rejected" - Ix.Benchmark.Results.writeRow path key status - [ ("constants", Lean.toJson totalConsts) - , ("shards", Lean.toJson results.size) - , ("check-time", Ix.Benchmark.Results.jsonRound 3 secs) - , ("throughput", Ix.Benchmark.Results.jsonRound 2 tput) - , ("peak-rss", Lean.toJson peakRss) ] - if failures == 0 then - IO.println s!"All {results.size} shard(s) passed" - return 0 - IO.eprintln s!"{failures} of {results.size} shard(s) FAILED" - -- Under `--json` a kernel rejection is the benchmark's `rejected` - -- exit (the row is already written), same contract as `check-rs`. - return if json?.isSome then Ix.Benchmark.Results.exitRejected else 1 + let cutLabeled (label : String) (blocks owned : Array Address) + (p : Nat) : Array (String × Array Address × Array Address) := + let parts := cutBlocks blocks p + (parts.zip (partitionOwned ixonEnv owned parts)).mapIdx + fun i (part, po) => (s!"{label}.{i}", part, po) + -- One loop over execution waves. Wave 0 is the planned partition; + -- with `--ram-budget` every over-budget shard is cut into the peak + -- model's suggested part count and the parts re-batched as the next + -- wave, until everything fits or is a single block. Without a + -- budget the FFI answers 1 part everywhere and the loop is a single + -- wave — the plain batch check. + -- The whole run's ownership assignment: one env pass here, and + -- every later wave's parts inherit their parent's owned consts via + -- `partitionOwned` — no wave ever rescans the env. + let mut wave : Array (String × Array Address × Array Address) := + (shards.zip (ownedConstsPer ixonEnv shards)).mapIdx + fun k (b, o) => (s!"{k}", b, o) + let mut waveNum := 0 + let mut failed : Array String := #[] + let mut final : Array (Array Address × Nat) := #[] + let mut totalConsts := 0 + let mut elapsedMs := 0 + while wave.size > 0 do + if waveNum == 0 then + IO.println s!"Typechecking {wave.size} shard(s) in one rayon \ + batch, {jobs} thread(s) (0 = all)" + else + IO.println s!"[wave {waveNum}] re-executing {wave.size} part(s)" + (← IO.getStdout).flush + let ownedPer := wave.map (·.2.2) + let start ← IO.monoMsNow + match compiled.bytecode.shardCheckBatchWithEnv funIdx envHandle + (addrListsBlob ownedPer) useBytecode jobs + Aiur.defaultCommitmentParameters Aiur.defaultFriParameters + maxRamBytes with + | .error e => IO.eprintln s!"shard batch (wave {waveNum}): {e}"; return 1 + | .ok rs => + if waveNum == 0 then + -- The bench row's measured window: the planned partition's + -- batch call, matching `check-rs`. Split waves are audit + -- extras outside the benchmarked engine window. + elapsedMs := (← IO.monoMsNow) - start + totalConsts := ownedPer.foldl (· + ·.size) 0 + let mut next : Array (String × Array Address × Array Address) := #[] + for (r, label, blocks, owned) in rs.zip wave do + let gib := toGib r.peakBytes + if !r.error.isEmpty then + IO.eprintln s!"[shard {label}] FAILED: {r.error}" + failed := failed.push s!"shard {label}: {r.error}" + final := final.push (blocks, r.peakBytes) + else if r.suggestedParts <= 1 then + IO.println s!"[shard {label}] ok, projected prover peak {gib} GiB" + final := final.push (blocks, r.peakBytes) + else if blocks.size <= 1 then + IO.eprintln s!"[shard {label}] peak {gib} GiB over budget and a \ + single block — cannot split" + failed := failed.push s!"shard {label}: single block over budget" + final := final.push (blocks, r.peakBytes) + else + IO.println s!"[shard {label}] peak {gib} GiB over budget — cut \ + into {r.suggestedParts}" + next := next ++ cutLabeled label blocks owned r.suggestedParts + wave := next + waveNum := waveNum + 1 + if maxRamBytes > 0 then + IO.println s!"[split-audit] {final.size} part(s) from {shards.size} \ + planned shard(s), {waveNum - 1} split wave(s), {failed.size} failure(s)" + -- Consolidation is arithmetic, never an execution: the measured + -- sum only shrinks as shards get coarser, so this count is a safe + -- under-count — and the next run's mandatory executions verify it + -- inline. 95%: margin for the model's +1.7% under-projection. + let sum := final.foldl (· + ·.2) 0 + let target := maxRamBytes * 95 / 100 + let n1 := max 1 ((sum + target - 1) / target) + if failed.isEmpty && n1 < final.size then + IO.println s!"[split-audit] measured total {toGib sum} GiB — \ + {n1} shard(s) would fit the budget: re-shard with --shards {n1}" + -- The recalibrated partition — what this run actually validated — + -- written as the manifest the next run should start from. + if let some out := outIxes then + emitCorrectedManifest "split-audit" envHandle out final failed.size + -- One env-keyed results row for `ix bench run --backend aiur-sharded-env`: + -- the measured window is the wave-0 batch FFI call (env load and + -- blob setup are excluded, matching what the benchmark tracks — the + -- execution engine, not the loader). + if let some (path, key) := json? then + let secs := elapsedMs.toFloat / 1000.0 + let tput := if elapsedMs > 0 + then totalConsts.toFloat * 1000.0 / elapsedMs.toFloat else 0.0 + let peakRss ← TracingTexray.peakTreeRssBytes + let status := if failed.isEmpty then "ok" else "rejected" + Ix.Benchmark.Results.writeRow path key status + [ ("constants", Lean.toJson totalConsts) + , ("shards", Lean.toJson final.size) + , ("check-time", Ix.Benchmark.Results.jsonRound 3 secs) + , ("throughput", Ix.Benchmark.Results.jsonRound 2 tput) + , ("peak-rss", Lean.toJson peakRss) ] + if failed.isEmpty then + IO.println s!"All {final.size} shard(s) passed" + return 0 + IO.eprintln s!"{failed.size} of {final.size} shard(s) FAILED:" + for f in failed do IO.eprintln s!" {f}" + -- Under `--json` a kernel rejection is the benchmark's `rejected` + -- exit (the row is already written), same contract as `check-rs`. + return if json?.isSome then Ix.Benchmark.Results.exitRejected else 1 /-- Run the shard operation over EVERY shard — the whole-partition behavior of `--ixes` with no `--shard` (used by `prove`). Loads the env once. Returns 1 @@ -825,7 +937,9 @@ def runCheckCmd (p : Cli.Parsed) : IO UInt32 := do let json? := (p.flag? "json").map fun f => (f.as! String, ((p.flag? "json-name").map (·.as! String)).getD "env") return (← runShardBatchNative manifest ixe - ((p.flag? "jobs").map (·.as! Nat)) compiled useBytecode json?) + ((p.flag? "jobs").map (·.as! Nat)) compiled useBytecode json? + (((p.flag? "ram-budget").map (·.as! Nat)).getD 0 * gibBytes) + ((p.flag? "out-ixes").map (·.as! String))) | _, _, _ => -- `--jobs N` (N ≠ 1) with an `--ixe` env and no `--claim` takes the -- parallel batch path: one FFI call, rayon over the target list, @@ -860,6 +974,8 @@ def checkCmd : Cli.Cmd := `[Cli| "jobs" : Nat; "Parallelism. With --ixes (no --shard): max shards checked concurrently (default: all at once). With --ixe alone and N ≠ 1: check the targeted constants on N Rust threads (0 = all cores), each claim over its own private record — peak RAM is bounded by N in-flight claim closures." "json" : String; "With --ixes (no --shard): append one env-keyed results row (see Ix.Benchmark.Results) for the batch to this file — check-time, throughput, peak-rss, constants, shards. Used by `ix bench run --backend aiur-sharded-env`." "json-name" : String; "Row key for the --json row (default: `env`)." + "ram-budget" : Nat; "The destination prove box's per-shard RAM budget, GiB (with --ixes, no --shard): after the batch, cut every shard whose projected prover peak exceeds the budget into the peak model's suggested part count and re-batch the parts, wave by wave, until everything fits — the exec-only split audit. Under-filled partitions get a printed suggestion (the shard count the measured total says would fit), never an extra execution: re-shard with --shards N and let the next run's mandatory executions verify it inline. Same unit and model as `ix prove --max-ram`, but no auto-detection: the budget describes the prove box the partition is destined for, not the machine running the check. Omit for a plain check." + "out-ixes" : String; "With --ram-budget: write the recalibrated partition — the block lists the wave loop actually validated, splits included — as a `.ixes` manifest to this path. Skipped if any shard failed. The manifest the next run of this env should start from." ARGS: ...names : String; "Fully-qualified Lean.Name(s) to check. With none, iterate every named constant in the env (sorted)." diff --git a/Ix/Cli/ProveCmd.lean b/Ix/Cli/ProveCmd.lean index f5802c09e..ff4e0fefa 100644 --- a/Ix/Cli/ProveCmd.lean +++ b/Ix/Cli/ProveCmd.lean @@ -86,7 +86,13 @@ def proveOne (aiurSystem : Aiur.AiurSystem) | .error e => IO.eprintln s!"{label}: shardProveWithEnv error: {e}" return 1 - | .ok (_claimBytes, proof, _outIO) => pure proof + | .ok r => match r.proof with + | some proof => pure proof + | none => + IO.eprintln s!"{label}: projected prover peak {r.peakBytes} bytes \ + exceeds the budget; this target is not a manifest shard, so it \ + cannot be split — raise --max-ram" + return 1 | .leanW witness, _ => match aiurSystem.proveIxVM funIdx witness.input witness.inputIOBuffer with | .error e => @@ -101,47 +107,136 @@ def proveOne (aiurSystem : Aiur.AiurSystem) IO.println (toString proofAddr) return 0 -/-- Per-shard prove via the end-to-end Rust path - (`shardProveIxVM`): witness build, `execute_ixvm`, and STARK - prove run in one FFI trip with the parallel Rust witness - builder. -/ -def runShardProveNative (manifestPath : String) (envHandle : Aiur.EnvHandle) - (ixonEnv : Ixon.Env) (shards : Array (Array Address)) (shardK : Nat) - (aiurSystem : Aiur.AiurSystem) (compiled : Aiur.CompiledToplevel) - (_printStats : Bool) : IO UInt32 := do - match shards[shardK]? with - | none => IO.eprintln s!"shard {shardK} out of range (0..{shards.size})"; return 1 - | some blocks => do - let owned := Ix.Cli.CheckCmd.ownedConstsForBlocks ixonEnv blocks - let mut blob := ByteArray.empty - for a in owned do - blob := blob ++ a.hash - let label := s!"shard {shardK}" - IO.println s!"Proving {label}" - (← IO.getStdout).flush - let funIdx := compiled.getFuncIdx `verify_claim |>.get! - match aiurSystem.shardProveWithEnv funIdx envHandle blob with - | .error e => - IO.eprintln s!"{label}: shardProveWithEnv error: {e}" - return 1 - | .ok (claimBytes, proof, _outIO) => +/-- Prove the constants owned by `blocks` as ONE shard, cutting it into + the peak model's suggested part count whenever the executed record's + projected prover peak exceeds the budget, and recursing on any part + that still misses. + + Cuts are contiguous equal-block-count runs: rows are what the RAM + model responds to, and equal-cumulative-vspan cuts measured worse + on every axis (parts, depth, executions) across a synthetic + fixture, init, and lean — vspan tracks time, not rows. + + A shard is only a plan, never part of a statement: its claim is + `checkEnv(ownedRoot, asmRoot)`, a pure function of `(env, owned)` + with the frontier recomputed from the owned set. So cutting `blocks` + yields claims exactly as valid as the parent's — each part's grown + frontier is discharged by its siblings — and the cut costs only the + parent's execution, the cheap half of the run. Splitting on BLOCKS + rather than constants is what keeps mutual-recursion groups intact, + since every constant maps to exactly one block. + + The suggested count is optimistic — parts re-execute dependencies + shared across the cut, and equal block counts are not equal row + counts — so each part is gated on its own executed record and + re-cut if it misses. That bias is deliberate: an under-cut costs + one cheap re-execution, an over-cut pays the per-proof floor on + every extra part for the life of the partition. + + Returns the parts actually proven, in proof order, each with its + measured projected prover peak: the parent's own when it fit, + otherwise its descendants'. That is the partition the run really + produced — what the corrected manifest describes, peaks included. -/ +partial def proveBlocksWithinBudget (envHandle : Aiur.EnvHandle) + (ixonEnv : Ixon.Env) (aiurSystem : Aiur.AiurSystem) + (funIdx : Aiur.Bytecode.FunIdx) (maxRamBytes : Nat) + (execOnly : Bool) (label : String) + (blocks owned : Array Address) : + IO (Except String (Array (Array Address × Nat))) := do + let mut blob := ByteArray.empty + for a in owned do + blob := blob ++ a.hash + IO.println s!"Proving {label} ({blocks.size} blocks, {owned.size} consts)" + (← IO.getStdout).flush + match aiurSystem.shardProveWithEnv funIdx envHandle blob maxRamBytes + execOnly with + | .error e => return .error s!"{label}: shardProveWithEnv error: {e}" + | .ok { claimBytes, proof, peakBytes, suggestedParts } => + let gib := toGib peakBytes + match proof with + | none => + -- The peak fit but nothing proved: reachable only in exec-only + -- mode (a budgeted prove either proves or suggests ≥ 2 parts). + if suggestedParts <= 1 then + IO.println s!"[{label}] prover peak {gib} GiB (exec-only)" + return .ok #[(blocks, peakBytes)] + -- A single block is the atom the kernel checks together; there is + -- no smaller shard to fall back to. + if blocks.size <= 1 then + return .error s!"{label}: projected prover peak {gib} GiB exceeds \ + the budget and the shard is a single block — raise --max-ram" + let cut := Ix.Cli.CheckCmd.cutBlocks blocks suggestedParts + IO.println s!"[{label}] peak {gib} GiB over budget — cutting \ + {blocks.size} blocks into {cut.size} parts" + let cutOwned := Ix.Cli.CheckCmd.partitionOwned ixonEnv owned cut + let mut proven : Array (Array Address × Nat) := #[] + for (i, part, po) in (cut.zip cutOwned).mapIdx + (fun i (part, po) => (i, part, po)) do + match ← proveBlocksWithinBudget envHandle ixonEnv aiurSystem funIdx + maxRamBytes execOnly s!"{label}.{i}" part po with + | .error e => return .error e + | .ok parts => proven := proven ++ parts + return .ok proven + | some proof => + IO.println s!"[{label}] prover peak {gib} GiB" -- Rust returns the canonical CheckEnv claim's wire bytes; deserialize -- back to `Ix.Claim` to persist alongside the proof. Avoids -- recomputing the closure walk + canonical AssumptionTree Lean-side. match Ixon.runGet Ix.Claim.get claimBytes with - | .error e => - IO.eprintln s!"{label}: Claim wire-decode failed: {e}" - return 1 + | .error e => return .error s!"{label}: Claim wire-decode failed: {e}" | .ok claim => do let _ ← StoreIO.toIO (Store.write (Ix.Claim.ser claim)) let wrapper : Ixon.Proof := { claim, proof := proof.toBytes } let proofAddr ← StoreIO.toIO (Store.write (Ixon.Proof.ser wrapper)) IO.println (toString proofAddr) - let _ := manifestPath -- kept for parity with previous signature - return 0 + return .ok #[(blocks, peakBytes)] + +/-- Prove shard `shardK` of the partition, splitting it as needed. + Returns the block sets actually proven for that shard. -/ +def runShardProveNative (envHandle : Aiur.EnvHandle) (ixonEnv : Ixon.Env) + (shards ownedPerShard : Array (Array Address)) (shardK : Nat) + (aiurSystem : Aiur.AiurSystem) (compiled : Aiur.CompiledToplevel) + (maxRamBytes : Nat) (execOnly : Bool) : + IO (Except String (Array (Array Address × Nat))) := do + match shards[shardK]?, ownedPerShard[shardK]? with + | some blocks, some owned => + let funIdx := compiled.getFuncIdx `verify_claim |>.get! + proveBlocksWithinBudget envHandle ixonEnv aiurSystem funIdx maxRamBytes + execOnly s!"shard {shardK}" blocks owned + | _, _ => return .error s!"shard {shardK} out of range (0..{shards.size})" + +/-- Report the partition a prove run actually produced against the one + its manifest planned. They differ exactly when a shard was split. -/ +def reportPartition (proven : Array (Array Address × Nat)) (planned : Nat) : IO Unit := do + if proven.size == planned then + IO.println s!"[prove] {proven.size} shard(s) proven, partition unchanged" + else + IO.println s!"[prove] {proven.size} shard(s) proven from {planned} planned \ + ({proven.size - planned} from splits) — re-shard with this partition to \ + skip the splits next run" 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 + -- breakdown (execute vs witness vs STARK) of every prove in the run. + if p.hasFlag "texray" then TracingTexray.init {} 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`). + let maxRamBytes := + ((p.flag? "max-ram").map (·.as! Nat)).getD 0 * gibBytes + let execOnly := p.hasFlag "exec-only" + let outIxes := (p.flag? "out-ixes").map (·.as! String) + -- Both `--ixes` branches end the same way: the partition the run + -- actually proved, written as the manifest to start the next run from. + let emitCorrected (envHandle : Aiur.EnvHandle) + (partition : Array (Array Address × Nat)) (failures : Nat) : + IO Unit := do + if let some out := outIxes then + Ix.Cli.CheckCmd.emitCorrectedManifest "prove" envHandle out partition + failures let ixePath : Option String := (p.flag? "ixe").map (·.as! String) let claimHex : Option String := (p.flag? "claim").map (·.as! String) let names := (p.variableArgsAs! String).toList @@ -163,7 +258,21 @@ def runProveCmd (p : Cli.Parsed) : IO UInt32 := do let envHandle ← match Aiur.EnvHandle.fromIxe ixe with | .error e => IO.eprintln s!"EnvHandle.fromIxe {ixe}: {e}"; return 1 | .ok h => pure h - runShardProveNative manifest envHandle ixonEnv shards k aiurSystem compiled false + match ← runShardProveNative envHandle ixonEnv shards + (Ix.Cli.CheckCmd.ownedConstsPer ixonEnv shards) k aiurSystem + compiled maxRamBytes execOnly with + | .error e => IO.eprintln e; return 1 + | .ok parts => + -- A single-shard run plans exactly one shard, not the whole + -- partition. + reportPartition parts 1 + -- The corrected whole partition: the plan with shard k replaced + -- by the parts this run actually proved (other shards stay + -- unmeasured — this run never executed them). + emitCorrected envHandle + ((shards.extract 0 k).map (·, 0) ++ parts + ++ (shards.extract (k + 1) shards.size).map (·, 0)) 0 + return 0 | some ixe, some manifest, none => -- IxVM-native all-shards prove. Same envHandle reused across -- every shard. @@ -173,12 +282,37 @@ def runProveCmd (p : Cli.Parsed) : IO UInt32 := do let envHandle ← match Aiur.EnvHandle.fromIxe ixe with | .error e => IO.eprintln s!"EnvHandle.fromIxe {ixe}: {e}"; return 1 | .ok h => pure h - let mut rc : UInt32 := 0 + -- Accumulate the partition the run actually produced. Splits stay + -- in this process, so the manifest describing them is written + -- once at the end rather than rewritten per split. Ownership is + -- assigned in ONE env pass here; splits inherit it. + let ownedPer := Ix.Cli.CheckCmd.ownedConstsPer ixonEnv shards + let mut proven : Array (Array Address × Nat) := #[] + let mut failed : Array String := #[] for k in [0 : shards.size] do - if (← runShardProveNative manifest envHandle ixonEnv shards k - aiurSystem compiled false) != 0 then - rc := 1 - pure rc + match ← runShardProveNative envHandle ixonEnv shards ownedPer k + aiurSystem compiled maxRamBytes execOnly with + | .error e => IO.eprintln e; failed := failed.push e + | .ok parts => proven := proven ++ parts + if failed.isEmpty then + reportPartition proven shards.size + -- Consolidation is arithmetic from the peaks this run already + -- measured, never an extra execution: the measured sum only + -- shrinks as shards get coarser, so the count is a safe + -- under-count, verified inline by the next run's own gate. + if maxRamBytes > 0 then + let sum := proven.foldl (· + ·.2) 0 + let n1 := max 1 ((sum + maxRamBytes * 95 / 100 - 1) + / (maxRamBytes * 95 / 100)) + if n1 < proven.size then + IO.println s!"[prove] measured total {toGib sum} GiB — {n1} \ + shard(s) would fit the budget: re-shard with --shards {n1}" + else + IO.eprintln s!"[prove] {failed.size} of {shards.size} shard(s) FAILED:" + for e in failed do + IO.eprintln s!" {e}" + emitCorrected envHandle proven failed.size + return if failed.isEmpty then 0 else 1 | _, _, _ => Ix.Cli.CheckCmd.forEachClaim ixePath claimHex names keepGoing "prove" false runOne @@ -191,10 +325,14 @@ def proveCmd : Cli.Cmd := `[Cli| FLAGS: "keep-going"; "Continue past failures and report them at the end instead of halting on the first." + "texray"; "Stream per-phase `[texray]` timing/RSS lines (execute, witness, STARK stages) to stderr as each span closes." "ixe" : String; "Path to a serialized `.ixe` env. When set, the binary reads the env from disk instead of using the compiled-in Lean env." "claim" : String; "32-byte hex address of a persisted `Ix.Claim` in `~/.ix/store/`. When set, proves the persisted claim against the `--ixe` env (single proof, skips per-const iteration)." "ixes" : String; "Path to a `.ixes` shard manifest (with --ixe). With --shard K: prove shard K. Without --shard: prove every shard in the partition." "shard" : Nat; "0-based shard index K (with --ixes and --ixe): prove that one shard's CheckEnv claim." + "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." + "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." 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/ShardCmd.lean b/Ix/Cli/ShardCmd.lean index e2f0e6d29..3e5a4b294 100644 --- a/Ix/Cli/ShardCmd.lean +++ b/Ix/Cli/ShardCmd.lean @@ -8,10 +8,12 @@ static walk-edge nets (the relation that generates each shard's thin frontier, i.e. its real ingress), then a global rebalance post-pass toward equal predicted Aiur FFT cost (fitted model — constants and - provenance in `ix_kernel::shard::STATIC_OWNED_PER_BYTE`). Requires - `--shards N`. Measured against the profiled strategy on the 8-shard - Init / 24-shard Std harnesses: mean shard FFT −30%, max shard −44/−51%, - stddev 17.7%→7.1% / 30.6%→8.9%. + provenance in `ix_kernel::shard::STATIC_OWNED_PER_BYTE`). `--shards N` + fixes the count; `--max-ram G` seeds it from the same static block-shape + score, mildly superlinear in both library scale and inverse RAM budget. + Measured against the profiled strategy on the 8-shard Init / 24-shard Std + harnesses: mean shard FFT −30%, max shard −44/−51%, stddev 17.7%→7.1% / + 30.6%→8.9%. - **Profiled (`--profile `)**: the original pipeline over an `ix profile` run, unchanged. Modes (precedence in `runShardCmd`): default / `--max-ram G` / `--max-cycles C` **bin-pack to a per-shard @@ -31,6 +33,7 @@ -/ module public import Cli +public import Ix.Common public import Ix.KernelCheck public import Ix.Cli.ConstsFile @@ -143,40 +146,33 @@ def runShardCmd (p : Cli.Parsed) : IO UInt32 := do | none => -- STATIC strategy (no out-of-circuit profiling): byte-balanced min-cut -- over the env's walk-edge nets + predicted-FFT rebalance post-pass. - -- Shard count comes from `--shards N`, or NAIVELY from `--max-ram G`: - -- without a profile the planner cannot know a shard's real prover - -- peak, so it estimates it from serialized env bytes via a measured - -- amplification factor and picks N so the estimated per-shard peak - -- fits the budget. The projections printed by an executed run - -- (`ix check --ixes --batch`) are the ground truth that corrects - -- the estimate — over-budget shards re-split, far-under merge. + -- `--shards N` fixes the count. Under `--max-ram G`, Rust chooses a + -- seed only after building the static block profile: the fitted owned + -- score sees both bytes and their per-block shape, unlike `.ixe` file + -- size. Boundaries and budget fitting are still measured where execution + -- happens anyway (the prove gate splits over-budget shards inline and the + -- corrected manifest remembers the leaves). if maxCycles.isSome then p.printError "error: --max-cycles requires --profile (its budget \ model is calibrated on profiled op counters)" return 1 - let n ← match shardsFlag, maxRam with - | some n, _ => pure n - | none, some gib => - if gib == 0 then - p.printError "error: --max-ram must be positive"; return 1 - let envBytes := (← System.FilePath.metadata envPath).byteSize.toNat - -- Prover-peak-per-env-byte, measured on init and FLT partitions - -- (2026-08-21): analytic prover peaks landed at ~20,000-21,000x - -- the shard's serialized bytes. Target ~2/3 of the budget so - -- the median sits under it with headroom for the measured - -- ~2.5x median-to-max spread across shards of one partition. - let amplification := 20000 - let targetBytes := gib * 1024 * 1024 * 1024 * 2 / 3 - let n := max 1 ((envBytes * amplification + targetBytes - 1) / targetBytes) - IO.println s!"[shard] naive RAM sizing: {envBytes} env bytes × \ - {amplification} / ({gib} GiB × 2/3) → {n} shard(s)" - pure n - | none, none => - p.printError "error: the static strategy (no --profile) requires \ - --shards N or --max-ram G" - return 1 - IO.println s!"Sharding {envPath} into {n} shards (static strategy, balance ±{balancePct}%)" - rsShardEnvStaticFFI envPath (toString n) (toString balancePct) outPath + match shardsFlag, maxRam with + | some n, _ => + if n == 0 then + p.printError "error: --shards must be positive"; return 1 + IO.println s!"Sharding {envPath} into {n} shards \ + (static strategy, balance ±{balancePct}%)" + rsShardEnvStaticFFI envPath (toString n) "0" (toString balancePct) outPath + | none, some gib => + if gib == 0 then + p.printError "error: --max-ram must be positive"; return 1 + IO.println s!"Sharding {envPath} for a {gib} GiB prover budget \ + (static block-shape seed, balance ±{balancePct}%)" + rsShardEnvStaticFFI envPath "0" (toString gib) (toString balancePct) outPath + | none, none => + p.printError "error: the static strategy (no --profile) requires \ + --shards N or --max-ram G" + return 1 | some espPath => -- Profiled strategy, unchanged: partition the `.ixprof`. -- Precedence: explicit --shards (fixed count) > explicit --max-cycles/--max-ram @@ -208,7 +204,7 @@ def shardCmd : Cli.Cmd := `[Cli| profile : String; "Path to a `.ixprof` from `ix profile`. When given, use the profiled strategy (cap budgeting / balanced min-cut over measured costs); when absent, the static strategy partitions the `.ixe` directly." shards : Nat; "Fixed number of shards N (static strategy; overrides --max-ram sizing and the profiled default budget sizing)" "max-cycles" : Nat; "Per-shard guest-cycle budget (profiled strategy only)" - "max-ram" : Nat; "Per-shard prover-RAM budget, GiB. Static strategy: NAIVE sizing — estimates per-shard prover peak from serialized env bytes (measured ~20,000x amplification) and picks the shard count so the estimate fits ~2/3 of the budget; executed projections then correct it (split/merge). Profiled strategy: budget from measured op counters (default: detected system RAM)." + "max-ram" : Nat; "Per-shard prover-RAM budget, GiB. Static strategy: seed the shard count from the `.ixe` block-shape score (serialized block size + a superlinear large-body term), with fitted scale and budget exponents anchored at Mathlib; the execution gate measures and corrects every boundary. Profiled strategy: budget from measured op counters (default: detected system RAM)." balance : Nat; "Per-bisection balance tolerance, percent (default 5)" parallelism : Nat; "Provers assumed for the prove-time estimate (profiled strategy only; default 1 = sequential)" out : String; "Output .ixes manifest path (default: env base name + `.ixes`, e.g. init.ixe → init.ixes)" diff --git a/Ix/Common.lean b/Ix/Common.lean index cb7bc3235..1980c644b 100644 --- a/Ix/Common.lean +++ b/Ix/Common.lean @@ -476,17 +476,24 @@ def Nat.formatMs (ms : Nat) : String := else "< 1ms" +/-- Bytes per GiB — the unit of the RAM-budget flags and every + prover-peak log line. -/ +def gibBytes : Nat := 1024 * 1024 * 1024 + +/-- Bytes rendered as GiB for log lines. -/ +def toGib (bytes : Nat) : Float := bytes.toFloat / gibBytes.toFloat + /-- Format a byte count with appropriate unit suffix (B, kB, MB, GB). -/ def fmtBytes (n : Nat) : String := if n < 1024 then s!"{n} B" else if n < 1024 * 1024 then let kb := n * 10 / 1024 s!"{kb / 10}.{kb % 10} kB" - else if n < 1024 * 1024 * 1024 then + else if n < gibBytes then let mb := n * 10 / (1024 * 1024) s!"{mb / 10}.{mb % 10} MB" else - let gb := n * 10 / (1024 * 1024 * 1024) + let gb := n * 10 / gibBytes s!"{gb / 10}.{gb % 10} GB" /-- ` · rss X.Y GiB (hwm Z.W)` from `/proc/self/status`; empty where diff --git a/Ix/KernelCheck.lean b/Ix/KernelCheck.lean index d403997a5..3c72b13e0 100644 --- a/Ix/KernelCheck.lean +++ b/Ix/KernelCheck.lean @@ -194,15 +194,17 @@ opaque rsEnvExtractFFI : @& Bool → -- quiet IO Unit -/-- FFI: partition a `.ixe` into `numShards` shards with the STATIC - strategy (no out-of-circuit profiling): byte-balanced min-cut over the - static walk-edge nets + a predicted-FFT-cost rebalance post-pass - (`ix_kernel::shard::shard_static`; model constants documented there). - Writes a `.ixes` manifest; prints the report to stderr. -/ +/-- FFI: partition a `.ixe` with the STATIC strategy (no out-of-circuit + profiling): byte-balanced min-cut over the static walk-edge nets + a + predicted-FFT-cost rebalance post-pass (`ix_kernel::shard::shard_static`; + model constants documented there). A nonzero `numShards` fixes the count; + otherwise `ramGib` selects the static block-shape seed after the `.ixe` has + been profiled. Writes a `.ixes` manifest; prints the report to stderr. -/ @[extern "rs_shard_env_static"] opaque rsShardEnvStaticFFI : @& String → -- .ixe path - @& String → -- num_shards (N) + @& String → -- num_shards (N; "0" = score seed) + @& String → -- max RAM GiB (used when N = 0) @& String → -- balance percent @& String → -- .ixes output path ("" = skip) IO Unit diff --git a/crates/aiur/src/synthesis.rs b/crates/aiur/src/synthesis.rs index 3e833a752..d2d6b4267 100644 --- a/crates/aiur/src/synthesis.rs +++ b/crates/aiur/src/synthesis.rs @@ -38,6 +38,22 @@ pub struct PeakProveBytes { pub peak: usize, } +/// Outcome of a budget-gated prove +/// ([`AiurSystem::prove_ixvm_within_budget`]). Only the case that ran a +/// STARK carries a proof; the other two report the measured peak the +/// caller decides with. +// One short-lived value per prove; the variant size gap is irrelevant. +#[allow(clippy::large_enum_variant)] +pub enum GatedProve { + /// Fit the budget; proven from the gating record. + Proved { claim: Vec, proof: AiurProof, peak: usize }, + /// Over budget: the record was dropped, and `parts` is the count + /// [`AiurSystem::suggested_split_parts`] projects will fit. + Split { peak: usize, parts: usize }, + /// `exec_only` with a fitting peak: measured, nothing left to do. + Measured { peak: usize }, +} + pub struct AiurSystem { toplevel: Toplevel, // perhaps remove the key from the system in verifier only mode? @@ -73,6 +89,27 @@ pub struct CircuitShape { pub preprocessed_height: usize, } +/// Raw row count of a circuit under `record`, ceil-divided into `parts` +/// even shares — `parts = 1` is the record's exact heights. The byte +/// 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, + parts: usize, +) -> impl Fn(usize, &CircuitType) -> usize + '_ { + move |_, ct| match ct { + CircuitType::Function { idx } => { + record.function_queries[*idx].len().div_ceil(parts) + }, + CircuitType::Memory { width } => { + record.memory_queries.get(width).map_or(0, |m| m.len().div_ceil(parts)) + }, + CircuitType::Bytes1 => 256, + CircuitType::Bytes2 => 65536, + } +} + impl AiurSystem { pub fn build( toplevel: Toplevel, @@ -230,14 +267,7 @@ impl AiurSystem { /// which per-fft models blur. pub fn peak_prove_bytes(&self, record: &QueryRecord) -> PeakProveBytes { self.peak_prove_bytes_by( - |_, ct| match ct { - CircuitType::Function { idx } => record.function_queries[*idx].len(), - CircuitType::Memory { width } => { - record.memory_queries.get(width).map_or(0, |m| m.len()) - }, - CircuitType::Bytes1 => 256, - CircuitType::Bytes2 => 65536, - }, + raw_of(record, 1), crate::execute::record_retained_bytes(record), ) } @@ -302,26 +332,67 @@ impl AiurSystem { } } - #[tracing::instrument(level = "info", skip_all, name = "aiur/prove")] - pub fn prove( + /// Smallest power-of-two part count whose projected per-part peak + /// fits `max_bytes`, assuming the record's rows divide evenly across + /// parts. The gadget circuits keep their constant heights — they are + /// the same size in every shard and are most of the model's fixed + /// floor, which dividing cannot shrink. + /// + /// The estimate is optimistic: a part re-executes dependencies shared + /// across the cut, so its real rows exceed its 1/n share. A caller + /// splitting on this number must still gate each part on its own + /// executed record and re-split the ones that miss. Optimism is the + /// right bias — an under-split costs one cheap re-execution, while an + /// over-split pays the per-proof floor on every extra part forever. + /// + /// Returns 1 when the record already fits. + pub fn suggested_split_parts( + &self, + record: &QueryRecord, + max_bytes: usize, + ) -> usize { + let record_bytes = crate::execute::record_retained_bytes(record); + let mut parts = 1usize; + // A shard still over budget at 2^20 parts is not splittable by row + // 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; + if peak <= max_bytes { + break; + } + parts *= 2; + } + parts + } + + /// Prove an execution that has ALREADY happened: everything from the + /// witness phase on, over a record the caller hands across. + /// + /// Taking `query_record` by value is the point. The record is the + /// witness phase's dominant residency, and it is dropped here the + /// instant the traces exist — before the LDE/commit/FRI phases that + /// actually set the prover's peak (see [`Self::peak_prove_bytes`]). + /// A caller that keeps its own copy alive past this call pays that + /// peak *on top of* the record, which at shard scale is the + /// difference between fitting in RAM and not. + /// + /// `input`, `io_buffer` and `output` must be the ones the execution + /// ran on: they reconstruct the claim the proof commits to, and the + /// witness reads the buffer the execution left behind. + /// + /// Deliberately not `#[tracing::instrument]`ed — the `aiur/witness` + /// span below stays directly under the caller's `aiur/prove*` span, + /// so the stage-scoped measurements keep their existing shape. + pub fn prove_from_execution( &self, fun_idx: FunIdx, input: &[G], - io_buffer: &mut IOBuffer, + io_buffer: &IOBuffer, + query_record: QueryRecord, + output: &[G], ) -> (Vec, AiurProof) { - tracing_texray::examine_current(); - - // Execute the Aiur bytecode. - let _g = tracing::info_span!("aiur/execute").entered(); - // Execute the Aiur bytecode. The prover assumes inputs are valid; any - // execution error here is a programmer bug, so we unwrap. - let (query_record, output) = self - .toplevel - .execute(fun_idx, input.to_vec(), io_buffer) - .expect("Aiur execution failed during prove"); - drop(_g); - - // Build the `SystemWitness` let _g = tracing::info_span!("aiur/witness").entered(); let circuit_types = self.circuit_types(); let witness_data = circuit_types @@ -363,6 +434,28 @@ impl AiurSystem { (claim, proof) } + #[tracing::instrument(level = "info", skip_all, name = "aiur/prove")] + pub fn prove( + &self, + fun_idx: FunIdx, + input: &[G], + io_buffer: &mut IOBuffer, + ) -> (Vec, AiurProof) { + tracing_texray::examine_current(); + + // Execute the Aiur bytecode. + let _g = tracing::info_span!("aiur/execute").entered(); + // Execute the Aiur bytecode. The prover assumes inputs are valid; any + // execution error here is a programmer bug, so we unwrap. + let (query_record, output) = self + .toplevel + .execute(fun_idx, input.to_vec(), io_buffer) + .expect("Aiur execution failed during prove"); + drop(_g); + + self.prove_from_execution(fun_idx, input, io_buffer, query_record, &output) + } + /// IxVM-native prove: identical to `prove` except the execute step /// is provided by the caller as `executor` (a closure that runs /// the codegen'd Rust kernel `ix::aiur_ixvm_runner::execute_ixvm` @@ -381,6 +474,51 @@ impl AiurSystem { io_buffer: &mut IOBuffer, executor: F, ) -> (Vec, AiurProof) + where + F: FnOnce( + &Toplevel, + FunIdx, + Vec, + &mut IOBuffer, + ) -> Result<(QueryRecord, Vec), ExecError>, + { + match self.prove_ixvm_within_budget( + fun_idx, input, io_buffer, executor, None, false, + ) { + GatedProve::Proved { claim, proof, .. } => (claim, proof), + _ => unreachable!("an unbudgeted prove always proves"), + } + } + + /// `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 + /// loop runs on executions alone, no STARK started. + /// + /// The peak is measured on the REAL record ([`Self::peak_prove_bytes`]), + /// not estimated from serialized bytes, so an over-budget shard is + /// caught in the gap between execution and the witness phase — before + /// the LDE/commit/FRI phases that would actually exhaust the box: the + /// record is dropped and [`GatedProve::Split`] carries the part count + /// [`Self::suggested_split_parts`] projects will fit, computed here + /// because this is the last moment the record exists to read counts + /// from. Every outcome carries the measured peak: proving a shard + /// measures it for free, so a prove run yields the same split/merge + /// signal a check run does without a second execution. + #[tracing::instrument( + level = "info", + skip_all, + name = "aiur/prove_ixvm_within_budget" + )] + pub fn prove_ixvm_within_budget( + &self, + fun_idx: FunIdx, + input: &[G], + io_buffer: &mut IOBuffer, + executor: F, + max_bytes: Option, + exec_only: bool, + ) -> GatedProve where F: FnOnce( &Toplevel, @@ -396,43 +534,24 @@ impl AiurSystem { .expect("IxVM-native Aiur execution failed during prove_ixvm"); drop(_g); - let _g = tracing::info_span!("aiur/witness").entered(); - let circuit_types = self.circuit_types(); - let witness_data = circuit_types - .into_par_iter() - .enumerate() - .map(|(circuit_idx, circuit_type)| { - let slot_arg_widths = self.slot_arg_widths(circuit_idx); - match circuit_type { - CircuitType::Function { idx } => self.toplevel.witness_data( - idx, - &query_record, - io_buffer, - &slot_arg_widths, - ), - CircuitType::Memory { width } => { - Memory::witness_data(width, &query_record, &slot_arg_widths) - }, - CircuitType::Bytes1 => { - Bytes1.witness_data(&query_record, &slot_arg_widths) - }, - CircuitType::Bytes2 => { - Bytes2.witness_data(&query_record, &slot_arg_widths) - }, - } - }) - .collect::>(); - drop(query_record); - let (traces, lookups) = witness_data.into_iter().unzip(); - let witness = SystemWitness { traces, lookups }; - drop(_g); - - let mut claim = vec![function_channel(), G::from_usize(fun_idx)]; - claim.extend(input); - claim.extend(output); - - let proof = self.system.prove(&self.key, &claim, witness); - (claim, proof) + let peak = self.peak_prove_bytes(&query_record).peak; + if let Some(max) = max_bytes + && peak > max + { + let parts = self.suggested_split_parts(&query_record, max); + return GatedProve::Split { peak, parts }; + } + if exec_only { + return GatedProve::Measured { peak }; + } + let (claim, proof) = self.prove_from_execution( + fun_idx, + input, + io_buffer, + query_record, + &output, + ); + GatedProve::Proved { claim, proof, peak } } #[inline] diff --git a/crates/ffi/src/aiur/protocol.rs b/crates/ffi/src/aiur/protocol.rs index 4247d342a..a304e6828 100644 --- a/crates/ffi/src/aiur/protocol.rs +++ b/crates/ffi/src/aiur/protocol.rs @@ -7,7 +7,7 @@ use std::sync::LazyLock; use lean_ffi::object::{ ExternalClass, LeanArray, LeanBorrowed, LeanByteArray, LeanExcept, - LeanExternal, LeanNat, LeanOwned, LeanProd, LeanRef, LeanString, + LeanExternal, LeanNat, LeanOption, LeanOwned, LeanProd, LeanRef, LeanString, }; use crate::{ @@ -15,13 +15,14 @@ use crate::{ lean::{ LeanAiurCircuitShape, LeanAiurCommitmentParameters, LeanAiurExecuteResult, LeanAiurFriParameters, LeanAiurIOKeyInfo, LeanAiurProveEnvResult, - LeanAiurProveResult, LeanAiurQueryCount, LeanAiurToplevel, + LeanAiurProveResult, LeanAiurQueryCount, LeanAiurShardProveResult, + LeanAiurShardResult, LeanAiurToplevel, }, }; use aiur::{ G, execute::{IOBuffer, IOKeyInfo, QueryRecord}, - synthesis::{AiurProof, AiurSystem, CircuitShape}, + synthesis::{AiurProof, AiurSystem, CircuitShape, GatedProve}, }; // ============================================================================= @@ -591,6 +592,33 @@ extern "C" fn rs_aiur_toplevel_check_addrs_with_env( LeanExcept::ok(arr) } +/// Decode a counted address-list blob (`Ix.Cli.CheckCmd.addrListsBlob`): +/// per list, a 4-byte LE count followed by that many 32-byte addresses. +/// The wire format of every partition crossing the FFI (the shard +/// batch, the manifest emit). +pub(crate) fn decode_addr_lists( + bytes: &[u8], +) -> Result>, String> { + let mut lists = Vec::new(); + let mut off = 0usize; + while off < bytes.len() { + if off + 4 > bytes.len() { + return Err("addr lists: truncated count".into()); + } + let n = + u32::from_le_bytes(bytes[off..off + 4].try_into().unwrap()) as usize; + off += 4; + if off + n * 32 > bytes.len() { + return Err("addr lists: truncated addresses".into()); + } + lists.push( + ix_common::address::Address::unpack(&bytes[off..off + n * 32]).collect(), + ); + off += n * 32; + } + Ok(lists) +} + /// Byte-weighted admission gate: a counting semaphore over estimated /// execution RSS, expressed with the std Mutex+Condvar construction. /// Bounds MEMORY in flight instead of shards in flight, so the rayon @@ -634,7 +662,42 @@ impl RamGate { /// RSS): 5.7 MB owned -> ~9.6 GB and 1.4 MB -> ~5.5 GB, giving /// ~4 GiB + ~1000x; both terms rounded up for cross-shard spread. const EXEC_RSS_FIXED_BYTES: usize = 9 * (1 << 29); // 4.5 GiB +// Two calibrations, each accurate in its own regime, combined as a MAX +// in `exec_rss_estimate` because the per-owned-byte execution footprint +// is env-dependent and the gate's contract is NEVER OOM: +// - The affine fit (4.5 GiB + 1100x) measured on ISLB's small shards +// (1.4-5.7 MB owned; a pure ratio under-reserved them and OOM'd a +// 128 GB box). +// - The pure ratio (2500x, ~2300x measured + margin) validated on +// Mathlib's full-width 233-shard batch at +2% of estimate; the +// affine slope alone under-reserved Mathlib-class shards and +// over-admitted a 132-shard full-width batch to 486/495 GB +// (OOM, 2026-08-29 — the first Mathlib full-width run under the +// affine constant). +// The max reproduces each fit where it was measured: small shards take +// the affine branch (ISLB bench reservations unchanged), large shards +// the ratio branch. const EXEC_RSS_PER_OWNED_BYTE: usize = 1100; +const EXEC_RSS_RATIO_PER_OWNED_BYTE: usize = 2500; + +/// Per-shard execution-RSS reserve: the max of the two measured fits +/// (see the constants above). +fn exec_rss_estimate(owned_bytes: usize) -> usize { + EXEC_RSS_FIXED_BYTES + .saturating_add(owned_bytes.saturating_mul(EXEC_RSS_PER_OWNED_BYTE)) + .max(owned_bytes.saturating_mul(EXEC_RSS_RATIO_PER_OWNED_BYTE)) +} + +/// Detected prover/execution RAM budget: +/// [`ix_kernel::shard::RAM_USABLE_FRAC`] of `MemAvailable`, reserving +/// the rest for the OS. `None` (no gate) when meminfo is unreadable — +/// disabling the check beats guessing at it. +#[allow(clippy::cast_precision_loss, clippy::cast_possible_truncation)] +#[allow(clippy::cast_sign_loss)] // MemAvailable and the fraction are positive +fn detected_ram_budget() -> Option { + available_ram_bytes() + .map(|b| (b as f64 * ix_kernel::shard::RAM_USABLE_FRAC) as usize) +} /// `MemAvailable` from `/proc/meminfo`, in bytes (Linux; includes /// reclaimable page cache). `None` if unreadable — the caller then @@ -646,25 +709,6 @@ fn available_ram_bytes() -> Option { Some(kib * 1024) } -/// `Bytecode.Toplevel.shardCheckBatchWithEnv`: check EVERY shard of a -/// partition in one call — rayon over the shard list with true -/// work-stealing (no chunk barriers), each shard through the exact -/// single-shard machinery (`build_shard_check_env_witness` + -/// `dispatch_execute`) over its own private record and witness io. -/// Parallel, not concurrent: tasks share only the read-only toplevel, -/// env, and the `AiurSystem` built once here for the prover RAM model. -/// -/// `shards_blob` encodes the partition as, per shard, a 4-byte LE -/// owned-constant count followed by that many 32-byte addresses. -/// Returns one `(error, peak_bytes)` pair PER SHARD in shard order: -/// an empty error string means the shard checked clean, and -/// `peak_bytes` is the analytic prover peak -/// ([`aiur::synthesis::AiurSystem::peak_prove_bytes`]) of its executed -/// record — the number split/merge decisions compare against a prover -/// budget (0 when the shard failed). `jobs = 0` uses rayon's default -/// pool width (all cores) — safe at full width because admission is -/// bounded by [`RamGate`], not by thread count; pass `jobs` only to -/// narrow CPU use. // cast_precision_loss: the [ram-gate] line renders byte counts in GiB // for humans; f64's 52-bit mantissa is exact far past any real budget. #[allow(clippy::cast_precision_loss)] @@ -681,36 +725,17 @@ extern "C" fn rs_aiur_toplevel_shard_check_batch( jobs: LeanNat>, commitment_parameters: LeanAiurCommitmentParameters>, fri_parameters: LeanAiurFriParameters>, + max_ram_bytes: LeanNat>, ) -> LeanExcept { use rayon::prelude::*; let toplevel = decode_toplevel(&toplevel_obj); let fun_idx = lean_unbox_nat_as_usize(fun_idx.inner()); let jobs = lean_unbox_nat_as_usize(jobs.inner()); - let mut shards: Vec> = Vec::new(); - { - let bytes = shards_blob.as_bytes(); - let mut off = 0usize; - while off < bytes.len() { - if off + 4 > bytes.len() { - return LeanExcept::error_string("shards_blob: truncated count"); - } - let n = - u32::from_le_bytes(bytes[off..off + 4].try_into().unwrap()) as usize; - off += 4; - if off + n * 32 > bytes.len() { - return LeanExcept::error_string("shards_blob: truncated addresses"); - } - shards.push( - bytes[off..off + n * 32] - .as_chunks::<32>() - .0 - .iter() - .map(|c| ix_common::address::Address::from_slice(c).unwrap()) - .collect(), - ); - off += n * 32; - } - } + let max_ram_bytes = lean_unbox_nat_as_usize(max_ram_bytes.inner()); + let shards = match decode_addr_lists(shards_blob.as_bytes()) { + Ok(v) => v, + Err(e) => return LeanExcept::error_string(&e), + }; let env = &env_handle.get().env; // One system build for the whole batch: the RAM model reads circuit // widths and lookup counts off the compiled circuits. @@ -727,19 +752,18 @@ extern "C" fn rs_aiur_toplevel_shard_check_batch( let estimates: Vec = shards .iter() .map(|owned| { - EXEC_RSS_FIXED_BYTES.saturating_add( + exec_rss_estimate( owned .iter() .filter_map(|a| env.get_const_bytes(a).map(|b| b.len())) - .sum::() - .saturating_mul(EXEC_RSS_PER_OWNED_BYTE), + .sum::(), ) }) .collect(); let gate = RamGate { reserved: std::sync::Mutex::new(0), cv: std::sync::Condvar::new(), - budget: available_ram_bytes().map_or(usize::MAX, |b| b / 100 * 85), + budget: detected_ram_budget().unwrap_or(usize::MAX), }; { let gib = 1024.0 * 1024.0 * 1024.0; @@ -753,7 +777,7 @@ extern "C" fn rs_aiur_toplevel_shard_check_batch( max as f64 / gib, ); } - let check_batch = || -> Vec<(String, usize)> { + let check_batch = || -> Vec<(String, usize, usize)> { shards .par_iter() .zip(estimates.par_iter()) @@ -763,7 +787,7 @@ extern "C" fn rs_aiur_toplevel_shard_check_batch( match ixvm_codegen::aiur_ixvm_witness::build_shard_check_env_witness( env, owned, ) { - Err(e) => (format!("witness build: {e}"), 0), + Err(e) => (format!("witness build: {e}"), 0, 1), Ok((_claim, input, mut io_buffer)) => match dispatch_execute( &toplevel, fun_idx, @@ -771,9 +795,19 @@ extern "C" fn rs_aiur_toplevel_shard_check_batch( &mut io_buffer, use_bytecode, ) { - Err(e) => (e, 0), + Err(e) => (e, 0, 1), + // Both reductions happen here, while the record is + // still owned by this task: it is dropped before + // `gate.release`, so admission keeps bounding peak RSS by + // the shards in flight rather than by the whole partition. Ok((record, _output)) => { - (String::new(), system.peak_prove_bytes(&record).peak) + let peak = system.peak_prove_bytes(&record).peak; + let parts = if max_ram_bytes > 0 && peak > max_ram_bytes { + system.suggested_split_parts(&record, max_ram_bytes) + } else { + 1 + }; + (String::new(), peak, parts) }, }, }; @@ -794,9 +828,12 @@ extern "C" fn rs_aiur_toplevel_shard_check_batch( pool.install(check_batch) }; let arr = LeanArray::alloc(results.len()); - for (i, (err, peak)) in results.iter().enumerate() { - arr - .set(i, LeanProd::new(LeanString::new(err), LeanOwned::box_usize(*peak))); + for (i, (err, peak, parts)) in results.iter().enumerate() { + let row = LeanAiurShardResult::alloc(0); + row.set_obj(0, LeanString::new(err)); + row.set_obj(1, LeanOwned::box_usize(*peak)); + row.set_obj(2, LeanOwned::box_usize(*parts)); + arr.set(i, row); } LeanExcept::ok(arr) } @@ -914,9 +951,32 @@ extern "C" fn rs_aiur_system_prove_addr_with_env( }) } -/// `AiurSystem.shardProveWithEnv`: per-shard prove against a -/// Rust-owned `EnvHandle`. Same `ProveEnvResult` return shape as -/// `proveAddrWithEnv`. +/// `AiurSystem.shardProveWithEnv`: per-shard prove against a Rust-owned +/// `EnvHandle`, executing ONCE and proving from that record. +/// +/// `max_ram_bytes` is a per-shard prover-RAM budget checked against the +/// executed record's projected peak, in the gap between execution and +/// the witness phase. `0` means "detect": 85% of `MemAvailable`, the +/// same policy the check batch's RAM gate uses. An unreadable +/// `/proc/meminfo` disables the check rather than guessing at it. +/// +/// Over budget, the record is dropped and `proof` comes back as `none` +/// with the measured `peakBytes` and `suggestedParts` — the part count +/// the peak model projects will fit the budget +/// ([`AiurSystem::suggested_split_parts`], measured on the record while +/// it still exists). That is a RESULT, not an error: the caller's +/// answer is to cut the shard into that many parts and prove those. +/// Learning it here costs one execution instead of an OOM part-way +/// through an FFT. `suggestedParts` is `1` whenever the prove ran. +/// +/// `exec_only` stops after execution + measurement: `proof` is `none` +/// either way, and `suggestedParts` is 1 exactly when the peak fits — +/// the split loop runs on executions alone, no STARK ever starts. +/// +/// Returns `(claimBytes, proof?, peakBytes, suggestedParts)`. The final IO buffer is +/// deliberately NOT returned: it is the shard's whole ingested byte +/// scope, both Lean callers discarded it, and marshalling it back is +/// pure cost on the largest buffers in the system. #[unsafe(no_mangle)] extern "C" fn rs_aiur_system_shard_prove_with_env( aiur_system_obj: LeanExternal>, @@ -926,9 +986,12 @@ extern "C" fn rs_aiur_system_shard_prove_with_env( LeanBorrowed<'_>, >, owned_blob: LeanByteArray>, + max_ram_bytes: LeanNat>, + exec_only: bool, ) -> LeanExcept { ffi_catch_unwind_except("AiurSystem.shardProveWithEnv", || { let fun_idx = lean_unbox_nat_as_usize(fun_idx.inner()); + let max_ram_bytes = lean_unbox_nat_as_usize(max_ram_bytes.inner()); let owned = match decode_owned_blob(&owned_blob) { Ok(v) => v, Err(e) => return LeanExcept::error_string(&e), @@ -945,14 +1008,40 @@ extern "C" fn rs_aiur_system_shard_prove_with_env( }, }; - let (_aiur_claim_arr, proof) = aiur_system_obj.get().prove_ixvm( + // 0 = detect. Matching the check batch's gate keeps one RAM policy in + // the system rather than two that can disagree. + let budget = if max_ram_bytes > 0 { + Some(max_ram_bytes) + } else { + detected_ram_budget() + }; + let proved = aiur_system_obj.get().prove_ixvm_within_budget( fun_idx, &input, &mut io_buffer, ixvm_codegen::aiur_ixvm_runner::execute_ixvm, + budget, + exec_only, ); - - LeanExcept::ok(build_prove_env_result(&claim, proof, &io_buffer)) + let (proof, peak, parts) = match proved { + GatedProve::Proved { proof, peak, .. } => ( + LeanOption::some(LeanExternal::alloc(&AIUR_PROOF_CLASS, proof)), + peak, + 1, + ), + GatedProve::Split { peak, parts } => (LeanOption::none(), peak, parts), + GatedProve::Measured { peak } => (LeanOption::none(), peak, 1), + }; + drop(io_buffer); + + let mut claim_bytes: Vec = Vec::new(); + claim.put(&mut claim_bytes); + let result = LeanAiurShardProveResult::alloc(0); + result.set_obj(0, LeanByteArray::from_bytes(&claim_bytes)); + result.set_obj(1, proof); + result.set_obj(2, LeanOwned::box_usize(peak)); + result.set_obj(3, LeanOwned::box_usize(parts)); + LeanExcept::ok(result) }) } diff --git a/crates/ffi/src/kernel.rs b/crates/ffi/src/kernel.rs index 2b7807124..72125cb1f 100644 --- a/crates/ffi/src/kernel.rs +++ b/crates/ffi/src/kernel.rs @@ -41,8 +41,8 @@ use lean_ffi::object::LeanNat; use rustc_hash::FxHashMap; use lean_ffi::object::{ - LeanArray, LeanBool, LeanBorrowed, LeanIOResult, LeanList, LeanOption, - LeanOwned, LeanProd, LeanRef, LeanString, + LeanArray, LeanBool, LeanBorrowed, LeanByteArray, LeanExternal, LeanIOResult, + LeanList, LeanOption, LeanOwned, LeanProd, LeanRef, LeanString, }; use crate::lean::LeanIxCheckError; @@ -2333,21 +2333,30 @@ fn static_block_profile(env: &IxonEnv) -> BlockProfile { builder.finish() } -/// FFI: partition a `.ixe` into `num_shards` shards with the STATIC -/// strategy — no out-of-circuit profiling run. Builds the static block -/// profile ([`static_block_profile`]), byte-balanced min-cut over the -/// walk-edge nets, then the predicted-cost rebalance post-pass +/// FFI: partition a `.ixe` with the STATIC strategy — no out-of-circuit +/// profiling run. An explicit nonzero `num_shards` fixes the count. Otherwise, +/// `ram_gib` seeds it from the static block-shape score after the profile is +/// loaded ([`ix_kernel::shard::static_seed_shards`]). Builds the static block +/// profile ([`static_block_profile`]), byte-balanced min-cut over the walk-edge +/// nets, then the predicted-cost rebalance post-pass /// (`ix_kernel::shard::shard_static`). Writes a `.ixes` manifest. #[allow(clippy::cast_precision_loss)] // balance_pct is a small percentage #[unsafe(no_mangle)] pub extern "C" fn rs_shard_env_static( env_path: LeanString>, num_shards: LeanString>, + ram_gib: LeanString>, balance_pct: LeanString>, out_path: LeanString>, ) -> LeanIOResult { let path = env_path.to_string(); - let num_shards = num_shards.to_string().parse::().unwrap_or(1); + let requested_shards = num_shards.to_string().parse::().unwrap_or(0); + let ram_gib = ram_gib.to_string().parse::().unwrap_or(0); + if requested_shards == 0 && ram_gib == 0 { + return LeanIOResult::error_string( + "rs_shard_env_static: pass a positive shard count or RAM budget", + ); + } let balance = (balance_pct.to_string().parse::().unwrap_or(5) as f64) / 100.0; let out = out_path.to_string(); @@ -2371,6 +2380,21 @@ pub extern "C" fn rs_shard_env_static( profile.num_blocks(), profile.num_edges() ); + let num_shards = if requested_shards > 0 { + requested_shards + } else { + let score = ix_kernel::shard::static_env_score(&profile); + let n = ix_kernel::shard::static_seed_shards(&profile, ram_gib); + eprintln!( + "[shard] static seed: score={score:.3e}; round({} x (score/{:.3e})^{:.2} x ({}/{ram_gib})^{:.2}) -> {n} shard(s) (heuristic; gated execution corrects every boundary)", + ix_kernel::shard::STATIC_SEED_REFERENCE_SHARDS, + ix_kernel::shard::STATIC_SEED_REFERENCE_SCORE, + ix_kernel::shard::STATIC_SEED_SCORE_EXPONENT, + ix_kernel::shard::STATIC_SEED_REFERENCE_RAM_GIB, + ix_kernel::shard::STATIC_SEED_BUDGET_EXPONENT, + ); + n + }; match ix_kernel::shard::shard_static(&profile, num_shards, balance, out_opt) { Ok(report) => { eprintln!("[rs_shard_static]\n{report}"); @@ -2380,6 +2404,97 @@ pub extern "C" fn rs_shard_env_static( } } +/// FFI: write a `.ixes` manifest describing an EXPLICIT partition — the +/// block lists a prove run actually produced (splits included), so the +/// emitted manifest is the one `ix verify --ixes` can check the run's +/// proofs against, and the one that seeds the next run with a +/// partition that already fits. +/// +/// `shards_blob` encodes, per shard, a 4-byte LE block count followed +/// by that many 32-byte block addresses. Every profile block must +/// appear in EXACTLY one shard: a prove partition covers the env by +/// construction (splits partition their parent), so a gap or duplicate +/// here is an input error, not a policy choice. +/// +/// `peaks_blob`: one 8-byte LE measured prover peak per shard in shard +/// order — the run's `peak_prove_bytes` measurements, recorded on each +/// manifest shard for schedulers to bin-pack on. +/// +/// Takes the caller's live `EnvHandle`: both call sites hold one for +/// this env already, so the emit costs a profile build, not an env +/// re-parse. +#[allow(clippy::cast_possible_truncation)] // block/shard ids are u32 by construction +#[unsafe(no_mangle)] +pub extern "C" fn rs_shard_manifest_from_partition( + env_handle: LeanExternal< + ixvm_codegen::env_handle::EnvHandle, + LeanBorrowed<'_>, + >, + shards_blob: LeanByteArray>, + peaks_blob: LeanByteArray>, + out_path: LeanString>, +) -> LeanIOResult { + let out = out_path.to_string(); + let profile = static_block_profile(&env_handle.get().env); + let block_id: FxHashMap = profile + .blocks() + .iter() + .enumerate() + .map(|(i, b)| (b.addr.clone(), i as u32)) + .collect(); + let lists = + match crate::aiur::protocol::decode_addr_lists(shards_blob.as_bytes()) { + Ok(v) => v, + Err(e) => return LeanIOResult::error_string(&e), + }; + let num_shards = lists.len(); + let mut shard_of = vec![u32::MAX; profile.num_blocks()]; + for (k, list) in lists.iter().enumerate() { + for addr in list { + let Some(&b) = block_id.get(addr) else { + return LeanIOResult::error_string(&format!( + "shard {k}: block {} not in the env's block profile", + addr.hex() + )); + }; + if shard_of[b as usize] != u32::MAX { + return LeanIOResult::error_string(&format!( + "block {} appears in shards {} and {k}", + addr.hex(), + shard_of[b as usize] + )); + } + shard_of[b as usize] = k as u32; + } + } + let uncovered = shard_of.iter().filter(|&&s| s == u32::MAX).count(); + if uncovered > 0 { + return LeanIOResult::error_string(&format!( + "partition covers {} of {} blocks ({uncovered} missing)", + profile.num_blocks() - uncovered, + profile.num_blocks() + )); + } + let peaks: Vec = peaks_blob + .as_bytes() + .as_chunks::<8>() + .0 + .iter() + .map(|c| u64::from_le_bytes(*c)) + .collect(); + match ix_kernel::shard::shard_manifest_explicit( + &profile, &shard_of, num_shards, &peaks, &out, + ) { + Ok(summary) => { + eprintln!("[shard-manifest] {out}: {summary}"); + LeanIOResult::ok(LeanOwned::box_usize(0)) + }, + Err(e) => LeanIOResult::error_string(&format!( + "rs_shard_manifest_from_partition: {e}" + )), + } +} + /// FFI: dump the static block-level reference graph of a `.ixe` as text — /// `block ` per ingress unit (a Muts block or a /// standalone constant) and `edge ` per deduped diff --git a/crates/ffi/src/lean.rs b/crates/ffi/src/lean.rs index 0cd0232a0..21db0f4a3 100644 --- a/crates/ffi/src/lean.rs +++ b/crates/ffi/src/lean.rs @@ -294,6 +294,10 @@ lean_ffi::lean_inductive! { LeanAiurProveResult [ { num_obj: 4 } ]; // claimBytes, proof, ioData, ioMap LeanAiurProveEnvResult [ { num_obj: 4 } ]; + // error, peakBytes, suggestedParts + LeanAiurShardResult [ { num_obj: 3 } ]; + // claimBytes, proof, peakBytes, suggestedParts + LeanAiurShardProveResult [ { num_obj: 4 } ]; // --- Block / comparison types --- diff --git a/crates/kernel/src/shard.rs b/crates/kernel/src/shard.rs index c06d1ae9a..670d2a195 100644 --- a/crates/kernel/src/shard.rs +++ b/crates/kernel/src/shard.rs @@ -2,15 +2,16 @@ //! cross-shard delta-unfold ingress (see `plans/sharding.md`). Two driving //! modes share the same min-cut machinery: //! -//! - **Fixed count** ([`Hypergraph::partition`], the `--shards N` CLI path): -//! recursive *balanced* min-cut bisection into exactly `N` shards — even -//! per-shard work, the bisection tree doubling as the aggregation tree. -//! - **Cap / packing** ([`partition_for_cycle_cap`], the default and -//! `--max-cycles`/`--max-ram` paths): **bin-pack to the cap** — the fewest -//! shards that each stay under a per-shard cycle (hence prover-RAM) budget, -//! each filled as full as the dependency structure allows. This does *not* -//! balance: uniformity over-shards (every shard left partly empty), whereas -//! packing hits `≈⌈total/cap⌉` shards. It still uses a fine min-cut +//! - **Fixed count** ([`Hypergraph::partition`], explicit `--shards N` and the +//! static strategy's `--max-ram` score seed): recursive *balanced* min-cut +//! bisection into exactly `N` shards — even per-shard work, the bisection tree +//! doubling as the aggregation tree. +//! - **Cap / packing** ([`partition_for_cycle_cap`], the profiled strategy's +//! default and `--max-cycles`/`--max-ram` paths): **bin-pack to the cap** — the +//! fewest shards that each stay under a per-shard cycle (hence prover-RAM) +//! budget, each filled as full as the dependency structure allows. This does +//! *not* balance: uniformity over-shards (every shard left partly empty), +//! whereas packing hits `≈⌈total/cap⌉` shards. It still uses a fine min-cut //! pre-partition for a **cut-coherent order** so dependency overlap packs into //! the same shard (overlap paid once, not re-ingressed per shard). //! @@ -1278,6 +1279,11 @@ pub struct ShardInfo { /// root over the foreign part of the shard's static reference closure. `None` /// until populated by the env-aware layer (the pure partitioner has no `Env`). pub assumption_root: Option
, + /// Measured projected prover peak of this shard's executed record + /// (`AiurSystem::peak_prove_bytes`), recorded when the manifest was + /// emitted from a run that executed the shard; 0 = unmeasured (planner + /// output). The scheduling signal: STARK prove wall is ~linear in it. + pub measured_peak_bytes: u64, } /// The sharding manifest: the partition plus its cost metrics. Assumption-tree @@ -1345,6 +1351,7 @@ impl ShardManifest { foreign_blocks, cross_ingress, assumption_root: None, + measured_peak_bytes: 0, }); } ShardManifest { @@ -1383,9 +1390,28 @@ impl ShardManifest { let empty = self.shards.iter().filter(|s| s.blocks.is_empty()).count(); let max_cross = self.shards.iter().map(|s| s.cross_ingress).max().unwrap_or(0); + let peaks: Vec = self + .shards + .iter() + .map(|s| s.measured_peak_bytes) + .filter(|&p| p != 0) + .collect(); + #[allow(clippy::cast_precision_loss)] + let measured = if peaks.is_empty() { + String::new() + } else { + let gib = |b: u64| b as f64 / (1u64 << 30) as f64; + format!( + " measured-peaks[{}/{} shards, min={:.1} max={:.1} GiB]", + peaks.len(), + self.shards.len(), + gib(peaks.iter().copied().min().unwrap_or(0)), + gib(peaks.iter().copied().max().unwrap_or(0)), + ) + }; format!( "shards={} (empty={}) heartbeats[min={} mean={} max={}] imbalance={:.2}x \ - cross_ingress_total={} max_shard_cross={}", + cross_ingress_total={} max_shard_cross={}{measured}", self.shards.len(), empty, min, @@ -1434,6 +1460,18 @@ impl ShardManifest { }, None => out.push(0), } + // Trailing measured-peaks section (same older-readers-ignore contract + // as the tree): presence byte, then one u64 LE per shard in order. + // Absent when nothing was measured, so planner manifests stay + // byte-identical to the pre-peaks format. + if self.shards.iter().any(|s| s.measured_peak_bytes != 0) { + out.push(1); + for sh in &self.shards { + out.extend_from_slice(&sh.measured_peak_bytes.to_le_bytes()); + } + } else { + out.push(0); + } out } @@ -1462,6 +1500,7 @@ impl ShardManifest { foreign_blocks, cross_ingress, assumption_root, + measured_peak_bytes: 0, }); } // Optional trailing tree section. Absent (end-of-input) on pre-tree @@ -1471,6 +1510,12 @@ impl ShardManifest { } else { None }; + // Optional trailing measured-peaks section; absent on older manifests. + if c.pos < c.buf.len() && c.u8()? == 1 { + for sh in &mut shards { + sh.measured_peak_bytes = c.u64()?; + } + } // The tree is the aggregation plan: a leaf set that is not exactly the // shard id set (each id once) would silently drop or duplicate proven // shards in the fold, so a manifest carrying such a tree is invalid. @@ -1573,12 +1618,73 @@ pub const STATIC_OWNED_SUPERLINEAR: f64 = 681.08; /// See [`STATIC_OWNED_PER_BYTE`]. pub const STATIC_FRONTIER_PER_BYTE: f64 = 38_000.0; +/// Reference point for the profile-free shard-count seed: Mathlib's static +/// owned-side score under this model. At a 400 GiB prover budget, 233 seed +/// shards gave the established Mathlib-class behavior (a small number of +/// heavy-tail splits, corrected by the execution gate). +pub const STATIC_SEED_REFERENCE_SCORE: f64 = 1.254e14; +/// Mathlib seed count at [`STATIC_SEED_REFERENCE_RAM_GIB`]. +pub const STATIC_SEED_REFERENCE_SHARDS: f64 = 233.0; +/// Prover-RAM budget of the reference calibration. +pub const STATIC_SEED_REFERENCE_RAM_GIB: f64 = 400.0; +/// Mild superlinear scale correction over the static owned-side score. +/// +/// A linear score model still over-counted the small environments: cross-shard +/// duplication and the heavy-tail risk both grow with library scale. The 1.10 +/// exponent preserves the Mathlib reference while fitting the measured knees +/// at Init (8 shards) and ISLB (23) at the 400 GiB reference budget. +pub const STATIC_SEED_SCORE_EXPONENT: f64 = 1.10; +/// Superlinear correction when scaling the seed away from 400 GiB. +/// +/// Fitted on the 2026-08-28 100 GiB calibration: linear scaling caused broad +/// correction (Init 32→48 leaves; ISLB 93→123), while 1.20 seeded Init at 43 +/// (42 measured clean) and ISLB at 122 (four initial misses, 128 leaves). +/// Out-of-sample Batteries seeded 74 with one initial miss and 79 leaves. The +/// discarded 1.36 candidate over-seeded ISLB at 152→156. Unlike the score +/// exponent, this changes no reference-budget prediction. +pub const STATIC_SEED_BUDGET_EXPONENT: f64 = 1.20; + /// Predicted owned-side cost of one block under the static model. fn static_owned_weight(size: u32) -> f64 { let s = f64::from(size); STATIC_OWNED_PER_BYTE * s + STATIC_OWNED_SUPERLINEAR * s * s.sqrt() } +/// Whole-environment owned-side score used only to seed a shard count. +/// +/// There is no frontier when the environment is viewed as one set, so this is +/// just the sum of [`static_owned_weight`] over blocks. Unlike raw `.ixe` file +/// bytes, it sees how bytes are packaged: the superlinear term distinguishes a +/// large reduction-heavy proof body from the same bytes spread over many small +/// declarations. The actual partition still includes frontier cost and is +/// always checked by the execution-time prover-RAM gate. +pub fn static_env_score(profile: &BlockProfile) -> f64 { + profile.blocks().iter().map(|b| static_owned_weight(b.serialized_size)).sum() +} + +/// Score/budget model for the profile-free shard-count seed. +/// +/// This deliberately rounds to the nearest shard instead of always rounding +/// upward: it is a seed, not a safety boundary, and the prove/check execution +/// gate measures and splits every over-budget shard before a STARK starts. +/// Avoiding an unconditional upward bias matters most for small environments. +fn static_seed_shards_for_score(score: f64, ram_gib: u64) -> usize { + let budget = ram_gib.max(1) as f64; + let scaled = STATIC_SEED_REFERENCE_SHARDS + * (score / STATIC_SEED_REFERENCE_SCORE).powf(STATIC_SEED_SCORE_EXPONENT) + * (STATIC_SEED_REFERENCE_RAM_GIB / budget) + .powf(STATIC_SEED_BUDGET_EXPONENT); + scaled.round().max(1.0) as usize +} + +/// Suggested static seed count for `profile` at a per-shard prover-RAM budget. +/// The result is capped at one shard per atomic block; finer partitioning could +/// only create empty shards and cannot make an oversized block splittable. +pub fn static_seed_shards(profile: &BlockProfile, ram_gib: u64) -> usize { + static_seed_shards_for_score(static_env_score(profile), ram_gib) + .min(profile.num_blocks().max(1)) +} + /// Greedy global rebalance toward equal predicted per-shard cost under the /// static model. Repeatedly moves the best block from the most expensive /// shard to the cheapest until the hot/cold gap is within 1% of the mean @@ -1797,14 +1903,7 @@ pub fn shard_static( ); let mut manifest = ShardManifest::build(profile, &shard_of, num_shards).with_tree(tree); - for shard in &mut manifest.shards { - shard.assumption_root = - ixon::merkle::merkle_root_canonical(&shard.foreign_blocks); - } - if let Some(op) = out_path { - std::fs::write(op, manifest.to_bytes()) - .map_err(|e| format!("write {op}: {e}"))?; - } + seal_and_write(&mut manifest, out_path)?; let costs = static_predicted_costs(profile, &shard_of, num_shards); let (mut lo, mut hi, mut sum) = (f64::INFINITY, 0.0f64, 0.0f64); for &c in &costs { @@ -1849,14 +1948,7 @@ pub fn shard_esp( let (shard_of, tree) = h.partition_with_tree(num_shards, balance); let mut manifest = ShardManifest::build(&profile, &shard_of, num_shards).with_tree(tree); - for shard in &mut manifest.shards { - shard.assumption_root = - ixon::merkle::merkle_root_canonical(&shard.foreign_blocks); - } - if let Some(op) = out_path { - std::fs::write(op, manifest.to_bytes()) - .map_err(|e| format!("write {op}: {e}"))?; - } + seal_and_write(&mut manifest, out_path)?; // The largest single block's heartbeats is the *floor* on achievable // per-shard balance: a mutual block is atomic and cannot be split, so no // partition can drive max-shard heartbeats below it. When the heaviest shard @@ -1885,6 +1977,57 @@ pub fn shard_esp( )) } +/// Seal every shard's assumption root (the canonical merkle root over +/// its foreign blocks) and write the manifest when a path is given. +/// Shared tail of every manifest-producing entry point. +fn seal_and_write( + manifest: &mut ShardManifest, + out_path: Option<&str>, +) -> Result<(), String> { + for shard in &mut manifest.shards { + shard.assumption_root = + ixon::merkle::merkle_root_canonical(&shard.foreign_blocks); + } + if let Some(op) = out_path { + std::fs::write(op, manifest.to_bytes()) + .map_err(|e| format!("write {op}: {e}"))?; + } + Ok(()) +} + +/// Write a `.ixes` manifest for an EXPLICIT shard assignment — the +/// partition a prove run actually produced (splits included) rather +/// than a planner's output. Same manifest construction as +/// [`shard_static`]'s tail: [`ShardManifest::build`] recomputes +/// own sizes, foreign blocks and cross-ingress from the profile, and +/// each shard's assumption root is the canonical merkle root of its +/// foreign blocks. No aggregation tree is attached (consumers fall +/// back to the flat tree-fold); heartbeat figures carry whatever proxy +/// the profile was built with (serialized size, on the static path). +/// `peaks` (empty, or one measured prover-peak per shard in order) +/// fills each shard's `measured_peak_bytes` — the run's measurement +/// riding along for schedulers to bin-pack on. +pub fn shard_manifest_explicit( + profile: &BlockProfile, + shard_of: &[u32], + num_shards: usize, + peaks: &[u64], + out_path: &str, +) -> Result { + if peaks.len() != num_shards { + return Err(format!( + "measured peaks: {} values for {num_shards} shards", + peaks.len() + )); + } + let mut manifest = ShardManifest::build(profile, shard_of, num_shards); + for (shard, &peak) in manifest.shards.iter_mut().zip(peaks) { + shard.measured_peak_bytes = peak; + } + seal_and_write(&mut manifest, Some(out_path))?; + Ok(manifest.summary()) +} + /// Like [`shard_esp`] but sized to a per-shard Zisk **cycle** budget /// (`max_cycles`) rather than a fixed shard count: grows `N` until the heaviest /// splittable shard fits the budget (see [`partition_for_cycle_cap`]). Use @@ -1904,14 +2047,7 @@ pub fn shard_esp_cap( let mut manifest = ShardManifest::build(&profile, &plan.shard_of, plan.num_shards) .with_tree(plan.tree); - for shard in &mut manifest.shards { - shard.assumption_root = - ixon::merkle::merkle_root_canonical(&shard.foreign_blocks); - } - if let Some(op) = out_path { - std::fs::write(op, manifest.to_bytes()) - .map_err(|e| format!("write {op}: {e}"))?; - } + seal_and_write(&mut manifest, out_path)?; let note = if plan.infeasible_atomic_floor { "\n [INFEASIBLE: a single atomic block exceeds the cap — split it upstream, raise the cap, or use a bigger box]" } else { @@ -2469,6 +2605,62 @@ mod tests { ); } + #[test] + fn static_seed_model_scales_by_score_and_budget() { + // 400 GiB calibration points from the 2026-08-28 sweep. Init and ISLB + // are the small/mid-size knees; Mathlib is the normalization point. FLT + // records the large-env extrapolation, not an optimum (its old corrected + // leaf count depended on the previous seed). These are scores, not raw + // env bytes. + assert_eq!(static_seed_shards_for_score(6.010e12, 400), 8); // Init + assert_eq!(static_seed_shards_for_score(1.536e13, 400), 23); // ISLB + assert_eq!( + static_seed_shards_for_score(STATIC_SEED_REFERENCE_SCORE, 400), + 233 + ); + assert_eq!(static_seed_shards_for_score(1.316e14, 400), 246); // FLT + + // The 100 GiB calibration keeps broad split correction out of Init and + // ISLB without the 1.36 experiment's over-sharding. + assert_eq!(static_seed_shards_for_score(6.010e12, 100), 43); // Init + assert_eq!(static_seed_shards_for_score(1.536e13, 100), 122); // ISLB + assert_eq!( + static_seed_shards_for_score(STATIC_SEED_REFERENCE_SCORE, 100), + 1230 + ); + + // An empty or vanishingly small score still produces a usable one-shard + // seed. + assert_eq!(static_seed_shards_for_score(0.0, 400), 1); + } + + #[test] + fn static_env_score_distinguishes_block_shape() { + let mut concentrated = ProfileBuilder::new(); + concentrated.block(addr(1), 0, 4000, 1, OpCounts::default()); + + let mut spread = ProfileBuilder::new(); + for i in 1..=4u8 { + spread.block(addr(i), 0, 1000, 1, OpCounts::default()); + } + + assert!( + static_env_score(&concentrated.finish()) + > static_env_score(&spread.finish()) + ); + } + + #[test] + fn static_seed_count_never_exceeds_atomic_blocks() { + let mut b = ProfileBuilder::new(); + for i in 1..=6u8 { + // Force the uncapped estimate far above six shards. + b.block(addr(i), 0, u32::MAX, 1, OpCounts::default()); + } + let p = b.finish(); + assert_eq!(static_seed_shards(&p, 1), p.num_blocks()); + } + #[test] fn objective_matches_manifest_cross_ingress() { let p = two_clusters();