Sb/cluster

Ix/Aiur/Protocol.lean

@@ -211,13 +211,16 @@ proving, and persistence. `proofHexes` is one store address per line;
 `reproveSlotCode` is zero for a full run and `slot + 1` for a targeted replay;
 the latter loads and verifies only the target's immediate cached children.
 When `writeOutputs` is false, proofs are hashed but neither the store nor cache
-is changed. Returns the root or replayed proof address. -/
+is changed. Returns the root or replayed proof address. With `verifyOnly`, the
+run stops after the parallel proof import — every shard claim reconstructed
+natively and every supplied proof bound to its shard and verified, exactly one
+per shard — returning the empty string (`ix verify --ixes` composed verdict). -/
 @[extern "rs_aiur_stage2_aggregate"]
 opaque aggregateStage2 (ixvmSystem aggrSystem : @& AiurSystem)
   (envHandle : @& EnvHandle) (manifestPath proofHexes : @& String)
   (verifyIdx aggrIdx jobs ramBudgetBytes structuralAbove reproveSlotCode : @& Nat)
   (directJoins planOnly : Bool) (cacheFriBytes : @& ByteArray)
-  (useCache writeOutputs : Bool) :
+  (useCache writeOutputs verifyOnly : Bool) :
     Except String String
 
 /-- Reconstruct and audit the manifest-relative aggregate root entirely in
@@ -311,6 +314,18 @@ opaque proofToAdviceBytes : @& AiurSystem →
 
 end AiurSystem
 
+/-- Native manifest-leaf proving with private split healing. `blocks` and
+`owned` use the counted address-list format; `ids` are newline-separated.
+`maxRam` is a positive byte count limiting the predicted proving peak of
+each executed query record. Oversized records are dropped and split.
+Only original leaf proofs are printed. An empty index path disables indexing.
+Every split is rejoined canonically and checked against the original claim. -/
+@[extern "rs_aiur_shard_pipeline"]
+opaque shardPipeline : @& AiurSystem → @& AiurSystem → @& EnvHandle →
+  @& ByteArray → @& ByteArray → @& String → @& Nat → @& Nat →
+  @& Nat → @& String → @& String → @& String →
+  Bool → Bool → Bool → IO (Except String String)
+
 /-- Write a `.ixes` manifest for an EXPLICIT partition — the block lists
     a run actually produced (splits included) rather than a planner's
     output. `shardsBlob`: per shard, a 4-byte LE block count followed by

Ix/Cli/AggregateCmd.lean

@@ -16,6 +16,7 @@
 -/
 module
 import Std.Sync
+import Ix.TracingTexray
 public import Cli
 public import Ix.Aggr
 public import Ix.Cli.CheckCmd
@@ -101,37 +102,52 @@ def aggregateLiftRamBytes : Nat := 195 * aggregateGiB
 benchmark/test union is migrated in M1-e/M1-f. -/
 def aggregateWrapRamBytes : Nat := aggregateLiftRamBytes
 
-/-- Structural joins are dominated by the same two recursive-proof checks as
-lifts. Keep the conservative lift reserve until the real E2E calibration. -/
+/-- Base reserve for self-pairs, before the flat or structural subject term. -/
 def aggregateStructuralJoinRamBytes : Nat := aggregateLiftRamBytes
 
 /-- Native verification and serialization of a raw shard proof in direct mode
 is charged to its consuming pair. -/
 def aggregateRawShardRamBytes : Nat := 4 * aggregateGiB
 
 /-- Measured upper envelope for an `IxVM + IxVM` pair (shapes 2/6). -/
-def aggregateDirectJoinRamBytes : Nat := 390 * aggregateGiB
+def aggregateDirectJoinRamBytes : Nat := 180 * aggregateGiB
 
-/-- Measured upper envelope for a mixed recursive/IxVM pair (shapes 3/4/7/8). -/
-def aggregateMixedJoinRamBytes : Nat := 340 * aggregateGiB
+/-- Mixed recursive/IxVM pairs (shapes 3/4/7/8) predicted up to 385 GiB in
+the 2026-09-09 Mathlib run, exceeding the previous 340 GiB reserve. -/
+def aggregateMixedJoinRamBytes : Nat := 180 * aggregateGiB
 
 /-- Flat joins add canonical subject-tree work to the recursive-proof base.
 One MiB per subject is a deliberately conservative placeholder: at Init's
 ~52k-subject root it adds ~51 GiB, consistent with the §11.4 estimate. The
 default structural threshold caps this term near 4 GiB in production. -/
 def aggregateFlatJoinRamPerSubjectBytes : Nat := 1024 * 1024
 
-/-- Per-shape RAM weight used by the Lean admission gate. Shape 5 retains the
-flat subject-count reserve; shape 9 is the O(1)-subject structural arm. The
-direct/mixed values are conservative round-ups of the §3.4 measurements. -/
+/-- Mathlib structural self-joins need a subject-dependent reserve despite
+their O(1) subject-root fold: assumption/path checks and child verification
+can grow. Use a subject term and a doubled base above 64k subjects to cover
+trace-size steps (380.5 GiB predicted at 91068 subjects). See
+`exp/design/numa-slot-pinning.md` §12. Keep these constants in sync with
+Rust's `STRUCTURAL_RAM_PER_SUBJECT` and `STRUCTURAL_LARGE_SUBJECTS`. -/
+def aggregateStructuralJoinRamPerSubjectBytes : Nat := 5 * 1024 * 1024 / 4
+
+def aggregateStructuralLargeSubjects : Nat := 64 * 1024
+
+/-- Per-shape RAM weight used by the Lean admission gate. Both self-pair
+shapes (5/9) reserve subject-dependent work. The direct/mixed values are
+conservative round-ups of the §3.4 measurements. -/
 def aggregateShapeRamBytes (shape subjectCount : Nat) : Nat :=
   match shape with
   | 0 | 1 => aggregateWrapRamBytes
   | 2 | 6 => aggregateDirectJoinRamBytes
   | 3 | 4 | 7 | 8 => aggregateMixedJoinRamBytes
   | 5 => aggregateStructuralJoinRamBytes +
       subjectCount * aggregateFlatJoinRamPerSubjectBytes
-  | 9 => aggregateStructuralJoinRamBytes
+  | 9 =>
+      let weight := aggregateStructuralJoinRamBytes +
+        subjectCount * aggregateStructuralJoinRamPerSubjectBytes
+      if subjectCount > aggregateStructuralLargeSubjects then
+        max weight (2 * aggregateStructuralJoinRamBytes)
+      else weight
   | _ => aggregateDirectJoinRamBytes
 
 /-- Calibration-pending per-slot RAM weight used by the Lean admission gate.
@@ -143,7 +159,7 @@ def aggregateSlotRamBytes (item : ScheduledFold) : Nat :=
     | .leaf _ => if item.kind == .ixvm then aggregateRawShardRamBytes
         else aggregateWrapRamBytes
     | .join _ _ =>
-      if item.structural then aggregateStructuralJoinRamBytes
+      if item.structural then aggregateShapeRamBytes 9 item.subjectCount
       else aggregateStructuralJoinRamBytes +
         item.subjectCount * aggregateFlatJoinRamPerSubjectBytes
 
@@ -972,13 +988,16 @@ private def runAggregateCmdNativeWith
     | .ok backend => pure backend
   let verifyIdx := ixvmBackend.compiled.getFuncIdx `verify_claim |>.get!
   let aggrIdx := aggrBackend.compiled.getFuncIdx `ix_aggr |>.get!
+  -- Streamed `[texray]` per-span lines (execute / witness / STARK) for every
+  -- Stage 2 slot, as `ix prove --texray` does for shards.
+  if p.hasFlag "texray" then TracingTexray.init {}
   let nativeResult ← IO.lazyPure fun _ =>
     ixvmBackend.system.aggregateStage2 aggrBackend.system envHandle
       manifestPath proofHexes verifyIdx aggrIdx jobs ramBudgetBytes
       structuralAbove reproveSlotCode (p.hasFlag "direct-joins")
       (p.hasFlag "plan-only")
       recursionParameters.cacheFriBytes (!(p.hasFlag "no-cache"))
-      (!(p.hasFlag "no-write"))
+      (!(p.hasFlag "no-write")) false
   match nativeResult with
   | .error e => IO.eprintln s!"aggregate failed: {e}"; return 1
   | .ok _ => return 0
@@ -1017,6 +1036,9 @@ private def runAggregateCmdLeanReferenceWith
   let structuralAbove := ((p.flag? "structural-above").map (·.as! Nat)).getD
     defaultStructuralAbove
   let directJoins := p.hasFlag "direct-joins"
+  -- Same streamed `[texray]` per-span lines as `ix prove --texray`: the
+  -- execute / witness / STARK split of every Stage 2 slot.
+  if p.hasFlag "texray" then TracingTexray.init {}
   let plan ← match schedulePlan view.aggregationTree.foldPlan shardCounts
       structuralAbove directJoins with
     | .error e => IO.eprintln e; return 1
@@ -1035,13 +1057,20 @@ private def runAggregateCmdLeanReferenceWith
   let budgetSource := if maxRamGb?.isSome then "--max-ram" else "92% MemTotal"
   IO.println s!"[aggregate] scheduler: jobs={jobsLabel}, RAM budget \
     {formatAggregateGiB ramBudgetBytes} GiB ({budgetSource}); \
-    wrap/self reserve {formatAggregateGiB aggregateWrapRamBytes} GiB, \
+    wrap/self base {formatAggregateGiB aggregateWrapRamBytes} GiB, \
     direct {formatAggregateGiB aggregateDirectJoinRamBytes} GiB, mixed \
-    {formatAggregateGiB aggregateMixedJoinRamBytes} GiB, flat +1 MiB/subject"
+    {formatAggregateGiB aggregateMixedJoinRamBytes} GiB, flat self +1 MiB/subject, \
+    structural self +1.25 MiB/subject (minimum \
+    {formatAggregateGiB (2 * aggregateStructuralJoinRamBytes)} GiB above \
+    {aggregateStructuralLargeSubjects} subjects)"
   if p.hasFlag "plan-only" then return 0
 
   let proofHexes := (p.variableArgsAs! String).toList
-  if proofHexes.length != view.shards.size then
+  -- `IX_AGGREGATE_SHARDS` (experiment knob, see `shard_selection` in
+  -- aggregate.rs) aggregates a subtree; the native side then requires one
+  -- proof per selected shard and ignores the rest.
+  let partialRun := (← IO.getEnv "IX_AGGREGATE_SHARDS").isSome
+  if !partialRun && proofHexes.length != view.shards.size then
     IO.eprintln s!"aggregate requires exactly {view.shards.size} shard proofs; got {proofHexes.length}"
     return 1
 
@@ -1210,6 +1239,7 @@ def aggregateCmd : Cli.Cmd := `[Cli|
     "max-ram" : Nat; "Aggregate in-flight RAM budget in GiB (default: 92% of MemTotal). An estimated-oversized slot runs alone."
     "structural-above" : Nat; "Use structural joins when a node contains more than N subject leaves (default 4096; 0 means every join)."
     "direct-joins";  "Keep IxVM leaves raw until their first pair instead of wrapping first (non-default; substantially higher RAM)."
+    "texray";        "Stream per-phase `[texray]` timing/RSS lines (execute, witness, STARK stages) for every slot to stderr, as `ix prove --texray` does."
 
   ARGS:
     ...proofs : String; "Persisted shard-proof wrapper addresses, in any order (one per nonempty shard, except --plan-only or replay with aggregate children)."

Ix/Cli/ProveCmd.lean

@@ -237,6 +237,58 @@ def reportPartition (proven : Array (Array Address × Nat)) (planned : Nat) : IO
       ({proven.size - planned} from splits) — re-shard with this partition to \
       skip the splits next run"
 
+private def runShardPipeline (p : Cli.Parsed) (ixe manifest : String)
+    (aiurSystem : Aiur.AiurSystem) (compiled : Aiur.CompiledToplevel)
+    (maxRamBytes : Nat) (indexDir? : Option System.FilePath) : IO UInt32 := do
+  let (ixonEnv, shards) ← match ← Ix.Cli.CheckCmd.loadEnvAndShards manifest ixe with
+    | .error e => IO.eprintln e; return 1
+    | .ok value => pure value
+  let selected ← match (p.flag? "shard").map (·.as! Nat), (p.flag? "shards").map (·.as! String) with
+    | some k, none => pure #[k]
+    | none, some s => match Ix.Cli.CheckCmd.parseShardSelection s with
+      | .error e => IO.eprintln e; return 1
+      | .ok ids => pure ids
+    | none, none => pure (Array.range shards.size)
+    | some _, some _ => IO.eprintln "use only one of --shard and --shards"; return 1
+  if let some k := selected.find? (· ≥ shards.size) then
+    IO.eprintln s!"shard {k} out of range ({shards.size} shards)"
+    return 1
+  let envHandle ← match Aiur.EnvHandle.fromIxe ixe with
+    | .error e => IO.eprintln e; return 1
+    | .ok handle => pure handle
+  let some verifyIdx := compiled.getFuncIdx `verify_claim
+    | IO.eprintln "verify_claim entrypoint missing"; return 1
+  let recursion ← match ShardProofIndex.buildRecursionBackend aiurSystem verifyIdx with
+    | .error e => IO.eprintln e; return 1
+    | .ok backend => pure backend
+  let owned := Ix.Cli.CheckCmd.ownedConstsPer ixonEnv shards
+  let storePath ← StoreIO.toIO Store.storeDir
+  let plans ← StoreIO.toIO (Store.cacheDir "shard-splits")
+  -- `id<TAB>measuredPeakBytes` per selected leaf: the native pipeline balances
+  -- its NUMA lanes by measured prover peak when the manifest carries one
+  -- (`ix shard refine`), and falls back to block counts otherwise.
+  let peaks : Array Nat ← do
+    match Ix.Cli.CheckCmd.parseIxesManifest (← IO.FS.readBinFile manifest) with
+    | .ok view => pure (selected.map fun k =>
+        match view.shardIds.findIdx? (· == k) with
+        | some i => (view.measuredPeakBytes[i]?).getD 0
+        | none => (view.measuredPeakBytes[k]?).getD 0)
+    | .error _ => pure (selected.map fun _ => 0)
+  let ids := String.intercalate "\n"
+    ((selected.zip peaks).toList.map fun (k, peak) => s!"{k}\t{peak}")
+  match ← Aiur.shardPipeline aiurSystem recursion.system envHandle
+      (Ix.Cli.CheckCmd.addrListsBlob (selected.map (shards[·]!)))
+      (Ix.Cli.CheckCmd.addrListsBlob (selected.map (owned[·]!))) ids
+      verifyIdx recursion.aggrIdx maxRamBytes storePath.toString
+      (indexDir?.map (·.toString) |>.getD "") plans.toString
+      (p.hasFlag "lookahead") (p.hasFlag "skip-proven") (p.hasFlag "keep-going") with
+  | .error e => IO.eprintln s!"[prove] {e}"; return 1
+  | .ok summary =>
+    IO.eprintln s!"[prove] {summary}; original manifest claims preserved"
+    if let some out := (p.flag? "out-ixes").map (·.as! String) then
+      IO.FS.writeBinFile out (← IO.FS.readBinFile manifest)
+    return 0
+
 def runProveCmd (p : Cli.Parsed) : IO UInt32 := do
   -- Streamed `[texray] <span>: <dur> ── RAM Δ/peak` lines on stderr as
   -- each `aiur/` / `stark/` span closes: the per-phase wall + RSS
@@ -245,10 +297,19 @@ def runProveCmd (p : Cli.Parsed) : IO UInt32 := do
   let keepGoing := p.hasFlag "keep-going"
   -- Same units as `ix shard --max-ram`: the per-shard prover budget the
   -- partition was sized against, re-checked here against each shard's
-  -- measured peak. 0 = detect (85% of `MemAvailable`, the check batch's
-  -- gate policy — see `shardProveWithEnv`).
+  -- predicted peak. The legacy path detects a budget when this is 0
+  -- (see `shardProveWithEnv`); split healing requires an explicit budget.
   let maxRamBytes :=
     ((p.flag? "max-ram").map (·.as! Nat)).getD 0 * gibBytes
+  let healSplits := p.hasFlag "heal-splits" || p.hasFlag "lookahead"
+  if healSplits then
+    if !(p.hasFlag "ixe" && p.hasFlag "ixes") || p.hasFlag "exec-only" ||
+        p.hasFlag "claim" || !(p.variableArgsAs! String).isEmpty then
+      IO.eprintln "--heal-splits/--lookahead require --ixe and --ixes, without --exec-only, --claim or names"
+      return 1
+    if maxRamBytes == 0 then
+      IO.eprintln "--heal-splits/--lookahead require explicit positive --max-ram (GiB)"
+      return 1
   let execOnly := p.hasFlag "exec-only"
   let outIxes := (p.flag? "out-ixes").map (·.as! String)
   -- The shard-proof index: written for every proof this run persists,
@@ -277,6 +338,9 @@ def runProveCmd (p : Cli.Parsed) : IO UInt32 := do
     | .error e => IO.eprintln s!"compilation failed: {e}"; return 1
     | .ok c => pure c
   let aiurSystem := Aiur.AiurSystem.build compiled.bytecode commitmentParameters friParameters
+  if healSplits then
+    return ← runShardPipeline p (p.flag! "ixe" |>.as! String)
+      (p.flag! "ixes" |>.as! String) aiurSystem compiled maxRamBytes indexDir?
   let runOne := proveOne aiurSystem compiled
   match ixePath, (p.flag? "ixes").map (·.as! String), (p.flag? "shard").map (·.as! Nat) with
   | some ixe, some manifest, some k =>
@@ -373,9 +437,11 @@ def proveCmd : Cli.Cmd := `[Cli|
     "shards" : String;  "With --ixes and no --shard: prove only these leaves — `K`, `a-b`, or a comma list of those — in one process (one env load); every other leaf is carried over unchanged by --out-ixes."
     "out-ixes" : String; "Write the partition this run actually proved — splits included — as a `.ixes` manifest to this path: the manifest `ix verify --ixes` checks these proofs against, and the one the next run of this env should start from. Skipped if any shard failed."
     "exec-only";        "Execute each shard and measure its projected prover peak, splitting over-budget shards as usual, but never start a STARK. The cheap way to audit a partition's split behavior at scale."
+    "heal-splits";      "With --ixe/--ixes: split oversized shards privately and prove flat joins back to each original CheckEnv claim. Requires explicit positive --max-ram; --out-ixes keeps the original manifest."
+    "lookahead";        "Enable split healing and execute at most one next shard while the current proof runs. Check each executed record's predicted proving peak against --max-ram before proving it."
     "skip-proven";      "With --ixes: before executing a leaf, look its claim up in the shard-proof index (`~/.ix/cache/shard-proofs/<claim-digest>`); a recorded proof that decodes, bundles exactly that claim and verifies natively is reused — its address printed, nothing executed — instead of proving again. How a partially proved partition resumes after a refinement."
     "no-index";         "Neither read nor write the shard-proof index (every persisted proof is normally recorded there under its claim digest)."
-    "max-ram" : Nat;    "Per-shard prover-RAM budget, GiB — normally the same value the partition was sized with (`ix shard --max-ram`). Each shard is executed, its projected prover peak measured on the resulting record, and the proof attempted only if it fits; an over-budget shard is cut into the part count the peak model projects will fit, and each part re-gated, instead of being taken into the FFT phases that would exhaust the box. Omit to detect: 85% of the machine's available RAM."
+    "max-ram" : Nat;    "Per-shard prover-RAM budget, GiB — normally the same value the partition was sized with (`ix shard --max-ram`). Each shard is executed, its projected prover peak measured on the resulting record, and the proof attempted only if it fits; an over-budget shard is cut into the part count the peak model projects will fit, and each part re-gated, instead of being taken into the FFT phases that would exhaust the box. Required with --heal-splits/--lookahead; otherwise omit to detect 85% of the machine's available RAM."
 
   ARGS:
     ...names : String; "Fully-qualified Lean.Name(s) to prove. With none, iterate every named constant in the env (sorted)."