proxima-primitives 0.1.0

//! The one source-neutral codec identity: which packed byte layout a
//! quantized/packed buffer holds. Payload-less by design -- a pipe or fn
//! pointer cannot serialize itself to a wire byte, so a sidecar wire format
//! (`proxima-model-interop`'s expert sidecar today) needs a compile-time-
//! exhaustive, serializable tag instead. [`Codec::tag`]/[`Codec::from_tag`]
//! are that wire mapping.
//!
//! This enum is the single identity meant to replace the parallel copies
//! that grew up independently: `proxima-model-interop`'s `PackedOwnedKind`
//! (migrated onto this type), `omega`'s `PackedCodec`, and
//! `proxima-tensor`'s `QuantizedBlock` tag (both still to migrate).

use crate::block_layout::BlockLayout;
use crate::float_layout::FloatLayout;

/// Which packed byte layout a codec-tagged buffer holds.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Codec {
    Q4K,
    Q5K,
    Q6K,
    Q8_0,
    Q3K,
    Q4_0,
    Float16,
    BFloat16,
    Q2K,
    Q5_1,
    Q5_0,
    Q4_1,
    Q8_1,
    Q8K,
    Iq1S,
    Iq1M,
    Iq2Xxs,
    Iq2Xs,
    Iq2S,
    Iq3Xxs,
    Iq3S,
    Iq4Nl,
    Iq4Xs,
    Tq10,
    Tq20,
    Mxfp4,
    Nvfp4,
    Q1_0,
    Q2_0,
    BFloat8,
    Bf8E5M2,
    Bf4E2M1,
    Fp6E3M2,
}

impl Codec {
    /// The wire byte a sidecar tags this codec with. Stable across
    /// releases -- an on-disk sidecar's bytes decode by this mapping.
    #[must_use]
    pub fn tag(self) -> u8 {
        match self {
            Codec::Q2K => 0,
            Codec::Q3K => 1,
            Codec::Q4K => 2,
            Codec::Q5K => 3,
            Codec::Q6K => 4,
            Codec::Q8_0 => 5,
            Codec::Q4_0 => 6,
            Codec::Float16 => 7,
            Codec::BFloat16 => 8,
            Codec::Q5_1 => 9,
            Codec::Q5_0 => 10,
            Codec::Q4_1 => 11,
            Codec::Q8_1 => 12,
            Codec::Q8K => 13,
            Codec::Iq1S => 14,
            Codec::Iq1M => 15,
            Codec::Iq2Xxs => 16,
            Codec::Iq2Xs => 17,
            Codec::Iq2S => 18,
            Codec::Iq3Xxs => 19,
            Codec::Iq3S => 20,
            Codec::Iq4Nl => 21,
            Codec::Iq4Xs => 22,
            Codec::Tq10 => 23,
            Codec::Tq20 => 24,
            Codec::Mxfp4 => 25,
            Codec::Nvfp4 => 26,
            Codec::Q1_0 => 27,
            Codec::Q2_0 => 28,
            Codec::BFloat8 => 29,
            Codec::Bf8E5M2 => 30,
            Codec::Bf4E2M1 => 31,
            Codec::Fp6E3M2 => 32,
        }
    }

    /// The inverse of [`Self::tag`]. `None` for a byte no codec claims.
    #[must_use]
    pub fn from_tag(tag: u8) -> Option<Self> {
        match tag {
            0 => Some(Codec::Q2K),
            1 => Some(Codec::Q3K),
            2 => Some(Codec::Q4K),
            3 => Some(Codec::Q5K),
            4 => Some(Codec::Q6K),
            5 => Some(Codec::Q8_0),
            6 => Some(Codec::Q4_0),
            7 => Some(Codec::Float16),
            8 => Some(Codec::BFloat16),
            9 => Some(Codec::Q5_1),
            10 => Some(Codec::Q5_0),
            11 => Some(Codec::Q4_1),
            12 => Some(Codec::Q8_1),
            13 => Some(Codec::Q8K),
            14 => Some(Codec::Iq1S),
            15 => Some(Codec::Iq1M),
            16 => Some(Codec::Iq2Xxs),
            17 => Some(Codec::Iq2Xs),
            18 => Some(Codec::Iq2S),
            19 => Some(Codec::Iq3Xxs),
            20 => Some(Codec::Iq3S),
            21 => Some(Codec::Iq4Nl),
            22 => Some(Codec::Iq4Xs),
            23 => Some(Codec::Tq10),
            24 => Some(Codec::Tq20),
            25 => Some(Codec::Mxfp4),
            26 => Some(Codec::Nvfp4),
            27 => Some(Codec::Q1_0),
            28 => Some(Codec::Q2_0),
            29 => Some(Codec::BFloat8),
            30 => Some(Codec::Bf8E5M2),
            31 => Some(Codec::Bf4E2M1),
            32 => Some(Codec::Fp6E3M2),
            _ => None,
        }
    }

    /// Weights per [`Codec::Fp6E3M2`] block, one f16 absmax scale each.
    pub const FP6_E3M2_BLOCK_ELEMENTS: usize = 32;

    /// Bytes per [`Codec::Fp6E3M2`] block: the 6-bit [`FloatLayout::E3M2`] fields
    /// packed, then the f16 scale.
    pub const FP6_E3M2_BLOCK_BYTES: usize =
        Self::FP6_E3M2_BLOCK_ELEMENTS * FloatLayout::E3M2.width() as usize / 8 + 2;

    /// The block packing of the linear block codecs, `None` for every other codec.
    #[must_use]
    pub const fn block_layout(self) -> Option<BlockLayout> {
        match self {
            Codec::Q4_0 => Some(BlockLayout::Q4_0),
            Codec::Q4_1 => Some(BlockLayout::Q4_1),
            Codec::Q5_0 => Some(BlockLayout::Q5_0),
            Codec::Q5_1 => Some(BlockLayout::Q5_1),
            Codec::Q8_0 => Some(BlockLayout::Q8_0),
            Codec::Q8_1 => Some(BlockLayout::Q8_1),
            Codec::Q1_0 => Some(BlockLayout::Q1_0),
            Codec::Q2_0 => Some(BlockLayout::Q2_0),
            Codec::Q2K => Some(BlockLayout::Q2_K),
            Codec::Q3K => Some(BlockLayout::Q3_K),
            Codec::Q4K => Some(BlockLayout::Q4_K),
            Codec::Q5K => Some(BlockLayout::Q5_K),
            Codec::Q6K => Some(BlockLayout::Q6_K),
            Codec::Q8K => Some(BlockLayout::Q8_K),
            _ => None,
        }
    }

    /// The element layout of the float codecs, `None` for every codec that is not a plain float
    /// layout (block-quantized integer codecs, codebooks, and the formats with no decoder).
    /// `Fp6E3M2` is the layout of its six-bit fields; the f16 scale per 32 elements is its block.
    #[must_use]
    pub const fn layout(self) -> Option<FloatLayout> {
        match self {
            Codec::Float16 => Some(FloatLayout::F16),
            Codec::BFloat16 => Some(FloatLayout::BF16),
            Codec::BFloat8 => Some(FloatLayout::E5M2),
            Codec::Bf8E5M2 => Some(FloatLayout::E5M2_SATURATING),
            Codec::Bf4E2M1 => Some(FloatLayout::E2M1),
            Codec::Fp6E3M2 => Some(FloatLayout::E3M2),
            Codec::Q4K
            | Codec::Q5K
            | Codec::Q6K
            | Codec::Q8_0
            | Codec::Q3K
            | Codec::Q4_0
            | Codec::Q2K
            | Codec::Q5_1
            | Codec::Q5_0
            | Codec::Q4_1
            | Codec::Q8_1
            | Codec::Q8K
            | Codec::Iq1S
            | Codec::Iq1M
            | Codec::Iq2Xxs
            | Codec::Iq2Xs
            | Codec::Iq2S
            | Codec::Iq3Xxs
            | Codec::Iq3S
            | Codec::Iq4Nl
            | Codec::Iq4Xs
            | Codec::Tq10
            | Codec::Tq20
            | Codec::Mxfp4
            | Codec::Nvfp4
            | Codec::Q1_0
            | Codec::Q2_0 => None,
        }
    }
}

#[cfg(test)]
mod tests {
    use super::Codec;

    #[test]
    fn card_02_bf8_codec_round_trips_all_stable_tags() {
        let variants = [
            (Codec::Q2K, 0),
            (Codec::Q3K, 1),
            (Codec::Q4K, 2),
            (Codec::Q5K, 3),
            (Codec::Q6K, 4),
            (Codec::Q8_0, 5),
            (Codec::Q4_0, 6),
            (Codec::Float16, 7),
            (Codec::BFloat16, 8),
            (Codec::Q5_1, 9),
            (Codec::Q5_0, 10),
            (Codec::Q4_1, 11),
            (Codec::Q8_1, 12),
            (Codec::Q8K, 13),
            (Codec::Iq1S, 14),
            (Codec::Iq1M, 15),
            (Codec::Iq2Xxs, 16),
            (Codec::Iq2Xs, 17),
            (Codec::Iq2S, 18),
            (Codec::Iq3Xxs, 19),
            (Codec::Iq3S, 20),
            (Codec::Iq4Nl, 21),
            (Codec::Iq4Xs, 22),
            (Codec::Tq10, 23),
            (Codec::Tq20, 24),
            (Codec::Mxfp4, 25),
            (Codec::Nvfp4, 26),
            (Codec::Q1_0, 27),
            (Codec::Q2_0, 28),
            (Codec::BFloat8, 29),
            (Codec::Bf8E5M2, 30),
            (Codec::Bf4E2M1, 31),
            (Codec::Fp6E3M2, 32),
        ];
        assert_eq!(variants.len(), 33);
        for (codec, stable_tag) in variants {
            assert_eq!(codec.tag(), stable_tag);
            assert_eq!(Codec::from_tag(stable_tag), Some(codec));
        }
    }

    #[test]
    fn fp6_e3m2_codec_tag_32_round_trip() {
        assert_eq!(Codec::Fp6E3M2.tag(), 32);
        assert_eq!(Codec::from_tag(32), Some(Codec::Fp6E3M2));
    }

    #[test]
    fn fp6_e3m2_codec_tag_33_rejected() {
        assert_eq!(Codec::from_tag(33), None);
    }

    #[test]
    fn fp6_e3m2_codec_block_geometry_26_bytes_32_elements() {
        assert_eq!(Codec::FP6_E3M2_BLOCK_BYTES, 26);
        assert_eq!(Codec::FP6_E3M2_BLOCK_ELEMENTS, 32);
    }

    #[test]
    fn card_02_bf8_codec_rejects_unassigned_tag() {
        assert_eq!(Codec::from_tag(33), None);
        assert_eq!(Codec::from_tag(255), None);
    }
}