onnx_runtime_loader/
encoder.rs

1//! ONNX protobuf **encoding** — the inverse of [`graph_builder`](crate::graph_builder)
2//! and [`weights`](crate::weights) (§19.1, §55.4 dump path).
3//!
4//! Serialises an [`onnx_runtime_ir::Graph`] (plus model-level metadata that the
5//! IR does not itself store) back into an ONNX `ModelProto`, then to protobuf
6//! bytes via `prost`. This is the foundational capability the EPContext writer
7//! (§55.4) builds on, but it is deliberately **model-agnostic**: it hardcodes no
8//! op type, vendor, or model name.
9//!
10//! ## Round-trip contract
11//!
12//! Everything the load path (`decode → build → weights`) preserves survives an
13//! `encode → decode` round-trip byte-for-byte:
14//!
15//! * nodes: `op_type`, `domain`, `input`/`output` order (incl. skipped optional
16//!   slots), attributes, `doc_string`;
17//! * graph inputs / outputs / interior `value_info` (dtype + static & symbolic
18//!   dims, symbols re-emitted by their interned name);
19//! * initializers: all supported dtypes, raw little-endian bytes byte-exact
20//!   (including `STRING` payloads);
21//! * opset imports, `ir_version`, producer fields, model `doc_string`,
22//!   `metadata_props`.
23//!
24//! ### STRING attributes are byte-preserving (§55.3/§55.4)
25//!
26//! ONNX `STRING` attributes are arbitrary byte strings — a compiled-vendor blob,
27//! a relative path, or text — that are not guaranteed to be valid UTF-8. The IR
28//! stores them as raw bytes ([`Attribute::String`](onnx_runtime_ir::Attribute)),
29//! so both decode and encode round-trip the bytes exactly with **zero** op or
30//! attribute-name knowledge. The `EPContext` `ep_cache_context` opaque blob is
31//! preserved purely by this generic mechanism — the encoder contains no
32//! op-specific branch.
33//!
34//! ## Fields deliberately not encoded
35//!
36//! The IR `Graph` does not model these, so they cannot be reproduced and are
37//! emitted empty / default:
38//!
39//! * per-node `name` (the IR [`Node`](onnx_runtime_ir::Node) has no name field);
40//! * `TrainingInfoProto`, `FunctionProto`, sparse initializers, quantization
41//!   annotations (not represented in the IR);
42//! * control-flow **subgraph** formal `input`/`output` lists — the load path
43//!   does not register them as graph I/O for nested graphs. Rather than silently
44//!   emit a subgraph with a truncated signature, [`encode_attribute`] returns a
45//!   clean "unsupported" error for any `Graph`/`Graphs` attribute (no Phase-1 op
46//!   uses control-flow subgraphs).
47
48use std::collections::HashSet;
49use std::path::Path;
50
51use prost::Message;
52
53use onnx_runtime_ir::{
54    Attribute, DataType, Dim, Graph, Node, Shape, TensorData, TypeProto, ValueId, WeightRef,
55};
56
57use crate::proto::onnx::{
58    self, attribute_proto::AttributeType, tensor_shape_proto, type_proto, AttributeProto,
59    GraphProto, ModelProto, NodeProto, OperatorSetIdProto, StringStringEntryProto, TensorProto,
60    TensorShapeProto, ValueInfoProto,
61};
62use crate::weights::WeightStore;
63use crate::LoaderError;
64
65/// Default ONNX `ir_version` stamped when [`ModelMetadata`] does not override it
66/// (IR version 10, the version paired with opset 21).
67pub const DEFAULT_IR_VERSION: i64 = 10;
68
69/// Model-level metadata that the IR [`Graph`] does not itself carry.
70///
71/// The load path drops these (it only keeps `opset_imports` on the `Graph`), so
72/// a caller that wants a faithful `ModelProto` supplies them here. All fields
73/// default to empty/zero except [`ir_version`](ModelMetadata::ir_version).
74#[derive(Clone, Debug, PartialEq, Eq)]
75pub struct ModelMetadata {
76    /// `ModelProto.ir_version`.
77    pub ir_version: i64,
78    /// `ModelProto.producer_name`.
79    pub producer_name: String,
80    /// `ModelProto.producer_version`.
81    pub producer_version: String,
82    /// `ModelProto.domain`.
83    pub domain: String,
84    /// `ModelProto.model_version`.
85    pub model_version: i64,
86    /// `ModelProto.doc_string`.
87    pub doc_string: Option<String>,
88    /// `GraphProto.name` for the top-level graph.
89    pub graph_name: String,
90    /// `ModelProto.metadata_props` (`key → value`), emitted in order.
91    pub metadata_props: Vec<(String, String)>,
92}
93
94impl Default for ModelMetadata {
95    fn default() -> Self {
96        Self {
97            ir_version: DEFAULT_IR_VERSION,
98            producer_name: String::new(),
99            producer_version: String::new(),
100            domain: String::new(),
101            model_version: 0,
102            doc_string: None,
103            graph_name: String::new(),
104            metadata_props: Vec::new(),
105        }
106    }
107}
108
109/// An IR [`Graph`] bundled with the model-level metadata and live weight bytes
110/// needed to encode a complete ONNX `ModelProto`.
111///
112/// Construct with [`Model::new`] and refine via [`Model::with_metadata`] /
113/// [`Model::with_weights`]. A [`WeightStore`] is required only when the graph
114/// has `External`-backed initializers (inline initializers carry their own
115/// bytes).
116pub struct Model<'a> {
117    /// The graph to encode.
118    pub graph: &'a Graph,
119    /// Model-level metadata (see [`ModelMetadata`]).
120    pub metadata: ModelMetadata,
121    /// Live weight store backing any [`WeightRef::External`] initializers.
122    pub weights: Option<&'a WeightStore>,
123}
124
125impl<'a> Model<'a> {
126    /// A model over `graph` with default metadata and no external weight store.
127    ///
128    /// **Note on defaulted metadata:** the load path does not preserve
129    /// model-level metadata (it keeps only `opset_imports` on the `Graph`), so a
130    /// model built this way stamps [`DEFAULT_IR_VERSION`] and empty producer /
131    /// version / `metadata_props` fields. A rewrite path that must reproduce the
132    /// original model faithfully should capture the source metadata and pass it
133    /// via [`Model::with_metadata`] rather than relying on these defaults.
134    pub fn new(graph: &'a Graph) -> Self {
135        Self {
136            graph,
137            metadata: ModelMetadata::default(),
138            weights: None,
139        }
140    }
141
142    /// Attach model-level metadata.
143    pub fn with_metadata(mut self, metadata: ModelMetadata) -> Self {
144        self.metadata = metadata;
145        self
146    }
147
148    /// Attach the live [`WeightStore`] backing external initializers.
149    pub fn with_weights(mut self, weights: &'a WeightStore) -> Self {
150        self.weights = Some(weights);
151        self
152    }
153}
154
155/// Encode `model` into serialized ONNX protobuf bytes.
156pub fn encode_model(model: &Model) -> Result<Vec<u8>, LoaderError> {
157    Ok(encode_model_proto(model)?.encode_to_vec())
158}
159
160/// Encode `model` and write the serialized bytes to `path`.
161pub fn write_model(model: &Model, path: impl AsRef<Path>) -> Result<(), LoaderError> {
162    let bytes = encode_model(model)?;
163    let path = path.as_ref();
164    std::fs::write(path, bytes).map_err(|source| LoaderError::Io {
165        path: path.to_path_buf(),
166        source,
167    })
168}
169
170/// Build the [`ModelProto`] for `model` without serialising it (useful when the
171/// caller wants to mutate the proto before encoding, e.g. the §55.4 writer
172/// splicing in `EPContext` nodes).
173pub fn encode_model_proto(model: &Model) -> Result<ModelProto, LoaderError> {
174    let meta = &model.metadata;
175    let graph = encode_graph_proto(model.graph, model.weights, true, &meta.graph_name)?;
176
177    // Opset imports sorted by domain for deterministic output.
178    let mut opset_import: Vec<OperatorSetIdProto> = model
179        .graph
180        .opset_imports
181        .iter()
182        .map(|(domain, &version)| OperatorSetIdProto {
183            domain: domain.clone(),
184            version: version as i64,
185        })
186        .collect();
187    opset_import.sort_by(|a, b| a.domain.cmp(&b.domain));
188
189    let metadata_props = meta
190        .metadata_props
191        .iter()
192        .map(|(key, value)| StringStringEntryProto {
193            key: key.clone(),
194            value: value.clone(),
195        })
196        .collect();
197
198    Ok(ModelProto {
199        ir_version: meta.ir_version,
200        opset_import,
201        producer_name: meta.producer_name.clone(),
202        producer_version: meta.producer_version.clone(),
203        domain: meta.domain.clone(),
204        model_version: meta.model_version,
205        doc_string: meta.doc_string.clone().unwrap_or_default(),
206        graph: Some(graph),
207        metadata_props,
208        ..Default::default()
209    })
210}
211
212/// Encode a [`Graph`] into a `GraphProto`. `is_top_level` reserved for future
213/// subgraph-specific behaviour; currently the same fields are emitted for both.
214fn encode_graph_proto(
215    graph: &Graph,
216    weights: Option<&WeightStore>,
217    _is_top_level: bool,
218    name: &str,
219) -> Result<GraphProto, LoaderError> {
220    // 1. Initializers, ordered by value id for determinism.
221    let mut init_ids: Vec<ValueId> = graph.initializers.keys().copied().collect();
222    init_ids.sort_by_key(|v| v.0);
223    let mut initializer = Vec::with_capacity(init_ids.len());
224    for vid in &init_ids {
225        let weight = &graph.initializers[vid];
226        let iname = value_name(graph, *vid).unwrap_or_default().to_string();
227        initializer.push(encode_weight(iname, weight, weights)?);
228    }
229
230    // 2. Graph inputs / outputs as ValueInfoProtos.
231    let input: Vec<ValueInfoProto> = graph
232        .inputs
233        .iter()
234        .map(|&vid| encode_value_info(graph, vid))
235        .collect();
236    let output: Vec<ValueInfoProto> = graph
237        .outputs
238        .iter()
239        .map(|&vid| encode_value_info(graph, vid))
240        .collect();
241
242    // 3. Interior value_info: every named value that is not a graph input,
243    //    output, or initializer. Anonymous values (skipped optional outputs,
244    //    unnamed SSA edges) carry no name and are omitted.
245    let mut excluded: HashSet<ValueId> = HashSet::new();
246    excluded.extend(graph.inputs.iter().copied());
247    excluded.extend(graph.outputs.iter().copied());
248    excluded.extend(init_ids.iter().copied());
249    let mut value_info = Vec::new();
250    for (vid, value) in graph.values.iter() {
251        if excluded.contains(&vid) {
252            continue;
253        }
254        if value.name.as_deref().is_some_and(|n| !n.is_empty()) {
255            value_info.push(encode_value_info(graph, vid));
256        }
257    }
258
259    // 4. Nodes, in ascending node-id (== load) order.
260    let mut node = Vec::with_capacity(graph.num_nodes());
261    for (_, n) in graph.nodes.iter() {
262        node.push(encode_node(graph, n)?);
263    }
264
265    Ok(GraphProto {
266        node,
267        name: name.to_string(),
268        initializer,
269        input,
270        output,
271        value_info,
272        ..Default::default()
273    })
274}
275
276/// Encode a single node into a `NodeProto`.
277fn encode_node(graph: &Graph, node: &Node) -> Result<NodeProto, LoaderError> {
278    let input: Vec<String> = node
279        .inputs
280        .iter()
281        .map(|slot| match slot {
282            Some(vid) => value_name(graph, *vid).unwrap_or_default().to_string(),
283            None => String::new(),
284        })
285        .collect();
286    let output: Vec<String> = node
287        .outputs
288        .iter()
289        .map(|&vid| value_name(graph, vid).unwrap_or_default().to_string())
290        .collect();
291
292    // Sort attributes by name so the encoding is deterministic (the IR stores
293    // them in a HashMap).
294    let mut keys: Vec<&String> = node.attributes.keys().collect();
295    keys.sort();
296    let mut attribute = Vec::with_capacity(keys.len());
297    for key in keys {
298        attribute.push(encode_attribute(graph, key, &node.attributes[key])?);
299    }
300
301    Ok(NodeProto {
302        input,
303        output,
304        name: String::new(),
305        op_type: node.op_type.clone(),
306        domain: node.domain.clone(),
307        attribute,
308        doc_string: node.doc_string.clone().unwrap_or_default(),
309        ..Default::default()
310    })
311}
312
313/// Encode one IR [`Attribute`] into an `AttributeProto`, setting the `type`
314/// discriminator to match the populated field (required for IR ≥ 0.0.2).
315fn encode_attribute(
316    graph: &Graph,
317    name: &str,
318    attr: &Attribute,
319) -> Result<AttributeProto, LoaderError> {
320    let mut ap = AttributeProto {
321        name: name.to_string(),
322        ..Default::default()
323    };
324    match attr {
325        Attribute::Int(v) => {
326            ap.i = *v;
327            ap.r#type = AttributeType::Int as i32;
328        }
329        Attribute::Float(v) => {
330            ap.f = *v;
331            ap.r#type = AttributeType::Float as i32;
332        }
333        Attribute::String(s) => {
334            ap.s = s.clone();
335            ap.r#type = AttributeType::String as i32;
336        }
337        Attribute::Ints(v) => {
338            ap.ints = v.clone();
339            ap.r#type = AttributeType::Ints as i32;
340        }
341        Attribute::Floats(v) => {
342            ap.floats = v.clone();
343            ap.r#type = AttributeType::Floats as i32;
344        }
345        Attribute::Strings(v) => {
346            ap.strings = v.clone();
347            ap.r#type = AttributeType::Strings as i32;
348        }
349        Attribute::Tensor(t) => {
350            ap.t = Some(encode_tensor(t));
351            ap.r#type = AttributeType::Tensor as i32;
352        }
353        Attribute::Graph(_) | Attribute::Graphs(_) => {
354            // The load path does not register a nested subgraph's formal
355            // `input`/`output` lists (only its nodes / value_info survive), so a
356            // control-flow subgraph cannot be re-emitted with a faithful
357            // signature. Fail cleanly rather than silently emit a subgraph with a
358            // truncated I/O contract. (No Phase-1 op uses control-flow
359            // subgraphs; see the module docs.)
360            return Err(LoaderError::GraphBuild(format!(
361                "attribute {name:?}: encoding control-flow subgraph attributes is \
362                 unsupported (nested-graph formal I/O is not round-tripped by the \
363                 load path)"
364            )));
365        }
366        Attribute::TypeProto(tp) => {
367            ap.tp = Some(encode_type_proto(graph, tp));
368            ap.r#type = AttributeType::TypeProto as i32;
369        }
370        // The IR carries a SparseTensor attribute variant, but the load path
371        // never builds one (it errors on sparse attribute kinds). Surface a
372        // clean error rather than emit a malformed proto.
373        Attribute::SparseTensor(_) => {
374            return Err(LoaderError::GraphBuild(format!(
375                "attribute {name:?}: SparseTensor encoding is unsupported"
376            )));
377        }
378    }
379    Ok(ap)
380}
381
382/// Encode a [`TensorData`] into a `TensorProto`, preserving raw little-endian
383/// bytes (or `STRING` payloads) byte-exactly.
384fn encode_tensor(t: &TensorData) -> TensorProto {
385    let mut tp = TensorProto {
386        dims: t.dims.iter().map(|&d| d as i64).collect(),
387        data_type: t.dtype.to_onnx(),
388        name: t.name.clone().unwrap_or_default(),
389        ..Default::default()
390    };
391    if t.dtype == DataType::String {
392        tp.string_data = t.strings.iter().map(|s| s.clone().into_bytes()).collect();
393    } else {
394        tp.raw_data = t.data.clone();
395    }
396    tp
397}
398
399/// Encode an initializer [`WeightRef`] into a `TensorProto` named `name`.
400///
401/// Inline weights are emitted directly; external weights are materialised inline
402/// (as `raw_data`) from the provided [`WeightStore`]. Preserving the external
403/// `data_location` reference on write is a follow-up (see the decision note).
404fn encode_weight(
405    name: String,
406    weight: &WeightRef,
407    weights: Option<&WeightStore>,
408) -> Result<TensorProto, LoaderError> {
409    match weight {
410        WeightRef::Inline(t) => {
411            let mut tp = encode_tensor(t);
412            tp.name = name;
413            Ok(tp)
414        }
415        WeightRef::External { dtype, dims, .. } => {
416            if *dtype == DataType::String {
417                return Err(LoaderError::GraphBuild(format!(
418                    "external initializer {name:?}: STRING external data is unsupported"
419                )));
420            }
421            let bytes = weights.and_then(|s| s.bytes(weight)).ok_or_else(|| {
422                LoaderError::GraphBuild(format!(
423                    "external initializer {name:?}: weight bytes unavailable \
424                     (attach a WeightStore via Model::with_weights)"
425                ))
426            })?;
427            Ok(TensorProto {
428                name,
429                data_type: dtype.to_onnx(),
430                dims: dims.iter().map(|&d| d as i64).collect(),
431                raw_data: bytes.to_vec(),
432                ..Default::default()
433            })
434        }
435    }
436}
437
438/// Encode a value's `(dtype, shape)` into a `ValueInfoProto`.
439fn encode_value_info(graph: &Graph, vid: ValueId) -> ValueInfoProto {
440    let value = graph.value(vid);
441    ValueInfoProto {
442        name: value.name.clone().unwrap_or_default(),
443        r#type: Some(encode_tensor_type(graph, value.dtype, &value.shape)),
444        ..Default::default()
445    }
446}
447
448/// Build a tensor `TypeProto` from an element type and shape.
449fn encode_tensor_type(graph: &Graph, dtype: DataType, shape: &Shape) -> onnx::TypeProto {
450    onnx::TypeProto {
451        value: Some(type_proto::Value::TensorType(type_proto::Tensor {
452            elem_type: dtype.to_onnx(),
453            shape: Some(encode_shape(graph, shape)),
454        })),
455        ..Default::default()
456    }
457}
458
459/// Encode an IR [`Shape`] into a `TensorShapeProto`. Static dims become
460/// `dim_value`; symbolic dims are re-emitted by their interned name as
461/// `dim_param`, or as a valueless (unknown) dimension when unnamed.
462fn encode_shape(graph: &Graph, shape: &Shape) -> TensorShapeProto {
463    use tensor_shape_proto::{dimension::Value as DV, Dimension};
464    let dim = shape
465        .iter()
466        .map(|d| {
467            let value = match d {
468                Dim::Static(n) => Some(DV::DimValue(*n as i64)),
469                Dim::Symbolic(sym) => graph
470                    .symbol_constraints
471                    .get(sym)
472                    .and_then(|c| c.name.clone())
473                    .map(DV::DimParam),
474            };
475            Dimension {
476                value,
477                ..Default::default()
478            }
479        })
480        .collect();
481    TensorShapeProto { dim }
482}
483
484/// Encode an IR [`TypeProto`] (the inverse of `graph_builder::convert_type_proto`).
485fn encode_type_proto(graph: &Graph, tp: &TypeProto) -> onnx::TypeProto {
486    let value = match tp {
487        TypeProto::Tensor { dtype, shape } => type_proto::Value::TensorType(type_proto::Tensor {
488            elem_type: dtype.to_onnx(),
489            shape: Some(encode_shape(graph, shape)),
490        }),
491        TypeProto::SparseTensor { dtype, shape } => {
492            type_proto::Value::SparseTensorType(type_proto::SparseTensor {
493                elem_type: dtype.to_onnx(),
494                shape: Some(encode_shape(graph, shape)),
495            })
496        }
497        TypeProto::Sequence(inner) => type_proto::Value::SequenceType(Box::new(type_proto::Sequence {
498            elem_type: Some(Box::new(encode_type_proto(graph, inner))),
499        })),
500        TypeProto::Optional(inner) => type_proto::Value::OptionalType(Box::new(type_proto::Optional {
501            elem_type: Some(Box::new(encode_type_proto(graph, inner))),
502        })),
503        TypeProto::Map { key, value } => type_proto::Value::MapType(Box::new(type_proto::Map {
504            key_type: key.to_onnx(),
505            value_type: Some(Box::new(encode_type_proto(graph, value))),
506        })),
507    };
508    onnx::TypeProto {
509        value: Some(value),
510        ..Default::default()
511    }
512}
513
514/// The name of a graph value, if it has one.
515fn value_name(graph: &Graph, vid: ValueId) -> Option<&str> {
516    graph.try_value(vid).and_then(|v| v.name.as_deref())
517}
onnx_runtime_loader/encoder.rs

onnx_runtime_loader/
encoder.rs
