Skip to main content

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}