use super::super::{broadcast_words, check_elements, swap_mask_avx2};
use crate::field::gf16::Elem;
use crate::kernel::gf16::mul_add_scalar;
use crate::kernel::tables::TowerCoeff;
use super::scale_gfni;
use super::single::mul_add_gfni;
#[cfg(target_arch = "x86")]
use core::arch::x86::*;
#[cfg(target_arch = "x86_64")]
use core::arch::x86_64::*;
const SOURCE_GROUP: usize = 4;
#[archmage::arcane(import_intrinsics)]
pub fn mul_add_gather_gfni(
token: archmage::X64V3GfniCryptoToken,
dst: &mut [u8],
coeffs: &[Elem],
srcs: &[&[u8]],
) {
check_elements("gf16::mul_add_gather_gfni", dst.len());
assert_eq!(
coeffs.len(),
srcs.len(),
"gf16::mul_add_gather_gfni: coefficients is {} but sources is {}",
coeffs.len(),
srcs.len(),
);
for (index, &src) in srcs.iter().enumerate() {
assert_eq!(
dst.len(),
src.len(),
"gf16::mul_add_gather_gfni: dst is {} bytes but source {index} is {} bytes",
dst.len(),
src.len(),
);
}
if dst.is_empty() || srcs.is_empty() {
return;
}
let swap = swap_mask_avx2();
let mut i = 0;
while i + SOURCE_GROUP <= coeffs.len() {
let mut group = [Elem(0); SOURCE_GROUP];
group.copy_from_slice(&coeffs[i..i + SOURCE_GROUP]);
let mut sources: [&[u8]; SOURCE_GROUP] = [&[]; SOURCE_GROUP];
sources.copy_from_slice(&srcs[i..i + SOURCE_GROUP]);
gather_group(dst, group, sources, swap);
i += SOURCE_GROUP;
}
if i + 2 <= coeffs.len() {
let mut group = [Elem(0); 2];
group.copy_from_slice(&coeffs[i..i + 2]);
let mut sources: [&[u8]; 2] = [&[]; 2];
sources.copy_from_slice(&srcs[i..i + 2]);
gather_group(dst, group, sources, swap);
i += 2;
}
if i < coeffs.len() {
mul_add_gfni(token, dst, TowerCoeff::new(coeffs[i]), srcs[i]);
}
}
#[archmage::rite(v3_gfni_crypto, import_intrinsics)]
fn gather_group<const N: usize>(
dst: &mut [u8],
coeffs: [Elem; N],
srcs: [&[u8]; N],
swap: __m256i,
) {
let tail_start = dst.len() & !31;
let mut same = [_mm256_setzero_si256(); N];
let mut cross = [_mm256_setzero_si256(); N];
for (k, &coeff) in coeffs.iter().enumerate() {
let (same_word, cross_word) = broadcast_words(TowerCoeff::new(coeff));
same[k] = _mm256_set1_epi16(same_word);
cross[k] = _mm256_set1_epi16(cross_word);
}
let (dst_lanes, dst_tail) = dst.as_chunks_mut::<32>();
for (w, dst_lane) in dst_lanes.iter_mut().enumerate() {
let mut acc = _mm256_loadu_si256(&*dst_lane);
for (k, &src) in srcs.iter().enumerate() {
let (src_lanes, _) = src.as_chunks::<32>();
let x = _mm256_loadu_si256(&src_lanes[w]);
let swapped = _mm256_shuffle_epi8(x, swap);
acc = _mm256_xor_si256(acc, scale_gfni(x, swapped, same[k], cross[k]));
}
_mm256_storeu_si256(dst_lane, acc);
}
for (k, &coeff) in coeffs.iter().enumerate() {
mul_add_scalar(dst_tail, coeff, &srcs[k][tail_start..]);
}
}