Expand description
The optimize_graph lock-step pass driver — Phase B PR-B3 of the
“plan IS the graph” rebuild
(../../docs/session-prompts/plan-is-graph-rebuild.md capability 5;
../../docs/architecture/04-optimization.md §“The two-stage
transformation”, §“Relationship to PR 3”).
§What this is
optimize_graph was a hardcoded sequence — execution_plan order →
compile_plan (per-node placement +
cost ranking + the PR-B2 per-ending-device Pareto frontier) →
seed_placement_fork_branches (the PR-A4 deliberate-fork
pathfinder). PR-B3 restructures that sequence into a lock-step
driver over three registered pass kinds, mirroring the shape of
fuel_graph::opt::RuleRegistry (a Vec<Box<dyn …>> of passes +
builder + a driver method that walks them) so future passes (fusion,
algebraic, dtype-lowering) plug in without rewriting the loop:
Pathfinders — ADD candidate paths / branches to the graph. The PR-A4PlacementForkPathfinder(deliberate-fork seed, the one pathfinder that records anOp::Branchper genuine multi-placement fork) is the first registered pathfinder.- Rankers — MEASURE each path on the
CostVector(PR-B1/B2rank_by_cost). Rankers are not a driver-visible trait in B3: ranking is applied per kernel-bearing node insidecrate::plan::compile_planas the per-nodeAlternativeSetis built, and the ranked result is carried into the driver viaOptimizationContext::plan. The driver re-runs no ranker; the measure step is the plan it is handed. (B3 is a structural refactor — when a later PR adds a pathfinder that emits new candidates the plan never priced, that is the point to lift ranking into a driver-visibleRankertrait.) Optimizers — MERGE / DISCARD paths. The PR-B2 per-ending-device Pareto frontier + crowding cap is applied per-node insidecrate::plan::compile_plan; the registeredFrontierConvergenceOptimizeris the in-graph counterpart that runs after each pathfinder to (a) collapse duplicate (forward-identical) arms a pathfinder may have proposed [duplicate-path convergence], and (b) assert the optimizer invariants the constitution requires — never strand the last(device, backend)path, never mutate a node in an active cycle / currently-executing region.
§Lock-step / prune-as-you-go (the constitution’s bound)
The driver runs the passes interleaved, not explode-then-extract:
for each registered pathfinder, it ADDs its candidate paths, then
immediately runs every registered optimizer to prune the region the
pathfinder just touched, before moving to the next pathfinder. This
is the PassRegistry::run_lockstep contract — the working set
never explodes-then-extracts, so it stays bounded by construction
(per ../../docs/architecture/04-optimization.md §“Bounding the
frontier”: “run the merge/discard optimizers after each pathfinder
step, not once at the end”).
§Behavior preservation (B3 is a structural refactor)
The optimized result — the arm-0 winner, the retained per-device
frontier (applied inside compile_plan, unchanged), and the emitted
Op::Branch nodes — is identical to the pre-B3 hardcoded
sequence on every graph. The driver is the same three operations in
the same order; the only change is that they are reached through the
registry rather than inlined. The load-bearing
driver_matches_legacy_sequence test in
crate::optimize proves it on a representative multi-placement
graph.
Structs§
- Frontier
Convergence Optimizer - The PR-B2 per-device-frontier optimizer, re-expressed as a registered
Optimizer. It runs after each pathfinder (lock-step) and is the in-graph MERGE/DISCARD counterpart ofAlternativeSet::retain_per_device_frontier, which is applied per kernel-bearing node insidecompile_plan. - Optimization
Context - Read-only context threaded to every
Pathfinder::proposeandOptimizer::prunecall during aPassRegistry::run_lockstepdrive. - Pass
Registry - A registry of
Pathfinders andOptimizers, driven lock-step bySelf::run_lockstep. Mirrorsfuel_graph::opt::RuleRegistry: aVec<Box<dyn …>>of passes built up withwith_*builders and run through one driver method. Keeping the same shape means the future passes the constitution names (fusion, algebraic, dtype-lowering) register exactly as the two shipped passes do. - Placement
Fork Pathfinder - The PR-A4 deliberate-fork pathfinder, re-expressed as a registered
Pathfinder. Scans the ranked plan for nodes the placement DP / ranker admitted with ≥2 distinct(backend, device)placements — a genuine placement choice — and records each as ONEOp::Branchvia thefuel-graphA1 builders (arm-0 = the DP winner, arm-1 = the runner-up clone, orphaned so the live data flow is untouched and the realize result is identical).
Traits§
- Optimizer
- An optimizer MERGEs / DISCARDs paths. The PR-B2 per-device
Pareto frontier is applied per-node inside
compile_plan; the registeredFrontierConvergenceOptimizeris its in-graph counterpart that runs after each pathfinder to collapse duplicate (forward-identical) arms and enforce the optimizer invariants. - Pathfinder
- A pathfinder ADDs candidate paths / branches to the graph. The
PR-A4 deliberate-fork seed (
PlacementForkPathfinder) is the first registered pathfinder; later pathfinders (fusion, algebraic, dtype-lowering under tolerance, layout fixups) implement the same trait and register the same way.