onnx_runtime_loader/
weights.rs

1//! Initializer weight resolution: inline data and external `mmap` (§19.2, §12).
2//!
3//! Turns each `TensorProto` initializer into an [`onnx_runtime_ir::WeightRef`]
4//! descriptor that the IR stores. Inline tensors keep their bytes; external
5//! tensors are described by `(path, offset, length)` and the referenced files
6//! are memory-mapped so downstream consumers get zero-copy access via
7//! [`WeightStore::bytes`].
8
9use std::collections::HashMap;
10use std::fs::File;
11use std::path::{Path, PathBuf};
12
13use memmap2::Mmap;
14use onnx_runtime_ir::{DataType, TensorData, ValueId, WeightRef};
15
16use crate::proto::onnx::{tensor_proto, ModelProto, TensorProto};
17use crate::{pathsafe::guarded_join, LoaderError};
18
19/// A resolved set of initializer weights, keyed by the value they populate,
20/// plus the live memory maps backing any external data files.
21#[derive(Debug, Default)]
22pub struct WeightStore {
23    /// IR weight descriptors, keyed by the graph value they initialize.
24    pub weights: HashMap<ValueId, WeightRef>,
25    /// Live memory maps for external-data files, keyed by absolute path.
26    mmaps: HashMap<PathBuf, Mmap>,
27}
28
29impl WeightStore {
30    /// An empty store.
31    pub fn new() -> Self {
32        Self::default()
33    }
34
35    /// Memory-map an external-data file and register it under `path`, so any
36    /// [`WeightRef::External`] whose `path` matches resolves zero-copy via
37    /// [`bytes`](Self::bytes). Idempotent: mapping the same path twice is a
38    /// no-op. This is the programmatic counterpart to the loader path (which
39    /// maps files while resolving `TensorProto` initializers), useful when
40    /// constructing a store by hand.
41    ///
42    /// The map is read-only and kept alive for the store's lifetime; callers
43    /// must not mutate or unlink the file while the store is live.
44    pub fn map_external(&mut self, path: impl AsRef<Path>) -> Result<(), LoaderError> {
45        self.mmap_file(path.as_ref())
46    }
47
48    /// Resolve a weight descriptor to its raw little-endian bytes.
49    ///
50    /// For inline weights this borrows the stored bytes; for external weights
51    /// it slices into the memory-mapped file. Returns `None` if an external
52    /// mapping is missing or the `[offset, offset+length)` window is out of
53    /// bounds.
54    pub fn bytes<'a>(&'a self, weight: &'a WeightRef) -> Option<&'a [u8]> {
55        match weight {
56            WeightRef::Inline(t) => Some(&t.data),
57            WeightRef::External {
58                path,
59                offset,
60                length,
61                ..
62            } => {
63                let mmap = self.mmaps.get(path)?;
64                mmap.get(*offset..offset.checked_add(*length)?)
65            }
66        }
67    }
68
69    fn mmap_file(&mut self, path: &Path) -> Result<(), LoaderError> {
70        if self.mmaps.contains_key(path) {
71            return Ok(());
72        }
73        let file = File::open(path).map_err(|_| LoaderError::ExternalDataNotFound {
74            path: path.to_path_buf(),
75        })?;
76        // SAFETY: we hold the `File` open for the duration of the map and never
77        // expose a mutable view. The mapped bytes are treated as immutable
78        // weight storage. This is the only `unsafe` in the loader; the IR crate
79        // stays `#![forbid(unsafe_code)]`.
80        let mmap = unsafe { Mmap::map(&file) }.map_err(|e| LoaderError::Mmap(e.to_string()))?;
81        self.mmaps.insert(path.to_path_buf(), mmap);
82        Ok(())
83    }
84}
85
86/// Resolve all initializers, memory-mapping external data relative to
87/// `model_dir`. `name_map` maps initializer names to the graph value ids
88/// created by the [`graph_builder`](crate::graph_builder).
89pub fn load_weights(
90    model: &ModelProto,
91    model_dir: &Path,
92    name_map: &HashMap<String, ValueId>,
93) -> Result<WeightStore, LoaderError> {
94    let mut store = WeightStore::new();
95    let Some(graph) = model.graph.as_ref() else {
96        return Ok(store);
97    };
98
99    for init in &graph.initializer {
100        let Some(&vid) = name_map.get(&init.name) else {
101            // An initializer with no corresponding graph value: skip it. The
102            // builder registers a value for every initializer, so this only
103            // happens for malformed models.
104            continue;
105        };
106        let weight = resolve_initializer(&mut store, init, model_dir)?;
107        store.weights.insert(vid, weight);
108    }
109
110    Ok(store)
111}
112
113fn resolve_initializer(
114    store: &mut WeightStore,
115    init: &TensorProto,
116    model_dir: &Path,
117) -> Result<WeightRef, LoaderError> {
118    let dtype = DataType::from_onnx(init.data_type).ok_or_else(|| {
119        LoaderError::UnsupportedDataType {
120            raw: init.data_type,
121            context: format!("initializer {:?}", init.name),
122        }
123    })?;
124    let dims: Vec<usize> = init.dims.iter().map(|&d| d.max(0) as usize).collect();
125
126    if init.data_location == tensor_proto::DataLocation::External as i32 {
127        let mut location = None;
128        let mut offset: usize = 0;
129        let mut length: Option<usize> = None;
130        for kv in &init.external_data {
131            match kv.key.as_str() {
132                "location" => location = Some(kv.value.clone()),
133                "offset" => offset = kv.value.parse().unwrap_or(0),
134                "length" => length = kv.value.parse().ok(),
135                _ => {}
136            }
137        }
138        let location = location.ok_or_else(|| {
139            LoaderError::GraphBuild(format!(
140                "external initializer {:?} missing 'location'",
141                init.name
142            ))
143        })?;
144        let path = resolve_external_path(model_dir, &location)?;
145        store.mmap_file(&path)?;
146        let numel: usize = dims.iter().product();
147        let length = length.unwrap_or_else(|| dtype.storage_bytes(numel));
148        // Validate the window lies within the mapped file (catches truncated or
149        // mis-described external data early).
150        if let Some(mmap) = store.mmaps.get(&path) {
151            let end = offset.checked_add(length);
152            if end.is_none_or(|e| e > mmap.len()) {
153                return Err(LoaderError::Mmap(format!(
154                    "external initializer {:?}: window [{offset}, {:?}) exceeds file {} ({} bytes)",
155                    init.name,
156                    end,
157                    path.display(),
158                    mmap.len()
159                )));
160            }
161        }
162        Ok(WeightRef::External {
163            path,
164            offset,
165            length,
166            dtype,
167            dims,
168        })
169    } else {
170        let data = tensor_data_from_proto(init, dtype, &dims)?;
171        Ok(WeightRef::Inline(data))
172    }
173}
174
175/// Join an external-data location onto `model_dir`, rejecting paths that can
176/// escape the model directory.
177fn resolve_external_path(model_dir: &Path, location: &str) -> Result<PathBuf, LoaderError> {
178    guarded_join(model_dir, location).map_err(|reason| LoaderError::ExternalDataPath {
179        path: location.to_string(),
180        reason,
181    })
182}
183
184/// Convert a `TensorProto`'s payload into an IR [`TensorData`] with raw
185/// little-endian bytes (or string payloads for `STRING` tensors).
186pub(crate) fn tensor_data_from_proto(
187    proto: &TensorProto,
188    dtype: DataType,
189    dims: &[usize],
190) -> Result<TensorData, LoaderError> {
191    let mut td = TensorData::from_raw(dtype, dims.to_vec(), Vec::new());
192    if !proto.name.is_empty() {
193        td.name = Some(proto.name.clone());
194    }
195
196    if dtype == DataType::String {
197        td.strings = proto
198            .string_data
199            .iter()
200            .map(|b| String::from_utf8_lossy(b).into_owned())
201            .collect();
202        return Ok(td);
203    }
204
205    // Prefer raw_data when present (the common, mmap-friendly encoding).
206    if !proto.raw_data.is_empty() {
207        td.data = proto.raw_data.clone();
208        return Ok(td);
209    }
210
211    // Otherwise serialise the type-specific repeated field to LE bytes.
212    td.data = match dtype {
213        DataType::Float32 => proto
214            .float_data
215            .iter()
216            .flat_map(|v| v.to_le_bytes())
217            .collect(),
218        DataType::Float64 => proto
219            .double_data
220            .iter()
221            .flat_map(|v| v.to_le_bytes())
222            .collect(),
223        DataType::Int64 => proto
224            .int64_data
225            .iter()
226            .flat_map(|v| v.to_le_bytes())
227            .collect(),
228        DataType::Uint64 | DataType::Uint32 => proto
229            .uint64_data
230            .iter()
231            .flat_map(|v| match dtype {
232                DataType::Uint32 => (*v as u32).to_le_bytes().to_vec(),
233                _ => v.to_le_bytes().to_vec(),
234            })
235            .collect(),
236        DataType::Int32 => proto
237            .int32_data
238            .iter()
239            .flat_map(|v| v.to_le_bytes())
240            .collect(),
241        // Types packed into int32_data at a narrower width.
242        DataType::Int16 | DataType::Uint16 | DataType::Float16 | DataType::BFloat16 => proto
243            .int32_data
244            .iter()
245            .flat_map(|v| (*v as u16).to_le_bytes())
246            .collect(),
247        DataType::Int8 | DataType::Uint8 | DataType::Bool => {
248            proto.int32_data.iter().map(|v| *v as u8).collect()
249        }
250        _ => Vec::new(),
251    };
252    Ok(td)
253}
onnx_runtime_loader/weights.rs

onnx_runtime_loader/
weights.rs
