onnx_runtime_memory/
validate.rs

1//! Post-hoc validation of an [`ActivationPlan`] against the graph's liveness.
2//!
3//! Recomputes liveness independently of the plan's slot assignments, then
4//! enforces the correctness invariants of buffer sharing. Tests call this on
5//! every produced plan, and it must *catch* a deliberately corrupted plan.
6
7use std::collections::HashMap;
8
9use onnx_runtime_ir::{Graph, ValueId};
10
11use crate::error::ValidateError;
12use crate::liveness::compute_liveness;
13use crate::options::PlanOptions;
14use crate::oracle::static_size_oracle;
15use crate::plan::{ActivationPlan, SlotId};
16use crate::view_map::ViewMap;
17
18/// Validate a plan against the graph, using `size_oracle` to check slot
19/// capacities. Enforces:
20///
21/// 1. **No overlap sharing** — two owners with overlapping live intervals never
22///    share a slot (the core safety invariant).
23/// 2. **Coverage & capacity** — every owner has a slot that exists and is large
24///    enough; a graph output's slot is never reused after its def (implied by
25///    (1), since an output's interval overlaps everything after its def).
26/// 3. **View folding** — every zero-copy view's last use falls within its
27///    source owner's live interval (the source outlives the alias).
28pub fn validate<F>(
29    plan: &ActivationPlan,
30    graph: &Graph,
31    view_map: &ViewMap,
32    size_oracle: F,
33    options: &PlanOptions,
34) -> Result<(), ValidateError>
35where
36    F: Fn(ValueId) -> Option<usize>,
37{
38    let live = compute_liveness(graph, view_map, options).map_err(|_| ValidateError::Cycle)?;
39
40    let slot_cap: HashMap<SlotId, usize> = plan
41        .slots
42        .iter()
43        .map(|s| (s.id, s.capacity_bytes))
44        .collect();
45
46    // (1) + (2): coverage, capacity, and overlap-free slot sharing.
47    let mut by_slot: HashMap<SlotId, Vec<ValueId>> = HashMap::new();
48    for &owner in live.intervals.keys() {
49        let Some(&sid) = plan.assignments.get(&owner) else {
50            return Err(ValidateError::MissingAssignment { value: owner });
51        };
52        let Some(&cap) = slot_cap.get(&sid) else {
53            return Err(ValidateError::UnknownSlot {
54                value: owner,
55                slot: sid,
56            });
57        };
58        if let Some(need) = size_oracle(owner)
59            && need > cap
60        {
61            return Err(ValidateError::UndersizedSlot {
62                value: owner,
63                slot: sid,
64                needed: need,
65                capacity: cap,
66            });
67        }
68        by_slot.entry(sid).or_default().push(owner);
69    }
70
71    for (&sid, members) in &by_slot {
72        for i in 0..members.len() {
73            for j in (i + 1)..members.len() {
74                let a = members[i];
75                let b = members[j];
76                if live.intervals[&a].overlaps(&live.intervals[&b]) {
77                    return Err(ValidateError::SlotConflict { a, b, slot: sid });
78                }
79            }
80        }
81    }
82
83    // (3): every view's last use is covered by its source owner's interval.
84    for vid in graph.values.keys() {
85        if !view_map.is_view(vid) {
86            continue;
87        }
88        let root = view_map.root(vid);
89        let Some(root_interval) = live.intervals.get(&root) else {
90            continue; // source not part of the arena (e.g. excluded input)
91        };
92        let view_use = graph
93            .consumers(vid)
94            .into_iter()
95            .filter_map(|consumer| live.order_index.get(&consumer).copied())
96            .max();
97        let view_use = match view_use {
98            Some(u) => u.max(view_output_end(graph, &live, vid)),
99            None => view_output_end(graph, &live, vid),
100        };
101        if view_use > root_interval.use_end {
102            return Err(ValidateError::ViewOutlivesSource {
103                view: vid,
104                source_owner: root,
105                view_use,
106                source_end: root_interval.use_end,
107            });
108        }
109    }
110
111    Ok(())
112}
113
114/// If `vid` is a graph output, its liveness extends to the last node.
115fn view_output_end(graph: &Graph, live: &crate::liveness::Liveness, vid: ValueId) -> usize {
116    if graph.outputs.contains(&vid) {
117        live.last_index
118    } else {
119        0
120    }
121}
122
123/// Convenience: validate a plan built from fully-static shapes.
124pub fn validate_static(
125    plan: &ActivationPlan,
126    graph: &Graph,
127    view_map: &ViewMap,
128    options: &PlanOptions,
129) -> Result<(), ValidateError> {
130    let oracle = static_size_oracle(graph);
131    validate(plan, graph, view_map, oracle, options)
132}
onnx_runtime_memory/validate.rs

onnx_runtime_memory/
validate.rs
