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> =
41        plan.slots.iter().map(|s| (s.id, s.capacity_bytes)).collect();
42
43    // (1) + (2): coverage, capacity, and overlap-free slot sharing.
44    let mut by_slot: HashMap<SlotId, Vec<ValueId>> = HashMap::new();
45    for &owner in live.intervals.keys() {
46        let Some(&sid) = plan.assignments.get(&owner) else {
47            return Err(ValidateError::MissingAssignment { value: owner });
48        };
49        let Some(&cap) = slot_cap.get(&sid) else {
50            return Err(ValidateError::UnknownSlot {
51                value: owner,
52                slot: sid,
53            });
54        };
55        if let Some(need) = size_oracle(owner)
56            && need > cap
57        {
58            return Err(ValidateError::UndersizedSlot {
59                value: owner,
60                slot: sid,
61                needed: need,
62                capacity: cap,
63            });
64        }
65        by_slot.entry(sid).or_default().push(owner);
66    }
67
68    for (&sid, members) in &by_slot {
69        for i in 0..members.len() {
70            for j in (i + 1)..members.len() {
71                let a = members[i];
72                let b = members[j];
73                if live.intervals[&a].overlaps(&live.intervals[&b]) {
74                    return Err(ValidateError::SlotConflict { a, b, slot: sid });
75                }
76            }
77        }
78    }
79
80    // (3): every view's last use is covered by its source owner's interval.
81    for vid in graph.values.keys() {
82        if !view_map.is_view(vid) {
83            continue;
84        }
85        let root = view_map.root(vid);
86        let Some(root_interval) = live.intervals.get(&root) else {
87            continue; // source not part of the arena (e.g. excluded input)
88        };
89        let view_use = graph
90            .value(vid)
91            .consumers
92            .iter()
93            .filter_map(|c| live.order_index.get(c).copied())
94            .max();
95        let view_use = match view_use {
96            Some(u) => u.max(view_output_end(graph, &live, vid)),
97            None => view_output_end(graph, &live, vid),
98        };
99        if view_use > root_interval.use_end {
100            return Err(ValidateError::ViewOutlivesSource {
101                view: vid,
102                source_owner: root,
103                view_use,
104                source_end: root_interval.use_end,
105            });
106        }
107    }
108
109    Ok(())
110}
111
112/// If `vid` is a graph output, its liveness extends to the last node.
113fn view_output_end(graph: &Graph, live: &crate::liveness::Liveness, vid: ValueId) -> usize {
114    if graph.outputs.contains(&vid) {
115        live.last_index
116    } else {
117        0
118    }
119}
120
121/// Convenience: validate a plan built from fully-static shapes.
122pub fn validate_static(
123    plan: &ActivationPlan,
124    graph: &Graph,
125    view_map: &ViewMap,
126    options: &PlanOptions,
127) -> Result<(), ValidateError> {
128    let oracle = static_size_oracle(graph);
129    validate(plan, graph, view_map, oracle, options)
130}
onnx_runtime_memory/validate.rs

onnx_runtime_memory/
validate.rs
