apr_format/v2/
tensor_index_impl.rs

1//! Tensor-index entry impl, `TensorDType`, and 64-byte alignment utilities
2//! (issue #2231). Formerly `include!`d into `v2/mod.rs`; now a real module.
3
4use super::{TensorIndexEntry, V2FormatError, ALIGNMENT, MAX_TENSOR_NAME_LEN};
5
6impl TensorIndexEntry {
7    /// Create new tensor index entry
8    #[must_use]
9    pub fn new(
10        name: impl Into<String>,
11        dtype: TensorDType,
12        shape: Vec<usize>,
13        offset: u64,
14        size: u64,
15    ) -> Self {
16        Self {
17            name: name.into(),
18            dtype,
19            shape,
20            offset,
21            size,
22        }
23    }
24
25    /// Calculate element count
26    #[must_use]
27    pub fn element_count(&self) -> usize {
28        self.shape.iter().product()
29    }
30
31    /// Serialize to bytes
32    #[must_use]
33    pub fn to_bytes(&self) -> Vec<u8> {
34        let mut buf = Vec::new();
35
36        // Name length (2 bytes) + name
37        let name_bytes = self.name.as_bytes();
38        let name_len = name_bytes.len().min(MAX_TENSOR_NAME_LEN) as u16;
39        buf.extend_from_slice(&name_len.to_le_bytes());
40        buf.extend_from_slice(&name_bytes[..name_len as usize]);
41
42        // Dtype (1 byte)
43        buf.push(self.dtype as u8);
44
45        // Shape: ndim (1 byte) + dims (8 bytes each)
46        let ndim = self.shape.len().min(8) as u8;
47        buf.push(ndim);
48        for &dim in self.shape.iter().take(8) {
49            buf.extend_from_slice(&(dim as u64).to_le_bytes());
50        }
51
52        // Offset (8 bytes)
53        buf.extend_from_slice(&self.offset.to_le_bytes());
54
55        // Size (8 bytes)
56        buf.extend_from_slice(&self.size.to_le_bytes());
57
58        buf
59    }
60
61    /// Deserialize from bytes
62    ///
63    /// # Errors
64    /// Returns error if buffer is invalid.
65    pub fn from_bytes(buf: &[u8]) -> Result<(Self, usize), V2FormatError> {
66        if buf.len() < 4 {
67            return Err(V2FormatError::InvalidTensorIndex(
68                "buffer too small".to_string(),
69            ));
70        }
71
72        let mut pos = 0;
73
74        // Name length + name
75        let name_len = u16::from_le_bytes([buf[pos], buf[pos + 1]]) as usize;
76        pos += 2;
77
78        if buf.len() < pos + name_len + 18 {
79            return Err(V2FormatError::InvalidTensorIndex(
80                "buffer too small for name".to_string(),
81            ));
82        }
83
84        let name = String::from_utf8_lossy(&buf[pos..pos + name_len]).to_string();
85        pos += name_len;
86
87        // Dtype
88        let dtype = TensorDType::from_u8(buf[pos])
89            .ok_or_else(|| V2FormatError::InvalidTensorIndex("invalid dtype".to_string()))?;
90        pos += 1;
91
92        // Shape
93        let ndim = buf[pos] as usize;
94        pos += 1;
95
96        let mut shape = Vec::with_capacity(ndim);
97        for _ in 0..ndim {
98            if buf.len() < pos + 8 {
99                return Err(V2FormatError::InvalidTensorIndex(
100                    "buffer too small for shape".to_string(),
101                ));
102            }
103            let dim = u64::from_le_bytes(buf[pos..pos + 8].try_into().unwrap_or([0; 8])) as usize;
104            shape.push(dim);
105            pos += 8;
106        }
107
108        // Offset
109        if buf.len() < pos + 16 {
110            return Err(V2FormatError::InvalidTensorIndex(
111                "buffer too small for offset/size".to_string(),
112            ));
113        }
114        let offset = u64::from_le_bytes(buf[pos..pos + 8].try_into().unwrap_or([0; 8]));
115        pos += 8;
116
117        // Size
118        let size = u64::from_le_bytes(buf[pos..pos + 8].try_into().unwrap_or([0; 8]));
119        pos += 8;
120
121        Ok((
122            Self {
123                name,
124                dtype,
125                shape,
126                offset,
127                size,
128            },
129            pos,
130        ))
131    }
132}
133
134/// Tensor data type for APR v2 format.
135///
136/// # GGML Standard Compliance (GH-438)
137///
138/// CRITICAL: IDs in the 0–31 range MUST match the GGML standard (llama.cpp ggml.h).
139/// realizar's `GgmlQuantType::from_id()` decodes these bytes directly.
140///
141/// GGML standard reference (authoritative: ggml.h `enum ggml_type`):
142///   F32=0, F16=1, Q4_0=2, Q4_1=3, Q5_0=6, Q5_1=7, Q8_0=8, Q8_1=9,
143///   Q2_K=10, Q3_K=11, Q4_K=12, Q5_K=13, Q6_K=14, Q8_K=15,
144///   IQ2_XXS=16, IQ2_XS=17, ..., BF16=30
145///
146/// APR-native quantization types (AprQ4, AprQ8) use IDs >= 128 to avoid
147/// collision with the GGML ID space. These have different block formats
148/// than any GGML type and must NOT share IDs with GGML types.
149///
150/// Legacy note: APR files written before GH-438 used IDs 8 (Q4) and 9 (Q8),
151/// which collide with GGML Q8_0=8 and Q8_1=9. `from_u8()` accepts both
152/// old and new IDs for backwards compatibility.
153#[derive(Debug, Clone, Copy, PartialEq, Eq)]
154#[repr(u8)]
155pub enum TensorDType {
156    /// 32-bit float (GGML type 0)
157    F32 = 0,
158    /// 16-bit float (GGML type 1)
159    F16 = 1,
160    /// Brain float 16 (GGML type 30)
161    BF16 = 30,
162    /// 64-bit float (APR extension, not in GGML)
163    F64 = 3,
164    /// 32-bit signed integer (APR extension, not in GGML)
165    I32 = 4,
166    /// 64-bit signed integer (APR extension, not in GGML)
167    I64 = 5,
168    /// 8-bit signed integer (APR extension, not in GGML)
169    I8 = 6,
170    /// 8-bit unsigned integer (APR extension, not in GGML)
171    U8 = 7,
172    /// APR-native 4-bit symmetric block quantization (NOT GGML Q4_0/Q4_K).
173    /// Format: per-32-block [scale: f16 (2B)] + [16 packed nibble bytes]
174    /// ID 128: outside GGML range to prevent collision.
175    /// Legacy: was ID 8 (collided with GGML Q8_0). See GH-438.
176    AprQ4 = 128,
177    /// APR-native 8-bit single-scale quantization (NOT GGML Q8_0/Q8_1).
178    /// Format: [scale: f32 (4B)] + [i8 x N] (single whole-tensor scale)
179    /// ID 129: outside GGML range to prevent collision.
180    /// Legacy: was ID 9 (collided with GGML Q8_1). See GH-438.
181    AprQ8 = 129,
182    /// GGUF Q4_K format (GGML type 12, raw super-blocks, ~4.5 bits/weight)
183    /// Format: 256-element blocks with super-block scales
184    Q4K = 12,
185    /// GGUF Q6_K format (GGML type 14, raw super-blocks, ~6.5 bits/weight)
186    Q6K = 14,
187}
188
189impl std::fmt::Display for TensorDType {
190    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
191        let name = match self {
192            Self::F32 => "F32",
193            Self::F16 => "F16",
194            Self::BF16 => "BF16",
195            Self::F64 => "F64",
196            Self::I32 => "I32",
197            Self::I64 => "I64",
198            Self::I8 => "I8",
199            Self::U8 => "U8",
200            Self::AprQ4 => "APR_Q4",
201            Self::AprQ8 => "APR_Q8",
202            Self::Q4K => "Q4_K",
203            Self::Q6K => "Q6_K",
204        };
205        f.write_str(name)
206    }
207}
208
209// ============================================================================
210// Compile-time assertions: GGML-aligned IDs must match the standard (GH-438)
211// ============================================================================
212const _: () = assert!(TensorDType::F32 as u8 == 0, "F32 must be GGML type 0");
213const _: () = assert!(TensorDType::F16 as u8 == 1, "F16 must be GGML type 1");
214const _: () = assert!(TensorDType::BF16 as u8 == 30, "BF16 must be GGML type 30");
215const _: () = assert!(
216    TensorDType::Q4K as u8 == 12,
217    "Q4K must be GGML type 12 (Q4_K)"
218);
219const _: () = assert!(
220    TensorDType::Q6K as u8 == 14,
221    "Q6K must be GGML type 14 (Q6_K)"
222);
223// APR-native types must be outside GGML range (>=128)
224const _: () = assert!(
225    TensorDType::AprQ4 as u8 >= 128,
226    "AprQ4 must be outside GGML range"
227);
228const _: () = assert!(
229    TensorDType::AprQ8 as u8 >= 128,
230    "AprQ8 must be outside GGML range"
231);
232
233impl TensorDType {
234    /// Convert from u8.
235    ///
236    /// Accepts both current IDs (128=AprQ4, 129=AprQ8) and legacy IDs
237    /// (8=AprQ4, 9=AprQ8) for backwards compatibility with pre-GH-438 APR files.
238    #[must_use]
239    pub fn from_u8(value: u8) -> Option<Self> {
240        match value {
241            0 => Some(Self::F32),
242            1 => Some(Self::F16),
243            30 => Some(Self::BF16),
244            3 => Some(Self::F64),
245            4 => Some(Self::I32),
246            5 => Some(Self::I64),
247            6 => Some(Self::I8),
248            7 => Some(Self::U8),
249            // GH-438: Legacy IDs 8/9 (collided with GGML Q8_0/Q8_1)
250            8 | 128 => Some(Self::AprQ4),
251            9 | 129 => Some(Self::AprQ8),
252            12 => Some(Self::Q4K),
253            14 => Some(Self::Q6K),
254            _ => None,
255        }
256    }
257
258    /// Get bytes per element (0 for packed types)
259    #[must_use]
260    pub const fn bytes_per_element(self) -> usize {
261        match self {
262            Self::F32 | Self::I32 => 4,
263            Self::F16 | Self::BF16 => 2,
264            Self::F64 | Self::I64 => 8,
265            Self::I8 | Self::U8 | Self::AprQ8 => 1,
266            Self::AprQ4 | Self::Q4K | Self::Q6K => 0, // Packed/block formats, need special handling
267        }
268    }
269
270    /// Get type name
271    #[must_use]
272    pub const fn name(self) -> &'static str {
273        match self {
274            Self::F32 => "f32",
275            Self::F16 => "f16",
276            Self::BF16 => "bf16",
277            Self::F64 => "f64",
278            Self::I32 => "i32",
279            Self::I64 => "i64",
280            Self::I8 => "i8",
281            Self::U8 => "u8",
282            Self::AprQ4 => "q4",
283            Self::AprQ8 => "q8",
284            Self::Q4K => "q4_k",
285            Self::Q6K => "q6_k",
286        }
287    }
288}
289
290// ============================================================================
291// Alignment Utilities
292// ============================================================================
293
294/// Align value up to the nearest multiple of alignment
295#[must_use]
296pub const fn align_up(value: usize, alignment: usize) -> usize {
297    (value + alignment - 1) & !(alignment - 1)
298}
299
300/// Align value up to 64-byte boundary
301#[must_use]
302pub const fn align_64(value: usize) -> usize {
303    align_up(value, ALIGNMENT)
304}
305
306/// Calculate padding needed to reach alignment
307#[must_use]
308pub const fn padding_to_align(value: usize, alignment: usize) -> usize {
309    let aligned = align_up(value, alignment);
310    aligned - value
311}
312
313/// Check if value is 64-byte aligned
314#[must_use]
315pub const fn is_aligned_64(value: usize) -> bool {
316    value.is_multiple_of(ALIGNMENT)
317}
318
319// The `AprV2Writer` / `AprV2StreamingWriter` / `AprV2Reader` / `AprV2ReaderRef`
320// struct declarations now live alongside their `impl` blocks (in `writer.rs` and
321// `streaming_writer.rs`) so private-field access stays module-local after the
322// include!()→mod split (issue #2231).
apr_format/v2/tensor_index_impl.rs

apr_format/v2/
tensor_index_impl.rs
