Skip to main content

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 &copy_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, &params);
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}