cubecl-tile 0.12.0-pre.1

//! How an operand's values sit in memory.

use cubecl::client::Client;
use cubecl::e2m1x2;
use cubecl::ir::features::TypeUsage;
use cubecl::ir::types::Fp8Format;
use cubecl::ir::{ElemType, FloatKind, IntKind};
use cubecl::post_processing::fp4::E2M1_F16_LIFT;
use cubecl::post_processing::minifloat::fp8_f16_lift;
use cubecl::prelude::Scalar;
use cubecl::quant::scheme::QuantValue;
use cubecl::quant::scheme::ScaleDtype;

/// How an operand's values sit in memory.
#[derive(Clone, Copy, PartialEq, Eq, Hash, Debug)]
pub(crate) enum Packing {
    /// Served as stored.
    Plain,
    /// Several `field`-wide values per stored `u32`, unpacked at the read.
    Packed {
        /// The slot one value occupies.
        field: Field,
    },
}

/// The slot one packed value occupies.
#[derive(Clone, Copy, PartialEq, Eq, Hash, Debug)]
pub enum Field {
    Quant(QuantValue),
    Fp8(Fp8Format),
    Float(FloatKind),
    /// An index `bits` wide into a table, its codebook
    /// ([`TileArgLaunch::codebook`](crate::TileArgLaunch::codebook)), read as raw bits.
    Index {
        bits: usize,
    },
    /// An `e2m1` code on a device that converts `e2m1x2` itself: a word is read as four of them
    /// and cast, the decode left to the compiler. What [`read_on`](Field::read_on) makes of
    /// `Quant(E2M1)` there.
    ConvertedE2M1,
    /// A float code (`e2m1`, `e4m3` or `e5m2`) on a device that emulates converting it: a word is
    /// decoded as `f16` pairs from its bits, lifted to its values, or, where `lifted` is false,
    /// each placed `lift` short of its value for a reader whose factor carries the
    /// lift (a tile read placed). What [`read_on`](Field::read_on) makes of `Quant(value)` there,
    /// lifted.
    Placed {
        value: QuantValue,
        lifted: bool,
    },
}

impl From<QuantValue> for Field {
    fn from(value: QuantValue) -> Self {
        Field::Quant(value)
    }
}

impl From<Fp8Format> for Field {
    fn from(format: Fp8Format) -> Self {
        Field::Fp8(format)
    }
}

/// An 8-bit float is a [`Fp8`](Field::Fp8) field, a wider one a [`Float`](Field::Float) field.
impl From<FloatKind> for Field {
    fn from(kind: FloatKind) -> Self {
        Field::of_float(kind)
    }
}

impl Field {
    /// The slot's width in bits.
    pub fn size_bits(self) -> usize {
        match self {
            Field::Quant(value) => value.size_bits(),
            Field::Fp8(_) => 8,
            Field::Float(kind) => Field::float_bits(kind),
            Field::Index { bits } => bits,
            Field::ConvertedE2M1 => QuantValue::E2M1.size_bits(),
            Field::Placed { value, .. } => value.size_bits(),
        }
    }

    /// Values one `u32` holds.
    pub fn per_word(self) -> usize {
        u32::BITS as usize / self.size_bits()
    }

    /// Whether the slot holds an `e2m1` code, however the device reads it
    /// ([`read_on`](Field::read_on)): what a reader that takes the codes' values whole asks.
    pub fn holds_e2m1(self) -> bool {
        matches!(
            self,
            Field::Quant(QuantValue::E2M1)
                | Field::ConvertedE2M1
                | Field::Placed {
                    value: QuantValue::E2M1,
                    ..
                }
        )
    }
}

impl Field {
    /// This field as the device `client` reads it: a float code (`e2m1`, `e4m3`, `e5m2`) decoded
    /// from its bits where the device emulates converting it anyway, so a reader's factor can
    /// carry the decode's lift; an `e2m1` code converted by the device where it converts
    /// `e2m1x2` itself; every other field as it is. A device's own conversion is the one its
    /// compiler lowers best, so where it exists a software decode is only ever slower.
    pub fn read_on(self, client: &Client) -> Field {
        let converted_as = match self {
            Field::Quant(QuantValue::E2M1) => Some(FloatKind::E2M1x2),
            Field::Quant(QuantValue::E4M3) => Some(FloatKind::E4M3),
            Field::Quant(QuantValue::E5M2) => Some(FloatKind::E5M2),
            _ => None,
        };
        let emulated = converted_as.is_some_and(|kind| {
            client
                .features()
                .conversion_is_emulated(ElemType::Float(kind))
        });
        let converts = e2m1x2::supported_uses(client).contains(TypeUsage::Conversion);
        match self {
            Field::Quant(value) if emulated => Field::Placed {
                value,
                lifted: true,
            },
            Field::Quant(QuantValue::E2M1) if converts => Field::ConvertedE2M1,
            other => other,
        }
    }

    /// This field read placed ([`Tile::placed`](crate::Tile::placed)): a float code decoded
    /// from its bits, its lift left to the reader's factor; every other field as it is.
    pub(crate) fn placed(self) -> Field {
        match self {
            Field::Placed { value, .. } => Field::Placed {
                value,
                lifted: false,
            },
            other => other,
        }
    }

    /// What a value read this way is short of: [`E2M1_F16_LIFT`] for an `e2m1` code placed, the
    /// format's [`fp8_f16_lift`] for an `e4m3` or `e5m2` one, one for every value read as itself.
    pub(crate) fn lift(self) -> f64 {
        match self.decode() {
            FieldDecode::Placed { lifted: false } => E2M1_F16_LIFT,
            FieldDecode::PlacedByte {
                format,
                lifted: false,
            } => fp8_f16_lift(format),
            _ => 1.0,
        }
    }

    /// The field a float of `kind` occupies.
    pub fn of_float(kind: FloatKind) -> Field {
        match kind {
            FloatKind::E4M3 => Field::Fp8(Fp8Format::E4M3),
            FloatKind::E5M2 => Field::Fp8(Fp8Format::E5M2),
            FloatKind::UE8M0 => Field::Fp8(Fp8Format::UE8M0),
            other => Field::Float(other),
        }
    }

    /// The field a scale of `dtype` occupies in its stored word.
    pub fn of_scale(dtype: ScaleDtype) -> Field {
        match ElemType::from_scale_dtype(dtype) {
            ElemType::Float(kind) => Field::of_float(kind),
            other => panic!("Field: a scale is stored as a float, and {dtype:?} is {other:?}"),
        }
    }

    /// The number format a value of this field is written in: what an encoding copy rounds it to
    /// ([`Tile::quantize_from`](crate::Tile::quantize_from)). A float code reads back from its
    /// bits or through the device's conversion, but is written the same way either way.
    ///
    /// # Panics
    ///
    /// For a field no value is quantized to: an index, a whole float, an 8-bit scale code.
    pub(crate) fn quantized_value(self) -> QuantValue {
        match self {
            Field::Quant(value) | Field::Placed { value, .. } => value,
            Field::ConvertedE2M1 => QuantValue::E2M1,
            other => panic!("Field: a value is quantized to a number format, not to {other:?}"),
        }
    }

    /// The scale dtype a scale of this field is stored as: the inverse of
    /// [`of_scale`](Field::of_scale).
    ///
    /// # Panics
    ///
    /// For a field no scale is stored as.
    pub(crate) fn scale_dtype(self) -> ScaleDtype {
        match self {
            Field::Fp8(Fp8Format::E4M3) => ScaleDtype::UE4M3,
            Field::Fp8(Fp8Format::UE8M0) => ScaleDtype::UE8M0,
            Field::Float(FloatKind::F32) => ScaleDtype::F32,
            Field::Float(FloatKind::F16) => ScaleDtype::F16,
            Field::Float(FloatKind::BF16) => ScaleDtype::BF16,
            other => panic!("Field: a scale is stored as a float, not as {other:?}"),
        }
    }

    /// The bits a whole float occupies.
    pub(crate) fn float_bits(kind: FloatKind) -> usize {
        match kind {
            FloatKind::F32 => 32,
            FloatKind::F16 | FloatKind::BF16 => 16,
            other => panic!(
                "Field::Float: a stored float is 16 or 32 bits wide, got {other:?}; an 8-bit \
                 code is Field::Fp8"
            ),
        }
    }

    /// How this field reads back.
    pub(crate) fn decode(self) -> FieldDecode {
        match self {
            Field::Quant(
                QuantValue::Q8F
                | QuantValue::Q8S
                | QuantValue::Q4F
                | QuantValue::Q4S
                | QuantValue::Q2F
                | QuantValue::Q2S,
            ) => FieldDecode::SignExtended,
            Field::Quant(QuantValue::E2M1) => FieldDecode::Reinterpreted,
            Field::Quant(QuantValue::E4M3) => FieldDecode::Byte(Fp8Format::E4M3),
            Field::Quant(QuantValue::E5M2) => FieldDecode::Byte(Fp8Format::E5M2),
            Field::Fp8(format) => FieldDecode::Byte(format),
            Field::Float(kind) => FieldDecode::Bits(kind),
            Field::Index { .. } => FieldDecode::Unsigned,
            Field::ConvertedE2M1 => FieldDecode::Converted,
            // A placed code is the field it was, decoded from its bits instead.
            Field::Placed { value, lifted } => match Field::Quant(value).decode() {
                FieldDecode::Reinterpreted => FieldDecode::Placed { lifted },
                FieldDecode::Byte(format) => FieldDecode::PlacedByte { format, lifted },
                other => {
                    panic!("Field::Placed: {value:?} reads back as {other:?}, not a float code")
                }
            },
        }
    }
}

/// How a stored field reads back ([`Field::decode`]).
#[derive(Clone, Copy, PartialEq, Eq, Hash, Debug)]
pub(crate) enum FieldDecode {
    /// An integer slot: the top bit is its sign, so the value sign-extends out of its bits.
    SignExtended,
    /// An index slot: its bits, unsigned.
    Unsigned,
    /// A 4-bit float code, read back by reinterpreting the byte two of them share.
    Reinterpreted,
    /// A 4-bit float code the device converts: a word cast as four `e2m1x2`.
    Converted,
    /// A 4-bit float code decoded from its bits as `f16` pairs, lifted to its value or placed
    /// [`E2M1_F16_LIFT`] short of it.
    Placed { lifted: bool },
    /// An 8-bit float code decoded from its bits as `f16` pairs, lifted to its value or placed
    /// [`fp8_f16_lift`] short of it.
    PlacedByte { format: Fp8Format, lifted: bool },
    /// An 8-bit float code, read back through the format's own decoder.
    Byte(Fp8Format),
    /// A whole float, read back by reinterpreting the bits of its slot.
    Bits(FloatKind),
}

impl Packing {
    /// The number format a value stored this way, as `element`, is written in: what an encoding
    /// copy rounds it to ([`Tile::quantize_from`](crate::Tile::quantize_from)). A packed field's
    /// ([`Field::quantized_value`]), or a plain `i8`'s, `e4m3`'s or `e5m2`'s.
    ///
    /// # Panics
    ///
    /// For a store no value is quantized to.
    pub(crate) fn quantized_value(self, element: ElemType) -> QuantValue {
        match (self, element) {
            (Packing::Packed { field }, _) => field.quantized_value(),
            (Packing::Plain, ElemType::Int(IntKind::I8)) => QuantValue::Q8S,
            (Packing::Plain, ElemType::Float(FloatKind::E4M3)) => QuantValue::E4M3,
            (Packing::Plain, ElemType::Float(FloatKind::E5M2)) => QuantValue::E5M2,
            (Packing::Plain, other) => {
                panic!("Packing: no value is quantized to a plain {other:?}")
            }
        }
    }

    /// The dtype a scale stored this way, as `element`, is rounded to: a packed field's
    /// ([`Field::scale_dtype`]), or a plain float's.
    ///
    /// # Panics
    ///
    /// For a store no scale is held in.
    pub(crate) fn scale_dtype(self, element: ElemType) -> ScaleDtype {
        match (self, element) {
            (Packing::Packed { field }, _) => field.scale_dtype(),
            (Packing::Plain, ElemType::Float(kind)) => Field::of_float(kind).scale_dtype(),
            (Packing::Plain, other) => panic!("Packing: a scale is a float, not {other:?}"),
        }
    }

    /// This packing read placed ([`Field::placed`]).
    pub(crate) fn placed(self) -> Packing {
        match self {
            Packing::Packed { field } => Packing::Packed {
                field: field.placed(),
            },
            Packing::Plain => Packing::Plain,
        }
    }

    /// What its values are short of ([`Field::lift`]).
    pub(crate) fn lift(self) -> f64 {
        match self {
            Packing::Packed { field } => field.lift(),
            Packing::Plain => 1.0,
        }
    }

    /// This packing as the device `client` reads it ([`Field::read_on`]).
    pub(crate) fn read_on(self, client: &Client) -> Packing {
        match self {
            Packing::Packed { field } => Packing::Packed {
                field: field.read_on(client),
            },
            Packing::Plain => Packing::Plain,
        }
    }
}

impl Packing {
    /// Values per stored element: one, unless a `u32` holds several fields.
    fn factor(&self) -> usize {
        match self {
            Packing::Plain => 1,
            Packing::Packed { field } => field.per_word(),
        }
    }

    /// The physical line a `served`-wide logical line occupies.
    pub(crate) fn physical(&self, served: usize) -> usize {
        served / self.factor()
    }

    /// The width a binding `bound` wide serves.
    pub(crate) fn served(&self, bound: usize) -> usize {
        bound * self.factor()
    }
}

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

    #[test]
    fn a_packing_narrows_the_line_it_stores() {
        assert_eq!(Packing::Plain.physical(16), 16);
        assert_eq!(
            Packing::Packed {
                field: QuantValue::Q4S.into()
            }
            .physical(16),
            2
        );
    }

    #[test]
    fn a_field_states_how_many_fit_in_a_word() {
        assert_eq!(
            Packing::Packed {
                field: QuantValue::Q4S.into()
            }
            .factor(),
            8
        );
        assert_eq!(
            Packing::Packed {
                field: QuantValue::Q8S.into()
            }
            .factor(),
            4
        );
        assert_eq!(
            Packing::Packed {
                field: QuantValue::Q2S.into()
            }
            .factor(),
            16
        );
        assert_eq!(
            Packing::Packed {
                field: Fp8Format::UE8M0.into()
            }
            .factor(),
            4
        );
    }

    #[test]
    fn a_float_element_names_its_own_field() {
        assert_eq!(
            Field::of_float(FloatKind::UE8M0),
            Field::Fp8(Fp8Format::UE8M0)
        );
        assert_eq!(
            Field::of_float(FloatKind::E4M3),
            Field::Fp8(Fp8Format::E4M3)
        );
        assert_eq!(
            Field::of_float(FloatKind::F16),
            Field::Float(FloatKind::F16)
        );
    }

    #[test]
    fn a_stored_float_is_a_field_of_its_own_width() {
        assert_eq!(Field::Float(FloatKind::F32).per_word(), 1);
        assert_eq!(Field::Float(FloatKind::F16).per_word(), 2);
        assert_eq!(Field::Float(FloatKind::BF16).size_bits(), 16);
        assert_eq!(
            Field::from(FloatKind::F32).decode(),
            FieldDecode::Bits(FloatKind::F32)
        );
    }

    #[test]
    fn a_byte_field_decodes_through_its_format() {
        assert_eq!(
            Field::from(QuantValue::E4M3).decode(),
            FieldDecode::Byte(Fp8Format::E4M3)
        );
        assert_eq!(
            Field::from(Fp8Format::UE8M0).decode(),
            FieldDecode::Byte(Fp8Format::UE8M0)
        );
        assert_eq!(
            Field::from(QuantValue::E2M1).decode(),
            FieldDecode::Reinterpreted
        );
        assert_eq!(Field::ConvertedE2M1.decode(), FieldDecode::Converted);
        assert_eq!(Field::ConvertedE2M1.per_word(), 8);
    }

    /// An `e2m1` code is one however a device reads it: converted, placed, or from its bits.
    #[test]
    fn every_read_of_an_e2m1_code_holds_e2m1() {
        let placed = |value, lifted| Field::Placed { value, lifted };
        for field in [
            Field::from(QuantValue::E2M1),
            Field::ConvertedE2M1,
            placed(QuantValue::E2M1, true),
            placed(QuantValue::E2M1, false),
        ] {
            assert!(field.holds_e2m1(), "{field:?}");
        }
        for field in [
            Field::from(QuantValue::E4M3),
            placed(QuantValue::E4M3, true),
            Field::from(Fp8Format::E4M3),
        ] {
            assert!(!field.holds_e2m1(), "{field:?}");
        }
    }

    /// A float code placed is the field it was, decoded from its bits, and read placed it is
    /// short of its value by the gap between its exponent bias and an `f16`'s: `2^14` for
    /// `e2m1`, `2^8` for `e4m3`, nothing for `e5m2`, which is an `f16`'s top byte. Lifted, it
    /// is short of nothing.
    #[test]
    fn a_placed_code_leaves_its_lift_to_the_reader() {
        let placed = |value| Field::Placed {
            value,
            lifted: true,
        };
        for (value, lift) in [
            (QuantValue::E2M1, 16384.0),
            (QuantValue::E4M3, 256.0),
            (QuantValue::E5M2, 1.0),
        ] {
            assert_eq!(placed(value).size_bits(), Field::from(value).size_bits());
            assert_eq!(placed(value).lift(), 1.0, "{value:?}");
            assert_eq!(placed(value).placed().lift(), lift, "{value:?}");
        }
        assert_eq!(
            placed(QuantValue::E4M3).placed().decode(),
            FieldDecode::PlacedByte {
                format: Fp8Format::E4M3,
                lifted: false
            }
        );
        assert_eq!(
            placed(QuantValue::E2M1).decode(),
            FieldDecode::Placed { lifted: true }
        );
    }
}