Skip to main content

Module opt

Module opt 

Source
Expand description

Graph-level, backend-agnostic optimization passes.

Transforms a fuel-graph computation graph before execution. Every backend benefits from these passes because they operate purely on the abstract op graph — no backend-specific knowledge required.

Passes today:

  • CSE (common subexpression elimination): if the graph contains two structurally-identical nodes — same op, same inputs (after prior simplification), same shape/dtype — consumers of the duplicate are redirected to the canonical node. The duplicate becomes unreferenced and is silently dropped by the executor’s topo-walk. Commutative ops (Add, Mul) are keyed on sorted input IDs so a + b and b + a fold to the same canonical node. Maximum/Minimum are NOT in that set: commutative in value, not in bits, because the NaN tie-break returns the FIRST operand. See GAP-271.

  • Algebraic simplification: a handful of identity/zero rules that eliminate no-op ops:

    • AddScalar(0.0)(x) → x
    • MulScalar(1.0)(x) → x
    • Neg(Neg(x)) → x
    • Reshape(Reshape(x, _), s) → Reshape(x, s)

    These rarely appear in hand-written user code but show up routinely in autograd-generated backward graphs and in generic transformer building blocks that sometimes collapse.

Design: graphs are append-only, so rewrites don’t mutate existing nodes. Instead, the pass walks topologically, builds a remap HashMap<NodeId, NodeId> of old → canonical, and appends newly-canonicalized nodes to the same graph. Unreferenced originals remain in the arena but are never visited during realize. This is effectively a combined CSE + simplification + free DCE pass.

Return value: the function takes the roots the caller cares about and returns the rewritten roots. Callers use these to update their NodeHandle handles.

Structs§

CastFusionRule
Cast(src→dst) → Op(dst) ≡ Op(src) when the consumer has a kernel for src.
FusionRule
Auto-generated fusion rule from a FusedOpEntry. Matches the canonical primitive-subgraph pattern carried by the entry’s SubgraphPattern and emits a single Op::Fused(id, params) node. Round-trips with LoweringRule for the same id.
LoweringRule
Auto-generated lowering rule from a FusedOpEntry. Matches Op::Fused(id, _) and decomposes via the entry’s decompose function. Each rule corresponds to one fused-op id in the registry.
OrderingEdges
Derived ordering edges beyond the data-flow graph. Each entry (nid, deps) means nid must run after every node in deps.
RuleRegistry
A registry of Rules. Drives optimize-to-fixpoint via Self::optimize_to_fixpoint: lowering rules run to fixpoint, then fusion rules run to fixpoint. The locked design (ROADMAP.md Phase 7.5) intends this entry point to be wrapped in transactions in subsequent PRs without rewriting the loop body — keep this API stable for that future evolution.

Enums§

RuleFamily
Which family a Rule belongs to. Drives the phase ordering in RuleRegistry::optimize_to_fixpoint: lowering runs to fixpoint first, then fusion. This ordering ensures inverse rule pairs converge cleanly — fusion has the last word, and its output is what the executor runs.

Traits§

Rule
One graph-rewrite rule. A rule pairs a structural matcher with a rewriter that emits replacement nodes and records the (old_id → new_id) remapping so the driver can rewrite consumer input edges to point at the replacement.

Functions§

base_map_hash
NodeId-independent content hash of the subgraph rooted at root.
derive_ordering
Derive ordering edges for every destructive op reachable from roots. Result: nid → list of other readers of nid’s destroyed input.
execution_plan
Build an execution plan that respects both data-flow edges (via topo_order_multi) and ordering edges (via derive_ordering). Returns a Vec<NodeId> in an order the executor can walk linearly, evaluating each node’s dependencies before it.
fuse_linear
Fuse MatMul → Add(rank-1 bias) patterns into Op::FusedLinear (Phase 6d Track 3). Walks the graph; for each Add whose LHS is a MatMul and whose RHS is a rank-1 bias whose length equals the matmul output’s last dim, emits a fresh FusedLinear node and remaps consumers of the Add to it.
insert_cast_fixups
Insert Op::Cast fixups that CHANGE a consumer’s input dtype — the sibling of insert_layout_fixups, for dtype rather than layout. Given precomputed (consumer, input, target_dtype) edges (an input whose dtype should be promoted so the consumer’s (op, per-operand-dtypes) key becomes servable), insert one Op::Cast(target) per distinct (input, target) — CSE-shared across consumers — and rewire each consumer’s edge to it.
insert_copies
Insert Op::Copy nodes so every op’s inputs live on the same device as the op itself. Returns the rewritten roots.
insert_cross_device_copies
Insert Op::Copy { target } nodes on every edge where the consumer’s intended DeviceLocation doesn’t share a storage substrate with the producer’s. Returns a remap of original NodeIds to their post-rewrite equivalents (most NodeIds map to themselves; only consumers of cross-device edges acquire new inputs in-place via rewrite_input).
insert_evict_reload
Insert an evict-chain around candidate to free its device storage during a gap between uses, with automatic reload before the post-gap consumers. Returns (move_id, reload_id).
insert_layout_fixups
Phase 2.2 of the picker arc — insert Op::Contiguize before any kernel-bearing node whose chosen kernel doesn’t accept strided inputs but whose live input layout is non-contiguous (strided, offset, or both).
insert_safety_copies
Insert defensive Op::Copy snapshots ahead of destructive ops that would otherwise produce a data-flow ↔ ordering-edge cycle.
lower_const_placement
Lower Const placements: if every consumer of an unplaced Const has the same placement hint, tag the Const with that placement.
lower_to_base_map
Lower every reachable fused op in graph (rooted at roots) to its primitive base map — the fixpoint of decompose over every node. A thin named wrapper over RuleRegistry::lowering_only + optimize_to_fixpoint (the machinery every fused op — static-registry AND runtime — already flows through). Never panics: a self-returning decompose is a clean fixpoint (the never-panic total-decompose contract), not a loop.
optimize
Run CSE + algebraic simplification on the graph reachable from roots. Returns the (possibly-rewritten) roots the caller should use afterward.
promote_views_for_liveness
Walk every multi-output producer in the graph and promote some of its Op::View consumers to Op::ViewOwned when slot lifetimes are sufficiently asymmetric that keeping the whole bundle alive for the longest-lived slot is wasteful.

Type Aliases§

CapabilityPredicate
Predicate answering “is there a registered kernel for this op with these input/output dtypes?” The dtypes vector is the consumer’s full input-dtype list (after the cast-source substitution) followed by the consumer’s output dtype. Format matches the binding-table key shape used by fuel-storage.