use std::borrow::Cow;
use std::collections::{HashMap, HashSet};
use onnx_runtime_ir::is_default_domain;
use crate::LoaderError;
use crate::proto::onnx::{
AttributeProto, FunctionProto, GraphProto, ModelProto, NodeProto, OperatorSetIdProto,
};
type FnKey = (String, String, String);
fn fn_key_of_function(f: &FunctionProto) -> FnKey {
(f.domain.clone(), f.name.clone(), f.overload.clone())
}
fn fn_key_of_call(n: &NodeProto) -> FnKey {
(n.domain.clone(), n.op_type.clone(), n.overload.clone())
}
pub fn inline_functions(model: &ModelProto) -> Result<Cow<'_, ModelProto>, LoaderError> {
if model.functions.is_empty() {
return Ok(Cow::Borrowed(model));
}
let mut funcs: HashMap<FnKey, &FunctionProto> = HashMap::new();
for f in &model.functions {
funcs.insert(fn_key_of_function(f), f);
}
let graph = model
.graph
.as_ref()
.ok_or_else(|| LoaderError::GraphBuild("ModelProto has no graph".into()))?;
let mut counter: usize = 0;
let mut stack: Vec<FnKey> = Vec::new();
let mut used: HashSet<String> = HashSet::new();
collect_used_names(graph, &mut used);
let mut synthesized_default = false;
let new_graph = inline_graph(
graph,
&funcs,
&mut counter,
&mut stack,
&mut used,
&mut synthesized_default,
)?;
let mut out = model.clone();
out.graph = Some(new_graph);
out.opset_import = merged_opset_imports(model);
if synthesized_default {
ensure_default_opset_import(&mut out.opset_import);
}
out.functions.clear();
Ok(Cow::Owned(out))
}
const DEFAULT_ONNX_OPSET_VERSION: i64 = 17;
const DEFAULT_DOMAIN_KEY: &str = "";
fn domain_key(domain: &str) -> String {
if is_default_domain(domain) {
DEFAULT_DOMAIN_KEY.to_string()
} else {
domain.to_string()
}
}
fn ensure_default_opset_import(imports: &mut Vec<OperatorSetIdProto>) {
let has_default = imports.iter().any(|o| is_default_domain(&o.domain));
if !has_default {
imports.push(OperatorSetIdProto {
domain: String::new(),
version: DEFAULT_ONNX_OPSET_VERSION,
});
}
}
fn merged_opset_imports(model: &ModelProto) -> Vec<OperatorSetIdProto> {
let mut order: Vec<String> = Vec::new();
let mut best: HashMap<String, i64> = HashMap::new();
let mut default_spelling: Option<String> = None;
let mut note = |domain: &str, version: i64, from_model: bool| {
if from_model && is_default_domain(domain) && default_spelling.is_none() {
default_spelling = Some(domain.to_string());
}
let key = domain_key(domain);
match best.entry(key.clone()) {
std::collections::hash_map::Entry::Occupied(mut e) => {
if version > *e.get() {
*e.get_mut() = version;
}
}
std::collections::hash_map::Entry::Vacant(e) => {
order.push(key);
e.insert(version);
}
}
};
for o in &model.opset_import {
note(&o.domain, o.version, true);
}
for f in &model.functions {
for o in &f.opset_import {
note(&o.domain, o.version, false);
}
}
order
.into_iter()
.map(|key| {
let version = best[&key];
let domain = if key == DEFAULT_DOMAIN_KEY {
default_spelling.clone().unwrap_or_default()
} else {
key
};
OperatorSetIdProto { domain, version }
})
.collect()
}
fn inline_graph(
gp: &GraphProto,
funcs: &HashMap<FnKey, &FunctionProto>,
counter: &mut usize,
stack: &mut Vec<FnKey>,
used: &mut HashSet<String>,
synthesized_default: &mut bool,
) -> Result<GraphProto, LoaderError> {
let mut out = gp.clone();
out.node = Vec::with_capacity(gp.node.len());
for node in &gp.node {
expand_node(
node,
funcs,
counter,
stack,
used,
synthesized_default,
&mut out.node,
)?;
}
Ok(out)
}
fn expand_node(
node: &NodeProto,
funcs: &HashMap<FnKey, &FunctionProto>,
counter: &mut usize,
stack: &mut Vec<FnKey>,
used: &mut HashSet<String>,
synthesized_default: &mut bool,
sink: &mut Vec<NodeProto>,
) -> Result<(), LoaderError> {
if let Some(func) = funcs.get(&fn_key_of_call(node)) {
instantiate(
node,
func,
funcs,
counter,
stack,
used,
synthesized_default,
sink,
)?;
} else {
sink.push(inline_subgraph_attrs(
node,
funcs,
counter,
stack,
used,
synthesized_default,
)?);
}
Ok(())
}
fn inline_subgraph_attrs(
node: &NodeProto,
funcs: &HashMap<FnKey, &FunctionProto>,
counter: &mut usize,
stack: &mut Vec<FnKey>,
used: &mut HashSet<String>,
synthesized_default: &mut bool,
) -> Result<NodeProto, LoaderError> {
let mut out = node.clone();
for attr in &mut out.attribute {
if let Some(g) = attr.g.as_mut() {
*g = inline_graph(g, funcs, counter, stack, used, synthesized_default)?;
}
for g in &mut attr.graphs {
*g = inline_graph(g, funcs, counter, stack, used, synthesized_default)?;
}
}
Ok(out)
}
#[allow(clippy::too_many_arguments)]
fn instantiate(
call: &NodeProto,
func: &FunctionProto,
funcs: &HashMap<FnKey, &FunctionProto>,
counter: &mut usize,
stack: &mut Vec<FnKey>,
used: &mut HashSet<String>,
synthesized_default: &mut bool,
sink: &mut Vec<NodeProto>,
) -> Result<(), LoaderError> {
let key = fn_key_of_function(func);
if stack.contains(&key) {
let mut chain: Vec<String> = stack.iter().map(fmt_key).collect();
chain.push(fmt_key(&key));
return Err(LoaderError::RecursiveFunction {
function: fmt_key(&key),
chain: chain.join(" -> "),
});
}
if call.input.len() > func.input.len() {
return Err(LoaderError::FunctionArityMismatch {
function: fmt_key(&key),
node: node_label(call),
kind: "input",
formal: func.input.len(),
actual: call.input.len(),
});
}
if call.output.len() > func.output.len() {
return Err(LoaderError::FunctionArityMismatch {
function: fmt_key(&key),
node: node_label(call),
kind: "output",
formal: func.output.len(),
actual: call.output.len(),
});
}
let inst_id = *counter;
*counter += 1;
let produced: HashSet<&str> = func
.node
.iter()
.flat_map(|n| n.output.iter())
.filter(|o| !o.is_empty())
.map(String::as_str)
.collect();
let mut rename: HashMap<String, String> = HashMap::new();
let mut aliases: Vec<(String, String)> = Vec::new();
for (i, formal) in func.input.iter().enumerate() {
if formal.is_empty() {
continue;
}
let actual = call.input.get(i).cloned().unwrap_or_default();
rename.insert(formal.clone(), actual);
}
for (j, formal) in func.output.iter().enumerate() {
if formal.is_empty() {
continue;
}
let actual = call.output.get(j).cloned().unwrap_or_default();
if produced.contains(formal.as_str()) {
rename.insert(formal.clone(), actual);
} else if let Some(src) = rename.get(formal) {
if !actual.is_empty() && src != &actual {
aliases.push((src.clone(), actual));
}
} else {
rename.insert(formal.clone(), actual);
}
}
for bn in &func.node {
for name in bn.input.iter().chain(bn.output.iter()) {
if name.is_empty() || rename.contains_key(name) {
continue;
}
let fresh = fresh_name(name, inst_id, used);
rename.insert(name.clone(), fresh);
}
}
stack.push(key.clone());
let result = (|| {
let mut instantiated: Vec<NodeProto> = Vec::with_capacity(func.node.len());
for (idx, bn) in func.node.iter().enumerate() {
let mut nn = bn.clone();
nn.name = if bn.name.is_empty() {
format!("__fn{inst_id}_n{idx}")
} else {
format!("__fn{inst_id}_{}", bn.name)
};
bind_node_attributes(&mut nn, call, func, &key)?;
rename_value_refs(&mut nn, &rename);
instantiated.push(nn);
}
for (k, (src, dst)) in aliases.iter().enumerate() {
*synthesized_default = true;
instantiated.push(NodeProto {
op_type: "Identity".to_string(),
input: vec![src.clone()],
output: vec![dst.clone()],
name: format!("__fn{inst_id}_alias{k}"),
..Default::default()
});
}
let mut expanded: Vec<NodeProto> = Vec::new();
for n in &instantiated {
expand_node(
n,
funcs,
counter,
stack,
used,
synthesized_default,
&mut expanded,
)?;
}
Ok::<Vec<NodeProto>, LoaderError>(expanded)
})();
stack.pop();
sink.extend(result?);
Ok(())
}
fn bind_node_attributes(
node: &mut NodeProto,
call: &NodeProto,
func: &FunctionProto,
key: &FnKey,
) -> Result<(), LoaderError> {
let mut bound: Vec<AttributeProto> = Vec::with_capacity(node.attribute.len());
for attr in &node.attribute {
if let Some(mut resolved) = bind_attribute(attr, call, func, key)? {
if let Some(g) = resolved.g.as_mut() {
for sub in &mut g.node {
bind_node_attributes(sub, call, func, key)?;
}
}
for g in &mut resolved.graphs {
for sub in &mut g.node {
bind_node_attributes(sub, call, func, key)?;
}
}
bound.push(resolved);
}
}
node.attribute = bound;
Ok(())
}
fn bind_attribute(
attr: &AttributeProto,
call: &NodeProto,
func: &FunctionProto,
key: &FnKey,
) -> Result<Option<AttributeProto>, LoaderError> {
if attr.ref_attr_name.is_empty() {
return Ok(Some(attr.clone()));
}
let a = &attr.ref_attr_name;
if let Some(supplied) = call.attribute.iter().find(|ca| &ca.name == a) {
let mut bound = supplied.clone();
bound.name = attr.name.clone();
bound.ref_attr_name.clear();
return Ok(Some(bound));
}
if let Some(default) = func.attribute_proto.iter().find(|d| &d.name == a) {
let mut bound = default.clone();
bound.name = attr.name.clone();
bound.ref_attr_name.clear();
return Ok(Some(bound));
}
if func.attribute.iter().any(|req| req == a) {
return Err(LoaderError::MissingRequiredFunctionAttribute {
function: fmt_key(key),
node: node_label(call),
attribute: a.clone(),
});
}
Ok(None)
}
fn rename_value_refs(node: &mut NodeProto, rename: &HashMap<String, String>) {
for name in node.input.iter_mut().chain(node.output.iter_mut()) {
if let Some(new) = rename.get(name.as_str()) {
*name = new.clone();
}
}
for attr in &mut node.attribute {
if let Some(g) = attr.g.as_mut() {
rename_subgraph_refs(g, rename);
}
for g in &mut attr.graphs {
rename_subgraph_refs(g, rename);
}
}
}
fn rename_subgraph_refs(gp: &mut GraphProto, rename: &HashMap<String, String>) {
let mut locals: HashSet<&str> = HashSet::new();
for i in &gp.input {
if !i.name.is_empty() {
locals.insert(i.name.as_str());
}
}
for init in &gp.initializer {
if !init.name.is_empty() {
locals.insert(init.name.as_str());
}
}
for sparse in &gp.sparse_initializer {
if let Some(values) = &sparse.values
&& !values.name.is_empty()
{
locals.insert(values.name.as_str());
}
}
for n in &gp.node {
for o in &n.output {
if !o.is_empty() {
locals.insert(o.as_str());
}
}
}
let effective: HashMap<String, String> = rename
.iter()
.filter(|(k, _)| !locals.contains(k.as_str()))
.map(|(k, v)| (k.clone(), v.clone()))
.collect();
for n in &mut gp.node {
for name in n.input.iter_mut().chain(n.output.iter_mut()) {
if let Some(new) = effective.get(name.as_str()) {
*name = new.clone();
}
}
for attr in &mut n.attribute {
if let Some(g) = attr.g.as_mut() {
rename_subgraph_refs(g, &effective);
}
for g in &mut attr.graphs {
rename_subgraph_refs(g, &effective);
}
}
}
for o in &mut gp.output {
if let Some(new) = effective.get(o.name.as_str()) {
o.name = new.clone();
}
}
}
fn collect_used_names(gp: &GraphProto, used: &mut HashSet<String>) {
for i in &gp.input {
if !i.name.is_empty() {
used.insert(i.name.clone());
}
}
for o in &gp.output {
if !o.name.is_empty() {
used.insert(o.name.clone());
}
}
for init in &gp.initializer {
if !init.name.is_empty() {
used.insert(init.name.clone());
}
}
for sparse in &gp.sparse_initializer {
if let Some(values) = &sparse.values
&& !values.name.is_empty()
{
used.insert(values.name.clone());
}
}
for vi in &gp.value_info {
if !vi.name.is_empty() {
used.insert(vi.name.clone());
}
}
for n in &gp.node {
for name in n.input.iter().chain(n.output.iter()) {
if !name.is_empty() {
used.insert(name.clone());
}
}
for attr in &n.attribute {
if let Some(g) = &attr.g {
collect_used_names(g, used);
}
for g in &attr.graphs {
collect_used_names(g, used);
}
}
}
}
fn fresh_name(base: &str, inst_id: usize, used: &mut HashSet<String>) -> String {
let mut candidate = format!("__fn{inst_id}_{base}");
let mut suffix = 0usize;
while used.contains(&candidate) {
suffix += 1;
candidate = format!("__fn{inst_id}_{base}__{suffix}");
}
used.insert(candidate.clone());
candidate
}
fn fmt_key(key: &FnKey) -> String {
let (domain, name, overload) = key;
let d = if domain.is_empty() { "ai.onnx" } else { domain };
if overload.is_empty() {
format!("{d}::{name}")
} else {
format!("{d}::{name}:{overload}")
}
}
fn node_label(n: &NodeProto) -> String {
if n.name.is_empty() {
format!("<{}::{} (unnamed)>", n.domain, n.op_type)
} else {
n.name.clone()
}
}