onnx_runtime_ir/
value.rs

1//! Graph values: the typed, shaped SSA edges between nodes.
2
3use std::collections::HashSet;
4use std::fmt;
5
6use crate::arena::ArenaKey;
7use crate::device::DeviceId;
8use crate::dtype::DataType;
9use crate::layout::TensorLayout;
10use crate::node::NodeId;
11use crate::shape::Shape;
12
13/// Unique identifier for a [`Value`] within a [`Graph`](crate::Graph).
14#[derive(Clone, Copy, PartialEq, Eq, Hash, Debug)]
15pub struct ValueId(pub u32);
16
17impl ArenaKey for ValueId {
18    fn from_raw(raw: u32) -> Self {
19        ValueId(raw)
20    }
21    fn to_raw(self) -> u32 {
22        self.0
23    }
24}
25
26/// A node input slot that uses a value.
27pub type Usage = (NodeId, u32);
28
29/// The unordered internal set of input slots that consume a value.
30///
31/// Iteration is intentionally exposed only through sorted snapshots so hash
32/// iteration order cannot affect graph rewrites or serialized output.
33#[derive(Clone, Default, PartialEq, Eq)]
34pub struct Consumers {
35    uses: HashSet<Usage>,
36}
37
38impl Consumers {
39    pub(crate) fn insert(&mut self, node: NodeId, input_index: u32) {
40        self.uses.insert((node, input_index));
41    }
42
43    pub(crate) fn remove(&mut self, node: NodeId, input_index: u32) -> bool {
44        self.uses.remove(&(node, input_index))
45    }
46
47    pub(crate) fn contains(&self, node: NodeId, input_index: u32) -> bool {
48        self.uses.contains(&(node, input_index))
49    }
50
51    /// Number of consuming input slots.
52    pub fn len(&self) -> usize {
53        self.uses.len()
54    }
55
56    /// Whether no node input slot consumes the value.
57    pub fn is_empty(&self) -> bool {
58        self.uses.is_empty()
59    }
60
61    /// Consuming input slots sorted by `(NodeId, input_index)`.
62    pub fn uses(&self) -> Vec<Usage> {
63        let mut uses: Vec<_> = self.uses.iter().copied().collect();
64        uses.sort_unstable_by_key(|&(node, input_index)| (node.0, input_index));
65        uses
66    }
67
68    /// Distinct consuming nodes sorted by ascending [`NodeId`].
69    pub fn nodes(&self) -> Vec<NodeId> {
70        let mut nodes: Vec<_> = self.uses.iter().map(|&(node, _)| node).collect();
71        nodes.sort_unstable_by_key(|node| node.0);
72        nodes.dedup();
73        nodes
74    }
75}
76
77impl fmt::Debug for Consumers {
78    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
79        f.debug_list().entries(self.uses()).finish()
80    }
81}
82
83/// A value flowing through the graph — the output of one node and the input of
84/// zero or more others (SSA form: at most one producer).
85///
86/// Carries a first-class [`TensorLayout`] and optional [`DeviceId`] placement,
87/// beyond ONNX's plain type+shape.
88#[derive(Clone, Debug, PartialEq)]
89pub struct Value {
90    pub id: ValueId,
91    /// Optional name (graph inputs/outputs and initializers are named; interior
92    /// SSA values may be anonymous).
93    pub name: Option<String>,
94    pub dtype: DataType,
95    pub shape: Shape,
96    pub layout: TensorLayout,
97    /// Device placement, filled in by the placement pass.
98    pub device: Option<DeviceId>,
99    /// The node that produces this value, or `None` for graph inputs and
100    /// initializers.
101    pub producer: Option<NodeId>,
102    /// Input slots that consume this value, keyed by `(node, input_index)`.
103    pub consumers: Consumers,
104    /// Whether this value is present in [`Graph::inputs`](crate::Graph::inputs).
105    pub is_graph_input: bool,
106    /// Whether this value is present in [`Graph::outputs`](crate::Graph::outputs).
107    pub is_graph_output: bool,
108}
109
110impl Value {
111    /// A new anonymous value with a contiguous default layout and no edges.
112    pub fn new(id: ValueId, dtype: DataType, shape: Shape) -> Self {
113        Self {
114            id,
115            name: None,
116            dtype,
117            shape,
118            layout: TensorLayout::contiguous(),
119            device: None,
120            producer: None,
121            consumers: Consumers::default(),
122            is_graph_input: false,
123            is_graph_output: false,
124        }
125    }
126
127    /// Whether this value has no producing node (graph input / initializer).
128    pub fn is_source(&self) -> bool {
129        self.producer.is_none()
130    }
131
132    /// The static rank (number of dimensions).
133    pub fn rank(&self) -> usize {
134        self.shape.len()
135    }
136}
onnx_runtime_ir/value.rs

onnx_runtime_ir/
value.rs
