onnx_runtime_memory/
view_map.rs

1//! Zero-copy **view → source** aliasing supplied by the executor.
2//!
3//! The executor treats the outputs of layout/movement ops
4//! (`Slice`/`Reshape`/`Squeeze`/`Unsqueeze`/`Transpose`/`Expand`, …) as
5//! *zero-copy views*: they own no buffer and instead borrow (alias) a source
6//! buffer with some strided geometry (see `ValueView` in
7//! `onnx-runtime-session/src/executor.rs`). The source buffer is therefore
8//! *pinned* — it must outlive every view that aliases it, or a reused buffer
9//! would clobber a still-live alias (a use-after-free / silent-corruption bug).
10//!
11//! The planner does **not** hardcode op names. Instead the caller (the
12//! executor, which already computes this) supplies a [`ViewMap`] of
13//! `view → source` edges. The planner folds those edges transitively to a
14//! *root* owner and extends the root's live interval to cover every view's last
15//! use. This is exactly what prevents the greedy allocator from recycling a
16//! buffer that a live view still points into.
17
18use std::collections::HashMap;
19
20use onnx_runtime_ir::ValueId;
21
22/// A set of zero-copy `view → source` aliasing relationships.
23///
24/// Each entry means "`view` owns no activation slot; its bytes live inside
25/// `source`'s buffer". A view of a view is allowed; [`ViewMap::root`] folds a
26/// chain down to the single real buffer owner.
27#[derive(Clone, Debug, Default)]
28pub struct ViewMap {
29    /// `view → immediate source`. The source may itself be a view.
30    edges: HashMap<ValueId, ValueId>,
31}
32
33impl ViewMap {
34    /// An empty map (no views — every value owns its own buffer).
35    pub fn new() -> Self {
36        Self::default()
37    }
38
39    /// Build a map from `(view, source)` pairs.
40    pub fn from_pairs(pairs: impl IntoIterator<Item = (ValueId, ValueId)>) -> Self {
41        let mut m = Self::new();
42        for (view, source) in pairs {
43            m.insert(view, source);
44        }
45        m
46    }
47
48    /// Record that `view` aliases `source` (owns no buffer of its own).
49    pub fn insert(&mut self, view: ValueId, source: ValueId) {
50        self.edges.insert(view, source);
51    }
52
53    /// Whether `value` is a zero-copy view (aliases another value's buffer and
54    /// therefore gets no slot of its own).
55    pub fn is_view(&self, value: ValueId) -> bool {
56        self.edges.contains_key(&value)
57    }
58
59    /// The immediate source `value` aliases, if it is a view.
60    pub fn source_of(&self, value: ValueId) -> Option<ValueId> {
61        self.edges.get(&value).copied()
62    }
63
64    /// Fold `value` to the **root** buffer owner by following `view → source`
65    /// edges transitively. Returns `value` itself when it is not a view.
66    ///
67    /// A cycle in the (malformed) view map is broken defensively: the last node
68    /// visited before the cycle closes is returned rather than looping forever.
69    pub fn root(&self, value: ValueId) -> ValueId {
70        let mut seen: Vec<ValueId> = Vec::new();
71        let mut cur = value;
72        while let Some(&next) = self.edges.get(&cur) {
73            if seen.contains(&cur) {
74                break;
75            }
76            seen.push(cur);
77            cur = next;
78        }
79        cur
80    }
81
82    /// Number of recorded view edges.
83    pub fn len(&self) -> usize {
84        self.edges.len()
85    }
86
87    /// Whether there are no view edges.
88    pub fn is_empty(&self) -> bool {
89        self.edges.is_empty()
90    }
91}
onnx_runtime_memory/view_map.rs

onnx_runtime_memory/
view_map.rs
