use crate::{Graph, NodeId, Op, topo_order_multi};
use fuel_ir::probe::BackendId;
use std::collections::{HashMap, HashSet};
pub type PickedRoute = HashMap<NodeId, usize>;
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Run {
pub entry: NodeId,
pub exit: NodeId,
pub members: Vec<NodeId>,
pub device: Option<BackendId>,
}
pub fn extract_runs(graph: &Graph, root: NodeId) -> Vec<Run> {
extract_runs_multi(graph, &[root])
}
pub fn extract_runs_multi(graph: &Graph, roots: &[NodeId]) -> Vec<Run> {
let eff_roots = effective_roots(graph, roots);
let order = topo_order_multi(graph, &eff_roots);
let reachable: HashSet<NodeId> = order.iter().copied().collect();
let reconverge_points: HashSet<NodeId> = order
.iter()
.filter_map(|&id| match graph.node(id).op {
Op::Branch { reconverge_at } => Some(reconverge_at),
_ => None,
})
.collect();
let arm_entries = compute_arm_entries(graph, &order, &reachable);
let mut pred_inside: HashMap<NodeId, Vec<NodeId>> = HashMap::new();
for &id in &order {
let preds: Vec<NodeId> = graph
.node(id)
.inputs
.iter()
.copied()
.filter(|p| reachable.contains(p))
.collect();
pred_inside.insert(id, preds);
}
let mut runs: Vec<Run> = Vec::new();
let mut current: Vec<NodeId> = Vec::new();
let mut current_device: Option<BackendId> = None;
for &id in &order {
if matches!(graph.node(id).op, Op::Branch { .. }) {
continue;
}
let preds = &pred_inside[&id];
let device = graph.target_backend(id);
let starts_new = {
let is_entry = preds.is_empty();
let is_reconverge = reconverge_points.contains(&id);
let is_arm = arm_entries.contains(&id);
let is_fan_in = preds.len() > 1;
let is_transfer = matches!(graph.node(id).op, Op::Copy { .. } | Op::Move { .. });
let is_seam = preds.len() == 1 && graph.target_backend(preds[0]) != device;
let breaks_chain = match current.last() {
None => true,
Some(&last) => !(preds.len() == 1 && preds[0] == last),
};
let device_mismatch = !current.is_empty() && device != current_device;
is_entry
|| is_reconverge
|| is_arm
|| is_fan_in
|| is_transfer
|| is_seam
|| breaks_chain
|| device_mismatch
};
if starts_new && !current.is_empty() {
runs.push(finish_run(¤t, current_device));
current = Vec::new();
}
if current.is_empty() {
current_device = device;
}
current.push(id);
}
if !current.is_empty() {
runs.push(finish_run(¤t, current_device));
}
runs
}
pub fn effective_roots(graph: &Graph, roots: &[NodeId]) -> Vec<NodeId> {
let mut seeds: Vec<NodeId> = roots.to_vec();
let mut seen: HashSet<NodeId> = seeds.iter().copied().collect();
loop {
let reachable: HashSet<NodeId> = topo_order_multi(graph, &seeds).into_iter().collect();
let mut added = false;
for idx in 0..graph.len() {
let id = NodeId(idx);
if seen.contains(&id) {
continue;
}
let Op::Branch { reconverge_at } = graph.node(id).op else {
continue;
};
let participates = reachable.contains(&reconverge_at)
|| graph.node(id).inputs.iter().any(|a| reachable.contains(a));
if participates {
seen.insert(id);
seeds.push(id);
added = true;
}
}
if !added {
break;
}
}
seeds
}
pub fn branches_in_topo_order(graph: &Graph, roots: &[NodeId]) -> Vec<NodeId> {
let eff_roots = effective_roots(graph, roots);
let order = topo_order_multi(graph, &eff_roots);
let position: HashMap<NodeId, usize> =
order.iter().enumerate().map(|(i, &id)| (id, i)).collect();
let mut branches: Vec<(usize, NodeId)> = Vec::new();
for idx in 0..graph.len() {
let id = NodeId(idx);
let Op::Branch { reconverge_at } = graph.node(id).op else {
continue;
};
if let Some(&pos) = position.get(&reconverge_at) {
branches.push((pos, id));
}
}
branches.sort_by_key(|&(pos, _)| pos);
branches.into_iter().map(|(_, id)| id).collect()
}
fn finish_run(members: &[NodeId], device: Option<BackendId>) -> Run {
Run {
entry: members[0],
exit: *members.last().expect("a run has at least one member"),
members: members.to_vec(),
device,
}
}
fn compute_arm_entries(
graph: &Graph,
order: &[NodeId],
reachable: &HashSet<NodeId>,
) -> HashSet<NodeId> {
let mut arm_entries: HashSet<NodeId> = HashSet::new();
for &id in order {
let Op::Branch { .. } = graph.node(id).op else {
continue;
};
let arm_exits: Vec<NodeId> = graph.node(id).inputs.clone();
if arm_exits.len() < 2 {
continue;
}
let cones: Vec<HashSet<NodeId>> = arm_exits
.iter()
.map(|&e| backward_cone(graph, e, reachable))
.collect();
let mut shared: HashSet<NodeId> = cones[0].clone();
for c in &cones[1..] {
shared = shared.intersection(c).copied().collect();
}
for cone in &cones {
for &n in cone {
if shared.contains(&n) {
continue;
}
let departs = graph.node(n).inputs.iter().any(|p| shared.contains(p));
if departs {
arm_entries.insert(n);
}
}
}
}
arm_entries
}
fn backward_cone(graph: &Graph, from: NodeId, reachable: &HashSet<NodeId>) -> HashSet<NodeId> {
let mut seen: HashSet<NodeId> = HashSet::new();
let mut stack = vec![from];
while let Some(n) = stack.pop() {
if !reachable.contains(&n) || !seen.insert(n) {
continue;
}
for &inp in &graph.node(n).inputs {
if !seen.contains(&inp) {
stack.push(inp);
}
}
}
seen
}
pub fn lower_run(run: &Run) -> &[NodeId] {
&run.members
}
pub fn lower_runs_arm0(graph: &Graph, roots: &[NodeId]) -> Vec<NodeId> {
lower_picked_route(graph, roots, &PickedRoute::new())
}
pub fn lower_picked_route(graph: &Graph, roots: &[NodeId], picked: &PickedRoute) -> Vec<NodeId> {
let runs = extract_runs_multi(graph, roots);
let skip = non_chosen_arm_nodes(graph, roots, picked);
let kept: Vec<Run> = runs
.into_iter()
.filter(|run| !skip.contains(&run.entry))
.collect();
let perm = device_alternating_order(graph, &kept);
let mut order = Vec::new();
for &ri in &perm {
order.extend_from_slice(lower_run(&kept[ri]));
}
order
}
pub fn device_alternating_order(graph: &Graph, runs: &[Run]) -> Vec<usize> {
let n = runs.len();
if n <= 1 {
return (0..n).collect();
}
let mut distinct: Vec<Option<BackendId>> = Vec::new();
for r in runs {
if !distinct.contains(&r.device) {
distinct.push(r.device);
}
}
if distinct.len() <= 1 {
return (0..n).collect();
}
let mut run_of: HashMap<NodeId, usize> = HashMap::new();
for (ri, r) in runs.iter().enumerate() {
for &m in &r.members {
run_of.insert(m, ri);
}
}
let mut preds: Vec<HashSet<usize>> = vec![HashSet::new(); n];
let mut succs: Vec<Vec<usize>> = vec![Vec::new(); n];
for (ri, r) in runs.iter().enumerate() {
for &m in &r.members {
for &inp in &graph.node(m).inputs {
if let Some(&pi) = run_of.get(&inp)
&& pi != ri
&& preds[ri].insert(pi)
{
succs[pi].push(ri);
}
}
}
}
let mut indeg: Vec<usize> = preds.iter().map(|p| p.len()).collect();
let compute_size = |ri: usize| -> u64 {
if matches!(
graph.node(runs[ri].entry).op,
Op::Copy { .. } | Op::Move { .. }
) {
0
} else {
runs[ri].members.len() as u64
}
};
let mut down_weight: Vec<u64> = vec![u64::MAX; n]; for start in 0..n {
if down_weight[start] != u64::MAX {
continue;
}
let mut stack: Vec<(usize, bool)> = vec![(start, false)];
while let Some((r, done)) = stack.pop() {
if done {
let mut w = compute_size(r);
for &s in &succs[r] {
w = w.max(down_weight[s]);
}
down_weight[r] = w;
continue;
}
if down_weight[r] != u64::MAX {
continue;
}
stack.push((r, true));
for &s in &succs[r] {
if down_weight[s] == u64::MAX {
stack.push((s, false));
}
}
}
}
let mut emitted: Vec<bool> = vec![false; n];
let mut ready: Vec<usize> = (0..n).filter(|&i| indeg[i] == 0).collect();
let mut order: Vec<usize> = Vec::with_capacity(n);
let mut last_device: Option<Option<BackendId>> = None;
while !ready.is_empty() {
let key_of = |ri: usize| -> (std::cmp::Reverse<u64>, bool, usize) {
let same_device = last_device.is_some_and(|ld| runs[ri].device == ld);
(std::cmp::Reverse(down_weight[ri]), same_device, ri)
};
let mut best_pos = 0usize;
let mut best_key = key_of(ready[0]);
for (pos, &ri) in ready.iter().enumerate().skip(1) {
let key = key_of(ri);
if key < best_key {
best_key = key;
best_pos = pos;
}
}
let ri = ready.swap_remove(best_pos);
emitted[ri] = true;
order.push(ri);
last_device = Some(runs[ri].device);
for &si in &succs[ri] {
if emitted[si] {
continue;
}
indeg[si] -= 1;
if indeg[si] == 0 {
ready.push(si);
}
}
}
if order.len() != n {
#[expect(
clippy::needless_range_loop,
reason = "the bound `n` is the node count checked above (order.len() != n), not `emitted.len()`; enumerate could change the iteration count -- a wrong-number bug"
)]
for i in 0..n {
if !emitted[i] {
order.push(i);
}
}
}
order
}
pub fn lower_picked_route_streaming<F>(
graph: &Graph,
roots: &[NodeId],
mut resolve: F,
) -> Vec<NodeId>
where
F: FnMut(NodeId) -> usize,
{
let mut order = Vec::new();
walk_picked_route_streaming(graph, roots, &mut resolve, |members| {
order.extend_from_slice(members);
});
order
}
pub fn walk_picked_route_streaming<R, E>(
graph: &Graph,
roots: &[NodeId],
mut resolve: R,
mut emit: E,
) where
R: FnMut(NodeId) -> usize,
E: FnMut(&[NodeId]),
{
let runs = extract_runs_multi(graph, roots);
let branch_meta = branch_arm_meta(graph, roots);
let mut entry_to_branches: HashMap<NodeId, Vec<NodeId>> = HashMap::new();
for meta in &branch_meta {
for &(entry, _arm) in &meta.arm_entries {
entry_to_branches
.entry(entry)
.or_default()
.push(meta.branch);
}
}
let meta_of: HashMap<NodeId, &BranchArmMeta> =
branch_meta.iter().map(|m| (m.branch, m)).collect();
let mut skip: HashSet<NodeId> = HashSet::new();
let mut resolved: HashSet<NodeId> = HashSet::new();
let mut segment: Vec<&Run> = Vec::new();
let flush = |segment: &mut Vec<&Run>, emit: &mut E| {
if segment.is_empty() {
return;
}
let owned: Vec<Run> = segment.iter().map(|r| (*r).clone()).collect();
let perm = device_alternating_order(graph, &owned);
for &ri in &perm {
emit(lower_run(&owned[ri]));
}
segment.clear();
};
for run in &runs {
if let Some(branches) = entry_to_branches.get(&run.entry) {
let opens_new = branches.iter().any(|b| !resolved.contains(b));
if opens_new {
flush(&mut segment, &mut emit);
}
for &branch in branches {
if !resolved.insert(branch) {
continue;
}
let meta = meta_of[&branch];
let chosen = resolve(branch);
add_non_chosen_skip(meta, chosen, &mut skip);
}
}
if skip.contains(&run.entry) {
continue;
}
segment.push(run);
}
flush(&mut segment, &mut emit);
}
struct BranchArmMeta {
branch: NodeId,
arm_entries: Vec<(NodeId, usize)>,
cones: Vec<HashSet<NodeId>>,
shared: HashSet<NodeId>,
}
fn branch_arm_meta(graph: &Graph, roots: &[NodeId]) -> Vec<BranchArmMeta> {
let eff_roots = effective_roots(graph, roots);
let order = topo_order_multi(graph, &eff_roots);
let reachable: HashSet<NodeId> = order.iter().copied().collect();
let mut metas = Vec::new();
for &id in &order {
let Op::Branch { .. } = graph.node(id).op else {
continue;
};
let arm_exits = graph.node(id).inputs.clone();
if arm_exits.len() < 2 {
continue;
}
let cones: Vec<HashSet<NodeId>> = arm_exits
.iter()
.map(|&e| backward_cone(graph, e, &reachable))
.collect();
let mut shared: HashSet<NodeId> = cones[0].clone();
for c in &cones[1..] {
shared = shared.intersection(c).copied().collect();
}
let mut arm_entries: Vec<(NodeId, usize)> = Vec::new();
for (arm_idx, cone) in cones.iter().enumerate() {
for &n in cone {
if shared.contains(&n) {
continue;
}
let departs = graph.node(n).inputs.iter().any(|p| shared.contains(p));
if departs {
arm_entries.push((n, arm_idx));
}
}
}
metas.push(BranchArmMeta {
branch: id,
arm_entries,
cones,
shared,
});
}
metas
}
fn add_non_chosen_skip(meta: &BranchArmMeta, chosen: usize, skip: &mut HashSet<NodeId>) {
let chosen = if chosen < meta.cones.len() { chosen } else { 0 };
let chosen_cone = &meta.cones[chosen];
for (i, cone) in meta.cones.iter().enumerate() {
if i == chosen {
continue;
}
for &n in cone {
if !meta.shared.contains(&n) && !chosen_cone.contains(&n) {
skip.insert(n);
}
}
}
}
fn non_chosen_arm_nodes(graph: &Graph, roots: &[NodeId], picked: &PickedRoute) -> HashSet<NodeId> {
let eff_roots = effective_roots(graph, roots);
let order = topo_order_multi(graph, &eff_roots);
let reachable: HashSet<NodeId> = order.iter().copied().collect();
let mut skip: HashSet<NodeId> = HashSet::new();
for &id in &order {
let Op::Branch { .. } = graph.node(id).op else {
continue;
};
let arm_exits = &graph.node(id).inputs;
if arm_exits.len() < 2 {
continue;
}
let chosen = picked
.get(&id)
.copied()
.filter(|&c| c < arm_exits.len())
.unwrap_or(0);
let cones: Vec<HashSet<NodeId>> = arm_exits
.iter()
.map(|&e| backward_cone(graph, e, &reachable))
.collect();
let mut shared: HashSet<NodeId> = cones[0].clone();
for c in &cones[1..] {
shared = shared.intersection(c).copied().collect();
}
let chosen_cone = &cones[chosen];
for (i, cone) in cones.iter().enumerate() {
if i == chosen {
continue;
}
for &n in cone {
if !shared.contains(&n) && !chosen_cone.contains(&n) {
skip.insert(n);
}
}
}
}
skip
}
pub fn branch_density(graph: &Graph, root: NodeId) -> f32 {
branch_density_multi(graph, &[root])
}
pub fn branch_density_multi(graph: &Graph, roots: &[NodeId]) -> f32 {
let eff_roots = effective_roots(graph, roots);
let order = topo_order_multi(graph, &eff_roots);
if order.is_empty() {
return 0.0;
}
let branches = order
.iter()
.filter(|&&id| matches!(graph.node(id).op, Op::Branch { .. }))
.count();
branches as f32 / order.len() as f32
}
pub const FEWNESS_THRESHOLD: f32 = 0.05;
pub fn passes_fewness_gate(graph: &Graph, root: NodeId) -> bool {
branch_density(graph, root) < FEWNESS_THRESHOLD
}
#[cfg(test)]
mod tests {
use super::*;
use crate::Node;
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,
})
}
fn diamond_with_branch() -> (Graph, NodeId, NodeId, 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]);
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 finalizes")
.expect("2 arms survive");
let post = f32_node(&mut g, Op::Tanh, vec![reconverge]);
(g, diverge, arm0, arm1, reconverge, branch, post)
}
#[test]
fn straight_line_extracts_to_one_run() {
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::Silu, vec![b]);
let d = f32_node(&mut g, Op::Tanh, vec![c]);
let runs = extract_runs(&g, d);
assert_eq!(runs.len(), 1, "a straight-line graph is exactly one run");
let run = &runs[0];
assert_eq!(run.entry, a, "the run enters at the source");
assert_eq!(run.exit, d, "the run exits at the root");
assert_eq!(
run.members,
vec![a, b, c, d],
"the single run covers every reachable node in topo order",
);
assert_eq!(run.device, None);
assert_eq!(lower_run(run), &[a, b, c, d]);
}
#[test]
fn two_arm_branch_extracts_to_four_runs() {
let (g, diverge, arm0, arm1, reconverge, _branch, post) = diamond_with_branch();
let runs = extract_runs(&g, post);
for r in &runs {
assert_eq!(r.entry, r.members[0]);
assert_eq!(r.exit, *r.members.last().unwrap());
assert!(!r.members.is_empty());
}
let run_of = |node: NodeId| -> usize {
runs.iter()
.position(|r| r.members.contains(&node))
.unwrap_or_else(|| panic!("no run contains Node#{}", node.0))
};
let pre_run = run_of(diverge);
let arm0_run = run_of(arm0);
let arm1_run = run_of(arm1);
let post_run = run_of(reconverge);
let distinct: HashSet<usize> = [pre_run, arm0_run, arm1_run, post_run]
.into_iter()
.collect();
assert_eq!(
distinct.len(),
4,
"a 2-arm branch yields exactly {{pre, arm0, arm1, post}} = 4 runs; got runs={runs:?}",
);
assert_eq!(runs.len(), 4, "no run beyond the four; got {runs:?}");
assert_ne!(arm0_run, arm1_run, "the two arms must be separate runs");
assert_ne!(pre_run, arm0_run);
assert_ne!(pre_run, arm1_run);
assert!(runs[post_run].members.contains(&post));
}
#[test]
fn lower_runs_arm0_follows_arm0_and_skips_other_arms() {
let (g, diverge, arm0, arm1, reconverge, _branch, post) = diamond_with_branch();
let order = lower_runs_arm0(&g, &[post]);
assert!(
!order.contains(&arm1),
"arm-0 lowering skips arm 1's run; order={order:?} arm1={arm1:?}",
);
assert!(order.contains(&diverge), "the pre-run (diverge) executes");
assert!(order.contains(&arm0), "arm 0 executes");
assert!(order.contains(&reconverge), "the reconverge executes");
assert!(order.contains(&post), "the post-run executes");
for &id in &order {
assert!(
!matches!(g.node(id).op, Op::Branch { .. }),
"the Branch node is never an executable member",
);
}
let mut g2 = Graph::new();
let a = f32_node(&mut g2, Op::Const, vec![]);
let b = f32_node(&mut g2, Op::Relu, vec![a]);
let c = f32_node(&mut g2, Op::Silu, vec![b]);
let d = f32_node(&mut g2, Op::Tanh, vec![c]);
let flat: Vec<NodeId> = extract_runs(&g2, d)
.iter()
.flat_map(|r| lower_run(r).to_vec())
.collect();
assert_eq!(
lower_runs_arm0(&g2, &[d]),
flat,
"on a branchless graph the arm-0 lowering equals the flat concatenation",
);
}
#[test]
fn lower_picked_route_follows_chosen_arm_and_skips_others() {
let (g, _diverge, arm0, arm1, _reconverge, branch, post) = diamond_with_branch();
let empty = PickedRoute::new();
let route_order = lower_picked_route(&g, &[post], &empty);
let arm0_order = lower_runs_arm0(&g, &[post]);
assert_eq!(
route_order, arm0_order,
"the empty route lowers byte-identically to lower_runs_arm0",
);
assert!(route_order.contains(&arm0), "empty route follows arm 0");
assert!(
!route_order.contains(&arm1),
"empty route skips arm 1's run; order={route_order:?}",
);
let mut picked = PickedRoute::new();
picked.insert(branch, 1);
let order = lower_picked_route(&g, &[post], &picked);
assert!(
order.contains(&arm1),
"picking arm 1 follows arm 1's run; order={order:?}",
);
assert!(
!order.contains(&arm0),
"picking arm 1 skips arm 0's interior run; order={order:?}",
);
for &id in &order {
assert!(
!matches!(g.node(id).op, Op::Branch { .. }),
"the Branch node is never an executable member",
);
}
let absent = PickedRoute::new();
let absent_order = lower_picked_route(&g, &[post], &absent);
assert_eq!(
absent_order, arm0_order,
"a branch the picker did not resolve defaults to arm 0",
);
let mut bad = PickedRoute::new();
bad.insert(branch, 99);
let bad_order = lower_picked_route(&g, &[post], &bad);
assert_eq!(
bad_order, arm0_order,
"an out-of-range arm index clamps to arm 0 rather than panicking",
);
}
#[test]
fn collapse_variant_branch_bakes_the_winner_into_the_route() {
let (mut g, _diverge, arm0, arm1, reconverge, branch, post) = diamond_with_branch();
let before = lower_runs_arm0(&g, &[post]);
assert!(
before.contains(&arm0) && !before.contains(&arm1),
"before the bake the default route is arm 0"
);
g.collapse_variant_branch(branch, 1)
.expect("collapse to arm 1");
assert!(
g.node(reconverge).inputs.contains(&arm1),
"reconverge rewired to read the baked winner (arm 1)"
);
assert!(
!g.node(reconverge).inputs.contains(&arm0),
"reconverge no longer reads the pruned arm 0"
);
assert_eq!(
g.node(branch).inputs,
vec![arm1],
"the branch collapsed to the single winning arm"
);
let after = lower_runs_arm0(&g, &[post]);
assert!(
after.contains(&arm1),
"after the bake the default route follows the winner (arm 1); order={after:?}"
);
assert!(
!after.contains(&arm0),
"after the bake arm 0's interior is pruned; order={after:?}"
);
let (mut g2, _d2, a0_2, a1_2, recon2, branch2, post2) = diamond_with_branch();
let recon2_inputs_before = g2.node(recon2).inputs.clone();
g2.collapse_variant_branch(branch2, 0)
.expect("collapse to arm 0");
assert_eq!(
g2.node(recon2).inputs,
recon2_inputs_before,
"arm-0 collapse leaves the merge unchanged"
);
assert_eq!(
g2.node(branch2).inputs,
vec![a0_2],
"arm-0 collapse reduces the branch to [arm0] (variant arm pruned)"
);
let order2 = lower_runs_arm0(&g2, &[post2]);
assert!(
order2.contains(&a0_2) && !order2.contains(&a1_2),
"arm-0 collapse realizes arm 0 (byte-identical to today)"
);
let (mut g3, _d3, _a0_3, _a1_3, _r3, branch3, post3) = diamond_with_branch();
assert!(
g3.collapse_variant_branch(NodeId(9999), 0).is_err(),
"out-of-bounds branch node ⇒ Err, not panic"
);
assert!(
g3.collapse_variant_branch(branch3, 7).is_err(),
"out-of-range winner_arm ⇒ Err, not panic"
);
assert!(
g3.collapse_variant_branch(post3, 0).is_err(),
"non-Op::Branch node ⇒ Err, not panic"
);
}
fn two_coupled_diamonds() -> (Graph, NodeId, NodeId, NodeId, NodeId, NodeId) {
let mut g = Graph::new();
let pre = f32_node(&mut g, Op::Const, vec![]);
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]);
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");
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]);
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]);
(g, branch1, branch2, post, a1_1, a1_2)
}
#[test]
fn streaming_lowering_equals_one_shot_on_every_pick() {
let (g, _diverge, _arm0, _arm1, _reconverge, branch, post) = diamond_with_branch();
for chosen in [0usize, 1, 99] {
let mut route = PickedRoute::new();
if chosen != 0 {
route.insert(branch, chosen);
}
let one_shot = lower_picked_route(&g, &[post], &route);
let streamed = lower_picked_route_streaming(&g, &[post], |b| {
assert_eq!(b, branch, "only the diamond's branch is resolved");
chosen
});
assert_eq!(
streamed, one_shot,
"streamed order must equal one-shot for branch pick {chosen}",
);
}
let (g2, branch1, branch2, post2, _a1_1, _a1_2) = two_coupled_diamonds();
for c1 in [0usize, 1] {
for c2 in [0usize, 1] {
let mut route = PickedRoute::new();
if c1 != 0 {
route.insert(branch1, c1);
}
if c2 != 0 {
route.insert(branch2, c2);
}
let one_shot = lower_picked_route(&g2, &[post2], &route);
let mut seen: Vec<NodeId> = Vec::new();
let streamed = lower_picked_route_streaming(&g2, &[post2], |b| {
assert!(
!seen.contains(&b),
"each branch is resolved at most once; b={b:?} seen={seen:?}",
);
seen.push(b);
if b == branch1 {
c1
} else if b == branch2 {
c2
} else {
panic!("unexpected branch {b:?}")
}
});
assert_eq!(
seen,
vec![branch1, branch2],
"coupled branches resolve upstream-first at the frontier",
);
assert_eq!(
streamed, one_shot,
"streamed == one-shot for picks ({c1},{c2})",
);
}
}
}
#[test]
fn residency_seam_starts_new_run() {
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::Silu, vec![b]);
let d = f32_node(&mut g, Op::Tanh, vec![c]);
g.set_target_backend(a, BackendId::Cpu);
g.set_target_backend(b, BackendId::Cpu);
g.set_target_backend(c, BackendId::Cuda);
g.set_target_backend(d, BackendId::Cuda);
let runs = extract_runs(&g, d);
assert_eq!(
runs.len(),
2,
"a residency change mid-chain cuts the chain into two runs; got {runs:?}",
);
assert_eq!(runs[0].members, vec![a, b]);
assert_eq!(runs[0].device, Some(BackendId::Cpu));
assert_eq!(runs[1].members, vec![c, d]);
assert_eq!(runs[1].device, Some(BackendId::Cuda));
}
#[test]
fn fan_in_starts_new_run() {
let mut g = Graph::new();
let a = f32_node(&mut g, Op::Const, vec![]);
let a1 = f32_node(&mut g, Op::Relu, vec![a]);
let b = f32_node(&mut g, Op::Const, vec![]);
let b1 = f32_node(&mut g, Op::Silu, vec![b]);
let sum = f32_node(&mut g, Op::Add, vec![a1, b1]);
let runs = extract_runs(&g, sum);
let sum_run = runs
.iter()
.find(|r| r.members.contains(&sum))
.expect("sum is in some run");
assert_eq!(
sum_run.entry, sum,
"a fan-in node starts a fresh run (it is that run's entry)",
);
assert!(
!sum_run.members.contains(&a1) && !sum_run.members.contains(&b1),
"the fan-in run does not absorb either feeder; got {sum_run:?}",
);
}
#[test]
fn fewness_gate_passes_on_sparse_branching() {
let mut g = Graph::new();
let mut prev = f32_node(&mut g, Op::Const, vec![]);
let mut body: Vec<NodeId> = vec![prev];
for _ in 0..40 {
prev = f32_node(&mut g, Op::Relu, vec![prev]);
body.push(prev);
}
let diverge = prev;
let arm0 = f32_node(&mut g, Op::Silu, vec![diverge]);
let arm1 = f32_node(&mut g, Op::GeluTanh, vec![diverge]);
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);
b.finalize_branches(&mut g, reconverge)
.expect("valid branch")
.expect("2 arms survive");
let root = f32_node(&mut g, Op::Tanh, vec![reconverge]);
let density = branch_density(&g, root);
assert!(
density < FEWNESS_THRESHOLD,
"sparse layer-boundary branching must pass the fewness gate; density={density}",
);
assert!(passes_fewness_gate(&g, root));
}
#[test]
fn fewness_gate_fails_on_per_op_branching() {
let mut g = Graph::new();
let mut prev = f32_node(&mut g, Op::Const, vec![]);
for _ in 0..12 {
let diverge = prev;
let arm0 = f32_node(&mut g, Op::Silu, vec![diverge]);
let arm1 = f32_node(&mut g, Op::GeluTanh, vec![diverge]);
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);
b.finalize_branches(&mut g, reconverge)
.expect("valid branch")
.expect("2 arms survive");
prev = reconverge;
}
let root = prev;
let density = branch_density(&g, root);
assert!(
density >= FEWNESS_THRESHOLD,
"per-op branching must FAIL the fewness gate; density={density}",
);
assert!(!passes_fewness_gate(&g, root));
}
fn two_device_reconverge(
cuda_len: usize,
vk_len: usize,
) -> (Graph, NodeId, Vec<NodeId>, Vec<NodeId>, NodeId) {
let mut g = Graph::new();
let cc = f32_node(&mut g, Op::Const, vec![]);
let vc = f32_node(&mut g, Op::Const, vec![]);
let mut cuda_nodes = Vec::new();
let mut cprev = cc;
for _ in 0..cuda_len {
let n = f32_node(&mut g, Op::Relu, vec![cprev]);
g.set_target_backend(n, BackendId::Cuda);
cuda_nodes.push(n);
cprev = n;
}
let mut vk_nodes = Vec::new();
let mut vprev = vc;
for _ in 0..vk_len {
let n = f32_node(&mut g, Op::Silu, vec![vprev]);
g.set_target_backend(n, BackendId::Vulkan);
vk_nodes.push(n);
vprev = n;
}
let copy = g.push(Node {
op: Op::Copy {
target: fuel_ir::DeviceLocation::Cuda { gpu_id: 0 },
},
inputs: vec![vprev],
shape: Shape::from_dims(&[2]),
dtype: DType::F32,
});
g.set_target_backend(copy, BackendId::Cuda);
let out = g.push(Node {
op: Op::Add,
inputs: vec![cprev, copy],
shape: Shape::from_dims(&[2]),
dtype: DType::F32,
});
g.set_target_backend(out, BackendId::Cuda);
(g, out, cuda_nodes, vk_nodes, copy)
}
fn assert_valid_topo_run_order(graph: &Graph, runs: &[Run], order: &[usize]) {
let mut run_of: HashMap<NodeId, usize> = HashMap::new();
for (ri, r) in runs.iter().enumerate() {
for &m in &r.members {
run_of.insert(m, ri);
}
}
let mut emitted: HashSet<usize> = HashSet::new();
for &ri in order {
for &m in &runs[ri].members {
for &inp in &graph.node(m).inputs {
if let Some(&pi) = run_of.get(&inp)
&& pi != ri
{
assert!(
emitted.contains(&pi),
"run {ri} emitted before its producer run {pi}",
);
}
}
}
emitted.insert(ri);
}
assert_eq!(order.len(), runs.len(), "every run emitted exactly once");
let distinct: HashSet<usize> = order.iter().copied().collect();
assert_eq!(distinct.len(), runs.len(), "the order is a permutation");
}
#[test]
fn device_alternating_reorder_interleaves_two_devices() {
let (g, out, cuda_nodes, vk_nodes, copy) = two_device_reconverge(4, 4);
let runs = extract_runs(&g, out);
let perm = device_alternating_order(&g, &runs);
assert_valid_topo_run_order(&g, &runs, &perm);
let flat: Vec<NodeId> = perm
.iter()
.flat_map(|&ri| runs[ri].members.clone())
.collect();
let pos = |n: NodeId| flat.iter().position(|&x| x == n).expect("node present");
let cuda_root = *cuda_nodes.last().unwrap();
let vk_root = *vk_nodes.last().unwrap();
assert!(
pos(cuda_root) < pos(copy),
"the CUDA chunk must be enqueued before the cross-device copy; \
flat={flat:?}",
);
assert!(
pos(vk_root) < pos(copy),
"the Vulkan chunk must be recorded before the cross-device copy; \
flat={flat:?}",
);
let dev_seq: Vec<Option<BackendId>> = perm.iter().map(|&ri| runs[ri].device).collect();
let switches = dev_seq.windows(2).filter(|w| w[0] != w[1]).count();
assert!(
switches >= 1,
"a two-device run list must contain at least one device switch; \
dev_seq={dev_seq:?}",
);
}
#[test]
fn device_alternating_reorder_single_device_is_identity() {
let mut g = Graph::new();
let mut prev = f32_node(&mut g, Op::Const, vec![]);
for _ in 0..6 {
prev = f32_node(&mut g, Op::Relu, vec![prev]);
g.set_target_backend(prev, BackendId::Cuda);
}
let runs = extract_runs(&g, prev);
let perm = device_alternating_order(&g, &runs);
assert_eq!(
perm,
(0..runs.len()).collect::<Vec<_>>(),
"a single-device run list reorders to the identity permutation",
);
let mut g2 = Graph::new();
let a = f32_node(&mut g2, Op::Const, vec![]);
let b = f32_node(&mut g2, Op::Relu, vec![a]);
let c = f32_node(&mut g2, Op::Silu, vec![b]);
let runs2 = extract_runs(&g2, c);
assert_eq!(
device_alternating_order(&g2, &runs2),
(0..runs2.len()).collect::<Vec<_>>(),
"an all-default (device=None) run list is identity",
);
}
#[test]
fn device_alternating_reorder_respects_dependencies() {
let mut g = Graph::new();
let a = f32_node(&mut g, Op::Const, vec![]);
g.set_target_backend(a, BackendId::Cpu);
let b = f32_node(&mut g, Op::Relu, vec![a]);
g.set_target_backend(b, BackendId::Cuda);
let c = f32_node(&mut g, Op::Silu, vec![b]);
g.set_target_backend(c, BackendId::Vulkan);
let d = f32_node(&mut g, Op::Tanh, vec![c]);
g.set_target_backend(d, BackendId::Cuda);
let runs = extract_runs(&g, d);
let perm = device_alternating_order(&g, &runs);
assert_valid_topo_run_order(&g, &runs, &perm);
let flat: Vec<NodeId> = perm
.iter()
.flat_map(|&ri| runs[ri].members.clone())
.collect();
assert_eq!(
flat,
vec![a, b, c, d],
"a strict cross-device chain keeps dependency order despite \
device alternation",
);
}
#[test]
fn lower_picked_route_multidevice_is_a_valid_reordering() {
let (g, out, _cuda, _vk, _copy) = two_device_reconverge(4, 4);
let runs = extract_runs(&g, out);
let reference: Vec<NodeId> = runs.iter().flat_map(|r| lower_run(r).to_vec()).collect();
let lowered = lower_picked_route(&g, &[out], &PickedRoute::new());
let mut a = reference.clone();
let mut b = lowered.clone();
a.sort();
b.sort();
assert_eq!(a, b, "the reorder is a permutation of the same node set");
let lowered_pos: HashMap<NodeId, usize> =
lowered.iter().enumerate().map(|(i, &n)| (n, i)).collect();
for (i, &n) in lowered.iter().enumerate() {
for &inp in &g.node(n).inputs {
if let Some(&j) = lowered_pos.get(&inp) {
assert!(j < i, "input {inp:?} of {n:?} must precede it");
}
}
}
assert_ne!(
lowered, reference,
"on this two-device graph the reorder must change the order \
(else auto-overlap did nothing)",
);
}
#[test]
fn streaming_multidevice_equals_one_shot_reordered() {
let (g, out, _cuda, _vk, _copy) = two_device_reconverge(4, 4);
let one_shot = lower_picked_route(&g, &[out], &PickedRoute::new());
let streamed = lower_picked_route_streaming(&g, &[out], |_b| 0);
assert_eq!(
streamed, one_shot,
"streamed multi-device order equals the one-shot reordered order",
);
}
#[test]
fn device_alternating_reorder_emits_heavier_chunk_first() {
let (g, out, cuda_nodes, vk_nodes, copy) = two_device_reconverge(8, 2);
let runs = extract_runs(&g, out);
let perm = device_alternating_order(&g, &runs);
assert_valid_topo_run_order(&g, &runs, &perm);
let flat: Vec<NodeId> = perm
.iter()
.flat_map(|&ri| runs[ri].members.clone())
.collect();
let pos = |n: NodeId| flat.iter().position(|&x| x == n).expect("present");
let cuda_root = *cuda_nodes.last().unwrap();
let vk_root = *vk_nodes.last().unwrap();
assert!(
pos(cuda_root) < pos(vk_root),
"the heavier CUDA chunk must be emitted before the lighter Vulkan \
chunk (the auto-submit device runs first); flat={flat:?}",
);
assert!(pos(cuda_root) < pos(copy), "heavy chunk before the drain");
assert!(pos(vk_root) < pos(copy), "light chunk before the drain");
let (g2, out2, cuda2, vk2, copy2) = two_device_reconverge(2, 8);
let runs2 = extract_runs(&g2, out2);
let perm2 = device_alternating_order(&g2, &runs2);
assert_valid_topo_run_order(&g2, &runs2, &perm2);
let flat2: Vec<NodeId> = perm2
.iter()
.flat_map(|&ri| runs2[ri].members.clone())
.collect();
let pos2 = |n: NodeId| flat2.iter().position(|&x| x == n).expect("present");
assert!(
pos2(*vk2.last().unwrap()) < pos2(*cuda2.last().unwrap()),
"when Vulkan is the heavier chunk it is emitted first; flat2={flat2:?}",
);
assert!(pos2(*vk2.last().unwrap()) < pos2(copy2));
assert!(pos2(*cuda2.last().unwrap()) < pos2(copy2));
}
}