hurray_core/quantization/
per_channel_affine.rs

1//! Per-channel affine quantization descriptor (scheme tag `0x02`).
2//!
3//! One `scale` and `zero_point` pair per slice along a specified axis. The scale
4//! and zero-point arrays are stored in separate buffers listed in the tensor
5//! descriptor's buffer table.
6//!
7//! See `docs/spec/quantization/per-channel-affine.md` for the normative definition.
8
9use crate::{ElementType, Error, Result};
10
11// ── Wire layout constants ─────────────────────────────────────────────────────
12
13/// Scheme tag byte for per-channel affine quantization.
14#[allow(dead_code)]
15pub(crate) const SCHEME_TAG: u8 = 0x02;
16
17/// Total descriptor length in bytes (header 4 + axis 4 + scale_buf 4 + zp_buf 4
18/// + scale_type_tag 1 + reserved 3).
19pub(crate) const ENCODED_LEN: usize = 20;
20
21/// Version this implementation supports.
22pub(crate) const SUPPORTED_VERSION: u8 = 0x01;
23
24/// Wire sentinel meaning "no zero-point buffer" (symmetric mode).
25const ZP_SENTINEL: u32 = 0xFFFF_FFFF;
26
27/// `scale_type_tag` is locked to `float32` (`0x03`) for `scheme_version = 0x01`.
28const SCALE_TYPE_TAG_FLOAT32: u8 = 0x03;
29
30/// Bit 0: symmetric flag — zero-point array is implicit (all zeros).
31pub const FLAG_SYMMETRIC: u16 = 0b1;
32/// All bits except bit 0 are reserved and must be zero.
33pub const RESERVED_MASK: u16 = !FLAG_SYMMETRIC;
34
35// Wire field offsets (relative to start of descriptor, including header).
36const OFFSET_AXIS: usize = 4;
37const OFFSET_SCALE_BUF: usize = 8;
38const OFFSET_ZP_BUF: usize = 12;
39const OFFSET_SCALE_TYPE: usize = 16;
40const OFFSET_RESERVED: usize = 17;
41
42// ── PerChannelAffine ──────────────────────────────────────────────────────────
43
44/// Quantization parameters for per-channel affine quantization.
45///
46/// Each slice along `axis` has its own `scale` (and optionally `zero_point`).
47/// The parameter arrays are stored in separate buffers referenced by index in
48/// the tensor descriptor's buffer table.
49///
50/// The dequantization formula for element `q` at logical index
51/// `[i_0, …, i_{rank-1}]` is:
52///
53/// ```text
54/// c      = i_axis
55/// x_real = scale[c] * (q - zero_point[c])
56/// ```
57///
58/// When the `SYMMETRIC` flag is set, `zero_point[c]` is treated as `0`.
59///
60/// # Wire format
61///
62/// Total descriptor length: **20 bytes** (including the 4-byte header).
63///
64/// | Offset | Field | Type |
65/// |--------|-------|------|
66/// | 4 | `axis` | `uint32` LE |
67/// | 8 | `scale_buffer_index` | `uint32` LE |
68/// | 12 | `zero_point_buffer_index` | `uint32` LE (`0xFFFFFFFF` if symmetric) |
69/// | 16 | `scale_type_tag` | `uint8` (must be `0x03` in v1) |
70/// | 17 | `_reserved` | `uint8[3]` (must be `0x00`) |
71///
72/// # Design notes
73///
74/// `PartialEq`, `Eq`, and `Hash` are all derived because this struct contains
75/// no floating-point fields — all fields are integers or `Option<u32>`.
76///
77/// `Copy` because the struct is ≤ 24 bytes with no `Drop` glue.
78///
79/// # Examples
80///
81/// ```
82/// use hurray_core::PerChannelAffine;
83///
84/// // Asymmetric: separate scale (buffer 1) and zero-point (buffer 2) arrays.
85/// let q = PerChannelAffine::new_asymmetric(0, 1, 2).unwrap();
86/// assert!(!q.is_symmetric());
87/// assert_eq!(q.zero_point_buffer_index(), Some(2));
88///
89/// // Symmetric: only a scale buffer, zero-point is implicitly zero.
90/// let s = PerChannelAffine::new_symmetric(0, 1).unwrap();
91/// assert!(s.is_symmetric());
92/// assert_eq!(s.zero_point_buffer_index(), None);
93/// ```
94// WHY Eq + Hash: no float fields, so the IEEE 754 NaN problem does not apply.
95// WHY Copy: ≤24 bytes, no Drop glue (design decision #7).
96#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
97#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
98pub struct PerChannelAffine {
99    axis: u32,
100    scale_buffer_index: u32,
101    /// `None` encodes the symmetric case (wire sentinel `0xFFFFFFFF`).
102    ///
103    /// WHY `Option<u32>`: `None` maps to wire sentinel `0xFFFFFFFF`; makes the
104    /// symmetric/asymmetric distinction unrepresentable-when-wrong at the type
105    /// level (design decision #6).
106    zero_point_buffer_index: Option<u32>,
107    scale_type: ElementType,
108}
109
110impl PerChannelAffine {
111    /// Creates an asymmetric [`PerChannelAffine`] descriptor.
112    ///
113    /// # Errors
114    ///
115    /// - [`Error::InvalidQuantization`] — `scale_buf` equals `0xFFFFFFFF`
116    ///   (the zero-point sentinel value, which is not a valid scale buffer index).
117    ///
118    /// # Examples
119    ///
120    /// ```
121    /// use hurray_core::PerChannelAffine;
122    ///
123    /// let q = PerChannelAffine::new_asymmetric(0, 1, 2).unwrap();
124    /// assert!(!q.is_symmetric());
125    /// assert_eq!(q.scale_buffer_index(), 1);
126    /// assert_eq!(q.zero_point_buffer_index(), Some(2));
127    /// ```
128    pub fn new_asymmetric(
129        axis: u32,
130        scale_buffer_index: u32,
131        zero_point_buffer_index: u32,
132    ) -> Result<Self> {
133        if scale_buffer_index == ZP_SENTINEL {
134            return Err(Error::InvalidQuantization(
135                "scale_buffer_index must not be 0xFFFFFFFF (reserved sentinel)".into(),
136            ));
137        }
138        Ok(Self {
139            axis,
140            scale_buffer_index,
141            zero_point_buffer_index: Some(zero_point_buffer_index),
142            scale_type: ElementType::Float32,
143        })
144    }
145
146    /// Creates a symmetric [`PerChannelAffine`] descriptor.
147    ///
148    /// In symmetric mode the zero-point array is implicit (all zeros); no
149    /// zero-point buffer entry is required in the tensor descriptor's buffer table.
150    ///
151    /// # Errors
152    ///
153    /// - [`Error::InvalidQuantization`] — `scale_buf` equals `0xFFFFFFFF`.
154    ///
155    /// # Examples
156    ///
157    /// ```
158    /// use hurray_core::PerChannelAffine;
159    ///
160    /// let q = PerChannelAffine::new_symmetric(1, 2).unwrap();
161    /// assert!(q.is_symmetric());
162    /// assert_eq!(q.zero_point_buffer_index(), None);
163    /// ```
164    pub fn new_symmetric(axis: u32, scale_buffer_index: u32) -> Result<Self> {
165        if scale_buffer_index == ZP_SENTINEL {
166            return Err(Error::InvalidQuantization(
167                "scale_buffer_index must not be 0xFFFFFFFF (reserved sentinel)".into(),
168            ));
169        }
170        Ok(Self {
171            axis,
172            scale_buffer_index,
173            zero_point_buffer_index: None,
174            scale_type: ElementType::Float32,
175        })
176    }
177
178    /// Returns `true` if this descriptor uses symmetric quantization (no zero point).
179    ///
180    /// # Examples
181    ///
182    /// ```
183    /// use hurray_core::PerChannelAffine;
184    ///
185    /// assert!(PerChannelAffine::new_symmetric(0, 1).unwrap().is_symmetric());
186    /// assert!(!PerChannelAffine::new_asymmetric(0, 1, 2).unwrap().is_symmetric());
187    /// ```
188    #[inline]
189    pub fn is_symmetric(&self) -> bool {
190        self.zero_point_buffer_index.is_none()
191    }
192
193    /// Returns the quantization axis index.
194    ///
195    /// # Examples
196    ///
197    /// ```
198    /// use hurray_core::PerChannelAffine;
199    ///
200    /// let q = PerChannelAffine::new_symmetric(3, 1).unwrap();
201    /// assert_eq!(q.axis(), 3);
202    /// ```
203    #[inline]
204    pub fn axis(&self) -> u32 {
205        self.axis
206    }
207
208    /// Returns the buffer table index of the per-channel scale array.
209    ///
210    /// # Examples
211    ///
212    /// ```
213    /// use hurray_core::PerChannelAffine;
214    ///
215    /// let q = PerChannelAffine::new_symmetric(0, 5).unwrap();
216    /// assert_eq!(q.scale_buffer_index(), 5);
217    /// ```
218    #[inline]
219    pub fn scale_buffer_index(&self) -> u32 {
220        self.scale_buffer_index
221    }
222
223    /// Returns the buffer table index of the per-channel zero-point array, or
224    /// `None` if this descriptor is symmetric.
225    ///
226    /// `None` maps to the wire sentinel `0xFFFFFFFF`.
227    ///
228    /// # Examples
229    ///
230    /// ```
231    /// use hurray_core::PerChannelAffine;
232    ///
233    /// let asym = PerChannelAffine::new_asymmetric(0, 1, 2).unwrap();
234    /// assert_eq!(asym.zero_point_buffer_index(), Some(2));
235    ///
236    /// let sym = PerChannelAffine::new_symmetric(0, 1).unwrap();
237    /// assert_eq!(sym.zero_point_buffer_index(), None);
238    /// ```
239    #[inline]
240    pub fn zero_point_buffer_index(&self) -> Option<u32> {
241        self.zero_point_buffer_index
242    }
243
244    /// Returns the element type used for scale values (always `Float32` in v1).
245    ///
246    /// # Examples
247    ///
248    /// ```
249    /// use hurray_core::{ElementType, PerChannelAffine};
250    ///
251    /// let q = PerChannelAffine::new_symmetric(0, 1).unwrap();
252    /// assert_eq!(q.scale_type(), ElementType::Float32);
253    /// ```
254    #[inline]
255    pub fn scale_type(&self) -> ElementType {
256        self.scale_type
257    }
258
259    /// Returns the set of storage [`ElementType`]s that are valid for this scheme.
260    ///
261    /// Per `docs/spec/quantization/per-channel-affine.md § Valid Storage Types`.
262    ///
263    /// WHY `&'static [ElementType]`: no allocation per call; `slice::contains`
264    /// over ≤10 items beats any hash structure (design decision #5).
265    ///
266    /// # Examples
267    ///
268    /// ```
269    /// use hurray_core::{ElementType, PerChannelAffine};
270    ///
271    /// assert!(PerChannelAffine::valid_storage_types().contains(&ElementType::Int8));
272    /// assert!(!PerChannelAffine::valid_storage_types().contains(&ElementType::Float32));
273    /// ```
274    pub fn valid_storage_types() -> &'static [ElementType] {
275        &[
276            ElementType::Int8,
277            ElementType::Uint8,
278            ElementType::Int16,
279            ElementType::Uint16,
280            ElementType::Int32,
281            ElementType::Uint32,
282            ElementType::Int4,
283            ElementType::Uint4,
284            ElementType::Int2,
285            ElementType::Uint2,
286        ]
287    }
288
289    // ── Crate-internal encode/decode ──────────────────────────────────────────
290
291    /// Decodes the scheme-specific payload from `bytes`.
292    ///
293    /// `bytes` is the full descriptor slice (including the 4-byte header).
294    /// `flags` are the header flags already read by the caller.
295    pub(crate) fn decode_payload(flags: u16, bytes: &[u8]) -> Result<Self> {
296        if bytes.len() < ENCODED_LEN {
297            return Err(Error::QuantizationDescriptorTooShort {
298                found: bytes.len(),
299                needed: ENCODED_LEN,
300            });
301        }
302        // Validate that no reserved flag bits are set.
303        if flags & RESERVED_MASK != 0 {
304            return Err(Error::ReservedQuantizationFlagsBits {
305                flags,
306                mask: RESERVED_MASK,
307            });
308        }
309        let symmetric = flags & FLAG_SYMMETRIC != 0;
310
311        let axis = u32::from_le_bytes(
312            bytes[OFFSET_AXIS..OFFSET_AXIS + 4]
313                .try_into()
314                .map_err(|_| Error::InvalidQuantization("axis slice error".into()))?,
315        );
316        let scale_buf = u32::from_le_bytes(
317            bytes[OFFSET_SCALE_BUF..OFFSET_SCALE_BUF + 4]
318                .try_into()
319                .map_err(|_| Error::InvalidQuantization("scale_buffer_index slice error".into()))?,
320        );
321        let zp_buf_raw =
322            u32::from_le_bytes(bytes[OFFSET_ZP_BUF..OFFSET_ZP_BUF + 4].try_into().map_err(
323                |_| Error::InvalidQuantization("zero_point_buffer_index slice error".into()),
324            )?);
325        let scale_type_byte = bytes[OFFSET_SCALE_TYPE];
326
327        // scale_buffer_index must not be the sentinel.
328        if scale_buf == ZP_SENTINEL {
329            return Err(Error::InvalidQuantization(
330                "scale_buffer_index must not be 0xFFFFFFFF".into(),
331            ));
332        }
333
334        // scale_type_tag must be 0x03 (float32) for scheme_version 0x01.
335        if scale_type_byte != SCALE_TYPE_TAG_FLOAT32 {
336            return Err(Error::InvalidQuantization(format!(
337                "per-channel affine scale_type_tag must be 0x03 (float32) in v1, got 0x{scale_type_byte:02X}"
338            )));
339        }
340
341        // Reserved bytes [17..20] must be 0x00.
342        if bytes[OFFSET_RESERVED..OFFSET_RESERVED + 3]
343            .iter()
344            .any(|&b| b != 0)
345        {
346            return Err(Error::InvalidQuantization(
347                "per-channel affine reserved bytes [17..20] must be 0x00".into(),
348            ));
349        }
350
351        // SYMMETRIC flag and zero_point_buffer_index sentinel must be consistent.
352        let zero_point_buffer_index = if symmetric {
353            if zp_buf_raw != ZP_SENTINEL {
354                return Err(Error::InvalidQuantization(
355                    "SYMMETRIC flag is set but zero_point_buffer_index is not 0xFFFFFFFF".into(),
356                ));
357            }
358            None
359        } else {
360            if zp_buf_raw == ZP_SENTINEL {
361                return Err(Error::InvalidQuantization(
362                    "SYMMETRIC flag is not set but zero_point_buffer_index is 0xFFFFFFFF".into(),
363                ));
364            }
365            Some(zp_buf_raw)
366        };
367
368        Ok(Self {
369            axis,
370            scale_buffer_index: scale_buf,
371            zero_point_buffer_index,
372            scale_type: ElementType::Float32,
373        })
374    }
375
376    /// Encodes the scheme-specific payload into `out`.
377    ///
378    /// `out` must be at least [`ENCODED_LEN`] bytes. The caller writes the 4-byte
379    /// header; this method writes bytes 4–19.
380    pub(crate) fn encode_payload(&self, out: &mut [u8]) {
381        out[OFFSET_AXIS..OFFSET_AXIS + 4].copy_from_slice(&self.axis.to_le_bytes());
382        out[OFFSET_SCALE_BUF..OFFSET_SCALE_BUF + 4]
383            .copy_from_slice(&self.scale_buffer_index.to_le_bytes());
384        let zp_wire = self.zero_point_buffer_index.unwrap_or(ZP_SENTINEL);
385        out[OFFSET_ZP_BUF..OFFSET_ZP_BUF + 4].copy_from_slice(&zp_wire.to_le_bytes());
386        out[OFFSET_SCALE_TYPE] = SCALE_TYPE_TAG_FLOAT32;
387        out[OFFSET_RESERVED..OFFSET_RESERVED + 3].fill(0);
388    }
389
390    /// Returns the flags word that encodes the symmetric/asymmetric state.
391    pub(crate) fn flags(&self) -> u16 {
392        if self.is_symmetric() {
393            FLAG_SYMMETRIC
394        } else {
395            0
396        }
397    }
398}
399
400// ── Tests ─────────────────────────────────────────────────────────────────────
401
402#[cfg(test)]
403mod tests {
404    use super::*;
405    use crate::{ElementType, Error};
406
407    // ── Constructors ──────────────────────────────────────────────────────────
408
409    #[test]
410    fn new_symmetric_sentinel_scale_buf_is_err() {
411        // scale_buffer_index == 0xFFFFFFFF is the ZP sentinel — must be rejected.
412        assert!(matches!(
413            PerChannelAffine::new_symmetric(0, 0xFFFF_FFFF),
414            Err(Error::InvalidQuantization(_))
415        ));
416    }
417
418    #[test]
419    fn new_asymmetric_sentinel_scale_buf_is_err() {
420        assert!(matches!(
421            PerChannelAffine::new_asymmetric(0, 0xFFFF_FFFF, 1),
422            Err(Error::InvalidQuantization(_))
423        ));
424    }
425
426    #[test]
427    fn new_symmetric_is_symmetric() {
428        let q = PerChannelAffine::new_symmetric(0, 1).unwrap();
429        assert!(q.is_symmetric());
430        assert_eq!(q.zero_point_buffer_index(), None);
431    }
432
433    #[test]
434    fn new_asymmetric_is_not_symmetric() {
435        let q = PerChannelAffine::new_asymmetric(0, 1, 2).unwrap();
436        assert!(!q.is_symmetric());
437        assert_eq!(q.zero_point_buffer_index(), Some(2));
438    }
439
440    #[test]
441    fn scale_type_is_always_float32() {
442        let q = PerChannelAffine::new_symmetric(0, 1).unwrap();
443        assert_eq!(q.scale_type(), ElementType::Float32);
444    }
445
446    // ── Round-trips ───────────────────────────────────────────────────────────
447
448    fn encode_decode(q: &PerChannelAffine) -> PerChannelAffine {
449        let mut buf = vec![0u8; ENCODED_LEN];
450        let flags = q.flags();
451        // Write header: scheme_tag=0x02, version=0x01, flags LE.
452        buf[0] = SCHEME_TAG;
453        buf[1] = SUPPORTED_VERSION;
454        buf[2] = (flags & 0xFF) as u8;
455        buf[3] = (flags >> 8) as u8;
456        q.encode_payload(&mut buf);
457        PerChannelAffine::decode_payload(flags, &buf).unwrap()
458    }
459
460    #[test]
461    fn round_trip_symmetric() {
462        let original = PerChannelAffine::new_symmetric(3, 5).unwrap();
463        let decoded = encode_decode(&original);
464        assert_eq!(decoded, original);
465    }
466
467    #[test]
468    fn round_trip_asymmetric() {
469        let original = PerChannelAffine::new_asymmetric(1, 2, 4).unwrap();
470        let decoded = encode_decode(&original);
471        assert_eq!(decoded, original);
472    }
473
474    // ── decode_payload error cases ────────────────────────────────────────────
475
476    #[test]
477    fn decode_payload_too_short_is_err() {
478        let short = vec![0u8; ENCODED_LEN - 1];
479        assert!(matches!(
480            PerChannelAffine::decode_payload(0, &short),
481            Err(Error::QuantizationDescriptorTooShort { .. })
482        ));
483    }
484
485    #[test]
486    fn decode_payload_reserved_flag_bits_is_err() {
487        let mut buf = vec![0u8; ENCODED_LEN];
488        buf[0] = SCHEME_TAG;
489        buf[1] = SUPPORTED_VERSION;
490        // Write a valid scale_buffer_index.
491        buf[8..12].copy_from_slice(&1u32.to_le_bytes());
492        // Write ZP sentinel (symmetric).
493        buf[12..16].copy_from_slice(&0xFFFF_FFFFu32.to_le_bytes());
494        buf[16] = 0x03; // valid scale_type_tag
495                        // flags = 0x0002 — reserved bit.
496        assert!(matches!(
497            PerChannelAffine::decode_payload(0x0002, &buf),
498            Err(Error::ReservedQuantizationFlagsBits { .. })
499        ));
500    }
501
502    #[test]
503    fn decode_payload_sentinel_scale_buf_is_err() {
504        let mut buf = vec![0u8; ENCODED_LEN];
505        buf[0] = SCHEME_TAG;
506        buf[1] = SUPPORTED_VERSION;
507        // scale_buffer_index = 0xFFFFFFFF.
508        buf[8..12].copy_from_slice(&0xFFFF_FFFFu32.to_le_bytes());
509        buf[12..16].copy_from_slice(&0xFFFF_FFFFu32.to_le_bytes());
510        buf[16] = 0x03;
511        assert!(matches!(
512            PerChannelAffine::decode_payload(FLAG_SYMMETRIC, &buf),
513            Err(Error::InvalidQuantization(_))
514        ));
515    }
516
517    #[test]
518    fn decode_payload_wrong_scale_type_tag_is_err() {
519        let mut buf = vec![0u8; ENCODED_LEN];
520        buf[0] = SCHEME_TAG;
521        buf[1] = SUPPORTED_VERSION;
522        buf[8..12].copy_from_slice(&1u32.to_le_bytes());
523        buf[12..16].copy_from_slice(&0xFFFF_FFFFu32.to_le_bytes());
524        // scale_type_tag = 0x01 (float16) — must be 0x03 in v1.
525        buf[16] = 0x01;
526        assert!(matches!(
527            PerChannelAffine::decode_payload(FLAG_SYMMETRIC, &buf),
528            Err(Error::InvalidQuantization(_))
529        ));
530    }
531
532    #[test]
533    fn decode_payload_symmetric_flag_set_but_zp_not_sentinel_is_err() {
534        let mut buf = vec![0u8; ENCODED_LEN];
535        buf[0] = SCHEME_TAG;
536        buf[1] = SUPPORTED_VERSION;
537        buf[8..12].copy_from_slice(&1u32.to_le_bytes());
538        // zp_buf = 2 (not sentinel), but SYMMETRIC flag is set — inconsistent.
539        buf[12..16].copy_from_slice(&2u32.to_le_bytes());
540        buf[16] = 0x03;
541        assert!(matches!(
542            PerChannelAffine::decode_payload(FLAG_SYMMETRIC, &buf),
543            Err(Error::InvalidQuantization(_))
544        ));
545    }
546
547    #[test]
548    fn decode_payload_symmetric_flag_not_set_but_zp_sentinel_is_err() {
549        let mut buf = vec![0u8; ENCODED_LEN];
550        buf[0] = SCHEME_TAG;
551        buf[1] = SUPPORTED_VERSION;
552        buf[8..12].copy_from_slice(&1u32.to_le_bytes());
553        // zp_buf = sentinel, but SYMMETRIC flag is NOT set — inconsistent.
554        buf[12..16].copy_from_slice(&0xFFFF_FFFFu32.to_le_bytes());
555        buf[16] = 0x03;
556        // flags = 0 (no SYMMETRIC bit).
557        assert!(matches!(
558            PerChannelAffine::decode_payload(0, &buf),
559            Err(Error::InvalidQuantization(_))
560        ));
561    }
562
563    #[test]
564    fn decode_payload_nonzero_reserved_bytes_is_err() {
565        let mut buf = vec![0u8; ENCODED_LEN];
566        buf[0] = SCHEME_TAG;
567        buf[1] = SUPPORTED_VERSION;
568        buf[8..12].copy_from_slice(&1u32.to_le_bytes());
569        buf[12..16].copy_from_slice(&0xFFFF_FFFFu32.to_le_bytes());
570        buf[16] = 0x03;
571        // Pollute reserved byte [17].
572        buf[17] = 0xAB;
573        assert!(matches!(
574            PerChannelAffine::decode_payload(FLAG_SYMMETRIC, &buf),
575            Err(Error::InvalidQuantization(_))
576        ));
577    }
578}
hurray_core/quantization/per_channel_affine.rs

hurray_core/quantization/
per_channel_affine.rs
