feat(aiur): Accurate whole-env sharding based on execution measurements

.github/workflows/bencher-thresholds-reset.yml

@@ -22,7 +22,7 @@ name: Bencher thresholds reset
 #     cancel by removing it before merge. Naming convention: one label per token,
 #     `bencher-thresholds-reset:<token>` where <token> is a workload (a backend
 #     testbed in Ix/Cli/BenchCmd.lean (backendSpecs) minus its runner-arch suffix:
-#     `ix-compile`, `ix-decompile`, `aiur-check-execute`, `aiur-check-prove`, `aiur-check-recursive`, `aiur-recursive`, `zisk-check-execute`, `sp1-check-execute`, `ooc-check`) or
+#     `ix-compile`, `ix-decompile`, `aiur-check-execute`, `aiur-check-prove`, `aiur-check-recursive`, `aiur-recursive`, `aiur-shard`, `zisk-check-execute`, `sp1-check-execute`, `ooc-check`) or
 #     `all` (the merge step expands an `all` label into every workload). Labeling
 #     requires Triage+, so PR authors from forks cannot self-queue a reset. The
 #     label shares the command/workflow name; the ref it moves is
@@ -44,7 +44,7 @@ on:
         # GitHub requires literal choice options, so this list stays static:
         # keep it (and the jobs' valid= lists below) in sync with the
         # backend testbeds in Ix/Cli/BenchCmd.lean (backendSpecs).
-        options: [ix-compile, ix-decompile, aiur-check-execute, aiur-check-prove, aiur-check-recursive, aiur-recursive, zisk-check-execute, sp1-check-execute, ooc-check, all]
+        options: [ix-compile, ix-decompile, aiur-check-execute, aiur-check-prove, aiur-check-recursive, aiur-recursive, aiur-shard, zisk-check-execute, sp1-check-execute, ooc-check, all]
       sha:
         description: "Commit to anchor to (default: HEAD)"
         required: false
@@ -77,7 +77,7 @@ jobs:
           # (backendSpecs) minus the runner-arch suffix. Static because this
           # job runs on a cheap runner with no built `ix`; keep in sync when
           # adding a backend.
-          valid="aiur-check-execute aiur-check-prove aiur-check-recursive aiur-recursive ix-compile ix-decompile ooc-check sp1-check-execute zisk-check-execute"
+          valid="aiur-check-execute aiur-check-prove aiur-check-recursive aiur-recursive aiur-shard ix-compile ix-decompile ooc-check sp1-check-execute zisk-check-execute"
           if [ "$EVENT" = workflow_dispatch ]; then
             # Reset the chosen workload(s) at the given commit; no PR scan.
             sha="${INPUT_SHA:-$HEAD_SHA}"
@@ -133,7 +133,7 @@ jobs:
           # which the merge job expands into every workload). Same static
           # list as the reset job; keep both in sync with backendSpecs in
           # Ix/Cli/BenchCmd.lean.
-          valid="aiur-check-execute aiur-check-prove aiur-check-recursive aiur-recursive ix-compile ix-decompile ooc-check sp1-check-execute zisk-check-execute"
+          valid="aiur-check-execute aiur-check-prove aiur-check-recursive aiur-recursive aiur-shard ix-compile ix-decompile ooc-check sp1-check-execute zisk-check-execute"
           accepted="$valid all"
           # Parse the workload token(s) after the command, lowercased.
           workloads=$(printf '%s' "$BODY" \

Ix/Aiur/Semantics/BytecodeFfi.lean

@@ -211,6 +211,24 @@ def shardCheckWithEnv (toplevel : @& Bytecode.Toplevel)
   (shardCheckWithEnv' toplevel funIdx envHandle ownedBlob useBytecode).map
     fun r => (r.output, .ofArrays r.ioData r.ioMap, r.queryCounts)
 
+/-- Scan-and-cut sharding: execute the env's check schedule through the
+    codegen'd kernel as incrementally-grown thin-frontier `CheckEnv`
+    claims — the same claim type the prover pays for — with a running
+    FFT-cost readout, cutting shard boundaries where the measured cost
+    reaches the RAM budget's FFT equivalent (see
+    `crates/ffi/src/aiur/scan.rs`). Witness bytes fault in lazily, so a
+    scanner's RAM tracks one segment's touched set. Writes the `.ixes`
+    manifest plus a `.costs.csv` sidecar carrying the MEASURED per-shard
+    FFT mapped through the calibrated resource lines. Numeric params are
+    decimal strings (ABI-simple): budget GiB, cut headroom percent,
+    parallel chunk scanners (0 = autoscale), fail-fast ("0" skips
+    kernel-rejected blocks, recording them in a `.failed.csv` sidecar;
+    anything else aborts on the first). -/
+@[extern "rs_aiur_scan_shards_with_env"]
+opaque scanShardsWithEnv : @& Bytecode.Toplevel →
+  @& Bytecode.FunIdx → @& EnvHandle → @& String → @& String → @& String →
+  @& String → @& String → Except String Unit
+
 end Bytecode.Toplevel
 
 end Aiur

Ix/Aiur/Stages/Codegen.lean

@@ -419,7 +419,9 @@ private def emitStore (out : Nat) (values : Array ValIdx) : Array RustStmt :=
     s!" result.output[0]" ++
     s!" } else \{" ++
     s!" let __ptr = G::from_usize(__mq.len());" ++
-    s!" __mq.insert(&__values[..], &[__ptr], G::from_bool(!unconstrained)); __ptr } }"
+    s!" __mq.insert(&__values[..], &[__ptr], G::from_bool(!unconstrained));" ++
+    s!" if record_budget_exceeded() \{ return Err(ExecError::RecordBudgetExceeded); }" ++
+    s!" __ptr } }"
   #[.letStmt false s!"__v_{out}" (some "G") (.lit blockExpr)]
 
 /-- `Op::Load`: mirror execute.rs lines 328-345. Look up by pointer
@@ -802,8 +804,14 @@ partial def emitCtrl (funIdx : FunIdx) (mcLabel? : Option String)
         #[.ref (.index (.var "inp") (.lit "..")),
           .ref (.index (.var "__ret") (.lit "..")),
           gFromBool (.lit "!unconstrained")])
+    -- The budget poll rides the memoization insert — the unit of record
+    -- growth — so a bounded execution aborts within one unique query of
+    -- its limit.
+    let budgetCheck : RustStmt :=
+      .exprStmt (.lit
+        "if record_budget_exceeded() { return Err(ExecError::RecordBudgetExceeded); }")
     -- Wrap in Ok(...) since fn now returns Result<[G; OUT_N], ExecError>.
-    return #[outArr, insertCall,
+    return #[outArr, insertCall, budgetCheck,
       .returnStmt (.call (.var "Ok") #[.var "__ret"])]
   | .match valIdx cases dflt? => do
     -- Each arm body executes from the SAME value-stack snapshot as
@@ -972,7 +980,8 @@ def optionalExecuteUses : Array (String × String) := #[
 /-- Build the `use aiur::execute::{...};` block, including only
     items whose search token appears in `body`. -/
 def emitConditionalImports (body : String) : String := Id.run do
-  let always : Array String := #["ExecError", "IOBuffer", "QueryRecord"]
+  let always : Array String :=
+    #["ExecError", "IOBuffer", "QueryRecord", "record_budget_exceeded"]
   let mut items : Array String := always
   for (path, token) in optionalExecuteUses do
     if (body.splitOn token).length > 1 then

Ix/Cli/BenchCmd.lean

@@ -38,6 +38,7 @@ module
 public import Cli
 public import Lean.Data.Json
 public import Ix.Benchmark.Results
+public import Ix.Cli.CheckCmd
 public import Ix.Cli.ConstsFile
 
 public section
@@ -131,6 +132,10 @@ structure EnvSpec where
   module : String
 
 def envSpecs : List EnvSpec := [
+  -- Init is the Aiur shard pipeline's deliverable env (the partition the
+  -- full-Init proof runs over), so its planner trend line matters on its
+  -- own — InitStd's does not subsume it.
+  { name := "Init",    module := "Benchmarks/Compile/CompileInit.lean" },
   { name := "InitStd", module := "Benchmarks/Compile/CompileInitStd.lean" },
   { name := "Lean",    module := "Benchmarks/Compile/CompileLean.lean" },
   { name := "Mathlib", module := "Benchmarks/Compile/CompileMathlib.lean" },
@@ -268,6 +273,28 @@ def backendSpecs : List BackendSpec := [
                    ("recursive-proof-size", "0.05", "_"),
                    ("prove-time", "0.10", "_"), ("proof-size", "0.05", "_"),
                    ("verify-time", "0.10", "_"), ("peak-rss", "0.10", "_")] },
+  -- The Aiur shard pipeline: one whole-env measured scan (`ix shard
+  -- scan`) at the pinned budget (`aiurShardBudgetGb`) — the FULL check
+  -- schedule executes through the codegen'd kernel as thin-frontier
+  -- CheckEnv claims with a running FFT readout, so the row tracks the
+  -- regime the per-constant aiur rows never enter (a single-constant run
+  -- faults a tiny closure). scan-time is the whole-env execution wall,
+  -- peak-rss the scan process tree's high-water; shards and the measured
+  -- max/total per-shard cost come from the manifest the scan writes.
+  -- fft-cost is the env's TOTAL measured fft — the same raw unit and
+  -- slug as the per-constant aiur rows, so the whole-env series joins
+  -- the existing Aiur FFT cost plot. The scan is deterministic — identical partitions across runs
+  -- on Init and InitStd — so the count and cost slugs ride 0/2% upper
+  -- bands: a regression can only move them up, and a legitimate kernel or
+  -- packer win lowers them freely.
+  { name := "aiur-shard", defaultMode := "execute", inputs := .perEnv,
+    testbeds := [("execute", "aiur-shard-x64-32x")],
+    metrics := [("execute", ["scan-time", "peak-rss", "shards",
+                             "max-shard-fft", "fft-cost"])],
+    thresholds := [("shards", "0", "_"), ("max-shard-fft", "0.02", "_"),
+                   ("fft-cost", "0.02", "_"),
+                   ("scan-time", "0.10", "_"),
+                   ("peak-rss", "0.10", "_")] },
   { name := "zisk", defaultMode := "execute", inputs := .perConstant,
     testbeds := [("execute", "zisk-check-execute-x64-32x")],
     metrics := [("execute", ["execute-time", "throughput", "peak-rss",
@@ -348,6 +375,14 @@ def recursiveConfigs : List (String × Array String) := [
   ("factorial-q100-b2", #["--queries", "100"])
 ]
 
+/-- `aiur-shard`'s pinned packing budget (GiB). Machine-independent on
+    purpose: a PR row and its bencher baseline must describe the same
+    partition problem, so the budget is the deliverable 256 GB prove
+    box — NOT the runner's own RAM (`--ceiling-gb` still guards the
+    runner). Executing a shard needs a small fraction of its predicted
+    PROVE RAM, so heavy-three execution fits hosts the prove would not. -/
+def aiurShardBudgetGb : Nat := 250
+
 def BackendSpec.testbedFor (b : BackendSpec) (mode : String) : Option String :=
   (b.testbeds.find? (·.1 == mode)).map (·.2)
 
@@ -793,6 +828,52 @@ def runBenchRunCmd (p : Cli.Parsed) : IO UInt32 := do
         let cfgArgs := ((recursiveConfigs.find? (·.1 == name)).map (·.2)).getD #[]
         runGuarded watchdog ceilingGb brv
           (cfgArgs ++ #["--json", out, "--json-name", name, "--texray"])
+  | "aiur-shard" =>
+    if mode != "execute" then
+      p.printError s!"error: {backend} supports only execute mode"
+      return exitUsage
+    let ixe ← ensureIxe repo info ((p.flag? "ixe").map (·.as! String))
+    let ix ← resolveBin repo "ix"
+    let manifest := s!"{repo}/{env}-shardbench.ixes"
+    let budget := (p.flag? "shard-budget").map (·.as! Nat)
+      |>.getD aiurShardBudgetGb
+    -- One whole-env measured scan: executes the entire check schedule as
+    -- thin-frontier CheckEnv claims and writes the manifest with MEASURED
+    -- per-shard fft. The tool self-reports scan-time/peak-rss through
+    -- the rows contract (env load excluded from the timed window, RSS =
+    -- process-tree high-water). A kill (≥128) records `oom`/`crash` — the
+    -- honest row for a box the scan no longer fits; any other nonzero
+    -- exit is a genuine kernel rejection and fails the run.
+    let exit ← runGuarded watchdog ceilingGb ix
+      #["shard", ixe, "--max-ram", toString budget,
+        "--out", manifest, "--json", out, "--json-name", info.name]
+    if exit ≥ 128 then
+      let status := killStatus exit
+      IO.eprintln s!"[bench] scan killed (exit {exit}); recording {status}"
+      markKilled out info.name status
+      return 0
+    if exit != 0 then
+      IO.eprintln s!"[bench] ix shard scan failed (exit {exit})"
+      return 1
+    -- Partition metrics straight from the manifest's tagged costs
+    -- (measured fft): count, heaviest shard, and the total — the
+    -- deterministic trend lines a kernel or packer change moves.
+    match Ix.Cli.CheckCmd.parseIxesShards (← IO.FS.readBinFile manifest) with
+    | .error e =>
+      IO.eprintln s!"[bench] manifest parse failed: {e}"
+      return 1
+    | .ok shardRows =>
+      if shardRows.isEmpty then
+        IO.eprintln s!"[bench] empty partition in {manifest}"
+        return 1
+      let maxFft := shardRows.foldl (fun m r => max m r.cost) 0
+      let totalFft := shardRows.foldl (fun s r => s + r.cost) 0
+      let rows ← readRows out
+      if let some envRow := (rows.getObjVal? info.name).toOption then
+        writeEntry out info.name <|
+          ((envRow.setObjVal! "shards" (Lean.toJson shardRows.size)).setObjVal!
+              "max-shard-fft" (Lean.toJson maxFft)).setObjVal!
+            "fft-cost" (Lean.toJson totalFft)
   | "zisk" | "sp1" =>
     if mode != "execute" then
       p.printError s!"error: {backend} supports only execute mode"
@@ -844,6 +925,7 @@ def runBenchRunCmd (p : Cli.Parsed) : IO UInt32 := do
   let expected := match backend with
     | "compile" => #[info.name]
     | "decompile" => #[info.name]
+    | "aiur-shard" => #[info.name]
     | "ooc" | "lean4lean" => #[info.name] ++ names
     | "aiur-recursive" => (recursiveConfigs.map (·.1)).toArray
     | _ => names
@@ -888,7 +970,7 @@ def benchRunCmd : Cli.Cmd := `[Cli|
   "Execute one benchmark run (backend × env × mode), writing benchmark results JSON. Exits 0 on success (rows saved as the local baseline), 3 when the kernel rejected any constant, 1 when no rows were produced."
 
   FLAGS:
-    backend      : String; "aiur | zisk | sp1 | ooc | lean4lean | compile | decompile | aiur-recursive"
+    backend      : String; "aiur | aiur-shard | zisk | sp1 | ooc | lean4lean | compile | decompile | aiur-recursive"
     env          : String; "Benchmark env from the registry (default: InitStd)"
     mode         : String; "prove | execute | recursive (default: the backend's defaultMode)"
     out          : String; "Benchmark results JSON output path (default: bench.json)"
@@ -900,6 +982,7 @@ def benchRunCmd : Cli.Cmd := `[Cli|
     "shard-only";          "Restrict to shard_target rows"
     ixe          : String; "Path to an existing .ixe env to use (default: compile <env> fresh; ignored by the compile backend)"
     "ceiling-gb" : Nat;    "RAM watchdog ceiling in GB (default: machine RAM minus 15 GB)"
+    "shard-budget" : Nat;  "aiur-shard only: pinned `ix shard <env.ixe> --max-ram` budget in GiB (default 250, the deliverable prove box). Change it only knowingly — the bencher baseline is keyed to the same partition problem."
     watchdog     : String; "Watchdog wrapper path (default: <repo>/.github/scripts/watchdog.sh; missing = run unguarded)"
 ]
 

Ix/Cli/BenchPlots.lean

@@ -81,6 +81,10 @@ def plotTitle (workload measure : String) : String :=
   | "aiur-recursive", "recursive-peak-rss"     => "Aiur Recursive Verifier Peak RAM Usage"
   | "aiur-recursive", "recursive-proof-size"   => "Aiur Recursive Verifier Proof Size"
   | "aiur-recursive", "recursive-fft-cost"     => "Aiur Recursive Verifier FFT Cost"
+  | "aiur-shard", "shards"              => "Aiur Predicted Shards"
+  | "aiur-shard", "union-bytes"         => "Aiur Shard Ingress Bytes"
+  | "aiur-shard", "heavy-execute-time"  => "Aiur Heavy-Shard Execute Time"
+  | "aiur-shard", "heavy-peak-rss"      => "Aiur Heavy-Shard Peak RAM Usage"
   | w, m => s!"{w}: {m}"
 
 /-- Tracked but not plotted solo. The two aiur runs re-measure each
@@ -97,13 +101,17 @@ def plotTitle (workload measure : String) : String :=
     prove-time / proof-size / verify-time / peak-rss measure the INNER
     toy-statement proof — tracked for the compare table, but the
     dashboard trend that matters is the recursion layer's own
-    `recursive-*` series, so the inner metrics aren't plotted. -/
+    `recursive-*` series, so the inner metrics aren't plotted. aiur-shard's
+    `pred-floor-ram` is a cap-overflow tripwire (thresholded), while the
+    heavy-three execution charts the real RAM — the prediction adds no
+    dashboard trend of its own. -/
 def plotSkips : List (String × String) :=
   [("aiur-check-prove", "execute-time"), ("aiur-check-execute", "fft-cost"),
    ("zisk-check-execute", "shards"), ("zisk-check-execute", "constants"),
    ("ix-decompile", "file-size"), ("ix-decompile", "constants"),
    ("aiur-recursive", "prove-time"), ("aiur-recursive", "proof-size"),
-   ("aiur-recursive", "verify-time"), ("aiur-recursive", "peak-rss")]
+   ("aiur-recursive", "verify-time"), ("aiur-recursive", "peak-rss"),
+   ("aiur-shard", "pred-floor-ram")]
 
 /-- Canonical units per measure slug, asserted on every sync: bencher
     auto-creates a measure with placeholder units ("Measure (units)") on
@@ -126,6 +134,10 @@ def unitsFor (slug : String) : Option String :=
    ("cycles", "cycles"),
    ("max-shard-cycles", "cycles"),
    ("shards", "shards"),
+   ("pred-floor-ram", "GiB"),
+   ("union-bytes", "bytes (B)"),
+   ("heavy-execute-time", "seconds (s)"),
+   ("heavy-peak-rss", "bytes (B)"),
    ("fft-cost", "FFTs"),
    ("recursive-execute-time", "seconds (s)"),
    ("recursive-prove-time", "seconds (s)"),
@@ -139,7 +151,7 @@ def unitsFor (slug : String) : Option String :=
     ooc); unranked workloads (a future backend) sort last. -/
 def workloadOrder : List String :=
   ["ix-compile", "ix-decompile", "aiur-check-prove", "aiur-check-execute",
-   "aiur-recursive", "zisk-check-execute", "ooc-check"]
+   "aiur-recursive", "aiur-shard", "zisk-check-execute", "ooc-check"]
 
 structure PlotSpec where
   testbed : String