perf: Tune ix shard counts and recalibrate for proving via --ram-budget

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
 

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)

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. -/

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})"