1use std::path::Path;
41use std::sync::Arc;
42
43use onnx_runtime_ir::Graph;
44use onnx_runtime_shape_inference::{InferenceRegistry, MergePolicy};
45
46use crate::graph_builder::BuiltGraph;
47
48pub mod encoder;
49pub mod epcontext;
50pub mod function_inline;
51pub(crate) mod graph_builder;
52pub mod proto;
53pub mod weights;
54pub mod writer;
55
56mod pathsafe;
57
58pub use encoder::{
59 DEFAULT_IR_VERSION, Model, ModelMetadata, encode_model, encode_model_proto, write_model,
60};
61pub use epcontext::{
62 EmbedMode, EpContextBlob, EpContextNode, ep_context_node_ids, ep_context_nodes,
63 is_ep_context_op, resolve_ep_context,
64};
65pub use error::LoaderError;
66pub use weights::WeightStore;
67pub use writer::{EpContextDumpConfig, EpContextPartition, dump_ep_context};
68
69mod error {
70 use std::path::PathBuf;
71
72 #[derive(Debug, thiserror::Error)]
74 pub enum LoaderError {
75 #[error("failed to read model file {path}: {source}")]
76 Io {
77 path: PathBuf,
78 #[source]
79 source: std::io::Error,
80 },
81
82 #[error("failed to parse ONNX protobuf: {0}")]
83 ProtobufParse(String),
84
85 #[error("unsupported opset: domain={domain}, version={version}")]
86 UnsupportedOpset { domain: String, version: u64 },
87
88 #[error(
89 "illegal ONNX model: operator {domain}::{op_type} at node {node} uses domain \
90 '{domain}' but no corresponding opset_import is declared. RULES #1: the model must \
91 declare an opset_import for domain '{domain}'; if you built this graph \
92 programmatically, add it before loading; if this is a file, the model is \
93 malformed/invalid per the ONNX spec"
94 )]
95 MissingOpsetImport {
96 op_type: String,
97 node: String,
98 domain: String,
99 },
100
101 #[error(
102 "unsupported ONNX model: operator {domain}::{op_type} at node {node} carries a \
103 subgraph attribute '{attr}' (control-flow / nested-graph op) that this runtime cannot \
104 execute. RULES #1: ep-cpu recursively executes the standard control-flow ops \
105 If/Loop/Scan (ai.onnx), but not {op_type}, so the model cannot be run as-is. \
106 Expected: express control flow with If/Loop/Scan, lower/unroll {op_type} into \
107 supported ops, or register a kernel able to execute its subgraph body"
108 )]
109 UnsupportedControlFlow {
110 op_type: String,
111 node: String,
112 domain: String,
113 attr: String,
114 },
115
116 #[error(
117 "illegal ONNX model: operator {domain}::{op_type} at node {node} consumes tensor \
118 '{tensor}', but no producer exists — it is not a graph input, not an initializer, and \
119 not produced by any upstream node. RULES #1: every consumed tensor must be sourced; \
120 the graph is structurally malformed. Expected: add '{tensor}' as a graph input or \
121 initializer, or add a node that produces it; if this is a file, the model is invalid \
122 per the ONNX spec"
123 )]
124 DanglingTensorRef {
125 op_type: String,
126 node: String,
127 domain: String,
128 tensor: String,
129 },
130
131 #[error(
132 "illegal ONNX model: tensor '{tensor}' is declared as an initializer but is also \
133 produced as an output of node {node} — an initializer must be a constant source with \
134 no producer. RULES #1: initializer names must be unique and must not collide with any \
135 node output name; a producer-backed initializer would let a kernel write through \
136 read-only weight storage. Expected: rename the node output or the initializer so they \
137 no longer share a name; if this is a file, the model is malformed per the ONNX spec"
138 )]
139 InitializerHasProducer { tensor: String, node: String },
140
141 #[error(
142 "illegal ONNX model: value '{tensor}' has multiple producers ({first} and {second}). \
143 RULES #1: ONNX graphs are in SSA form, so a value name may be assigned only once. \
144 Expected: give each graph input and node output a unique name"
145 )]
146 DuplicateValueProducer {
147 tensor: String,
148 first: String,
149 second: String,
150 },
151
152 #[error(
153 "illegal ONNX model: operator {domain}::{op_type} at node {node} has attribute \
154 '{attr}' referring to function attribute '{ref_attr_name}' outside a FunctionProto. \
155 RULES #1: ref_attr_name is only bound while inlining a FunctionProto; it has no \
156 executable value in a main graph or control-flow subgraph. Expected: replace it with \
157 a concrete attribute value or move the node into a FunctionProto"
158 )]
159 RefAttributeOutsideFunction {
160 op_type: String,
161 node: String,
162 domain: String,
163 attr: String,
164 ref_attr_name: String,
165 },
166
167 #[error(
168 "illegal ONNX model: ir_version {ir_version} is invalid. RULES #1: ir_version is \
169 required and ONNX IR versions start at 1. Expected: emit a model with ir_version >= 1"
170 )]
171 InvalidIrVersion { ir_version: i64 },
172
173 #[error(
174 "illegal ONNX model: ir_version {ir_version} requires at least one opset_import \
175 (ONNX IR>=3). Expected: add an opset_import for every operator domain used by the \
176 model"
177 )]
178 MissingModelOpsetImport { ir_version: i64 },
179
180 #[error(
181 "illegal ONNX model: initializer '{tensor}' in an outer graph is shadowed by a \
182 subgraph input of the same name. RULES #1: this runtime does not permit ambiguous \
183 initializer/subgraph binding. Expected: rename the subgraph formal input or the \
184 outer initializer"
185 )]
186 SubgraphInputShadowsInitializer { tensor: String },
187
188 #[error(
189 "illegal ONNX model: graph output '{tensor}' has no producer in its graph. RULES #1: \
190 every output must be a graph input, initializer, or node output in the same scope. \
191 Expected: produce '{tensor}' locally or declare it as an input/initializer"
192 )]
193 GraphOutputMissingProducer { tensor: String },
194
195 #[error("external data file not found: {path}")]
196 ExternalDataNotFound { path: PathBuf },
197
198 #[error("external data path rejected ({reason}): {path}")]
199 ExternalDataPath { path: String, reason: &'static str },
200
201 #[error("weight mmap failed: {0}")]
202 Mmap(String),
203
204 #[error("EPContext node error: {0}")]
205 EpContext(String),
206
207 #[error("EPContext external path rejected ({reason}): {path}")]
208 EpContextPath { path: String, reason: &'static str },
209
210 #[error("graph construction failed: {0}")]
211 GraphBuild(String),
212
213 #[error(
214 "illegal ONNX model: model-local function {function} is recursive (call chain: \
215 {chain}). RULES #1: ONNX function bodies may reference other model-local functions \
216 but MUST NOT be recursive — inlining cannot terminate. Expected: break the cycle so \
217 no function transitively calls itself"
218 )]
219 RecursiveFunction { function: String, chain: String },
220
221 #[error(
222 "illegal ONNX model: call to model-local function {function} at node {node} passes \
223 {actual} {kind}(s) but the function declares only {formal}. RULES #1: a function \
224 call may omit trailing optional {kind}s but must not supply more than are declared. \
225 Expected: remove the extra {kind}(s) or fix the function signature"
226 )]
227 FunctionArityMismatch {
228 function: String,
229 node: String,
230 kind: &'static str,
231 formal: usize,
232 actual: usize,
233 },
234
235 #[error(
236 "illegal ONNX model: call to model-local function {function} at node {node} is missing \
237 required attribute '{attribute}', and the function declares no default for it. \
238 RULES #1: an attribute listed in FunctionProto.attribute has no default and must be \
239 supplied at every call site. Expected: set '{attribute}' on the call node, or give \
240 the function a default via attribute_proto"
241 )]
242 MissingRequiredFunctionAttribute {
243 function: String,
244 node: String,
245 attribute: String,
246 },
247
248 #[error("unsupported ONNX data_type {raw} at {context}")]
249 UnsupportedDataType { raw: i32, context: String },
250
251 #[error("shape inference failed: {0}")]
252 ShapeInference(#[from] onnx_runtime_shape_inference::ShapeInferError),
253
254 #[error(transparent)]
255 Ir(#[from] onnx_runtime_ir::IrError),
256 }
257}
258
259pub fn load_model(path: impl AsRef<Path>) -> Result<Graph, LoaderError> {
274 Ok(load_model_with_weights(path)?.0)
275}
276
277pub fn load_model_bytes(bytes: &[u8]) -> Result<Graph, LoaderError> {
288 Ok(load_model_bytes_with_weights(bytes, Path::new("."))?.0)
289}
290
291pub fn load_model_with_weights(
303 path: impl AsRef<Path>,
304) -> Result<(Graph, Arc<WeightStore>), LoaderError> {
305 let path = path.as_ref();
306 let bytes = std::fs::read(path).map_err(|source| LoaderError::Io {
307 path: path.to_path_buf(),
308 source,
309 })?;
310 let model_dir = path.parent().unwrap_or_else(|| Path::new("."));
311 build_from_bytes_with_weights(&bytes, model_dir)
312}
313
314pub fn load_model_bytes_with_weights(
321 bytes: &[u8],
322 base_dir: impl AsRef<Path>,
323) -> Result<(Graph, Arc<WeightStore>), LoaderError> {
324 build_from_bytes_with_weights(bytes, base_dir.as_ref())
325}
326
327fn build_from_bytes_with_weights(
328 bytes: &[u8],
329 model_dir: &Path,
330) -> Result<(Graph, Arc<WeightStore>), LoaderError> {
331 let model = proto::decode_model(bytes)?;
332 validate_model_proto(&model)?;
333 let BuiltGraph {
334 mut graph,
335 name_map,
336 } = graph_builder::build_graph(&model)?;
337
338 validate_opset_imports(&graph)?;
341
342 let store = weights::load_weights(&model, model_dir, &name_map)?;
343 for (&vid, weight) in &store.weights {
345 graph.set_initializer(vid, weight.clone());
346 }
347
348 graph
356 .validate()
357 .map_err(|errs| LoaderError::GraphBuild(format!("{errs:?}")))?;
358
359 validate_model(&graph)?;
364
365 let registry = InferenceRegistry::default_registry();
374 let opset_imports = graph.opset_imports.clone();
375 registry.infer_graph(&mut graph, &opset_imports, MergePolicy::Permissive)?;
376
377 Ok((graph, Arc::new(store)))
378}
379
380pub fn validate_model(graph: &Graph) -> Result<(), LoaderError> {
413 validate_opset_imports(graph)?;
414 validate_no_control_flow(graph)?;
415 validate_no_dangling_refs(graph)?;
416 validate_no_initializer_producer(graph)?;
417 Ok(())
418}
419
420pub fn validate_model_proto(model: &proto::onnx::ModelProto) -> Result<(), LoaderError> {
427 use std::collections::HashSet;
428
429 use proto::onnx::GraphProto;
430
431 if model.ir_version < 1 {
434 return Err(LoaderError::InvalidIrVersion {
435 ir_version: model.ir_version,
436 });
437 }
438 if model.ir_version >= 3 && model.opset_import.is_empty() {
445 return Err(LoaderError::MissingModelOpsetImport {
446 ir_version: model.ir_version,
447 });
448 }
449
450 fn node_description(node: &proto::onnx::NodeProto, index: usize) -> String {
451 if node.name.is_empty() {
452 format!("<unnamed node #{index}>")
453 } else {
454 format!("{:?}", node.name)
455 }
456 }
457
458 fn check_graph(graph: &GraphProto) -> Result<(), LoaderError> {
459 let mut producers = std::collections::HashMap::new();
460 for input in &graph.input {
461 if !input.name.is_empty() {
462 producers.insert(input.name.clone(), "graph input".to_string());
463 }
464 }
465 for (index, node) in graph.node.iter().enumerate() {
466 let node_description = node_description(node, index);
467 for output in &node.output {
468 if output.is_empty() {
469 continue;
470 }
471 let producer = format!("output of {node_description}");
472 if let Some(first) = producers.insert(output.clone(), producer.clone()) {
473 return Err(LoaderError::DuplicateValueProducer {
474 tensor: output.clone(),
475 first,
476 second: producer,
477 });
478 }
479 }
480 for attribute in &node.attribute {
481 if !attribute.ref_attr_name.is_empty() {
482 return Err(LoaderError::RefAttributeOutsideFunction {
483 op_type: node.op_type.clone(),
484 node: node_description.clone(),
485 domain: display_domain(&node.domain),
486 attr: attribute.name.clone(),
487 ref_attr_name: attribute.ref_attr_name.clone(),
488 });
489 }
490 }
491 }
492
493 let sources: HashSet<&str> = graph
494 .input
495 .iter()
496 .map(|input| input.name.as_str())
497 .chain(
498 graph
499 .initializer
500 .iter()
501 .map(|initializer| initializer.name.as_str()),
502 )
503 .chain(
504 graph
505 .node
506 .iter()
507 .flat_map(|node| node.output.iter().map(String::as_str)),
508 )
509 .collect();
510 for output in &graph.output {
511 if !output.name.is_empty() && !sources.contains(output.name.as_str()) {
512 return Err(LoaderError::GraphOutputMissingProducer {
513 tensor: output.name.clone(),
514 });
515 }
516 }
517
518 let outer_initializers: HashSet<&str> = graph
519 .initializer
520 .iter()
521 .map(|initializer| initializer.name.as_str())
522 .collect();
523 for node in &graph.node {
524 for attribute in &node.attribute {
525 let subgraphs = attribute.g.iter().chain(attribute.graphs.iter());
526 for subgraph in subgraphs {
527 if let Some(input) = subgraph
528 .input
529 .iter()
530 .find(|input| outer_initializers.contains(input.name.as_str()))
531 {
532 return Err(LoaderError::SubgraphInputShadowsInitializer {
533 tensor: input.name.clone(),
534 });
535 }
536 check_graph(subgraph)?;
537 }
538 }
539 }
540 Ok(())
541 }
542
543 if let Some(graph) = &model.graph {
544 check_graph(graph)?;
545 }
546 Ok(())
547}
548
549fn node_label(node: &onnx_runtime_ir::Node) -> String {
552 if node.name.is_empty() {
553 format!("<unnamed node #{}>", node.id.0)
554 } else {
555 format!("{:?}", node.name)
556 }
557}
558
559fn display_domain(domain: &str) -> String {
561 if domain.is_empty() {
562 "ai.onnx".to_string()
563 } else {
564 domain.to_string()
565 }
566}
567
568pub fn validate_no_control_flow(graph: &Graph) -> Result<(), LoaderError> {
582 use onnx_runtime_ir::Attribute;
583
584 fn is_default_domain(domain: &str) -> bool {
585 domain.is_empty() || domain == "ai.onnx"
586 }
587
588 fn is_implemented_control_flow(op_type: &str, domain: &str) -> bool {
590 is_default_domain(domain) && matches!(op_type, "If" | "Loop" | "Scan")
591 }
592
593 fn check_graph(graph: &Graph) -> Result<(), LoaderError> {
594 for (_, node) in graph.nodes.iter() {
595 let mut subgraph_attrs: Vec<&String> = node
597 .attributes
598 .iter()
599 .filter(|(_, v)| matches!(v, Attribute::Graph(_) | Attribute::Graphs(_)))
600 .map(|(k, _)| k)
601 .collect();
602 subgraph_attrs.sort();
603 if let Some(attr) = subgraph_attrs.first() {
604 if !is_implemented_control_flow(&node.op_type, &node.domain) {
607 return Err(LoaderError::UnsupportedControlFlow {
608 op_type: node.op_type.clone(),
609 node: node_label(node),
610 domain: display_domain(&node.domain),
611 attr: (*attr).clone(),
612 });
613 }
614 }
615 }
616 for subgraph in graph.subgraphs.values() {
619 check_graph(subgraph)?;
620 }
621 Ok(())
622 }
623
624 check_graph(graph)
625}
626
627pub fn validate_no_dangling_refs(graph: &Graph) -> Result<(), LoaderError> {
640 use std::collections::HashSet;
641
642 let graph_inputs: HashSet<_> = graph.inputs.iter().copied().collect();
643
644 for (_, node) in graph.nodes.iter() {
645 for vid in node.input_values() {
646 let Some(value) = graph.values.get(vid) else {
647 continue;
650 };
651 let is_sourced = value.producer.is_some()
652 || graph_inputs.contains(&vid)
653 || graph.initializers.contains_key(&vid);
654 if !is_sourced {
655 let tensor = value
656 .name
657 .clone()
658 .unwrap_or_else(|| format!("<anonymous value #{}>", vid.0));
659 return Err(LoaderError::DanglingTensorRef {
660 op_type: node.op_type.clone(),
661 node: node_label(node),
662 domain: display_domain(&node.domain),
663 tensor,
664 });
665 }
666 }
667 }
668 Ok(())
669}
670
671pub fn validate_no_initializer_producer(graph: &Graph) -> Result<(), LoaderError> {
689 for &vid in graph.initializers.keys() {
690 let Some(value) = graph.values.get(vid) else {
691 continue;
692 };
693 if let Some(producer) = value.producer {
694 let tensor = value
695 .name
696 .clone()
697 .unwrap_or_else(|| format!("<anonymous value #{}>", vid.0));
698 let node = if graph.nodes.contains(producer) {
699 node_label(graph.node(producer))
700 } else {
701 format!("<node #{}>", producer.0)
702 };
703 return Err(LoaderError::InitializerHasProducer { tensor, node });
704 }
705 }
706 Ok(())
707}
708pub fn validate_opset_imports(graph: &Graph) -> Result<(), LoaderError> {
712 fn has_import(imports: &std::collections::HashMap<String, u64>, domain: &str) -> bool {
713 imports.contains_key(domain)
714 || (domain.is_empty() && imports.contains_key("ai.onnx"))
715 || (domain == "ai.onnx" && imports.contains_key(""))
716 }
717
718 fn validate_graph(
719 graph: &Graph,
720 imports: &std::collections::HashMap<String, u64>,
721 ) -> Result<(), LoaderError> {
722 for (_, node) in graph.nodes.iter() {
723 if !has_import(imports, &node.domain) {
724 let domain = if node.domain.is_empty() {
725 "ai.onnx".to_string()
726 } else {
727 node.domain.clone()
728 };
729 let node_name = if node.name.is_empty() {
730 format!("<unnamed node #{}>", node.id.0)
731 } else {
732 format!("{:?}", node.name)
733 };
734 return Err(LoaderError::MissingOpsetImport {
735 op_type: node.op_type.clone(),
736 node: node_name,
737 domain,
738 });
739 }
740 }
741 for subgraph in graph.subgraphs.values() {
742 validate_graph(subgraph, imports)?;
743 }
744 Ok(())
745 }
746
747 validate_graph(graph, &graph.opset_imports)
748}