kopitiam_tensor/storage.rs
1//! Owned CPU element buffers.
2//!
3//! `Storage` is the layer below [`crate::Tensor`] that actually owns bytes.
4//! It knows its [`DType`] and how many elements it holds, but nothing about
5//! shape, strides, or views — that is `Tensor`'s job, layered on top so that
6//! multiple tensors (a reshape, a transpose, a narrowed slice) can share one
7//! `Storage` via `Arc` without copying.
8//!
9//! One variant per non-quantized [`DType`], plus a single [`Storage::Quantized`]
10//! variant for every block-quantized format. Quantized formats share a
11//! variant — rather than getting one each — because none of them can be
12//! addressed as a `Vec<T>` of individually meaningful elements; they are all
13//! "a scale (and maybe a min) plus packed sub-byte weights", so the only
14//! thing that differs between them is how [`crate::quant`] decodes a block,
15//! not how the bytes are stored.
16
17use kopitiam_core::{DType, Error, Result, Shape};
18
19/// An owned CPU buffer of tensor elements.
20#[derive(Debug, Clone)]
21pub enum Storage {
22 F32(Vec<f32>),
23 /// Raw IEEE 754 binary16 bits. Stored as `u16` rather than a float type
24 /// because Rust's `f32`/`f64` are the only floats with hardware and
25 /// library support; every element must go through
26 /// [`crate::half::f16_to_f32`] to be computed on.
27 F16(Vec<u16>),
28 /// Raw bfloat16 bits, likewise `u16`; see [`crate::half::bf16_to_f32`].
29 BF16(Vec<u16>),
30 I8(Vec<i8>),
31 I32(Vec<i32>),
32 /// Any block-quantized format — the classic 32-element blocks
33 /// ([`DType::Q4_0`], [`DType::Q4_1`], [`DType::Q5_0`], [`DType::Q5_1`],
34 /// [`DType::Q8_0`]) and the 256-element K-quant super-blocks
35 /// ([`DType::Q4_K`], [`DType::Q6_K`], ...): raw on-disk block bytes plus
36 /// the tag saying how to decode them. Never indexed elementwise directly
37 /// — see [`DType::is_quantized`] and [`crate::quant`].
38 Quantized { dtype: DType, bytes: Vec<u8> },
39}
40
41impl Storage {
42 /// The element type this storage holds.
43 pub fn dtype(&self) -> DType {
44 match self {
45 Self::F32(_) => DType::F32,
46 Self::F16(_) => DType::F16,
47 Self::BF16(_) => DType::BF16,
48 Self::I8(_) => DType::I8,
49 Self::I32(_) => DType::I32,
50 Self::Quantized { dtype, .. } => *dtype,
51 }
52 }
53
54 /// Number of logical elements this storage holds.
55 ///
56 /// For quantized storage this is derived from the byte length (every
57 /// byte belongs to some block, and every block decodes to exactly
58 /// [`DType::block_size`] elements) rather than tracked separately,
59 /// which is only sound because [`Self::new_quantized`] is the sole
60 /// constructor and it guarantees the byte length is a whole number of
61 /// blocks.
62 pub fn len(&self) -> usize {
63 match self {
64 Self::F32(v) => v.len(),
65 Self::F16(v) | Self::BF16(v) => v.len(),
66 Self::I8(v) => v.len(),
67 Self::I32(v) => v.len(),
68 Self::Quantized { dtype, bytes } => bytes.len() / dtype.block_bytes() * dtype.block_size(),
69 }
70 }
71
72 pub fn is_empty(&self) -> bool {
73 self.len() == 0
74 }
75
76 /// Builds quantized storage for `elem_count` elements of `dtype` from
77 /// raw block bytes.
78 ///
79 /// Rejects two distinct failure modes with distinct errors, both
80 /// already defined by `kopitiam-core` for exactly this purpose:
81 ///
82 /// * `elem_count` is not a whole number of blocks (e.g. 33 elements of
83 /// a 32-wide format) -> [`Error::PartialQuantizedBlock`]. This is not
84 /// a rounding question — there is no valid byte layout for a partial
85 /// block, so it is rejected before ever looking at `bytes`.
86 /// * `bytes` does not contain exactly the number of bytes that
87 /// `elem_count` blocks require -> [`Error::StorageTooSmall`]. This
88 /// catches truncated reads and mismatched shape/data pairs at
89 /// construction time instead of an out-of-bounds panic on first use.
90 pub fn new_quantized(dtype: DType, bytes: Vec<u8>, elem_count: usize) -> Result<Self> {
91 let expected = dtype
92 .storage_bytes(elem_count)
93 .ok_or(Error::PartialQuantizedBlock {
94 dtype,
95 count: elem_count,
96 block_size: dtype.block_size(),
97 })?;
98 if bytes.len() != expected {
99 return Err(Error::StorageTooSmall {
100 shape: Shape::new([elem_count]),
101 dtype,
102 expected,
103 actual: bytes.len(),
104 });
105 }
106 Ok(Self::Quantized { dtype, bytes })
107 }
108}
109
110#[cfg(test)]
111mod tests {
112 use super::*;
113
114 #[test]
115 fn dtype_and_len_match_the_variant() {
116 assert_eq!(Storage::F32(vec![1.0, 2.0, 3.0]).dtype(), DType::F32);
117 assert_eq!(Storage::F32(vec![1.0, 2.0, 3.0]).len(), 3);
118 assert_eq!(Storage::I32(vec![1, 2]).dtype(), DType::I32);
119 }
120
121 #[test]
122 fn quantized_storage_reports_decoded_element_count() {
123 // 2 blocks of Q4_0 (18 bytes each) = 64 elements.
124 let storage = Storage::new_quantized(DType::Q4_0, vec![0u8; 36], 64).unwrap();
125 assert_eq!(storage.len(), 64);
126 assert_eq!(storage.dtype(), DType::Q4_0);
127 }
128
129 #[test]
130 fn a_partial_block_element_count_is_rejected() {
131 // 33 is not a multiple of Q4_0's 32-element block.
132 let err = Storage::new_quantized(DType::Q4_0, vec![0u8; 18], 33).unwrap_err();
133 assert!(matches!(err, Error::PartialQuantizedBlock { count: 33, .. }));
134 }
135
136 #[test]
137 fn mismatched_byte_length_is_rejected() {
138 // 64 elements of Q4_0 need 36 bytes; give it 18.
139 let err = Storage::new_quantized(DType::Q4_0, vec![0u8; 18], 64).unwrap_err();
140 assert!(matches!(
141 err,
142 Error::StorageTooSmall {
143 expected: 36,
144 actual: 18,
145 ..
146 }
147 ));
148 }
149
150 #[test]
151 fn zero_elements_is_a_valid_empty_quantized_tensor() {
152 let storage = Storage::new_quantized(DType::Q8_0, vec![], 0).unwrap();
153 assert!(storage.is_empty());
154 }
155}