fgf 3.0.0

//! Mersenne31 (`u32` lanes) AVX2 kernels.
//!
//! Eight lanes per vector. `2^31 ≡ 1 (mod p)`, so a multiply folds the
//! 62-bit product with a mask/shift/add and one conditional subtract; the
//! products come from `vpmuludq` on the even and odd 32-bit lanes. Serves
//! the `V3`/`V3GfniCrypto` backends: each entry is a safe capability-token
//! function taking an `X64V3Token` and validating its own geometry.
//!
//! The lane helpers are shared with the `QuadMersenne31` kernels, whose
//! limbs are these lanes.

use super::P;
use crate::field::Elem;
use crate::field::mersenne31::Mersenne31;
use crate::kernel::proven_checks::{check_elem_multiple, check_equal};
use crate::kernel::{prime, scalar};

#[cfg(target_arch = "x86")]
use core::arch::x86::*;
#[cfg(target_arch = "x86_64")]
use core::arch::x86_64::*;

/// Canonicalize arbitrary raw `u32` lanes below `p` in two unconditional
/// subtract-and-min steps. Every lane value fits `0..=2p + 1`
/// (`u32::MAX == 2p + 1`): `min(x, x - p)` lands in `0..=p + 1`, and the
/// second step lands in `0..=p - 1`.
#[archmage::rite(v3)]
pub(in crate::kernel::x86) fn min_chain_avx2(x: __m256i, p: __m256i) -> __m256i {
    let u = _mm256_min_epu32(x, _mm256_sub_epi32(x, p));
    _mm256_min_epu32(u, _mm256_sub_epi32(u, p))
}

/// `a * b (mod p)` over eight 31-bit lanes (`0..=p` in, canonical out).
///
/// Doubled-odd-limb form: `vpmuludq` on the even lanes and on the odd lanes
/// shifted down 31 (doubling them), blend-based recombination, one Mersenne
/// fold and one conditional subtract. Twelve arithmetic ops per vector; the
/// odd-lane gather is a port-5 `movehdup`, not a shift.
#[archmage::rite(v3)]
#[must_use]
pub(in crate::kernel::x86) fn mulmod_avx2(a: __m256i, b: __m256i, p: __m256i) -> __m256i {
    let a_odd = _mm256_srli_epi64::<31>(a); // odd lanes, doubled
    let b_odd = _mm256_castps_si256(_mm256_movehdup_ps(_mm256_castsi256_ps(b)));
    let prod_odd_dbl = _mm256_mul_epu32(a_odd, b_odd);
    let prod_evn = _mm256_mul_epu32(a, b);
    let prod_odd_lo = _mm256_slli_epi64::<31>(prod_odd_dbl);
    let prod_evn_hi = _mm256_srli_epi64::<31>(prod_evn);
    let lo_dirty = _mm256_blend_epi32::<0b1010_1010>(prod_evn, prod_odd_lo);
    let hi = _mm256_blend_epi32::<0b1010_1010>(prod_evn_hi, prod_odd_dbl);
    let lo = _mm256_and_si256(lo_dirty, p);
    let t = _mm256_add_epi32(lo, hi); // per dword < 2^32; t = p ⇒ prod ≡ 0
    _mm256_min_epu32(t, _mm256_sub_epi32(t, p))
}

/// `dst += src (mod p)`, Mersenne31, AVX2.
///
/// Every input lane bit pattern is legal: lanes canonicalize on load, and
/// every stored lane is canonical. Canonical operands sum to at most
/// `2p - 2`, so one final subtract-and-min completes the reduction.
///
/// # Panics
/// Panics if the slices differ in length or hold a partial lane.
#[allow(clippy::used_underscore_binding)]
#[archmage::arcane(import_intrinsics)]
pub fn add_assign_avx2(_token: archmage::X64V3Token, dst: &mut [u8], src: &[u8]) {
    check_equal(
        "mersenne31::add_assign_avx2",
        "dst",
        dst.len(),
        "src",
        src.len(),
    );
    check_elem_multiple("mersenne31::add_assign_avx2", dst.len(), 4);
    let p = _mm256_set1_epi32(P);
    let (dst_tiles, dst_rest) = dst.as_chunks_mut::<128>();
    let (src_tiles, src_rest) = src.as_chunks::<128>();
    for (dst_tile, src_tile) in dst_tiles.iter_mut().zip(src_tiles) {
        let (dst_lanes, _) = dst_tile.as_chunks_mut::<32>();
        let (src_lanes, _) = src_tile.as_chunks::<32>();
        for (dst_lane, src_lane) in dst_lanes.iter_mut().zip(src_lanes) {
            let d = min_chain_avx2(_mm256_loadu_si256(&*dst_lane), p);
            let s = min_chain_avx2(_mm256_loadu_si256(src_lane), p);
            let sum = _mm256_add_epi32(d, s); // <= 2p - 2
            _mm256_storeu_si256(dst_lane, _mm256_min_epu32(sum, _mm256_sub_epi32(sum, p)));
        }
    }
    let (dst_lanes, dst_tail) = dst_rest.as_chunks_mut::<32>();
    let (src_lanes, src_tail) = src_rest.as_chunks::<32>();
    for (dst_lane, src_lane) in dst_lanes.iter_mut().zip(src_lanes) {
        let d = min_chain_avx2(_mm256_loadu_si256(&*dst_lane), p);
        let s = min_chain_avx2(_mm256_loadu_si256(src_lane), p);
        let sum = _mm256_add_epi32(d, s); // <= 2p - 2
        _mm256_storeu_si256(dst_lane, _mm256_min_epu32(sum, _mm256_sub_epi32(sum, p)));
    }
    prime::add_assign::<Mersenne31>(dst_tail, src_tail);
}

/// `dst -= src (mod p)`, Mersenne31, AVX2.
///
/// Every input lane bit pattern is legal: lanes canonicalize on load, and
/// every stored lane is canonical. The wrapping difference of canonical
/// operands is already canonical after one unsigned min with its `+p`
/// fixup: a nonnegative difference stays below `p`, a negative one wraps
/// and the fixup restores the residue.
///
/// # Panics
/// Panics if the slices differ in length or hold a partial lane.
#[allow(clippy::used_underscore_binding)]
#[archmage::arcane(import_intrinsics)]
pub fn sub_assign_avx2(_token: archmage::X64V3Token, dst: &mut [u8], src: &[u8]) {
    check_equal(
        "mersenne31::sub_assign_avx2",
        "dst",
        dst.len(),
        "src",
        src.len(),
    );
    check_elem_multiple("mersenne31::sub_assign_avx2", dst.len(), 4);
    let p = _mm256_set1_epi32(P);
    let (dst_tiles, dst_rest) = dst.as_chunks_mut::<128>();
    let (src_tiles, src_rest) = src.as_chunks::<128>();
    for (dst_tile, src_tile) in dst_tiles.iter_mut().zip(src_tiles) {
        let (dst_lanes, _) = dst_tile.as_chunks_mut::<32>();
        let (src_lanes, _) = src_tile.as_chunks::<32>();
        for (dst_lane, src_lane) in dst_lanes.iter_mut().zip(src_lanes) {
            let d = min_chain_avx2(_mm256_loadu_si256(&*dst_lane), p);
            let s = min_chain_avx2(_mm256_loadu_si256(src_lane), p);
            let diff = _mm256_sub_epi32(d, s); // wraps iff negative
            _mm256_storeu_si256(dst_lane, _mm256_min_epu32(diff, _mm256_add_epi32(diff, p)));
        }
    }
    let (dst_lanes, dst_tail) = dst_rest.as_chunks_mut::<32>();
    let (src_lanes, src_tail) = src_rest.as_chunks::<32>();
    for (dst_lane, src_lane) in dst_lanes.iter_mut().zip(src_lanes) {
        let d = min_chain_avx2(_mm256_loadu_si256(&*dst_lane), p);
        let s = min_chain_avx2(_mm256_loadu_si256(src_lane), p);
        let diff = _mm256_sub_epi32(d, s); // wraps iff negative
        _mm256_storeu_si256(dst_lane, _mm256_min_epu32(diff, _mm256_add_epi32(diff, p)));
    }
    prime::sub_assign::<Mersenne31>(dst_tail, src_tail);
}

/// `dst += coeff * src (mod p)`, Mersenne31, AVX2.
///
/// A zero coefficient leaves `dst` untouched, mirroring the portable
/// reference. Every other lane is folded to canonical form on load, so
/// every input lane bit pattern is legal and every stored lane is
/// canonical.
///
/// # Panics
/// Panics if the slices differ in length or hold a partial lane.
#[allow(clippy::cast_possible_wrap, clippy::used_underscore_binding)]
#[archmage::arcane(import_intrinsics)]
pub fn mul_add_avx2(_token: archmage::X64V3Token, dst: &mut [u8], coeff: u32, src: &[u8]) {
    check_equal(
        "mersenne31::mul_add_avx2",
        "dst",
        dst.len(),
        "src",
        src.len(),
    );
    check_elem_multiple("mersenne31::mul_add_avx2", dst.len(), 4);
    // The coefficient is loop-invariant: reduce the raw word once. Raw
    // (non-canonical) coefficient words stay legal inputs, and the zero
    // check below runs on the value, matching the portable reference.
    let folded = (coeff & crate::field::mersenne31::MODULUS) + (coeff >> 31);
    let c = if folded >= crate::field::mersenne31::MODULUS {
        folded - crate::field::mersenne31::MODULUS
    } else {
        folded
    };
    if c == 0 {
        return;
    }
    let p = _mm256_set1_epi32(P);
    let cred = _mm256_set1_epi32(c as i32);
    let (dst_tiles, dst_rest) = dst.as_chunks_mut::<128>();
    let (src_tiles, src_rest) = src.as_chunks::<128>();
    for (dst_tile, src_tile) in dst_tiles.iter_mut().zip(src_tiles) {
        let (dst_lanes, _) = dst_tile.as_chunks_mut::<32>();
        let (src_lanes, _) = src_tile.as_chunks::<32>();
        for (dst_lane, src_lane) in dst_lanes.iter_mut().zip(src_lanes) {
            let prod = mulmod_avx2(cred, min_chain_avx2(_mm256_loadu_si256(src_lane), p), p);
            let d = min_chain_avx2(_mm256_loadu_si256(&*dst_lane), p);
            let sum = _mm256_add_epi32(d, prod); // <= 2p - 2
            let u = _mm256_min_epu32(sum, _mm256_sub_epi32(sum, p));
            _mm256_storeu_si256(dst_lane, u);
        }
    }
    let (dst_lanes, dst_tail) = dst_rest.as_chunks_mut::<32>();
    let (src_lanes, src_tail) = src_rest.as_chunks::<32>();
    for (dst_lane, src_lane) in dst_lanes.iter_mut().zip(src_lanes) {
        let prod = mulmod_avx2(cred, min_chain_avx2(_mm256_loadu_si256(src_lane), p), p);
        let d = min_chain_avx2(_mm256_loadu_si256(&*dst_lane), p);
        let sum = _mm256_add_epi32(d, prod);
        let u = _mm256_min_epu32(sum, _mm256_sub_epi32(sum, p));
        _mm256_storeu_si256(dst_lane, u);
    }
    prime::mul_add::<Mersenne31>(dst_tail, Elem::<Mersenne31>::from_raw(coeff), src_tail);
}

/// `dst *= coeff (mod p)`, Mersenne31, AVX2.
///
/// Every input lane bit pattern is legal: lanes canonicalize on load,
/// and every stored lane is canonical. `coeff` is canonicalized on
/// entry.
///
/// # Panics
/// Panics on a partial trailing lane.
#[allow(clippy::cast_possible_wrap, clippy::used_underscore_binding)]
#[archmage::arcane(import_intrinsics)]
pub fn mul_assign_avx2(_token: archmage::X64V3Token, dst: &mut [u8], coeff: u32) {
    check_elem_multiple("mersenne31::mul_assign_avx2", dst.len(), 4);
    // Reduce the raw word once; the unit check runs on the value, matching
    // the portable reference.
    let folded = (coeff & crate::field::mersenne31::MODULUS) + (coeff >> 31);
    let c = if folded >= crate::field::mersenne31::MODULUS {
        folded - crate::field::mersenne31::MODULUS
    } else {
        folded
    };
    if c == 1 {
        return;
    }
    let p = _mm256_set1_epi32(P);
    let cred = _mm256_set1_epi32(c as i32);
    let (dst_tiles, dst_rest) = dst.as_chunks_mut::<128>();
    for dst_tile in dst_tiles {
        let (dst_lanes, _) = dst_tile.as_chunks_mut::<32>();
        for dst_lane in dst_lanes {
            let d = min_chain_avx2(_mm256_loadu_si256(&*dst_lane), p);
            _mm256_storeu_si256(dst_lane, mulmod_avx2(cred, d, p));
        }
    }
    let (dst_lanes, dst_tail) = dst_rest.as_chunks_mut::<32>();
    for dst_lane in dst_lanes {
        let d = min_chain_avx2(_mm256_loadu_si256(&*dst_lane), p);
        _mm256_storeu_si256(dst_lane, mulmod_avx2(cred, d, p));
    }
    prime::mul_assign::<Mersenne31>(dst_tail, Elem::<Mersenne31>::from_raw(coeff));
}

/// `dst = coeff * src (mod p)`, Mersenne31, AVX2.
///
/// Every input lane bit pattern is legal: lanes canonicalize on load,
/// and every stored lane is canonical. `coeff` is canonicalized on
/// entry.
///
/// # Panics
/// Panics if the slices differ in length or hold a partial lane.
#[allow(clippy::cast_possible_wrap, clippy::used_underscore_binding)]
#[archmage::arcane(import_intrinsics)]
pub fn mul_into_avx2(_token: archmage::X64V3Token, dst: &mut [u8], coeff: u32, src: &[u8]) {
    check_equal(
        "mersenne31::mul_into_avx2",
        "dst",
        dst.len(),
        "src",
        src.len(),
    );
    check_elem_multiple("mersenne31::mul_into_avx2", dst.len(), 4);
    let p = _mm256_set1_epi32(P);
    // The coefficient is loop-invariant: canonicalize the raw broadcast
    // once. Raw (non-canonical) coefficient lanes stay legal inputs.
    let cred = min_chain_avx2(_mm256_set1_epi32(coeff as i32), p);
    let (dst_tiles, dst_rest) = dst.as_chunks_mut::<128>();
    let (src_tiles, src_rest) = src.as_chunks::<128>();
    for (dst_tile, src_tile) in dst_tiles.iter_mut().zip(src_tiles) {
        let (dst_lanes, _) = dst_tile.as_chunks_mut::<32>();
        let (src_lanes, _) = src_tile.as_chunks::<32>();
        for (dst_lane, src_lane) in dst_lanes.iter_mut().zip(src_lanes) {
            let s = min_chain_avx2(_mm256_loadu_si256(src_lane), p);
            _mm256_storeu_si256(dst_lane, mulmod_avx2(cred, s, p));
        }
    }
    let (dst_lanes, dst_tail) = dst_rest.as_chunks_mut::<32>();
    let (src_lanes, src_tail) = src_rest.as_chunks::<32>();
    for (dst_lane, src_lane) in dst_lanes.iter_mut().zip(src_lanes) {
        let s = min_chain_avx2(_mm256_loadu_si256(src_lane), p);
        _mm256_storeu_si256(dst_lane, mulmod_avx2(cred, s, p));
    }
    // Scalar tail.
    for (d, s) in dst_tail.chunks_exact_mut(4).zip(src_tail.chunks_exact(4)) {
        let v = Elem::<Mersenne31>::from_raw(coeff)
            * Elem::<Mersenne31>::from_raw(u32::from_le_bytes(s.try_into().unwrap()));
        d.copy_from_slice(&v.to_raw().to_le_bytes());
    }
}

/// `dst[i] = a[i] * b[i] (mod p)`, Mersenne31, AVX2.
///
/// Every input lane bit pattern is legal: lanes canonicalize on load,
/// and every stored lane is canonical.
///
/// # Panics
/// Panics unless all three buffers match in length (whole lanes).
#[allow(clippy::used_underscore_binding)]
#[archmage::arcane(import_intrinsics)]
pub fn mul_elementwise_avx2(_token: archmage::X64V3Token, dst: &mut [u8], a: &[u8], b: &[u8]) {
    check_equal(
        "mersenne31::mul_elementwise_avx2",
        "dst",
        dst.len(),
        "a",
        a.len(),
    );
    check_equal(
        "mersenne31::mul_elementwise_avx2",
        "dst",
        dst.len(),
        "b",
        b.len(),
    );
    check_elem_multiple("mersenne31::mul_elementwise_avx2", dst.len(), 4);
    let p = _mm256_set1_epi32(P);
    let (dst_tiles, dst_rest) = dst.as_chunks_mut::<128>();
    let (a_tiles, a_rest) = a.as_chunks::<128>();
    let (b_tiles, b_rest) = b.as_chunks::<128>();
    for (dst_tile, (a_tile, b_tile)) in dst_tiles.iter_mut().zip(a_tiles.iter().zip(b_tiles)) {
        let (dst_lanes, _) = dst_tile.as_chunks_mut::<32>();
        let (a_lanes, _) = a_tile.as_chunks::<32>();
        let (b_lanes, _) = b_tile.as_chunks::<32>();
        for ((dst_lane, a_lane), b_lane) in dst_lanes.iter_mut().zip(a_lanes).zip(b_lanes) {
            let va = min_chain_avx2(_mm256_loadu_si256(a_lane), p);
            let vb = min_chain_avx2(_mm256_loadu_si256(b_lane), p);
            _mm256_storeu_si256(dst_lane, mulmod_avx2(va, vb, p));
        }
    }
    let (dst_lanes, dst_tail) = dst_rest.as_chunks_mut::<32>();
    let (a_lanes, a_tail) = a_rest.as_chunks::<32>();
    let (b_lanes, b_tail) = b_rest.as_chunks::<32>();
    for ((dst_lane, a_lane), b_lane) in dst_lanes.iter_mut().zip(a_lanes).zip(b_lanes) {
        let va = min_chain_avx2(_mm256_loadu_si256(a_lane), p);
        let vb = min_chain_avx2(_mm256_loadu_si256(b_lane), p);
        _mm256_storeu_si256(dst_lane, mulmod_avx2(va, vb, p));
    }
    prime::mul_elementwise::<Mersenne31>(dst_tail, a_tail, b_tail);
}

/// `dst[i] = dst[i] * src[i] (mod p)`, Mersenne31, AVX2.
///
/// Every input lane bit pattern is legal: lanes canonicalize on load,
/// and every stored lane is canonical.
///
/// # Panics
/// Panics if the slices differ in length or hold a partial lane.
#[allow(clippy::used_underscore_binding)]
#[archmage::arcane(import_intrinsics)]
pub fn mul_elementwise_assign_avx2(_token: archmage::X64V3Token, dst: &mut [u8], src: &[u8]) {
    check_equal(
        "mersenne31::mul_elementwise_assign_avx2",
        "dst",
        dst.len(),
        "src",
        src.len(),
    );
    check_elem_multiple("mersenne31::mul_elementwise_assign_avx2", dst.len(), 4);
    let p = _mm256_set1_epi32(P);
    let (dst_tiles, dst_rest) = dst.as_chunks_mut::<128>();
    let (src_tiles, src_rest) = src.as_chunks::<128>();
    for (dst_tile, src_tile) in dst_tiles.iter_mut().zip(src_tiles) {
        let (dst_lanes, _) = dst_tile.as_chunks_mut::<32>();
        let (src_lanes, _) = src_tile.as_chunks::<32>();
        for (dst_lane, src_lane) in dst_lanes.iter_mut().zip(src_lanes) {
            let vd = min_chain_avx2(_mm256_loadu_si256(&*dst_lane), p);
            let vs = min_chain_avx2(_mm256_loadu_si256(src_lane), p);
            _mm256_storeu_si256(dst_lane, mulmod_avx2(vd, vs, p));
        }
    }
    let (dst_lanes, dst_tail) = dst_rest.as_chunks_mut::<32>();
    let (src_lanes, src_tail) = src_rest.as_chunks::<32>();
    for (dst_lane, src_lane) in dst_lanes.iter_mut().zip(src_lanes) {
        let vd = min_chain_avx2(_mm256_loadu_si256(&*dst_lane), p);
        let vs = min_chain_avx2(_mm256_loadu_si256(src_lane), p);
        _mm256_storeu_si256(dst_lane, mulmod_avx2(vd, vs, p));
    }
    scalar::mul_elementwise_assign::<Mersenne31>(dst_tail, src_tail);
}

/// `dst[i] += value (mod p)`, Mersenne31, AVX2. The raw `value` word is
/// canonicalized once on entry.
///
/// The destination lanes canonicalize on load; the canonical sum is at
/// most `2p - 2`, so one final subtract-and-min completes it. Every stored
/// lane is canonical.
///
/// # Panics
/// Panics on a partial trailing lane.
#[allow(clippy::used_underscore_binding)]
#[archmage::arcane(import_intrinsics)]
pub fn add_assign_scalar_avx2(_token: archmage::X64V3Token, dst: &mut [u8], value: u32) {
    check_elem_multiple("mersenne31::add_assign_scalar_avx2", dst.len(), 4);
    let p = _mm256_set1_epi32(P);
    // The broadcast value is loop-invariant: canonicalize the raw word once.
    let svec = _mm256_set1_epi32(Elem::<Mersenne31>::from_raw(value).to_raw().cast_signed());
    let (dst_tiles, dst_rest) = dst.as_chunks_mut::<128>();
    for dst_tile in dst_tiles {
        let (dst_lanes, _) = dst_tile.as_chunks_mut::<32>();
        for dst_lane in dst_lanes {
            let d = min_chain_avx2(_mm256_loadu_si256(&*dst_lane), p);
            let sum = _mm256_add_epi32(d, svec); // <= 2p - 2
            _mm256_storeu_si256(dst_lane, _mm256_min_epu32(sum, _mm256_sub_epi32(sum, p)));
        }
    }
    let (dst_lanes, dst_tail) = dst_rest.as_chunks_mut::<32>();
    for dst_lane in dst_lanes {
        let d = min_chain_avx2(_mm256_loadu_si256(&*dst_lane), p);
        let sum = _mm256_add_epi32(d, svec); // <= 2p - 2
        _mm256_storeu_si256(dst_lane, _mm256_min_epu32(sum, _mm256_sub_epi32(sum, p)));
    }
    prime::add_assign_scalar::<Mersenne31>(dst_tail, Elem::<Mersenne31>::from_raw(value));
}

/// `dst[i] -= value (mod p)`, Mersenne31, AVX2. The raw `value` word is
/// canonicalized once on entry.
///
/// The destination lanes canonicalize on load; the canonical difference is
/// canonical after one unsigned min with its `+p` fixup, matching
/// [`sub_assign_avx2`].
///
/// # Panics
/// Panics on a partial trailing lane.
#[allow(clippy::used_underscore_binding)]
#[archmage::arcane(import_intrinsics)]
pub fn sub_assign_scalar_avx2(_token: archmage::X64V3Token, dst: &mut [u8], value: u32) {
    check_elem_multiple("mersenne31::sub_assign_scalar_avx2", dst.len(), 4);
    let p = _mm256_set1_epi32(P);
    let svec = _mm256_set1_epi32(Elem::<Mersenne31>::from_raw(value).to_raw().cast_signed());
    let (dst_tiles, dst_rest) = dst.as_chunks_mut::<128>();
    for dst_tile in dst_tiles {
        let (dst_lanes, _) = dst_tile.as_chunks_mut::<32>();
        for dst_lane in dst_lanes {
            let d = min_chain_avx2(_mm256_loadu_si256(&*dst_lane), p);
            let diff = _mm256_sub_epi32(d, svec); // wraps iff negative
            _mm256_storeu_si256(dst_lane, _mm256_min_epu32(diff, _mm256_add_epi32(diff, p)));
        }
    }
    let (dst_lanes, dst_tail) = dst_rest.as_chunks_mut::<32>();
    for dst_lane in dst_lanes {
        let d = min_chain_avx2(_mm256_loadu_si256(&*dst_lane), p);
        let diff = _mm256_sub_epi32(d, svec); // wraps iff negative
        _mm256_storeu_si256(dst_lane, _mm256_min_epu32(diff, _mm256_add_epi32(diff, p)));
    }
    prime::sub_assign_scalar::<Mersenne31>(dst_tail, Elem::<Mersenne31>::from_raw(value));
}