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 ([`DType::Q4_0`], [`DType::Q4_1`],
33 /// [`DType::Q5_0`], [`DType::Q5_1`], [`DType::Q8_0`]): raw on-disk block
34 /// bytes plus the tag saying how to decode them. Never indexed
35 /// elementwise directly — see [`DType::is_quantized`] and
36 /// [`crate::quant`].
37 Quantized { dtype: DType, bytes: Vec<u8> },
38}
39
40impl Storage {
41 /// The element type this storage holds.
42 pub fn dtype(&self) -> DType {
43 match self {
44 Self::F32(_) => DType::F32,
45 Self::F16(_) => DType::F16,
46 Self::BF16(_) => DType::BF16,
47 Self::I8(_) => DType::I8,
48 Self::I32(_) => DType::I32,
49 Self::Quantized { dtype, .. } => *dtype,
50 }
51 }
52
53 /// Number of logical elements this storage holds.
54 ///
55 /// For quantized storage this is derived from the byte length (every
56 /// byte belongs to some block, and every block decodes to exactly
57 /// [`DType::block_size`] elements) rather than tracked separately,
58 /// which is only sound because [`Self::new_quantized`] is the sole
59 /// constructor and it guarantees the byte length is a whole number of
60 /// blocks.
61 pub fn len(&self) -> usize {
62 match self {
63 Self::F32(v) => v.len(),
64 Self::F16(v) | Self::BF16(v) => v.len(),
65 Self::I8(v) => v.len(),
66 Self::I32(v) => v.len(),
67 Self::Quantized { dtype, bytes } => bytes.len() / dtype.block_bytes() * dtype.block_size(),
68 }
69 }
70
71 pub fn is_empty(&self) -> bool {
72 self.len() == 0
73 }
74
75 /// Builds quantized storage for `elem_count` elements of `dtype` from
76 /// raw block bytes.
77 ///
78 /// Rejects two distinct failure modes with distinct errors, both
79 /// already defined by `kopitiam-core` for exactly this purpose:
80 ///
81 /// * `elem_count` is not a whole number of blocks (e.g. 33 elements of
82 /// a 32-wide format) -> [`Error::PartialQuantizedBlock`]. This is not
83 /// a rounding question — there is no valid byte layout for a partial
84 /// block, so it is rejected before ever looking at `bytes`.
85 /// * `bytes` does not contain exactly the number of bytes that
86 /// `elem_count` blocks require -> [`Error::StorageTooSmall`]. This
87 /// catches truncated reads and mismatched shape/data pairs at
88 /// construction time instead of an out-of-bounds panic on first use.
89 pub fn new_quantized(dtype: DType, bytes: Vec<u8>, elem_count: usize) -> Result<Self> {
90 let expected = dtype
91 .storage_bytes(elem_count)
92 .ok_or(Error::PartialQuantizedBlock {
93 dtype,
94 count: elem_count,
95 block_size: dtype.block_size(),
96 })?;
97 if bytes.len() != expected {
98 return Err(Error::StorageTooSmall {
99 shape: Shape::new([elem_count]),
100 dtype,
101 expected,
102 actual: bytes.len(),
103 });
104 }
105 Ok(Self::Quantized { dtype, bytes })
106 }
107}
108
109#[cfg(test)]
110mod tests {
111 use super::*;
112
113 #[test]
114 fn dtype_and_len_match_the_variant() {
115 assert_eq!(Storage::F32(vec![1.0, 2.0, 3.0]).dtype(), DType::F32);
116 assert_eq!(Storage::F32(vec![1.0, 2.0, 3.0]).len(), 3);
117 assert_eq!(Storage::I32(vec![1, 2]).dtype(), DType::I32);
118 }
119
120 #[test]
121 fn quantized_storage_reports_decoded_element_count() {
122 // 2 blocks of Q4_0 (18 bytes each) = 64 elements.
123 let storage = Storage::new_quantized(DType::Q4_0, vec![0u8; 36], 64).unwrap();
124 assert_eq!(storage.len(), 64);
125 assert_eq!(storage.dtype(), DType::Q4_0);
126 }
127
128 #[test]
129 fn a_partial_block_element_count_is_rejected() {
130 // 33 is not a multiple of Q4_0's 32-element block.
131 let err = Storage::new_quantized(DType::Q4_0, vec![0u8; 18], 33).unwrap_err();
132 assert!(matches!(err, Error::PartialQuantizedBlock { count: 33, .. }));
133 }
134
135 #[test]
136 fn mismatched_byte_length_is_rejected() {
137 // 64 elements of Q4_0 need 36 bytes; give it 18.
138 let err = Storage::new_quantized(DType::Q4_0, vec![0u8; 18], 64).unwrap_err();
139 assert!(matches!(
140 err,
141 Error::StorageTooSmall {
142 expected: 36,
143 actual: 18,
144 ..
145 }
146 ));
147 }
148
149 #[test]
150 fn zero_elements_is_a_valid_empty_quantized_tensor() {
151 let storage = Storage::new_quantized(DType::Q8_0, vec![], 0).unwrap();
152 assert!(storage.is_empty());
153 }
154}