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