perf: Tune ix shard counts and recalibrate for proving via --ram-budget - #601
Draft
samuelburnham wants to merge 7 commits into
Draft
perf: Tune ix shard counts and recalibrate for proving via --ram-budget#601samuelburnham wants to merge 7 commits into
ix shard counts and recalibrate for proving via --ram-budget#601samuelburnham wants to merge 7 commits into
Conversation
…filing Adds an execution-time virtual-gas meter to the Aiur runtime: every constrained query touch is priced at its circuit width, and a memo hit replays the callee's recorded span instead of just the touch, so each function-query entry records its standalone cost — the cost as if nothing were memoized, independent of execution order. - querymap: per-map weight/profile fields, a vspans u64 arena, and replay_at/finish methods that maintain the meter for both engines (the interpreter and the codegen'd kernel share them). - execute + codegen: Call/Store/Load/Return sites routed through the new methods; generated fns snapshot record.virt at frame entry (aiur_ixvm.rs / aiur_multi_stark.rs regenerated). - ffi: every ExecuteResult now carries an opaque QueryRecordHandle owning the execution's QueryRecord, inspected from Lean via the rs_aiur_qr_* accessors (fnLen, fnKey, fnMult, fnVspan, memKey). A profile : Bool is plumbed end-to-end (Lean Full wrappers -> externs -> executors -> QueryRecord::new); prove paths never profile. - ix check --profile <file>: appends one "<addr_hex> <vspan> <mult> [name]" line per check_const query, the address recovered by dereferencing the Addr argument (last input slot) through memory[32]; names joined via ixonEnv.addrToName on the shard paths. Intended as order-independent per-constant weights for smarter re-sharding. Meters are side-band (multiplicities and outputs untouched): the kernel FFT cost is unchanged (lake exe ix check Nat.add_comm pin exact), native and bytecode-interpreter profile dumps are byte-identical, and the aiur-cross suite passes.
Per-constant weights come out of the rayon shard batch without round-tripping records to Lean: `const_weights` reduces each record in the task that owns it (48-byte `(addr, vspan, mult)` rows, one per checked constant), next to the peak reduction and before `gate.release` — admission keeps bounding peak RSS by the shards in flight rather than the whole partition. `prove_from_execution` is hoisted out of the two prove paths, taking the record BY VALUE: it is the witness phase's dominant residency and dies the instant the traces exist, before the LDE/commit/FRI phases that set the prover's peak. `prove_ixvm_within_budget` then measures `peak_prove_bytes` on the REAL record in the gap between execution and the witness phase: over budget, the record is dropped and the peak comes back instead of an OOM part-way into an FFT. Over budget is a RESULT, not an error — `proveBlocksWithinBudget` halves the block list and recurses (a shard is a plan, never part of a statement: halving yields claims exactly as valid as the parent's, each half's grown frontier discharged by its sibling), bottoming out at a single block, the atom the kernel checks together. The budget defaults to 85% of `MemAvailable`, the check batch's gate policy; `ix prove --max-ram G` takes the same units `ix shard` sized the partition with, and the run returns the partition it actually produced. Rejection costs one execution (0.65s) against a full prove (4.7s) on the 360-constant fixture; a 3.45 GiB shard against a 2 GiB budget recursed to four parts of 1.51-1.81 GiB in 11s, the parts partition the parent exactly (82+86+78+78 consts), and all proofs verify against unchanged claim digests.
An over-budget shard is now cut into the part count the peak model projects will fit, in one step, instead of bisecting blindly: suggested_split_parts re-evaluates peak_prove_bytes_by with each circuit's raw counts scaled 1/n (n = 2, 4, ...) until it fits — the model's power-of-two padding, fixed gadget floor and max-of-phases handled by construction. The estimate is optimistic (parts re-execute dependencies shared across the cut), so every part is still gated on its own executed record; optimism is the right bias, since an under-split costs one cheap re-execution while an over-split pays the per-proof floor forever. Measured on the 241-shard FLT/ram400 partition: 6 shards over budget (437-648 GiB vs 400), every one resolved by a single suggested halving — max recursion depth 1, final partition maxing at 99.7% of budget. Equal-cumulative-vspan cut placement was implemented, benchmarked against equal block counts on a 360-const fixture, init and lean, lost on parts/depth/executions on every env (vspan tracks time, not rows), and is not kept. ix prove --exec-only runs the same split loop on executions alone (execute_peak_ixvm: execute, measure, suggest — no STARK), and ix check --ixes --max-ram G is its batch twin: one wave loop where wave 0 is the plain batch check, over-budget shards are cut and their parts re-batched under the rayon scheduler + RAM gate, and the wave count is the split depth. --profile records vspans from wave 0 only — they are partition-invariant, so the first full pass already carries every constant's row. Both drivers now finish by enumerating exactly which shards failed, and can emit the partition the run ACTUALLY produced as a .ixes via rs_shard_manifest_from_partition (--out-ixes): coverage validated, own/foreign/cross-ingress and assumption roots rebuilt from the static profile, and each shard's measured prover peak recorded in a trailing manifest section (older readers ignore it; 0 = unmeasured planner output) — the scheduling signal, since STARK wall is ~linear in peak (92-112s at 229-278 GiB vs 217s at 472 GiB). Emitted manifests are fixed points: re-running on one produces zero splits and re-emits it byte-identically. Also: ix prove --texray streams the per-phase timing/RSS lines (execute / witness / stark stages) — Mathlib shards measure ~30s execute vs ~100s STARK on this engine, retiring the arena-port urgency; the prove-path --max-ram docs now match the FFI (0 = detect 85% of MemAvailable, not unchecked); the checkAddrsWithEnv docstring now matches the batch-index return.
Quality/dedup across the splitting and manifest work: GatedProve (Proved/Split/Measured) replaces the Result whose Err arm meant success under exec_only; one raw_of record->rows mapping serves the peak model and the split suggester; detected_ram_budget() carries the 85% policy once, on the kernel's RAM_USABLE_FRAC; decode_addr_lists uses Address::unpack and the blob encoders use toLEBytes; pure helpers are structural folds. Ownership is assigned in ONE env pass per run (ownedConstsPer) with split parts inheriting their parent's constants (partitionOwned) — removing a full env rescan per shard and per split node (~1 h of setup at Mathlib scale). The manifest emit takes the caller's live EnvHandle; the measured-peaks section is presence-gated, strictly validated, and read back by summary(). The uncorrelated Lean RAM projection (projectedProverBytes) is deleted: predicted 135/258/462 GiB where measured RSS was 278/229/265, while the Rust three-phase model tracked +1.2-4.5%. The check audit's budget flag is --ram-budget: it names the destination prove box the partition is destined for, which the audit machine cannot detect — unlike ix prove's --max-ram, which defends the local box and detects 85% of MemAvailable when omitted. gibBytes/toGib live in Ix/Common. Seed sizing settles on 20,000x env bytes against 2/3 of the budget, Mathlib-class-normative: amplification is env-dependent (init 12.5-14k, FLT 19.6k, Mathlib 20.2-21.7k) and Mathlib's shape is the common case for real Lean libraries — a 17k midpoint seed put 129/132 Mathlib shards over a 400 GiB budget. Measured at 233 seeds: 8 over, depth-2 cascade, 242 measured shards at 161-397 GiB, zero failing constants, 17:20 wall. Small envs over-shard instead (init: ~20% extra STARK, first run only, reclaimed by the printed consolidation count). EXEC_RSS_PER_OWNED_BYTE returns to 2500 — the exec footprint is also env-dependent (~1000x ISLB, ~2300x Mathlib) and the gate's contract is never-OOM, so the slope carries the worst measured env (the 1100x refit OOM'd full-width Mathlib at 486/495 GB). Explorations implemented, measured, and removed en route (results are the keepable part): re-seed/calibration execution phases (re-execute the env to learn what the prove's own gated execution measures free); streaming cut-and-prove fusion (hint-row promotion is a full constrained replay under accumulation-based multiplicities); static byte->peak regressions (12-60% error, non-transferable); raw-sum group planning (bound held at 0.92-1.00x truth but fill equaled naive byte-sharding — summed fine segments count the shared core K times, and the claim spine is not the union of its parts'). Standing conclusion: planning signals must see row unions, only execution sees unions, so the prove run's gated execution is the planner and the seed only has to land in the right decade.
The backend passes the same 100 GiB runner budget to both stages: ix shard seeds for it and ix check --ram-budget cuts any seed the spread pushes over it in place, so the row's shards metric counts the leaf partition the run actually validated. The measured check-time window stays the wave-0 batch call, keeping the timing metric comparable; a shards increase on first run reflects real splits at this budget and takes a one-time re-pin.
samuelburnham
marked this pull request as draft
August 28, 2026 19:12
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Aiur cost profiling and budget-aware sharding
Summary
This branch adds execution-derived per-constant Aiur cost profiles and uses
measured prover-RAM estimates to make sharding self-correcting. Over-budget
shards are split before proving, audits can emit the partition they actually
validated (including measured peaks), and the
aiur-sharded-envbenchmark nowaudits the same 100-GiB budget used to seed its manifests.
The CLI distinguishes the local prove limit (
ix prove --max-ram, defaultingto 85% of available RAM) from the target machine used by an execution-only
audit (
ix check --ram-budget). Staticix shard --max-ram Gsizing now usesthe environment's block-size distribution instead of raw
.ixefile size.Explicit
--shards Nbehavior is unchanged.Updated shard model
For each atomic block
b, lets_bbe its serialized owned bytes (not.ixefile bytes):
Gis the requested per-shard RAM budget in GiB. The400is only thecalibration reference, so its factor is exactly one at 400 GiB. The per-block
s_b^(3/2)term captures block shape, the1.10exponent scales acrossenvironment sizes, and the
1.20budget exponent is fitted from the 100-GiBruns below. This remains a seed: measured execution splits any over-budget tail
before proving.
400 GiB seed counts
At the reference budget, the budget exponent has no effect; this table shows
the size-axis correction relative to the previous raw-file-size model.
The eight-shard Init partition measured 266–323 GiB with no split waves, and
Batteries needed one heavy-tail correction (14 → 15). Mathlib's unchanged
233-shard seed produced 242 measured leaves: eight initial shards were over
budget, all final shards landed at 161–397 GiB, and the run completed in 17:20
with zero failing constants. An older 241-shard FLT audit found six shards at
437–648 GiB; each fit after one model-selected halving, with the final maximum
at 99.7% of budget. Mathlib is the normalization point; FLT's new 246 count is
an extrapolation.
100 GiB calibration
* Init used the adjacent 42-shard calibration probe; it was clean, and the
final implementation was verified to emit 43. ISLB and Batteries' remaining
splits were localized heavy tails rather than broad under-seeding. For
comparison, the linear-scaling probes (Init at 32 and ISLB at 93) caused 16 and
28 initial misses respectively; a tested
1.36exponent instead over-seededISLB at 152 shards. The model selects 1,230 shards for Mathlib at 100 GiB; that
extrapolation was not run.
Validation
Focused Rust model tests, the full
ixbuild, and isolated 100-GiB Init, ISLB,and Batteries calibration runs pass.