burn_pack/
base.rs

1//! Core types and constants for the Burnpack file format.
2//!
3//! See the [crate root](crate) for the complete file format specification.
4
5use alloc::collections::BTreeMap;
6use alloc::format;
7use alloc::string::String;
8use alloc::vec::Vec;
9use burn_std::DType;
10use byteorder::{ByteOrder, LittleEndian};
11use serde::{Deserialize, Serialize};
12
13/// Magic number identifying a Burnpack file: "BURN" in ASCII (0x4255524E)
14/// When written to file in little-endian format, appears as "NRUB" bytes
15pub const MAGIC_NUMBER: u32 = 0x4255524E;
16
17/// Current format version
18pub const FORMAT_VERSION: u16 = 0x0001;
19
20/// Size of the magic number in bytes
21pub const MAGIC_SIZE: usize = 4;
22
23/// Size of the format version in bytes
24pub const VERSION_SIZE: usize = 2;
25
26/// Size of the metadata size field in bytes
27pub const METADATA_SIZE_FIELD_SIZE: usize = 4;
28
29/// Total header size (computed from components)
30pub const HEADER_SIZE: usize = MAGIC_SIZE + VERSION_SIZE + METADATA_SIZE_FIELD_SIZE;
31
32/// Alignment for tensor data in bytes.
33///
34/// All tensor data is aligned to 256-byte boundaries to enable efficient
35/// memory-mapped (mmap) zero-copy loading. This alignment ensures:
36/// - Proper pointer alignment for all tensor element types (f64 requires 8-byte alignment)
37/// - Cache-line friendly access (most CPUs use 64-byte cache lines)
38/// - GPU memory alignment (CUDA prefers 256-byte for coalesced access)
39/// - Future-proofing for wider SIMD (AVX-512 = 64 bytes, future AVX-1024 = 128 bytes)
40///
41/// Industry alignment choices:
42/// - 256-byte: GGUF, MLX, ncnn, MNN, TNN, vLLM-AWQ, Marlin (15+ formats)
43/// - 64-byte: SafeTensors (minimum for AVX-512)
44/// - 4096-byte: Core ML
45///
46/// 256-byte alignment has negligible overhead for typical tensor sizes while
47/// providing maximum compatibility with current and future hardware.
48pub const TENSOR_ALIGNMENT: u64 = 256;
49
50/// Calculate the byte offset where the tensor data section starts.
51///
52/// The data section is padded to start at a 256-byte aligned position
53/// so that all tensor offsets (which are relative to data section) result
54/// in properly aligned absolute file positions for mmap zero-copy access.
55///
56/// This function must be used consistently by both writer and reader.
57#[inline]
58pub fn aligned_data_section_start(metadata_size: usize) -> usize {
59    let unaligned_start = (HEADER_SIZE + metadata_size) as u64;
60    // Keep multiplication in u64 space to avoid overflow on 32-bit systems
61    (unaligned_start.div_ceil(TENSOR_ALIGNMENT) * TENSOR_ALIGNMENT) as usize
62}
63
64// Security limits to prevent DoS attacks via resource exhaustion
65// These can be adjusted based on your use case
66
67/// Maximum allowed metadata size (100 MB)
68/// Prevents memory exhaustion attacks via oversized metadata claims
69pub const MAX_METADATA_SIZE: u32 = 100 * 1024 * 1024;
70
71/// Maximum allowed tensor size per tensor
72/// Prevents memory exhaustion attacks via oversized tensor claims
73/// 32-bit platforms: 2 GB limit (to fit within usize range)
74/// 64-bit platforms: 10 GB limit
75#[cfg(target_pointer_width = "32")]
76pub const MAX_TENSOR_SIZE: usize = 2 * 1024 * 1024 * 1024;
77#[cfg(not(target_pointer_width = "32"))]
78pub const MAX_TENSOR_SIZE: usize = 10 * 1024 * 1024 * 1024;
79
80/// Maximum allowed number of tensors (100,000)
81/// Prevents resource exhaustion via excessive tensor counts
82pub const MAX_TENSOR_COUNT: usize = 100_000;
83
84/// Maximum CBOR deserialization recursion depth (128 levels)
85/// Prevents stack overflow attacks via deeply nested CBOR structures
86pub const MAX_CBOR_RECURSION_DEPTH: usize = 128;
87
88/// Maximum allowed file size (100 GB)
89/// Prevents resource exhaustion from extremely large files
90/// This limit applies to file-based loading (mmap and buffered)
91#[cfg(feature = "std")]
92pub const MAX_FILE_SIZE: u64 = 100 * 1024 * 1024 * 1024;
93
94/// Byte range for magic number in header
95pub const fn magic_range() -> core::ops::Range<usize> {
96    let start = 0;
97    let end = start + MAGIC_SIZE;
98    start..end
99}
100
101/// Byte range for format version in header
102pub const fn version_range() -> core::ops::Range<usize> {
103    let start = MAGIC_SIZE;
104    let end = start + VERSION_SIZE;
105    start..end
106}
107
108/// Byte range for metadata size field in header
109pub const fn metadata_size_range() -> core::ops::Range<usize> {
110    let start = MAGIC_SIZE + VERSION_SIZE;
111    let end = start + METADATA_SIZE_FIELD_SIZE;
112    start..end
113}
114
115// Compile-time validation that ranges are correct
116const _: () = assert!(MAGIC_SIZE + VERSION_SIZE + METADATA_SIZE_FIELD_SIZE == HEADER_SIZE);
117
118/// Header structure for Burnpack files
119#[derive(Debug, Clone, Copy)]
120pub struct Header {
121    /// Magic number (4 bytes): 0x4255524E ("BURN")
122    pub magic: u32,
123    /// Format version (2 bytes)
124    pub version: u16,
125    /// Size of CBOR metadata in bytes (4 bytes)
126    pub metadata_size: u32,
127}
128
129impl Header {
130    /// Create a new header with the given metadata size
131    pub fn new(metadata_size: u32) -> Self {
132        Self {
133            magic: MAGIC_NUMBER,
134            version: FORMAT_VERSION,
135            metadata_size,
136        }
137    }
138
139    /// Serialize header into bytes
140    pub fn into_bytes(self) -> [u8; HEADER_SIZE] {
141        let mut bytes = [0u8; HEADER_SIZE];
142        LittleEndian::write_u32(&mut bytes[magic_range()], self.magic);
143        LittleEndian::write_u16(&mut bytes[version_range()], self.version);
144        LittleEndian::write_u32(&mut bytes[metadata_size_range()], self.metadata_size);
145        bytes
146    }
147
148    /// Deserialize header from bytes
149    pub fn from_bytes(bytes: &[u8]) -> Result<Self, Error> {
150        if bytes.len() < HEADER_SIZE {
151            return Err(Error::InvalidHeader);
152        }
153
154        let magic = LittleEndian::read_u32(&bytes[magic_range()]);
155        if magic != MAGIC_NUMBER {
156            return Err(Error::InvalidMagicNumber);
157        }
158
159        let version = LittleEndian::read_u16(&bytes[version_range()]);
160        let metadata_size = LittleEndian::read_u32(&bytes[metadata_size_range()]);
161
162        Ok(Self {
163            magic,
164            version,
165            metadata_size,
166        })
167    }
168}
169
170/// A typed scalar value stored alongside tensors in a burnpack container.
171///
172/// Scalars are kept in the CBOR metadata section (not the tensor data section), so they carry
173/// no alignment cost. The field is optional in the format: files written before scalar support
174/// simply omit it, and readers default it to empty.
175///
176/// Convert to/from the primitive numeric and boolean types with [`From`] / [`TryFrom`]
177/// (e.g. `Scalar::from(3usize)`, `u32::try_from(scalar)`).
178#[derive(Debug, Clone, Copy, PartialEq, Serialize, Deserialize)]
179pub enum Scalar {
180    /// A signed integer.
181    Int(i64),
182    /// An unsigned integer.
183    UInt(u64),
184    /// A floating-point number.
185    Float(f64),
186    /// A boolean.
187    Bool(bool),
188}
189
190/// Error returned when a [`Scalar`] cannot be converted to a requested primitive type
191/// (wrong variant or out of range).
192#[derive(Debug, Clone, Copy, PartialEq, Eq)]
193pub struct ScalarConversionError;
194
195impl core::fmt::Display for ScalarConversionError {
196    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
197        write!(f, "scalar value does not fit the requested type")
198    }
199}
200
201impl core::error::Error for ScalarConversionError {}
202
203macro_rules! impl_scalar_int {
204    ($($t:ty => $variant:ident),* $(,)?) => {
205        $(
206            impl From<$t> for Scalar {
207                fn from(value: $t) -> Self {
208                    Scalar::$variant(value as _)
209                }
210            }
211
212            impl TryFrom<Scalar> for $t {
213                type Error = ScalarConversionError;
214                fn try_from(scalar: Scalar) -> Result<Self, Self::Error> {
215                    match scalar {
216                        Scalar::Int(v) => v.try_into().map_err(|_| ScalarConversionError),
217                        Scalar::UInt(v) => v.try_into().map_err(|_| ScalarConversionError),
218                        _ => Err(ScalarConversionError),
219                    }
220                }
221            }
222        )*
223    };
224}
225
226impl_scalar_int!(
227    i8 => Int, i16 => Int, i32 => Int, i64 => Int, isize => Int,
228    u8 => UInt, u16 => UInt, u32 => UInt, u64 => UInt, usize => UInt,
229);
230
231impl From<f64> for Scalar {
232    fn from(value: f64) -> Self {
233        Scalar::Float(value)
234    }
235}
236
237impl From<f32> for Scalar {
238    fn from(value: f32) -> Self {
239        Scalar::Float(value as f64)
240    }
241}
242
243impl From<bool> for Scalar {
244    fn from(value: bool) -> Self {
245        Scalar::Bool(value)
246    }
247}
248
249impl TryFrom<Scalar> for f64 {
250    type Error = ScalarConversionError;
251    fn try_from(scalar: Scalar) -> Result<Self, Self::Error> {
252        match scalar {
253            Scalar::Float(v) => Ok(v),
254            Scalar::Int(v) => Ok(v as f64),
255            Scalar::UInt(v) => Ok(v as f64),
256            _ => Err(ScalarConversionError),
257        }
258    }
259}
260
261impl TryFrom<Scalar> for f32 {
262    type Error = ScalarConversionError;
263    fn try_from(scalar: Scalar) -> Result<Self, Self::Error> {
264        // Mirror `f64`'s acceptance of integer variants; float reads may be lossy (documented).
265        match scalar {
266            Scalar::Float(v) => Ok(v as f32),
267            Scalar::Int(v) => Ok(v as f32),
268            Scalar::UInt(v) => Ok(v as f32),
269            _ => Err(ScalarConversionError),
270        }
271    }
272}
273
274impl TryFrom<Scalar> for bool {
275    type Error = ScalarConversionError;
276    fn try_from(scalar: Scalar) -> Result<Self, Self::Error> {
277        match scalar {
278            Scalar::Bool(v) => Ok(v),
279            _ => Err(ScalarConversionError),
280        }
281    }
282}
283
284/// Metadata structure serialized with CBOR
285#[derive(Debug, Clone, Serialize, Deserialize)]
286pub(crate) struct Metadata {
287    /// Tensor descriptors mapped by name for efficient lookup
288    pub tensors: BTreeMap<String, TensorDescriptor>,
289    /// Optional additional metadata
290    #[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
291    pub metadata: BTreeMap<String, String>,
292    /// Optional typed scalars mapped by name.
293    ///
294    /// Defaulted on read for backward compatibility with files written before scalar support.
295    #[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
296    pub scalars: BTreeMap<String, Scalar>,
297}
298
299/// Individual tensor descriptor
300#[derive(Debug, Clone, Serialize, Deserialize)]
301pub(crate) struct TensorDescriptor {
302    /// Data type of the tensor
303    pub dtype: DType,
304    /// Tensor shape dimensions
305    pub shape: Vec<u64>,
306    /// Byte offsets in data section (start, end)
307    pub data_offsets: (u64, u64),
308    /// Parameter ID for training state persistence matching.
309    /// Generated automatically if not present during loading.
310    #[serde(default, skip_serializing_if = "Option::is_none")]
311    pub param_id: Option<u64>,
312}
313
314/// Validate that a tensor's byte range agrees with its declared shape and dtype.
315///
316/// Quantized tensors are exempt because their packed values and inline scales don't follow the
317/// ordinary `num_elements * dtype_size` layout.
318pub(crate) fn validate_tensor_byte_len(
319    name: &str,
320    dtype: DType,
321    shape: &[u64],
322    data_len: u64,
323) -> Result<(), Error> {
324    if matches!(dtype, DType::QFloat(_)) {
325        return Ok(());
326    }
327
328    let num_elements = shape
329        .iter()
330        .try_fold(1u64, |total, &dimension| total.checked_mul(dimension));
331    let expected = num_elements
332        .and_then(|num_elements| num_elements.checked_mul(dtype.size() as u64))
333        .ok_or_else(|| {
334            Error::ValidationError(format!(
335                "tensor '{name}' byte length calculation overflows u64 for shape {shape:?} and dtype {dtype:?}"
336            ))
337        })?;
338
339    if data_len != expected {
340        return Err(Error::ValidationError(format!(
341            "tensor '{name}' declares {data_len} bytes but its shape {shape:?} and dtype {dtype:?} need {expected}"
342        )));
343    }
344
345    Ok(())
346}
347
348/// Error types for Burnpack operations
349#[derive(Debug)]
350pub enum Error {
351    InvalidHeader,
352    InvalidMagicNumber,
353    InvalidVersion,
354    MetadataSerializationError(String),
355    MetadataDeserializationError(String),
356    IoError(String),
357    TensorNotFound(String),
358    TensorBytesSizeMismatch(String),
359    ValidationError(String),
360    /// A write that must not replace an existing file found one at this path.
361    AlreadyExists(String),
362}
363
364impl core::fmt::Display for Error {
365    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
366        match self {
367            Error::InvalidHeader => write!(f, "Invalid header: insufficient bytes"),
368            Error::InvalidMagicNumber => write!(f, "Invalid magic number"),
369            Error::InvalidVersion => write!(f, "Unsupported version"),
370            Error::MetadataSerializationError(e) => {
371                write!(f, "Metadata serialization error: {}", e)
372            }
373            Error::MetadataDeserializationError(e) => {
374                write!(f, "Metadata deserialization error: {}", e)
375            }
376            Error::IoError(e) => write!(f, "I/O error: {}", e),
377            Error::TensorNotFound(name) => write!(f, "Tensor not found: {}", name),
378            Error::TensorBytesSizeMismatch(e) => {
379                write!(f, "Tensor bytes size mismatch: {}", e)
380            }
381            Error::ValidationError(e) => write!(f, "Validation error: {}", e),
382            Error::AlreadyExists(path) => write!(
383                f,
384                "File already exists: {}. To replace it, use .overwrite(true), or \
385                 AtomicFile::commit instead of commit_new.",
386                path
387            ),
388        }
389    }
390}
391
392impl Error {
393    /// Name the tensor an error came from, keeping the variant intact.
394    ///
395    /// Used on the write path, where a tensor's bytes are produced by caller code partway
396    /// through the container: without the name, a failure on one tensor of many says only
397    /// that something went wrong somewhere.
398    pub(crate) fn in_tensor(mut self, name: &str) -> Self {
399        match &mut self {
400            Error::MetadataSerializationError(message)
401            | Error::MetadataDeserializationError(message)
402            | Error::IoError(message)
403            | Error::TensorBytesSizeMismatch(message)
404            | Error::ValidationError(message) => *message = format!("tensor '{name}': {message}"),
405            // Header failures carry no message to annotate (and are not expected from an
406            // entry), while `TensorNotFound` and `AlreadyExists` carry a name or path, not a
407            // sentence - prefixing either would garble its Display output.
408            Error::InvalidHeader
409            | Error::InvalidMagicNumber
410            | Error::InvalidVersion
411            | Error::TensorNotFound(_)
412            | Error::AlreadyExists(_) => {}
413        }
414        self
415    }
416}
417
418impl core::error::Error for Error {}
419
420#[cfg(test)]
421mod scalar_tests {
422    use super::*;
423
424    #[test]
425    fn int_round_trips_through_checked_conversion() {
426        assert_eq!(i32::try_from(Scalar::from(-5i32)).unwrap(), -5);
427        assert_eq!(u8::try_from(Scalar::from(200u8)).unwrap(), 200);
428        assert_eq!(usize::try_from(Scalar::from(42usize)).unwrap(), 42);
429    }
430
431    #[test]
432    fn out_of_range_int_conversion_is_rejected() {
433        // u64 value beyond i32::MAX cannot become i32.
434        let big = Scalar::from(5_000_000_000u64);
435        assert!(i32::try_from(big).is_err());
436        // Negative value cannot become u32.
437        assert!(u32::try_from(Scalar::from(-1i32)).is_err());
438        // 300 does not fit in u8.
439        assert!(u8::try_from(Scalar::from(300u32)).is_err());
440    }
441
442    #[test]
443    fn float_and_bool_variant_mismatches_are_rejected() {
444        // An integer field must not read a stored float.
445        assert!(i64::try_from(Scalar::Float(1.5)).is_err());
446        // A bool field must not read a stored int.
447        assert!(bool::try_from(Scalar::Int(1)).is_err());
448        // A float field must not read a stored bool.
449        assert!(f64::try_from(Scalar::Bool(true)).is_err());
450    }
451
452    #[test]
453    fn float_accepts_int_variants_symmetrically() {
454        assert_eq!(f64::try_from(Scalar::Int(3)).unwrap(), 3.0);
455        assert_eq!(f32::try_from(Scalar::Int(3)).unwrap(), 3.0);
456        assert_eq!(f64::try_from(Scalar::Float(2.5)).unwrap(), 2.5);
457        assert_eq!(f32::try_from(Scalar::Float(2.5)).unwrap(), 2.5);
458    }
459}
burn_pack/base.rs

burn_pack/
base.rs
