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(
57    graph: &Graph,
58    view_map: &ViewMap,
59    options: &PlanOptions,
60    value: ValueId,
61) -> bool {
62    if graph.initializers.contains_key(&value) {
63        return false; // streamed / caller-owned weight
64    }
65    if view_map.is_view(value) {
66        return false; // zero-copy alias, no buffer of its own
67    }
68    let Some(val) = graph.try_value(value) else {
69        return false;
70    };
71    if val.producer.is_some() {
72        return true; // an activation produced by a node
73    }
74    // No producer: only a graph input can qualify, and only if opted in.
75    options.include_graph_inputs && graph.inputs.contains(&value)
76}
77
78/// Compute the live interval of every activation buffer owner.
79///
80/// Views are folded to their root owner: a view's consumers (and graph-output
81/// status) extend the root's `use_end`, guaranteeing the source outlives every
82/// alias. Returns [`PlanError::Cycle`] if the graph is not schedulable.
83pub fn compute_liveness(
84    graph: &Graph,
85    view_map: &ViewMap,
86    options: &PlanOptions,
87) -> Result<Liveness, PlanError> {
88    let order = graph.topological_order().map_err(|_| PlanError::Cycle)?;
89    let order_index: HashMap<NodeId, usize> =
90        order.iter().enumerate().map(|(i, &n)| (n, i)).collect();
91    let last_index = order.len().saturating_sub(1);
92
93    let outputs: std::collections::HashSet<ValueId> = graph.outputs.iter().copied().collect();
94
95    // 1. Seed intervals for every buffer owner at its definition point.
96    let mut intervals: HashMap<ValueId, Interval> = HashMap::new();
97    for vid in graph.values.keys() {
98        if !is_buffer_owner(graph, view_map, options, vid) {
99            continue;
100        }
101        let def = graph
102            .value(vid)
103            .producer
104            .and_then(|p| order_index.get(&p).copied())
105            .unwrap_or(0); // included graph input: live from the start
106        intervals.insert(vid, Interval { def, use_end: def });
107    }
108
109    // 2. Extend each root owner's interval by every activation value (the owner
110    //    itself or any view folding to it): its consumers and output status.
111    for vid in graph.values.keys() {
112        // Skip values that are neither owners nor views (e.g. excluded graph
113        // inputs and initializers); their root is not in `intervals`.
114        let root = view_map.root(vid);
115        let Some(interval) = intervals.get_mut(&root) else {
116            continue;
117        };
118        for consumer in graph.consumers(vid) {
119            if let Some(&idx) = order_index.get(&consumer) {
120                interval.use_end = interval.use_end.max(idx);
121            }
122        }
123        // A graph output (or a view that is a graph output) pins its root to the
124        // end of execution: it must never be overwritten.
125        if outputs.contains(&vid) {
126            interval.use_end = interval.use_end.max(last_index);
127        }
128    }
129
130    // 3. Deterministic owner ordering: ascending def, then value id.
131    let mut owners: Vec<ValueId> = intervals.keys().copied().collect();
132    owners.sort_by_key(|v| (intervals[v].def, v.0));
133
134    Ok(Liveness {
135        order_index,
136        order,
137        last_index,
138        intervals,
139        owners,
140    })
141}
onnx_runtime_memory/liveness.rs

onnx_runtime_memory/
liveness.rs
