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    pub op_type: String,
34    /// Operator domain (`""` == the default ONNX domain).
35    pub domain: String,
36    pub inputs: Vec<Option<ValueId>>,
37    pub outputs: Vec<ValueId>,
38    pub attributes: HashMap<String, Attribute>,
39    pub doc_string: Option<String>,
40    /// Device placement, filled in by the placement pass.
41    pub device: Option<DeviceId>,
42    /// Position in the final execution schedule, filled in by the scheduler.
43    pub exec_order: Option<usize>,
44}
45
46impl Node {
47    /// A new node with the given op type and edges, and no attributes.
48    pub fn new(
49        id: NodeId,
50        op_type: impl Into<String>,
51        inputs: Vec<Option<ValueId>>,
52        outputs: Vec<ValueId>,
53    ) -> Self {
54        Self {
55            id,
56            op_type: op_type.into(),
57            domain: String::new(),
58            inputs,
59            outputs,
60            attributes: HashMap::new(),
61            doc_string: None,
62            device: None,
63            exec_order: None,
64        }
65    }
66
67    /// Iterate over the present (non-skipped) input value ids.
68    pub fn input_values(&self) -> impl Iterator<Item = ValueId> + '_ {
69        self.inputs.iter().filter_map(|slot| *slot)
70    }
71
72    /// Look up an attribute by name.
73    pub fn attr(&self, name: &str) -> Option<&Attribute> {
74        self.attributes.get(name)
75    }
76}
77
78/// An ONNX operator attribute. Covers all attribute value kinds.
79#[derive(Clone, Debug)]
80pub enum Attribute {
81    Int(i64),
82    Float(f32),
83    /// An ONNX `STRING` attribute. Stored as **raw bytes**, not `String`, so
84    /// that the load/dump path round-trips the payload byte-exactly: ONNX
85    /// `STRING` attributes are arbitrary byte strings (e.g. an opaque compiled
86    /// blob) that are not guaranteed to be valid UTF-8. Use [`Attribute::as_str`]
87    /// to view the bytes as UTF-8 text when that is meaningful.
88    String(Vec<u8>),
89    Ints(Vec<i64>),
90    Floats(Vec<f32>),
91    /// An ONNX `STRINGS` attribute — a list of raw byte strings (see
92    /// [`Attribute::String`] for why bytes rather than `String`).
93    Strings(Vec<Vec<u8>>),
94    Tensor(TensorData),
95    SparseTensor(SparseTensorData),
96    /// A subgraph body (control-flow ops: If/Loop/Scan). Stored inline; the
97    /// owning [`Graph`] also indexes it in `subgraphs` for traversal.
98    Graph(Box<Graph>),
99    Graphs(Vec<Graph>),
100    TypeProto(TypeProto),
101}
102
103impl Attribute {
104    /// The `i64` value, if this is an [`Attribute::Int`].
105    pub fn as_int(&self) -> Option<i64> {
106        match self {
107            Attribute::Int(v) => Some(*v),
108            _ => None,
109        }
110    }
111
112    /// The `f32` value, if this is an [`Attribute::Float`].
113    pub fn as_float(&self) -> Option<f32> {
114        match self {
115            Attribute::Float(v) => Some(*v),
116            _ => None,
117        }
118    }
119
120    /// The value as UTF-8 text, if this is an [`Attribute::String`] whose bytes
121    /// are valid UTF-8. Returns `None` for a non-string attribute or for string
122    /// bytes that are not valid UTF-8 (e.g. an opaque binary payload).
123    pub fn as_str(&self) -> Option<&str> {
124        match self {
125            Attribute::String(v) => std::str::from_utf8(v).ok(),
126            _ => None,
127        }
128    }
129
130    /// The raw bytes of an [`Attribute::String`], regardless of whether they are
131    /// valid UTF-8. Returns `None` for any other attribute kind.
132    pub fn as_bytes(&self) -> Option<&[u8]> {
133        match self {
134            Attribute::String(v) => Some(v),
135            _ => None,
136        }
137    }
138
139    /// The `&[i64]` slice, if this is an [`Attribute::Ints`].
140    pub fn as_ints(&self) -> Option<&[i64]> {
141        match self {
142            Attribute::Ints(v) => Some(v),
143            _ => None,
144        }
145    }
146
147    /// Interpret an `Ints` attribute as a shape of static dims.
148    pub fn as_shape(&self) -> Option<Shape> {
149        self.as_ints()
150            .map(|v| v.iter().map(|&d| (d as usize).into()).collect())
151    }
152}
onnx_runtime_ir/node.rs

onnx_runtime_ir/
node.rs
