taconite_bundle/
lib.rs

1// SPDX-FileCopyrightText: Copyright (C) 2026 Brishen Hawkins
2// SPDX-License-Identifier: Apache-2.0
3
4//! Reads the bundles IRON's exporters write for the Rust runtimes.
5//!
6//! IRON compiles ahead of time: an exporter (`iron/common/bundle.py` has the
7//! writer side) builds every NPU kernel a model needs and writes a bundle
8//! directory a Rust runtime replays:
9//!
10//! ```text
11//! manifest.txt   one record a line: `<tag> <field>...`; a field `k=v` is
12//!                also looked up by key. Blank lines and `#` lines are skipped.
13//! tensors.txt    <name> <dtype> <d0,d1,...> <offset> <bytes>, one a line,
14//! tensors.bin    into this blob (offsets 64-byte aligned; dtypes f32, bf16,
15//!                u8, i32, little-endian)
16//! kernels/       xclbins and instruction streams
17//! ```
18//!
19//! Three manifest records mean the same thing in every bundle, and
20//! [`Manifest`] interprets them:
21//!
22//! ```text
23//! version <n>                        format version; must match the runtime's
24//! param <name> <value...>            a model constant
25//! xclbin <key> <file> <kernel name>  a hardware context kernels refer to by key
26//! ```
27//!
28//! Every other record (kernels, steps, test cases) belongs to the model, and
29//! its runtime reads them from [`Manifest::records`], in file order.
30
31use std::collections::HashMap;
32use std::fmt;
33use std::fs;
34use std::io::Read;
35use std::path::{Path, PathBuf};
36use std::str::FromStr;
37
38#[derive(Debug)]
39pub enum Error {
40    /// A bundle file could not be read.
41    Io(PathBuf, std::io::Error),
42    /// A bundle file is malformed or inconsistent (the message says where).
43    Format(String),
44}
45
46impl fmt::Display for Error {
47    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
48        match self {
49            Error::Io(p, e) => write!(f, "{}: {e}", p.display()),
50            Error::Format(m) => f.write_str(m),
51        }
52    }
53}
54
55impl std::error::Error for Error {}
56
57fn bad(what: impl Into<String>) -> Error {
58    Error::Format(what.into())
59}
60
61// ----------------------------------------------------------------------------
62// manifest.txt
63// ----------------------------------------------------------------------------
64
65/// One manifest line: a tag, then whitespace-separated fields.
66#[derive(Debug, Clone)]
67pub struct Record {
68    /// 1-based line number in manifest.txt.
69    pub line: usize,
70    pub tag: String,
71    /// Every field after the tag, `k=v` ones included.
72    pub fields: Vec<String>,
73    kv: HashMap<String, String>,
74    text: String,
75}
76
77impl Record {
78    fn parse(line: usize, text: &str) -> Option<Record> {
79        let mut it = text.split_whitespace();
80        let tag = it.next().filter(|t| !t.starts_with('#'))?.to_string();
81        let fields: Vec<String> = it.map(str::to_string).collect();
82        let kv = fields.iter().filter_map(|f| f.split_once('=')).map(|(k, v)| (k.to_string(), v.to_string())).collect();
83        Some(Record { line, tag, fields, kv, text: text.to_string() })
84    }
85
86    /// The line as written.
87    pub fn text(&self) -> &str {
88        &self.text
89    }
90
91    /// An error pointing at this line.
92    pub fn error(&self, msg: impl fmt::Display) -> Error {
93        bad(format!("{}:{}: {msg}, in: {}", Manifest::FILE, self.line, self.text))
94    }
95
96    /// The `i`th field after the tag.
97    pub fn field(&self, i: usize) -> Result<&str, Error> {
98        self.fields.get(i).map(String::as_str).ok_or_else(|| self.error(format!("expected at least {} fields", i + 1)))
99    }
100
101    pub fn field_as<T: FromStr>(&self, i: usize) -> Result<T, Error> {
102        let f = self.field(i)?;
103        f.parse().map_err(|_| self.error(format!("field {} ({f}) is malformed", i + 1)))
104    }
105
106    /// The text after the first `n` fields (after the tag), spacing inside
107    /// it intact -- for a free-text last field, such as a prompt.
108    pub fn rest(&self, n: usize) -> String {
109        let mut s = self.text.as_str();
110        for _ in 0..=n {
111            s = s.trim_start();
112            s = s.find(char::is_whitespace).map_or("", |i| &s[i..]);
113        }
114        s.strip_prefix(' ').unwrap_or(s).to_string()
115    }
116
117    pub fn has(&self, key: &str) -> bool {
118        self.kv.contains_key(key)
119    }
120
121    /// The value of field `key=`.
122    pub fn str(&self, key: &str) -> Result<&str, Error> {
123        self.kv.get(key).map(String::as_str).ok_or_else(|| self.error(format!("`{key}=` missing")))
124    }
125
126    pub fn get<T: FromStr>(&self, key: &str) -> Result<T, Error> {
127        let v = self.str(key)?;
128        v.parse().map_err(|_| self.error(format!("`{key}={v}` is malformed")))
129    }
130
131    /// A `key=0` / `key=1` field.
132    pub fn flag(&self, key: &str) -> Result<bool, Error> {
133        match self.str(key)? {
134            "0" => Ok(false),
135            "1" => Ok(true),
136            v => Err(self.error(format!("`{key}={v}` is not 0 or 1"))),
137        }
138    }
139}
140
141/// A hardware context: an xclbin and the kernel name inside it.
142#[derive(Debug, Clone)]
143pub struct Xclbin {
144    pub key: String,
145    pub path: PathBuf,
146    pub kernel: String,
147}
148
149#[derive(Debug, Clone)]
150pub struct Manifest {
151    pub dir: PathBuf,
152    pub version: u32,
153    records: Vec<Record>,
154    params: HashMap<String, String>,
155    xclbins: HashMap<String, Xclbin>,
156}
157
158impl Manifest {
159    pub const FILE: &'static str = "manifest.txt";
160
161    /// Read `<dir>/manifest.txt`, which must be format `version`.
162    pub fn load(dir: &Path, version: u32) -> Result<Self, Error> {
163        let path = dir.join(Self::FILE);
164        let text = fs::read_to_string(&path).map_err(|e| Error::Io(path.clone(), e))?;
165        let mut m = Manifest {
166            dir: dir.to_path_buf(),
167            version: 0,
168            records: Vec::new(),
169            params: HashMap::new(),
170            xclbins: HashMap::new(),
171        };
172        let mut seen_version = None;
173        for (i, line) in text.lines().enumerate() {
174            let Some(r) = Record::parse(i + 1, line) else { continue };
175            match r.tag.as_str() {
176                "version" => {
177                    if seen_version.is_some() {
178                        return Err(r.error("second version record"));
179                    }
180                    seen_version = Some(r.field_as::<u32>(0)?);
181                }
182                "param" => {
183                    let k = r.field(0)?.to_string();
184                    r.field(1)?;
185                    if m.params.insert(k.clone(), r.rest(1)).is_some() {
186                        return Err(r.error(format!("param {k} defined twice")));
187                    }
188                }
189                "xclbin" => {
190                    let x = Xclbin {
191                        key: r.field(0)?.to_string(),
192                        path: dir.join(r.field(1)?),
193                        kernel: r.field(2)?.to_string(),
194                    };
195                    if m.xclbins.contains_key(&x.key) {
196                        return Err(r.error(format!("xclbin {} defined twice", x.key)));
197                    }
198                    m.xclbins.insert(x.key.clone(), x);
199                }
200                _ => m.records.push(r),
201            }
202        }
203        match seen_version {
204            Some(v) if v == version => m.version = v,
205            v => {
206                let found = v.map_or("no version record".to_string(), |v| format!("version {v}"));
207                return Err(bad(format!(
208                    "{}: {found}, this runtime reads version {version}; re-export the bundle",
209                    path.display()
210                )));
211            }
212        }
213        Ok(m)
214    }
215
216    /// Every model-specific record (not `version` / `param` / `xclbin`), in
217    /// file order.
218    pub fn records(&self) -> &[Record] {
219        &self.records
220    }
221
222    pub fn tagged<'a>(&'a self, tag: &'a str) -> impl Iterator<Item = &'a Record> {
223        self.records.iter().filter(move |r| r.tag == tag)
224    }
225
226    /// A bundle-relative path.
227    pub fn path(&self, rel: &str) -> PathBuf {
228        self.dir.join(rel)
229    }
230
231    pub fn has_param(&self, k: &str) -> bool {
232        self.params.contains_key(k)
233    }
234
235    pub fn param(&self, k: &str) -> Result<&str, Error> {
236        self.params.get(k).map(String::as_str).ok_or_else(|| bad(format!("param {k} missing from the manifest")))
237    }
238
239    pub fn param_as<T: FromStr>(&self, k: &str) -> Result<T, Error> {
240        let v = self.param(k)?;
241        v.parse().map_err(|_| bad(format!("param {k} ({v}) is malformed")))
242    }
243
244    /// A comma-separated param (empty items skipped, so `""` is `[]`).
245    pub fn list<T: FromStr>(&self, k: &str) -> Result<Vec<T>, Error> {
246        self.param(k)?
247            .split(',')
248            .filter(|s| !s.is_empty())
249            .map(|s| s.parse().map_err(|_| bad(format!("param {k}: bad list item {s}"))))
250            .collect()
251    }
252
253    pub fn xclbin(&self, key: &str) -> Result<&Xclbin, Error> {
254        self.xclbins.get(key).ok_or_else(|| bad(format!("xclbin {key} missing from the manifest")))
255    }
256
257    pub fn xclbins(&self) -> impl Iterator<Item = &Xclbin> {
258        self.xclbins.values()
259    }
260}
261
262// ----------------------------------------------------------------------------
263// tensors.txt + tensors.bin
264// ----------------------------------------------------------------------------
265
266#[derive(Debug, Clone, Copy, PartialEq, Eq)]
267pub enum DType {
268    F32,
269    Bf16,
270    U8,
271    I32,
272}
273
274impl DType {
275    pub fn size(self) -> usize {
276        match self {
277            DType::F32 | DType::I32 => 4,
278            DType::Bf16 => 2,
279            DType::U8 => 1,
280        }
281    }
282}
283
284#[derive(Debug, Clone)]
285pub struct Entry {
286    pub dtype: DType,
287    pub shape: Vec<usize>,
288    off: usize,
289    len: usize,
290}
291
292impl Entry {
293    /// Number of elements.
294    pub fn elems(&self) -> usize {
295        self.len / self.dtype.size()
296    }
297}
298
299/// Every tensor of the bundle. On Unix `tensors.bin` is memory-mapped
300/// read-only: a tensor's pages are read when it is first touched and stay
301/// reclaimable page cache, so loading a multi-GB bundle costs no heap and
302/// a runtime that uploads its weights and drops the store never holds them
303/// twice. Elsewhere (or if mapping fails) the file is read into a `u64`
304/// buffer. Either way the 64-byte-aligned offsets `tensors.txt` records are
305/// aligned in memory (mappings are page-aligned), so the typed views below
306/// are sound. Don't rewrite a bundle's `tensors.bin` while a runtime has it
307/// loaded: exporters write new directories.
308pub struct Store {
309    data: Backing,
310    map: HashMap<String, Entry>,
311}
312
313enum Backing {
314    Heap(Vec<u64>),
315    #[cfg(unix)]
316    Mapped(mmap::Map),
317}
318
319impl Backing {
320    fn as_ptr(&self) -> *const u8 {
321        match self {
322            Backing::Heap(v) => v.as_ptr() as *const u8,
323            #[cfg(unix)]
324            Backing::Mapped(m) => m.ptr,
325        }
326    }
327}
328
329#[cfg(unix)]
330mod mmap {
331    //! A read-only private file mapping, through libc's `mmap` (std has no
332    //! wrapper; the constants are the same on Linux and macOS).
333    use std::ffi::{c_int, c_void};
334    use std::os::fd::AsRawFd;
335
336    unsafe extern "C" {
337        fn mmap(addr: *mut c_void, len: usize, prot: c_int, flags: c_int, fd: c_int, off: i64) -> *mut c_void;
338        fn munmap(addr: *mut c_void, len: usize) -> c_int;
339    }
340    const PROT_READ: c_int = 1;
341    const MAP_PRIVATE: c_int = 2;
342
343    pub struct Map {
344        pub ptr: *const u8,
345        len: usize,
346    }
347
348    // SAFETY: the mapping is read-only and owned by the Map for its life.
349    unsafe impl Send for Map {}
350    unsafe impl Sync for Map {}
351
352    impl Map {
353        /// `file`'s first `len` bytes (`len > 0`), or None if mmap fails.
354        pub fn new(file: &std::fs::File, len: usize) -> Option<Map> {
355            // SAFETY: a fresh read-only mapping of a file we hold open; the
356            // result is checked against MAP_FAILED (-1).
357            let p = unsafe { mmap(std::ptr::null_mut(), len, PROT_READ, MAP_PRIVATE, file.as_raw_fd(), 0) };
358            if p as isize == -1 { None } else { Some(Map { ptr: p as *const u8, len }) }
359        }
360    }
361
362    impl Drop for Map {
363        fn drop(&mut self) {
364            // SAFETY: the mapping this Map created, unmapped once.
365            unsafe { munmap(self.ptr as *mut c_void, self.len) };
366        }
367    }
368}
369
370impl Store {
371    pub fn load(dir: &Path) -> Result<Self, Error> {
372        let idx = dir.join("tensors.txt");
373        let text = fs::read_to_string(&idx).map_err(|e| Error::Io(idx.clone(), e))?;
374        let mut map = HashMap::new();
375        for (i, line) in text.lines().enumerate() {
376            let at = |msg: &str| bad(format!("tensors.txt:{}: {msg}, in: {line}", i + 1));
377            let f: Vec<&str> = line.split_whitespace().collect();
378            if f.len() != 5 {
379                return Err(at("expected <name> <dtype> <shape> <offset> <bytes>"));
380            }
381            let dtype = match f[1] {
382                "f32" => DType::F32,
383                "bf16" => DType::Bf16,
384                "u8" => DType::U8,
385                "i32" => DType::I32,
386                _ => return Err(at("unknown dtype")),
387            };
388            let shape: Vec<usize> =
389                f[2].split(',').map(|s| s.parse().map_err(|_| at("bad shape"))).collect::<Result<_, _>>()?;
390            let off: usize = f[3].parse().map_err(|_| at("bad offset"))?;
391            let len: usize = f[4].parse().map_err(|_| at("bad byte count"))?;
392            if !off.is_multiple_of(8) {
393                return Err(at("offset is not 8-byte aligned"));
394            }
395            if !len.is_multiple_of(dtype.size()) {
396                return Err(at("byte count is not a whole number of elements"));
397            }
398            if map.insert(f[0].to_string(), Entry { dtype, shape, off, len }).is_some() {
399                return Err(at("tensor defined twice"));
400            }
401        }
402        let bin = dir.join("tensors.bin");
403        let mut file = fs::File::open(&bin).map_err(|e| Error::Io(bin.clone(), e))?;
404        let bytes = file.metadata().map_err(|e| Error::Io(bin.clone(), e))?.len() as usize;
405        #[cfg(unix)]
406        let mapped = if bytes > 0 { mmap::Map::new(&file, bytes).map(Backing::Mapped) } else { None };
407        #[cfg(not(unix))]
408        let mapped = None;
409        let data = match mapped {
410            Some(m) => m,
411            None => {
412                let mut data = vec![0u64; bytes.div_ceil(8)];
413                // SAFETY: a u64 buffer viewed as its bytes.
414                let raw = unsafe { std::slice::from_raw_parts_mut(data.as_mut_ptr() as *mut u8, bytes) };
415                file.read_exact(raw).map_err(|e| Error::Io(bin.clone(), e))?;
416                Backing::Heap(data)
417            }
418        };
419        for (name, e) in &map {
420            if e.off + e.len > bytes {
421                return Err(bad(format!("tensor {name} runs past the end of tensors.bin")));
422            }
423        }
424        Ok(Store { data, map })
425    }
426
427    pub fn has(&self, name: &str) -> bool {
428        self.map.contains_key(name)
429    }
430
431    pub fn names(&self) -> impl Iterator<Item = &str> {
432        self.map.keys().map(String::as_str)
433    }
434
435    pub fn entry(&self, name: &str) -> Result<&Entry, Error> {
436        self.map.get(name).ok_or_else(|| bad(format!("tensor {name} missing from the bundle")))
437    }
438
439    pub fn shape(&self, name: &str) -> Result<&[usize], Error> {
440        Ok(&self.entry(name)?.shape)
441    }
442
443    /// `name`'s entry, checked to be `dtype` with `elems` elements.
444    pub fn expect(&self, name: &str, dtype: DType, elems: usize) -> Result<&Entry, Error> {
445        let e = self.entry(name)?;
446        if e.dtype != dtype || e.elems() != elems {
447            return Err(bad(format!("tensor {name} is {} x {:?}, wanted {elems} x {dtype:?}", e.elems(), e.dtype)));
448        }
449        Ok(e)
450    }
451
452    fn typed<T: Copy>(&self, name: &str, want: DType) -> Result<&[T], Error> {
453        let e = self.entry(name)?;
454        if e.dtype != want {
455            return Err(bad(format!("tensor {name} is {:?}, wanted {want:?}", e.dtype)));
456        }
457        // SAFETY: in bounds (checked at load), aligned (8-byte offsets into a
458        // u64 buffer or a page-aligned mapping), and every bit pattern is a
459        // valid f32/u16/u8/i32.
460        Ok(unsafe {
461            std::slice::from_raw_parts(self.data.as_ptr().add(e.off) as *const T, e.len / std::mem::size_of::<T>())
462        })
463    }
464
465    /// Any tensor's raw bytes.
466    pub fn bytes(&self, name: &str) -> Result<&[u8], Error> {
467        let e = self.entry(name)?;
468        // SAFETY: in bounds (checked at load).
469        Ok(unsafe { std::slice::from_raw_parts(self.data.as_ptr().add(e.off), e.len) })
470    }
471
472    pub fn f32(&self, name: &str) -> Result<&[f32], Error> {
473        self.typed(name, DType::F32)
474    }
475
476    /// bf16 as raw bits.
477    pub fn bf16(&self, name: &str) -> Result<&[u16], Error> {
478        self.typed(name, DType::Bf16)
479    }
480
481    pub fn u8(&self, name: &str) -> Result<&[u8], Error> {
482        self.typed(name, DType::U8)
483    }
484
485    pub fn i32(&self, name: &str) -> Result<&[i32], Error> {
486        self.typed(name, DType::I32)
487    }
488}
489
490#[cfg(test)]
491mod tests {
492    use super::*;
493    use std::sync::atomic::{AtomicUsize, Ordering};
494
495    fn tmpdir() -> PathBuf {
496        static N: AtomicUsize = AtomicUsize::new(0);
497        let d = std::env::temp_dir().join(format!(
498            "taconite-bundle-test-{}-{}",
499            std::process::id(),
500            N.fetch_add(1, Ordering::Relaxed)
501        ));
502        fs::create_dir_all(&d).unwrap();
503        d
504    }
505
506    fn manifest(text: &str) -> Result<Manifest, Error> {
507        let d = tmpdir();
508        fs::write(d.join("manifest.txt"), text).unwrap();
509        Manifest::load(&d, 3)
510    }
511
512    #[test]
513    fn records_params_xclbins() {
514        let m = manifest(
515            "version 3\n\
516             # a comment\n\
517             \n\
518             param grid 72\n\
519             param splits 1,2,,3\n\
520             param name IR 18\n\
521             xclbin ctx0 kernels/a.xclbin MLIR_AIE\n\
522             gemm g0 ctx=ctx0 M=256 bias=1\n\
523             case 0 cases/cats.jpg Two  dogs, playing!\n",
524        )
525        .unwrap();
526        assert_eq!(m.param_as::<usize>("grid").unwrap(), 72);
527        assert_eq!(m.list::<usize>("splits").unwrap(), [1, 2, 3]);
528        assert_eq!(m.param("name").unwrap(), "IR 18");
529        let x = m.xclbin("ctx0").unwrap();
530        assert_eq!((x.path.ends_with("kernels/a.xclbin"), x.kernel.as_str()), (true, "MLIR_AIE"));
531        assert_eq!(m.records().len(), 2);
532        let g = m.tagged("gemm").next().unwrap();
533        assert_eq!((g.line, g.field(0).unwrap(), g.str("ctx").unwrap()), (8, "g0", "ctx0"));
534        assert_eq!(g.get::<usize>("M").unwrap(), 256);
535        assert!(g.flag("bias").unwrap());
536        assert!(g.get::<usize>("N").unwrap_err().to_string().starts_with("manifest.txt:8: `N=` missing"));
537        let c = m.tagged("case").next().unwrap();
538        assert_eq!(c.rest(2), "Two  dogs, playing!");
539        assert_eq!(c.rest(0), "0 cases/cats.jpg Two  dogs, playing!");
540    }
541
542    #[test]
543    fn version_is_checked() {
544        let e = manifest("version 2\n").unwrap_err().to_string();
545        assert!(e.contains("version 2, this runtime reads version 3"), "{e}");
546        let e = manifest("param a 1\n").unwrap_err().to_string();
547        assert!(e.contains("no version record"), "{e}");
548        assert!(manifest("version 3\nversion 3\n").is_err());
549        assert!(manifest("version 3\nparam a 1\nparam a 2\n").is_err());
550        assert!(manifest("version 3\nxclbin k f\n").is_err());
551    }
552
553    fn write_store(d: &Path, idx: &str, bin: &[u8]) -> Result<Store, Error> {
554        fs::write(d.join("tensors.txt"), idx).unwrap();
555        fs::write(d.join("tensors.bin"), bin).unwrap();
556        Store::load(d)
557    }
558
559    #[test]
560    fn store() {
561        let d = tmpdir();
562        let mut bin = vec![0u8; 72];
563        bin[..8].copy_from_slice(&[0, 0, 128, 63, 0, 0, 0, 64]); // f32 1.0, 2.0
564        bin[64..68].copy_from_slice(&[0x80, 0x3f, 0x00, 0x40]); // bf16 1.0, 2.0
565        let s = write_store(&d, "a f32 2 0 8\nb bf16 1,2 64 4\n", &bin).unwrap();
566        assert_eq!(s.f32("a").unwrap(), [1.0, 2.0]);
567        assert_eq!(s.bf16("b").unwrap(), [0x3f80, 0x4000]);
568        assert_eq!(s.shape("b").unwrap(), [1, 2]);
569        assert!(s.f32("b").is_err());
570        assert!(s.expect("b", DType::Bf16, 2).is_ok());
571        assert!(s.expect("b", DType::Bf16, 3).is_err());
572        assert!(s.entry("c").is_err());
573        assert!(write_store(&d, "a f32 2 0 8\nb bf16 1,6 64 12\n", &bin).is_err()); // past the end
574        assert!(write_store(&d, "a f32 2 4 8\n", &bin).is_err()); // misaligned
575        assert!(write_store(&d, "a f32 2 0 6\n", &bin).is_err()); // partial element
576        assert!(write_store(&d, "a f32 2 0 8\na f32 2 0 8\n", &bin).is_err()); // twice
577    }
578}
taconite_bundle/lib.rs

taconite_bundle/
lib.rs
