prism-inference 0.1.17

//! Exact integer reductions over one decoded block shared by a prompt tile.

use super::super::{MatvecBackend, pq2_values, ptq1_values, q8_native_lanes};
use crate::Result;

#[cfg(all(
    feature = "experimental-neon",
    target_arch = "aarch64",
    target_os = "windows",
    target_env = "msvc",
    target_feature = "neon"
))]
#[path = "prefill_neon.rs"]
mod neon;

/// Exactly one block, scoped to the packed row loop; no heap allocation.
pub(super) enum Block {
    Scalar([i16; 128]),
    #[cfg(all(
        feature = "experimental-neon",
        target_arch = "aarch64",
        target_os = "windows",
        target_env = "msvc",
        target_feature = "neon"
    ))]
    Neon(neon::Block),
}

impl Block {
    #[cfg(test)]
    #[allow(
        clippy::missing_const_for_fn,
        reason = "Only the scalar fallback is constant; retain one signature across targets and feature sets."
    )]
    pub(super) fn new(weights: [i16; 128], backend: MatvecBackend) -> Self {
        #[cfg(all(
            feature = "experimental-neon",
            target_arch = "aarch64",
            target_os = "windows",
            target_env = "msvc",
            target_feature = "neon"
        ))]
        if backend != MatvecBackend::Scalar
            && let Some(block) = neon::Block::new(&weights)
        {
            return Self::Neon(block);
        }
        // Unsupported builds retain the scalar result for every request policy.
        let _ = backend;
        Self::Scalar(weights)
    }

    pub(super) fn from_pq2(packed: &[u8; 32], backend: MatvecBackend) -> Self {
        #[cfg(all(
            feature = "experimental-neon",
            target_arch = "aarch64",
            target_os = "windows",
            target_env = "msvc",
            target_feature = "neon"
        ))]
        if backend != MatvecBackend::Scalar
            && let Some(block) = neon::Block::from_pq2(packed)
        {
            return Self::Neon(block);
        }
        let _ = backend;
        let mut full = [0_u8; 34];
        full[2..].copy_from_slice(packed);
        Self::Scalar(pq2_values(&full))
    }

    pub(super) fn from_ptq1(packed: &[u8; 28], backend: MatvecBackend) -> Result<Self> {
        #[cfg(all(
            feature = "experimental-neon",
            target_arch = "aarch64",
            target_os = "windows",
            target_env = "msvc",
            target_feature = "neon"
        ))]
        if backend != MatvecBackend::Scalar
            && let Some(block) = neon::Block::from_ptq1(packed)
        {
            return Ok(Self::Neon(block));
        }
        let _ = backend;
        Ok(Self::Scalar(ptq1_values(packed)?))
    }

    #[inline]
    pub(super) fn dot128(&self, values: &[i8; 128]) -> i32 {
        match self {
            Self::Scalar(weights) => scalar(weights, values),
            #[cfg(all(
                feature = "experimental-neon",
                target_arch = "aarch64",
                target_os = "windows",
                target_env = "msvc",
                target_feature = "neon"
            ))]
            Self::Neon(block) => block.dot128(values),
        }
    }

    #[inline]
    pub(super) fn dot128_four(&self, values: [&[i8; 128]; 4]) -> [i32; 4] {
        match self {
            Self::Scalar(weights) => values.map(|token| scalar(weights, token)),
            #[cfg(all(
                feature = "experimental-neon",
                target_arch = "aarch64",
                target_os = "windows",
                target_env = "msvc",
                target_feature = "neon"
            ))]
            Self::Neon(block) => block.dot128_four(values),
        }
    }

    #[inline]
    pub(super) fn dot128_eight(&self, values: [&[i8; 128]; 8]) -> [i32; 8] {
        match self {
            Self::Scalar(weights) => values.map(|token| scalar(weights, token)),
            #[cfg(all(
                feature = "experimental-neon",
                target_arch = "aarch64",
                target_os = "windows",
                target_env = "msvc",
                target_feature = "neon"
            ))]
            Self::Neon(block) => block.dot128_eight(values),
        }
    }

    #[inline]
    pub(super) fn dot32(&self, group: usize, values: &[i8; 32]) -> i32 {
        match self {
            Self::Scalar(weights) => scalar(&weights[group * 32..group * 32 + 32], values),
            #[cfg(all(
                feature = "experimental-neon",
                target_arch = "aarch64",
                target_os = "windows",
                target_env = "msvc",
                target_feature = "neon"
            ))]
            Self::Neon(block) => block.dot32(group, values),
        }
    }

    #[inline]
    pub(super) fn dot32_four(&self, group: usize, values: [&[i8; 32]; 4]) -> [i32; 4] {
        match self {
            Self::Scalar(weights) => {
                let slice = &weights[group * 32..group * 32 + 32];
                values.map(|token| scalar(slice, token))
            }
            #[cfg(all(
                feature = "experimental-neon",
                target_arch = "aarch64",
                target_os = "windows",
                target_env = "msvc",
                target_feature = "neon"
            ))]
            Self::Neon(block) => block.dot32_four(group, values),
        }
    }

    #[inline]
    pub(super) fn dot32_eight(&self, group: usize, values: [&[i8; 32]; 8]) -> [i32; 8] {
        match self {
            Self::Scalar(weights) => {
                let slice = &weights[group * 32..group * 32 + 32];
                values.map(|token| scalar(slice, token))
            }
            #[cfg(all(
                feature = "experimental-neon",
                target_arch = "aarch64",
                target_os = "windows",
                target_env = "msvc",
                target_feature = "neon"
            ))]
            Self::Neon(block) => block.dot32_eight(group, values),
        }
    }

    #[inline]
    pub(super) fn lanes32(&self, group: usize, values: &[i8; 32]) -> [i32; 4] {
        match self {
            Self::Scalar(weights) => q8_native_lanes(&weights[group * 32..group * 32 + 32], values),
            #[cfg(all(
                feature = "experimental-neon",
                target_arch = "aarch64",
                target_os = "windows",
                target_env = "msvc",
                target_feature = "neon"
            ))]
            Self::Neon(block) => block.lanes32(group, values),
        }
    }

    #[inline]
    pub(super) fn lanes32_four(&self, group: usize, values: [&[i8; 32]; 4]) -> [[i32; 4]; 4] {
        match self {
            Self::Scalar(weights) => {
                let slice = &weights[group * 32..group * 32 + 32];
                values.map(|token| q8_native_lanes(slice, token))
            }
            #[cfg(all(
                feature = "experimental-neon",
                target_arch = "aarch64",
                target_os = "windows",
                target_env = "msvc",
                target_feature = "neon"
            ))]
            Self::Neon(block) => block.lanes32_four(group, values),
        }
    }

    #[inline]
    pub(super) fn lanes32_eight(&self, group: usize, values: [&[i8; 32]; 8]) -> [[i32; 4]; 8] {
        match self {
            Self::Scalar(weights) => {
                let slice = &weights[group * 32..group * 32 + 32];
                values.map(|token| q8_native_lanes(slice, token))
            }
            #[cfg(all(
                feature = "experimental-neon",
                target_arch = "aarch64",
                target_os = "windows",
                target_env = "msvc",
                target_feature = "neon"
            ))]
            Self::Neon(block) => block.lanes32_eight(group, values),
        }
    }
}

fn scalar(weights: &[i16], values: &[i8]) -> i32 {
    weights
        .iter()
        .zip(values)
        .map(|(&weight, &value)| i32::from(weight) * i32::from(value))
        .sum()
}

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

    #[test]
    fn tiled_integer_reductions_preserve_extreme_values_and_native_lanes() {
        for seed in 0_u8..=255 {
            let weights = std::array::from_fn(|index: usize| {
                i16::from(seed.wrapping_add(index.to_le_bytes()[0]) & 3) - 1
            });
            let values = std::array::from_fn(|index: usize| {
                i8::from_ne_bytes([seed.wrapping_mul(197).wrapping_add(index.to_le_bytes()[0])])
            });
            for backend in [
                MatvecBackend::Scalar,
                MatvecBackend::Auto,
                MatvecBackend::ExperimentalNeon,
            ] {
                let block = Block::new(weights, backend);
                if backend == MatvecBackend::Scalar {
                    assert!(matches!(block, Block::Scalar(_)));
                }
                #[cfg(all(
                    feature = "experimental-neon",
                    target_arch = "aarch64",
                    target_os = "windows",
                    target_env = "msvc",
                    target_feature = "neon"
                ))]
                if backend != MatvecBackend::Scalar {
                    assert!(matches!(block, Block::Neon(_)));
                }
                let expected128 = scalar(&weights, &values);
                assert_eq!(block.dot128(&values), expected128);
                assert_eq!(block.dot128_four([&values; 4]), [expected128; 4]);
                assert_eq!(block.dot128_eight([&values; 8]), [expected128; 8]);
                for (group, values) in values.as_chunks::<32>().0.iter().enumerate() {
                    let weights = &weights[group * 32..group * 32 + 32];
                    let expected32 = scalar(weights, values);
                    let expected_lanes = q8_native_lanes(weights, values);
                    assert_eq!(block.dot32(group, values), expected32);
                    assert_eq!(block.dot32_four(group, [values; 4]), [expected32; 4]);
                    assert_eq!(block.dot32_eight(group, [values; 8]), [expected32; 8]);
                    assert_eq!(block.lanes32(group, values), expected_lanes);
                    assert_eq!(block.lanes32_four(group, [values; 4]), [expected_lanes; 4]);
                    assert_eq!(block.lanes32_eight(group, [values; 8]), [expected_lanes; 8]);
                }
            }
        }
    }
}