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    Tensors(Vec<TensorData>),
99    SparseTensor(SparseTensorData),
100    SparseTensors(Vec<SparseTensorData>),
101    /// A subgraph body (control-flow ops: If/Loop/Scan). Stored inline; the
102    /// owning [`Graph`] also indexes it in `subgraphs` for traversal.
103    Graph(Box<Graph>),
104    Graphs(Vec<Graph>),
105    TypeProto(TypeProto),
106    TypeProtos(Vec<TypeProto>),
107}
108
109impl Attribute {
110    /// The `i64` value, if this is an [`Attribute::Int`].
111    pub fn as_int(&self) -> Option<i64> {
112        match self {
113            Attribute::Int(v) => Some(*v),
114            _ => None,
115        }
116    }
117
118    /// The `f32` value, if this is an [`Attribute::Float`].
119    pub fn as_float(&self) -> Option<f32> {
120        match self {
121            Attribute::Float(v) => Some(*v),
122            _ => None,
123        }
124    }
125
126    /// The value as UTF-8 text, if this is an [`Attribute::String`] whose bytes
127    /// are valid UTF-8. Returns `None` for a non-string attribute or for string
128    /// bytes that are not valid UTF-8 (e.g. an opaque binary payload).
129    pub fn as_str(&self) -> Option<&str> {
130        match self {
131            Attribute::String(v) => std::str::from_utf8(v).ok(),
132            _ => None,
133        }
134    }
135
136    /// The raw bytes of an [`Attribute::String`], regardless of whether they are
137    /// valid UTF-8. Returns `None` for any other attribute kind.
138    pub fn as_bytes(&self) -> Option<&[u8]> {
139        match self {
140            Attribute::String(v) => Some(v),
141            _ => None,
142        }
143    }
144
145    /// The `&[i64]` slice, if this is an [`Attribute::Ints`].
146    pub fn as_ints(&self) -> Option<&[i64]> {
147        match self {
148            Attribute::Ints(v) => Some(v),
149            _ => None,
150        }
151    }
152
153    /// Interpret an `Ints` attribute as a shape of static dims.
154    pub fn as_shape(&self) -> Option<Shape> {
155        self.as_ints()
156            .map(|v| v.iter().map(|&d| (d as usize).into()).collect())
157    }
158}
onnx_runtime_ir/node.rs

onnx_runtime_ir/
node.rs
