fuel_dispatch/optimize.rs
1// SPDX-License-Identifier: MIT OR Apache-2.0
2//! `optimize_graph` — the new "plan IS the graph" entry point.
3//!
4//! Phase A PR-A3a/A4 of the "plan IS the graph" rebuild
5//! ([`../../docs/session-prompts/plan-is-graph-rebuild.md`](
6//! ../../docs/session-prompts/plan-is-graph-rebuild.md), the PR-A3/A4
7//! lines + the "Representation — DECIDED 2026-06-15" blockquote, incl.
8//! the architect-approved arm-0 temporary lowering).
9//!
10//! ## What this is
11//!
12//! [`optimize_graph`] is the **new optimization entry point**: it
13//! transforms a [`Graph`] *in place* into the bounded multi-path form
14//! and returns a transient [`OptimizedGraph`] *view* whose dispatch
15//! order is the **arm-0 single-route lowering** of `fuel-graph`'s
16//! `extract_runs` (PR-A2/A4). The optimized form lives **in the
17//! graph** — a graph with zero [`Op::Branch`] nodes is exactly today's
18//! single-route graph.
19//!
20//! ## PR-A4 — the first pathfinder (deliberate-fork seed)
21//!
22//! [`PlacementForkPathfinder`](crate::driver::PlacementForkPathfinder) is the first real pathfinder: where
23//! the placement DP admitted a kernel-bearing node with **≥2 distinct
24//! `(backend, device)` placements** that has a producer (the diverge)
25//! and exactly one consumer (the reconverge), it records ONE
26//! `Op::Branch` — arm-0 = the DP winner (the route realize uses today),
27//! arm-1 = the runner-up clone (orphaned, read only by the Branch). It
28//! emits a branch only at a genuine placement choice, never at ordinary
29//! DAG fan-out, so the fewness gate holds. A CPU-only build (one
30//! placement per node) emits zero branches ⇒ today's single-route graph.
31//!
32//! [`OptimizedGraph::dispatch_order`] is the architect-approved
33//! **arm-0 single-route lowering**: it follows arm 0 through every
34//! branch and skips the other arms' runs, so a branched graph realizes
35//! to the same result as before (arm-0 = the winner). The Phase-C
36//! runtime picker is what later selects non-arm-0 arms.
37//!
38//! ## A3a origin scope (ADDITIVE — proved equivalence, deleted nothing)
39//!
40//! This entry point landed in the **split** A3a: introduced *alongside*
41//! the old path and proved equivalent before anything old was removed.
42//! PR-A3b-1 then made it the default realize path, and PR-A3b-2 deleted
43//! `PlanStore` and the legacy `compile_plan` route-picking dispatch
44//! entirely (`optimize_graph` is now the only path). As originally
45//! introduced in A3a:
46//!
47//! - There were **no pathfinders yet** — the first
48//! ([`PlacementForkPathfinder`](crate::driver::PlacementForkPathfinder)) landed in PR-A4. A graph with no
49//! competing routes is already its own single-route plan, so
50//! [`optimize_graph`] introduces **zero [`Op::Branch`] nodes** for it.
51//! - It **reuses** the existing placement / cost / `target_backend`
52//! annotation machinery wholesale by driving
53//! [`crate::plan::compile_plan`] over the same `execution_plan`
54//! order the bridge uses — the placement DP, cost composer, filter
55//! chain, and fail-fast missing-binding diagnostics are *not*
56//! reinvented. The point of A3a is to establish the new entry point
57//! + in-place form, not new optimization.
58//! - [`crate::plan::compile_plan`] + [`crate::plan::ExecutionPlan`]
59//! are reused (the latter as a transient by-product). PR-A3b-1 wired
60//! `optimize_graph` in as the default realize path; PR-A3b-2 made it
61//! the ONLY path — the legacy `compile_plan`/`PlanStore` route-picking
62//! dispatch and the identity-keyed plan store are deleted, and
63//! `optimize_graph` now surfaces its internal `ExecutionPlan` so the
64//! bridge's stamp/residency/layout passes reuse it (one `compile_plan`
65//! per realize).
66//!
67//! ## The equivalence gate
68//!
69//! For a graph with no competing routes, [`optimize_graph`] leaves
70//! zero `Op::Branch` nodes and the `extract_runs` / `lower_run`
71//! dispatch order equals today's
72//! [`crate::plan::compile_plan`]`(...).order` — the exact sequence the
73//! production executor walks. The order today is simply the
74//! `execution_plan` topo order `compile_plan` is handed (it copies it
75//! into `ExecutionPlan::order` verbatim); for a branchless,
76//! single-residency graph the run extractor produces exactly one run
77//! whose members are that same topo order. The gate asserts that
78//! equality exactly (same `NodeId`s, same order) — see
79//! `tests::equivalence_gate_branchless_order_matches_compile_plan`.
80
81use std::collections::{HashMap, HashSet};
82
83use fuel_graph::opt::insert_cross_device_copies;
84use fuel_graph::{Graph, NodeId, Op, extract_runs_multi, topo_order_multi};
85use fuel_ir::probe::BackendId;
86use fuel_ir::{DType, DeviceLocation, Result};
87
88use crate::driver::{OptimizationContext, PassRegistry};
89use crate::kernel::KernelBindingTable;
90use crate::plan::{ExecutionPlan, PlanOptions, compile_plan};
91use crate::topology::SystemTopology;
92
93/// The transient *view* [`optimize_graph`] returns — the realize-roots
94/// it optimized for plus the topology generation it ran under. It is
95/// **not** stored on the graph; the optimized form is the graph itself
96/// (with its `Op::Branch` decision points, of which A3a emits none).
97///
98/// The dispatch order is derived on demand from `fuel-graph`'s run
99/// extraction so the run view is always recomputed from the current
100/// arena — never a stale snapshot. This is the shape PR-A3b swaps the
101/// bridge onto in place of `ExecutionPlan::order`.
102#[derive(Debug, Clone)]
103pub struct OptimizedGraph {
104 /// The realize roots this optimization targeted (the same
105 /// `targets` the bridge passes today).
106 pub roots: Vec<NodeId>,
107 /// `SystemTopology` generation snapshotted at optimize time —
108 /// mirrors [`crate::plan::ExecutionPlan::generation`] so a later
109 /// chunk-boundary check can detect a topology shift exactly as
110 /// the executor does today.
111 pub generation: u64,
112 /// Final effective placement per node, as the optimizer resolved it.
113 /// Read via [`OptimizedGraph::placement_of`].
114 ///
115 /// `None` when the residency pass did not run — it is gated on
116 /// `PlanOptions::pinned_device`, which the production realize path always
117 /// sets and bare-graph test callers often do not. `None` therefore means
118 /// **"not computed"**, never "placed nowhere".
119 pub(crate) placements: Option<HashMap<NodeId, DeviceLocation>>,
120}
121
122impl OptimizedGraph {
123 /// Extract the runs of the optimized graph, in topological order
124 /// of their entries (delegates to [`extract_runs_multi`]). A
125 /// branchless, single-residency graph yields exactly one run.
126 pub fn runs(&self, graph: &Graph) -> Vec<fuel_graph::Run> {
127 extract_runs_multi(graph, &self.roots)
128 }
129
130 /// **The device the optimizer decided this node runs on** — the observable
131 /// form of the placement decision, for callers that need to check *where*
132 /// work landed rather than infer it from timing.
133 ///
134 /// This is the resolved *effective* placement, not a hint. It is the same
135 /// map the residency pass uses to decide where to stitch cross-device
136 /// copies, resolved by the priority in `effective_placements`:
137 /// residency-declaring ops → explicit `Graph::placement` → input residency
138 /// → **plan winner** → backend stamp → view pass-through.
139 ///
140 /// **Do not substitute [`fuel_graph::Graph::placement`] for this.** That is
141 /// only *level 2* of the six — a hint set by inserted copies and caller
142 /// overrides. The planner's actual choice for a compute node arrives at
143 /// level 4 ("plan winner"), so `Graph::placement` returns `None` for
144 /// precisely the nodes whose placement is interesting, and reading that
145 /// `None` as "unplaced" is a wrong answer that looks like a null one.
146 ///
147 /// Returns `None` if the node is unknown to this optimization **or** if the
148 /// residency pass did not run (no `pinned_device`; see the field docs) —
149 /// the two are deliberately not distinguished here because neither means
150 /// "placed nowhere". Use [`OptimizedGraph::has_placements`] to tell them
151 /// apart before drawing a conclusion from a `None`.
152 pub fn placement_of(&self, node: NodeId) -> Option<DeviceLocation> {
153 self.placements.as_ref()?.get(&node).copied()
154 }
155
156 /// Whether placement was computed at all. Distinguishes "this node has no
157 /// placement" from "no node does, because the pass was skipped" — the
158 /// difference between a real observation and an absent instrument.
159 pub fn has_placements(&self) -> bool {
160 self.placements.is_some()
161 }
162
163 /// Every resolved `(node, device)` pair, for dumping a whole graph's
164 /// placement. Empty iterator when placement was not computed — check
165 /// [`OptimizedGraph::has_placements`] first if that distinction matters.
166 pub fn placements(&self) -> impl Iterator<Item = (NodeId, DeviceLocation)> + '_ {
167 self.placements
168 .iter()
169 .flat_map(|m| m.iter().map(|(&id, &loc)| (id, loc)))
170 }
171
172 /// The flat executable dispatch order — the **arm-0 single-route
173 /// lowering** (PR-A4). It follows arm 0 through every `Op::Branch`
174 /// (pre-run, arm-0's run, post-run) and **skips every non-arm-0
175 /// arm's run**, via [`fuel_graph::lower_runs_arm0`].
176 ///
177 /// For a graph with **zero `Op::Branch` nodes** this is the exact
178 /// `NodeId` sequence today's executor walks via
179 /// [`crate::plan::compile_plan`]`(...).order` — the equivalence gate
180 /// proves it (no arm to skip ⇒ identical to concatenating
181 /// [`fuel_graph::run::lower_run`] over the runs). For a branched graph it is the
182 /// single-route lowering on **arm 0 = the DP winner** (the route
183 /// realize used before the branch was recorded), so a branched graph
184 /// realizes to the same result. The Phase-C runtime picker is what
185 /// will later select non-arm-0 arms at the decision points.
186 pub fn dispatch_order(&self, graph: &Graph) -> Vec<NodeId> {
187 fuel_graph::lower_runs_arm0(graph, &self.roots)
188 }
189
190 /// The number of `Op::Branch` decision points in the arena. Counting
191 /// over the whole arena (not just the roots' reachable set) is the
192 /// stronger claim. Zero for a single-route graph (CPU-only build, or
193 /// any graph with no genuine ≥2-placement fork); PR-A4's
194 /// deliberate-fork pathfinder emits one per genuine placement fork.
195 pub fn branch_count(&self, graph: &Graph) -> usize {
196 (0..graph.len())
197 .map(NodeId)
198 .filter(|&id| matches!(graph.node(id).op, Op::Branch { .. }))
199 .count()
200 }
201}
202
203/// Optimize `graph` **in place** into the "plan IS the graph" form and
204/// return the transient [`OptimizedGraph`] lowering view.
205///
206/// `roots` are the realize targets (the bridge's `targets`);
207/// `bindings_table` is the kernel registry the placement/cost
208/// machinery queries (production passes [`crate::dispatch::global_bindings`];
209/// tests pass a local table); `opts` is the same [`PlanOptions`] the
210/// bridge builds for `compile_plan`.
211///
212/// ## Behavior
213///
214/// `optimize_graph`:
215///
216/// 1. Derives the dispatch order via `fuel_graph::opt::execution_plan`
217/// — exactly the order the bridge feeds `compile_plan` today
218/// (data-flow topo refined by destructive-op ordering edges).
219/// 2. Drives [`compile_plan`] over that order to **reuse** the
220/// existing placement DP / cost composer / filter chain and to
221/// fail-fast at build time on a missing binding or absent device
222/// context (validate-at-graph-build-time per the working
223/// agreement). Its [`crate::plan::ExecutionPlan`] is a transient
224/// by-product here — the source of truth is the graph.
225/// 3. Runs the PR-A4 deliberate-fork pathfinder
226/// ([`PlacementForkPathfinder`](crate::driver::PlacementForkPathfinder)): records ONE `Op::Branch` per
227/// genuine ≥2-placement fork (arm-0 = winner, arm-1 = runner-up). A
228/// graph with no competing routes gets **zero** branches — it is
229/// already its own single-route plan.
230///
231/// The returned view's [`OptimizedGraph::dispatch_order`] is the arm-0
232/// single-route lowering; for any branchless graph it reproduces
233/// `compile_plan(...).order` exactly — the equivalence gate.
234///
235/// ## The internal `ExecutionPlan` (Step D — no longer returned)
236///
237/// `optimize_graph` drives [`compile_plan`] internally for
238/// placement/cost/validation; the resulting `ExecutionPlan` is a
239/// **transient, optimizer-internal** accumulator the stamp / residency /
240/// layout-fixup passes + the lock-step rankers read below. It is NOT
241/// returned or threaded anywhere — the durable output is the graph's
242/// `target_backend` stamps + `Op::Branch` arms. The executor reads the
243/// graph + re-derives any per-arm candidate from the binding registry at
244/// realize time, so the plan never becomes a dispatch-time authority.
245pub fn optimize_graph(
246 graph: &mut Graph,
247 roots: &[NodeId],
248 bindings_table: &KernelBindingTable,
249 opts: &PlanOptions<'_>,
250) -> Result<OptimizedGraph> {
251 optimize_graph_with_passes(
252 graph,
253 roots,
254 bindings_table,
255 opts,
256 PassRegistry::default_passes(),
257 )
258}
259
260/// [`optimize_graph`] with the runtime-fusion pathfinder registered
261/// ([`PassRegistry::default_passes_with_runtime_fusion`]): adopted (Tier-2 /
262/// JIT-synthesized) fused ops get their gated `Op::Branch` arms emitted before
263/// placement runs. This is the **production realize path's** entry once a
264/// process can adopt runtime kernels; `optimize_graph` stays runtime-op-free so
265/// bare-graph callers and tests never scan the process-global sidecar.
266pub fn optimize_graph_with_runtime_fusion(
267 graph: &mut Graph,
268 roots: &[NodeId],
269 bindings_table: &KernelBindingTable,
270 opts: &PlanOptions<'_>,
271) -> Result<OptimizedGraph> {
272 optimize_graph_with_passes(
273 graph,
274 roots,
275 bindings_table,
276 opts,
277 PassRegistry::default_passes_with_runtime_fusion(),
278 )
279}
280
281fn optimize_graph_with_passes(
282 graph: &mut Graph,
283 roots: &[NodeId],
284 bindings_table: &KernelBindingTable,
285 opts: &PlanOptions<'_>,
286 passes: PassRegistry,
287) -> Result<OptimizedGraph> {
288 // (1) The dispatch order today: data-flow topo refined by ordering
289 // edges. `compile_plan` copies this verbatim into
290 // `ExecutionPlan::order`, so it IS the executor's walk order.
291 let order = fuel_graph::opt::execution_plan(graph, roots);
292
293 // (2) Reuse the placement/cost/validation machinery. The plan is a
294 // transient, optimizer-INTERNAL by-product — we drive `compile_plan`
295 // so the same fail-fast diagnostics (missing binding, no device
296 // context) fire at optimize time, and so the placement DP / cost
297 // composer run unchanged. The plan is the working accumulator the
298 // stamp / residency / layout-fixup passes + the lock-step rankers
299 // read below; it is NOT returned or threaded (Step D — the graph
300 // stamps + arms are the only durable output).
301 // (2a) Dtype-reconciliation fixup — insert promoting `Op::Cast`s so a node
302 // whose per-operand key no compiled-in backend serves natively (e.g. a
303 // mixed-precision `[F32, BF16, F32]` matmul) becomes servable, instead of
304 // dead-ending in `compile_plan` below. Runs BEFORE the plan because it
305 // changes what's enumerable; if it rewrites, re-derive the execution
306 // order over the mutated graph. The layout/residency fixups are its
307 // post-plan siblings.
308 let order = if insert_dtype_fixups(graph, roots, bindings_table, opts) {
309 fuel_graph::opt::execution_plan(graph, roots)
310 } else {
311 order
312 };
313
314 let plan = compile_plan(graph, &order, bindings_table, opts)?;
315 let generation = plan.generation;
316
317 // (3) PR-B3: the lock-step pass driver. The pre-B3 hardcoded
318 // `seed_placement_fork_branches(...)` call is replaced by a
319 // registry of pathfinders + optimizers run lock-step
320 // (prune-as-you-go): for each registered pathfinder, ADD its
321 // candidate paths, then immediately run every registered
322 // optimizer to MERGE/DISCARD over the region just touched. The
323 // shipped configuration ([`PassRegistry::default_passes`]) is the
324 // PR-A4 `PlacementForkPathfinder` (the deliberate-fork seed) +
325 // the `FrontierConvergenceOptimizer` (duplicate-path convergence
326 // + the never-strand / no-active-cycle invariant guards) — which
327 // is exactly the pre-B3 sequence, re-expressed. A CPU-only build
328 // has one placement per node ⇒ the pathfinder proposes zero
329 // branches ⇒ today's single-route graph (unchanged).
330 //
331 // The *ranker* + the PR-B2 per-device Pareto frontier are applied
332 // per kernel-bearing node *inside* `compile_plan` (the MEASURE +
333 // per-node PRUNE), and the ranked, frontier-pruned result is the
334 // `plan` the driver reads via `OptimizationContext::plan`. Batch
335 // optimize has no executing region, so the cycle guard is empty.
336 let cycle_guard: HashSet<NodeId> = HashSet::new();
337 let ctx = OptimizationContext {
338 order: &order,
339 plan: &plan,
340 cycle_guard: &cycle_guard,
341 bindings: bindings_table,
342 };
343 passes.run_lockstep(graph, &ctx)?;
344
345 // "Plan IS the graph": commit the placement decision INTO the graph's
346 // `target_backend` side-table here, so a fully-optimized graph carries
347 // its own backend stamps and downstream consumers read the graph, not a
348 // threaded plan. Guarded on a pinned device — the production realize
349 // path always sets one (`PlanOptions::with_pinned_device`); bare-graph
350 // test callers that don't are unaffected. (Cleanup Step A: this is
351 // idempotent with the bridge's transitional `stamp_plan_backends`, which
352 // re-stamps the identical result; Step A2 removes the bridge copy.)
353 // Stays `None` when the residency pass below is skipped (no pinned device),
354 // so a caller can tell "not computed" from "no placement" — see
355 // `OptimizedGraph::has_placements`.
356 let mut resolved_placements: Option<HashMap<NodeId, DeviceLocation>> = None;
357 if let Some(pinned) = opts.pinned_device {
358 stamp_plan_backends(graph, roots, &plan, pinned);
359 // Optimize-time kernel-variant bake: resolve **same-device**
360 // kernel-variant branches (a decomposed region vs. a fused kernel on
361 // the SAME device — e.g. the CUDA flash decode arm) to their cost
362 // winner and COLLAPSE them, so the winning variant is selectable
363 // without a runtime pick (04-optimization: kernel-variant choice is
364 // "largely baked at optimize time"). Runs AFTER `stamp_plan_backends`
365 // (the arms now carry `target_backend`, which the same-device gate
366 // reads) and BEFORE `insert_residency_copies` (so the residency pass
367 // stitches inbound copies only for the surviving winner arm, not a
368 // pruned one). A **placement** branch (arms on ≥2 devices) is left
369 // LIVE for the runtime route picker; a graph with no same-device
370 // variant branch — every CPU/Vulkan build today, since no CUDA flash
371 // arm is offered there — is a strict no-op. Ties / unknown costs /
372 // capability-missing default to arm 0 (the decomposed oracle).
373 // The cost provider consults the Judge's MEASURED latency (Layer-2)
374 // FIRST per arm node, falling back to the Layer-1 composite where the
375 // Judge has no cell (the hybrid sum). With `opts.judge == None` this is
376 // byte-identical to the Layer-1-only provider, so the no-Judge path is
377 // unchanged. The same oracle `cost.rs`/`plan.rs` read is threaded here.
378 let judge = opts.judge;
379 crate::variant_bake::bake_variant_branches(graph, roots, &|g, branch, arm, interior| {
380 let backend = g.target_backend(g.node(branch).inputs.first().copied()?)?;
381 crate::variant_bake::decode_arm_composite_ns_judged(
382 g, branch, arm, interior, backend, judge,
383 )
384 });
385 // Cleanup Step B (residency): the optimizer-side cross-device copy pass
386 // (was the bridge's `insert_resident_input_copies`). Runs AFTER stamping
387 // (it reads `target_backend`) and BEFORE layout-fixup, preserving the
388 // prior residency→layout order. Cache residency (where persistent /
389 // const inputs already live) arrives via the `opts.input_residency`
390 // provider the realize path already threads in — the optimizer reads
391 // that fact, it does not touch realize-time storage.
392 // Capture the resolved placement map on the way past — it is the pass's
393 // own input, and surfacing it is what makes `OptimizedGraph::placement_of`
394 // an observation rather than an inference.
395 resolved_placements = Some(insert_residency_copies(
396 graph,
397 roots,
398 &plan,
399 pinned,
400 opts.input_residency,
401 ));
402 // Cleanup Step B (layout): insert `Op::Contiguize` before any kernel
403 // whose chosen winner rejects strided inputs and whose input layout is
404 // non-contiguous — the optimizer writing the layout-fixup decision INTO
405 // the graph (was the bridge's `apply_layout_fixups`, now retired). No
406 // runtime deps: reads only the in-scope `plan` winner caps + graph
407 // layout. CSE-deduped + idempotent. Gated on `pinned_device` so it runs
408 // on the realize path (which always sets one), matching the prior
409 // bridge behavior; bare-graph test callers without a pin are unaffected.
410 fuel_graph::opt::insert_layout_fixups(graph, roots, |id| {
411 plan.alternatives(id)
412 .and_then(|set| set.winner())
413 .map(|cand| cand.caps.strided_input)
414 .unwrap_or(true)
415 });
416 }
417
418 Ok(OptimizedGraph {
419 roots: roots.to_vec(),
420 generation,
421 placements: resolved_placements,
422 })
423}
424
425/// Dtype-reconciliation fixup (runs BEFORE `compile_plan`): for any kernel-
426/// bearing node whose per-operand `(op, dtypes)` key is not servable by any
427/// compiled-in backend, insert `Op::Cast`(s) promoting the mismatched input(s)
428/// to the node's OUTPUT dtype — IF that lossless upcast yields a key some backend
429/// DOES support. This is "backends advertise; the optimizer adapts"
430/// ([04-optimization], [05-backend-contract]) applied to dtype, and the sibling
431/// of `insert_layout_fixups` (`Op::Contiguize`) and `insert_residency_copies`
432/// (`Op::Copy`).
433///
434/// It is the non-native-kernel fallback the mixed-precision matmul path already
435/// ASSUMED exists — `vulkan_dispatch`'s matmul route picker documents "otherwise
436/// falls back to f32 matmul (after a Cast)" — but which was never built as a
437/// pass, so a graph builder emitting a mixed `[F32, BF16] → F32` matmul
438/// dead-ended at `NoBackendForOp` on every backend without the native mixed key
439/// (all but Vulkan; and CPU-only default builds have none).
440///
441/// SAFE + CONSERVATIVE: promotes ONLY by a value-LOSSLESS upcast to the output
442/// dtype ([`lossless_upcast`]), so it never changes op semantics (a comparison's
443/// `U8` output is not a lossless upcast target, so comparisons are left alone) and
444/// never corrupts a value.
445///
446/// **NUMERIC CAVEAT:** the upcast is value-lossless but the promoted op is NOT
447/// numerically identical to a native mixed-precision kernel — it accumulates in
448/// the higher precision, so the *same graph* can differ by which arm ran (a CPU
449/// upcast vs a native `[F32,BF16,F32]` kernel; and it already differs
450/// CPU-vs-CUDA). The pass REPORTS every promotion via `tracing` (visibility, per
451/// C-6's "never silently deoptimize"). Making a promoting cast **forbiddable
452/// under a C-5 determinism/tolerance constraint** (so a byte-exact-oracle or
453/// CPU-vs-CUDA-parity consumer gets the native path or a hard fail, not a silent
454/// upcast) and **accounting its extra resident bytes under C-4** (BF16→F32
455/// doubles a weight) are tracked ROADMAP follow-ups.
456///
457/// When no lossless promotion yields a supported key, the node is left unchanged
458/// and fails cleanly at `compile_plan` with the truthful coverage diagnostic — a
459/// genuine gap, surfaced honestly. Returns whether it changed the graph (so the
460/// caller re-derives the execution order).
461///
462/// The broader fix — expanding per-op dtype coverage across kernels (starting
463/// CPU) so fewer combos need reconciling — is a tracked ROADMAP item; this pass
464/// makes the graph builder's mixed-precision promise executable everywhere today.
465fn insert_dtype_fixups(
466 graph: &mut Graph,
467 roots: &[NodeId],
468 bindings: &KernelBindingTable,
469 options: &PlanOptions<'_>,
470) -> bool {
471 let order = topo_order_multi(graph, roots);
472 // Decision only (immutable graph reads): the `(consumer, input, target)` edges
473 // to promote. The mechanical Cast-insert + rewire is
474 // `fuel_graph::opt::insert_cast_fixups` (the dtype-changing sibling of the
475 // layout fixup). Built in topo order ⇒ deterministic cast-node assignment.
476 let mut casts: Vec<(NodeId, NodeId, DType)> = Vec::new();
477
478 for &id in &order {
479 let (op, inputs, out_dtype) = {
480 let node = graph.node(id);
481 (node.op.clone(), node.inputs.clone(), node.dtype)
482 };
483 let op_kind = match crate::pipelined::op_to_op_kind(&op) {
484 Some(k) => k,
485 None => continue, // view / Const / not-dispatched → no binding key
486 };
487 // The exact per-operand key the dispatcher will look up.
488 let native_key = crate::pipelined::build_lookup_dtypes(graph, graph.node(id));
489 // Only reconcile the generic `[in0, in1, …, out]` key shape. Ops whose
490 // key is canonicalized (Concat, WriteSlice*) share a uniform input dtype
491 // by construction, so they never present a mixed-input miss here.
492 if native_key.len() != inputs.len() + 1 {
493 continue;
494 }
495 // The backends THIS node may actually be placed on. Asked via the
496 // planner's own resolver so the two cannot drift: a weaker predicate
497 // here (e.g. "does any registered backend serve it?") lets a backend
498 // the node can never run on suppress the reconciliation the node needs,
499 // and `compile_plan` below then dead-ends at `NoBackendForOp`.
500 let candidates = crate::plan::candidate_backends_for(graph, options, id, graph.node(id));
501 let servable = |key: &[DType]| {
502 candidates
503 .iter()
504 .any(|&backend| bindings.backend_supports(op_kind, key, backend))
505 };
506
507 if servable(&native_key) {
508 continue; // a backend this node can run on serves it natively
509 }
510 // Promotion target = the output dtype ("compute dtype"), taken ONLY when
511 // every input can losslessly upcast to it — else decline.
512 let target = out_dtype;
513 if !inputs
514 .iter()
515 .all(|&i| lossless_upcast(graph.node(i).dtype, target))
516 {
517 continue;
518 }
519 // Uniform promoted key `[target, …, target, out]` (target == out).
520 let promoted_key: Vec<DType> = inputs
521 .iter()
522 .map(|_| target)
523 .chain(std::iter::once(out_dtype))
524 .collect();
525 if promoted_key == native_key {
526 continue; // inputs already == target; nothing to promote
527 }
528 if !servable(&promoted_key) {
529 continue; // the lossless promotion still isn't servable — genuine gap
530 }
531 for &input_id in &inputs {
532 if graph.node(input_id).dtype != target {
533 casts.push((id, input_id, target));
534 }
535 }
536 }
537
538 if casts.is_empty() {
539 return false;
540 }
541 // Visibility (C-6 "fuel never silently deoptimizes without saying so", applied
542 // to precision rather than fusion): a promoting cast is NOT numerically
543 // identical to a native mixed-precision kernel — it accumulates the upcast in
544 // higher precision, so the same graph can diverge by which arm ran (CPU upcast
545 // vs CUDA native mixed) — and it increases the promoted operand's resident
546 // bytes (e.g. BF16→F32 doubles a weight). A byte-exact-oracle or determinism
547 // consumer must be able to SEE it. Reported via `tracing`; forbidding it under
548 // a C-5 constraint + full C-4 cost accounting are tracked ROADMAP follow-ups.
549 tracing::warn!(
550 promoting_casts = casts.len(),
551 "dtype-reconciliation inserted {} promoting Cast(s): a node's per-operand \
552 dtype key had no native kernel on any backend, so input(s) were upcast to \
553 the output dtype. This CHANGES numerics vs a native mixed-precision kernel \
554 and increases resident bytes for the promoted operand(s).",
555 casts.len(),
556 );
557 fuel_graph::opt::insert_cast_fixups(graph, &casts);
558 true
559}
560
561/// Lossless dtype upcast predicate for [`insert_dtype_fixups`] — `true` ONLY for
562/// conversions that preserve every value exactly (identity + float widenings
563/// where the target's exponent AND mantissa cover the source's). Deliberately
564/// conservative: an unlisted pair returns `false`, so the pass declines rather
565/// than risk a lossy or semantics-changing promotion.
566fn lossless_upcast(from: DType, to: DType) -> bool {
567 use DType::*;
568 if from == to {
569 return true;
570 }
571 matches!(
572 (from, to),
573 (BF16, F32) | (BF16, F64) | (F16, F32) | (F16, F64) | (F32, F64)
574 )
575}
576
577/// Commit the plan's per-node winner backend to the graph's
578/// `target_backend` side-table — the optimizer writing its placement
579/// decision into the graph.
580///
581/// Per kernel-bearing node: stamp `winner.backend` if the plan has an
582/// `AlternativeSet` for it, else the pinned device's backend (structural
583/// ops the planner skips — `Op::Copy`/`Op::Move`/`Op::Alloc`/`Op::ZeroFill`
584/// — plus any op without an `OpKind` mapping). `Op::Const`/`Op::Release`/
585/// `Op::Contiguize`/view ops/`Op::Reshape` inherit or don't need a stamp
586/// and are skipped. Always overwrites, so re-optimizing after a
587/// `TopologyChanged` retry re-stamps consistently from the fresh plan.
588fn stamp_plan_backends(
589 graph: &mut Graph,
590 roots: &[NodeId],
591 plan: &ExecutionPlan,
592 pinned_loc: DeviceLocation,
593) {
594 let pinned_backend = location_to_backend_id(pinned_loc);
595 let order = topo_order_multi(graph, roots);
596 for &id in &order {
597 let node = graph.node(id);
598 if matches!(node.op, Op::Const | Op::Release | Op::Contiguize)
599 || node.op.is_view_op()
600 || matches!(node.op, Op::Reshape(_))
601 {
602 continue;
603 }
604 let stamp = plan
605 .alternatives(id)
606 .and_then(|set| set.winner())
607 .map(|c| c.backend)
608 .unwrap_or(pinned_backend);
609 graph.set_target_backend(id, stamp);
610 }
611}
612
613/// Cleanup Step B (residency): the optimizer-side cross-device copy pass —
614/// insert `Op::Copy { target }` on every edge whose producer's resident
615/// location doesn't share a storage substrate with the consumer's placement,
616/// then stamp each inserted copy's `target_backend` with the SOURCE backend
617/// (the pipelined executor's `Op::Copy` convention: the transfer kernel runs
618/// on the backend the bytes come FROM). Was the bridge's
619/// `insert_resident_input_copies`; moved here so the optimizer — not the
620/// realize-time bridge — owns the residency decision and writes it into the
621/// graph.
622///
623/// Placements come from graph-knowable facts only (no realize-time storage):
624/// residency-declaring ops, explicit `Graph::placement` (the bridge converts
625/// cache residency into placement stamps before optimize — cleanup Step B),
626/// the plan winner's device, the `target_backend` stamp (→ pinned), and view
627/// pass-through. Substrate sharing is queried via the process-global
628/// `SystemTopology` (available pre-realize once backends are registered).
629///
630/// Runs after `stamp_plan_backends` (it reads `target_backend`); the re-stamp
631/// sweep below restores the source-backend stamp on copies/moves that the
632/// stamping pass overwrote with the pinned backend.
633/// Returns the effective-placement map it resolved, so the optimizer can hand
634/// it out via [`OptimizedGraph::placement_of`]. The map is the pass's own input
635/// — surfacing it costs nothing and makes the placement decision observable
636/// instead of inferable.
637///
638/// Note the snapshot is taken BEFORE copy insertion (see the `placements`
639/// comment below), which is the right vantage for "where did this node land":
640/// it reports the placement of the original nodes, with the inserted copies
641/// being the consequence rather than part of the answer.
642fn insert_residency_copies(
643 graph: &mut Graph,
644 roots: &[NodeId],
645 plan: &ExecutionPlan,
646 pinned_loc: DeviceLocation,
647 input_residency: Option<&dyn Fn(NodeId) -> Option<DeviceLocation>>,
648) -> HashMap<NodeId, DeviceLocation> {
649 // Step E Phase C, PR C-0: walk the ARM-INCLUSIVE reachable set, not bare
650 // `roots`. A finalized `Op::Branch` is orphaned (its `reconverge_at` reads
651 // arm-0 directly, per the PR-A1 runnability invariant), so a plain
652 // `topo_order_multi(roots)` never reaches it — and therefore never reaches
653 // its non-arm-0 arms. The residency pass MUST stitch inbound cross-device
654 // copies for EVERY surviving arm so the executor can legally re-pick
655 // arm-1+ at runtime by live load (C2): whichever arm it picks, that arm's
656 // device-inputs are already resident. `effective_roots` (the same seeding
657 // the run extractor + route picker use) pulls every Branch's arms into the
658 // walk. On a branchless graph it returns exactly `roots` ⇒ byte-identical.
659 let eff_roots = fuel_graph::effective_roots(graph, roots);
660 let placements = effective_placements(graph, &eff_roots, plan, pinned_loc, input_residency);
661 let topology = SystemTopology::current();
662 // Identical locations short-circuit to `true` BEFORE the topology lookup
663 // (a location trivially shares bytes with itself; `shares_storage` returns
664 // `false` for unknown backends, so without this an unprobed topology would
665 // wrap every same-device edge in a copy).
666 let shares = |a: DeviceLocation, b: DeviceLocation| -> bool {
667 if a == b {
668 return true;
669 }
670 topology.shares_storage(
671 (location_to_backend_id(a), a),
672 (location_to_backend_id(b), b),
673 )
674 };
675 // Hop-count oracle: an edge whose transfer path is `HostStaging` can't
676 // be done as a single `Op::Copy` (no direct device-to-device copy
677 // kernel between distinct GPU substrates — the cross-VENDOR CUDA↔Vulkan
678 // case), so the residency pass routes it through a CPU intermediate as
679 // TWO hops. `SystemTopology::transfer_path` already returns `HostStaging`
680 // for such pairs and `SameDevice` for `a == b`; CPU↔CUDA / CPU↔Vulkan
681 // resolve to a direct path (or the universal host-staging fallback only
682 // when CPU is NOT one of the endpoints), so any CPU-touching edge stays
683 // single-hop. We additionally guard on neither endpoint being CPU: a CPU
684 // intermediate for a CPU-touching edge would be nonsensical (and the
685 // single-hop CPU↔GPU Copy kernels exist and are byte-identical to
686 // today). Same-vendor GPU pairs that don't share storage (e.g. CUDA gpu0
687 // ↔ gpu1) fall back to `HostStaging` too — correctly two-hop, since
688 // there's no direct cross-gpu_id copy kernel today either.
689 let needs_host_staging = |a: DeviceLocation, b: DeviceLocation| -> bool {
690 if a == DeviceLocation::Cpu || b == DeviceLocation::Cpu {
691 return false;
692 }
693 matches!(
694 topology.transfer_path(a, b),
695 fuel_ir::backend::TransferPath::HostStaging
696 )
697 };
698 let inserted = insert_cross_device_copies(
699 graph,
700 &eff_roots,
701 |id| placements.get(&id).copied(),
702 shares,
703 needs_host_staging,
704 );
705 // Source-location resolver for stamping inserted copies. The
706 // `placements` snapshot was computed BEFORE insertion, so it does NOT
707 // contain the freshly-inserted CPU intermediate of a two-hop edge. That
708 // intermediate's placement was written into the graph by the pass
709 // (`set_placement(cpu_id, Cpu)`), so fall back to `graph.placement(src)`
710 // when the snapshot lacks the node — this is what stamps a two-hop's
711 // consumer-side hop with `Cpu` (its bytes come FROM the CPU intermediate,
712 // so its Copy kernel is the CPU `copy_from_cpu_wrapper`).
713 let src_location = |graph: &Graph, src: NodeId| -> Option<DeviceLocation> {
714 placements
715 .get(&src)
716 .copied()
717 .or_else(|| graph.placement(src))
718 };
719 // Stamp the new copies: target_backend = SOURCE backend. The pass only
720 // inserts a copy when the producer's placement resolved to Some.
721 for ©_id in &inserted {
722 if let Some(&src) = graph.node(copy_id).inputs.first()
723 && let Some(src_loc) = src_location(graph, src)
724 {
725 graph.set_target_backend(copy_id, location_to_backend_id(src_loc));
726 }
727 }
728 // Re-stamp ALL copies/moves with their SOURCE backend — graph rewrites are
729 // sticky and `stamp_plan_backends` just overwrote pre-existing ones (e.g.
730 // realize-root splices) with the pinned backend; the transfer kernel runs
731 // where the bytes come from. The freshly-inserted copies are visited too,
732 // but their input's placement resolves to the same source location, so the
733 // sweep re-applies the identical stamp (idempotent). Copies whose source
734 // placement is absent keep whatever stamp they already have. The two-hop
735 // CPU intermediate's source is the original GPU producer (stamped to that
736 // GPU); the consumer-side hop's source is the CPU intermediate (stamped
737 // Cpu via the `graph.placement` fallback in `src_location`).
738 // Arm-inclusive (C-0): re-walk from the arm-seeded roots so copies just
739 // inserted inside non-arm-0 arms are re-stamped too. The `eff_roots` seed
740 // set (Branch nodes are unchanged by insertion) reaches the fresh arm
741 // copies via the now-rewired arm-input edges.
742 let order = topo_order_multi(graph, &eff_roots);
743 for &id in &order {
744 if !matches!(graph.node(id).op, Op::Copy { .. } | Op::Move { .. }) {
745 continue;
746 }
747 let Some(&src) = graph.node(id).inputs.first() else {
748 continue;
749 };
750 let Some(src_loc) = src_location(graph, src) else {
751 continue;
752 };
753 graph.set_target_backend(id, location_to_backend_id(src_loc));
754 }
755
756 placements
757}
758
759/// Compute every reachable node's *effective placement* for the residency
760/// pass, mirroring the bridge's old `effective_placements`. Priority per node:
761///
762/// 1. **Residency-declaring ops** — `Op::Copy`/`Op::Move`/`Op::Alloc` carry
763/// their output location in the variant (definitional).
764/// 2. **Explicit `Graph::placement`** — set by inserted copies (and any
765/// caller-provided placement hints).
766/// 3. **Input residency** — where a persistent / const input already lives,
767/// supplied by the realize path via `PlanOptions::input_residency`. This is
768/// the only runtime fact, read through the provider closure — the optimizer
769/// never touches the storage itself.
770/// 4. **Plan winner** — a planned node runs on its winner's device.
771/// 5. **Backend stamp** — a node with `target_backend` but no plan entry
772/// (structural ops) follows the pinned device.
773/// 6. **View pass-throughs** — view ops / `Reshape` / `Contiguize` follow
774/// their data input (already resolved; `order` is topological).
775fn effective_placements(
776 g: &Graph,
777 roots: &[NodeId],
778 plan: &ExecutionPlan,
779 pinned_loc: DeviceLocation,
780 input_residency: Option<&dyn Fn(NodeId) -> Option<DeviceLocation>>,
781) -> HashMap<NodeId, DeviceLocation> {
782 let order = topo_order_multi(g, roots);
783 let mut map: HashMap<NodeId, DeviceLocation> = HashMap::with_capacity(order.len());
784 for &id in &order {
785 let node = g.node(id);
786 match node.op {
787 Op::Copy { target } | Op::Move { target } | Op::Alloc { target } => {
788 map.insert(id, target);
789 continue;
790 }
791 _ => {}
792 }
793 if let Some(loc) = g.placement(id) {
794 map.insert(id, loc);
795 continue;
796 }
797 if let Some(loc) = input_residency.and_then(|f| f(id)) {
798 map.insert(id, loc);
799 continue;
800 }
801 if let Some(winner) = plan.alternatives(id).and_then(|set| set.winner()) {
802 map.insert(id, winner.device);
803 continue;
804 }
805 if g.target_backend(id).is_some() {
806 map.insert(id, pinned_loc);
807 continue;
808 }
809 if (node.op.is_view_op() || matches!(node.op, Op::Reshape(_) | Op::Contiguize))
810 && let Some(&loc) = node.inputs.first().and_then(|i| map.get(i))
811 {
812 map.insert(id, loc);
813 }
814 }
815 map
816}
817
818fn location_to_backend_id(loc: DeviceLocation) -> BackendId {
819 match loc {
820 DeviceLocation::Cpu => BackendId::Cpu,
821 DeviceLocation::Cuda { .. } => BackendId::Cuda,
822 DeviceLocation::Vulkan { .. } => BackendId::Vulkan,
823 DeviceLocation::Metal { .. } => BackendId::Metal,
824 }
825}
826
827/// The PR-A4 deliberate-fork pathfinder + the PR-B2 frontier
828/// convergence optimizer now live in [`crate::driver`] as registered
829/// [`crate::driver::Pathfinder`] / [`crate::driver::Optimizer`] impls
830/// ([`crate::driver::PlacementForkPathfinder`] /
831/// [`crate::driver::FrontierConvergenceOptimizer`]). PR-B3 replaced the
832/// hardcoded `seed_placement_fork_branches(...)` call in
833/// [`optimize_graph`] with a lock-step
834/// [`crate::driver::PassRegistry::run_lockstep`] drive over those
835/// registered passes (see the module docs). The pathfinder body is
836/// unchanged — it moved verbatim behind the trait.
837#[cfg(test)]
838mod tests {
839 use super::*;
840 use crate::fused::PrecisionGuarantee;
841 use crate::kernel::{KernelCaps, OpParams, unknown_cost};
842 use fuel_graph::opt::execution_plan;
843 use fuel_graph::{Node, Op};
844 use fuel_ir::dispatch::OpKind;
845 use fuel_ir::probe::BackendId;
846 use fuel_ir::{DType, DeviceLocation, Layout, Result as FuelResult, Shape};
847 use fuel_memory::Storage;
848 use std::sync::{Arc, RwLock};
849
850 fn noop_kernel(
851 _i: &[Arc<RwLock<Storage>>],
852 _o: &mut [Arc<RwLock<Storage>>],
853 _l: &[Layout],
854 _p: &OpParams,
855 ) -> FuelResult<()> {
856 Ok(())
857 }
858
859 fn register_elementwise(table: &mut KernelBindingTable, op: OpKind, n_in: usize) {
860 let mut dtypes = vec![DType::F32; n_in];
861 dtypes.push(DType::F32); // output dtype
862 table.register_full(
863 op,
864 &dtypes,
865 BackendId::Cpu,
866 noop_kernel,
867 KernelCaps::empty(),
868 PrecisionGuarantee::PRIMITIVE_DETERMINISTIC_CPU,
869 unknown_cost,
870 );
871 }
872
873 fn f32_node(g: &mut Graph, op: Op, inputs: Vec<NodeId>) -> NodeId {
874 // No pre-stamp: post-Step-A `optimize_graph` writes `target_backend`
875 // itself, so pre-stamping here is redundant for the tests that run the
876 // optimizer — and it would falsify the `*_stamps_backends_*` tests'
877 // "unstamped before optimize" / "Const leaf not stamped" preconditions.
878 g.push(Node {
879 op,
880 inputs,
881 shape: Shape::from_dims(&[4]),
882 dtype: DType::F32,
883 })
884 }
885
886 /// A representative branchless CPU graph the optimizer can fully
887 /// place: a fan-in (`Add` over two unary chains) feeding a unary
888 /// tail. Every kernel-bearing op has a CPU binding registered. The
889 /// fan-in exercises the run extractor's multi-predecessor boundary,
890 /// so the equivalence holds across more than a single straight
891 /// chain.
892 fn build_branchless_graph(table: &mut KernelBindingTable) -> (Graph, NodeId) {
893 register_elementwise(table, OpKind::ReluElementwise, 1);
894 register_elementwise(table, OpKind::SiluElementwise, 1);
895 register_elementwise(table, OpKind::AddElementwise, 2);
896 register_elementwise(table, OpKind::TanhElementwise, 1);
897
898 let mut g = Graph::new();
899 let a = f32_node(&mut g, Op::Const, vec![]);
900 let a1 = f32_node(&mut g, Op::Relu, vec![a]);
901 let b = f32_node(&mut g, Op::Const, vec![]);
902 let b1 = f32_node(&mut g, Op::Silu, vec![b]);
903 let sum = f32_node(&mut g, Op::Add, vec![a1, b1]);
904 let out = f32_node(&mut g, Op::Tanh, vec![sum]);
905 (g, out)
906 }
907
908 /// A pure straight-line CPU graph — the simplest no-competing-route
909 /// case (exactly one run by construction).
910 fn build_straight_line_graph(table: &mut KernelBindingTable) -> (Graph, NodeId) {
911 register_elementwise(table, OpKind::ReluElementwise, 1);
912 register_elementwise(table, OpKind::SiluElementwise, 1);
913 register_elementwise(table, OpKind::TanhElementwise, 1);
914
915 let mut g = Graph::new();
916 let a = f32_node(&mut g, Op::Const, vec![]);
917 let b = f32_node(&mut g, Op::Relu, vec![a]);
918 let c = f32_node(&mut g, Op::Silu, vec![b]);
919 let d = f32_node(&mut g, Op::Tanh, vec![c]);
920 (g, d)
921 }
922
923 fn cpu_opts() -> PlanOptions<'static> {
924 // No cost machinery: A3a only needs placement/validation, not
925 // ranking. Pin the realize device so every node resolves.
926 PlanOptions::new()
927 .without_cost_population()
928 .with_pinned_device(DeviceLocation::Cpu)
929 }
930
931 /// Cleanup Step A: `optimize_graph` stamps each computational node's
932 /// chosen backend onto the graph (the optimizer writes its placement
933 /// decision INTO the graph). On a pinned-CPU straight-line graph every
934 /// kernel-bearing node is stamped `Cpu`; the leaf `Op::Const` is
935 /// skipped (it inherits / needs no stamp). When `pinned_device` is
936 /// unset, no stamping occurs (the bare-graph test contract).
937 #[test]
938 fn optimize_graph_stamps_backends_onto_the_graph() {
939 let mut table = KernelBindingTable::new();
940 let (mut g, root) = build_straight_line_graph(&mut table);
941 // a=Const(0), b=Relu(1), c=Silu(2), d=Tanh(3) by construction.
942 let (a, b, c, d) = (NodeId(0), NodeId(1), NodeId(2), NodeId(3));
943 assert!(g.target_backend(b).is_none(), "unstamped before optimize");
944
945 let _optimized = optimize_graph(&mut g, &[root], &table, &cpu_opts())
946 .expect("optimize_graph on a straight-line CPU graph");
947
948 assert_eq!(
949 g.target_backend(b),
950 Some(BackendId::Cpu),
951 "Relu stamped Cpu"
952 );
953 assert_eq!(
954 g.target_backend(c),
955 Some(BackendId::Cpu),
956 "Silu stamped Cpu"
957 );
958 assert_eq!(
959 g.target_backend(d),
960 Some(BackendId::Cpu),
961 "Tanh stamped Cpu"
962 );
963 assert!(
964 g.target_backend(a).is_none(),
965 "Op::Const leaf is not stamped"
966 );
967 }
968
969 /// `OptimizedGraph::placement_of` makes the placement DECISION observable,
970 /// so a consumer can check where work landed instead of inferring it from
971 /// timing (or from `Graph::placement`, which is only level 2 of the six-way
972 /// resolution and returns `None` for exactly the compute nodes that matter).
973 ///
974 /// Two halves, and the second is the **positive control**: an accessor that
975 /// silently returned `None` for everything would satisfy any test that only
976 /// asserted `is_none()` on the unpinned path. So the pinned path asserts
977 /// concrete `Some(Cpu)` values AND a non-empty map — if placement were never
978 /// captured, that half fails loudly rather than passing vacuously.
979 #[test]
980 fn placement_of_observes_the_decision_and_marks_not_computed() {
981 // --- pinned: the residency pass runs, placements are captured ---
982 let mut table = KernelBindingTable::new();
983 let (mut g, root) = build_straight_line_graph(&mut table);
984 let (_a, b, c, d) = (NodeId(0), NodeId(1), NodeId(2), NodeId(3));
985
986 let optimized = optimize_graph(&mut g, &[root], &table, &cpu_opts())
987 .expect("optimize_graph on a straight-line CPU graph");
988
989 assert!(
990 optimized.has_placements(),
991 "a pinned device runs the residency pass, so placements must be computed",
992 );
993 // POSITIVE CONTROL: concrete values, not merely "not None". An empty or
994 // never-populated map would fail here instead of passing quietly.
995 assert_eq!(
996 optimized.placement_of(b),
997 Some(DeviceLocation::Cpu),
998 "Relu placed Cpu"
999 );
1000 assert_eq!(
1001 optimized.placement_of(c),
1002 Some(DeviceLocation::Cpu),
1003 "Silu placed Cpu"
1004 );
1005 assert_eq!(
1006 optimized.placement_of(d),
1007 Some(DeviceLocation::Cpu),
1008 "Tanh placed Cpu"
1009 );
1010 assert!(
1011 optimized.placements().count() >= 3,
1012 "the map must actually be populated, got {} entries",
1013 optimized.placements().count(),
1014 );
1015
1016 // --- unpinned: the pass is skipped, and that is DISTINGUISHABLE ---
1017 // This is the property that keeps a `None` from being misread as
1018 // "placed nowhere" — the failure mode this accessor exists to prevent.
1019 //
1020 // Device context has to come from SOMEWHERE or `compile_plan` refuses
1021 // outright ("node has no device context — set PlanOptions::pinned_device,
1022 // a graph placement, or the node's target_backend"). So supply it via
1023 // per-node graph placement instead of a pin: the graph optimizes fine,
1024 // but `pinned_device` is unset so the residency pass never runs — which
1025 // is exactly the "optimized successfully, placement not computed" state
1026 // a caller must be able to tell apart from "placed nowhere".
1027 let mut table2 = KernelBindingTable::new();
1028 let (mut g2, root2) = build_straight_line_graph(&mut table2);
1029 for n in [NodeId(0), NodeId(1), NodeId(2), NodeId(3)] {
1030 g2.set_placement(n, DeviceLocation::Cpu);
1031 }
1032 let bare = PlanOptions::new().without_cost_population();
1033 let unpinned = optimize_graph(&mut g2, &[root2], &table2, &bare)
1034 .expect("optimize_graph with graph placements but no pinned device");
1035
1036 assert!(
1037 !unpinned.has_placements(),
1038 "no pinned device => residency pass skipped => placements NOT computed",
1039 );
1040 assert_eq!(
1041 unpinned.placement_of(b),
1042 None,
1043 "and every lookup is None — which `has_placements()` explains as \
1044 'not computed', never 'placed nowhere'",
1045 );
1046 }
1047
1048 /// Dtype-reconciliation: a mixed-precision `[F32, BF16, F32]` matmul (F32
1049 /// activation × BF16 weight) — which no compiled-in backend serves natively in
1050 /// a CPU-only build — is made runnable by inserting a `Cast(BF16→F32)` on the
1051 /// weight, so the matmul becomes uniform `[F32, F32, F32]` (which CPU serves).
1052 /// Born-red: before this pass, `optimize_graph` failed with `NoBackendForOp`
1053 /// (the exact llama-lazy failure). This is the "f32 matmul after a Cast"
1054 /// fallback the design assumed but never built.
1055 #[test]
1056 fn optimize_graph_reconciles_mixed_precision_matmul_via_cast() {
1057 let bindings = crate::dispatch::global_bindings(); // real populated CPU table
1058 let mut g = Graph::new();
1059 let lhs = g.push(Node {
1060 op: Op::Const,
1061 inputs: vec![],
1062 shape: Shape::from_dims(&[2, 3]),
1063 dtype: DType::F32,
1064 });
1065 let rhs = g.push(Node {
1066 op: Op::Const,
1067 inputs: vec![],
1068 shape: Shape::from_dims(&[3, 2]),
1069 dtype: DType::BF16,
1070 });
1071 let mm = g.push(Node {
1072 op: Op::MatMul,
1073 inputs: vec![lhs, rhs],
1074 shape: Shape::from_dims(&[2, 2]),
1075 dtype: DType::F32,
1076 });
1077
1078 // Born-red baseline: without the fixup this errors with NoBackendForOp.
1079 optimize_graph(&mut g, &[mm], &bindings, &cpu_opts())
1080 .expect("mixed-precision matmul reconciled + optimized on CPU");
1081
1082 // The matmul's rhs is now a Cast(F32) of the original BF16 weight; lhs
1083 // (already F32) is untouched.
1084 let mm_inputs = g.node(mm).inputs.clone();
1085 assert_eq!(mm_inputs[0], lhs, "lhs (already F32) untouched");
1086 let cast_id = mm_inputs[1];
1087 assert_ne!(cast_id, rhs, "rhs was rewired to a Cast");
1088 assert!(
1089 matches!(g.node(cast_id).op, Op::Cast(DType::F32)),
1090 "rhs cast to F32"
1091 );
1092 assert_eq!(g.node(cast_id).dtype, DType::F32);
1093 assert_eq!(
1094 g.node(cast_id).inputs,
1095 vec![rhs],
1096 "cast reads the original BF16 weight"
1097 );
1098 }
1099
1100 fn register_matmul(table: &mut KernelBindingTable, dtypes: &[DType], backend: BackendId) {
1101 table.register_full(
1102 OpKind::MatMul,
1103 dtypes,
1104 backend,
1105 noop_kernel,
1106 KernelCaps::empty(),
1107 PrecisionGuarantee::PRIMITIVE_DETERMINISTIC_CPU,
1108 unknown_cost,
1109 );
1110 }
1111
1112 /// The reconciling `Op::Cast` the pass inserts needs a binding of its own to
1113 /// plan. The production table has these (`(Cpu, Cast, [BF16, F32])`); a
1114 /// hand-built test table must register them explicitly.
1115 fn register_cast(table: &mut KernelBindingTable, from: DType, to: DType, backend: BackendId) {
1116 table.register_full(
1117 OpKind::Cast,
1118 &[from, to],
1119 backend,
1120 noop_kernel,
1121 KernelCaps::empty(),
1122 PrecisionGuarantee::PRIMITIVE_DETERMINISTIC_CPU,
1123 unknown_cost,
1124 );
1125 }
1126
1127 /// `[F32] @ [BF16] -> [F32]` on a graph pinned to `pinned`.
1128 fn mixed_matmul_graph() -> (Graph, NodeId, NodeId, NodeId) {
1129 let mut g = Graph::new();
1130 let lhs = g.push(Node {
1131 op: Op::Const,
1132 inputs: vec![],
1133 shape: Shape::from_dims(&[2, 3]),
1134 dtype: DType::F32,
1135 });
1136 let rhs = g.push(Node {
1137 op: Op::Const,
1138 inputs: vec![],
1139 shape: Shape::from_dims(&[3, 2]),
1140 dtype: DType::BF16,
1141 });
1142 let mm = g.push(Node {
1143 op: Op::MatMul,
1144 inputs: vec![lhs, rhs],
1145 shape: Shape::from_dims(&[2, 2]),
1146 dtype: DType::F32,
1147 });
1148 (g, lhs, rhs, mm)
1149 }
1150
1151 /// **Born-red regression guard: a backend must never break a graph that
1152 /// isn't running on it.**
1153 ///
1154 /// `insert_dtype_fixups` used to guard on `any_backend_supports` — a
1155 /// predicate over the WHOLE binding table that ignores the key's backend
1156 /// component. `compile_plan` enumerates over the node's *candidate
1157 /// placement set*. So when a non-candidate backend owned the only native
1158 /// mixed-dtype kernel, the pass saw "someone serves it", declined to
1159 /// reconcile, and the planner then dead-ended at `NoBackendForOp`.
1160 ///
1161 /// In-tree that surfaced as `--features vulkan` breaking CPU-pinned
1162 /// mixed-precision matmul, because Vulkan happens to be the only backend
1163 /// registering `MatMul[F32, BF16, F32]`. This test builds that situation
1164 /// from a LOCAL table so it reproduces in **every** build configuration —
1165 /// the feature-gated version is why the defect shipped unnoticed.
1166 #[test]
1167 fn dtype_reconciliation_ignores_kernels_on_non_candidate_backends() {
1168 let mut table = KernelBindingTable::new();
1169 // The pinned backend serves only the uniform key...
1170 register_matmul(
1171 &mut table,
1172 &[DType::F32, DType::F32, DType::F32],
1173 BackendId::Cpu,
1174 );
1175 // ...while a backend this graph will never run on owns the mixed key.
1176 register_matmul(
1177 &mut table,
1178 &[DType::F32, DType::BF16, DType::F32],
1179 BackendId::Vulkan,
1180 );
1181 register_cast(&mut table, DType::BF16, DType::F32, BackendId::Cpu);
1182
1183 let (mut g, lhs, rhs, mm) = mixed_matmul_graph();
1184 optimize_graph(&mut g, &[mm], &table, &cpu_opts()).expect(
1185 "CPU-pinned mixed-precision matmul must reconcile via Cast even though \
1186 a NON-candidate backend owns the native mixed kernel",
1187 );
1188
1189 let mm_inputs = g.node(mm).inputs.clone();
1190 assert_eq!(mm_inputs[0], lhs, "lhs (already F32) untouched");
1191 let cast_id = mm_inputs[1];
1192 assert_ne!(cast_id, rhs, "rhs was rewired to a Cast");
1193 assert!(
1194 matches!(g.node(cast_id).op, Op::Cast(DType::F32)),
1195 "rhs cast to F32"
1196 );
1197 assert_eq!(
1198 g.node(cast_id).inputs,
1199 vec![rhs],
1200 "cast reads the original BF16 weight"
1201 );
1202 }
1203
1204 /// The converse guard — the fix must not over-cast. When the node's own
1205 /// candidate backend serves the mixed key natively, reconciliation must
1206 /// decline, so the native mixed-precision kernel is preserved (casting
1207 /// would change numerics and double the promoted operand's bytes).
1208 #[test]
1209 fn dtype_reconciliation_declines_when_the_candidate_backend_serves_it_natively() {
1210 let mut table = KernelBindingTable::new();
1211 register_matmul(
1212 &mut table,
1213 &[DType::F32, DType::BF16, DType::F32],
1214 BackendId::Cpu,
1215 );
1216
1217 let (mut g, lhs, rhs, mm) = mixed_matmul_graph();
1218 optimize_graph(&mut g, &[mm], &table, &cpu_opts()).expect("native mixed kernel plans");
1219
1220 assert_eq!(
1221 g.node(mm).inputs,
1222 vec![lhs, rhs],
1223 "no Cast: the candidate backend has the native mixed kernel",
1224 );
1225 }
1226
1227 /// Multi-device behavior must be PRESERVED, not sacrificed to fix the bug
1228 /// above. When an off-device fallback is admitted, the foreign backend IS
1229 /// a candidate for this node — so its native mixed kernel is reachable and
1230 /// reconciliation must still decline. This is why the fix mirrors the
1231 /// planner's candidate set rather than narrowing to the pinned device.
1232 #[test]
1233 fn dtype_reconciliation_counts_admitted_off_device_fallback_as_a_candidate() {
1234 let mut table = KernelBindingTable::new();
1235 register_matmul(
1236 &mut table,
1237 &[DType::F32, DType::F32, DType::F32],
1238 BackendId::Cpu,
1239 );
1240 register_matmul(
1241 &mut table,
1242 &[DType::F32, DType::BF16, DType::F32],
1243 BackendId::Vulkan,
1244 );
1245
1246 let fallback = |_from: DeviceLocation| -> Vec<(BackendId, DeviceLocation)> {
1247 vec![(BackendId::Vulkan, DeviceLocation::Vulkan { gpu_id: 0 })]
1248 };
1249 let opts = PlanOptions::new()
1250 .without_cost_population()
1251 .with_pinned_device(DeviceLocation::Cpu)
1252 .with_fallback_placements_for(&fallback);
1253
1254 let (mut g, lhs, rhs, mm) = mixed_matmul_graph();
1255 optimize_graph(&mut g, &[mm], &table, &opts).expect("off-device native mixed kernel plans");
1256
1257 assert_eq!(
1258 g.node(mm).inputs,
1259 vec![lhs, rhs],
1260 "no Cast: the admitted off-device fallback backend serves the mixed key natively",
1261 );
1262 }
1263
1264 /// The stamping is guarded on a pinned device — an `optimize_graph`
1265 /// with no `pinned_device` leaves the graph unstamped (so the bare
1266 /// optimize.rs test callers below are unaffected by Step A).
1267 #[test]
1268 fn optimize_graph_without_pinned_device_does_not_stamp() {
1269 let mut table = KernelBindingTable::new();
1270 let (mut g, root) = build_straight_line_graph(&mut table);
1271 // Give each node a device context via a graph PLACEMENT (a distinct
1272 // side-table from `target_backend`), so `compile_plan` resolves without
1273 // a pinned device — then assert `optimize_graph` performs no
1274 // `target_backend` stamping (the stamping arm is guarded on
1275 // `pinned_device`, so without a pin nothing is stamped).
1276 for id in [NodeId(0), NodeId(1), NodeId(2), NodeId(3)] {
1277 g.set_placement(id, DeviceLocation::Cpu);
1278 }
1279 let opts = PlanOptions::new().without_cost_population(); // no pinned device
1280 let _ = optimize_graph(&mut g, &[root], &table, &opts).expect("optimize_graph");
1281 assert!(
1282 g.target_backend(NodeId(1)).is_none(),
1283 "no pinned device ⇒ no target_backend stamping",
1284 );
1285 }
1286
1287 /// Cleanup Step B (layout): `optimize_graph` inserts an `Op::Contiguize`
1288 /// before a kernel whose chosen winner rejects strided inputs
1289 /// (`caps.strided_input == false`, e.g. the `KernelCaps::empty()`
1290 /// elementwise bindings) and whose input layout is non-contiguous — the
1291 /// layout-fixup pass moved from the realize-time bridge into the optimizer.
1292 /// The pass itself is exhaustively tested in `fuel-graph::opt`; this guards
1293 /// the optimizer-side wiring (callback reads the plan winner's caps, runs
1294 /// inside the pinned guard).
1295 #[test]
1296 fn optimize_graph_inserts_layout_fixup_for_strided_input() {
1297 let mut table = KernelBindingTable::new();
1298 register_elementwise(&mut table, OpKind::ReluElementwise, 1); // empty caps ⇒ rejects strided
1299 let mut g = Graph::new();
1300 let a = f32_node(&mut g, Op::Const, vec![]);
1301 let r = f32_node(&mut g, Op::Relu, vec![a]);
1302 // Make `a`'s effective layout non-contiguous (non-zero start_offset),
1303 // so the strided-rejecting Relu needs a Contiguize on its input.
1304 g.set_layout(a, Layout::contiguous_with_offset(Shape::from_dims(&[4]), 1));
1305
1306 let _optimized = optimize_graph(&mut g, &[r], &table, &cpu_opts())
1307 .expect("optimize_graph with a strided input");
1308
1309 // Relu's input was `a`; after the fixup it must read through an inserted
1310 // Op::Contiguize wrapping `a`.
1311 let relu_input = g.node(r).inputs[0];
1312 assert!(
1313 matches!(g.node(relu_input).op, Op::Contiguize),
1314 "strided input to a strided-rejecting kernel must be fixed up via \
1315 Op::Contiguize, got {:?}",
1316 g.node(relu_input).op,
1317 );
1318 assert_eq!(
1319 g.node(relu_input).inputs,
1320 vec![a],
1321 "the inserted Op::Contiguize wraps the original strided input",
1322 );
1323 }
1324
1325 // ---- Cleanup Step B (residency) — insert_residency_copies tests ----
1326 // (migrated from the bridge's insert_resident_input_copies unit tests;
1327 // cache residency is supplied via an `input_residency` closure instead of a
1328 // live StorageCache. Placement metadata only — no GPU needed; the unprobed
1329 // test `SystemTopology` reports Cpu/Cuda as non-sharing, so a crossing is
1330 // detected and a copy inserted.)
1331
1332 /// Co-located graph ⇒ `insert_residency_copies` is a no-op (no crossings).
1333 #[test]
1334 fn residency_noop_when_colocated() {
1335 let mut g = Graph::new();
1336 let c1 = f32_node(&mut g, Op::Const, vec![]);
1337 let c2 = f32_node(&mut g, Op::Const, vec![]);
1338 let add = f32_node(&mut g, Op::Add, vec![c1, c2]);
1339 g.set_target_backend(add, BackendId::Cpu);
1340 let residency = |id: NodeId| (id == c1 || id == c2).then_some(DeviceLocation::Cpu);
1341 let pre = g.len();
1342 insert_residency_copies(
1343 &mut g,
1344 &[add],
1345 &ExecutionPlan::empty(),
1346 DeviceLocation::Cpu,
1347 Some(&residency),
1348 );
1349 assert_eq!(g.len(), pre, "co-located graph must be a no-op");
1350 assert_eq!(g.node(add).inputs, vec![c1, c2], "edges untouched");
1351 }
1352
1353 /// A CPU-resident input feeding two CUDA-pinned consumers ⇒ exactly ONE
1354 /// `Op::Copy` bridges the crossing (CSE-deduped), targeting the consumer
1355 /// device, stamped `target_backend = SOURCE` (Cpu, the H2D wrapper) with
1356 /// its output placed on the consumer device.
1357 #[test]
1358 fn residency_one_copy_per_crossing_deduped() {
1359 let cuda0 = DeviceLocation::Cuda { gpu_id: 0 };
1360 let mut g = Graph::new();
1361 let c1 = f32_node(&mut g, Op::Const, vec![]);
1362 let neg = f32_node(&mut g, Op::Neg, vec![c1]);
1363 let sqr = f32_node(&mut g, Op::Sqr, vec![c1]);
1364 g.set_target_backend(neg, BackendId::Cuda);
1365 g.set_target_backend(sqr, BackendId::Cuda);
1366 let residency = |id: NodeId| (id == c1).then_some(DeviceLocation::Cpu);
1367
1368 let pre = g.len();
1369 insert_residency_copies(
1370 &mut g,
1371 &[neg, sqr],
1372 &ExecutionPlan::empty(),
1373 cuda0,
1374 Some(&residency),
1375 );
1376
1377 assert_eq!(g.len(), pre + 1, "one crossing → one copy, CSE-deduped");
1378 let neg_in = g.node(neg).inputs[0];
1379 let sqr_in = g.node(sqr).inputs[0];
1380 assert_eq!(neg_in, sqr_in, "both consumers share the one copy");
1381 assert_ne!(neg_in, c1, "consumers rewired off the raw input");
1382 let copy = g.node(neg_in);
1383 assert!(
1384 matches!(copy.op, Op::Copy { target } if target == cuda0),
1385 "copy targets the consumer device; got {:?}",
1386 copy.op,
1387 );
1388 assert_eq!(copy.inputs, vec![c1], "copy reads the resident slot");
1389 assert_eq!(
1390 g.target_backend(neg_in),
1391 Some(BackendId::Cpu),
1392 "stamped with the SOURCE backend (H2D runs on the CPU wrapper)",
1393 );
1394 assert_eq!(
1395 g.placement(neg_in),
1396 Some(cuda0),
1397 "copy output placed on the consumer device",
1398 );
1399 }
1400
1401 /// Idempotence + re-stamp: a second pass on the rewritten graph inserts
1402 /// nothing and restores the source-backend stamp even after a clobber
1403 /// (mimicking `stamp_plan_backends` overwriting it with the pinned backend).
1404 #[test]
1405 fn residency_idempotent_and_restamped_on_second_call() {
1406 let cuda0 = DeviceLocation::Cuda { gpu_id: 0 };
1407 let mut g = Graph::new();
1408 let c1 = f32_node(&mut g, Op::Const, vec![]);
1409 let neg = f32_node(&mut g, Op::Neg, vec![c1]);
1410 g.set_target_backend(neg, BackendId::Cuda);
1411 let residency = |id: NodeId| (id == c1).then_some(DeviceLocation::Cpu);
1412
1413 insert_residency_copies(
1414 &mut g,
1415 &[neg],
1416 &ExecutionPlan::empty(),
1417 cuda0,
1418 Some(&residency),
1419 );
1420 let copy_id = g.node(neg).inputs[0];
1421 // Simulate stamp_plan_backends clobbering the copy with the pinned backend.
1422 g.set_target_backend(copy_id, BackendId::Cuda);
1423
1424 let pre = g.len();
1425 insert_residency_copies(
1426 &mut g,
1427 &[neg],
1428 &ExecutionPlan::empty(),
1429 cuda0,
1430 Some(&residency),
1431 );
1432 assert_eq!(g.len(), pre, "re-run inserts nothing");
1433 assert_eq!(
1434 g.target_backend(copy_id),
1435 Some(BackendId::Cpu),
1436 "re-stamp sweep restores the source-backend stamp",
1437 );
1438 }
1439
1440 // ---- Step E Phase C, PR C-0: arm residency-copy audit ----
1441 //
1442 // C2 lets the executor re-pick an `Op::Branch` arm at runtime by live
1443 // device load — it may choose arm-1+ instead of the static arm-0 winner.
1444 // That is only legal if the optimizer's residency pass has ALREADY
1445 // stitched the inbound cross-device `Op::Copy` for EVERY surviving arm,
1446 // so whichever arm the executor picks, that arm's device-inputs are
1447 // resident on its device. This test builds a 2-arm diamond whose two arms
1448 // live on DIFFERENT devices (arm-0 CPU, arm-1 CUDA) sharing one CPU
1449 // producer, runs the residency pass, and asserts BOTH arms' cross-device
1450 // inputs got their copy — not just arm-0's.
1451 //
1452 // The finalized `Op::Branch` is orphaned (its `reconverge_at` reads arm-0
1453 // directly, per the PR-A1 runnability invariant), so arm-1 is reachable
1454 // ONLY through the Branch node. `insert_residency_copies` walks
1455 // `topo_order_multi(roots)` directly (NOT `effective_roots`), so it never
1456 // reaches an orphaned Branch — and therefore never reaches arm-1. The
1457 // run/route machinery (`extract_runs_multi`, `branches_in_topo_order`)
1458 // does use `effective_roots`; the residency pass does not. This is the
1459 // C-0 prerequisite gap.
1460
1461 /// Build a 2-arm diamond: a CPU producer `diverge` fans into arm-0 (CPU)
1462 /// and arm-1 (`arm1_loc`); `reconverge` reads arm-0 (the runnability
1463 /// invariant); the Branch merges {arm0, arm1} at `reconverge`. Returns
1464 /// `(graph, branch, diverge, arm0, arm1, post)`. Placements are set via
1465 /// `set_placement` so the residency pass reads them with no plan/winner.
1466 fn residency_diamond(
1467 arm1_loc: DeviceLocation,
1468 ) -> (Graph, NodeId, NodeId, NodeId, NodeId, NodeId) {
1469 let mut g = Graph::new();
1470 let pre = f32_node(&mut g, Op::Const, vec![]);
1471 let diverge = f32_node(&mut g, Op::Relu, vec![pre]);
1472 let arm0 = f32_node(&mut g, Op::Silu, vec![diverge]);
1473 let arm1 = f32_node(&mut g, Op::GeluTanh, vec![diverge]);
1474 let reconverge = f32_node(&mut g, Op::Relu, vec![arm0]);
1475 let mut b = g.open_branch(diverge);
1476 b.add_arm(arm0);
1477 b.add_arm(arm1);
1478 let branch = b
1479 .finalize_branches(&mut g, reconverge)
1480 .expect("well-formed 2-arm branch")
1481 .expect("2 arms survive");
1482 let post = f32_node(&mut g, Op::Tanh, vec![reconverge]);
1483 // Producer + arm-0 + the post-merge region on CPU; arm-1 elsewhere.
1484 g.set_placement(pre, DeviceLocation::Cpu);
1485 g.set_placement(diverge, DeviceLocation::Cpu);
1486 g.set_placement(arm0, DeviceLocation::Cpu);
1487 g.set_placement(reconverge, DeviceLocation::Cpu);
1488 g.set_placement(post, DeviceLocation::Cpu);
1489 g.set_placement(arm1, arm1_loc);
1490 (g, branch, diverge, arm0, arm1, post)
1491 }
1492
1493 /// C-0 PREREQUISITE: the residency pass must insert the inbound
1494 /// cross-device `Op::Copy` for BOTH arms of a 2-device branch — arm-0
1495 /// (CPU, same device as the producer ⇒ no copy needed) AND arm-1 (CUDA,
1496 /// crosses ⇒ MUST get a copy). If only arm-0 were stitched, re-picking
1497 /// arm-1 at runtime (C2) would dispatch a CUDA kernel whose input was
1498 /// never copied to the GPU → the un-bridged-mixed-edge error / a UAF.
1499 #[test]
1500 fn residency_stitches_every_branch_arm_not_just_arm0() {
1501 let cuda0 = DeviceLocation::Cuda { gpu_id: 0 };
1502 let (mut g, _branch, diverge, arm0, arm1, post) = residency_diamond(cuda0);
1503
1504 insert_residency_copies(
1505 &mut g,
1506 &[post],
1507 &ExecutionPlan::empty(),
1508 DeviceLocation::Cpu,
1509 None,
1510 );
1511
1512 // arm-0 is CPU, same substrate as the CPU producer ⇒ its input edge
1513 // does NOT cross ⇒ no copy (and must stay reading `diverge`).
1514 assert_eq!(
1515 g.node(arm0).inputs,
1516 vec![diverge],
1517 "arm-0 is co-located with its producer ⇒ no copy, edge untouched",
1518 );
1519
1520 // arm-1 is CUDA, its producer `diverge` is CPU ⇒ the edge crosses ⇒
1521 // the residency pass MUST have inserted an Op::Copy targeting CUDA on
1522 // arm-1's input. THIS is the C-0 invariant: every surviving arm's
1523 // device-inputs are made resident, not just arm-0's.
1524 let arm1_in = g.node(arm1).inputs[0];
1525 assert_ne!(
1526 arm1_in, diverge,
1527 "arm-1 (CUDA) reads the CPU producer directly — its inbound \
1528 cross-device Op::Copy was NEVER inserted. Re-picking arm-1 at \
1529 runtime (C2) would dispatch a CUDA kernel over un-copied CPU \
1530 bytes. The residency pass stitched only arm-0's route.",
1531 );
1532 let copy = g.node(arm1_in);
1533 assert!(
1534 matches!(copy.op, Op::Copy { target } if target == cuda0),
1535 "arm-1's inbound copy must target CUDA; got {:?}",
1536 copy.op,
1537 );
1538 assert_eq!(
1539 copy.inputs,
1540 vec![diverge],
1541 "arm-1's copy reads the shared CPU producer",
1542 );
1543 }
1544
1545 // ---- PR-A4 deliberate-fork-seed test substrate ----
1546
1547 fn noop_kernel_b(
1548 _i: &[Arc<RwLock<Storage>>],
1549 _o: &mut [Arc<RwLock<Storage>>],
1550 _l: &[Layout],
1551 _p: &OpParams,
1552 ) -> FuelResult<()> {
1553 Ok(())
1554 }
1555
1556 /// Register a CPU AND a (synthetic) CUDA binding for `op` so a node
1557 /// enumerated against both placements gets two viable
1558 /// `(backend, device)` candidates — the genuine placement fork the
1559 /// A4 pathfinder seeds a branch from.
1560 fn register_two_backend(table: &mut KernelBindingTable, op: OpKind, n_in: usize) {
1561 let mut dtypes = vec![DType::F32; n_in];
1562 dtypes.push(DType::F32);
1563 table.register_full(
1564 op,
1565 &dtypes,
1566 BackendId::Cpu,
1567 noop_kernel,
1568 KernelCaps::empty(),
1569 PrecisionGuarantee::PRIMITIVE_DETERMINISTIC_CPU,
1570 unknown_cost,
1571 );
1572 table.register_full(
1573 op,
1574 &dtypes,
1575 BackendId::Cuda,
1576 noop_kernel_b,
1577 KernelCaps::empty(),
1578 PrecisionGuarantee::PRIMITIVE_DETERMINISTIC_CPU,
1579 unknown_cost,
1580 );
1581 }
1582
1583 fn two_backend_placements(dev: DeviceLocation) -> Vec<BackendId> {
1584 if dev == DeviceLocation::Cpu {
1585 vec![BackendId::Cpu, BackendId::Cuda]
1586 } else {
1587 vec![]
1588 }
1589 }
1590
1591 /// Two backends co-located at the realize device, so every
1592 /// kernel-bearing node enumerates ≥2 placements — the multi-backend
1593 /// build the A4 pathfinder forks on. CPU-only builds (one placement
1594 /// per node) never hit this; the equivalence/idempotence gates above
1595 /// cover that.
1596 fn two_backend_opts() -> PlanOptions<'static> {
1597 PlanOptions::new()
1598 .without_cost_population()
1599 .with_pinned_device(DeviceLocation::Cpu)
1600 .with_placements_for_device(&two_backend_placements)
1601 }
1602
1603 /// A graph with a single kernel-bearing node that has a genuine
1604 /// ≥2-placement choice AND exactly one consumer (so it is a
1605 /// deliberate fork, not plain fan-out). Only `fork`'s op (`Silu`) is
1606 /// registered on two backends, so it is the ONLY node with a real
1607 /// ≥2-placement choice — exactly the "most ops are CPU-only, the
1608 /// matmul has cuBLAS too" shape.
1609 ///
1610 /// A long CPU-only `Relu` body surrounds the fork so the single
1611 /// deliberate branch sits far below the fewness threshold (a tiny
1612 /// graph would be all-branch by node count): `c -> body* -> prod ->
1613 /// fork -> tail -> body*`. `prod` is the diverge; `tail` is the
1614 /// reconverge (it reads `fork`).
1615 fn build_single_fork_graph(
1616 table: &mut KernelBindingTable,
1617 ) -> (Graph, NodeId, NodeId, NodeId, NodeId) {
1618 register_elementwise(table, OpKind::ReluElementwise, 1);
1619 register_two_backend(table, OpKind::SiluElementwise, 1); // the fork
1620 register_elementwise(table, OpKind::TanhElementwise, 1);
1621
1622 let mut g = Graph::new();
1623 let mut prev = f32_node(&mut g, Op::Const, vec![]);
1624 // Straight-line CPU body before the fork.
1625 for _ in 0..20 {
1626 prev = f32_node(&mut g, Op::Relu, vec![prev]);
1627 }
1628 let prod = prev; // the diverge (fork's producer)
1629 let fork = f32_node(&mut g, Op::Silu, vec![prod]);
1630 let tail = f32_node(&mut g, Op::Tanh, vec![fork]); // reconverge
1631 // Straight-line CPU body after the reconverge — the realize root
1632 // is the body's tail, so the whole graph is reachable and the
1633 // single fork sits far below the fewness threshold.
1634 let mut prev = tail;
1635 for _ in 0..20 {
1636 prev = f32_node(&mut g, Op::Relu, vec![prev]);
1637 }
1638 let root = prev;
1639 (g, prod, fork, tail, root)
1640 }
1641
1642 /// THE A4 DELIBERATE-FORK GATE (born-red until the pathfinder
1643 /// emits a Branch).
1644 ///
1645 /// (a) A node with two viable `(backend, device)` placements and a
1646 /// single consumer ⇒ `optimize_graph` emits exactly ONE 2-arm
1647 /// `Op::Branch` whose arm-0 is the DP winner, and the result passes
1648 /// the fewness gate.
1649 #[test]
1650 fn deliberate_fork_emits_one_two_arm_branch() {
1651 use fuel_graph::{branch_density, passes_fewness_gate};
1652 let mut table = KernelBindingTable::new();
1653 let (mut g, _prod, fork, _tail, root) = build_single_fork_graph(&mut table);
1654 let opts = two_backend_opts();
1655
1656 let nodes_before = g.len();
1657 let optimized = optimize_graph(&mut g, &[root], &table, &opts)
1658 .expect("optimize_graph succeeds on a 2-placement graph");
1659
1660 // Exactly one Op::Branch in the arena.
1661 assert_eq!(
1662 optimized.branch_count(&g),
1663 1,
1664 "a single genuine placement fork emits exactly one Op::Branch",
1665 );
1666
1667 // It is a 2-arm branch (arm-0 = the winner = `fork`).
1668 let branch_id = (nodes_before..g.len())
1669 .map(NodeId)
1670 .find(|&id| matches!(g.node(id).op, Op::Branch { .. }))
1671 .expect("a Branch node was appended");
1672 let branch = g.node(branch_id);
1673 assert_eq!(branch.inputs.len(), 2, "the branch has exactly two arms");
1674 assert_eq!(
1675 branch.inputs[0], fork,
1676 "arm-0 is the DP winner node (the route realize uses today)",
1677 );
1678 assert_ne!(
1679 branch.inputs[1], fork,
1680 "arm-1 is a distinct (runner-up placement) node",
1681 );
1682
1683 // The fewness gate holds — one branch among many nodes.
1684 assert!(
1685 passes_fewness_gate(&g, root),
1686 "a single deliberate fork passes the fewness gate; density={}",
1687 branch_density(&g, root),
1688 );
1689 }
1690
1691 /// (b) An ordinary DAG fan-out (one result, two distinct consumers)
1692 /// with the SAME 2-placement freedom is NOT flagged as a branch —
1693 /// fan-out is not a decision point.
1694 #[test]
1695 fn plain_fan_out_is_not_a_branch() {
1696 let mut table = KernelBindingTable::new();
1697 // Only the fan-out node's op (`Relu`) is dual-backend, so it is
1698 // the sole ≥2-placement candidate — and it is excluded purely
1699 // because it fans out (two consumers), not because it lacks a
1700 // placement choice. The consumers are CPU-only.
1701 register_two_backend(&mut table, OpKind::ReluElementwise, 1);
1702 register_elementwise(&mut table, OpKind::SiluElementwise, 1);
1703 register_elementwise(&mut table, OpKind::TanhElementwise, 1);
1704 register_elementwise(&mut table, OpKind::AddElementwise, 2);
1705
1706 // `shared` has two distinct consumers (c0, c1) that join at
1707 // `out` — plain fan-out, not a fork — even though it has a real
1708 // 2-placement choice.
1709 let mut g = Graph::new();
1710 let c = f32_node(&mut g, Op::Const, vec![]);
1711 let shared = f32_node(&mut g, Op::Relu, vec![c]);
1712 let c0 = f32_node(&mut g, Op::Silu, vec![shared]);
1713 let c1 = f32_node(&mut g, Op::Tanh, vec![shared]);
1714 let out = f32_node(&mut g, Op::Add, vec![c0, c1]);
1715 let opts = two_backend_opts();
1716
1717 let optimized =
1718 optimize_graph(&mut g, &[out], &table, &opts).expect("optimize_graph succeeds");
1719 assert_eq!(
1720 optimized.branch_count(&g),
1721 0,
1722 "plain fan-out (multiple consumers) is not a deliberate fork",
1723 );
1724 let _ = shared;
1725 }
1726
1727 /// (c) A branched graph realizes correctly via arm-0: the dispatch
1728 /// order follows {pre, arm-0, post} and SKIPS arm-1's run.
1729 #[test]
1730 fn branched_graph_realizes_via_arm0() {
1731 let mut table = KernelBindingTable::new();
1732 let (mut g, prod, fork, tail, root) = build_single_fork_graph(&mut table);
1733 let opts = two_backend_opts();
1734
1735 // Capture the single-route order BEFORE the branch is emitted —
1736 // this is the order realize must reproduce on arm-0.
1737 let pre_order = execution_plan(&g, &[root]);
1738
1739 let optimized =
1740 optimize_graph(&mut g, &[root], &table, &opts).expect("optimize_graph succeeds");
1741 assert_eq!(optimized.branch_count(&g), 1, "exactly one branch");
1742
1743 // The arm-1 node is the branch's second input — it must NOT
1744 // appear in the arm-0 dispatch order.
1745 let branch_id = (0..g.len())
1746 .map(NodeId)
1747 .find(|&id| matches!(g.node(id).op, Op::Branch { .. }))
1748 .expect("a Branch node exists");
1749 let arm1 = g.node(branch_id).inputs[1];
1750
1751 let order = optimized.dispatch_order(&g);
1752 assert!(
1753 !order.contains(&arm1),
1754 "arm-0 lowering must skip arm-1's run; order={order:?} arm1={arm1:?}",
1755 );
1756 // arm-0 (fork) and the pre/post nodes ARE executed.
1757 assert!(order.contains(&prod), "pre-run (diverge producer) executes");
1758 assert!(order.contains(&fork), "arm-0 (the winner) executes");
1759 assert!(order.contains(&tail), "post-run (reconverge) executes");
1760 assert!(order.contains(&root), "the realize root executes");
1761
1762 // Behavior preserving: arm-0 lowering reproduces exactly the
1763 // pre-branch single-route order (arm-0 = the DP winner = the
1764 // graph realize used before the branch was recorded).
1765 assert_eq!(
1766 order, pre_order,
1767 "arm-0 dispatch order equals the pre-branch single-route order",
1768 );
1769 }
1770
1771 /// Return type of [`cached_env_test_fixture`] — named purely to dodge
1772 /// `clippy::type_complexity`, not because any caller needs the alias.
1773 type CachedEnvTestFixture = (
1774 Arc<RwLock<Graph>>,
1775 crate::pipelined::StorageCache,
1776 OptimizedGraph,
1777 std::sync::OnceLock<Vec<NodeId>>,
1778 Arc<dyn crate::ranker::RuntimeSelector>,
1779 NodeId,
1780 );
1781
1782 /// Shared setup for the two `cached_env_*` tests below: a single-fork
1783 /// graph, a two-backend `OptimizedGraph`, a seeded `StorageCache`, a
1784 /// fresh `order_cache`, and a real (non-stub) `WinnerSelector`.
1785 /// Factored out purely to keep each test under Codacy's method-length
1786 /// bound — no behavior lives here that either test's own body needs to
1787 /// see inline.
1788 fn cached_env_test_fixture() -> CachedEnvTestFixture {
1789 use crate::pipelined::StorageCache;
1790 use crate::ranker::runtime_selector::WinnerSelector;
1791
1792 let mut table = KernelBindingTable::new();
1793 let (mut g, _prod, _fork, _tail, root) = build_single_fork_graph(&mut table);
1794 let opts = two_backend_opts();
1795 let optimized =
1796 optimize_graph(&mut g, &[root], &table, &opts).expect("optimize_graph succeeds");
1797 assert_eq!(optimized.branch_count(&g), 1, "exactly one branch");
1798
1799 let graph = Arc::new(RwLock::new(g));
1800 let mut cache = StorageCache::new();
1801 cache.insert(
1802 NodeId(0), // the leaf Const build_single_fork_graph pushes first.
1803 Arc::new(RwLock::new(fuel_memory::from_slice_cpu(&[
1804 1.0_f32, 2.0, 3.0, 4.0,
1805 ]))),
1806 );
1807 let order_cache: std::sync::OnceLock<Vec<NodeId>> = std::sync::OnceLock::new();
1808 let selector: Arc<dyn crate::ranker::RuntimeSelector> = Arc::new(WinnerSelector);
1809 (graph, cache, optimized, order_cache, selector, root)
1810 }
1811
1812 /// `order_for` fix track (`docs/design/incremental-consumer-index.md`
1813 /// §5), answering the PM's review question directly: on a graph with a
1814 /// REAL `Op::Branch` and a REAL selector (so the executor's picking
1815 /// entry point resolves a genuine pick, not just "no branches to
1816 /// pick"), `realize_with_optimized_picking_cached_env` must route
1817 /// through `OrderSource::Streaming`, NEVER through `OrderSource::Cached`
1818 /// — proven here by checking `order_cache` stays EMPTY across two
1819 /// separate picking realizes, not merely by reading the routing code.
1820 /// This is the stronger property PR review asked for: the fix doesn't
1821 /// need a cache key that "covers the pick" because picked graphs never
1822 /// reach the cache in the first place — `realize_with_optimized_picking_env`'s
1823 /// OWN routing (`streaming_pick_for` returning `Some` whenever a
1824 /// selector is given over a branched graph) already excludes them, and
1825 /// this test is what proves that exclusion actually holds for the
1826 /// cached sibling too, not just the uncached original.
1827 ///
1828 /// `result`/`result2` tolerate `Error::TopologyChanged` as a passing
1829 /// outcome, not just `Ok`: `optimized.generation` is stamped once, from
1830 /// the process-wide `topology_generation()` counter, at `optimize_graph`
1831 /// time above — and that counter is shared by every test in this binary
1832 /// (`register_runtime_kernel`/`extend_global_bindings`'s own doc admits
1833 /// this; see `register_runtime_kernel_into`'s test-only escape hatch).
1834 /// A sibling test registering a runtime-fused kernel between here and
1835 /// either realize call bumps it, and `realize_inner`'s chunk-boundary
1836 /// check (pipelined.rs) deliberately does NOT retry at this layer —
1837 /// only `pipelined_bridge`'s wrapper does (by design, so a low-level
1838 /// `PipelinedExecutor` caller sees the typed error, not a silent stall).
1839 /// A real retry would need a fresh `optimize_graph` call, which would
1840 /// defeat this test's actual subject (whether the same `OptimizedGraph`
1841 /// and `order_cache` stay consistent across repeat calls), so instead
1842 /// of retrying we just don't fail on the one outcome this test was
1843 /// never about. The `order_cache`-emptiness assertions stay
1844 /// unconditional: the `Streaming` path never touches `order_cache`
1845 /// regardless of whether the realize that invoked it succeeded or hit
1846 /// this race.
1847 #[test]
1848 fn cached_env_never_caches_a_branched_graph_with_a_real_selector() {
1849 use crate::PipelinedExecutor;
1850 use fuel_ir::{Error, SymEnv};
1851
1852 let (graph, cache, optimized, order_cache, selector, root) = cached_env_test_fixture();
1853
1854 let result = PipelinedExecutor::realize_with_optimized_picking_cached_env(
1855 graph.clone(),
1856 root,
1857 cache.clone(),
1858 &optimized,
1859 &order_cache,
1860 Some(selector.clone()),
1861 None,
1862 SymEnv::default(),
1863 );
1864 assert!(
1865 result.is_ok() || matches!(result, Err(Error::TopologyChanged { .. })),
1866 "a picked realize over a branched graph must either succeed or hit the \
1867 (unrelated, expected-under-concurrency) TopologyChanged race: {:?}",
1868 result.err()
1869 );
1870 assert!(
1871 order_cache.get().is_none(),
1872 "a branched graph with a real selector pick must NEVER populate order_cache — \
1873 it must route through OrderSource::Streaming, not OrderSource::Cached",
1874 );
1875
1876 // A second call, same cache slot: still must not populate it.
1877 let result2 = PipelinedExecutor::realize_with_optimized_picking_cached_env(
1878 graph,
1879 root,
1880 cache,
1881 &optimized,
1882 &order_cache,
1883 Some(selector),
1884 None,
1885 SymEnv::default(),
1886 );
1887 assert!(
1888 result2.is_ok() || matches!(result2, Err(Error::TopologyChanged { .. })),
1889 "second picked realize must either succeed or hit the same expected-under- \
1890 concurrency TopologyChanged race: {:?}",
1891 result2.err()
1892 );
1893 assert!(
1894 order_cache.get().is_none(),
1895 "order_cache must still be empty after a second picked realize",
1896 );
1897 }
1898
1899 /// Deterministic sibling of the test above: forces the EXACT race
1900 /// `Error::TopologyChanged` tolerance was added for, instead of relying
1901 /// on another test in the binary to collide by chance. Proves two
1902 /// things the flaky-under-`-j4` observation alone couldn't: (1) that a
1903 /// stale `optimized.generation` is SUFFICIENT to reproduce the exact
1904 /// failure this test's sibling hit (confirming the root cause, not
1905 /// just a plausible story), and (2) that `order_cache` stays empty
1906 /// even on the `Err` branch (the property this test file is actually
1907 /// about, independent of the race).
1908 ///
1909 /// Deliberately does NOT call the real, process-wide
1910 /// `bump_topology_generation()` to force this — that would inflict the
1911 /// exact `-j4` cross-test hazard this test documents onto every OTHER
1912 /// concurrently-running test in the binary that assumes a stable
1913 /// generation (an earlier draft of this test did exactly that, and
1914 /// broke `realize_with_optimized_picking_cached_env_matches_arm0_cold_and_warm`
1915 /// under `-j4`). `OptimizedGraph::generation` is `pub`, so mutating
1916 /// this test's own LOCAL copy to a value that can't match the live
1917 /// counter reproduces the identical code path with zero blast radius.
1918 #[test]
1919 fn cached_env_tolerates_a_concurrent_topology_generation_bump() {
1920 use crate::PipelinedExecutor;
1921 use fuel_ir::{Error, SymEnv};
1922
1923 let (graph, cache, mut optimized, order_cache, selector, root) = cached_env_test_fixture();
1924
1925 // Simulate a sibling test's legitimate `register_runtime_kernel`/
1926 // `extend_global_bindings` call landing between `optimize_graph`
1927 // (which stamped `optimized.generation` above) and this realize —
1928 // exactly the collision `cached_env_never_caches_a_branched_graph_with_a_real_selector`
1929 // hit under `-j4` — by staling THIS test's own copy of the stamp,
1930 // not the shared live counter every other test also reads.
1931 optimized.generation = optimized.generation.wrapping_add(1);
1932
1933 let result = PipelinedExecutor::realize_with_optimized_picking_cached_env(
1934 graph,
1935 root,
1936 cache,
1937 &optimized,
1938 &order_cache,
1939 Some(selector),
1940 None,
1941 SymEnv::default(),
1942 );
1943 assert!(
1944 matches!(result, Err(Error::TopologyChanged { .. })),
1945 "a forced generation bump before realize must surface as TopologyChanged, \
1946 not silently succeed or fail some other way: {:?}",
1947 result
1948 );
1949 assert!(
1950 order_cache.get().is_none(),
1951 "order_cache must stay empty even when the realize that would have \
1952 populated it (if it were ever going to) failed with TopologyChanged",
1953 );
1954 }
1955
1956 /// (d) No `DEFAULT_MAX_N` truncation: the fork records the winner +
1957 /// a runner-up without a fixed top-N cap stranding placements. Both
1958 /// arms survive; the per-device frontier — not a fixed N — bounds
1959 /// the arms.
1960 #[test]
1961 fn no_default_max_n_truncation() {
1962 let mut table = KernelBindingTable::new();
1963 let (mut g, _prod, fork, _tail, root) = build_single_fork_graph(&mut table);
1964 let opts = two_backend_opts();
1965
1966 let optimized =
1967 optimize_graph(&mut g, &[root], &table, &opts).expect("optimize_graph succeeds");
1968 assert_eq!(optimized.branch_count(&g), 1);
1969 let branch_id = (0..g.len())
1970 .map(NodeId)
1971 .find(|&id| matches!(g.node(id).op, Op::Branch { .. }))
1972 .expect("a Branch node exists");
1973 // Two distinct placements ⇒ two arms survive; the runner-up was
1974 // not dropped by any top-N cap.
1975 assert_eq!(g.node(branch_id).inputs.len(), 2);
1976 assert_eq!(g.node(branch_id).inputs[0], fork);
1977 }
1978
1979 /// THE EQUIVALENCE GATE (born-red until `optimize_graph` exists and
1980 /// its lowered order matches `compile_plan`).
1981 ///
1982 /// On representative no-competing-route graphs, assert:
1983 /// (a) `optimize_graph` leaves ZERO `Op::Branch` nodes, and
1984 /// (b) the concatenated `extract_runs`/`lower_run` dispatch order
1985 /// EQUALS today's `compile_plan(...).order` exactly (same
1986 /// NodeIds, same order).
1987 #[test]
1988 fn equivalence_gate_branchless_order_matches_compile_plan() {
1989 for build in [
1990 build_straight_line_graph as fn(&mut KernelBindingTable) -> (Graph, NodeId),
1991 build_branchless_graph as fn(&mut KernelBindingTable) -> (Graph, NodeId),
1992 ] {
1993 let mut table = KernelBindingTable::new();
1994 let (mut g, root) = build(&mut table);
1995 let opts = cpu_opts();
1996
1997 // Today's path: the exact order the executor walks.
1998 let order = execution_plan(&g, &[root]);
1999 let plan = compile_plan(&g, &order, &table, &opts)
2000 .expect("today's compile_plan succeeds on a placeable graph");
2001
2002 // New path.
2003 let optimized = optimize_graph(&mut g, &[root], &table, &opts)
2004 .expect("optimize_graph succeeds on the same graph");
2005
2006 // (a) zero competing routes => zero Branch nodes.
2007 assert_eq!(
2008 optimized.branch_count(&g),
2009 0,
2010 "a no-competing-route graph optimizes to zero Op::Branch nodes",
2011 );
2012
2013 // (b) EXACT order equality against compile_plan(...).order.
2014 let lowered = optimized.dispatch_order(&g);
2015 assert_eq!(
2016 lowered, plan.order,
2017 "the extract_runs/lower_run dispatch order must equal \
2018 compile_plan(...).order exactly (same NodeIds, same order)",
2019 );
2020 // And it covers every reachable node — the executor walks
2021 // the whole graph, runs partition it with no gaps/dupes.
2022 assert_eq!(
2023 lowered.len(),
2024 order.len(),
2025 "the lowered order covers every node compile_plan ordered",
2026 );
2027 }
2028 }
2029
2030 /// `optimize_graph` on a branchless graph is idempotent and adds no
2031 /// nodes: the node count is unchanged across repeated calls, the
2032 /// branch count stays zero, and the lowered order is stable.
2033 #[test]
2034 fn optimize_graph_branchless_is_idempotent_and_adds_no_nodes() {
2035 let mut table = KernelBindingTable::new();
2036 let (mut g, root) = build_branchless_graph(&mut table);
2037 let opts = cpu_opts();
2038
2039 let nodes_before = g.len();
2040 let first =
2041 optimize_graph(&mut g, &[root], &table, &opts).expect("first optimize succeeds");
2042 let order_first = first.dispatch_order(&g);
2043 let nodes_after_first = g.len();
2044
2045 let second =
2046 optimize_graph(&mut g, &[root], &table, &opts).expect("second optimize succeeds");
2047 let order_second = second.dispatch_order(&g);
2048 let nodes_after_second = g.len();
2049
2050 assert_eq!(
2051 nodes_before, nodes_after_first,
2052 "optimize_graph adds no nodes (A3a is single-route-only)",
2053 );
2054 assert_eq!(
2055 nodes_after_first, nodes_after_second,
2056 "a second optimize_graph adds no further nodes — idempotent",
2057 );
2058 assert_eq!(first.branch_count(&g), 0, "no branches after first");
2059 assert_eq!(second.branch_count(&g), 0, "no branches after second");
2060 assert_eq!(
2061 order_first, order_second,
2062 "the lowered dispatch order is stable across repeated optimize",
2063 );
2064 }
2065
2066 /// Identity proof for `docs/design/incremental-consumer-index.md` §5's
2067 /// `order_for` fix track: that fix proposes memoizing `dispatch_order`'s
2068 /// output WHOLESALE once per plan-once session, instead of recomputing
2069 /// it (via `extract_runs_multi` + `non_chosen_arm_nodes`, measured
2070 /// ~56ms/step, the single largest named sub-cost in a live decode step)
2071 /// on every `realize_inner` call. That's only valid if calling
2072 /// `dispatch_order` repeatedly on the SAME `OptimizedGraph` view with NO
2073 /// intervening graph mutation — exactly the plan-once decode shape,
2074 /// where only leaf values are rebound between tokens, never the graph's
2075 /// structure — yields a byte-identical `Vec<NodeId>` every time. Proven
2076 /// here, not assumed: five repeated calls, two structurally different
2077 /// branchless graphs (a fan-in boundary and a plain chain), zero
2078 /// mutation between calls (deliberately — this is NOT the existing
2079 /// `optimize_graph_branchless_is_idempotent_and_adds_no_nodes` test
2080 /// above, which re-runs `optimize_graph` itself between its two
2081 /// `dispatch_order` calls; this one proves `dispatch_order` is pure on
2082 /// its own, independent of re-optimizing).
2083 #[test]
2084 fn dispatch_order_is_stable_across_repeated_calls_with_no_mutation() {
2085 for build in [
2086 build_branchless_graph as fn(&mut KernelBindingTable) -> (Graph, NodeId),
2087 build_straight_line_graph,
2088 ] {
2089 let mut table = KernelBindingTable::new();
2090 let (mut g, root) = build(&mut table);
2091 let opts = cpu_opts();
2092 let optimized = optimize_graph(&mut g, &[root], &table, &opts).expect("optimize_graph");
2093
2094 let orders: Vec<Vec<NodeId>> = (0..5).map(|_| optimized.dispatch_order(&g)).collect();
2095 for (i, order) in orders.iter().enumerate().skip(1) {
2096 assert_eq!(
2097 &orders[0], order,
2098 "dispatch_order call {i} diverged from call 0 with no intervening \
2099 mutation — the order_for memoization fix would be UNSAFE if this failed",
2100 );
2101 }
2102 }
2103 }
2104
2105 /// Build-time validation: optimize_graph fails fast (Result, never
2106 /// panic) when a kernel-bearing node has no registered binding —
2107 /// reusing compile_plan's missing-binding diagnostic.
2108 #[test]
2109 fn optimize_graph_fails_fast_on_missing_binding() {
2110 let table = KernelBindingTable::new(); // empty — no bindings.
2111 let mut g = Graph::new();
2112 let a = f32_node(&mut g, Op::Const, vec![]);
2113 let relu = f32_node(&mut g, Op::Relu, vec![a]);
2114 let opts = cpu_opts();
2115
2116 let err = optimize_graph(&mut g, &[relu], &table, &opts)
2117 .map(|_| ())
2118 .unwrap_err();
2119 match err {
2120 fuel_ir::Error::NoBackendForOp { op, .. } => {
2121 assert_eq!(op, OpKind::ReluElementwise);
2122 }
2123 other => panic!("expected NoBackendForOp, got {other:?}"),
2124 }
2125 }
2126
2127 // ===== Phase B PR-B3: the lock-step pathfinder/ranker/optimizer
2128 // driver. =====
2129
2130 use crate::driver::{
2131 FrontierConvergenceOptimizer, OptimizationContext, Optimizer, PassRegistry, Pathfinder,
2132 PlacementForkPathfinder,
2133 };
2134 use std::collections::HashSet;
2135
2136 /// Snapshot of a graph's `Op::Branch` decision points, keyed by
2137 /// arena order: `(branch_id, reconverge_at, arm NodeIds)`. Two
2138 /// optimization paths that produce the same snapshot produced the
2139 /// same multi-path graph (behavior-preservation).
2140 fn branch_snapshot(g: &Graph) -> Vec<(NodeId, NodeId, Vec<NodeId>)> {
2141 (0..g.len())
2142 .map(NodeId)
2143 .filter_map(|id| match g.node(id).op {
2144 Op::Branch { reconverge_at } => {
2145 Some((id, reconverge_at, g.node(id).inputs.clone()))
2146 }
2147 _ => None,
2148 })
2149 .collect()
2150 }
2151
2152 /// THE LOAD-BEARING BEHAVIOR-PRESERVATION GATE (born-red until
2153 /// `optimize_graph` routes through the lock-step driver).
2154 ///
2155 /// On a representative multi-placement graph, the driver
2156 /// (`PassRegistry::default_passes().run_lockstep`) produces the
2157 /// **identical** multi-path graph as the old hardcoded sequence
2158 /// (pathfinder then optimizer, in that order): same branch count,
2159 /// same arm-0 winner, same arms, same reconverge points, and the
2160 /// same arm-0 dispatch order. B3 is a structural refactor — the
2161 /// optimized result must not change.
2162 #[test]
2163 fn driver_matches_legacy_sequence() {
2164 // (A) The reference: drive the registered passes BY HAND in the
2165 // legacy order (propose, then prune) onto a freshly compiled
2166 // plan — exactly what the pre-B3 hardcoded sequence did.
2167 let mut ref_table = KernelBindingTable::new();
2168 let (mut ref_g, _prod, ref_fork, _tail, ref_root) = build_single_fork_graph(&mut ref_table);
2169 let opts = two_backend_opts();
2170 let ref_order = execution_plan(&ref_g, &[ref_root]);
2171 let ref_plan = compile_plan(&ref_g, &ref_order, &ref_table, &opts)
2172 .expect("reference compile_plan succeeds");
2173 {
2174 let guard: HashSet<NodeId> = HashSet::new();
2175 let ctx = OptimizationContext {
2176 order: &ref_order,
2177 plan: &ref_plan,
2178 cycle_guard: &guard,
2179 bindings: &ref_table,
2180 };
2181 // Legacy order: the pathfinder ADDs, then the optimizer
2182 // prunes — the same two operations the hardcoded sequence ran.
2183 PlacementForkPathfinder
2184 .propose(&mut ref_g, &ctx)
2185 .expect("pathfinder proposes");
2186 FrontierConvergenceOptimizer
2187 .prune(&mut ref_g, &ctx)
2188 .expect("optimizer prunes");
2189 }
2190 let ref_snapshot = branch_snapshot(&ref_g);
2191 let ref_dispatch = fuel_graph::lower_runs_arm0(&ref_g, &[ref_root]);
2192
2193 // (B) The driver path: optimize_graph routes through
2194 // PassRegistry::default_passes().run_lockstep.
2195 let mut table = KernelBindingTable::new();
2196 let (mut g, _p, fork, _t, root) = build_single_fork_graph(&mut table);
2197 let optimized =
2198 optimize_graph(&mut g, &[root], &table, &opts).expect("optimize_graph succeeds");
2199 let got_snapshot = branch_snapshot(&g);
2200 let got_dispatch = optimized.dispatch_order(&g);
2201
2202 // Identical multi-path graph: same branches, arms, reconverge.
2203 assert_eq!(
2204 got_snapshot, ref_snapshot,
2205 "driver produces the identical Op::Branch structure as the legacy \
2206 pathfinder→optimizer sequence",
2207 );
2208 assert_eq!(
2209 got_snapshot.len(),
2210 1,
2211 "the representative multi-placement graph produces exactly one branch",
2212 );
2213 assert_eq!(
2214 got_snapshot[0].2[0], fork,
2215 "arm-0 is the DP winner (the route realize uses)",
2216 );
2217 // Identical dispatch order (arm-0 single-route lowering).
2218 assert_eq!(
2219 got_dispatch, ref_dispatch,
2220 "driver yields the identical arm-0 dispatch order as the legacy sequence",
2221 );
2222 let _ = ref_fork;
2223 }
2224
2225 /// A no-op registry (no registered pathfinders) leaves a branchless
2226 /// graph unchanged: no nodes added, zero branches, identical
2227 /// dispatch order. Proves the driver itself introduces no structure.
2228 #[test]
2229 fn noop_registry_leaves_branchless_graph_unchanged() {
2230 let mut table = KernelBindingTable::new();
2231 let (mut g, root) = build_branchless_graph(&mut table);
2232 let opts = cpu_opts();
2233
2234 let order = execution_plan(&g, &[root]);
2235 let plan = compile_plan(&g, &order, &table, &opts).expect("compile_plan succeeds");
2236
2237 let dispatch_before = fuel_graph::lower_runs_arm0(&g, &[root]);
2238 let nodes_before = g.len();
2239
2240 let registry = PassRegistry::new(); // empty — no passes.
2241 assert_eq!(registry.pathfinder_count(), 0);
2242 assert_eq!(registry.optimizer_count(), 0);
2243
2244 let guard: HashSet<NodeId> = HashSet::new();
2245 let ctx = OptimizationContext {
2246 order: &order,
2247 plan: &plan,
2248 cycle_guard: &guard,
2249 bindings: &table,
2250 };
2251 registry
2252 .run_lockstep(&mut g, &ctx)
2253 .expect("empty registry runs cleanly");
2254
2255 assert_eq!(g.len(), nodes_before, "no-op registry adds no nodes");
2256 let branches = (0..g.len())
2257 .map(NodeId)
2258 .filter(|&id| matches!(g.node(id).op, Op::Branch { .. }))
2259 .count();
2260 assert_eq!(branches, 0, "no-op registry adds no branches");
2261 let dispatch_after = fuel_graph::lower_runs_arm0(&g, &[root]);
2262 assert_eq!(
2263 dispatch_before, dispatch_after,
2264 "no-op registry leaves the dispatch order unchanged",
2265 );
2266 }
2267
2268 /// Lock-step ordering: a pathfinder ADDs before its dependent
2269 /// optimizer PRUNEs. We register a recording pathfinder and a
2270 /// recording optimizer that each stamp a shared log when invoked;
2271 /// `run_lockstep` must invoke the pathfinder strictly before the
2272 /// optimizer (prune-as-you-go, never optimizer-first).
2273 #[test]
2274 fn lockstep_runs_pathfinder_before_dependent_optimizer() {
2275 use std::sync::{Arc, Mutex};
2276
2277 #[derive(Clone)]
2278 struct RecordPathfinder(Arc<Mutex<Vec<&'static str>>>);
2279 impl Pathfinder for RecordPathfinder {
2280 fn name(&self) -> &'static str {
2281 "RecordPathfinder"
2282 }
2283 fn propose(
2284 &self,
2285 _g: &mut Graph,
2286 _ctx: &OptimizationContext<'_>,
2287 ) -> fuel_ir::Result<()> {
2288 self.0.lock().unwrap().push("propose");
2289 Ok(())
2290 }
2291 }
2292 #[derive(Clone)]
2293 struct RecordOptimizer(Arc<Mutex<Vec<&'static str>>>);
2294 impl Optimizer for RecordOptimizer {
2295 fn name(&self) -> &'static str {
2296 "RecordOptimizer"
2297 }
2298 fn prune(&self, _g: &mut Graph, _ctx: &OptimizationContext<'_>) -> fuel_ir::Result<()> {
2299 self.0.lock().unwrap().push("prune");
2300 Ok(())
2301 }
2302 }
2303
2304 let log = Arc::new(Mutex::new(Vec::new()));
2305 let registry = PassRegistry::new()
2306 .with_pathfinder(Box::new(RecordPathfinder(log.clone())))
2307 .with_optimizer(Box::new(RecordOptimizer(log.clone())));
2308
2309 let mut g = Graph::new();
2310 let _a = f32_node(&mut g, Op::Const, vec![]);
2311 let plan = ExecutionPlan::empty();
2312 let order: Vec<NodeId> = Vec::new();
2313 let guard: HashSet<NodeId> = HashSet::new();
2314 let table = KernelBindingTable::new();
2315 let ctx = OptimizationContext {
2316 order: &order,
2317 plan: &plan,
2318 cycle_guard: &guard,
2319 bindings: &table,
2320 };
2321 registry.run_lockstep(&mut g, &ctx).expect("drive runs");
2322
2323 let seq = log.lock().unwrap().clone();
2324 assert_eq!(
2325 seq,
2326 vec!["propose", "prune"],
2327 "lock-step drive runs the pathfinder ADD strictly before its \
2328 dependent optimizer PRUNE (prune-as-you-go)",
2329 );
2330 }
2331
2332 /// **Where does `PagedAttn`'s DECOMPOSED form land on CUDA?** — the
2333 /// measurement that gates whether offering a decomposed arm is worth
2334 /// building, and what to ask Baracuda for.
2335 ///
2336 /// Established already: the FUSED `Op::PagedAttn` node executes on the CPU
2337 /// under CUDA decode (`paged_decode_node_placement_report_cuda`), because
2338 /// no CUDA backend advertises it — so "runs on host" is not a decision the
2339 /// planner made, it is the only move it had. `decompose` exists and is
2340 /// closed-basis, but the planner is never offered it.
2341 ///
2342 /// This asks the question that decides the next build, and it is a
2343 /// GPU-vs-GPU question that the 186× host-ward figure cannot answer —
2344 /// **any** on-device path removes that round-trip equally, so its magnitude
2345 /// says something must be built without saying what:
2346 ///
2347 /// - **every recipe node CUDA-placed** → the round-trip is removable with
2348 /// NO kernel work at all; the ask to Baracuda may be nothing.
2349 /// - **some primitive falls back to host** → *that primitive* is the ask,
2350 /// far smaller and better-specified than "paged attention".
2351 ///
2352 /// A grep cannot answer this. I tried: it reported `Permute` in zero
2353 /// CUDA-backend files while the measured placement dump showed
2354 /// `Permute Cuda×6` — it undercounts ops served by generic dispatch. Only
2355 /// the planner's own resolved map settles it.
2356 ///
2357 /// REPORTS rather than asserting a verdict, for the same reason as the
2358 /// fused-node report: pre-committing to an answer is the error being
2359 /// corrected. What it DOES assert is that the instrument ran
2360 /// (`has_placements()`) and that decomposition actually happened (the root
2361 /// moved off the fused node) — an all-`None` dump would read as "nothing
2362 /// is on CUDA", a wrong answer wearing a null answer's clothes.
2363 #[test]
2364 #[ignore = "live GPU (CUDA); run via scripts/gpu-run.ps1"]
2365 #[cfg(feature = "cuda")]
2366 fn paged_attn_decomposed_arm_placement_report_cuda() {
2367 use fuel_graph::registry::{FusedOpParams, FusedOps};
2368 use std::collections::BTreeMap;
2369
2370 const B: usize = 1;
2371 const HQ: usize = 4;
2372 const HKV: usize = 4;
2373 const SQ: usize = 1;
2374 const D: usize = 4;
2375 const NUM_BLOCKS: usize = 32;
2376 const BLOCK_SIZE: usize = 4;
2377 const MAX_BLK: usize = 8;
2378
2379 let mut g = Graph::new();
2380 let leaf = |g: &mut Graph, dims: &[usize], dtype: DType| {
2381 g.push(Node {
2382 op: Op::Const,
2383 inputs: vec![],
2384 shape: Shape::from_dims(dims),
2385 dtype,
2386 })
2387 };
2388 let q = leaf(&mut g, &[B, HQ, SQ, D], DType::F32);
2389 let kc = leaf(&mut g, &[NUM_BLOCKS, BLOCK_SIZE, HKV, D], DType::F32);
2390 let vc = leaf(&mut g, &[NUM_BLOCKS, BLOCK_SIZE, HKV, D], DType::F32);
2391 let bt = leaf(&mut g, &[B, MAX_BLK], DType::U32);
2392 let cl = leaf(&mut g, &[B], DType::U32);
2393
2394 let params = FusedOpParams::PagedAttn {
2395 softmax_scale: 1.0 / (D as f32).sqrt(),
2396 block_size: BLOCK_SIZE,
2397 softcap: None,
2398 };
2399 let fused = g.push(Node {
2400 op: Op::Fused(FusedOps::PAGED_ATTN, params.clone()),
2401 inputs: vec![q, kc, vc, bt, cl],
2402 shape: Shape::from_dims(&[B, HQ, SQ, D]),
2403 dtype: DType::F32,
2404 });
2405
2406 // Lower the fused node to its primitive recipe — the arm the planner is
2407 // never currently offered.
2408 let root = fuel_graph::registry::paged_attn::decompose(&mut g, fused, ¶ms);
2409 assert_ne!(
2410 root, fused,
2411 "PRECONDITION: decompose must actually lower. A self-return is the \
2412 G2 fixpoint signal (malformed node / declined recipe), and every \
2413 placement below would then describe the FUSED node, silently \
2414 answering a different question than the one asked.",
2415 );
2416
2417 // THE PRODUCTION TABLE, not a hand-assembled one. An earlier version of
2418 // this test built the table with `register_cuda_kernels` alone and got
2419 // "no backend supports index_select" — which was an artifact of the
2420 // HARNESS, not a fact about Fuel: that function makes ~5 registrations
2421 // and the bulk of CUDA lives in `register_baracuda_cuda_kernels`, which
2422 // the real path (`register_optional_backends`) also calls, on top of
2423 // `register_cpu_kernels` and a cost fill. Reproducing that assembly by
2424 // hand is precisely the kind of thing that silently drifts, so ask for
2425 // the table the realize path actually uses.
2426 let table = crate::dispatch::global_bindings();
2427
2428 // POSITIVE CONTROL — the instrument must be able to see a CUDA binding
2429 // at all. Without this, an under-populated table reports "no CUDA
2430 // support" for EVERY op, which reads exactly like a real finding: it is
2431 // specific, it names a primitive, and it is completely wrong. That is
2432 // how the first run of this test nearly shipped a false Baracuda ask.
2433 // MatMul is the control because the fuel-core placement report measured
2434 // `MatMul Cuda×15` on a real decode graph — an independently observed
2435 // fact this table must reproduce.
2436 assert!(
2437 table
2438 .lookup(
2439 OpKind::MatMul,
2440 &[DType::F32, DType::F32, DType::F32],
2441 BackendId::Cuda
2442 )
2443 .is_ok(),
2444 "INSTRUMENT BROKEN: the binding table has no CUDA MatMul, so every \
2445 'no CUDA support' conclusion below would be an artifact of the \
2446 harness. Fix the table before reading any verdict from this test.",
2447 );
2448
2449 // ---- THE COMPLETE MISSING-BINDING SET, asked directly ----
2450 //
2451 // `optimize_graph` fails at the FIRST unplaceable node, so driving this
2452 // off its error yields one op per run — a trickle, and each run costs a
2453 // CUDA rebuild. Worse, it conflates two different things: "CUDA has no
2454 // kernel for this op" and "the planner could not place it under THIS
2455 // pin". Asking the binding table per node separates them and returns the
2456 // whole set at once.
2457 let mut missing_cuda: Vec<String> = Vec::new();
2458 let mut present_cuda: Vec<String> = Vec::new();
2459 for i in 0..g.len() {
2460 let id = NodeId(i);
2461 let node = g.node(id);
2462 let Some(kind) = crate::pipelined::op_to_op_kind(&node.op) else {
2463 continue; // Const / view / fused — not a binding-table lookup
2464 };
2465 let mut dtypes: Vec<DType> = node.inputs.iter().map(|&i| g.node(i).dtype).collect();
2466 dtypes.push(node.dtype); // binding key is [inputs.., output]
2467 let sig = format!(
2468 "{kind:?}{:?}",
2469 dtypes.iter().map(|d| format!("{d:?}")).collect::<Vec<_>>()
2470 );
2471 if table.lookup(kind, &dtypes, BackendId::Cuda).is_ok() {
2472 present_cuda.push(sig);
2473 } else {
2474 missing_cuda.push(sig);
2475 }
2476 }
2477 missing_cuda.sort();
2478 missing_cuda.dedup();
2479 present_cuda.sort();
2480 present_cuda.dedup();
2481
2482 println!("\n=== PagedAttn recipe: CUDA binding coverage (asked per node) ===");
2483 println!("PRESENT on CUDA ({}):", present_cuda.len());
2484 for s in &present_cuda {
2485 println!(" {s}");
2486 }
2487 println!("MISSING on CUDA ({}):", missing_cuda.len());
2488 for s in &missing_cuda {
2489 println!(" {s}");
2490 }
2491 // A run where NOTHING resolves would mean the table is empty again — the
2492 // MatMul control above already guards that, but state the invariant that
2493 // makes the "missing" list meaningful rather than vacuous.
2494 assert!(
2495 !present_cuda.is_empty(),
2496 "INSTRUMENT BROKEN: not one recipe node resolved on CUDA, so the \
2497 MISSING list is a property of the lookup, not of the backend",
2498 );
2499 println!(
2500 "\n>>> THE ASK: the {} signature(s) above are what CUDA lacks for the \
2501 §1 dense recipe. Each is a small, named primitive — not \"paged \
2502 attention\".\n",
2503 missing_cuda.len()
2504 );
2505
2506 let opts = PlanOptions::new().with_pinned_device(DeviceLocation::Cuda { gpu_id: 0 });
2507 let optimized = match optimize_graph(&mut g, &[root], &table, &opts) {
2508 Ok(o) => o,
2509 Err(e) => {
2510 println!("\n=== DECOMPOSED ARM DOES NOT PLAN ON CUDA ===");
2511 println!("optimize_graph error: {e:?}");
2512 println!("That is itself the answer: some recipe primitive has no CUDA");
2513 println!("binding, and the error names it. THAT is the Baracuda ask.");
2514 return;
2515 }
2516 };
2517
2518 assert!(
2519 optimized.has_placements(),
2520 "placement NOT COMPUTED — every placement_of would be None and the \
2521 report below would describe the missing instrument, not the graph",
2522 );
2523
2524 let mut by_op: BTreeMap<String, BTreeMap<String, usize>> = BTreeMap::new();
2525 let mut host_ops: Vec<String> = Vec::new();
2526 for i in 0..g.len() {
2527 let id = NodeId(i);
2528 let op_name = match &g.node(id).op {
2529 Op::Fused(fid, _) => format!("Fused({fid:?})"),
2530 other => format!("{other:?}")
2531 .split(' ')
2532 .next()
2533 .unwrap_or("?")
2534 .to_string(),
2535 };
2536 let place = match optimized.placement_of(id) {
2537 Some(d) => format!("{d:?}"),
2538 None => "None".to_string(),
2539 };
2540 if place.starts_with("Cpu") {
2541 host_ops.push(op_name.clone());
2542 }
2543 *by_op.entry(op_name).or_default().entry(place).or_insert(0) += 1;
2544 }
2545
2546 println!("\n=== PagedAttn DECOMPOSED arm: node placement (CUDA-pinned) ===");
2547 println!(
2548 "total nodes: {} (fused node #{} lowered to root #{})",
2549 g.len(),
2550 fused.0,
2551 root.0
2552 );
2553 for (op, places) in &by_op {
2554 let rendered: Vec<String> = places.iter().map(|(d, n)| format!("{d}×{n}")).collect();
2555 println!(" {op:<40} {}", rendered.join(" "));
2556 }
2557
2558 println!("\n--- VERDICT ---");
2559 if host_ops.is_empty() {
2560 println!(
2561 "EVERY recipe node is off-host. The decomposed paged attention runs \
2562 ON DEVICE, so the 186x round-trip is removable by OFFERING THIS ARM \
2563 — no Baracuda kernel required to eliminate it. A kernel would then \
2564 be a throughput question (the recipe materialises padded capacity, \
2565 not live context_len), not a round-trip one."
2566 );
2567 } else {
2568 host_ops.sort();
2569 host_ops.dedup();
2570 println!(
2571 "HOST-PLACED recipe primitives: {host_ops:?}\n\
2572 These are the ask — each is a far smaller and better-specified \
2573 request than \"paged attention\", and until they exist on device \
2574 the decomposed arm cannot win."
2575 );
2576 }
2577 println!("--- END VERDICT ---\n");
2578 }
2579}