onnx_runtime_ir/
node.rs

1//! Graph nodes (operations) and their attributes.
2
3use std::collections::HashMap;
4
5use crate::arena::ArenaKey;
6use crate::device::DeviceId;
7use crate::graph::Graph;
8use crate::shape::Shape;
9use crate::tensor::{SparseTensorData, TensorData, TypeProto};
10use crate::value::ValueId;
11
12/// Unique identifier for a [`Node`] within a [`Graph`](crate::Graph).
13#[derive(Clone, Copy, PartialEq, Eq, Hash, Debug)]
14pub struct NodeId(pub u32);
15
16impl ArenaKey for NodeId {
17    fn from_raw(raw: u32) -> Self {
18        NodeId(raw)
19    }
20    fn to_raw(self) -> u32 {
21        self.0
22    }
23}
24
25/// An operation in the graph.
26///
27/// Inputs are `Option<ValueId>` because ONNX ops may have optional (skipped)
28/// inputs represented by empty names; a `None` slot preserves positional
29/// arity. Outputs are always present (SSA values).
30#[derive(Clone, Debug)]
31pub struct Node {
32    pub id: NodeId,
33    /// Optional ONNX node name (`""` means unnamed).
34    pub name: String,
35    pub op_type: String,
36    /// Operator domain (`""` == the default ONNX domain).
37    pub domain: String,
38    pub inputs: Vec<Option<ValueId>>,
39    pub outputs: Vec<ValueId>,
40    pub attributes: HashMap<String, Attribute>,
41    pub doc_string: Option<String>,
42    /// Device placement, filled in by the placement pass.
43    pub device: Option<DeviceId>,
44    /// Position in the final execution schedule, filled in by the scheduler.
45    pub exec_order: Option<usize>,
46}
47
48impl Node {
49    /// A new node with the given op type and edges, and no attributes.
50    pub fn new(
51        id: NodeId,
52        op_type: impl Into<String>,
53        inputs: Vec<Option<ValueId>>,
54        outputs: Vec<ValueId>,
55    ) -> Self {
56        Self {
57            id,
58            name: String::new(),
59            op_type: op_type.into(),
60            domain: String::new(),
61            inputs,
62            outputs,
63            attributes: HashMap::new(),
64            doc_string: None,
65            device: None,
66            exec_order: None,
67        }
68    }
69
70    /// Iterate over the present (non-skipped) input value ids.
71    pub fn input_values(&self) -> impl Iterator<Item = ValueId> + '_ {
72        self.inputs.iter().filter_map(|slot| *slot)
73    }
74
75    /// Look up an attribute by name.
76    pub fn attr(&self, name: &str) -> Option<&Attribute> {
77        self.attributes.get(name)
78    }
79
80    /// Whether this node belongs to the default ONNX operator domain.
81    ///
82    /// Relies on the post-load invariant that the loader canonicalizes the
83    /// default domain to `""` (see [`crate::normalize_domain`]), so this is a
84    /// simple emptiness test — the `"ai.onnx"` spelling never reaches loaded IR.
85    #[inline]
86    #[must_use]
87    pub fn is_default_domain(&self) -> bool {
88        self.domain.is_empty()
89    }
90}
91
92/// An ONNX operator attribute. Covers all attribute value kinds.
93#[derive(Clone, Debug)]
94pub enum Attribute {
95    Int(i64),
96    Float(f32),
97    /// An ONNX `STRING` attribute. Stored as **raw bytes**, not `String`, so
98    /// that the load/dump path round-trips the payload byte-exactly: ONNX
99    /// `STRING` attributes are arbitrary byte strings (e.g. an opaque compiled
100    /// blob) that are not guaranteed to be valid UTF-8. Use [`Attribute::as_str`]
101    /// to view the bytes as UTF-8 text when that is meaningful.
102    String(Vec<u8>),
103    Ints(Vec<i64>),
104    Floats(Vec<f32>),
105    /// An ONNX `STRINGS` attribute — a list of raw byte strings (see
106    /// [`Attribute::String`] for why bytes rather than `String`).
107    Strings(Vec<Vec<u8>>),
108    Tensor(TensorData),
109    Tensors(Vec<TensorData>),
110    SparseTensor(SparseTensorData),
111    SparseTensors(Vec<SparseTensorData>),
112    /// A subgraph body (control-flow ops: If/Loop/Scan). Stored inline; the
113    /// owning [`Graph`] also indexes it in `subgraphs` for traversal.
114    Graph(Box<Graph>),
115    Graphs(Vec<Graph>),
116    TypeProto(TypeProto),
117    TypeProtos(Vec<TypeProto>),
118}
119
120impl Attribute {
121    /// The `i64` value, if this is an [`Attribute::Int`].
122    pub fn as_int(&self) -> Option<i64> {
123        match self {
124            Attribute::Int(v) => Some(*v),
125            _ => None,
126        }
127    }
128
129    /// The `f32` value, if this is an [`Attribute::Float`].
130    pub fn as_float(&self) -> Option<f32> {
131        match self {
132            Attribute::Float(v) => Some(*v),
133            _ => None,
134        }
135    }
136
137    /// The value as UTF-8 text, if this is an [`Attribute::String`] whose bytes
138    /// are valid UTF-8. Returns `None` for a non-string attribute or for string
139    /// bytes that are not valid UTF-8 (e.g. an opaque binary payload).
140    pub fn as_str(&self) -> Option<&str> {
141        match self {
142            Attribute::String(v) => std::str::from_utf8(v).ok(),
143            _ => None,
144        }
145    }
146
147    /// The raw bytes of an [`Attribute::String`], regardless of whether they are
148    /// valid UTF-8. Returns `None` for any other attribute kind.
149    pub fn as_bytes(&self) -> Option<&[u8]> {
150        match self {
151            Attribute::String(v) => Some(v),
152            _ => None,
153        }
154    }
155
156    /// The `&[i64]` slice, if this is an [`Attribute::Ints`].
157    pub fn as_ints(&self) -> Option<&[i64]> {
158        match self {
159            Attribute::Ints(v) => Some(v),
160            _ => None,
161        }
162    }
163
164    /// Interpret an `Ints` attribute as a shape of static dims.
165    pub fn as_shape(&self) -> Option<Shape> {
166        self.as_ints()
167            .map(|v| v.iter().map(|&d| (d as usize).into()).collect())
168    }
169}
onnx_runtime_ir/node.rs

onnx_runtime_ir/
node.rs
