zsfm_checkpoint/
onnx.rs

1//! ONNX checkpoint loading: weights are the graph's initializer tensors.
2//!
3//! ONNX is protobuf, but the weight layout only needs a handful of fields
4//! (ModelProto.graph → GraphProto.initializer → TensorProto), so this module
5//! decodes the protobuf wire format directly — no protoc / prost build-time
6//! dependency. Both storage layouts are handled: `raw_data` bytes and the typed
7//! repeated fields (`float_data`, `int32_data`, …), packed or unpacked.
8//! F32/F16/BF16 are preserved exactly; DOUBLE and integer/bool types are cast
9//! to F32 with a note. Models using external data files are rejected with
10//! guidance.
11
12use std::path::Path;
13
14use anyhow::{bail, Context, Result};
15
16use crate::cast::{f32_to_bytes, SrcDtype};
17use crate::read::{Checkpoint, RawTensor};
18
19// onnx.TensorProto.DataType values.
20const DT_FLOAT: i32 = 1;
21const DT_UINT8: i32 = 2;
22const DT_INT8: i32 = 3;
23const DT_UINT16: i32 = 4;
24const DT_INT16: i32 = 5;
25const DT_INT32: i32 = 6;
26const DT_INT64: i32 = 7;
27const DT_BOOL: i32 = 9;
28const DT_FLOAT16: i32 = 10;
29const DT_DOUBLE: i32 = 11;
30const DT_UINT32: i32 = 12;
31const DT_UINT64: i32 = 13;
32const DT_BFLOAT16: i32 = 16;
33
34const LOCATION_EXTERNAL: u64 = 1;
35
36pub fn load_onnx(path: &Path) -> Result<Checkpoint> {
37    const MAX_ONNX_BYTES: u64 = 8 << 30;
38    const MAX_INITIALIZERS: usize = 100_000;
39    let meta = std::fs::metadata(path).with_context(|| format!("stat {}", path.display()))?;
40    anyhow::ensure!(
41        meta.len() <= MAX_ONNX_BYTES,
42        "ONNX file too large: {} bytes (max {MAX_ONNX_BYTES})",
43        meta.len()
44    );
45    let bytes = std::fs::read(path).with_context(|| format!("read {}", path.display()))?;
46    let initializers = parse_model(&bytes).context("parse ONNX protobuf")?;
47    anyhow::ensure!(
48        !initializers.is_empty(),
49        "ONNX graph has no initializer tensors — the weights may live in external \
50         data files or be supplied as runtime inputs"
51    );
52    anyhow::ensure!(
53        initializers.len() <= MAX_INITIALIZERS,
54        "ONNX has too many initializers: {} (max {MAX_INITIALIZERS})",
55        initializers.len()
56    );
57
58    let mut tensors = Vec::with_capacity(initializers.len());
59    for t in &initializers {
60        tensors.push(tensor_proto_to_raw(t)?);
61    }
62    Ok(Checkpoint {
63        tensors,
64        metadata: Vec::new(),
65    })
66}
67
68// ---------------------------------------------------------------------------
69// Minimal protobuf wire-format reader
70// ---------------------------------------------------------------------------
71
72struct ProtoReader<'a> {
73    buf: &'a [u8],
74    pos: usize,
75}
76
77impl<'a> ProtoReader<'a> {
78    fn new(buf: &'a [u8]) -> Self {
79        Self { buf, pos: 0 }
80    }
81
82    fn done(&self) -> bool {
83        self.pos >= self.buf.len()
84    }
85
86    fn read_varint(&mut self) -> Result<u64> {
87        let mut value = 0u64;
88        let mut shift = 0u32;
89        loop {
90            let byte = *self
91                .buf
92                .get(self.pos)
93                .context("protobuf: truncated varint")?;
94            self.pos += 1;
95            if shift < 64 {
96                value |= u64::from(byte & 0x7F) << shift;
97            }
98            if byte & 0x80 == 0 {
99                return Ok(value);
100            }
101            shift += 7;
102            anyhow::ensure!(shift <= 70, "protobuf: varint too long");
103        }
104    }
105
106    /// Returns (field_number, wire_type).
107    fn read_tag(&mut self) -> Result<(u64, u8)> {
108        let tag = self.read_varint()?;
109        Ok((tag >> 3, (tag & 0x7) as u8))
110    }
111
112    fn read_len_delimited(&mut self) -> Result<&'a [u8]> {
113        let len = self.read_varint()? as usize;
114        let end = self
115            .pos
116            .checked_add(len)
117            .filter(|&e| e <= self.buf.len())
118            .context("protobuf: truncated length-delimited field")?;
119        let slice = &self.buf[self.pos..end];
120        self.pos = end;
121        Ok(slice)
122    }
123
124    fn read_fixed32(&mut self) -> Result<u32> {
125        let end = self.pos + 4;
126        anyhow::ensure!(end <= self.buf.len(), "protobuf: truncated fixed32");
127        let v = u32::from_le_bytes(self.buf[self.pos..end].try_into().unwrap());
128        self.pos = end;
129        Ok(v)
130    }
131
132    fn read_fixed64(&mut self) -> Result<u64> {
133        let end = self.pos + 8;
134        anyhow::ensure!(end <= self.buf.len(), "protobuf: truncated fixed64");
135        let v = u64::from_le_bytes(self.buf[self.pos..end].try_into().unwrap());
136        self.pos = end;
137        Ok(v)
138    }
139
140    fn skip(&mut self, wire_type: u8) -> Result<()> {
141        match wire_type {
142            0 => {
143                self.read_varint()?;
144            }
145            1 => {
146                self.read_fixed64()?;
147            }
148            2 => {
149                self.read_len_delimited()?;
150            }
151            5 => {
152                self.read_fixed32()?;
153            }
154            other => bail!("protobuf: unsupported wire type {other}"),
155        }
156        Ok(())
157    }
158}
159
160#[derive(Default)]
161struct TensorProto {
162    dims: Vec<i64>,
163    data_type: i32,
164    float_data: Vec<f32>,
165    int32_data: Vec<i32>,
166    int64_data: Vec<i64>,
167    uint64_data: Vec<u64>,
168    double_data: Vec<f64>,
169    name: String,
170    raw_data: Vec<u8>,
171    data_location: u64,
172}
173
174/// ModelProto: field 7 = graph (GraphProto).
175fn parse_model(buf: &[u8]) -> Result<Vec<TensorProto>> {
176    let mut r = ProtoReader::new(buf);
177    let mut initializers = Vec::new();
178    while !r.done() {
179        let (field, wire) = r.read_tag()?;
180        if field == 7 && wire == 2 {
181            let graph = r.read_len_delimited()?;
182            initializers.extend(parse_graph(graph)?);
183        } else {
184            r.skip(wire)?;
185        }
186    }
187    Ok(initializers)
188}
189
190/// GraphProto: field 5 = initializer (repeated TensorProto).
191fn parse_graph(buf: &[u8]) -> Result<Vec<TensorProto>> {
192    let mut r = ProtoReader::new(buf);
193    let mut initializers = Vec::new();
194    while !r.done() {
195        let (field, wire) = r.read_tag()?;
196        if field == 5 && wire == 2 {
197            let tensor = r.read_len_delimited()?;
198            initializers.push(parse_tensor(tensor)?);
199        } else {
200            r.skip(wire)?;
201        }
202    }
203    Ok(initializers)
204}
205
206fn parse_tensor(buf: &[u8]) -> Result<TensorProto> {
207    let mut r = ProtoReader::new(buf);
208    let mut t = TensorProto::default();
209    while !r.done() {
210        let (field, wire) = r.read_tag()?;
211        match (field, wire) {
212            // dims: repeated int64 (packed or unpacked)
213            (1, 0) => t.dims.push(r.read_varint()? as i64),
214            (1, 2) => {
215                let mut p = ProtoReader::new(r.read_len_delimited()?);
216                while !p.done() {
217                    t.dims.push(p.read_varint()? as i64);
218                }
219            }
220            (2, 0) => t.data_type = r.read_varint()? as i32,
221            // float_data: repeated float (packed or unpacked)
222            (4, 5) => t.float_data.push(f32::from_bits(r.read_fixed32()?)),
223            (4, 2) => {
224                let mut p = ProtoReader::new(r.read_len_delimited()?);
225                while !p.done() {
226                    t.float_data.push(f32::from_bits(p.read_fixed32()?));
227                }
228            }
229            // int32_data: repeated int32
230            (5, 0) => t.int32_data.push(r.read_varint()? as i64 as i32),
231            (5, 2) => {
232                let mut p = ProtoReader::new(r.read_len_delimited()?);
233                while !p.done() {
234                    t.int32_data.push(p.read_varint()? as i64 as i32);
235                }
236            }
237            // int64_data: repeated int64
238            (7, 0) => t.int64_data.push(r.read_varint()? as i64),
239            (7, 2) => {
240                let mut p = ProtoReader::new(r.read_len_delimited()?);
241                while !p.done() {
242                    t.int64_data.push(p.read_varint()? as i64);
243                }
244            }
245            (8, 2) => {
246                t.name = String::from_utf8_lossy(r.read_len_delimited()?).into_owned();
247            }
248            (9, 2) => t.raw_data = r.read_len_delimited()?.to_vec(),
249            // double_data: repeated double
250            (10, 1) => t.double_data.push(f64::from_bits(r.read_fixed64()?)),
251            (10, 2) => {
252                let mut p = ProtoReader::new(r.read_len_delimited()?);
253                while !p.done() {
254                    t.double_data.push(f64::from_bits(p.read_fixed64()?));
255                }
256            }
257            // uint64_data: repeated uint64
258            (11, 0) => t.uint64_data.push(r.read_varint()?),
259            (11, 2) => {
260                let mut p = ProtoReader::new(r.read_len_delimited()?);
261                while !p.done() {
262                    t.uint64_data.push(p.read_varint()?);
263                }
264            }
265            (14, 0) => t.data_location = r.read_varint()?,
266            (_, w) => r.skip(w)?,
267        }
268    }
269    Ok(t)
270}
271
272// ---------------------------------------------------------------------------
273// TensorProto → RawTensor
274// ---------------------------------------------------------------------------
275
276fn tensor_proto_to_raw(t: &TensorProto) -> Result<RawTensor> {
277    let name = t.name.clone();
278    if t.data_location == LOCATION_EXTERNAL {
279        bail!(
280            "tensor {name}: stored as external data — re-export the model with weights \
281             embedded (onnx.external_data_helper.load_external_data_for_model + save) \
282             and convert again"
283        );
284    }
285    let shape: Vec<u64> = t.dims.iter().map(|&d| d as u64).collect();
286
287    let (dtype, data): (SrcDtype, Vec<u8>) = match t.data_type {
288        DT_FLOAT => {
289            if !t.raw_data.is_empty() {
290                (SrcDtype::F32, t.raw_data.clone())
291            } else {
292                (SrcDtype::F32, f32_to_bytes(&t.float_data))
293            }
294        }
295        DT_FLOAT16 | DT_BFLOAT16 => {
296            let dtype = if t.data_type == DT_FLOAT16 {
297                SrcDtype::F16
298            } else {
299                SrcDtype::BF16
300            };
301            if !t.raw_data.is_empty() {
302                (dtype, t.raw_data.clone())
303            } else {
304                // Spec: 16-bit floats without raw_data are stored as the low 16
305                // bits of int32_data entries.
306                let bytes = t
307                    .int32_data
308                    .iter()
309                    .flat_map(|&v| (v as u16).to_le_bytes())
310                    .collect();
311                (dtype, bytes)
312            }
313        }
314        DT_DOUBLE => {
315            eprintln!("note: tensor {name}: casting DOUBLE to F32");
316            let vals: Vec<f32> = if !t.raw_data.is_empty() {
317                t.raw_data
318                    .chunks_exact(8)
319                    .map(|c| f64::from_le_bytes(c.try_into().unwrap()) as f32)
320                    .collect()
321            } else {
322                t.double_data.iter().map(|&v| v as f32).collect()
323            };
324            (SrcDtype::F32, f32_to_bytes(&vals))
325        }
326        DT_INT64 => {
327            eprintln!("note: tensor {name}: casting INT64 to F32");
328            let vals: Vec<f32> = if !t.raw_data.is_empty() {
329                t.raw_data
330                    .chunks_exact(8)
331                    .map(|c| i64::from_le_bytes(c.try_into().unwrap()) as f32)
332                    .collect()
333            } else {
334                t.int64_data.iter().map(|&v| v as f32).collect()
335            };
336            (SrcDtype::F32, f32_to_bytes(&vals))
337        }
338        DT_UINT64 => {
339            eprintln!("note: tensor {name}: casting UINT64 to F32");
340            let vals: Vec<f32> = if !t.raw_data.is_empty() {
341                t.raw_data
342                    .chunks_exact(8)
343                    .map(|c| u64::from_le_bytes(c.try_into().unwrap()) as f32)
344                    .collect()
345            } else {
346                t.uint64_data.iter().map(|&v| v as f32).collect()
347            };
348            (SrcDtype::F32, f32_to_bytes(&vals))
349        }
350        DT_INT32 | DT_UINT32 | DT_INT16 | DT_UINT16 | DT_INT8 | DT_UINT8 | DT_BOOL => {
351            eprintln!("note: tensor {name}: casting integer/bool data to F32");
352            let vals: Vec<f32> = if !t.raw_data.is_empty() {
353                decode_small_ints(&t.raw_data, t.data_type)?
354            } else {
355                // The typed field for every sub-32-bit integer type is int32_data.
356                t.int32_data.iter().map(|&v| v as f32).collect()
357            };
358            (SrcDtype::F32, f32_to_bytes(&vals))
359        }
360        other => bail!("tensor {name}: ONNX data_type {other} not supported"),
361    };
362
363    Ok(RawTensor {
364        name,
365        shape,
366        dtype,
367        data,
368    })
369}
370
371fn decode_small_ints(raw: &[u8], data_type: i32) -> Result<Vec<f32>> {
372    Ok(match data_type {
373        DT_INT32 => raw
374            .chunks_exact(4)
375            .map(|c| i32::from_le_bytes(c.try_into().unwrap()) as f32)
376            .collect(),
377        DT_UINT32 => raw
378            .chunks_exact(4)
379            .map(|c| u32::from_le_bytes(c.try_into().unwrap()) as f32)
380            .collect(),
381        DT_INT16 => raw
382            .chunks_exact(2)
383            .map(|c| i16::from_le_bytes(c.try_into().unwrap()) as f32)
384            .collect(),
385        DT_UINT16 => raw
386            .chunks_exact(2)
387            .map(|c| u16::from_le_bytes(c.try_into().unwrap()) as f32)
388            .collect(),
389        DT_INT8 => raw.iter().map(|&b| b as i8 as f32).collect(),
390        DT_UINT8 => raw.iter().map(|&b| b as f32).collect(),
391        DT_BOOL => raw.iter().map(|&b| (b != 0) as u8 as f32).collect(),
392        other => bail!("unexpected small-int data_type {other}"),
393    })
394}
395
396#[cfg(test)]
397mod tests {
398    use super::*;
399
400    // Hand-encoded protobuf helpers for building fixture bytes.
401    fn varint(mut v: u64, out: &mut Vec<u8>) {
402        loop {
403            let byte = (v & 0x7F) as u8;
404            v >>= 7;
405            if v == 0 {
406                out.push(byte);
407                break;
408            }
409            out.push(byte | 0x80);
410        }
411    }
412
413    fn tag(field: u64, wire: u8, out: &mut Vec<u8>) {
414        varint(field << 3 | wire as u64, out);
415    }
416
417    fn len_delim(field: u64, payload: &[u8], out: &mut Vec<u8>) {
418        tag(field, 2, out);
419        varint(payload.len() as u64, out);
420        out.extend_from_slice(payload);
421    }
422
423    fn build_model(tensor: &[u8]) -> Vec<u8> {
424        let mut graph = Vec::new();
425        len_delim(5, tensor, &mut graph); // GraphProto.initializer
426        let mut model = Vec::new();
427        tag(1, 0, &mut model); // ModelProto.ir_version
428        varint(8, &mut model);
429        len_delim(7, &graph, &mut model); // ModelProto.graph
430        model
431    }
432
433    #[test]
434    fn parses_f32_raw_data_tensor() {
435        let vals = [1.5f32, -2.25, 3.0];
436        let mut t = Vec::new();
437        // dims: packed [3]
438        let mut dims = Vec::new();
439        varint(3, &mut dims);
440        len_delim(1, &dims, &mut t);
441        tag(2, 0, &mut t); // data_type = FLOAT
442        varint(DT_FLOAT as u64, &mut t);
443        len_delim(8, b"w", &mut t); // name
444        let raw: Vec<u8> = vals.iter().flat_map(|v| v.to_le_bytes()).collect();
445        len_delim(9, &raw, &mut t); // raw_data
446
447        let model = build_model(&t);
448        let init = parse_model(&model).unwrap();
449        assert_eq!(init.len(), 1);
450        let rt = tensor_proto_to_raw(&init[0]).unwrap();
451        assert_eq!(rt.name, "w");
452        assert_eq!(rt.shape, vec![3]);
453        assert_eq!(rt.dtype, SrcDtype::F32);
454        let got: Vec<f32> = rt
455            .data
456            .chunks_exact(4)
457            .map(|c| f32::from_le_bytes(c.try_into().unwrap()))
458            .collect();
459        assert_eq!(got, vals);
460    }
461
462    #[test]
463    fn parses_packed_float_data_and_int64_dims_unpacked() {
464        let mut t = Vec::new();
465        tag(1, 0, &mut t); // dims unpacked: 2
466        varint(2, &mut t);
467        tag(1, 0, &mut t); // dims unpacked: 2
468        varint(2, &mut t);
469        tag(2, 0, &mut t);
470        varint(DT_FLOAT as u64, &mut t);
471        len_delim(8, b"packed", &mut t);
472        let mut fd = Vec::new();
473        for v in [0.5f32, 1.0, 1.5, 2.0] {
474            fd.extend_from_slice(&v.to_le_bytes());
475        }
476        len_delim(4, &fd, &mut t); // float_data packed
477
478        let init = parse_model(&build_model(&t)).unwrap();
479        let rt = tensor_proto_to_raw(&init[0]).unwrap();
480        assert_eq!(rt.shape, vec![2, 2]);
481        assert_eq!(rt.data.len(), 16);
482    }
483
484    #[test]
485    fn rejects_external_data() {
486        let mut t = Vec::new();
487        tag(2, 0, &mut t);
488        varint(DT_FLOAT as u64, &mut t);
489        len_delim(8, b"ext", &mut t);
490        tag(14, 0, &mut t); // data_location = EXTERNAL
491        varint(LOCATION_EXTERNAL, &mut t);
492
493        let init = parse_model(&build_model(&t)).unwrap();
494        let err = match tensor_proto_to_raw(&init[0]) {
495            Ok(_) => panic!("expected external-data error"),
496            Err(e) => e,
497        };
498        assert!(err.to_string().contains("external data"));
499    }
500}
zsfm_checkpoint/onnx.rs

zsfm_checkpoint/
onnx.rs
