use crate::plan::backend_for_device;
use fuel_graph::{Graph, NodeId, SharedGraph, opt, topo_order_multi};
use fuel_ir::{DeviceLocation, Error, Result};
#[cfg(test)]
use fuel_graph::Op;
use std::collections::HashMap;
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct LiveRange {
pub node: NodeId,
pub bytes: usize,
pub first_use: usize,
pub last_use: usize,
pub inactive_gap: usize,
}
#[derive(Debug, Clone)]
pub struct ResidencyReport {
pub total_bytes: usize,
pub peak_bytes: usize,
pub peak_op_index: usize,
pub eviction_candidates: Vec<LiveRange>,
}
impl ResidencyReport {
pub fn fits_in(&self, budget: usize) -> bool {
self.peak_bytes <= budget
}
pub fn overage(&self, budget: usize) -> usize {
self.peak_bytes.saturating_sub(budget)
}
}
pub struct ResidencyPlanner;
impl ResidencyPlanner {
pub fn analyze(g: &Graph, roots: &[NodeId]) -> ResidencyReport {
let order = topo_order_multi(g, roots);
let n = order.len();
let mut byte_of: HashMap<NodeId, usize> = HashMap::with_capacity(n);
let mut first_use: HashMap<NodeId, usize> = HashMap::with_capacity(n);
let mut last_use: HashMap<NodeId, usize> = HashMap::with_capacity(n);
let mut reads_at: Vec<Vec<NodeId>> = vec![Vec::new(); n];
for (op_idx, &nid) in order.iter().enumerate() {
let node = g.node(nid);
let elems = node.shape.elem_count();
let bytes = elems * node.dtype.size_in_bytes();
byte_of.insert(nid, bytes);
first_use.entry(nid).or_insert(op_idx);
for &input in &node.inputs {
last_use.insert(input, op_idx);
reads_at[op_idx].push(input);
}
}
for &r in roots {
if let Some(pos) = order.iter().position(|&x| x == r) {
last_use.entry(r).or_insert(pos);
}
}
let mut live_bytes: usize = 0;
let mut peak_bytes: usize = 0;
let mut peak_op_index: usize = 0;
let mut total_bytes: usize = 0;
for (op_idx, &nid) in order.iter().enumerate() {
let bytes = *byte_of.get(&nid).unwrap_or(&0);
live_bytes += bytes;
total_bytes += bytes;
if live_bytes > peak_bytes {
peak_bytes = live_bytes;
peak_op_index = op_idx;
}
for &input in &reads_at[op_idx] {
if last_use.get(&input) == Some(&op_idx) {
let b = *byte_of.get(&input).unwrap_or(&0);
live_bytes = live_bytes.saturating_sub(b);
}
}
}
let mut gap_of: HashMap<NodeId, usize> = HashMap::with_capacity(n);
for (nid, &first) in &first_use {
let reads: Vec<usize> = reads_at
.iter()
.enumerate()
.filter_map(|(i, rs)| if rs.contains(nid) { Some(i) } else { None })
.collect();
let last = last_use.get(nid).copied().unwrap_or(first);
let mut points = vec![first];
points.extend(reads.iter().copied());
points.push(last);
points.sort();
let max_gap = points.windows(2).map(|w| w[1] - w[0]).max().unwrap_or(0);
gap_of.insert(*nid, max_gap);
}
let mut cands: Vec<LiveRange> = order
.iter()
.map(|&nid| LiveRange {
node: nid,
bytes: *byte_of.get(&nid).unwrap_or(&0),
first_use: *first_use.get(&nid).unwrap_or(&0),
last_use: *last_use.get(&nid).unwrap_or(&0),
inactive_gap: *gap_of.get(&nid).unwrap_or(&0),
})
.collect();
cands.sort_by(|a, b| {
let a_score = (a.bytes as u128) * (a.inactive_gap as u128);
let b_score = (b.bytes as u128) * (b.inactive_gap as u128);
b_score.cmp(&a_score).then(b.bytes.cmp(&a.bytes))
});
cands.retain(|c| c.bytes > 0 && c.inactive_gap > 0);
cands.truncate(32);
ResidencyReport {
total_bytes,
peak_bytes,
peak_op_index,
eviction_candidates: cands,
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct EvictReload {
pub candidate: NodeId,
pub move_node: NodeId,
pub reload: NodeId,
pub src_device: DeviceLocation,
}
pub fn insert_residency_evictions(
graph: &SharedGraph,
roots: &[NodeId],
budget_bytes: usize,
max_rounds: usize,
placement_of: impl Fn(NodeId) -> Option<DeviceLocation>,
) -> Result<Vec<EvictReload>> {
let mut evicted: std::collections::HashSet<NodeId> = std::collections::HashSet::new();
let mut chains: Vec<EvictReload> = Vec::new();
for _round in 0..max_rounds {
let report = {
let g = read_graph(graph)?;
ResidencyPlanner::analyze(&g, roots)
};
if report.fits_in(budget_bytes) {
return Ok(chains);
}
let mut progressed = false;
for cand in &report.eviction_candidates {
if evicted.contains(&cand.node) {
continue;
}
let post_gap_consumers = {
let g = read_graph(graph)?;
match find_gap_consumers(&g, cand.node) {
Some((_pre_gap, post)) if !post.is_empty() => post,
_ => continue,
}
};
let src_device = placement_of(cand.node).unwrap_or(DeviceLocation::Cpu);
let (move_id, reload_id) =
opt::insert_evict_reload(graph, cand.node, src_device, &post_gap_consumers);
{
let mut g = write_graph(graph)?;
g.set_target_backend(move_id, backend_for_device(src_device));
g.set_target_backend(reload_id, backend_for_device(DeviceLocation::Cpu));
g.set_placement(move_id, DeviceLocation::Cpu);
g.set_placement(reload_id, src_device);
}
evicted.insert(cand.node);
chains.push(EvictReload {
candidate: cand.node,
move_node: move_id,
reload: reload_id,
src_device,
});
progressed = true;
break;
}
if !progressed {
return Ok(chains);
}
}
Ok(chains)
}
fn find_gap_consumers(g: &Graph, candidate: NodeId) -> Option<(NodeId, Vec<NodeId>)> {
let all: Vec<NodeId> = (0..g.len()).map(NodeId).collect();
let order = topo_order_multi(g, &all);
let mut positions: Vec<(usize, NodeId)> = Vec::new();
for (i, &nid) in order.iter().enumerate() {
if g.node(nid).inputs.contains(&candidate) {
positions.push((i, nid));
}
}
if positions.len() < 2 {
return None;
}
let mut best_gap = 0usize;
let mut gap_idx = 0usize;
for i in 0..positions.len() - 1 {
let gap = positions[i + 1].0 - positions[i].0;
if gap > best_gap {
best_gap = gap;
gap_idx = i;
}
}
if best_gap < 2 {
return None;
}
let last_before = positions[gap_idx].1;
let after_gap: Vec<NodeId> = positions[gap_idx + 1..].iter().map(|(_, n)| *n).collect();
Some((last_before, after_gap))
}
fn read_graph(graph: &SharedGraph) -> Result<std::sync::RwLockReadGuard<'_, Graph>> {
graph
.read()
.map_err(|_| Error::Msg("residency: graph RwLock poisoned".into()).bt())
}
fn write_graph(graph: &SharedGraph) -> Result<std::sync::RwLockWriteGuard<'_, Graph>> {
graph
.write()
.map_err(|_| Error::Msg("residency: graph RwLock poisoned".into()).bt())
}
#[cfg(test)]
mod tests {
use super::*;
use crate::pipelined::{PipelinedExecutor, StorageCache};
use fuel_graph::Node;
use fuel_ir::{DType, Shape, probe::BackendId};
use std::sync::{Arc, RwLock};
fn shared() -> SharedGraph {
Arc::new(RwLock::new(Graph::new()))
}
fn push(g: &mut Graph, op: Op, inputs: Vec<NodeId>, dims: &[usize]) -> NodeId {
g.push(Node {
op,
inputs,
shape: Shape::from_dims(dims),
dtype: DType::F32,
})
}
fn count_op(graph: &SharedGraph, pred: impl Fn(&Op) -> bool) -> usize {
let g = graph.read().unwrap();
(0..g.len())
.filter(|i| pred(&g.node(NodeId(*i)).op))
.count()
}
#[test]
fn chain_graph_peak_is_two_live_tensors() {
let graph = shared();
let c = {
let mut g = graph.write().unwrap();
let a = push(&mut g, Op::Const, vec![], &[1024]);
let b = push(&mut g, Op::Relu, vec![a], &[1024]);
push(&mut g, Op::Neg, vec![b], &[1024])
};
let g = graph.read().unwrap();
let report = ResidencyPlanner::analyze(&g, &[c]);
assert_eq!(report.total_bytes, 4096 * 3);
assert_eq!(report.peak_bytes, 4096 * 2);
}
#[test]
fn shared_const_stays_live_across_multiple_uses() {
let graph = shared();
let (a, sum) = {
let mut g = graph.write().unwrap();
let a = push(&mut g, Op::Const, vec![], &[256]);
let b = push(&mut g, Op::Const, vec![], &[256]);
let ab = push(&mut g, Op::Add, vec![a, b], &[256]);
let na = push(&mut g, Op::Neg, vec![a], &[256]);
let sum = push(&mut g, Op::Add, vec![ab, na], &[256]);
(a, sum)
};
let g = graph.read().unwrap();
let report = ResidencyPlanner::analyze(&g, &[sum]);
let a_range = report.eviction_candidates.iter().find(|c| c.node == a);
assert!(
a_range.is_some() || report.peak_bytes > 0,
"expected `a` tracked with a gap OR at least some peak bytes",
);
}
#[test]
fn report_fits_in_and_overage() {
let graph = shared();
let b = {
let mut g = graph.write().unwrap();
let a = push(&mut g, Op::Const, vec![], &[256]);
push(&mut g, Op::Relu, vec![a], &[256])
};
let g = graph.read().unwrap();
let report = ResidencyPlanner::analyze(&g, &[b]);
assert!(report.fits_in(10_000));
assert!(!report.fits_in(1));
assert_eq!(report.overage(10_000), 0);
assert!(report.overage(1) > 0);
}
fn over_budget_graph() -> (SharedGraph, NodeId, NodeId) {
let graph = shared();
let (a, out) = {
let mut g = graph.write().unwrap();
let a = push(&mut g, Op::Const, vec![], &[4]);
let b = push(&mut g, Op::Relu, vec![a], &[4]);
let pad = push(&mut g, Op::Neg, vec![b], &[4]);
let pad2 = push(&mut g, Op::Neg, vec![pad], &[4]);
let c = push(&mut g, Op::Mul, vec![pad2, a], &[4]);
let out = push(&mut g, Op::Add, vec![b, c], &[4]);
for n in [b, pad, pad2, c, out] {
g.set_target_backend(n, BackendId::Cpu);
}
(a, out)
};
(graph, a, out)
}
#[test]
fn eviction_is_noop_under_budget() {
let (graph, _a, out) = over_budget_graph();
let moves_before = count_op(&graph, |op| matches!(op, Op::Move { .. }));
let copies_before = count_op(&graph, |op| matches!(op, Op::Copy { .. }));
let chains = insert_residency_evictions(&graph, &[out], 1_000_000_000, 16, |_| None)
.expect("eviction pass");
assert!(
chains.is_empty(),
"under-budget graph should emit no chains"
);
assert_eq!(
count_op(&graph, |op| matches!(op, Op::Move { .. })),
moves_before
);
assert_eq!(
count_op(&graph, |op| matches!(op, Op::Copy { .. })),
copies_before
);
}
#[test]
fn eviction_emits_move_reload_chain_when_over_budget() {
let (graph, a, out) = over_budget_graph();
let chains =
insert_residency_evictions(&graph, &[out], 1, 16, |_| Some(DeviceLocation::Cpu))
.expect("eviction pass");
assert!(
!chains.is_empty(),
"over-budget graph should emit at least one chain"
);
assert!(
count_op(&graph, |op| matches!(op, Op::Move { .. })) >= 1,
"over-budget graph should emit at least one Op::Move",
);
assert!(
count_op(&graph, |op| matches!(op, Op::Copy { .. })) >= 1,
"each chain carries a reload Op::Copy",
);
let chain = chains
.iter()
.find(|c| c.candidate == a)
.unwrap_or(&chains[0]);
let g = graph.read().unwrap();
assert!(matches!(
g.node(chain.move_node).op,
Op::Move {
target: DeviceLocation::Cpu
},
));
assert_eq!(g.node(chain.reload).inputs, vec![chain.move_node]);
assert_eq!(g.target_backend(chain.move_node), Some(BackendId::Cpu));
assert_eq!(g.target_backend(chain.reload), Some(BackendId::Cpu));
}
#[test]
fn eviction_preserves_correctness_end_to_end_on_pipelined() {
let realize = |evict: bool| -> Vec<f32> {
let (graph, a, out) = over_budget_graph();
if evict {
let chains = insert_residency_evictions(&graph, &[out], 1, 16, |_| {
Some(DeviceLocation::Cpu)
})
.expect("eviction pass");
assert!(!chains.is_empty(), "budget=1 must trigger eviction");
}
let mut inputs = StorageCache::new();
inputs.insert(
a,
Arc::new(RwLock::new(fuel_memory::from_slice_cpu(&[
-1.0_f32, 2.0, -3.0, 4.0,
]))),
);
let (storage, _layout) =
PipelinedExecutor::realize(graph, out, inputs).expect("realize");
let guard = storage.read().unwrap();
let fuel_memory::BackendStorage::Cpu(c) = &guard.inner else {
panic!("expected CPU output");
};
c.as_slice::<f32>().unwrap().to_vec()
};
let baseline = realize(false);
let evicted = realize(true);
assert_eq!(
baseline, evicted,
"eviction must preserve the output bit-exactly",
);
assert_eq!(baseline, vec![0.0_f32, 6.0, 0.0, 20.0]);
}
#[test]
fn eviction_terminates_when_budget_unreachable() {
let (graph, _a, out) = over_budget_graph();
let chains =
insert_residency_evictions(&graph, &[out], 0, 16, |_| None).expect("eviction pass");
let len_after = graph.read().unwrap().len();
let chains2 = insert_residency_evictions(&graph, &[out], 0, 16, |_| None)
.expect("second eviction pass");
let len_final = graph.read().unwrap().len();
assert!(chains.len() <= 16);
assert!(chains2.len() <= 16);
assert!(len_final >= len_after);
assert!(len_final - len_after <= 2 * chains2.len());
}
}