onnx_runtime_memory/
liveness.rs

1//! Liveness analysis: the live interval of every activation buffer in
2//! execution (topological) order.
3//!
4//! The unit of allocation is a *buffer owner* — a value that needs a real
5//! activation slot. Graph inputs (supplied by the caller) and initializers /
6//! streamed weights (caller-owned, memory-mapped) are **excluded** from the
7//! activation arena. Zero-copy views own no buffer either; their liveness is
8//! folded into the root owner they alias (see [`crate::ViewMap`]).
9
10use std::collections::HashMap;
11
12use onnx_runtime_ir::{Graph, NodeId, ValueId};
13
14use crate::error::PlanError;
15use crate::options::PlanOptions;
16use crate::view_map::ViewMap;
17
18/// The half-open-agnostic live interval of a buffer owner, in units of
19/// topological node index.
20///
21/// * `def` — the index of the node that produces the value (or `0` for an
22///   included graph input, which is live from the start of execution).
23/// * `use_end` — the index of the value's last use. Two owners whose intervals
24///   touch even at a single index (`a.def <= b.use_end && b.def <= a.use_end`)
25///   are considered overlapping and must not share a slot, because at that
26///   node the producer's output and the still-needed input coexist.
27#[derive(Clone, Copy, Debug, PartialEq, Eq)]
28pub struct Interval {
29    pub def: usize,
30    pub use_end: usize,
31}
32
33impl Interval {
34    /// Whether two intervals overlap (share at least one node index).
35    pub fn overlaps(&self, other: &Interval) -> bool {
36        self.def <= other.use_end && other.def <= self.use_end
37    }
38}
39
40/// Result of liveness analysis over a graph.
41#[derive(Clone, Debug)]
42pub struct Liveness {
43    /// Topological position of each live node.
44    pub order_index: HashMap<NodeId, usize>,
45    /// Node ids in topological order.
46    pub order: Vec<NodeId>,
47    /// Index of the last node (0 for an empty/single-node schedule).
48    pub last_index: usize,
49    /// Live interval of each buffer *owner* (never a view, input, or weight).
50    pub intervals: HashMap<ValueId, Interval>,
51    /// Owners in deterministic allocation order: ascending `def`, then id.
52    pub owners: Vec<ValueId>,
53}
54
55/// Whether `value` needs its own activation slot.
56fn is_buffer_owner(graph: &Graph, view_map: &ViewMap, options: &PlanOptions, value: ValueId) -> bool {
57    if graph.initializers.contains_key(&value) {
58        return false; // streamed / caller-owned weight
59    }
60    if view_map.is_view(value) {
61        return false; // zero-copy alias, no buffer of its own
62    }
63    let Some(val) = graph.try_value(value) else {
64        return false;
65    };
66    if val.producer.is_some() {
67        return true; // an activation produced by a node
68    }
69    // No producer: only a graph input can qualify, and only if opted in.
70    options.include_graph_inputs && graph.inputs.contains(&value)
71}
72
73/// Compute the live interval of every activation buffer owner.
74///
75/// Views are folded to their root owner: a view's consumers (and graph-output
76/// status) extend the root's `use_end`, guaranteeing the source outlives every
77/// alias. Returns [`PlanError::Cycle`] if the graph is not schedulable.
78pub fn compute_liveness(
79    graph: &Graph,
80    view_map: &ViewMap,
81    options: &PlanOptions,
82) -> Result<Liveness, PlanError> {
83    let order = graph.topological_order().map_err(|_| PlanError::Cycle)?;
84    let order_index: HashMap<NodeId, usize> =
85        order.iter().enumerate().map(|(i, &n)| (n, i)).collect();
86    let last_index = order.len().saturating_sub(1);
87
88    let outputs: std::collections::HashSet<ValueId> = graph.outputs.iter().copied().collect();
89
90    // 1. Seed intervals for every buffer owner at its definition point.
91    let mut intervals: HashMap<ValueId, Interval> = HashMap::new();
92    for vid in graph.values.keys() {
93        if !is_buffer_owner(graph, view_map, options, vid) {
94            continue;
95        }
96        let def = graph
97            .value(vid)
98            .producer
99            .and_then(|p| order_index.get(&p).copied())
100            .unwrap_or(0); // included graph input: live from the start
101        intervals.insert(vid, Interval { def, use_end: def });
102    }
103
104    // 2. Extend each root owner's interval by every activation value (the owner
105    //    itself or any view folding to it): its consumers and output status.
106    for vid in graph.values.keys() {
107        // Skip values that are neither owners nor views (e.g. excluded graph
108        // inputs and initializers); their root is not in `intervals`.
109        let root = view_map.root(vid);
110        let Some(interval) = intervals.get_mut(&root) else {
111            continue;
112        };
113        for &consumer in &graph.value(vid).consumers {
114            if let Some(&idx) = order_index.get(&consumer) {
115                interval.use_end = interval.use_end.max(idx);
116            }
117        }
118        // A graph output (or a view that is a graph output) pins its root to the
119        // end of execution: it must never be overwritten.
120        if outputs.contains(&vid) {
121            interval.use_end = interval.use_end.max(last_index);
122        }
123    }
124
125    // 3. Deterministic owner ordering: ascending def, then value id.
126    let mut owners: Vec<ValueId> = intervals.keys().copied().collect();
127    owners.sort_by_key(|v| (intervals[v].def, v.0));
128
129    Ok(Liveness {
130        order_index,
131        order,
132        last_index,
133        intervals,
134        owners,
135    })
136}
onnx_runtime_memory/liveness.rs

onnx_runtime_memory/
liveness.rs
