use std::path::Path;
use std::sync::Arc;
use onnx_runtime_ir::Graph;
use onnx_runtime_shape_inference::{InferenceRegistry, MergePolicy};
use crate::graph_builder::BuiltGraph;
pub mod encoder;
pub mod epcontext;
pub mod function_inline;
pub(crate) mod graph_builder;
pub mod proto;
pub mod weights;
pub mod writer;
mod pathsafe;
pub use encoder::{
DEFAULT_IR_VERSION, Model, ModelMetadata, encode_model, encode_model_proto, write_model,
};
pub use epcontext::{
EmbedMode, EpContextBlob, EpContextNode, ep_context_node_ids, ep_context_nodes,
is_ep_context_op, resolve_ep_context,
};
pub use error::LoaderError;
pub use weights::WeightStore;
pub use writer::{EpContextDumpConfig, EpContextPartition, dump_ep_context};
mod error {
use std::path::PathBuf;
#[derive(Debug, thiserror::Error)]
pub enum LoaderError {
#[error("failed to read model file {path}: {source}")]
Io {
path: PathBuf,
#[source]
source: std::io::Error,
},
#[error("failed to parse ONNX protobuf: {0}")]
ProtobufParse(String),
#[error("unsupported opset: domain={domain}, version={version}")]
UnsupportedOpset { domain: String, version: u64 },
#[error(
"illegal ONNX model: operator {domain}::{op_type} at node {node} uses domain \
'{domain}' but no corresponding opset_import is declared. RULES #1: the model must \
declare an opset_import for domain '{domain}'; if you built this graph \
programmatically, add it before loading; if this is a file, the model is \
malformed/invalid per the ONNX spec"
)]
MissingOpsetImport {
op_type: String,
node: String,
domain: String,
},
#[error(
"unsupported ONNX model: operator {domain}::{op_type} at node {node} carries a \
subgraph attribute '{attr}' (control-flow / nested-graph op) that this runtime cannot \
execute. RULES #1: ep-cpu recursively executes the standard control-flow ops \
If/Loop/Scan (ai.onnx), but not {op_type}, so the model cannot be run as-is. \
Expected: express control flow with If/Loop/Scan, lower/unroll {op_type} into \
supported ops, or register a kernel able to execute its subgraph body"
)]
UnsupportedControlFlow {
op_type: String,
node: String,
domain: String,
attr: String,
},
#[error(
"illegal ONNX model: operator {domain}::{op_type} at node {node} consumes tensor \
'{tensor}', but no producer exists — it is not a graph input, not an initializer, and \
not produced by any upstream node. RULES #1: every consumed tensor must be sourced; \
the graph is structurally malformed. Expected: add '{tensor}' as a graph input or \
initializer, or add a node that produces it; if this is a file, the model is invalid \
per the ONNX spec"
)]
DanglingTensorRef {
op_type: String,
node: String,
domain: String,
tensor: String,
},
#[error(
"illegal ONNX model: tensor '{tensor}' is declared as an initializer but is also \
produced as an output of node {node} — an initializer must be a constant source with \
no producer. RULES #1: initializer names must be unique and must not collide with any \
node output name; a producer-backed initializer would let a kernel write through \
read-only weight storage. Expected: rename the node output or the initializer so they \
no longer share a name; if this is a file, the model is malformed per the ONNX spec"
)]
InitializerHasProducer { tensor: String, node: String },
#[error(
"illegal ONNX model: value '{tensor}' has multiple producers ({first} and {second}). \
RULES #1: ONNX graphs are in SSA form, so a value name may be assigned only once. \
Expected: give each graph input and node output a unique name"
)]
DuplicateValueProducer {
tensor: String,
first: String,
second: String,
},
#[error(
"illegal ONNX model: operator {domain}::{op_type} at node {node} has attribute \
'{attr}' referring to function attribute '{ref_attr_name}' outside a FunctionProto. \
RULES #1: ref_attr_name is only bound while inlining a FunctionProto; it has no \
executable value in a main graph or control-flow subgraph. Expected: replace it with \
a concrete attribute value or move the node into a FunctionProto"
)]
RefAttributeOutsideFunction {
op_type: String,
node: String,
domain: String,
attr: String,
ref_attr_name: String,
},
#[error(
"illegal ONNX model: ir_version {ir_version} is invalid. RULES #1: ir_version is \
required and ONNX IR versions start at 1. Expected: emit a model with ir_version >= 1"
)]
InvalidIrVersion { ir_version: i64 },
#[error(
"illegal ONNX model: ir_version {ir_version} requires at least one opset_import \
(ONNX IR>=3). Expected: add an opset_import for every operator domain used by the \
model"
)]
MissingModelOpsetImport { ir_version: i64 },
#[error(
"illegal ONNX model: initializer '{tensor}' in an outer graph is shadowed by a \
subgraph input of the same name. RULES #1: this runtime does not permit ambiguous \
initializer/subgraph binding. Expected: rename the subgraph formal input or the \
outer initializer"
)]
SubgraphInputShadowsInitializer { tensor: String },
#[error(
"illegal ONNX model: graph output '{tensor}' has no producer in its graph. RULES #1: \
every output must be a graph input, initializer, or node output in the same scope. \
Expected: produce '{tensor}' locally or declare it as an input/initializer"
)]
GraphOutputMissingProducer { tensor: String },
#[error("external data file not found: {path}")]
ExternalDataNotFound { path: PathBuf },
#[error("external data path rejected ({reason}): {path}")]
ExternalDataPath { path: String, reason: &'static str },
#[error("weight mmap failed: {0}")]
Mmap(String),
#[error("EPContext node error: {0}")]
EpContext(String),
#[error("EPContext external path rejected ({reason}): {path}")]
EpContextPath { path: String, reason: &'static str },
#[error("graph construction failed: {0}")]
GraphBuild(String),
#[error(
"illegal ONNX model: model-local function {function} is recursive (call chain: \
{chain}). RULES #1: ONNX function bodies may reference other model-local functions \
but MUST NOT be recursive — inlining cannot terminate. Expected: break the cycle so \
no function transitively calls itself"
)]
RecursiveFunction { function: String, chain: String },
#[error(
"illegal ONNX model: call to model-local function {function} at node {node} passes \
{actual} {kind}(s) but the function declares only {formal}. RULES #1: a function \
call may omit trailing optional {kind}s but must not supply more than are declared. \
Expected: remove the extra {kind}(s) or fix the function signature"
)]
FunctionArityMismatch {
function: String,
node: String,
kind: &'static str,
formal: usize,
actual: usize,
},
#[error(
"illegal ONNX model: call to model-local function {function} at node {node} is missing \
required attribute '{attribute}', and the function declares no default for it. \
RULES #1: an attribute listed in FunctionProto.attribute has no default and must be \
supplied at every call site. Expected: set '{attribute}' on the call node, or give \
the function a default via attribute_proto"
)]
MissingRequiredFunctionAttribute {
function: String,
node: String,
attribute: String,
},
#[error("unsupported ONNX data_type {raw} at {context}")]
UnsupportedDataType { raw: i32, context: String },
#[error("shape inference failed: {0}")]
ShapeInference(#[from] onnx_runtime_shape_inference::ShapeInferError),
#[error(transparent)]
Ir(#[from] onnx_runtime_ir::IrError),
}
}
pub fn load_model(path: impl AsRef<Path>) -> Result<Graph, LoaderError> {
Ok(load_model_with_weights(path)?.0)
}
pub fn load_model_bytes(bytes: &[u8]) -> Result<Graph, LoaderError> {
Ok(load_model_bytes_with_weights(bytes, Path::new("."))?.0)
}
pub fn load_model_with_weights(
path: impl AsRef<Path>,
) -> Result<(Graph, Arc<WeightStore>), LoaderError> {
let path = path.as_ref();
let bytes = std::fs::read(path).map_err(|source| LoaderError::Io {
path: path.to_path_buf(),
source,
})?;
let model_dir = path.parent().unwrap_or_else(|| Path::new("."));
build_from_bytes_with_weights(&bytes, model_dir)
}
pub fn load_model_bytes_with_weights(
bytes: &[u8],
base_dir: impl AsRef<Path>,
) -> Result<(Graph, Arc<WeightStore>), LoaderError> {
build_from_bytes_with_weights(bytes, base_dir.as_ref())
}
fn build_from_bytes_with_weights(
bytes: &[u8],
model_dir: &Path,
) -> Result<(Graph, Arc<WeightStore>), LoaderError> {
let model = proto::decode_model(bytes)?;
validate_model_proto(&model)?;
let BuiltGraph {
mut graph,
name_map,
} = graph_builder::build_graph(&model)?;
validate_opset_imports(&graph)?;
let store = weights::load_weights(&model, model_dir, &name_map)?;
for (&vid, weight) in &store.weights {
graph.set_initializer(vid, weight.clone());
}
graph
.validate()
.map_err(|errs| LoaderError::GraphBuild(format!("{errs:?}")))?;
validate_model(&graph)?;
let registry = InferenceRegistry::default_registry();
let opset_imports = graph.opset_imports.clone();
registry.infer_graph(&mut graph, &opset_imports, MergePolicy::Permissive)?;
Ok((graph, Arc::new(store)))
}
pub fn validate_model(graph: &Graph) -> Result<(), LoaderError> {
validate_opset_imports(graph)?;
validate_no_control_flow(graph)?;
validate_no_dangling_refs(graph)?;
validate_no_initializer_producer(graph)?;
Ok(())
}
pub fn validate_model_proto(model: &proto::onnx::ModelProto) -> Result<(), LoaderError> {
use std::collections::HashSet;
use proto::onnx::GraphProto;
if model.ir_version < 1 {
return Err(LoaderError::InvalidIrVersion {
ir_version: model.ir_version,
});
}
if model.ir_version >= 3 && model.opset_import.is_empty() {
return Err(LoaderError::MissingModelOpsetImport {
ir_version: model.ir_version,
});
}
fn node_description(node: &proto::onnx::NodeProto, index: usize) -> String {
if node.name.is_empty() {
format!("<unnamed node #{index}>")
} else {
format!("{:?}", node.name)
}
}
fn check_graph(graph: &GraphProto) -> Result<(), LoaderError> {
let mut producers = std::collections::HashMap::new();
for input in &graph.input {
if !input.name.is_empty() {
producers.insert(input.name.clone(), "graph input".to_string());
}
}
for (index, node) in graph.node.iter().enumerate() {
let node_description = node_description(node, index);
for output in &node.output {
if output.is_empty() {
continue;
}
let producer = format!("output of {node_description}");
if let Some(first) = producers.insert(output.clone(), producer.clone()) {
return Err(LoaderError::DuplicateValueProducer {
tensor: output.clone(),
first,
second: producer,
});
}
}
for attribute in &node.attribute {
if !attribute.ref_attr_name.is_empty() {
return Err(LoaderError::RefAttributeOutsideFunction {
op_type: node.op_type.clone(),
node: node_description.clone(),
domain: display_domain(&node.domain),
attr: attribute.name.clone(),
ref_attr_name: attribute.ref_attr_name.clone(),
});
}
}
}
let sources: HashSet<&str> = graph
.input
.iter()
.map(|input| input.name.as_str())
.chain(
graph
.initializer
.iter()
.map(|initializer| initializer.name.as_str()),
)
.chain(
graph
.node
.iter()
.flat_map(|node| node.output.iter().map(String::as_str)),
)
.collect();
for output in &graph.output {
if !output.name.is_empty() && !sources.contains(output.name.as_str()) {
return Err(LoaderError::GraphOutputMissingProducer {
tensor: output.name.clone(),
});
}
}
let outer_initializers: HashSet<&str> = graph
.initializer
.iter()
.map(|initializer| initializer.name.as_str())
.collect();
for node in &graph.node {
for attribute in &node.attribute {
let subgraphs = attribute.g.iter().chain(attribute.graphs.iter());
for subgraph in subgraphs {
if let Some(input) = subgraph
.input
.iter()
.find(|input| outer_initializers.contains(input.name.as_str()))
{
return Err(LoaderError::SubgraphInputShadowsInitializer {
tensor: input.name.clone(),
});
}
check_graph(subgraph)?;
}
}
}
Ok(())
}
if let Some(graph) = &model.graph {
check_graph(graph)?;
}
Ok(())
}
fn node_label(node: &onnx_runtime_ir::Node) -> String {
if node.name.is_empty() {
format!("<unnamed node #{}>", node.id.0)
} else {
format!("{:?}", node.name)
}
}
fn display_domain(domain: &str) -> String {
if domain.is_empty() {
"ai.onnx".to_string()
} else {
domain.to_string()
}
}
pub fn validate_no_control_flow(graph: &Graph) -> Result<(), LoaderError> {
use onnx_runtime_ir::Attribute;
fn is_default_domain(domain: &str) -> bool {
domain.is_empty() || domain == "ai.onnx"
}
fn is_implemented_control_flow(op_type: &str, domain: &str) -> bool {
is_default_domain(domain) && matches!(op_type, "If" | "Loop" | "Scan")
}
fn check_graph(graph: &Graph) -> Result<(), LoaderError> {
for (_, node) in graph.nodes.iter() {
let mut subgraph_attrs: Vec<&String> = node
.attributes
.iter()
.filter(|(_, v)| matches!(v, Attribute::Graph(_) | Attribute::Graphs(_)))
.map(|(k, _)| k)
.collect();
subgraph_attrs.sort();
if let Some(attr) = subgraph_attrs.first() {
if !is_implemented_control_flow(&node.op_type, &node.domain) {
return Err(LoaderError::UnsupportedControlFlow {
op_type: node.op_type.clone(),
node: node_label(node),
domain: display_domain(&node.domain),
attr: (*attr).clone(),
});
}
}
}
for subgraph in graph.subgraphs.values() {
check_graph(subgraph)?;
}
Ok(())
}
check_graph(graph)
}
pub fn validate_no_dangling_refs(graph: &Graph) -> Result<(), LoaderError> {
use std::collections::HashSet;
let graph_inputs: HashSet<_> = graph.inputs.iter().copied().collect();
for (_, node) in graph.nodes.iter() {
for vid in node.input_values() {
let Some(value) = graph.values.get(vid) else {
continue;
};
let is_sourced = value.producer.is_some()
|| graph_inputs.contains(&vid)
|| graph.initializers.contains_key(&vid);
if !is_sourced {
let tensor = value
.name
.clone()
.unwrap_or_else(|| format!("<anonymous value #{}>", vid.0));
return Err(LoaderError::DanglingTensorRef {
op_type: node.op_type.clone(),
node: node_label(node),
domain: display_domain(&node.domain),
tensor,
});
}
}
}
Ok(())
}
pub fn validate_no_initializer_producer(graph: &Graph) -> Result<(), LoaderError> {
for &vid in graph.initializers.keys() {
let Some(value) = graph.values.get(vid) else {
continue;
};
if let Some(producer) = value.producer {
let tensor = value
.name
.clone()
.unwrap_or_else(|| format!("<anonymous value #{}>", vid.0));
let node = if graph.nodes.contains(producer) {
node_label(graph.node(producer))
} else {
format!("<node #{}>", producer.0)
};
return Err(LoaderError::InitializerHasProducer { tensor, node });
}
}
Ok(())
}
pub fn validate_opset_imports(graph: &Graph) -> Result<(), LoaderError> {
fn has_import(imports: &std::collections::HashMap<String, u64>, domain: &str) -> bool {
imports.contains_key(domain)
|| (domain.is_empty() && imports.contains_key("ai.onnx"))
|| (domain == "ai.onnx" && imports.contains_key(""))
}
fn validate_graph(
graph: &Graph,
imports: &std::collections::HashMap<String, u64>,
) -> Result<(), LoaderError> {
for (_, node) in graph.nodes.iter() {
if !has_import(imports, &node.domain) {
let domain = if node.domain.is_empty() {
"ai.onnx".to_string()
} else {
node.domain.clone()
};
let node_name = if node.name.is_empty() {
format!("<unnamed node #{}>", node.id.0)
} else {
format!("{:?}", node.name)
};
return Err(LoaderError::MissingOpsetImport {
op_type: node.op_type.clone(),
node: node_name,
domain,
});
}
}
for subgraph in graph.subgraphs.values() {
validate_graph(subgraph, imports)?;
}
Ok(())
}
validate_graph(graph, &graph.opset_imports)
}