onnx_runtime_loader/
graph_builder.rs

1//! Build an [`onnx_runtime_ir::Graph`] from a decoded `ModelProto` (§19.1).
2//!
3//! Responsible for the graph-construction invariants of `docs/ORT2.md` §3.5:
4//! stable value ids, unique node outputs (SSA), source values for inputs and
5//! initializers, and interning symbolic dims that share a protobuf name.
6
7use std::collections::HashMap;
8
9use onnx_runtime_ir::{
10    Attribute, DataType, Dim, Graph, Node, NodeId, Shape, TensorData, TypeProto, ValueId,
11};
12
13use crate::proto::onnx::{
14    self, attribute_proto, tensor_shape_proto, type_proto, AttributeProto, GraphProto, ModelProto,
15    TensorShapeProto,
16};
17use crate::weights::tensor_data_from_proto;
18use crate::LoaderError;
19
20/// The result of building a graph: the IR graph plus the mapping from ONNX
21/// tensor names to the value ids they were assigned (needed by the weight
22/// loader).
23pub struct BuiltGraph {
24    pub graph: Graph,
25    pub name_map: HashMap<String, ValueId>,
26}
27
28/// Build the IR graph (nodes, values, symbols, opsets) from a `ModelProto`.
29///
30/// Weights and shape inference are applied by later pipeline stages.
31pub fn build_graph(model: &ModelProto) -> Result<BuiltGraph, LoaderError> {
32    let mut graph = Graph::new();
33
34    // Opset imports: domain -> version.
35    for opset in &model.opset_import {
36        if opset.version > 0 {
37            graph
38                .opset_imports
39                .insert(opset.domain.clone(), opset.version as u64);
40        }
41    }
42
43    let graph_proto = model
44        .graph
45        .as_ref()
46        .ok_or_else(|| LoaderError::GraphBuild("ModelProto has no graph".into()))?;
47
48    let name_map = build_graph_proto(&mut graph, graph_proto, true)?;
49
50    graph
51        .validate()
52        .map_err(|errs| LoaderError::GraphBuild(format!("{errs:?}")))?;
53
54    Ok(BuiltGraph { graph, name_map })
55}
56
57/// Populate `graph` from a `GraphProto`. When `is_top_level` is true, inputs
58/// and outputs are registered as graph I/O. Returns the name→value map for the
59/// values created in this graph scope.
60fn build_graph_proto(
61    graph: &mut Graph,
62    gp: &GraphProto,
63    is_top_level: bool,
64) -> Result<HashMap<String, ValueId>, LoaderError> {
65    let mut names: HashMap<String, ValueId> = HashMap::new();
66
67    // 1. Initializers: fully-typed source values (producer = None).
68    for init in &gp.initializer {
69        if init.name.is_empty() {
70            continue;
71        }
72        let dtype = decode_dtype(init.data_type, || {
73            format!("initializer '{}'", init.name)
74        })?;
75        let dims: Vec<usize> = init.dims.iter().map(|&d| d.max(0) as usize).collect();
76        let shape: Shape = dims.into_iter().map(Dim::Static).collect();
77        let vid = graph.create_named_value(init.name.clone(), dtype, shape);
78        names.insert(init.name.clone(), vid);
79    }
80
81    // 2. Graph inputs. Names that are also initializers are constants, not
82    //    real graph inputs (invariant §3.5.3).
83    for vi in &gp.input {
84        if vi.name.is_empty() {
85            continue;
86        }
87        if names.contains_key(&vi.name) {
88            continue; // initializer-backed constant input
89        }
90        let (dtype, shape) = value_info_type(graph, vi)?;
91        let vid = graph.create_named_value(vi.name.clone(), dtype, shape);
92        names.insert(vi.name.clone(), vid);
93        if is_top_level {
94            graph.add_input(vid);
95        }
96    }
97
98    // 3. Declared value_info type hints for interior values.
99    for vi in &gp.value_info {
100        if vi.name.is_empty() || names.contains_key(&vi.name) {
101            continue;
102        }
103        let (dtype, shape) = value_info_type(graph, vi)?;
104        let vid = graph.create_named_value(vi.name.clone(), dtype, shape);
105        names.insert(vi.name.clone(), vid);
106    }
107
108    // 4. Graph outputs: typed values that a node will later produce by name.
109    for vi in &gp.output {
110        if vi.name.is_empty() {
111            continue;
112        }
113        if !names.contains_key(&vi.name) {
114            let (dtype, shape) = value_info_type(graph, vi)?;
115            let vid = graph.create_named_value(vi.name.clone(), dtype, shape);
116            names.insert(vi.name.clone(), vid);
117        }
118    }
119
120    // 5. Nodes: wire inputs/outputs, converting attributes.
121    for np in &gp.node {
122        let inputs: Vec<Option<ValueId>> = np
123            .input
124            .iter()
125            .map(|name| {
126                if name.is_empty() {
127                    None
128                } else {
129                    Some(get_or_create(graph, &mut names, name))
130                }
131            })
132            .collect();
133
134        let outputs: Vec<ValueId> = np
135            .output
136            .iter()
137            .map(|name| {
138                if name.is_empty() {
139                    // An omitted (unused) optional output: give it an anonymous
140                    // value to preserve positional arity.
141                    graph.create_value(DataType::Float32, Vec::new())
142                } else {
143                    get_or_create(graph, &mut names, name)
144                }
145            })
146            .collect();
147
148        let mut node = Node::new(NodeId(0), np.op_type.clone(), inputs, outputs);
149        node.domain = np.domain.clone();
150        if !np.doc_string.is_empty() {
151            node.doc_string = Some(np.doc_string.clone());
152        }
153        for ap in &np.attribute {
154            if let Some((key, attr)) = convert_attribute(graph, ap)? {
155                node.attributes.insert(key, attr);
156            }
157        }
158
159        let nid = graph.insert_node(node);
160        register_subgraphs(graph, nid);
161    }
162
163    // 6. Register graph outputs in order.
164    if is_top_level {
165        for vi in &gp.output {
166            if let Some(&vid) = names.get(&vi.name) {
167                graph.add_output(vid);
168            }
169        }
170    }
171
172    Ok(names)
173}
174
175/// After a node is inserted, move any `Graph`/`Graphs` attribute bodies into the
176/// graph's `subgraphs` index so traversal/validation can reach them (§3.3).
177fn register_subgraphs(graph: &mut Graph, nid: NodeId) {
178    let attrs: Vec<(String, usize)> = graph
179        .node(nid)
180        .attributes
181        .iter()
182        .filter_map(|(k, v)| match v {
183            Attribute::Graph(_) => Some((k.clone(), 1)),
184            Attribute::Graphs(gs) => Some((k.clone(), gs.len())),
185            _ => None,
186        })
187        .collect();
188    for (key, count) in attrs {
189        match graph.node(nid).attributes.get(&key) {
190            Some(Attribute::Graph(g)) => {
191                let sub = (**g).clone();
192                graph.subgraphs.insert((nid, key), sub);
193            }
194            Some(Attribute::Graphs(_)) => {
195                for i in 0..count {
196                    if let Some(Attribute::Graphs(gs)) = graph.node(nid).attributes.get(&key) {
197                        let sub = gs[i].clone();
198                        graph.subgraphs.insert((nid, format!("{key}[{i}]")), sub);
199                    }
200                }
201            }
202            _ => {}
203        }
204    }
205}
206
207/// Fetch the value id for `name`, creating a placeholder value if it does not
208/// exist yet (interior SSA value with as-yet-unknown type).
209fn get_or_create(
210    graph: &mut Graph,
211    names: &mut HashMap<String, ValueId>,
212    name: &str,
213) -> ValueId {
214    if let Some(&vid) = names.get(name) {
215        return vid;
216    }
217    let vid = graph.create_named_value(name.to_string(), DataType::Float32, Vec::new());
218    names.insert(name.to_string(), vid);
219    vid
220}
221
222/// Decode a raw ONNX `TensorProto.DataType` integer into an IR [`DataType`],
223/// failing closed when the runtime does not model the type. This prevents an
224/// unmodeled dtype (e.g. `COMPLEX64` = 14) from being silently mislabeled as
225/// `Float32` at any tensor-type decode site (§19.1).
226fn decode_dtype(
227    raw: i32,
228    context: impl FnOnce() -> String,
229) -> Result<DataType, LoaderError> {
230    DataType::from_onnx(raw).ok_or_else(|| LoaderError::UnsupportedDataType {
231        raw,
232        context: context(),
233    })
234}
235
236/// Extract `(dtype, shape)` from a `ValueInfoProto`'s tensor type, interning
237/// symbolic dims into `graph` by name.
238fn value_info_type(
239    graph: &mut Graph,
240    vi: &onnx::ValueInfoProto,
241) -> Result<(DataType, Shape), LoaderError> {
242    match vi.r#type.as_ref() {
243        Some(tp) => type_proto_to_dtype_shape(graph, tp, &vi.name),
244        // A value-info with no type at all is genuinely untyped, not an
245        // unmodeled dtype: keep the tensor-centric placeholder default.
246        None => Ok((DataType::Float32, Vec::new())),
247    }
248}
249
250fn type_proto_to_dtype_shape(
251    graph: &mut Graph,
252    tp: &onnx::TypeProto,
253    name: &str,
254) -> Result<(DataType, Shape), LoaderError> {
255    match tp.value.as_ref() {
256        Some(type_proto::Value::TensorType(t)) => {
257            let dtype = decode_dtype(t.elem_type, || format!("value-info '{name}'"))?;
258            let shape = t
259                .shape
260                .as_ref()
261                .map(|s| tensor_shape_to_shape(graph, s))
262                .unwrap_or_default();
263            Ok((dtype, shape))
264        }
265        Some(type_proto::Value::SparseTensorType(t)) => {
266            let dtype = decode_dtype(t.elem_type, || format!("value-info '{name}'"))?;
267            let shape = t
268                .shape
269                .as_ref()
270                .map(|s| tensor_shape_to_shape(graph, s))
271                .unwrap_or_default();
272            Ok((dtype, shape))
273        }
274        // Non-tensor containers (sequence/map/optional): the IR value model is
275        // tensor-centric; record a placeholder type. These do not occur in the
276        // Phase-1 (BERT) op set.
277        _ => Ok((DataType::Float32, Vec::new())),
278    }
279}
280
281/// Convert an ONNX `TensorShapeProto` to an IR [`Shape`], interning dim-params
282/// by name (invariant §3.5.4) and allocating fresh anonymous symbols for
283/// unknown dims.
284fn tensor_shape_to_shape(graph: &mut Graph, tsp: &TensorShapeProto) -> Shape {
285    tsp.dim
286        .iter()
287        .map(|d| match d.value.as_ref() {
288            Some(tensor_shape_proto::dimension::Value::DimValue(v)) if *v >= 0 => {
289                Dim::Static(*v as usize)
290            }
291            Some(tensor_shape_proto::dimension::Value::DimParam(name)) if !name.is_empty() => {
292                Dim::Symbolic(graph.intern_symbol(name))
293            }
294            // Unknown dim (no value, negative, or empty param): fresh symbol.
295            _ => Dim::Symbolic(graph.create_symbol(None)),
296        })
297        .collect()
298}
299
300/// Convert an `AttributeProto` to an IR `(name, Attribute)`. Returns `Ok(None)`
301/// for empty/absent attributes, and errors (rather than silently dropping) on
302/// tensor/sparse-tensor/type-proto **list** kinds the IR does not model, which
303/// do not appear in the Phase-1 op set.
304fn convert_attribute(
305    graph: &mut Graph,
306    ap: &AttributeProto,
307) -> Result<Option<(String, Attribute)>, LoaderError> {
308    use attribute_proto::AttributeType as AT;
309
310    // Determine the attribute kind, falling back to field-presence heuristics
311    // for IR<0.0.2 protos where `type` may be unset.
312    let ty = AT::try_from(ap.r#type).unwrap_or(AT::Undefined);
313
314    let attr = match ty {
315        AT::Float => Attribute::Float(ap.f),
316        AT::Int => Attribute::Int(ap.i),
317        // STRING attributes are arbitrary byte strings on the wire (an opaque
318        // blob, a path, or text). Preserve the exact bytes rather than lossily
319        // decoding as UTF-8, so encode is a byte-exact inverse of decode.
320        AT::String => Attribute::String(ap.s.clone()),
321        AT::Floats => Attribute::Floats(ap.floats.clone()),
322        AT::Ints => Attribute::Ints(ap.ints.clone()),
323        AT::Strings => Attribute::Strings(ap.strings.clone()),
324        AT::Tensor => match ap.t.as_ref() {
325            Some(t) => Attribute::Tensor(convert_tensor(t)?),
326            None => return Ok(None),
327        },
328        AT::Graph => match ap.g.as_ref() {
329            Some(g) => {
330                let mut sub = Graph::new();
331                build_graph_proto(&mut sub, g, false)?;
332                Attribute::Graph(Box::new(sub))
333            }
334            None => return Ok(None),
335        },
336        AT::Graphs => {
337            let mut subs = Vec::new();
338            for g in &ap.graphs {
339                let mut sub = Graph::new();
340                build_graph_proto(&mut sub, g, false)?;
341                subs.push(sub);
342            }
343            Attribute::Graphs(subs)
344        }
345        AT::TypeProto => match ap.tp.as_ref() {
346            Some(tp) => Attribute::TypeProto(convert_type_proto(graph, tp)?),
347            None => return Ok(None),
348        },
349        // Field-presence fallback when `type` is UNDEFINED.
350        AT::Undefined => {
351            if let Some(t) = ap.t.as_ref() {
352                Attribute::Tensor(convert_tensor(t)?)
353            } else if !ap.floats.is_empty() {
354                Attribute::Floats(ap.floats.clone())
355            } else if !ap.ints.is_empty() {
356                Attribute::Ints(ap.ints.clone())
357            } else if !ap.strings.is_empty() {
358                Attribute::Strings(ap.strings.clone())
359            } else if !ap.s.is_empty() {
360                Attribute::String(ap.s.clone())
361            } else if ap.i != 0 {
362                Attribute::Int(ap.i)
363            } else if ap.f != 0.0 {
364                Attribute::Float(ap.f)
365            } else {
366                return Ok(None);
367            }
368        }
369        // Tensor / sparse-tensor / type-proto list attributes have no IR
370        // variant. Surface a clean error rather than silently dropping the
371        // attribute, which could otherwise mis-model an op (§19.1).
372        AT::Tensors | AT::SparseTensor | AT::SparseTensors | AT::TypeProtos => {
373            return Err(LoaderError::GraphBuild(format!(
374                "attribute {:?} has unmodeled type {:?}",
375                ap.name, ty
376            )));
377        }
378    };
379    Ok(Some((ap.name.clone(), attr)))
380}
381
382fn convert_tensor(t: &onnx::TensorProto) -> Result<TensorData, LoaderError> {
383    let dtype = decode_dtype(t.data_type, || format!("attribute tensor '{}'", t.name))?;
384    let dims: Vec<usize> = t.dims.iter().map(|&d| d.max(0) as usize).collect();
385    tensor_data_from_proto(t, dtype, &dims)
386}
387
388fn convert_type_proto(
389    graph: &mut Graph,
390    tp: &onnx::TypeProto,
391) -> Result<TypeProto, LoaderError> {
392    let ty = match tp.value.as_ref() {
393        Some(type_proto::Value::TensorType(t)) => {
394            let dtype = decode_dtype(t.elem_type, || "type-proto attribute (tensor)".to_string())?;
395            let shape = t
396                .shape
397                .as_ref()
398                .map(|s| tensor_shape_to_shape(graph, s))
399                .unwrap_or_default();
400            TypeProto::Tensor { dtype, shape }
401        }
402        Some(type_proto::Value::SparseTensorType(t)) => {
403            let dtype =
404                decode_dtype(t.elem_type, || "type-proto attribute (sparse tensor)".to_string())?;
405            let shape = t
406                .shape
407                .as_ref()
408                .map(|s| tensor_shape_to_shape(graph, s))
409                .unwrap_or_default();
410            TypeProto::SparseTensor { dtype, shape }
411        }
412        Some(type_proto::Value::SequenceType(s)) => {
413            let inner = s
414                .elem_type
415                .as_ref()
416                .map(|e| convert_type_proto(graph, e))
417                .transpose()?
418                .unwrap_or(TypeProto::Tensor {
419                    dtype: DataType::Float32,
420                    shape: Vec::new(),
421                });
422            TypeProto::Sequence(Box::new(inner))
423        }
424        Some(type_proto::Value::OptionalType(o)) => {
425            let inner = o
426                .elem_type
427                .as_ref()
428                .map(|e| convert_type_proto(graph, e))
429                .transpose()?
430                .unwrap_or(TypeProto::Tensor {
431                    dtype: DataType::Float32,
432                    shape: Vec::new(),
433                });
434            TypeProto::Optional(Box::new(inner))
435        }
436        Some(type_proto::Value::MapType(m)) => {
437            let key = decode_dtype(m.key_type, || "type-proto attribute (map key)".to_string())?;
438            let value = m
439                .value_type
440                .as_ref()
441                .map(|e| convert_type_proto(graph, e))
442                .transpose()?
443                .unwrap_or(TypeProto::Tensor {
444                    dtype: DataType::Float32,
445                    shape: Vec::new(),
446                });
447            TypeProto::Map {
448                key,
449                value: Box::new(value),
450            }
451        }
452        None => TypeProto::Tensor {
453            dtype: DataType::Float32,
454            shape: Vec::new(),
455        },
456    };
457    Ok(ty)
458}
onnx_runtime_loader/graph_builder.rs

onnx_runtime_loader/
graph_builder.rs
