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    /// Resolve a weight descriptor to its raw little-endian bytes.
36    ///
37    /// For inline weights this borrows the stored bytes; for external weights
38    /// it slices into the memory-mapped file. Returns `None` if an external
39    /// mapping is missing or the `[offset, offset+length)` window is out of
40    /// bounds.
41    pub fn bytes<'a>(&'a self, weight: &'a WeightRef) -> Option<&'a [u8]> {
42        match weight {
43            WeightRef::Inline(t) => Some(&t.data),
44            WeightRef::External {
45                path,
46                offset,
47                length,
48                ..
49            } => {
50                let mmap = self.mmaps.get(path)?;
51                mmap.get(*offset..offset.checked_add(*length)?)
52            }
53        }
54    }
55
56    fn mmap_file(&mut self, path: &Path) -> Result<(), LoaderError> {
57        if self.mmaps.contains_key(path) {
58            return Ok(());
59        }
60        let file = File::open(path).map_err(|_| LoaderError::ExternalDataNotFound {
61            path: path.to_path_buf(),
62        })?;
63        // SAFETY: we hold the `File` open for the duration of the map and never
64        // expose a mutable view. The mapped bytes are treated as immutable
65        // weight storage. This is the only `unsafe` in the loader; the IR crate
66        // stays `#![forbid(unsafe_code)]`.
67        let mmap = unsafe { Mmap::map(&file) }.map_err(|e| LoaderError::Mmap(e.to_string()))?;
68        self.mmaps.insert(path.to_path_buf(), mmap);
69        Ok(())
70    }
71}
72
73/// Resolve all initializers, memory-mapping external data relative to
74/// `model_dir`. `name_map` maps initializer names to the graph value ids
75/// created by the [`graph_builder`](crate::graph_builder).
76pub fn load_weights(
77    model: &ModelProto,
78    model_dir: &Path,
79    name_map: &HashMap<String, ValueId>,
80) -> Result<WeightStore, LoaderError> {
81    let mut store = WeightStore::new();
82    let Some(graph) = model.graph.as_ref() else {
83        return Ok(store);
84    };
85
86    for init in &graph.initializer {
87        let Some(&vid) = name_map.get(&init.name) else {
88            // An initializer with no corresponding graph value: skip it. The
89            // builder registers a value for every initializer, so this only
90            // happens for malformed models.
91            continue;
92        };
93        let weight = resolve_initializer(&mut store, init, model_dir)?;
94        store.weights.insert(vid, weight);
95    }
96
97    Ok(store)
98}
99
100fn resolve_initializer(
101    store: &mut WeightStore,
102    init: &TensorProto,
103    model_dir: &Path,
104) -> Result<WeightRef, LoaderError> {
105    let dtype = DataType::from_onnx(init.data_type).ok_or_else(|| {
106        LoaderError::UnsupportedDataType {
107            raw: init.data_type,
108            context: format!("initializer {:?}", init.name),
109        }
110    })?;
111    let dims: Vec<usize> = init.dims.iter().map(|&d| d.max(0) as usize).collect();
112
113    if init.data_location == tensor_proto::DataLocation::External as i32 {
114        let mut location = None;
115        let mut offset: usize = 0;
116        let mut length: Option<usize> = None;
117        for kv in &init.external_data {
118            match kv.key.as_str() {
119                "location" => location = Some(kv.value.clone()),
120                "offset" => offset = kv.value.parse().unwrap_or(0),
121                "length" => length = kv.value.parse().ok(),
122                _ => {}
123            }
124        }
125        let location = location.ok_or_else(|| {
126            LoaderError::GraphBuild(format!(
127                "external initializer {:?} missing 'location'",
128                init.name
129            ))
130        })?;
131        let path = resolve_external_path(model_dir, &location)?;
132        store.mmap_file(&path)?;
133        let numel: usize = dims.iter().product();
134        let length = length.unwrap_or_else(|| dtype.storage_bytes(numel));
135        // Validate the window lies within the mapped file (catches truncated or
136        // mis-described external data early).
137        if let Some(mmap) = store.mmaps.get(&path) {
138            let end = offset.checked_add(length);
139            if end.is_none_or(|e| e > mmap.len()) {
140                return Err(LoaderError::Mmap(format!(
141                    "external initializer {:?}: window [{offset}, {:?}) exceeds file {} ({} bytes)",
142                    init.name,
143                    end,
144                    path.display(),
145                    mmap.len()
146                )));
147            }
148        }
149        Ok(WeightRef::External {
150            path,
151            offset,
152            length,
153            dtype,
154            dims,
155        })
156    } else {
157        let data = tensor_data_from_proto(init, dtype, &dims)?;
158        Ok(WeightRef::Inline(data))
159    }
160}
161
162/// Join an external-data location onto `model_dir`, rejecting paths that can
163/// escape the model directory.
164fn resolve_external_path(model_dir: &Path, location: &str) -> Result<PathBuf, LoaderError> {
165    guarded_join(model_dir, location).map_err(|reason| LoaderError::ExternalDataPath {
166        path: location.to_string(),
167        reason,
168    })
169}
170
171/// Convert a `TensorProto`'s payload into an IR [`TensorData`] with raw
172/// little-endian bytes (or string payloads for `STRING` tensors).
173pub(crate) fn tensor_data_from_proto(
174    proto: &TensorProto,
175    dtype: DataType,
176    dims: &[usize],
177) -> Result<TensorData, LoaderError> {
178    let mut td = TensorData::from_raw(dtype, dims.to_vec(), Vec::new());
179    if !proto.name.is_empty() {
180        td.name = Some(proto.name.clone());
181    }
182
183    if dtype == DataType::String {
184        td.strings = proto
185            .string_data
186            .iter()
187            .map(|b| String::from_utf8_lossy(b).into_owned())
188            .collect();
189        return Ok(td);
190    }
191
192    // Prefer raw_data when present (the common, mmap-friendly encoding).
193    if !proto.raw_data.is_empty() {
194        td.data = proto.raw_data.clone();
195        return Ok(td);
196    }
197
198    // Otherwise serialise the type-specific repeated field to LE bytes.
199    td.data = match dtype {
200        DataType::Float32 => proto
201            .float_data
202            .iter()
203            .flat_map(|v| v.to_le_bytes())
204            .collect(),
205        DataType::Float64 => proto
206            .double_data
207            .iter()
208            .flat_map(|v| v.to_le_bytes())
209            .collect(),
210        DataType::Int64 => proto
211            .int64_data
212            .iter()
213            .flat_map(|v| v.to_le_bytes())
214            .collect(),
215        DataType::Uint64 | DataType::Uint32 => proto
216            .uint64_data
217            .iter()
218            .flat_map(|v| match dtype {
219                DataType::Uint32 => (*v as u32).to_le_bytes().to_vec(),
220                _ => v.to_le_bytes().to_vec(),
221            })
222            .collect(),
223        DataType::Int32 => proto
224            .int32_data
225            .iter()
226            .flat_map(|v| v.to_le_bytes())
227            .collect(),
228        // Types packed into int32_data at a narrower width.
229        DataType::Int16 | DataType::Uint16 | DataType::Float16 | DataType::BFloat16 => proto
230            .int32_data
231            .iter()
232            .flat_map(|v| (*v as u16).to_le_bytes())
233            .collect(),
234        DataType::Int8 | DataType::Uint8 | DataType::Bool => {
235            proto.int32_data.iter().map(|v| *v as u8).collect()
236        }
237        _ => Vec::new(),
238    };
239    Ok(td)
240}
onnx_runtime_loader/weights.rs

onnx_runtime_loader/
weights.rs
