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            .consumers(vid)
91            .into_iter()
92            .filter_map(|consumer| live.order_index.get(&consumer).copied())
93            .max();
94        let view_use = match view_use {
95            Some(u) => u.max(view_output_end(graph, &live, vid)),
96            None => view_output_end(graph, &live, vid),
97        };
98        if view_use > root_interval.use_end {
99            return Err(ValidateError::ViewOutlivesSource {
100                view: vid,
101                source_owner: root,
102                view_use,
103                source_end: root_interval.use_end,
104            });
105        }
106    }
107
108    Ok(())
109}
110
111/// If `vid` is a graph output, its liveness extends to the last node.
112fn view_output_end(graph: &Graph, live: &crate::liveness::Liveness, vid: ValueId) -> usize {
113    if graph.outputs.contains(&vid) {
114        live.last_index
115    } else {
116        0
117    }
118}
119
120/// Convenience: validate a plan built from fully-static shapes.
121pub fn validate_static(
122    plan: &ActivationPlan,
123    graph: &Graph,
124    view_map: &ViewMap,
125    options: &PlanOptions,
126) -> Result<(), ValidateError> {
127    let oracle = static_size_oracle(graph);
128    validate(plan, graph, view_map, oracle, options)
129}
onnx_runtime_memory/validate.rs

onnx_runtime_memory/
validate.rs
