use super::super::gfni::{mul_add_gfni_impl, mul_assign_gfni_impl, mul_into_gfni_impl};
use super::super::{BodySite, EntryBody, PeelBody};
use super::{bfactor_avx512, bmul_avx512, peel_avx512};
use crate::kernel::gf8::Prepared;
#[allow(clippy::used_underscore_binding)]
#[archmage::arcane(import_intrinsics)]
pub fn mul_add_avx512(_token: archmage::X64V4xToken, dst: &mut [u8], coeff: Prepared, src: &[u8]) {
assert_eq!(
dst.len(),
src.len(),
"mul_add_avx512: dst and src differ in length"
);
mul_add_avx512_impl::<EntryBody>(dst, coeff, src);
}
#[archmage::rite(v4x, import_intrinsics)]
pub(super) fn mul_add_avx512_impl<S: BodySite>(dst: &mut [u8], coeff: Prepared, src: &[u8]) {
let factor = bfactor_avx512(coeff.affine);
let head = peel_avx512(dst.as_ptr(), dst.len());
if head > 0 {
mul_add_gfni_impl::<false, PeelBody>(&mut dst[..head], coeff, &src[..head]);
}
let dst = &mut dst[head..];
let src = &src[head..];
let (dst_tiles, dst_mid) = dst.as_chunks_mut::<256>();
let (src_tiles, src_mid) = src.as_chunks::<256>();
for (dtile, stile) in dst_tiles.iter_mut().zip(src_tiles) {
let (d, _) = dtile.as_chunks_mut::<64>();
let (s, _) = stile.as_chunks::<64>();
let x0 = _mm512_loadu_si512(&s[0]);
let x1 = _mm512_loadu_si512(&s[1]);
let x2 = _mm512_loadu_si512(&s[2]);
let x3 = _mm512_loadu_si512(&s[3]);
let d0 = _mm512_loadu_si512(&d[0]);
let d1 = _mm512_loadu_si512(&d[1]);
let d2 = _mm512_loadu_si512(&d[2]);
let d3 = _mm512_loadu_si512(&d[3]);
let r0 = _mm512_xor_si512(d0, bmul_avx512(x0, factor));
let r1 = _mm512_xor_si512(d1, bmul_avx512(x1, factor));
let r2 = _mm512_xor_si512(d2, bmul_avx512(x2, factor));
let r3 = _mm512_xor_si512(d3, bmul_avx512(x3, factor));
_mm512_storeu_si512(&mut d[0], r0);
_mm512_storeu_si512(&mut d[1], r1);
_mm512_storeu_si512(&mut d[2], r2);
_mm512_storeu_si512(&mut d[3], r3);
}
let (dst_lanes, dst_rest) = dst_mid.as_chunks_mut::<64>();
let (src_lanes, src_rest) = src_mid.as_chunks::<64>();
for (dlane, slane) in dst_lanes.iter_mut().zip(src_lanes) {
let x = _mm512_loadu_si512(slane);
let d = _mm512_loadu_si512(dlane);
let r = _mm512_xor_si512(d, bmul_avx512(x, factor));
_mm512_storeu_si512(dlane, r);
}
mul_add_gfni_impl::<false, PeelBody>(dst_rest, coeff, src_rest);
}
#[allow(clippy::used_underscore_binding)]
#[archmage::arcane(import_intrinsics)]
pub fn mul_assign_avx512(_token: archmage::X64V4xToken, dst: &mut [u8], coeff: Prepared) {
let factor = bfactor_avx512(coeff.affine);
let head = peel_avx512(dst.as_ptr(), dst.len());
if head > 0 {
mul_assign_gfni_impl::<false, PeelBody>(&mut dst[..head], coeff);
}
let dst = &mut dst[head..];
let (tiles, rest) = dst.as_chunks_mut::<256>();
for tile in tiles {
let (d, _) = tile.as_chunks_mut::<64>();
let d0 = _mm512_loadu_si512(&d[0]);
let d1 = _mm512_loadu_si512(&d[1]);
let d2 = _mm512_loadu_si512(&d[2]);
let d3 = _mm512_loadu_si512(&d[3]);
_mm512_storeu_si512(&mut d[0], bmul_avx512(d0, factor));
_mm512_storeu_si512(&mut d[1], bmul_avx512(d1, factor));
_mm512_storeu_si512(&mut d[2], bmul_avx512(d2, factor));
_mm512_storeu_si512(&mut d[3], bmul_avx512(d3, factor));
}
let (lanes, tail) = rest.as_chunks_mut::<64>();
for lane in lanes {
let d = _mm512_loadu_si512(&*lane);
_mm512_storeu_si512(lane, bmul_avx512(d, factor));
}
mul_assign_gfni_impl::<false, PeelBody>(tail, coeff);
}
#[allow(clippy::used_underscore_binding)]
#[archmage::arcane(import_intrinsics)]
pub fn mul_into_avx512(_token: archmage::X64V4xToken, dst: &mut [u8], coeff: Prepared, src: &[u8]) {
assert_eq!(
dst.len(),
src.len(),
"mul_into_avx512: dst and src differ in length"
);
let head = peel_avx512(dst.as_ptr(), dst.len());
if head > 0 {
mul_into_gfni_impl::<false, PeelBody, false>(&mut dst[..head], coeff, &src[..head]);
}
let (dst, src) = (&mut dst[head..], &src[head..]);
let factor = bfactor_avx512(coeff.affine);
let prefetch = crate::kernel::x86::prefetch_dst(dst, false);
let tiles = dst.len() / 256 * 256;
let (mut drest, dst_mid) = dst.split_at_mut(tiles);
let (mut srest, src_mid) = src.split_at(tiles);
while !drest.is_empty() {
if prefetch && drest.len() >= crate::kernel::x86::PREFETCH_AHEAD + 256 {
crate::kernel::x86::prefetch_tile(&drest[crate::kernel::x86::PREFETCH_AHEAD..]);
crate::kernel::x86::prefetch_tile(&drest[crate::kernel::x86::PREFETCH_AHEAD + 128..]);
}
let (dtile, dnext) = drest.split_at_mut(256);
let (stile, snext) = srest.split_at(256);
let (d, _) = dtile.as_chunks_mut::<64>();
let (s, _) = stile.as_chunks::<64>();
let x0 = _mm512_loadu_si512(&s[0]);
let x1 = _mm512_loadu_si512(&s[1]);
let x2 = _mm512_loadu_si512(&s[2]);
let x3 = _mm512_loadu_si512(&s[3]);
_mm512_storeu_si512(&mut d[0], bmul_avx512(x0, factor));
_mm512_storeu_si512(&mut d[1], bmul_avx512(x1, factor));
_mm512_storeu_si512(&mut d[2], bmul_avx512(x2, factor));
_mm512_storeu_si512(&mut d[3], bmul_avx512(x3, factor));
drest = dnext;
srest = snext;
}
let (dst_lanes, dst_rest) = dst_mid.as_chunks_mut::<64>();
let (src_lanes, src_rest) = src_mid.as_chunks::<64>();
for (dlane, slane) in dst_lanes.iter_mut().zip(src_lanes) {
let x = _mm512_loadu_si512(slane);
_mm512_storeu_si512(dlane, bmul_avx512(x, factor));
}
mul_into_gfni_impl::<false, PeelBody, false>(dst_rest, coeff, src_rest);
}