goosedump 0.12.1

// 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::{Context as _, Result, bail, ensure};
use num_traits::ToPrimitive as _;
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)
}

pub(super) struct Q8Activation {
    scales: Vec<f32>,
    values: Vec<i8>,
}

impl Q8Activation {
    pub(super) 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);
            for value in block {
                let quantized = (value * inverse).round().clamp(-127.0, 127.0);
                values.push(quantized.to_i8().context("convert Q8 activation")?);
            }
        }
        Ok(Self { scales, values })
    }
}

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

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!("avx512f") && std::is_x86_feature_detected!("avx512bw")
            {
                return Self::Avx512;
            }
            #[cfg(target_arch = "x86_64")]
            if std::is_x86_feature_detected!("avx2") {
                return Self::Avx2;
            }
            #[cfg(target_arch = "aarch64")]
            {
                return Self::Neon;
            }
            #[allow(unreachable_code)]
            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) },
            #[cfg(target_arch = "x86_64")]
            Self::Avx512 => unsafe { dot_q8_avx512(row, activation) },
            #[cfg(target_arch = "aarch64")]
            Self::Neon => unsafe { dot_q8_neon(row, activation) },
        }
    }
}

#[derive(Clone, Copy)]
enum Mxfp4Kernel {
    Scalar,
    #[cfg(target_arch = "x86_64")]
    Avx2,
    #[cfg(target_arch = "x86_64")]
    Avx512,
    #[cfg(target_arch = "aarch64")]
    Neon,
}

impl Mxfp4Kernel {
    fn detect() -> Self {
        static KERNEL: OnceLock<Mxfp4Kernel> = OnceLock::new();
        *KERNEL.get_or_init(|| {
            #[cfg(target_arch = "x86_64")]
            if std::is_x86_feature_detected!("avx512f") && std::is_x86_feature_detected!("avx512bw")
            {
                return Self::Avx512;
            }
            #[cfg(target_arch = "x86_64")]
            if std::is_x86_feature_detected!("avx2") {
                return Self::Avx2;
            }
            #[cfg(target_arch = "aarch64")]
            {
                return Self::Neon;
            }
            #[allow(unreachable_code)]
            Self::Scalar
        })
    }

    fn dot(self, row: &[u8], activation: &Q8Activation) -> f32 {
        match self {
            Self::Scalar => dot_mxfp4_scalar(row, activation),
            #[cfg(target_arch = "x86_64")]
            Self::Avx2 => unsafe { dot_mxfp4_avx2(row, activation) },
            #[cfg(target_arch = "x86_64")]
            Self::Avx512 => unsafe { dot_mxfp4_avx512(row, activation) },
            #[cfg(target_arch = "aarch64")]
            Self::Neon => unsafe { dot_mxfp4_neon(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_scalar(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;

    let mut blocks = row.chunks_exact(BLOCK_BYTES);
    let mut values_chunks = activation.values.chunks_exact(BLOCK_VALUES);
    let mut scales = activation.scales.iter().copied();

    while let (Some(b0), Some(v0), Some(s0), Some(b1), Some(v1), Some(s1)) = (
        blocks.next(),
        values_chunks.next(),
        scales.next(),
        blocks.next(),
        values_chunks.next(),
        scales.next(),
    ) {
        let w0: __m256i = bytemuck::pod_read_unaligned(&b0[2..]);
        let a0: __m256i = bytemuck::pod_read_unaligned(bytemuck::cast_slice(v0));
        let w1: __m256i = bytemuck::pod_read_unaligned(&b1[2..]);
        let a1: __m256i = bytemuck::pod_read_unaligned(bytemuck::cast_slice(v1));

        let signed0 = _mm256_sign_epi8(a0, w0);
        let signed1 = _mm256_sign_epi8(a1, w1);
        let mag0 = _mm256_abs_epi8(w0);
        let mag1 = _mm256_abs_epi8(w1);
        let pairs0 = _mm256_maddubs_epi16(mag0, signed0);
        let pairs1 = _mm256_maddubs_epi16(mag1, signed1);
        let prod0 = _mm256_madd_epi16(pairs0, ones);
        let prod1 = _mm256_madd_epi16(pairs1, ones);

        let low0 = _mm256_castsi256_si128(prod0);
        let high0 = _mm256_extracti128_si256::<1>(prod0);
        let lanes0 = _mm_add_epi32(low0, high0);
        let p0 = _mm_add_epi32(lanes0, _mm_unpackhi_epi64(lanes0, lanes0));
        let tot0 = _mm_add_epi32(p0, _mm_shuffle_epi32::<0x55>(p0));

        let low1 = _mm256_castsi256_si128(prod1);
        let high1 = _mm256_extracti128_si256::<1>(prod1);
        let lanes1 = _mm_add_epi32(low1, high1);
        let p1 = _mm_add_epi32(lanes1, _mm_unpackhi_epi64(lanes1, lanes1));
        let tot1 = _mm_add_epi32(p1, _mm_shuffle_epi32::<0x55>(p1));

        let ws0 = fp16_to_f32(u16::from_le_bytes([b0[0], b0[1]]));
        let ws1 = fp16_to_f32(u16::from_le_bytes([b1[0], b1[1]]));

        sum += ws0 * s0 * _mm_cvtss_f32(_mm_cvtepi32_ps(tot0));
        sum += ws1 * s1 * _mm_cvtss_f32(_mm_cvtepi32_ps(tot1));
    }

    while let (Some(b0), Some(v0), Some(s0)) = (blocks.next(), values_chunks.next(), scales.next())
    {
        let w0: __m256i = bytemuck::pod_read_unaligned(&b0[2..]);
        let a0: __m256i = bytemuck::pod_read_unaligned(bytemuck::cast_slice(v0));
        let signed0 = _mm256_sign_epi8(a0, w0);
        let mag0 = _mm256_abs_epi8(w0);
        let pairs0 = _mm256_maddubs_epi16(mag0, signed0);
        let prod0 = _mm256_madd_epi16(pairs0, ones);
        let low0 = _mm256_castsi256_si128(prod0);
        let high0 = _mm256_extracti128_si256::<1>(prod0);
        let lanes0 = _mm_add_epi32(low0, high0);
        let p0 = _mm_add_epi32(lanes0, _mm_unpackhi_epi64(lanes0, lanes0));
        let tot0 = _mm_add_epi32(p0, _mm_shuffle_epi32::<0x55>(p0));
        let ws0 = fp16_to_f32(u16::from_le_bytes([b0[0], b0[1]]));
        sum += ws0 * s0 * _mm_cvtss_f32(_mm_cvtepi32_ps(tot0));
    }

    sum
}

#[cfg(target_arch = "x86_64")]
#[target_feature(enable = "avx2")]
unsafe fn dot_mxfp4_avx2(row: &[u8], activation: &Q8Activation) -> f32 {
    use std::arch::x86_64::{
        __m128i, __m256i, _mm_add_epi32, _mm_and_si128, _mm_cvtepi32_ps, _mm_cvtss_f32,
        _mm_set1_epi8, _mm_shuffle_epi8, _mm_shuffle_epi32, _mm_srli_epi16, _mm_unpackhi_epi64,
        _mm256_abs_epi8, _mm256_castsi256_si128, _mm256_extracti128_si256, _mm256_madd_epi16,
        _mm256_maddubs_epi16, _mm256_set_m128i, _mm256_set1_epi16, _mm256_sign_epi8,
    };

    const LUT: [i8; 16] = [0, 1, 2, 3, 4, 6, 8, 12, 0, -1, -2, -3, -4, -6, -8, -12];
    let lut_128: __m128i = bytemuck::cast(LUT);
    let mask_0f = _mm_set1_epi8(0x0F);
    let ones = _mm256_set1_epi16(1);

    let mut sum = 0.0;

    let mut blocks = row.chunks_exact(MXFP4_BLOCK_BYTES);
    let mut values_chunks = activation.values.chunks_exact(BLOCK_VALUES);
    let mut scales = activation.scales.iter().copied();

    while let (Some(b0), Some(v0), Some(s0), Some(b1), Some(v1), Some(s1)) = (
        blocks.next(),
        values_chunks.next(),
        scales.next(),
        blocks.next(),
        values_chunks.next(),
        scales.next(),
    ) {
        let p128_0: __m128i = bytemuck::pod_read_unaligned(&b0[1..17]);
        let low0 = _mm_and_si128(p128_0, mask_0f);
        let high0 = _mm_and_si128(_mm_srli_epi16(p128_0, 4), mask_0f);
        let w_low0 = _mm_shuffle_epi8(lut_128, low0);
        let w_high0 = _mm_shuffle_epi8(lut_128, high0);
        let w0 = _mm256_set_m128i(w_high0, w_low0);
        let a0: __m256i = bytemuck::pod_read_unaligned(bytemuck::cast_slice(v0));

        let p128_1: __m128i = bytemuck::pod_read_unaligned(&b1[1..17]);
        let low1 = _mm_and_si128(p128_1, mask_0f);
        let high1 = _mm_and_si128(_mm_srli_epi16(p128_1, 4), mask_0f);
        let w_low1 = _mm_shuffle_epi8(lut_128, low1);
        let w_high1 = _mm_shuffle_epi8(lut_128, high1);
        let w1 = _mm256_set_m128i(w_high1, w_low1);
        let a1: __m256i = bytemuck::pod_read_unaligned(bytemuck::cast_slice(v1));

        let signed0 = _mm256_sign_epi8(a0, w0);
        let signed1 = _mm256_sign_epi8(a1, w1);
        let mag0 = _mm256_abs_epi8(w0);
        let mag1 = _mm256_abs_epi8(w1);
        let pairs0 = _mm256_maddubs_epi16(mag0, signed0);
        let pairs1 = _mm256_maddubs_epi16(mag1, signed1);
        let prod0 = _mm256_madd_epi16(pairs0, ones);
        let prod1 = _mm256_madd_epi16(pairs1, ones);

        let l0 = _mm256_castsi256_si128(prod0);
        let h0 = _mm256_extracti128_si256::<1>(prod0);
        let lanes0 = _mm_add_epi32(l0, h0);
        let p0 = _mm_add_epi32(lanes0, _mm_unpackhi_epi64(lanes0, lanes0));
        let tot0 = _mm_add_epi32(p0, _mm_shuffle_epi32::<0x55>(p0));

        let l1 = _mm256_castsi256_si128(prod1);
        let h1 = _mm256_extracti128_si256::<1>(prod1);
        let lanes1 = _mm_add_epi32(l1, h1);
        let p1 = _mm_add_epi32(lanes1, _mm_unpackhi_epi64(lanes1, lanes1));
        let tot1 = _mm_add_epi32(p1, _mm_shuffle_epi32::<0x55>(p1));

        let sc0 = e8m0_to_f32_half(b0[0]);
        let sc1 = e8m0_to_f32_half(b1[0]);

        sum += sc0 * s0 * _mm_cvtss_f32(_mm_cvtepi32_ps(tot0));
        sum += sc1 * s1 * _mm_cvtss_f32(_mm_cvtepi32_ps(tot1));
    }

    while let (Some(b0), Some(v0), Some(s0)) = (blocks.next(), values_chunks.next(), scales.next())
    {
        let p128_0: __m128i = bytemuck::pod_read_unaligned(&b0[1..17]);
        let low0 = _mm_and_si128(p128_0, mask_0f);
        let high0 = _mm_and_si128(_mm_srli_epi16(p128_0, 4), mask_0f);
        let w_low0 = _mm_shuffle_epi8(lut_128, low0);
        let w_high0 = _mm_shuffle_epi8(lut_128, high0);
        let w0 = _mm256_set_m128i(w_high0, w_low0);
        let a0: __m256i = bytemuck::pod_read_unaligned(bytemuck::cast_slice(v0));

        let signed0 = _mm256_sign_epi8(a0, w0);
        let mag0 = _mm256_abs_epi8(w0);
        let pairs0 = _mm256_maddubs_epi16(mag0, signed0);
        let prod0 = _mm256_madd_epi16(pairs0, ones);

        let l0 = _mm256_castsi256_si128(prod0);
        let h0 = _mm256_extracti128_si256::<1>(prod0);
        let lanes0 = _mm_add_epi32(l0, h0);
        let p0 = _mm_add_epi32(lanes0, _mm_unpackhi_epi64(lanes0, lanes0));
        let tot0 = _mm_add_epi32(p0, _mm_shuffle_epi32::<0x55>(p0));

        let sc0 = e8m0_to_f32_half(b0[0]);
        sum += sc0 * s0 * _mm_cvtss_f32(_mm_cvtepi32_ps(tot0));
    }

    sum
}

#[cfg(target_arch = "x86_64")]
#[target_feature(enable = "avx512f,avx512bw")]
unsafe fn dot_q8_avx512(row: &[u8], activation: &Q8Activation) -> f32 {
    unsafe { dot_q8_avx2(row, activation) }
}

#[cfg(target_arch = "x86_64")]
#[target_feature(enable = "avx512f,avx512bw")]
unsafe fn dot_mxfp4_avx512(row: &[u8], activation: &Q8Activation) -> f32 {
    unsafe { dot_mxfp4_avx2(row, activation) }
}

#[cfg(target_arch = "aarch64")]
#[target_feature(enable = "neon")]
unsafe fn dot_q8_neon(row: &[u8], activation: &Q8Activation) -> f32 {
    use std::arch::aarch64::{
        vaddq_s32, vaddvq_s32, vget_low_s8, vget_low_s16, vld1q_s8, vmovl_high_s8, vmovl_s8,
        vmull_high_s16, vmull_s16,
    };
    let mut sum = 0.0_f32;
    let blocks = row.chunks_exact(BLOCK_BYTES);
    let values_chunks = activation.values.chunks_exact(BLOCK_VALUES);
    let scales = activation.scales.iter().copied();

    for ((block, values), activation_scale) in blocks.zip(values_chunks).zip(scales) {
        let weight_scale = fp16_to_f32(u16::from_le_bytes([block[0], block[1]]));
        let w_ptr = block[2..].as_ptr().cast::<i8>();
        let v_ptr = values.as_ptr();

        let block_sum = unsafe {
            let w0 = vld1q_s8(w_ptr);
            let w1 = vld1q_s8(w_ptr.add(16));
            let v0 = vld1q_s8(v_ptr);
            let v1 = vld1q_s8(v_ptr.add(16));

            let w0_low = vmovl_s8(vget_low_s8(w0));
            let w0_high = vmovl_high_s8(w0);
            let v0_low = vmovl_s8(vget_low_s8(v0));
            let v0_high = vmovl_high_s8(v0);

            let w1_low = vmovl_s8(vget_low_s8(w1));
            let w1_high = vmovl_high_s8(w1);
            let v1_low = vmovl_s8(vget_low_s8(v1));
            let v1_high = vmovl_high_s8(v1);

            let p0 = vmull_s16(vget_low_s16(w0_low), vget_low_s16(v0_low));
            let p1 = vmull_high_s16(w0_low, v0_low);
            let p2 = vmull_s16(vget_low_s16(w0_high), vget_low_s16(v0_high));
            let p3 = vmull_high_s16(w0_high, v0_high);

            let p4 = vmull_s16(vget_low_s16(w1_low), vget_low_s16(v1_low));
            let p5 = vmull_high_s16(w1_low, v1_low);
            let p6 = vmull_s16(vget_low_s16(w1_high), vget_low_s16(v1_high));
            let p7 = vmull_high_s16(w1_high, v1_high);

            let acc0 = vaddq_s32(vaddq_s32(p0, p1), vaddq_s32(p2, p3));
            let acc1 = vaddq_s32(vaddq_s32(p4, p5), vaddq_s32(p6, p7));
            let acc = vaddq_s32(acc0, acc1);

            vaddvq_s32(acc)
        };

        #[allow(clippy::cast_precision_loss)]
        let block_sum_f32 = block_sum as f32;
        sum += weight_scale * activation_scale * block_sum_f32;
    }
    sum
}

#[cfg(target_arch = "aarch64")]
unsafe fn dot_mxfp4_neon(row: &[u8], activation: &Q8Activation) -> f32 {
    let mut sum = 0.0_f32;
    let blocks = row.chunks_exact(MXFP4_BLOCK_BYTES);
    let values_chunks = activation.values.chunks_exact(BLOCK_VALUES);
    let scales = activation.scales.iter().copied();

    for ((block, values), activation_scale) in blocks.zip(values_chunks).zip(scales) {
        let mut block_sum = 0.0_f32;
        for (index, packed) in block[1..].iter().copied().enumerate() {
            block_sum +=
                f32::from(values[index]) * f32::from(MXFP4_VALUES[usize::from(packed & 0x0f)]);
            block_sum += f32::from(values[index + BLOCK_VALUES / 2])
                * f32::from(MXFP4_VALUES[usize::from(packed >> 4)]);
        }
        sum += e8m0_to_f32_half(block[0]) * activation_scale * block_sum;
    }
    sum
}

pub(super) fn matrix_vector(matrix: &Tensor<'_>, vector: &[f32]) -> Result<Vec<f32>> {
    matrix_vector_with_activation(matrix, vector, None)
}

pub(super) fn matrix_vector_with_activation(
    matrix: &Tensor<'_>,
    vector: &[f32],
    activation: Option<&Q8Activation>,
) -> 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 => {
            if let Some(act) = activation {
                project_q8(matrix, act, &mut result)?;
            } else {
                let act = Q8Activation::new(vector)?;
                project_q8(matrix, &act, &mut result)?;
            }
        }
        TensorType::Mxfp4 => {
            if let Some(act) = activation {
                project_mxfp4(matrix, act, &mut result)?;
            } else {
                let act = Q8Activation::new(vector)?;
                project_mxfp4(matrix, &act, &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 mxfp4_kernel = Mxfp4Kernel::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 => mxfp4_kernel.dot(
                    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");
    let kernel = Mxfp4Kernel::detect();
    output
        .par_iter_mut()
        .enumerate()
        .try_for_each(|(row, value)| -> Result<()> {
            *value = kernel.dot(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(())
}