diff --git a/Ix/Cli/ShardCmd.lean b/Ix/Cli/ShardCmd.lean index 8bdee963..8fe2f890 100644 --- a/Ix/Cli/ShardCmd.lean +++ b/Ix/Cli/ShardCmd.lean @@ -1,19 +1,25 @@ /- - `ix shard `: partition a profiled environment into shards, - minimizing cross-shard delta-unfold ingress (see `plans/sharding.md`). - - Two modes (precedence in `runShardCmd`): - - default / `--max-ram G` / `--max-cycles C`: **bin-pack to a per-shard - cycle/RAM cap** — the fewest shards that each stay under the budget, each - packed as full as the dependency structure allows (no `--max-ram` ⇒ sized to - detected system RAM). Not balanced: packing yields the minimal shard count. - - `--shards N`: force exactly `N` **balanced** min-cut shards (manual override). - - Reads the `.ixprof` produced by `ix profile` (pure offline graph work, so the - budget/`N` is cheap to re-tune without re-running the kernel). Writes a `.ixes` - manifest and prints a what-if report (per-shard cost + total cross-shard - ingress). The partitioner is self-contained — no external graph-library - dependency. + `ix shard [--profile ]`: partition an environment + into shards for the Aiur (IxVM) checking pipeline. + + Two strategies: + - **Static (default, no `--profile`)**: computed from the `.ixe` alone — + no out-of-circuit kernel run. Byte-balanced min-cut over the env's + 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%. + - **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 + cycle/RAM cap** (no `--max-ram` ⇒ sized to detected system RAM); + `--shards N` forces exactly `N` balanced min-cut shards. + + Both write the same `.ixes` manifest and print a what-if report. The + partitioner is self-contained — no external graph-library dependency. `ix shard extract --consts `: the pipeline's scoping step — extract the named constants' dependency closure from a serialized @@ -71,11 +77,46 @@ def shardExtractCmd : Cli.Cmd := `[Cli| path : String; "Path to the source `.ixe` (e.g. from `ix compile`)." ] +def runShardGraphCmd (p : Cli.Parsed) : IO UInt32 := do + let some pathArg := p.positionalArg? "path" + | p.printError "error: must specify to a .ixe file" + return 1 + let envPath := pathArg.as! String + let outPath : String := + match p.flag? "out" with + | some flag => flag.as! String + -- `init.ixe` → `init.graph`, mirroring the `.ixprof` default naming. + | none => + let base := if envPath.endsWith ".ixe" then (envPath.dropEnd 4).toString else envPath + base ++ ".graph" + rsShardStaticGraphFFI envPath outPath + IO.println s!"[graph] wrote {outPath}" + return 0 + +def shardGraphCmd : Cli.Cmd := `[Cli| + "graph" VIA runShardGraphCmd; + "Dump a `.ixe`'s static block-level reference graph (`block`/`edge` text lines) for offline partitioner prototyping" + + FLAGS: + out : String; "Output text path. Defaults to the env's base name with `.graph` (e.g. `init.ixe` → `init.graph`)." + + ARGS: + path : String; "Path to the source `.ixe`." +] + def runShardCmd (p : Cli.Parsed) : IO UInt32 := do let some pathArg := p.positionalArg? "path" - | p.printError "error: must specify to a .ixprof file" + | p.printError "error: must specify to a .ixe file" return 1 - let espPath := pathArg.as! String + let envPath := pathArg.as! String + let profileFlag : Option String := (p.flag? "profile").map (·.as! String) + -- Old CLI shape took the `.ixprof` positionally; catch it so scripts + -- fail loudly instead of parsing a profile as an env. + if profileFlag.isNone && envPath.endsWith ".ixprof" then + p.printError "error: the positional argument is now the `.ixe` env; \ + pass the profile via --profile (or drop it for the \ + static strategy)" + return 1 let balancePct : Nat := match p.flag? "balance" with | some flag => flag.as! Nat @@ -83,10 +124,10 @@ def runShardCmd (p : Cli.Parsed) : IO UInt32 := do let outPath : String := match p.flag? "out" with | some flag => flag.as! String - -- Default manifest mirrors the profile's base name: `init.ixprof` → - -- `init.ixes` (not `init.ixprof.ixes`). + -- Default manifest mirrors the env's base name: `init.ixe` → + -- `init.ixes` (not `init.ixe.ixes`). | none => - let base := if espPath.endsWith ".ixprof" then (espPath.dropEnd 7).toString else espPath + let base := if envPath.endsWith ".ixe" then (envPath.dropEnd 4).toString else envPath base ++ ".ixes" let shardsFlag : Option Nat := (p.flag? "shards").map (·.as! Nat) let maxCycles : Option Nat := (p.flag? "max-cycles").map (·.as! Nat) @@ -98,20 +139,39 @@ def runShardCmd (p : Cli.Parsed) : IO UInt32 := do | some flag => max 1 (flag.as! Nat) | none => 1 - -- Precedence: explicit --shards (fixed count) > explicit --max-cycles/--max-ram - -- (budget) > default (size to detected system RAM). - match shardsFlag with - | some n => - IO.println s!"Sharding {espPath} into {n} shards (balance ±{balancePct}%)" - rsShardEspFFI espPath (toString n) (toString balancePct) (toString parallelism) - outPath + match profileFlag with | none => - if maxCycles.isNone && maxRam.isNone then - IO.println s!"Sharding {espPath} to detected system RAM (balance ±{balancePct}%)" - else - IO.println s!"Sharding {espPath} to budget (max-cycles={maxCycles.getD 0}, max-ram={maxRam.getD 0} GiB, balance ±{balancePct}%)" - rsShardEspCapFFI espPath (toString (maxCycles.getD 0)) (toString (maxRam.getD 0)) - (toString balancePct) (toString parallelism) outPath + -- STATIC strategy (no out-of-circuit profiling): byte-balanced min-cut + -- over the env's walk-edge nets + predicted-FFT rebalance post-pass. + -- Fixed shard count only for now — the cap modes' cycle/RAM budgeting + -- is calibrated against the profiled op counters, which the static + -- profile does not carry. + let some n := shardsFlag + | p.printError "error: the static strategy (no --profile) requires \ + --shards N; --max-cycles/--max-ram budgeting needs --profile" + return 1 + if maxCycles.isSome || maxRam.isSome then + p.printError "error: --max-cycles/--max-ram require --profile (their \ + budget model is calibrated on profiled op counters)" + return 1 + IO.println s!"Sharding {envPath} into {n} shards (static strategy, balance ±{balancePct}%)" + rsShardEnvStaticFFI envPath (toString n) (toString balancePct) outPath + | some espPath => + -- Profiled strategy, unchanged: partition the `.ixprof`. + -- Precedence: explicit --shards (fixed count) > explicit --max-cycles/--max-ram + -- (budget) > default (size to detected system RAM). + match shardsFlag with + | some n => + IO.println s!"Sharding {espPath} into {n} shards (balance ±{balancePct}%)" + rsShardEspFFI espPath (toString n) (toString balancePct) (toString parallelism) + outPath + | none => + if maxCycles.isNone && maxRam.isNone then + IO.println s!"Sharding {espPath} to detected system RAM (balance ±{balancePct}%)" + else + IO.println s!"Sharding {espPath} to budget (max-cycles={maxCycles.getD 0}, max-ram={maxRam.getD 0} GiB, balance ±{balancePct}%)" + rsShardEspCapFFI espPath (toString (maxCycles.getD 0)) (toString (maxRam.getD 0)) + (toString balancePct) (toString parallelism) outPath if !outPath.isEmpty then IO.println s!"[shard] wrote {outPath}" return 0 @@ -121,21 +181,23 @@ end Ix.Cli.ShardCmd open Ix.Cli.ShardCmd in def shardCmd : Cli.Cmd := `[Cli| "shard" VIA runShardCmd; - "Partition a `.ixprof` into shards: pack to a RAM/cycle cap (default) or N balanced shards" + "Partition a `.ixe` env into shards: static strategy by default (`--shards N`, no kernel run), or the profiled pipeline via `--profile `" FLAGS: - shards : Nat; "Fixed number of shards N (overrides the default budget sizing)" - "max-cycles" : Nat; "Per-shard guest-cycle budget (overrides the default RAM sizing)" - "max-ram" : Nat; "Per-shard host-RAM budget, GiB (default: detected system RAM)" + 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 (required for the static strategy; overrides the profiled default budget sizing)" + "max-cycles" : Nat; "Per-shard guest-cycle budget (profiled strategy only)" + "max-ram" : Nat; "Per-shard host-RAM budget, GiB (profiled strategy only; default: detected system RAM)" balance : Nat; "Per-bisection balance tolerance, percent (default 5)" - parallelism : Nat; "Provers assumed for the prove-time estimate (default 1 = sequential)" - out : String; "Output .ixes manifest path (default: .ixes, e.g. init.ixprof → init.ixes)" + 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)" ARGS: - path : String; "Path to a .ixprof produced by `ix profile`" + path : String; "Path to a serialized `.ixe` environment" SUBCOMMANDS: - shardExtractCmd + shardExtractCmd; + shardGraphCmd ] end diff --git a/Ix/KernelCheck.lean b/Ix/KernelCheck.lean index 5449c0a5..d403997a 100644 --- a/Ix/KernelCheck.lean +++ b/Ix/KernelCheck.lean @@ -194,6 +194,31 @@ 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. -/ +@[extern "rs_shard_env_static"] +opaque rsShardEnvStaticFFI : + @& String → -- .ixe path + @& String → -- num_shards (N) + @& String → -- balance percent + @& String → -- .ixes output path ("" = skip) + IO Unit + +/-- FFI: dump the static block-level reference graph of a `.ixe` as text: + `block ` per ingress unit and + `edge ` per deduped claim-walk edge at block + granularity. The walk-edge relation is what generates a shard's thin + frontier, so a partition's ingress cost is computable offline from + this file alone (partitioner-prototype input; no kernel run). -/ +@[extern "rs_shard_static_graph"] +opaque rsShardStaticGraphFFI : + @& String → -- .ixe path + @& String → -- output text path + IO Unit + /-- FFI: profile a `.ixe` out of circuit, writing a `.ixprof` sidecar with per-block heartbeats + the delta-unfold graph (the sharding cost model, see `plans/sharding.md`). Runs the anon kernel over every checkable target. diff --git a/crates/ffi/src/kernel.rs b/crates/ffi/src/kernel.rs index f548a1ba..60fd0f45 100644 --- a/crates/ffi/src/kernel.rs +++ b/crates/ffi/src/kernel.rs @@ -2235,6 +2235,189 @@ fn profile_block_size(env: &IxonEnv, block: &Address) -> u32 { .map_or(0, |b| b.len().min(u32::MAX as usize) as u32) } +/// Build the STATIC block profile of an env — the profile-free analog of +/// [`build_block_profile`]: same [`BlockProfile`] shape, but derived from +/// the serialized constants alone, no kernel run. +/// +/// - Vertices: ingress units (a Muts block or a standalone constant), +/// `serialized_size` real, `const_count` = member constants, +/// `heartbeats` := size (report-only). +/// - Balance weight: the partitioner's vertex weight is `block_step_cost` +/// (a linear model over the op counters), so the byte weight is carried +/// in `subst` — `vweight ∝ size` exactly, matching the measured-best +/// bisection weight (`ix_kernel::shard::STATIC_OWNED_PER_BYTE` docs). +/// - Edges: the claim-walk relation (`ixon::shard_claim::walk_edges` — +/// positively-used refs + projection block pointers) at block +/// granularity, deduped, self-edges dropped. These are exactly the +/// edges that generate a shard's thin frontier, i.e. its ingress. +#[allow(clippy::cast_possible_truncation)] // block sizes clamped to u32::MAX +fn static_block_profile(env: &IxonEnv) -> BlockProfile { + use rayon::prelude::*; + let addrs: Vec
= + env.consts.iter().map(|e| e.key().clone()).collect(); + // Phase 1 (parallel): parse each constant exactly ONCE. Lazy mmap + // entries re-parse on every `get()` (see `LazyConstant`'s cache + // policy), so the home block is derived from the constant we already + // hold instead of `profile_block_of` (which would parse again), and + // edge targets are kept as raw addresses to be resolved in phase 2 + // through the addr→home table this pass produces. + struct ConstRow { + addr: Address, + home: Address, + targets: Vec
, + } + let rows: Vec = addrs + .par_iter() + .filter_map(|addr| { + let c = env.get_const(addr)?; + let home = match &c.info { + IxonCI::IPrj(p) => p.block.clone(), + IxonCI::CPrj(p) => p.block.clone(), + IxonCI::RPrj(p) => p.block.clone(), + IxonCI::DPrj(p) => p.block.clone(), + _ => addr.clone(), + }; + let mut targets = Vec::new(); + ixon::shard_claim::walk_edges(&c, &mut targets); + Some(ConstRow { addr: addr.clone(), home, targets }) + }) + .collect(); + // Phase 2 (serial, no parsing): addr→home table, per-block member + // counts, and block-level edges. A target absent from the table is a + // blob payload (or an unparseable constant) — skipped, matching the + // claim walk's blob sentinel. + let home_of: FxHashMap<&Address, &Address> = + rows.iter().map(|r| (&r.addr, &r.home)).collect(); + let mut counts: FxHashMap<&Address, u32> = FxHashMap::default(); + for r in &rows { + *counts.entry(&r.home).or_insert(0) += 1; + } + let mut builder = ProfileBuilder::new(); + for r in &rows { + for t in &r.targets { + if let Some(&pb) = home_of.get(t) + && *pb != r.home + { + builder.delta_edge(r.home.clone(), pb.clone()); + } + } + } + for (home, cnt) in counts { + // Size without materializing: `raw_bytes().len()` on the lazy entry + // (`profile_block_size` would copy the bytes into a fresh Arc). + let size = env + .consts + .get(home) + .map_or(0, |e| e.value().raw_bytes().len().min(u32::MAX as usize)) + as u32; + let ops = OpCounts { subst_nodes: u64::from(size), ..OpCounts::default() }; + builder.block(home.clone(), u64::from(size), size, cnt, ops); + } + 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 +/// (`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>, + 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 balance = + (balance_pct.to_string().parse::().unwrap_or(5) as f64) / 100.0; + let out = out_path.to_string(); + let out_opt = if out.is_empty() { None } else { Some(out.as_str()) }; + let t0 = Instant::now(); + let env = match IxonEnv::get_anon_mmap(std::path::Path::new(&path)) { + Ok(e) => e, + Err(e) => { + return LeanIOResult::error_string(&format!( + "rs_shard_env_static: failed to load {path}: {e}" + )); + }, + }; + let t_load = t0.elapsed(); + let t1 = Instant::now(); + let profile = static_block_profile(&env); + eprintln!( + "[rs_shard_static] env load {:.1?}, static graph {:.1?} ({} blocks, {} edges)", + t_load, + t1.elapsed(), + profile.num_blocks(), + profile.num_edges() + ); + match ix_kernel::shard::shard_static(&profile, num_shards, balance, out_opt) { + Ok(report) => { + eprintln!("[rs_shard_static]\n{report}"); + LeanIOResult::ok(LeanOwned::box_usize(0)) + }, + Err(e) => LeanIOResult::error_string(&format!("rs_shard_env_static: {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 +/// claim-walk edge (`ixon::shard_claim::walk_edges`, mapped to block +/// granularity, self-edges dropped). The graph is the input for offline +/// partitioner prototypes in the sharding campaign: the walk-edge +/// relation is exactly what generates a shard's thin frontier, so a +/// partition's ingress cost is computable from this file alone. Same +/// data as [`static_block_profile`], serialized for offline use. +#[allow(clippy::cast_possible_truncation)] // block ids are u32 by construction +#[unsafe(no_mangle)] +pub extern "C" fn rs_shard_static_graph( + env_path: LeanString>, + out_path: LeanString>, +) -> LeanIOResult { + let path = env_path.to_string(); + let out = out_path.to_string(); + let env = match IxonEnv::get_anon_mmap(std::path::Path::new(&path)) { + Ok(e) => e, + Err(e) => { + return LeanIOResult::error_string(&format!( + "rs_shard_static_graph: failed to load {path}: {e}" + )); + }, + }; + let profile = static_block_profile(&env); + let mut buf = String::new(); + for b in profile.blocks() { + buf.push_str(&format!( + "block {} {} {}\n", + b.addr.hex(), + b.serialized_size, + b.const_count + )); + } + for c in 0..profile.num_blocks() as u32 { + let ca = profile.block(c).addr.hex(); + for &p in profile.producers(c) { + buf.push_str(&format!("edge {} {}\n", ca, profile.block(p).addr.hex())); + } + } + if let Err(e) = std::fs::write(&out, buf) { + return LeanIOResult::error_string(&format!( + "rs_shard_static_graph: failed to write {out}: {e}" + )); + } + eprintln!( + "[rs_shard_static_graph] {} blocks, {} edges → {}", + profile.num_blocks(), + profile.num_edges(), + out + ); + LeanIOResult::ok(LeanOwned::box_usize(0)) +} + /// Print the general-purpose cost breakdown for `ix profile` — the kernel-work /// metrics plus the predicted Zisk leaf cost/RAM (à la `cargo-zisk … -p summary`). // `steps as f64` is a display-only cast for `{:.2e}` formatting; precision loss diff --git a/crates/kernel/src/shard.rs b/crates/kernel/src/shard.rs index abdf4cc9..c06d1ae9 100644 --- a/crates/kernel/src/shard.rs +++ b/crates/kernel/src/shard.rs @@ -1541,6 +1541,291 @@ impl<'a> Cur<'a> { } } +// ============================================================================ +// Static (profile-free) sharding — the `ix shard ` path +// ============================================================================ + +/// Fitted per-shard Aiur FFT cost model driving the static strategy: +/// +/// ```text +/// cost(S) ≈ STATIC_OWNED_PER_BYTE · owned_bytes(S) +/// + STATIC_OWNED_SUPERLINEAR · Σ_{b∈S} size(b)^1.5 +/// + STATIC_FRONTIER_PER_BYTE · frontier_bytes(S) +/// ``` +/// +/// where `size(b)` is a block's serialized byte length and +/// `frontier_bytes(S)` sums the sizes of foreign producer blocks referenced +/// by S's members (the thin frontier the kernel ingresses). A frontier byte +/// pays ingress only (blake3 + parse); an owned byte additionally pays +/// typechecking, which concentrates superlinearly in large proof bodies — +/// hence the per-block `^1.5` term (NOT `(Σ size)^1.5`: how bytes are +/// packaged into constants is exactly what it measures). +/// +/// Fit: least squares over 56 measured shards from 7 different 8-shard +/// partitions of Init (native IxVM, real blake3, `Total FFT cost` from +/// `ix check --stats-out`), mean |err| 4.7% on shards ≥ 3e11. Validated +/// without refit on Std at 24 shards (4.5% aggregate bias) — the constants +/// are Aiur-kernel physics, not env-specific. The fit's intercept +/// (−2.3e10) cancels in balance comparisons and is omitted. +pub const STATIC_OWNED_PER_BYTE: f64 = 28_201.0; +/// See [`STATIC_OWNED_PER_BYTE`]. +pub const STATIC_OWNED_SUPERLINEAR: f64 = 681.08; +/// See [`STATIC_OWNED_PER_BYTE`]. +pub const STATIC_FRONTIER_PER_BYTE: f64 = 38_000.0; + +/// 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() +} + +/// 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 +/// or no single move lowers the pair's max. Returns the move count. +/// +/// Why a post-pass at all: recursive bisection enforces balance only per +/// cut (±ε), which compounds to ~1.6–1.9× max/min over log₂(N) levels of +/// proportional integer budget splits; and no vertex-weight balance can +/// see the frontier term, which depends on the cut itself. Measured on the +/// 8-shard Init harness this pass took realized FFT stddev 14.7% → 7.1% +/// and the max shard down 6% on top of the byte-balanced min-cut. +/// +/// Move evaluation is exact under the model: moving `b` from `src` to +/// `dst` shifts its owned weight and updates both shards' frontiers — +/// producers of `b` may enter `dst`'s frontier or leave `src`'s (when `b` +/// was their only consumer there), and `b` itself may enter `src`'s +/// frontier (if consumers remain) or leave `dst`'s. Other shards are +/// unaffected: their view of `b` depends only on `b` being outside them. +pub fn rebalance_static( + profile: &BlockProfile, + shard_of: &mut [u32], + num_shards: usize, +) -> usize { + use rustc_hash::FxHashMap; + let n = profile.num_blocks(); + if num_shards <= 1 || n == 0 { + return 0; + } + let size = |b: u32| profile.block(b).serialized_size; + let (crow, ccol) = profile.consumers_csr(); + let consumers = |b: u32| &ccol[crow[b as usize]..crow[b as usize + 1]]; + + let mut owned = vec![0.0f64; num_shards]; + for b in 0..n { + owned[shard_of[b] as usize] += static_owned_weight(size(b as u32)); + } + // fcnt[k][p] = number of k-owned consumers of foreign producer p; + // fbytes[k] = Σ size(p) over k's frontier (fcnt keys). + let mut fcnt: Vec> = + vec![FxHashMap::default(); num_shards]; + let mut fbytes = vec![0.0f64; num_shards]; + for c in 0..n as u32 { + let kc = shard_of[c as usize]; + for &p in profile.producers(c) { + if shard_of[p as usize] != kc { + let e = fcnt[kc as usize].entry(p).or_insert(0); + if *e == 0 { + fbytes[kc as usize] += f64::from(size(p)); + } + *e += 1; + } + } + } + + let mut moves = 0usize; + loop { + let costs: Vec = (0..num_shards) + .map(|k| owned[k] + STATIC_FRONTIER_PER_BYTE * fbytes[k]) + .collect(); + let hot = + (0..num_shards).max_by(|&a, &b| costs[a].total_cmp(&costs[b])).unwrap(); + let cold = + (0..num_shards).min_by(|&a, &b| costs[a].total_cmp(&costs[b])).unwrap(); + let gap = costs[hot] - costs[cold]; + let mean = costs.iter().sum::() / num_shards as f64; + if gap <= 0.01 * mean { + break; + } + // Best move: lexicographically minimize (new max(hot', cold'), total + // cost increase). Deterministic: blocks scanned in id order, strict <. + let mut best: Option<(u32, f64, f64)> = None; + for b in 0..n as u32 { + if shard_of[b as usize] != hot as u32 { + continue; + } + let w = static_owned_weight(size(b)); + if w > gap { + continue; // moving it would overshoot the gap + } + // Δcost(src): lose w(b); b's producers leave the frontier when b was + // their only src consumer; b enters it if src still consumes it. + let mut df_src = 0.0f64; + for &p in profile.producers(b) { + if shard_of[p as usize] != hot as u32 && fcnt[hot].get(&p) == Some(&1) { + df_src -= f64::from(size(p)); + } + } + if consumers(b) + .iter() + .any(|&c| c != b && shard_of[c as usize] == hot as u32) + { + df_src += f64::from(size(b)); + } + // Δcost(dst): gain w(b); b leaves dst's frontier if present; b's + // foreign producers enter it if not already there. + let mut df_dst = 0.0f64; + if fcnt[cold].contains_key(&b) { + df_dst -= f64::from(size(b)); + } + for &p in profile.producers(b) { + if shard_of[p as usize] != cold as u32 && !fcnt[cold].contains_key(&p) { + df_dst += f64::from(size(p)); + } + } + let ds = -w + STATIC_FRONTIER_PER_BYTE * df_src; + let dd = w + STATIC_FRONTIER_PER_BYTE * df_dst; + let key = ((costs[hot] + ds).max(costs[cold] + dd), ds + dd); + if best + .is_none_or(|(_, k0, k1)| key.0 < k0 || (key.0 == k0 && key.1 < k1)) + { + best = Some((b, key.0, key.1)); + } + } + let Some((b, new_max, _)) = best else { break }; + if new_max >= costs[hot] { + break; // no move improves the pair + } + // Apply: owned weights, then frontier bookkeeping for src and dst. + let (src, dst) = (hot, cold); + shard_of[b as usize] = dst as u32; + let wb = static_owned_weight(size(b)); + owned[src] -= wb; + owned[dst] += wb; + for &p in profile.producers(b) { + if shard_of[p as usize] != src as u32 + && let Some(e) = fcnt[src].get_mut(&p) + { + *e -= 1; + if *e == 0 { + fcnt[src].remove(&p); + fbytes[src] -= f64::from(size(p)); + } + } + if shard_of[p as usize] != dst as u32 { + let e = fcnt[dst].entry(p).or_insert(0); + if *e == 0 { + fbytes[dst] += f64::from(size(p)); + } + *e += 1; + } + } + let src_consumers = consumers(b) + .iter() + .filter(|&&c| c != b && shard_of[c as usize] == src as u32) + .count() as u32; + if src_consumers > 0 { + let e = fcnt[src].entry(b).or_insert(0); + if *e == 0 { + fbytes[src] += f64::from(size(b)); + } + *e = src_consumers; + } + if fcnt[dst].remove(&b).is_some() { + fbytes[dst] -= f64::from(size(b)); + } + moves += 1; + } + moves +} + +/// Predicted per-shard costs under the static model for a given assignment +/// (fresh recomputation — used for reports and tests). +pub fn static_predicted_costs( + profile: &BlockProfile, + shard_of: &[u32], + num_shards: usize, +) -> Vec { + use rustc_hash::FxHashMap; + let mut owned = vec![0.0f64; num_shards]; + let mut fr: Vec> = vec![FxHashMap::default(); num_shards]; + for b in 0..profile.num_blocks() as u32 { + let k = shard_of[b as usize] as usize; + owned[k] += static_owned_weight(profile.block(b).serialized_size); + for &p in profile.producers(b) { + if shard_of[p as usize] != k as u32 { + fr[k].insert(p, ()); + } + } + } + (0..num_shards) + .map(|k| { + let fb: f64 = fr[k] + .keys() + .map(|&p| f64::from(profile.block(p).serialized_size)) + .sum(); + owned[k] + STATIC_FRONTIER_PER_BYTE * fb + }) + .collect() +} + +/// Partition a STATIC block profile (reference graph + sizes, no kernel +/// run — see `ix shard graph` / the FFI static-profile builder) into +/// `num_shards` shards: byte-balanced min-cut over the walk-edge nets, +/// then the [`rebalance_static`] post-pass toward equal predicted cost. +/// Writes the same `.ixes` manifest as [`shard_esp`]. This is the +/// `ix shard ` (no `--profile`) strategy; measured against the +/// profiled strategy it cut mean shard FFT ~30% and the max shard +/// 44–51% on the Init(8)/Std(24) harnesses. +pub fn shard_static( + profile: &BlockProfile, + num_shards: usize, + balance: f64, + out_path: Option<&str>, +) -> Result { + let t0 = Instant::now(); + let h = Hypergraph::from_profile(profile); + let t_graph = t0.elapsed(); + let t1 = Instant::now(); + let (mut shard_of, tree) = h.partition_with_tree(num_shards, balance); + let t_part = t1.elapsed(); + let t2 = Instant::now(); + let moves = rebalance_static(profile, &mut shard_of, num_shards); + eprintln!( + "[shard_static] hypergraph {t_graph:.1?}, partition {t_part:.1?}, rebalance {:.1?} ({moves} moves)", + t2.elapsed() + ); + 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}"))?; + } + 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 { + lo = lo.min(c); + hi = hi.max(c); + sum += c; + } + Ok(format!( + "blocks={} static_edges={} nets={}\n{}\nrebalance moves={} predicted FFT/shard mean={:.3e} min={:.3e} max={:.3e} spread={:.2}x", + profile.num_blocks(), + profile.num_edges(), + h.num_nets(), + manifest.summary(), + moves, + sum / num_shards as f64, + lo, + hi, + hi / lo.max(1.0), + )) +} + /// Read a `.ixprof`, partition into `num_shards` shards, and emit a manifest with /// per-shard cost metrics, foreign-block sets, and (delta-based) assumption /// roots. Optionally writes the manifest (`.ixes`). Returns a what-if report. @@ -2138,6 +2423,52 @@ mod tests { assert_eq!(h.connectivity_objective(&shard_of), 1000); } + /// `rebalance_static` must (a) keep its incremental owned/frontier + /// bookkeeping consistent with a from-scratch recomputation (the greedy + /// loop trusts those numbers for every move decision) and (b) actually + /// close the hot/cold gap on a skewed assignment. + #[test] + fn rebalance_static_balances_and_bookkeeping_is_exact() { + // 12 blocks, sizes 1000..12000, a producer chain i -> i+1, all + // initially dumped into shard 0 of 3. + let mut b = ProfileBuilder::new(); + for i in 1..=12u8 { + b.block(addr(i), 0, 1000 * u32::from(i), 1, OpCounts::default()); + } + for i in 1..12u8 { + b.delta_edge(addr(i), addr(i + 1)); + } + let p = b.finish(); + let mut shard_of = vec![0u32; 12]; + let moves = rebalance_static(&p, &mut shard_of, 3); + assert!(moves > 0, "skewed assignment must produce moves"); + let costs = static_predicted_costs(&p, &shard_of, 3); + let mean = costs.iter().sum::() / 3.0; + let gap = costs.iter().fold(0.0f64, |m, &c| m.max(c)) + - costs.iter().fold(f64::INFINITY, |m, &c| m.min(c)); + // The loop exits either within 1% of the mean or when no single move + // helps; on this fixture the largest block (12000 bytes) bounds the + // achievable gap — assert we got under that, not stuck near the start. + assert!( + gap <= static_owned_weight(12000) + STATIC_FRONTIER_PER_BYTE * 12000.0, + "gap {gap} vs mean {mean}: rebalance made no progress" + ); + // Every shard non-empty (nothing collapsed to zero). + for k in 0..3u32 { + assert!(shard_of.contains(&k), "shard {k} emptied"); + } + // Bookkeeping exactness via idempotence: a second pass rebuilds its + // owned/frontier state from scratch, so if the first pass's incremental + // accounting was truthful, its stopping condition also holds on fresh + // state and the second pass has nothing to do. Drifted bookkeeping + // would leave the second pass a real gap to close. + let again = rebalance_static(&p, &mut shard_of, 3); + assert_eq!( + again, 0, + "second pass moved {again} blocks — bookkeeping drift" + ); + } + #[test] fn objective_matches_manifest_cross_ingress() { let p = two_clusters();