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

zsfm_checkpoint/
onnx.rs
