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