Skip to main content

Module driver

Module driver 

Source
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-A4 PlacementForkPathfinder (deliberate-fork seed, the one pathfinder that records an Op::Branch per genuine multi-placement fork) is the first registered pathfinder.
  • Rankers — MEASURE each path on the CostVector (PR-B1/B2 rank_by_cost). Rankers are not a driver-visible trait in B3: ranking is applied per kernel-bearing node inside crate::plan::compile_plan as the per-node AlternativeSet is built, and the ranked result is carried into the driver via OptimizationContext::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-visible Ranker trait.)
  • Optimizers — MERGE / DISCARD paths. The PR-B2 per-ending-device Pareto frontier + crowding cap is applied per-node inside crate::plan::compile_plan; the registered FrontierConvergenceOptimizer is 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§

FrontierConvergenceOptimizer
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 of AlternativeSet::retain_per_device_frontier, which is applied per kernel-bearing node inside compile_plan.
OptimizationContext
Read-only context threaded to every Pathfinder::propose and Optimizer::prune call during a PassRegistry::run_lockstep drive.
PassRegistry
A registry of Pathfinders and Optimizers, driven lock-step by Self::run_lockstep. Mirrors fuel_graph::opt::RuleRegistry: a Vec<Box<dyn …>> of passes built up with with_* 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.
PlacementForkPathfinder
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 ONE Op::Branch via the fuel-graph A1 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 registered FrontierConvergenceOptimizer is 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.