prism-inference 0.1.17

//! Exact baseline NEON value tiles with adjacent channels in vector lanes.

use fearless_simd::{Neon, f32x4, prelude::*};

use crate::{Result, invalid};

use super::{VALUE_TILE, padded_dot_steps};

/// Consumes a validated tile whose width is a positive multiple of four.
///
/// The sixteen native temporal FMA lanes stay independent. Each vector groups
/// four adjacent value channels, making an update contiguous without changing
/// any channel's token order. The accumulator remains exactly 2 KiB; decoded
/// FP16 values exist only for one four-channel load. No FP16 arithmetic or
/// optional instruction set is required.
pub(super) fn value_tile<'a>(
    neon: Neon,
    mut past: impl Iterator<Item = &'a [u16]>,
    current: &[f32],
    weights: &[f32],
    range: &std::ops::Range<usize>,
    output: &mut [f32],
) -> Result<()> {
    let zero = f32x4::splat(neon, 0.0);
    let mut sums = [[zero; VALUE_TILE / 4]; 16];
    let groups = output.len() / 4;
    let last = weights.len() - 1;
    for (index, &weight) in weights[..last].iter().enumerate() {
        let values = past
            .next()
            .and_then(|row| row.get(range.start..range.end))
            .ok_or_else(|| invalid("pinned attention value row is short"))?;
        let weight = f32x4::splat(neon, weight);
        let (sum_pairs, sum_tail) = sums[index % 16][..groups].as_chunks_mut::<2>();
        let (val_chunks8, val_rem) = values.as_chunks::<8>();
        for (pair, chunk8) in sum_pairs.iter_mut().zip(val_chunks8) {
            let (v0, v1) = super::super::neon::f16_bits_to_f32_x8(neon, chunk8);
            pair[0] = v0.mul_add_precise(weight, pair[0]);
            pair[1] = v1.mul_add_precise(weight, pair[1]);
        }
        for (sum, values) in sum_tail.iter_mut().zip(val_rem.as_chunks::<4>().0) {
            let values = super::super::neon::f16_bits_to_f32_x4(neon, *values);
            *sum = values.mul_add_precise(weight, *sum);
        }
    }
    if past.next().is_some() {
        return Err(invalid("pinned attention value row is long"));
    }
    let weight = f32x4::splat(neon, weights[last]);
    for (sum, values) in sums[last % 16][..groups]
        .iter_mut()
        .zip(current.as_chunks::<4>().0)
    {
        *sum = f32x4::load_array(neon, *values).mul_add_precise(weight, *sum);
    }
    // Preserve the first masked FMA in every temporal lane, including -0.
    for index in weights.len()..padded_dot_steps(weights.len()) {
        for sum in &mut sums[index % 16][..groups] {
            *sum = zero.mul_add_precise(zero, *sum);
        }
    }
    for (group, out) in output.as_chunks_mut::<4>().0.iter_mut().enumerate() {
        let lanes: [f32x4<Neon>; 4] = std::array::from_fn(|lane| {
            (sums[lane][group] + sums[8 + lane][group])
                + (sums[4 + lane][group] + sums[12 + lane][group])
        });
        let result = ((lanes[0] + lanes[1]) + (lanes[2] + lanes[3])).to_array();
        // Publication and finite rejection retain increasing channel order.
        for (out, value) in out.iter_mut().zip(result) {
            *out = super::super::finite(value)?;
        }
    }
    Ok(())
}

/// Reduces a validated F16 key slice and F32 query slice with four NEON FMA registers.
pub(super) fn dot_f16(neon: Neon, key: &[u16], query: &[f32]) -> Result<f32> {
    let mut sums = [f32x4::splat(neon, 0.0); 4];
    let (key_chunks8, key_rem) = key.as_chunks::<8>();
    let (query_chunks8, query_rem) = query.as_chunks::<8>();
    for (pair_idx, (k8, q8)) in key_chunks8.iter().zip(query_chunks8).enumerate() {
        let base = (pair_idx * 2) % 4;
        let (kv0, kv1) = super::super::neon::f16_bits_to_f32_x8(neon, k8);
        let qv0 = f32x4::from_slice(neon, &q8[..4]);
        let qv1 = f32x4::from_slice(neon, &q8[4..]);
        sums[base] = kv0.mul_add_precise(qv0, sums[base]);
        sums[base + 1] = kv1.mul_add_precise(qv1, sums[base + 1]);
    }
    let rem_base = key_chunks8.len() * 2;
    let key_chunks4 = key_rem.as_chunks::<4>().0;
    let query_chunks4 = query_rem.as_chunks::<4>().0;
    for (offset, (k, q)) in key_chunks4.iter().zip(query_chunks4).enumerate() {
        let acc = &mut sums[(rem_base + offset) % 4];
        let kv = super::super::neon::f16_bits_to_f32_x4(neon, *k);
        let qv = f32x4::load_array(neon, *q);
        *acc = kv.mul_add_precise(qv, *acc);
    }
    let vector_steps = key_chunks8.len() * 8 + key_chunks4.len() * 4;
    let total_steps = padded_dot_steps(query.len());
    if vector_steps == total_steps {
        let lanes = ((sums[0] + sums[2]) + (sums[1] + sums[3])).to_array();
        return super::super::finite((lanes[0] + lanes[1]) + (lanes[2] + lanes[3]));
    }
    let mut scalar_sums = sums.map(f32x4::to_array);
    for index in vector_steps..total_steps {
        let (value, weight) = if index < query.len() {
            (bonsai_ternary::half_to_float(key[index]), query[index])
        } else {
            (0.0, 0.0)
        };
        let sum = &mut scalar_sums[index % 16 / 4][index % 4];
        *sum = value.mul_add(weight, *sum);
    }
    let mut lanes = [0.0; 4];
    for (index, lane) in lanes.iter_mut().enumerate() {
        *lane = (scalar_sums[0][index] + scalar_sums[2][index])
            + (scalar_sums[1][index] + scalar_sums[3][index]);
    }
    super::super::finite((lanes[0] + lanes[1]) + (lanes[2] + lanes[3]))
}

/// Reduces a validated F32 key slice and F32 query slice with four NEON FMA registers.
pub(super) fn dot_f32(neon: Neon, key: &[f32], query: &[f32]) -> Result<f32> {
    let mut sums = [f32x4::splat(neon, 0.0); 4];
    let key_chunks = key.as_chunks::<4>().0;
    let query_chunks = query.as_chunks::<4>().0;
    for (chunk, (k, q)) in key_chunks.iter().zip(query_chunks).enumerate() {
        let acc = &mut sums[chunk % 4];
        let kv = f32x4::load_array(neon, *k);
        let qv = f32x4::load_array(neon, *q);
        *acc = kv.mul_add_precise(qv, *acc);
    }
    let vector_steps = key_chunks.len() * 4;
    let total_steps = padded_dot_steps(query.len());
    if vector_steps == total_steps {
        let lanes = ((sums[0] + sums[2]) + (sums[1] + sums[3])).to_array();
        return super::super::finite((lanes[0] + lanes[1]) + (lanes[2] + lanes[3]));
    }
    let mut scalar_sums = sums.map(f32x4::to_array);
    for index in vector_steps..total_steps {
        let (value, weight) = if index < query.len() {
            (key[index], query[index])
        } else {
            (0.0, 0.0)
        };
        let sum = &mut scalar_sums[index % 16 / 4][index % 4];
        *sum = value.mul_add(weight, *sum);
    }
    let mut lanes = [0.0; 4];
    for (index, lane) in lanes.iter_mut().enumerate() {
        *lane = (scalar_sums[0][index] + scalar_sums[2][index])
            + (scalar_sums[1][index] + scalar_sums[3][index]);
    }
    super::super::finite((lanes[0] + lanes[1]) + (lanes[2] + lanes[3]))
}

/// Rotates validated 256-wide Swift heads (with 64-element rotary prefix) in NEON `f32x4`.
pub(super) fn rope_with_aliasing(
    neon: Neon,
    input: &mut [f32],
    frequencies: &[(f32, f32); 32],
    in_place: bool,
) {
    let cos_vecs: [f32x4<Neon>; 8] = std::array::from_fn(|chunk| {
        let base = chunk * 4;
        f32x4::load_array(
            neon,
            [
                frequencies[base].0,
                frequencies[base + 1].0,
                frequencies[base + 2].0,
                frequencies[base + 3].0,
            ],
        )
    });
    let sin_vecs: [f32x4<Neon>; 8] = std::array::from_fn(|chunk| {
        let base = chunk * 4;
        f32x4::load_array(
            neon,
            [
                frequencies[base].1,
                frequencies[base + 1].1,
                frequencies[base + 2].1,
                frequencies[base + 3].1,
            ],
        )
    });
    for head in input.as_chunks_mut::<256>().0 {
        let (first_half, second_half) = head[..64].split_at_mut(32);
        for (((first_chunk, second_chunk), &cos), &sin) in first_half
            .as_chunks_mut::<4>()
            .0
            .iter_mut()
            .zip(second_half.as_chunks_mut::<4>().0)
            .zip(&cos_vecs)
            .zip(&sin_vecs)
        {
            let first = f32x4::from_slice(neon, first_chunk);
            let second = f32x4::from_slice(neon, second_chunk);
            let rotated_first = if in_place {
                first.mul_add_precise(cos, -(second * sin))
            } else {
                (-second).mul_add_precise(sin, first * cos)
            };
            let rotated_second = second.mul_add_precise(cos, first * sin);
            rotated_first.store_slice(first_chunk);
            rotated_second.store_slice(second_chunk);
        }
    }
}

#[cfg(test)]
mod tests {
    use bonsai_ternary::half_to_float;
    use fearless_simd::{Neon, f32x4, prelude::*};

    use super::{VALUE_TILE, value_tile};

    #[derive(serde::Deserialize)]
    struct Capture {
        weighted_value: WeightedValue,
    }

    #[derive(serde::Deserialize)]
    #[allow(
        clippy::struct_field_names,
        reason = "Field names match the immutable JSON fixture schema."
    )]
    struct WeightedValue {
        value_half_bits: Vec<u16>,
        probability_half_bits: Vec<u16>,
        native_bits: u32,
    }

    #[test]
    fn neon_value_tile_baseline_is_available_and_accumulator_remains_two_kib() -> crate::Result<()>
    {
        let neon = fearless_simd::Level::baseline()
            .as_neon()
            .ok_or_else(|| crate::invalid("required attention baseline NEON token is absent"))?;
        assert_eq!(size_of::<[[f32x4<Neon>; VALUE_TILE / 4]; 16]>(), 2048);
        let capture: Capture = serde_json::from_str(include_str!(
            "../../tests/fixtures/swift-attention-dot.json"
        ))?;
        let captured = capture.weighted_value;
        assert_ne!(captured.value_half_bits, [0_u16; 0]);
        let last = captured.value_half_bits.len() - 1;
        let rows: Vec<_> = captured.value_half_bits[..last]
            .iter()
            .map(|&value| [value; VALUE_TILE])
            .collect();
        let current = [half_to_float(captured.value_half_bits[last]); VALUE_TILE];
        let weights: Vec<_> = captured
            .probability_half_bits
            .into_iter()
            .map(half_to_float)
            .collect();
        let mut output = [f32::NAN; VALUE_TILE];
        value_tile(
            neon,
            rows.iter().map(<[u16; VALUE_TILE]>::as_slice),
            &current,
            &weights,
            &(0..VALUE_TILE),
            &mut output,
        )?;
        assert!(
            output
                .iter()
                .all(|value| value.to_bits() == captured.native_bits)
        );
        Ok(())
    }

    #[test]
    fn neon_f16_conversions_preserve_all_binary16_bit_patterns() -> crate::Result<()> {
        let neon = fearless_simd::Level::baseline()
            .as_neon()
            .ok_or_else(|| crate::invalid("required attention baseline NEON token is absent"))?;
        let all_bits: Vec<u16> = (0..=u16::MAX).collect();
        for chunk8 in all_bits.as_chunks::<8>().0 {
            let (v0, v1) = super::super::super::neon::f16_bits_to_f32_x8(neon, chunk8);
            let a0 = v0.to_array();
            let a1 = v1.to_array();
            let c0 = super::super::super::neon::f16_bits_to_f32_x4(
                neon,
                [chunk8[0], chunk8[1], chunk8[2], chunk8[3]],
            )
            .to_array();
            let c1 = super::super::super::neon::f16_bits_to_f32_x4(
                neon,
                [chunk8[4], chunk8[5], chunk8[6], chunk8[7]],
            )
            .to_array();
            for lane in 0..4 {
                let expected0 = half_to_float(chunk8[lane]).to_bits();
                let expected1 = half_to_float(chunk8[lane + 4]).to_bits();
                assert_eq!(a0[lane].to_bits(), expected0);
                assert_eq!(c0[lane].to_bits(), expected0);
                assert_eq!(a1[lane].to_bits(), expected1);
                assert_eq!(c1[lane].to_bits(), expected1);
            }
        }
        Ok(())
    }
}