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
81/// An ONNX operator attribute. Covers all attribute value kinds.
82#[derive(Clone, Debug)]
83pub enum Attribute {
84    Int(i64),
85    Float(f32),
86    /// An ONNX `STRING` attribute. Stored as **raw bytes**, not `String`, so
87    /// that the load/dump path round-trips the payload byte-exactly: ONNX
88    /// `STRING` attributes are arbitrary byte strings (e.g. an opaque compiled
89    /// blob) that are not guaranteed to be valid UTF-8. Use [`Attribute::as_str`]
90    /// to view the bytes as UTF-8 text when that is meaningful.
91    String(Vec<u8>),
92    Ints(Vec<i64>),
93    Floats(Vec<f32>),
94    /// An ONNX `STRINGS` attribute — a list of raw byte strings (see
95    /// [`Attribute::String`] for why bytes rather than `String`).
96    Strings(Vec<Vec<u8>>),
97    Tensor(TensorData),
98    SparseTensor(SparseTensorData),
99    /// A subgraph body (control-flow ops: If/Loop/Scan). Stored inline; the
100    /// owning [`Graph`] also indexes it in `subgraphs` for traversal.
101    Graph(Box<Graph>),
102    Graphs(Vec<Graph>),
103    TypeProto(TypeProto),
104}
105
106impl Attribute {
107    /// The `i64` value, if this is an [`Attribute::Int`].
108    pub fn as_int(&self) -> Option<i64> {
109        match self {
110            Attribute::Int(v) => Some(*v),
111            _ => None,
112        }
113    }
114
115    /// The `f32` value, if this is an [`Attribute::Float`].
116    pub fn as_float(&self) -> Option<f32> {
117        match self {
118            Attribute::Float(v) => Some(*v),
119            _ => None,
120        }
121    }
122
123    /// The value as UTF-8 text, if this is an [`Attribute::String`] whose bytes
124    /// are valid UTF-8. Returns `None` for a non-string attribute or for string
125    /// bytes that are not valid UTF-8 (e.g. an opaque binary payload).
126    pub fn as_str(&self) -> Option<&str> {
127        match self {
128            Attribute::String(v) => std::str::from_utf8(v).ok(),
129            _ => None,
130        }
131    }
132
133    /// The raw bytes of an [`Attribute::String`], regardless of whether they are
134    /// valid UTF-8. Returns `None` for any other attribute kind.
135    pub fn as_bytes(&self) -> Option<&[u8]> {
136        match self {
137            Attribute::String(v) => Some(v),
138            _ => None,
139        }
140    }
141
142    /// The `&[i64]` slice, if this is an [`Attribute::Ints`].
143    pub fn as_ints(&self) -> Option<&[i64]> {
144        match self {
145            Attribute::Ints(v) => Some(v),
146            _ => None,
147        }
148    }
149
150    /// Interpret an `Ints` attribute as a shape of static dims.
151    pub fn as_shape(&self) -> Option<Shape> {
152        self.as_ints()
153            .map(|v| v.iter().map(|&d| (d as usize).into()).collect())
154    }
155}
onnx_runtime_ir/node.rs

onnx_runtime_ir/
node.rs
