bb_ir/proto/onnx_view.rs
1//! Helpers that adapt ONNX proto types between equivalent shapes.
2
3use std::collections::HashMap;
4
5use super::onnx::{FunctionProto, GraphProto, ValueInfoProto};
6
7/// Build a `GraphProto` view from a `FunctionProto`'s body - most
8/// compiler passes operate on a `GraphProto` per `docs/COMPILER.md`,
9/// and the engine's backend-subgraph dispatcher needs the same view
10///
11///
12/// The view resolves `function.input` / `output` (which are
13/// `repeated string` at the FunctionProto level) into `ValueInfoProto`s
14/// by looking up the matching entry in `function.value_info`. Missing
15/// entries surface as `ValueInfoProto { name, ..default() }`; the
16/// compiler's `validate` pass then catches them as `MissingTypeInfo`.
17pub fn function_to_graph_view(function: &FunctionProto) -> GraphProto {
18 let lookup: HashMap<&str, &ValueInfoProto> = function
19 .value_info
20 .iter()
21 .map(|v| (v.name.as_str(), v))
22 .collect();
23 let resolve = |name: &str| -> ValueInfoProto {
24 lookup
25 .get(name)
26 .map(|v| (*v).clone())
27 .unwrap_or(ValueInfoProto {
28 name: name.to_string(),
29 ..Default::default()
30 })
31 };
32 GraphProto {
33 node: function.node.clone(),
34 name: function.name.clone(),
35 input: function.input.iter().map(|n| resolve(n)).collect(),
36 output: function.output.iter().map(|n| resolve(n)).collect(),
37 value_info: function.value_info.clone(),
38 ..Default::default()
39 }
40}