Skip to main content

onnx_runtime_loader/
lib.rs

1//! # `onnx-runtime-loader`
2//!
3//! Loads ONNX models from disk into the [`onnx_runtime_ir::Graph`] IR
4//! (see `docs/ORT2.md` §19).
5//!
6//! Pipeline ([`load_model`] / [`load_model_with_weights`]):
7//! 1. [`proto`] — decode the ONNX protobuf (`prost` types generated from the
8//!    vendored `onnx.proto3`) into a `ModelProto`.
9//! 2. [`graph_builder`] — build an [`onnx_runtime_ir::Graph`] (nodes, values,
10//!    symbolic dim interning, opset imports), upholding the §3.5 invariants.
11//! 3. [`weights`] — resolve inline and external initializer data (external
12//!    files are memory-mapped into a [`WeightStore`]).
13//! 4. Static/symbolic shape inference via
14//!    [`onnx-runtime-shape-inference`](onnx_runtime_shape_inference): the loader
15//!    owns the "loader = shape-inference" seam, so after the [`Graph`] is built
16//!    (with initializers applied) it runs the extensible per-op registry to
17//!    populate every value's shape and dtype. Values that cannot be resolved
18//!    statically (genuinely data-dependent extents) are left symbolic for the
19//!    session to resolve just-in-time.
20//!
21//! ## Obtaining weight bytes at session time
22//!
23//! Use [`load_model_with_weights`] (or [`load_model_bytes_with_weights`]) to
24//! receive both the [`Graph`] and an [`Arc<WeightStore>`]. Then, given any
25//! [`onnx_runtime_ir::WeightRef`] stored in `graph.initializers`, call
26//! [`WeightStore::bytes`] to get the raw little-endian byte slice:
27//!
28//! ```ignore
29//! let (graph, store) = load_model_with_weights("model.onnx")?;
30//! for (vid, weight_ref) in &graph.initializers {
31//!     let bytes: &[u8] = store.bytes(weight_ref).expect("weight bytes live");
32//!     // ... hand bytes to a kernel
33//! }
34//! ```
35//!
36//! The `Arc` keeps all memory maps alive as long as any clone of it exists, so
37//! kernel dispatch can store `Arc<WeightStore>` alongside the `Graph` without
38//! lifetime coupling.
39
40use 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    /// Errors produced while loading an ONNX model.
76    #[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
262/// Load a model from a filesystem path, producing a fully-built [`Graph`].
263///
264/// Runs the full pipeline: parse → build → load weights → shape inference.
265/// External initializer data is resolved relative to the model file's
266/// directory.
267///
268/// # Note on external weights
269///
270/// The returned `Graph` stores [`onnx_runtime_ir::WeightRef::External`]
271/// descriptors (path / offset / length) for weights held in external data
272/// files, but the memory maps that back those bytes are **dropped** when this
273/// function returns. Callers that need to read external weight bytes must
274/// either re-map the files themselves or use [`load_model_with_weights`] which
275/// keeps the maps alive via the returned [`Arc<WeightStore>`].
276pub fn load_model(path: impl AsRef<Path>) -> Result<Graph, LoaderError> {
277    Ok(load_model_with_weights(path)?.0)
278}
279
280/// Load a model from an in-memory protobuf buffer, producing a [`Graph`].
281///
282/// External initializer data (if any) is resolved relative to the current
283/// working directory.
284///
285/// # Note on external weights
286///
287/// Same caveat as [`load_model`]: external weight bytes are not accessible
288/// from the returned `Graph` alone. Use [`load_model_bytes_with_weights`] to
289/// keep them live.
290pub fn load_model_bytes(bytes: &[u8]) -> Result<Graph, LoaderError> {
291    Ok(load_model_bytes_with_weights(bytes, Path::new("."))?.0)
292}
293
294/// Load a model from a filesystem path, returning both the [`Graph`] and the
295/// live [`WeightStore`] that backs all initializer data.
296///
297/// The [`Arc<WeightStore>`] keeps every external-data memory map alive for as
298/// long as any clone of the `Arc` exists. At session time, given a
299/// [`onnx_runtime_ir::WeightRef`] from `graph.initializers`, call
300/// [`WeightStore::bytes`] to obtain the raw little-endian byte slice — this
301/// works for both [`WeightRef::Inline`] and [`WeightRef::External`] weights.
302///
303/// External initializer data is resolved relative to the model file's
304/// directory.
305pub 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
317/// Load a model from an in-memory protobuf buffer, returning both the
318/// [`Graph`] and the live [`WeightStore`] that backs all initializer data.
319///
320/// External initializer data (if any) is resolved relative to `base_dir`.
321/// The [`Arc<WeightStore>`] keeps every memory map alive for as long as any
322/// clone of the `Arc` exists.
323pub 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    // Fail-fast legality check that needs no weights: reject illegal opset
342    // imports before we touch the (potentially large) weight files.
343    validate_opset_imports(&graph)?;
344
345    let store = weights::load_weights(&model, model_dir, &name_map)?;
346    // Copy descriptors into the graph; the store's mmaps stay alive via Arc.
347    for (&vid, weight) in &store.weights {
348        graph.set_initializer(vid, weight.clone());
349    }
350
351    // Structural IR validation runs here — *after* initializers are attached —
352    // rather than inside `graph_builder::build_graph`. A top-level initializer
353    // is only recorded in `graph.initializers` by the weight-loading path above,
354    // so validating earlier would mis-flag a legal initializer that is also a
355    // graph output (constant pass-through) or a pre-IR-4 graph input that is
356    // also an initializer as a producer-less `MissingProducer`. Validating the
357    // fully-assembled graph recognizes those values as initializer sources.
358    graph
359        .validate()
360        .map_err(|errs| LoaderError::GraphBuild(format!("{errs:?}")))?;
361
362    // Full fail-fast validation once initializers are attached (so
363    // initializer-backed values are recognized as sourced). Rejects
364    // statically-knowable unsupported/illegal constructs before shape
365    // inference or execution — see [`validate_model`].
366    validate_model(&graph)?;
367
368    // Static/symbolic shape inference (the loader owns this seam). Run the
369    // extensible per-op registry over the fully-built graph — inputs,
370    // initializers, and node outputs — to populate every value's shape and
371    // dtype. `Permissive`: prefer the more specific dim on a benign
372    // disagreement and keep going, and reconcile graph outputs with their
373    // declared shapes rather than clobbering them. Values that stay symbolic
374    // (genuinely data-dependent extents) are left for the session's JIT
375    // fallback to resolve at run time.
376    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
383/// Fail-fast, load-time validation of everything statically knowable to be
384/// illegal or unsupported (RULES #1: fail at *load*, never via a silent
385/// sentinel at run time).
386///
387/// This is the single cohesive entry point wired into **both** load paths — the
388/// disk/bytes loader ([`build_from_bytes_with_weights`]) and the session's
389/// programmatic entry ([`onnx_runtime_session`]'s `from_parts`/`from_graph`) —
390/// so the checks cannot drift between the two. It runs, in order:
391///
392/// Protobuf-only invariants (`ir_version`, raw SSA names, `ref_attr_name`,
393/// subgraph shadows, and output names) run earlier in
394/// [`validate_model_proto`], before graph construction coalesces names or drops
395/// protobuf-only fields. This IR-level phase then runs:
396///
397/// 1. [`validate_opset_imports`] — every node's domain must declare an opset.
398/// 2. [`validate_no_control_flow`] — allow the implemented subgraph-bearing ops
399///    (`If`/`Loop`/`Scan`) and reject any other op carrying a `GraphProto`
400///    attribute the executor cannot run.
401/// 3. [`validate_no_dangling_refs`] — every consumed tensor must be sourced
402///    (graph input, initializer, or an upstream node output).
403/// 4. [`validate_no_initializer_producer`] — an initializer must be a constant
404///    source; reject any initializer value that is also a node output (shares a
405///    `ValueId` with a producer), which the IR structural check does not cover.
406///
407/// Each rejection names the offending node/op/tensor and explains what is
408/// expected. No sentinel defaults, no silent skips.
409///
410/// Structural invariants that the IR builder already enforces via
411/// [`onnx_runtime_ir::Graph::validate`] at build time — duplicate output
412/// names, dangling value ids, producer/consumer link consistency, and data
413/// dependency cycles — are intentionally *not* re-checked here to avoid drift;
414/// this function adds the checks that path does not cover.
415pub 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
423/// Validate ONNX model metadata and protobuf-level graph invariants that are
424/// intentionally not preserved by the runtime IR (notably `ir_version` and
425/// `AttributeProto::ref_attr_name`).
426///
427/// This runs before graph construction, ensuring invalid names cannot be
428/// coalesced into a single IR value and attribute references cannot be dropped.
429pub fn validate_model_proto(model: &proto::onnx::ModelProto) -> Result<(), LoaderError> {
430    use std::collections::HashSet;
431
432    use proto::onnx::GraphProto;
433
434    // Lower sanity bound only: `ir_version` is a required ONNX field and IR
435    // versions start at 1, so reject an absent (0) or negative version.
436    if model.ir_version < 1 {
437        return Err(LoaderError::InvalidIrVersion {
438            ir_version: model.ir_version,
439        });
440    }
441    // No upper bound. Per the maintainer directive, new ONNX IR versions are
442    // effectively always backward-compatible (they add fields/metadata rather
443    // than breaking existing model semantics), so gating on a version ceiling
444    // only produces false-positive rejections of otherwise-valid newer models.
445    // If a genuinely unsupported construct ever ships, gate on that specific
446    // FEATURE at load time — never on the IR version number.
447    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
552/// Human-readable node label for diagnostics: the quoted ONNX node name, or a
553/// synthetic `<unnamed node #id>` when the model left it blank.
554fn 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
562/// Canonical display domain for a node (`""` renders as `ai.onnx`).
563fn display_domain(domain: &str) -> String {
564    if domain.is_empty() {
565        "ai.onnx".to_string()
566    } else {
567        domain.to_string()
568    }
569}
570
571/// Reject subgraph-bearing (control-flow) ops the runtime cannot execute.
572///
573/// The CPU executor implements the three standard subgraph-bearing control-flow
574/// ops — `If`, `Loop`, and `Scan` (default `ai.onnx` domain) — by recursively
575/// executing their nested [`onnx_runtime_ir::Attribute::Graph`]/`Graphs` bodies.
576/// Any *other* op that smuggles a subgraph attribute (a control-flow construct
577/// this runtime does not implement, or a custom op hiding a nested graph) is
578/// still rejected fast: the executor has no path to run it, so a silent skip or
579/// a late panic would be worse than a clear load-time error.
580///
581/// The check descends into every nested subgraph as well, so an unimplemented
582/// control-flow op buried inside an `If`/`Loop`/`Scan` body is caught at load
583/// rather than surfacing only when that branch/iteration executes.
584pub 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    /// The standard subgraph-bearing ops the CPU executor can run recursively.
592    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            // Report attributes in a deterministic order for stable diagnostics.
599            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                // A subgraph body is fine when its owner is an implemented
608                // control-flow op; otherwise fail fast.
609                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        // Descend into nested bodies so an unimplemented construct inside an
620        // implemented op's subgraph is still caught at load time.
621        for subgraph in graph.subgraphs.values() {
622            check_graph(subgraph)?;
623        }
624        Ok(())
625    }
626
627    check_graph(graph)
628}
629
630/// Reject graphs with a node input that has no source.
631///
632/// The graph builder materializes an unresolved input name as a fresh named
633/// value with no producer (see `graph_builder::get_or_create`); such a value is
634/// legal only if it is a graph input or an initializer. Any other producer-less
635/// consumed value is a dangling reference — a structurally malformed graph that
636/// [`onnx_runtime_ir::Graph::validate`] does not catch (it only requires graph
637/// *outputs* to be sourced, not node inputs). We reject it at load, naming the
638/// offending node and tensor.
639///
640/// Must run after initializers are attached to `graph.initializers` so
641/// initializer-backed inputs are recognized as sourced.
642pub 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                // A dangling value id is caught by IR-level structural
651                // validation; nothing to report here.
652                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
674/// Reject graphs where an initializer value is also produced by a node.
675///
676/// The graph builder maps tensor *names* → [`onnx_runtime_ir::ValueId`] for both
677/// node inputs and node outputs (see `graph_builder::get_or_create`). If a node
678/// output name collides with an initializer name, the node output reuses the
679/// initializer's `ValueId` and `connect_edges` then sets `producer = Some(node)`
680/// on that shared value. [`onnx_runtime_ir::Graph::validate`] rejects a *graph
681/// input* with a producer but has no equivalent check for an *initializer*, so
682/// such a malformed graph passes structural validation.
683///
684/// This matters for memory-safety: the session's weight-streaming path borrows
685/// an initializer's read-only mmap bytes zero-copy. A producer-backed
686/// initializer would let a kernel write through that read-only storage
687/// (SIGSEGV on external data, aliasing UB inline). The executor already refuses
688/// to borrow producer-backed initializers, but rejecting the graph here fails
689/// fast and cleanly regardless of the execution path. We name the tensor and
690/// the offending producing node.
691pub 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}
711///
712/// ONNX treats `""` and `"ai.onnx"` as equivalent spellings of the default
713/// domain. Model-level imports also govern nodes nested in subgraphs.
714pub 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}