burn_pack/
base.rs

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

burn_pack/
base.rs
