goosedump 0.11.2

// SPDX-License-Identifier: LGPL-2.1-or-later
// Copyright (C) Jarkko Sakkinen 2026

//! Scalar and runtime-selected CPU inference kernels.
use std::sync::OnceLock;

use anyhow::{Result, bail, ensure};
use num_traits::ToPrimitive;
use rayon::prelude::*;

use super::gguf::{Tensor, TensorType};

const BLOCK_VALUES: usize = 32;
const BLOCK_BYTES: usize = 34;
const MXFP4_BLOCK_BYTES: usize = 17;
const MXFP4_VALUES: [i8; 16] = [0, 1, 2, 3, 4, 6, 8, 12, 0, -1, -2, -3, -4, -6, -8, -12];

/// Lossless `usize -> f32` for model dimensions/indexes that fit in `u16`.
/// `f32::from(u16)` is exact (24-bit mantissa holds any u16); `try_from` is the
/// clippy-recommended narrowing. Use only when the value provably fits `u16`.
pub(super) fn dim_to_f32(value: usize) -> f32 {
    f32::from(u16::try_from(value).expect("model dimension/index exceeds u16"))
}

pub(super) fn fp16_to_f32(value: u16) -> f32 {
    let sign = u32::from(value & 0x8000) << 16;
    let exponent = (value >> 10) & 0x1f;
    let fraction = value & 0x03ff;
    match exponent {
        0 if fraction == 0 => f32::from_bits(sign),
        0 => {
            let sign_multiplier = if sign == 0 { 1.0 } else { -1.0 };
            sign_multiplier * f32::from(fraction) * 2.0_f32.powi(-24)
        }
        0x1f => f32::from_bits(sign | 0x7f80_0000 | (u32::from(fraction) << 13)),
        _ => {
            f32::from_bits(sign | ((u32::from(exponent) + 112) << 23) | (u32::from(fraction) << 13))
        }
    }
}

pub(super) fn f32_to_fp16(value: f32) -> u16 {
    let bits = value.to_bits();
    let sign = u16::try_from(bits >> 16).expect("top 16 bits fit u16") & 0x8000;
    let exponent = ((bits >> 23) & 0xff).cast_signed();
    let fraction = bits & 0x007f_ffff;
    if exponent == 0xff {
        return if fraction == 0 {
            sign | 0x7c00
        } else {
            sign | 0x7e00 | (u16::try_from(fraction >> 13).expect("fraction fits u16") & 0x01ff)
        };
    }
    let half_exponent = exponent - 112;
    if half_exponent >= 0x1f {
        return sign | 0x7c00;
    }
    if half_exponent <= 0 {
        if half_exponent < -10 {
            return sign;
        }
        let significand = fraction | 0x0080_0000;
        let shift = u32::try_from(14 - half_exponent).expect("shift is non-negative");
        let mut rounded = significand >> shift;
        let remainder = significand & ((1_u32 << shift) - 1);
        let halfway = 1_u32 << (shift - 1);
        if remainder > halfway || (remainder == halfway && rounded & 1 != 0) {
            rounded += 1;
        }
        return sign | u16::try_from(rounded).expect("rounded fits u16");
    }
    let mut rounded = fraction >> 13;
    let remainder = fraction & 0x1fff;
    if remainder > 0x1000 || (remainder == 0x1000 && rounded & 1 != 0) {
        rounded += 1;
    }
    let mut encoded_exponent = u16::try_from(half_exponent).expect("half_exponent in 1..=0x1e");
    if rounded == 0x400 {
        rounded = 0;
        encoded_exponent += 1;
        if encoded_exponent == 0x1f {
            return sign | 0x7c00;
        }
    }
    sign | (encoded_exponent << 10) | u16::try_from(rounded).expect("rounded fits u16")
}

pub(super) fn dequantize_row(row: &[u8], output: &mut [f32]) -> Result<()> {
    ensure!(
        output.len().is_multiple_of(BLOCK_VALUES),
        "Q8_0 output width is not divisible by 32"
    );
    ensure!(
        row.len() == output.len() / BLOCK_VALUES * BLOCK_BYTES,
        "invalid Q8_0 row size"
    );
    for (block, values) in row
        .chunks_exact(BLOCK_BYTES)
        .zip(output.chunks_exact_mut(BLOCK_VALUES))
    {
        let scale = fp16_to_f32(u16::from_le_bytes([block[0], block[1]]));
        for (value, quantized) in values.iter_mut().zip(&block[2..]) {
            *value = scale * f32::from(i8::from_ne_bytes([*quantized]));
        }
    }
    Ok(())
}

fn e8m0_to_f32_half(value: u8) -> f32 {
    let bits = if value < 2 {
        0x0020_0000_u32 << value
    } else {
        u32::from(value - 1) << 23
    };
    f32::from_bits(bits)
}

struct Q8Activation {
    scales: Vec<f32>,
    values: Vec<i8>,
}

impl Q8Activation {
    // `round()` + `clamp(-127.0, 127.0)` keeps the value in `i8` range, so the
    // checked `to_i8()` is total here and rejects only a contract bug.
    fn new(vector: &[f32]) -> Result<Self> {
        ensure!(
            vector.len().is_multiple_of(BLOCK_VALUES),
            "Q8 activation width is not divisible by 32"
        );
        ensure!(
            vector.iter().all(|value| value.is_finite()),
            "Q8 activation is not finite"
        );
        let mut scales = Vec::with_capacity(vector.len() / BLOCK_VALUES);
        let mut values = Vec::with_capacity(vector.len());
        for block in vector.chunks_exact(BLOCK_VALUES) {
            let maximum = block.iter().copied().map(f32::abs).fold(0.0, f32::max);
            let scale = maximum / 127.0;
            let inverse = if scale == 0.0 { 0.0 } else { scale.recip() };
            scales.push(scale);
            values.extend(block.iter().map(|value| {
                (value * inverse)
                    .round()
                    .clamp(-127.0, 127.0)
                    .to_i8()
                    .expect("scale is clamped into i8 range")
            }));
        }
        Ok(Self { scales, values })
    }
}

#[derive(Clone, Copy)]
enum Q8Kernel {
    Scalar,
    #[cfg(target_arch = "x86_64")]
    Avx2,
}

impl Q8Kernel {
    fn detect() -> Self {
        static KERNEL: OnceLock<Q8Kernel> = OnceLock::new();
        *KERNEL.get_or_init(|| {
            #[cfg(target_arch = "x86_64")]
            if std::is_x86_feature_detected!("avx2") {
                return Self::Avx2;
            }
            Self::Scalar
        })
    }

    fn dot(self, row: &[u8], activation: &Q8Activation) -> f32 {
        match self {
            Self::Scalar => dot_q8_scalar(row, activation),
            #[cfg(target_arch = "x86_64")]
            Self::Avx2 => unsafe { dot_q8_avx2(row, activation) },
        }
    }
}

fn dot_q8_scalar(row: &[u8], activation: &Q8Activation) -> f32 {
    row.chunks_exact(BLOCK_BYTES)
        .zip(activation.values.chunks_exact(BLOCK_VALUES))
        .zip(&activation.scales)
        .map(|((block, values), activation_scale)| {
            let weight_scale = fp16_to_f32(u16::from_le_bytes([block[0], block[1]]));
            // i8 weights and u8 activations each fit exactly in f32; the
            // 32-wide product sum stays below f32's 24-bit exact range.
            let sum = block[2..]
                .iter()
                .zip(values)
                .map(|(weight, value)| f32::from(i8::from_ne_bytes([*weight])) * f32::from(*value))
                .sum::<f32>();
            weight_scale * activation_scale * sum
        })
        .sum()
}

fn dot_mxfp4(row: &[u8], activation: &Q8Activation) -> f32 {
    row.chunks_exact(MXFP4_BLOCK_BYTES)
        .zip(activation.values.chunks_exact(BLOCK_VALUES))
        .zip(&activation.scales)
        .map(|((block, values), activation_scale)| {
            // u8 activations and i8 MXFP4 lattice points each fit exactly in
            // f32; the 32-wide sum stays below f32's 24-bit exact range.
            let mut sum = 0.0_f32;
            for (index, packed) in block[1..].iter().copied().enumerate() {
                sum +=
                    f32::from(values[index]) * f32::from(MXFP4_VALUES[usize::from(packed & 0x0f)]);
                sum += f32::from(values[index + BLOCK_VALUES / 2])
                    * f32::from(MXFP4_VALUES[usize::from(packed >> 4)]);
            }
            e8m0_to_f32_half(block[0]) * activation_scale * sum
        })
        .sum()
}

fn dot_bf16(row: &[u8], vector: &[f32]) -> f32 {
    row.chunks_exact(2)
        .zip(vector)
        .map(|(bytes, value)| {
            let weight = f32::from_bits(u32::from(u16::from_le_bytes([bytes[0], bytes[1]])) << 16);
            weight * value
        })
        .sum()
}

#[cfg(target_arch = "x86_64")]
#[target_feature(enable = "avx2")]
unsafe fn dot_q8_avx2(row: &[u8], activation: &Q8Activation) -> f32 {
    use std::arch::x86_64::{
        __m256i, _mm_add_epi32, _mm_cvtepi32_ps, _mm_cvtss_f32, _mm_shuffle_epi32,
        _mm_unpackhi_epi64, _mm256_abs_epi8, _mm256_castsi256_si128, _mm256_extracti128_si256,
        _mm256_madd_epi16, _mm256_maddubs_epi16, _mm256_set1_epi16, _mm256_sign_epi8,
    };
    let ones = _mm256_set1_epi16(1);
    let mut sum = 0.0;
    for ((block, values), activation_scale) in row
        .chunks_exact(BLOCK_BYTES)
        .zip(activation.values.chunks_exact(BLOCK_VALUES))
        .zip(&activation.scales)
    {
        // `read_unaligned` is clippy-clean (it is the intended unaligned read);
        // codegen is identical to `_mm256_loadu_si256` (`vmovups`).
        let weights = unsafe { std::ptr::read_unaligned(block[2..].as_ptr().cast::<__m256i>()) };
        let activations = unsafe { std::ptr::read_unaligned(values.as_ptr().cast::<__m256i>()) };
        let signed = _mm256_sign_epi8(activations, weights);
        let magnitudes = _mm256_abs_epi8(weights);
        let pairs = _mm256_maddubs_epi16(magnitudes, signed);
        let products = _mm256_madd_epi16(pairs, ones);
        let low = _mm256_castsi256_si128(products);
        let high = _mm256_extracti128_si256::<1>(products);
        let lanes = _mm_add_epi32(low, high);
        let pairs = _mm_add_epi32(lanes, _mm_unpackhi_epi64(lanes, lanes));
        let total = _mm_add_epi32(pairs, _mm_shuffle_epi32::<0x55>(pairs));
        let weight_scale = fp16_to_f32(u16::from_le_bytes([block[0], block[1]]));
        sum += weight_scale * activation_scale * _mm_cvtss_f32(_mm_cvtepi32_ps(total));
    }
    sum
}

pub(super) fn matrix_vector(matrix: &Tensor<'_>, vector: &[f32]) -> Result<Vec<f32>> {
    let [input, output] = matrix_dimensions(matrix)?;
    ensure!(vector.len() == input, "matrix input width differs");
    let mut result = vec![0.0; output];
    match matrix.tensor_type() {
        TensorType::F32 => {
            result
                .par_iter_mut()
                .enumerate()
                .try_for_each(|(row, value)| -> Result<()> {
                    *value = dot_f32(matrix.f32_row(row)?, vector);
                    Ok(())
                })?;
        }
        TensorType::Bf16 => {
            result
                .par_iter_mut()
                .enumerate()
                .try_for_each(|(row, value)| -> Result<()> {
                    *value = dot_bf16(matrix.bf16_row(row)?, vector);
                    Ok(())
                })?;
        }
        TensorType::Q8_0 => {
            let activation = Q8Activation::new(vector)?;
            project_q8(matrix, &activation, &mut result)?;
        }
        TensorType::Mxfp4 => {
            let activation = Q8Activation::new(vector)?;
            project_mxfp4(matrix, &activation, &mut result)?;
        }
    }
    Ok(result)
}

pub(super) fn matrix_vector_triple(
    first: &Tensor<'_>,
    second: &Tensor<'_>,
    third: &Tensor<'_>,
    vector: &[f32],
) -> Result<(Vec<f32>, Vec<f32>, Vec<f32>)> {
    if first.tensor_type() == TensorType::Q8_0
        && second.tensor_type() == TensorType::Q8_0
        && third.tensor_type() == TensorType::Q8_0
    {
        let activation = Q8Activation::new(vector)?;
        let first_rows = matrix_dimensions(first)?[1];
        let second_rows = matrix_dimensions(second)?[1];
        let third_rows = matrix_dimensions(third)?[1];
        let mut first_output = vec![0.0; first_rows];
        let mut second_output = vec![0.0; second_rows];
        let mut third_output = vec![0.0; third_rows];
        let kernel = Q8Kernel::detect();
        first_output
            .par_iter_mut()
            .chain(second_output.par_iter_mut())
            .chain(third_output.par_iter_mut())
            .enumerate()
            .try_for_each(|(row, value)| -> Result<()> {
                if row < first_rows {
                    *value = kernel.dot(first.q8_row(row)?, &activation);
                } else if row < first_rows + second_rows {
                    *value = kernel.dot(second.q8_row(row - first_rows)?, &activation);
                } else {
                    *value = kernel.dot(third.q8_row(row - first_rows - second_rows)?, &activation);
                }
                Ok(())
            })?;
        return Ok((first_output, second_output, third_output));
    }
    Ok((
        matrix_vector(first, vector)?,
        matrix_vector(second, vector)?,
        matrix_vector(third, vector)?,
    ))
}

pub(super) fn matrix_argmax(matrix: &Tensor<'_>, vector: &[f32]) -> Result<usize> {
    let [input, output] = matrix_dimensions(matrix)?;
    ensure!(vector.len() == input, "matrix input width differs");
    let q8_activation = match matrix.tensor_type() {
        TensorType::F32 | TensorType::Bf16 => None,
        TensorType::Q8_0 | TensorType::Mxfp4 => Some(Q8Activation::new(vector)?),
    };
    let kernel = Q8Kernel::detect();
    let best = (0..output)
        .into_par_iter()
        .map(|index| -> Result<(usize, f32)> {
            let value = match matrix.tensor_type() {
                TensorType::F32 => dot_f32(matrix.f32_row(index)?, vector),
                TensorType::Bf16 => dot_bf16(matrix.bf16_row(index)?, vector),
                TensorType::Q8_0 => kernel.dot(
                    matrix.q8_row(index)?,
                    q8_activation.as_ref().expect("quantized activation"),
                ),
                TensorType::Mxfp4 => dot_mxfp4(
                    matrix.mxfp4_row(index)?,
                    q8_activation.as_ref().expect("quantized activation"),
                ),
            };
            if !value.is_finite() {
                bail!("matrix output {index} is not finite");
            }
            Ok((index, value))
        })
        .try_reduce_with(|left, right| {
            Ok(match right.1.total_cmp(&left.1) {
                std::cmp::Ordering::Greater => right,
                std::cmp::Ordering::Equal if right.0 < left.0 => right,
                _ => left,
            })
        })
        .transpose()?
        .expect("validated matrix has output rows");
    Ok(best.0)
}

fn project_q8(matrix: &Tensor<'_>, activation: &Q8Activation, output: &mut [f32]) -> Result<()> {
    let [input, rows] = matrix_dimensions(matrix)?;
    ensure!(
        matrix.tensor_type() == TensorType::Q8_0,
        "projection is not Q8_0"
    );
    ensure!(
        activation.values.len() == input,
        "matrix input width differs"
    );
    ensure!(output.len() == rows, "matrix output height differs");
    let kernel = Q8Kernel::detect();
    output
        .par_iter_mut()
        .enumerate()
        .try_for_each(|(row, value)| -> Result<()> {
            *value = kernel.dot(matrix.q8_row(row)?, activation);
            Ok(())
        })?;
    Ok(())
}

fn project_mxfp4(matrix: &Tensor<'_>, activation: &Q8Activation, output: &mut [f32]) -> Result<()> {
    let [input, rows] = matrix_dimensions(matrix)?;
    ensure!(
        matrix.tensor_type() == TensorType::Mxfp4,
        "projection is not MXFP4"
    );
    ensure!(
        activation.values.len() == input,
        "matrix input width differs"
    );
    ensure!(output.len() == rows, "matrix output height differs");
    output
        .par_iter_mut()
        .enumerate()
        .try_for_each(|(row, value)| -> Result<()> {
            *value = dot_mxfp4(matrix.mxfp4_row(row)?, activation);
            Ok(())
        })?;
    Ok(())
}

fn matrix_dimensions(matrix: &Tensor<'_>) -> Result<[usize; 2]> {
    match matrix.dimensions() {
        [input, output] => Ok([*input, *output]),
        dimensions => bail!("matrix has dimensions {dimensions:?}, expected two"),
    }
}

fn dot_f32(left: &[f32], right: &[f32]) -> f32 {
    debug_assert_eq!(left.len(), right.len());
    #[cfg(target_arch = "x86_64")]
    if std::is_x86_feature_detected!("avx2") && std::is_x86_feature_detected!("fma") {
        return unsafe { dot_f32_avx2(left, right) };
    }
    left.iter()
        .zip(right)
        .map(|(left, right)| left * right)
        .sum()
}

#[cfg(target_arch = "x86_64")]
#[target_feature(enable = "avx2,fma")]
unsafe fn dot_f32_avx2(left: &[f32], right: &[f32]) -> f32 {
    use std::arch::x86_64::{
        _mm256_fmadd_ps, _mm256_loadu_ps, _mm256_setzero_ps, _mm256_storeu_ps,
    };
    let vectorized = left.len() / 8 * 8;
    let mut sums = _mm256_setzero_ps();
    for index in (0..vectorized).step_by(8) {
        let left = unsafe { _mm256_loadu_ps(left.as_ptr().add(index)) };
        let right = unsafe { _mm256_loadu_ps(right.as_ptr().add(index)) };
        sums = _mm256_fmadd_ps(left, right, sums);
    }
    let mut lanes = [0.0; 8];
    unsafe { _mm256_storeu_ps(lanes.as_mut_ptr(), sums) };
    lanes.into_iter().sum::<f32>()
        + left[vectorized..]
            .iter()
            .zip(&right[vectorized..])
            .map(|(left, right)| left * right)
            .sum::<f32>()
}

pub(super) fn rms_norm(
    values: &[f32],
    width: usize,
    weight: &[f32],
    epsilon: f32,
) -> Result<Vec<f32>> {
    ensure!(
        width != 0 && values.len().is_multiple_of(width),
        "invalid RMS norm shape"
    );
    ensure!(weight.len() == width, "invalid RMS norm weight");
    let mut output = vec![0.0; values.len()];
    for (input, output) in values
        .chunks_exact(width)
        .zip(output.chunks_exact_mut(width))
    {
        let width_f32 = dim_to_f32(width);
        let mean_square = input.iter().map(|value| value * value).sum::<f32>() / width_f32;
        let scale = (mean_square + epsilon).sqrt().recip();
        for index in 0..width {
            output[index] = input[index] * scale * weight[index];
        }
    }
    Ok(output)
}

pub(super) fn softmax(values: &mut [f32]) {
    let maximum = values.iter().copied().fold(f32::NEG_INFINITY, f32::max);
    let sum = values
        .iter_mut()
        .map(|value| {
            *value = (*value - maximum).exp();
            *value
        })
        .sum::<f32>();
    for value in values {
        *value /= sum;
    }
}

pub(super) fn vector_add(left: &mut [f32], right: &[f32]) -> Result<()> {
    ensure!(left.len() == right.len(), "vector lengths differ");
    left.iter_mut()
        .zip(right)
        .for_each(|(left, right)| *left += right);
    Ok(())
}