fuel_dispatch/ranker/

route_picker.rs

// SPDX-License-Identifier: MIT OR Apache-2.0
//! The **runtime route picker** ("Picker 2") — selects an arm at each
//! `Op::Branch` decision point by **live telemetry**.
//!
//! Phase C PR-C1 of the "plan IS the graph" rebuild
//! ([`../../docs/session-prompts/plan-is-graph-rebuild.md`](
//! ../../docs/session-prompts/plan-is-graph-rebuild.md) Phase C,
//! capability 7; [`../../docs/architecture/06-runtime.md`](
//! ../../docs/architecture/06-runtime.md) §"Route picker (the runtime
//! selector / Picker 2)").
//!
//! ## What this is
//!
//! Phases A+B emit the optimized form as an in-graph multi-path
//! structure: `optimize_graph` records `Op::Branch { reconverge_at }`
//! decision points whose arms (`inputs`) are competing `(device, backend)`
//! placements — arm-0 the cost-vector winner, arm-1+ runner-up
//! placements bounded by the per-device Pareto frontier. Phase B realize
//! followed **arm-0 statically** (`lower_runs_arm0`). This module is the
//! runtime selector that **chooses an arm per branch by live telemetry**,
//! replacing the static arm-0 lowering with [`fuel_graph::lower_picked_route`].
//!
//! ## How it picks
//!
//! [`pick_route`] walks the reachable branches in **topological order**
//! ([`fuel_graph::branches_in_topo_order`], keyed by each branch's
//! `reconverge_at` merge position) so coupled upstream decisions are
//! committed before a downstream branch is reached. At each branch it
//! builds a per-branch [`AlternativeSet`] over the arms — one
//! [`Candidate`] per arm carrying the arm's real `(backend, device)` and
//! its plan cost — and consults the supplied [`RuntimeSelector`] (the
//! production [`crate::ranker::ChainedSelector`]: VRAM-pressure guard
//! over the [`BackendRuntimeLookup`]'s per-tier free memory, then
//! Judge-measured rank, then the static winner). The picked candidate's
//! `(backend, device)` maps back to an arm index, stored in the
//! [`PickedRoute`].
//!
//! Under **no pressure / no telemetry** the chained selector degrades to
//! `set.winner()` = arm-0 (the cost-ranked first candidate), so the route
//! is arm-0 everywhere and [`fuel_graph::lower_picked_route`] reproduces
//! Phase B's `lower_runs_arm0` exactly. Under **VRAM pressure** the guard
//! demotes/skips the GPU arm and the picker takes a host-RAM arm
//! instead — exactly the "multiple paths per device" payoff
//! (06-runtime §"Memory pressure as the parallelism limit").
//!
//! ## Caching + re-resolution
//!
//! [`RouteCache`] holds the last resolved [`PickedRoute`] plus a
//! [`TelemetryFingerprint`] (per-tier free-memory buckets across the
//! candidate backends). [`RouteCache::resolve`] re-resolves only when the
//! fingerprint changes meaningfully (a bucket shift); otherwise it reuses
//! the cached route. In steady state the picker is a table lookup;
//! re-resolution fires on transitions (pressure rising, a device
//! appearing). A CPU-only build has zero branches ⇒ the route is empty ⇒
//! the picker is a no-op and realize is unchanged.

use std::sync::Arc;

use fuel_graph::{Graph, NodeId, PickedRoute, branches_in_topo_order};
use fuel_ir::dispatch::{OpKind, SizeClass};
use fuel_ir::probe::BackendId;
use fuel_ir::{DType, DeviceLocation, Shape};

use crate::fused::{CostEstimate, PrecisionGuarantee};
use crate::kernel::{KernelBindingTable, KernelCaps, OpParams};
use crate::pipelined::{build_lookup_dtypes, op_to_op_kind};
use crate::plan::default_device_for;
use crate::ranker::{
    AlternativeSet, BackendRuntimeLookup, Candidate, DecisionContext, RuntimeSelector,
    enumerate_candidates,
};

/// Bounded lookahead for adversarially-coupled branches (06-runtime
/// §"Resolution order matters when decisions are coupled"). The default
/// resolution is locally-greedy in topological order — by the time the
/// picker reaches a branch its upstream decisions are already committed.
/// Rare adversarial coupling is bounded by considering at most this many
/// branches jointly. PR-C1 commits upstream-first in topo order (the
/// greedy default); `K` is the documented ceiling the joint-resolution
/// window must never exceed.
pub const LOOKAHEAD_K: usize = 3;

/// A coarse fingerprint of the live per-tier free memory across the
/// backends a route's branches care about — the cache key for
/// [`RouteCache`]. Re-resolution fires only on a *meaningful* delta: we
/// bucket each backend's free bytes into coarse log-scale tiers so noise
/// (a few MB jitter) does not invalidate the cached route, but a genuine
/// pressure transition (free memory dropping an order of magnitude) does.
#[derive(Clone, Debug, Default, PartialEq, Eq)]
pub struct TelemetryFingerprint {
    /// `(backend, free-memory bucket)` pairs, sorted by backend for a
    /// stable equality check. An absent backend (no live handle) is
    /// omitted — its bucket is "unknown", which never forces a
    /// re-resolve on its own.
    buckets: Vec<(BackendId, u8)>,
}

impl TelemetryFingerprint {
    /// Fingerprint the backends in `backends` against the live lookup.
    /// Each backend's free bytes are bucketed via [`free_bytes_bucket`].
    fn sample(backends: &[BackendId], lookup: Option<&BackendRuntimeLookup>) -> Self {
        let mut buckets: Vec<(BackendId, u8)> = Vec::new();
        if let Some(lookup) = lookup {
            let mut seen: Vec<BackendId> = Vec::new();
            for &b in backends {
                if seen.contains(&b) {
                    continue;
                }
                seen.push(b);
                if let Some(handle) = lookup(b, default_device_for(b))
                    && let Some(free) = handle.available_bytes()
                {
                    buckets.push((b, free_bytes_bucket(free)));
                }
            }
        }
        buckets.sort_by_key(|&(b, _)| b as u8);
        Self { buckets }
    }
}

/// Bucket free-memory bytes into coarse log-scale tiers (powers of ~16×)
/// so the [`TelemetryFingerprint`] is stable under small jitter but
/// flips on an order-of-magnitude pressure transition. `0` for zero free;
/// otherwise `floor(log2(free)/4) + 1` clamped to a byte.
fn free_bytes_bucket(free: u64) -> u8 {
    if free == 0 {
        return 0;
    }
    let log2 = 63 - free.leading_zeros(); // floor(log2(free))
    ((log2 / 4) + 1).min(u8::MAX as u32) as u8
}

/// Per-realize cache of the resolved route + the telemetry it was
/// resolved under (06-runtime §"Telemetry caching for picker speed").
///
/// Live across realizes by being held by the caller (the bridge). The
/// first [`resolve`](RouteCache::resolve) populates it; subsequent calls
/// reuse the cached route while the [`TelemetryFingerprint`] is stable,
/// and re-resolve on a meaningful delta.
#[derive(Clone, Debug, Default)]
pub struct RouteCache {
    cached: Option<(TelemetryFingerprint, PickedRoute)>,
    /// Diagnostics: how many times [`resolve`](RouteCache::resolve) hit
    /// the cache vs. re-resolved. Tests assert the re-resolve discipline.
    pub hits: u64,
    pub resolves: u64,
}

impl RouteCache {
    /// A fresh, empty cache.
    pub fn new() -> Self {
        Self::default()
    }

    /// Resolve the route for `roots`, reusing the cached route while the
    /// live telemetry fingerprint is unchanged. Re-resolves (and updates
    /// the cache) on the first call or a meaningful telemetry delta.
    ///
    /// `bindings` is the runtime kernel registry (the per-arm candidates are
    /// re-enumerated from it + the graph); `selector` is the production runtime
    /// selector; `lookup` is the live per-tier free-memory closure (`None` ⇒ no
    /// pressure signal ⇒ arm-0).
    pub fn resolve(
        &mut self,
        graph: &Graph,
        roots: &[NodeId],
        bindings: &KernelBindingTable,
        selector: &dyn RuntimeSelector,
        lookup: Option<&BackendRuntimeLookup>,
    ) -> PickedRoute {
        let backends = route_backends(graph, roots);
        let fingerprint = TelemetryFingerprint::sample(&backends, lookup);

        if let Some((cached_fp, cached_route)) = &self.cached
            && *cached_fp == fingerprint
        {
            self.hits += 1;
            return cached_route.clone();
        }

        let route = pick_route(graph, roots, bindings, selector, lookup);
        self.resolves += 1;
        self.cached = Some((fingerprint, route.clone()));
        route
    }
}

/// The set of `(backend)`s the route's branches' arms target — the
/// backends whose live free memory the picker (and the cache
/// fingerprint) consults.
fn route_backends(graph: &Graph, roots: &[NodeId]) -> Vec<BackendId> {
    let mut backends: Vec<BackendId> = Vec::new();
    for branch in branches_in_topo_order(graph, roots) {
        for &arm in &graph.node(branch).inputs {
            if let Some(b) = graph.target_backend(arm)
                && !backends.contains(&b)
            {
                backends.push(b);
            }
        }
    }
    backends
}

/// Walk the multi-path graph's `Op::Branch` decision points in
/// **topological order** and choose one arm at each via the
/// `RuntimeSelector` over the arms' `(device, backend)` + plan cost +
/// live per-tier free memory. Returns the [`PickedRoute`]
/// ([`fuel_graph::lower_picked_route`] consumes it).
///
/// Topological order means a downstream branch is resolved only after its
/// upstream branches are committed (the picks recorded in `picked`), so
/// coupled decisions resolve consistently (06-runtime §"Resolution order
/// matters when decisions are coupled"). The greedy upstream-first pass
/// is the default; [`LOOKAHEAD_K`] bounds any joint-resolution window for
/// adversarial coupling.
///
/// With no pressure / no telemetry the selector degrades to `set.winner()`
/// = arm-0 for every branch, so the returned route is empty-equivalent
/// (every pick is arm 0) and realize is identical to Phase B.
pub fn pick_route(
    graph: &Graph,
    roots: &[NodeId],
    bindings: &KernelBindingTable,
    selector: &dyn RuntimeSelector,
    lookup: Option<&BackendRuntimeLookup>,
) -> PickedRoute {
    // The live per-tier free-memory lookup is consulted *through* the
    // selector — the production `ChainedSelector` holds its own
    // `BackendRuntimeLookup` and reads each arm's `(backend, device)` fit
    // status in `select`. `lookup` is part of the picker's documented
    // contract (06-runtime §"Route picker") and is the same handle the
    // `RouteCache` fingerprints telemetry through; the per-arm guard
    // itself lives in the selector, so `pick_route` does not re-query it.
    let _ = lookup;
    let mut picked = PickedRoute::new();
    let branches = branches_in_topo_order(graph, roots);

    for branch in branches {
        // Only record a non-arm-0 pick; arm 0 is the default in
        // `lower_picked_route`, so keeping the route minimal makes the
        // "no pressure ⇒ empty route ⇒ Phase B behavior" contract direct.
        let chosen = resolve_branch(graph, branch, bindings, selector);
        if chosen != 0 {
            picked.insert(branch, chosen);
        }
    }

    picked
}

/// Resolve **one** `Op::Branch` to its chosen arm index — the body of
/// [`pick_route`]'s per-branch loop, lifted out so the **streaming
/// run-walk** (Step E Phase C, PR C1) can resolve branches **lazily, as
/// the frontier reaches each one**, rather than resolving the whole route
/// up front (06-runtime §"selecting arms by live device load … per-
/// decision-point during the dispatch walk").
///
/// This is the SINGLE arm-decision primitive: [`pick_route`] (the eager
/// lane — [`RouteCache`] + the route-picker tests) calls it in
/// topological order over every branch; the streaming lowering
/// ([`fuel_graph::lower_picked_route_streaming`]) calls it one branch at a
/// time at the branch's arm-entry boundary. Because it consults the **same
/// [`RuntimeSelector`]** either way — and the production VRAM-pressure
/// chain reads only free-memory state (not walk progress) — the lazy and
/// eager resolutions are byte-identical on every input. C1 changes *when*
/// a branch resolves, never *which* arm it picks; the load signal that
/// will make the moment matter is C2.
///
/// A single-arm (or zero-arm) branch has no decision — returns arm 0
/// (`lower_picked_route`'s default), so it never enters the route map.
pub fn resolve_branch(
    graph: &Graph,
    branch: NodeId,
    bindings: &KernelBindingTable,
    selector: &dyn RuntimeSelector,
) -> usize {
    let arms = graph.node(branch).inputs.clone();
    if arms.len() < 2 {
        return 0;
    }
    pick_arm(graph, &arms, bindings, selector)
}

/// Choose one arm index for a single branch. Builds a per-arm
/// [`AlternativeSet`] (one [`Candidate`] per arm carrying the arm's real
/// `(backend, device)` + plan cost + the fork's decision context), runs
/// the selector, and maps the picked candidate's `(backend, device)`
/// back to its arm index.
fn pick_arm(
    graph: &Graph,
    arms: &[NodeId],
    bindings: &KernelBindingTable,
    selector: &dyn RuntimeSelector,
) -> usize {
    // Step D: re-derive the fork's candidate set from the graph + the runtime
    // binding registry (was the cached `plan.alternatives(fork)`). The fork
    // node = arm-0's exit; enumerating its `(op, dtypes)` across the arms'
    // placements reproduces the per-placement Candidates + the DecisionContext
    // the selector's VRAM-pressure guard + Judge-aware rank read. We mirror
    // them into a per-arm set in ARM ORDER so the picked index maps straight
    // back to an arm.
    let fork = arms[0];
    let fork_set = enumerate_fork_set(graph, fork, arms, bindings);

    let mut arm_set = AlternativeSet::empty();
    for &arm in arms {
        let backend = graph.target_backend(arm).unwrap_or(BackendId::Cpu);
        let device = default_device_for(backend);
        let candidate = candidate_for_placement(&fork_set, backend, device)
            .cloned()
            .unwrap_or_else(|| synthetic_candidate(backend, device));
        arm_set.push(candidate);
    }
    // Carry the fork's decision context so the Judge-aware rank leg can
    // re-query measurements per arm.
    if let Some(ctx) = fork_set.context() {
        arm_set.set_context(*ctx);
    }

    // The selector picks one candidate; arm index = its position in the
    // arm-ordered set. `select` returns `None` only on an empty set
    // (impossible here — arms.len() >= 2), in which case arm 0.
    match selector.select(&arm_set) {
        Some(picked) => arm_index_of(&arm_set, picked),
        None => 0,
    }
}

/// Re-derive the fork node's candidate set from the graph + the runtime
/// binding registry — the dispatch-time replacement for the cached
/// `ExecutionPlan::alternatives(fork)` (Step D). Enumerates the fork's
/// `(op, dtypes)` across the arms' `(backend, device)` placements and stamps
/// the fork's [`DecisionContext`] so the selector's Judge leg keys on the same
/// `(op, dtype, size_class)` triple the planner used. A non-kernel fork (no
/// `OpKind` mapping) yields an empty set ⇒ the per-arm loop falls back to
/// [`synthetic_candidate`], preserving the prior behavior when the plan
/// carried no entry for the fork. Layer-1 cost only (the candidates' default
/// `static_cost`); the selector's own Judge leg supplies the Layer-2
/// measurement, so no cost composition is needed here.
fn enumerate_fork_set(
    graph: &Graph,
    fork: NodeId,
    arms: &[NodeId],
    bindings: &KernelBindingTable,
) -> AlternativeSet {
    let node = graph.node(fork);
    let Some(op_kind) = op_to_op_kind(&node.op) else {
        return AlternativeSet::empty();
    };
    let dtypes = build_lookup_dtypes(graph, node);
    let placements: Vec<(BackendId, DeviceLocation)> = arms
        .iter()
        .map(|&arm| {
            let backend = graph.target_backend(arm).unwrap_or(BackendId::Cpu);
            (backend, default_device_for(backend))
        })
        .collect();
    // `op_params` is unread by the selector (it ranks on backend/device/cost/
    // kernel_source), so `None` suffices — the chosen arm is lowered to real
    // runs and dispatched through the normal binding-table path regardless.
    let mut set = enumerate_candidates(op_kind, &dtypes, &placements, &OpParams::None, bindings);
    if let Some(ctx) = fork_decision_context(graph, node, op_kind, &dtypes) {
        set.set_context(ctx);
    }
    set
}

/// The fork's [`DecisionContext`] — `(op, principal dtype, size class)`,
/// the triple the Judge keys measurements on. Mirrors `compile_plan`'s
/// stamping (plan.rs) exactly — `SizeClass::for_op` over the input
/// operand shapes — so dispatch-time re-derivation hits the same Judge
/// entries (a matmul keys on its `(m,n,k)` aspect, aligned with the
/// planner and producer).
fn fork_decision_context(
    graph: &Graph,
    node: &fuel_graph::Node,
    op_kind: OpKind,
    dtypes: &[DType],
) -> Option<DecisionContext> {
    let &principal_dtype = dtypes.first()?;
    let input_shapes: Vec<Shape> = node
        .inputs
        .iter()
        .map(|&input_id| graph.node(input_id).shape.clone())
        .collect();
    let size_class = SizeClass::for_op(op_kind, &input_shapes);
    Some(DecisionContext {
        op: op_kind,
        principal_dtype,
        size_class,
    })
}

/// Find a plan candidate matching `(backend, device)` (the arm's real
/// placement). Returns the first match — sibling kernels at the same
/// `(backend, device)` are interchangeable for the arm-selection purpose
/// (they differ only by `kernel_source`, which the selector still ranks).
fn candidate_for_placement(
    set: &AlternativeSet,
    backend: BackendId,
    device: DeviceLocation,
) -> Option<&Candidate> {
    set.alternatives()
        .iter()
        .find(|c| c.backend == backend && c.device == device)
        .or_else(|| {
            // Fall back to a backend-only match (the plan may carry a
            // different GPU ordinal than the arm's default device).
            set.alternatives().iter().find(|c| c.backend == backend)
        })
}

/// The arm index of `picked` within the arm-ordered `set` — by
/// `(backend, device)` identity. The set was built in arm order, so the
/// first matching position is the chosen arm.
fn arm_index_of(set: &AlternativeSet, picked: &Candidate) -> usize {
    set.alternatives()
        .iter()
        .position(|c| c.backend == picked.backend && c.device == picked.device)
        .unwrap_or(0)
}

/// A minimal stand-in [`Candidate`] for an arm with no matching plan
/// entry: real `(backend, device)`, zero cost, no-op kernel. The
/// selector reads only `backend` / `device` / `static_cost` /
/// `inbound_transfer_ns` / `kernel_source` for its pick, so the kernel
/// pointer is never invoked — this candidate exists purely to carry the
/// arm's placement to the VRAM-pressure guard.
fn synthetic_candidate(backend: BackendId, device: DeviceLocation) -> Candidate {
    Candidate {
        kernel: noop_kernel,
        caps: KernelCaps::empty(),
        backend,
        device,
        precision: PrecisionGuarantee::PRIMITIVE_DETERMINISTIC_CPU,
        static_cost: CostEstimate {
            flops: 0,
            bytes_moved: 0,
            kernel_overhead_ns: 0,
        },
        inbound_transfer_ns: 0,
        op_params: OpParams::None,
        coupling: Vec::new(),
        kernel_source: "route-picker-synthetic",
        kernel_revision_hash: 0,
    }
}

/// Never-invoked kernel for [`synthetic_candidate`] (see its docs). The
/// picker reads placement metadata only; the chosen arm is lowered to
/// runs and dispatched through the normal binding-table path, not via
/// this pointer.
fn noop_kernel(
    _i: &[Arc<std::sync::RwLock<fuel_memory::Storage>>],
    _o: &mut [Arc<std::sync::RwLock<fuel_memory::Storage>>],
    _l: &[fuel_ir::Layout],
    _p: &OpParams,
) -> fuel_ir::Result<()> {
    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::ranker::{BackendRuntimeHandle, ChainedSelector, WinnerSelector};
    use fuel_backend_contract::backend::BackendRuntime;
    use fuel_graph::{Node, Op};
    use fuel_ir::{DType, Shape};

    fn f32_node(g: &mut Graph, op: Op, inputs: Vec<NodeId>) -> NodeId {
        g.push(Node {
            op,
            inputs,
            shape: Shape::from_dims(&[2]),
            dtype: DType::F32,
        })
    }

    /// Build a 2-arm diamond with arm-0 on `arm0_backend` and arm-1 on
    /// `arm1_backend`. Returns `(graph, branch, arm0, arm1, post)`.
    ///
    /// Topology mirrors the run.rs diamond: `diverge -> {arm0, arm1}`;
    /// `reconverge` reads arm0 (the runnability invariant); the Branch
    /// merges {arm0, arm1} at `reconverge`. The arm exits carry their
    /// `target_backend` so the picker can read each arm's placement.
    fn diamond(
        arm0_backend: BackendId,
        arm1_backend: BackendId,
    ) -> (Graph, NodeId, NodeId, NodeId, NodeId) {
        let mut g = Graph::new();
        let pre = f32_node(&mut g, Op::Const, vec![]);
        let diverge = f32_node(&mut g, Op::Relu, vec![pre]);
        let arm0 = f32_node(&mut g, Op::Silu, vec![diverge]);
        let arm1 = f32_node(&mut g, Op::GeluTanh, vec![diverge]);
        g.set_target_backend(arm0, arm0_backend);
        g.set_target_backend(arm1, arm1_backend);
        let reconverge = f32_node(&mut g, Op::Relu, vec![arm0]);
        let mut b = g.open_branch(diverge);
        b.add_arm(arm0);
        b.add_arm(arm1);
        let branch = b
            .finalize_branches(&mut g, reconverge)
            .expect("well-formed 2-arm branch")
            .expect("2 arms survive");
        let post = f32_node(&mut g, Op::Tanh, vec![reconverge]);
        (g, branch, arm0, arm1, post)
    }

    /// A binding table with the diamond's fork op (`Op::Silu` ⇒
    /// `SiluElementwise`) registered at both arm backends, so `pick_arm`
    /// re-enumerates a real candidate per placement (Step D — was a hand-built
    /// `ExecutionPlan`). No-op kernels: the picker reads only placement
    /// metadata + `kernel_source`, never invoking the pointer.
    fn diamond_bindings() -> KernelBindingTable {
        let mut t = KernelBindingTable::new();
        for backend in [BackendId::Cuda, BackendId::Cpu] {
            t.register(
                OpKind::SiluElementwise,
                &[DType::F32, DType::F32],
                backend,
                noop_kernel,
            );
        }
        t
    }

    struct MockRuntime {
        available: Option<u64>,
        total: Option<u64>,
    }
    impl BackendRuntime for MockRuntime {
        fn available_bytes(&self) -> Option<u64> {
            self.available
        }
        fn total_bytes(&self) -> Option<u64> {
            self.total
        }
    }

    /// Per-backend `(available, total)` lookup; backends not listed
    /// resolve to `None` (= Unknown / no signal).
    fn lookup_for(entries: Vec<(BackendId, Option<u64>, Option<u64>)>) -> BackendRuntimeLookup {
        Arc::new(move |b, _d| {
            entries
                .iter()
                .find(|(eb, _, _)| *eb == b)
                .map(|&(_, a, t)| {
                    Box::new(MockRuntime {
                        available: a,
                        total: t,
                    }) as BackendRuntimeHandle
                })
        })
    }

    /// (a) NO PRESSURE / NO TELEMETRY ⇒ the picker chooses arm-0 (the
    /// winner) at every branch ⇒ an empty route ⇒ realize is identical to
    /// Phase B. Tested with both the bare `WinnerSelector` and the
    /// production `ChainedSelector` with no signals.
    #[test]
    fn no_telemetry_picks_arm0_empty_route() {
        let (g, _branch, _fork, _arm1, post) = diamond(BackendId::Cuda, BackendId::Cpu);
        let bindings = diamond_bindings();

        // Bare winner selector: always arm 0.
        let winner = WinnerSelector;
        let route = pick_route(&g, &[post], &bindings, &winner, None);
        assert!(
            route.is_empty(),
            "no telemetry ⇒ arm-0 everywhere ⇒ empty route; got {route:?}",
        );

        // Production chained selector with NO signals (no judge, no
        // lookup) degrades to the static winner = arm 0.
        let chained = ChainedSelector::with_default_estimator(None, None);
        let route = pick_route(&g, &[post], &bindings, &chained, None);
        assert!(
            route.is_empty(),
            "ChainedSelector with no signals ⇒ arm-0 everywhere; got {route:?}",
        );
    }

    /// (b) SIMULATED VRAM PRESSURE: the GPU arm's backend reports tiny
    /// free VRAM (WontFit), so the chained selector's guard skips it and
    /// the picker takes the host-RAM (CPU) arm instead — arm 1.
    #[test]
    fn vram_pressure_picks_host_ram_arm() {
        let (g, branch, _fork, _arm1, post) = diamond(BackendId::Cuda, BackendId::Cpu);
        // Both placements resolve to the same no-op kernel; the only signal is
        // memory pressure, so the guard alone must flip the pick.
        let bindings = diamond_bindings();

        // CUDA has 1 byte free of 10_000 → WontFit for the default
        // bytes_moved estimate; CPU has ample room → Comfortable.
        let lookup = lookup_for(vec![
            (BackendId::Cuda, Some(1), Some(10_000)),
            (BackendId::Cpu, Some(8_000), Some(10_000)),
        ]);
        let chained = ChainedSelector::with_default_estimator(None, Some(lookup.clone()));

        let route = pick_route(&g, &[post], &bindings, &chained, Some(&lookup));
        assert_eq!(
            route.get(&branch).copied(),
            Some(1),
            "under VRAM pressure the picker takes the host-RAM (CPU) arm; \
             route={route:?}",
        );
    }

    /// (c) TWO COUPLED BRANCHES resolve consistently in topological
    /// order: both GPU arms are under pressure, so the picker takes the
    /// CPU arm at BOTH branches — the downstream branch sees the same
    /// pressure the upstream one did and resolves the same way.
    #[test]
    fn two_coupled_branches_resolve_consistently() {
        // Two diamonds chained: post0 feeds the second diverge.
        let mut g = Graph::new();
        let pre = f32_node(&mut g, Op::Const, vec![]);
        // --- branch 1 ---
        let div1 = f32_node(&mut g, Op::Relu, vec![pre]);
        let a0_1 = f32_node(&mut g, Op::Silu, vec![div1]);
        let a1_1 = f32_node(&mut g, Op::GeluTanh, vec![div1]);
        g.set_target_backend(a0_1, BackendId::Cuda);
        g.set_target_backend(a1_1, BackendId::Cpu);
        let recon1 = f32_node(&mut g, Op::Relu, vec![a0_1]);
        let mut b1 = g.open_branch(div1);
        b1.add_arm(a0_1);
        b1.add_arm(a1_1);
        let branch1 = b1
            .finalize_branches(&mut g, recon1)
            .expect("branch1 valid")
            .expect("2 arms");
        // --- branch 2 (downstream of branch 1's merge) ---
        let div2 = f32_node(&mut g, Op::Tanh, vec![recon1]);
        let a0_2 = f32_node(&mut g, Op::Silu, vec![div2]);
        let a1_2 = f32_node(&mut g, Op::GeluTanh, vec![div2]);
        g.set_target_backend(a0_2, BackendId::Cuda);
        g.set_target_backend(a1_2, BackendId::Cpu);
        let recon2 = f32_node(&mut g, Op::Relu, vec![a0_2]);
        let mut b2 = g.open_branch(div2);
        b2.add_arm(a0_2);
        b2.add_arm(a1_2);
        let branch2 = b2
            .finalize_branches(&mut g, recon2)
            .expect("branch2 valid")
            .expect("2 arms");
        let post = f32_node(&mut g, Op::Tanh, vec![recon2]);

        // Topo order must place branch1 before branch2.
        let order = branches_in_topo_order(&g, &[post]);
        assert_eq!(
            order,
            vec![branch1, branch2],
            "coupled branches resolve upstream-first in topo order",
        );

        // Both forks (`Op::Silu`) re-enumerate from the same registered
        // bindings (Silu at both arm backends).
        let bindings = diamond_bindings();

        // GPU under pressure ⇒ both branches pick the CPU arm (arm 1).
        let lookup = lookup_for(vec![
            (BackendId::Cuda, Some(1), Some(10_000)),
            (BackendId::Cpu, Some(8_000), Some(10_000)),
        ]);
        let chained = ChainedSelector::with_default_estimator(None, Some(lookup.clone()));
        let route = pick_route(&g, &[post], &bindings, &chained, Some(&lookup));

        assert_eq!(route.get(&branch1).copied(), Some(1), "branch1 → CPU arm");
        assert_eq!(route.get(&branch2).copied(), Some(1), "branch2 → CPU arm");
    }

    /// (d) The route is CACHED and re-resolved only on a meaningful
    /// telemetry delta. Same fingerprint ⇒ cache hit (no re-resolve); a
    /// genuine pressure transition ⇒ re-resolve + a possibly different
    /// route.
    #[test]
    fn route_cached_and_reresolved_on_telemetry_delta() {
        let (g, branch, _fork, _arm1, post) = diamond(BackendId::Cuda, BackendId::Cpu);
        let bindings = diamond_bindings();

        // Start comfortable: both backends have ample free memory.
        let comfy = lookup_for(vec![
            (BackendId::Cuda, Some(8_000), Some(10_000)),
            (BackendId::Cpu, Some(8_000), Some(10_000)),
        ]);
        let chained_comfy = ChainedSelector::with_default_estimator(None, Some(comfy.clone()));

        let mut cache = RouteCache::new();
        let r1 = cache.resolve(&g, &[post], &bindings, &chained_comfy, Some(&comfy));
        assert!(r1.is_empty(), "comfortable ⇒ arm-0 (winner) route");
        assert_eq!(cache.resolves, 1, "first resolve populates the cache");
        assert_eq!(cache.hits, 0);

        // Same telemetry ⇒ cache hit, no re-resolve.
        let r2 = cache.resolve(&g, &[post], &bindings, &chained_comfy, Some(&comfy));
        assert_eq!(r2, r1, "stable telemetry reuses the cached route");
        assert_eq!(cache.resolves, 1, "no re-resolve on a stable fingerprint");
        assert_eq!(cache.hits, 1, "stable telemetry is a cache hit");

        // Meaningful delta: CUDA free memory collapses to 1 byte
        // (WontFit). Re-resolve fires and the route flips to the CPU arm.
        let pressured = lookup_for(vec![
            (BackendId::Cuda, Some(1), Some(10_000)),
            (BackendId::Cpu, Some(8_000), Some(10_000)),
        ]);
        let chained_pressured =
            ChainedSelector::with_default_estimator(None, Some(pressured.clone()));
        let r3 = cache.resolve(&g, &[post], &bindings, &chained_pressured, Some(&pressured));
        assert_eq!(
            cache.resolves, 2,
            "a meaningful telemetry delta re-resolves"
        );
        assert_eq!(cache.hits, 1, "the delta is not counted as a hit");
        assert_eq!(
            r3.get(&branch).copied(),
            Some(1),
            "after the pressure transition the route picks the CPU arm",
        );
    }

    /// A graph with ZERO branches (a CPU-only build) ⇒ the picker is a
    /// no-op ⇒ the route is empty ⇒ realize is unchanged. The cache
    /// resolves it once and caches the empty route.
    #[test]
    fn branchless_graph_is_a_noop() {
        let mut g = Graph::new();
        let a = f32_node(&mut g, Op::Const, vec![]);
        let b = f32_node(&mut g, Op::Relu, vec![a]);
        let c = f32_node(&mut g, Op::Tanh, vec![b]);
        let bindings = KernelBindingTable::new();
        let winner = WinnerSelector;
        let route = pick_route(&g, &[c], &bindings, &winner, None);
        assert!(route.is_empty(), "no branches ⇒ empty route (picker no-op)");
    }

    /// Step D: the fork's candidate set is re-enumerated from the runtime
    /// registry (not a cached plan). With the fork op registered at both arm
    /// backends, the set carries a real candidate per placement + the fork's
    /// `DecisionContext` (so the Judge leg keys on the same triple the planner
    /// used).
    #[test]
    fn enumerate_fork_set_pulls_registered_candidates() {
        let (g, _branch, fork, arm1, _post) = diamond(BackendId::Cuda, BackendId::Cpu);
        let bindings = diamond_bindings();
        let set = enumerate_fork_set(&g, fork, &[fork, arm1], &bindings);

        assert_eq!(
            set.alternatives().len(),
            2,
            "Silu registered at both arm backends ⇒ one candidate per placement",
        );
        assert!(
            set.alternatives()
                .iter()
                .any(|c| c.backend == BackendId::Cuda)
        );
        assert!(
            set.alternatives()
                .iter()
                .any(|c| c.backend == BackendId::Cpu)
        );
        let ctx = set.context().expect("fork DecisionContext stamped");
        assert_eq!(ctx.op, OpKind::SiluElementwise);
        assert_eq!(ctx.principal_dtype, DType::F32);
    }

    /// A fork op with no registered binding falls back to placement-only
    /// synthetic candidates — they still carry each arm's `(backend, device)`,
    /// so the VRAM-pressure guard still flips to the CPU arm. Preserves the
    /// prior no-plan-entry behavior.
    #[test]
    fn synthetic_fallback_when_fork_op_unregistered() {
        let (g, branch, _fork, _arm1, post) = diamond(BackendId::Cuda, BackendId::Cpu);
        let empty = KernelBindingTable::new();
        let lookup = lookup_for(vec![
            (BackendId::Cuda, Some(1), Some(10_000)),
            (BackendId::Cpu, Some(8_000), Some(10_000)),
        ]);
        let chained = ChainedSelector::with_default_estimator(None, Some(lookup.clone()));
        let route = pick_route(&g, &[post], &empty, &chained, Some(&lookup));
        assert_eq!(
            route.get(&branch).copied(),
            Some(1),
            "no binding ⇒ synthetic placement candidates ⇒ guard still picks CPU",
        );
    }

    // ===== Step E Phase C / C2: live-load arm re-pick (the headline gate) =====

    /// Silu bindings at CUDA + Vulkan (the two streaming arms of the
    /// load-gate diamond). No-op kernels — the picker reads only placement
    /// metadata for the arm decision.
    fn diamond_bindings_cuda_vulkan() -> KernelBindingTable {
        let mut t = KernelBindingTable::new();
        for backend in [BackendId::Cuda, BackendId::Vulkan] {
            t.register(
                OpKind::SiluElementwise,
                &[DType::F32, DType::F32],
                backend,
                noop_kernel,
            );
        }
        t
    }

    /// THE C2 HEADLINE GATE (selection logic, CPU — no GPU needed). A
    /// 2-device branched graph (arm-0 on CUDA, arm-1 on Vulkan, BOTH
    /// VRAM-fit) with a FAKE `BackendStreams` lookup reporting HIGH
    /// pending_work_count on arm-0's device (CUDA, saturated) and LOW on
    /// arm-1's (Vulkan, idle): the load-aware `ChainedSelector` re-picks
    /// the UNLOADED arm-1, and the chosen arm's `target_backend` is the
    /// unloaded device (Vulkan). Mirrors `vram_pressure_picks_host_ram_arm`
    /// but with a LOAD signal instead of a memory signal — the live-load
    /// re-pick that is the whole point of Step E, proven deterministically
    /// with a fake signal (the real signal is B1, already verified).
    #[test]
    fn load_pressure_picks_unloaded_arm() {
        let (g, branch, _arm0, arm1, post) = diamond(BackendId::Cuda, BackendId::Vulkan);
        let bindings = diamond_bindings_cuda_vulkan();

        // Both arms VRAM-fit (ample free / total). CUDA saturated (4 in
        // flight / 1 slot), Vulkan idle (0). The ONLY differentiator is
        // load — VRAM ties, cost ties (no-op kernels, same static cost),
        // so the load leg alone must flip the pick to arm-1.
        let lookup = crate::ranker::mock_combined_lookup(vec![
            (BackendId::Cuda, Some(10_000), Some(10_000), Some(4), 1),
            (BackendId::Vulkan, Some(10_000), Some(10_000), Some(0), 1),
        ]);
        let chained = ChainedSelector::with_default_estimator(None, Some(lookup.clone()));

        let route = pick_route(&g, &[post], &bindings, &chained, Some(&lookup));
        assert_eq!(
            route.get(&branch).copied(),
            Some(1),
            "under simulated CUDA saturation the picker takes the UNLOADED \
             Vulkan arm (arm-1); route={route:?}",
        );
        // The chosen arm's node lands on the unloaded device.
        assert_eq!(
            g.target_backend(arm1),
            Some(BackendId::Vulkan),
            "the picked arm's target_backend is the unloaded device",
        );
    }

    /// No-load mirror of the headline gate (the degenerate-fallback): the
    /// SAME 2-device diamond with BOTH devices reporting 0 in flight ⇒
    /// every load_tier is 0 ⇒ the route is arm-0 everywhere ⇒ empty route,
    /// byte-identical to `no_telemetry_picks_arm0_empty_route`. Load
    /// changes nothing unless devices genuinely contend.
    #[test]
    fn no_load_picks_arm0_empty_route() {
        let (g, _branch, _arm0, _arm1, post) = diamond(BackendId::Cuda, BackendId::Vulkan);
        let bindings = diamond_bindings_cuda_vulkan();

        // Both VRAM-fit, both idle (0 in flight) → no load signal anywhere.
        let lookup = crate::ranker::mock_combined_lookup(vec![
            (BackendId::Cuda, Some(10_000), Some(10_000), Some(0), 1),
            (BackendId::Vulkan, Some(10_000), Some(10_000), Some(0), 1),
        ]);
        let chained = ChainedSelector::with_default_estimator(None, Some(lookup.clone()));

        let route = pick_route(&g, &[post], &bindings, &chained, Some(&lookup));
        assert!(
            route.is_empty(),
            "no load anywhere ⇒ arm-0 everywhere ⇒ empty route; got {route:?}",
        );
    }
}