use super::{TERM_TILE, check_elements, swap_mask_ssse3};
use crate::kernel::Matrix;
use crate::kernel::gf16::{
Coeffs, mul_add_scalar, mul_assign_scalar, mul_into_scalar, tail_product,
};
use crate::kernel::proven_checks::check_terms;
use crate::kernel::tables::{NibbleFactors, ScaleTable, TowerTables};
use crate::kernel::x86::gf8;
#[cfg(target_arch = "x86")]
use core::arch::x86::*;
#[cfg(target_arch = "x86_64")]
use core::arch::x86_64::*;
pub(crate) struct NibbleSsse3 {
pub(crate) lo: [__m128i; 4],
pub(crate) hi: [__m128i; 4],
pub(crate) nibble: __m128i,
pub(crate) even: __m128i,
pub(crate) swap: __m128i,
}
#[archmage::rite(v2, import_intrinsics)]
pub(crate) fn nibble_ssse3<T: NibbleFactors>(tables: &T) -> NibbleSsse3 {
let mut lo = [_mm_setzero_si128(); 4];
let mut hi = [_mm_setzero_si128(); 4];
for slot in 0..4 {
lo[slot] = _mm_loadu_si128(tables.lo(slot));
hi[slot] = _mm_loadu_si128(tables.hi(slot));
}
NibbleSsse3 {
lo,
hi,
nibble: _mm_set1_epi8(0x0f),
even: _mm_set1_epi16(0x00ff),
swap: swap_mask_ssse3(),
}
}
#[archmage::rite(v2)]
pub(crate) fn split(value: __m128i, nibble: __m128i) -> (__m128i, __m128i) {
(
_mm_and_si128(value, nibble),
_mm_and_si128(_mm_srli_epi16(value, 4), nibble),
)
}
#[archmage::rite(v2)]
pub(crate) fn lookup(lo: __m128i, hi: __m128i, split: (__m128i, __m128i)) -> __m128i {
_mm_xor_si128(_mm_shuffle_epi8(lo, split.0), _mm_shuffle_epi8(hi, split.1))
}
#[archmage::rite(v2)]
pub(crate) fn scale_ssse3(src: __m128i, tables: &NibbleSsse3) -> __m128i {
let swapped = _mm_shuffle_epi8(src, tables.swap);
let direct = split(src, tables.nibble);
let crossed = split(swapped, tables.nibble);
let even = _mm_xor_si128(
lookup(tables.lo[0], tables.hi[0], direct),
lookup(tables.lo[2], tables.hi[2], crossed),
);
let odd = _mm_xor_si128(
lookup(tables.lo[1], tables.hi[1], direct),
lookup(tables.lo[3], tables.hi[3], crossed),
);
_mm_xor_si128(
_mm_and_si128(even, tables.even),
_mm_andnot_si128(tables.even, odd),
)
}
#[allow(clippy::used_underscore_binding)]
#[archmage::arcane(import_intrinsics)]
pub fn mul_add_ssse3(
_token: archmage::X64V2Token,
dst: &mut [u8],
tables: &TowerTables,
src: &[u8],
) {
check_elements("gf16::mul_add_ssse3", dst.len());
assert_eq!(
dst.len(),
src.len(),
"gf16::mul_add_ssse3: dst is {} bytes but src is {} bytes",
dst.len(),
src.len(),
);
let vectors = nibble_ssse3(tables);
let (dst_lanes, dst_tail) = dst.as_chunks_mut::<16>();
let (src_lanes, src_tail) = src.as_chunks::<16>();
for (dst_lane, src_lane) in dst_lanes.iter_mut().zip(src_lanes) {
let x = _mm_loadu_si128(src_lane);
let d = _mm_loadu_si128(&*dst_lane);
_mm_storeu_si128(dst_lane, _mm_xor_si128(d, scale_ssse3(x, &vectors)));
}
mul_add_scalar(dst_tail, tables, src_tail);
}
#[allow(clippy::used_underscore_binding)]
#[archmage::arcane(import_intrinsics)]
pub fn mul_assign_ssse3(_token: archmage::X64V2Token, dst: &mut [u8], tables: &TowerTables) {
check_elements("gf16::mul_assign_ssse3", dst.len());
let vectors = nibble_ssse3(tables);
let (dst_lanes, dst_tail) = dst.as_chunks_mut::<16>();
for dst_lane in dst_lanes {
let x = _mm_loadu_si128(&*dst_lane);
_mm_storeu_si128(dst_lane, scale_ssse3(x, &vectors));
}
mul_assign_scalar(dst_tail, tables);
}
#[allow(clippy::used_underscore_binding)]
#[archmage::arcane(import_intrinsics)]
pub fn mul_into_ssse3(
_token: archmage::X64V2Token,
dst: &mut [u8],
tables: &TowerTables,
src: &[u8],
) {
check_elements("gf16::mul_into_ssse3", dst.len());
assert_eq!(
dst.len(),
src.len(),
"gf16::mul_into_ssse3: dst is {} bytes but src is {} bytes",
dst.len(),
src.len(),
);
let vectors = nibble_ssse3(tables);
let (dst_lanes, dst_tail) = dst.as_chunks_mut::<16>();
let (src_lanes, src_tail) = src.as_chunks::<16>();
for (dst_lane, src_lane) in dst_lanes.iter_mut().zip(src_lanes) {
let x = _mm_loadu_si128(src_lane);
_mm_storeu_si128(dst_lane, scale_ssse3(x, &vectors));
}
mul_into_scalar(dst_tail, tables, src_tail);
}
#[allow(clippy::used_underscore_binding)]
#[allow(unsafe_code)]
#[archmage::arcane(import_intrinsics)]
pub fn mul_add_scatter_ssse3(
_token: archmage::X64V2Token,
rows: &mut [u8],
row_len: usize,
coeffs: &(impl Coeffs + ?Sized),
src: &[u8],
) {
check_elements("gf16::mul_add_scatter_ssse3", row_len);
assert_eq!(
row_len,
src.len(),
"gf16::mul_add_scatter_ssse3: row_len is {row_len} bytes but src is {} bytes",
src.len(),
);
let used = coeffs
.count()
.checked_mul(row_len)
.expect("gf16::mul_add_scatter_ssse3: row geometry overflows");
assert!(
rows.len() >= used,
"gf16::mul_add_scatter_ssse3: rows is {} bytes but {} rows of {row_len} bytes need {used}",
rows.len(),
coeffs.count(),
);
if row_len == 0 || coeffs.count() == 0 {
return;
}
unsafe { scatter_rows(rows, row_len, coeffs, src) };
}
#[allow(unsafe_code)]
#[archmage::rite(v2)]
unsafe fn scatter_rows(
rows: &mut [u8],
row_len: usize,
coeffs: &(impl Coeffs + ?Sized),
src: &[u8],
) {
let base = rows.as_mut_ptr();
let src_ptr = src.as_ptr();
for group in (0..coeffs.count()).step_by(4) {
let count = (coeffs.count() - group).min(4);
let vectors: [NibbleSsse3; 4] =
core::array::from_fn(|i| nibble_ssse3(&coeffs.resolved(group + i.min(count - 1))));
let mut offset = 0;
while offset + 16 <= row_len {
unsafe {
let source = _mm_loadu_si128(src_ptr.add(offset).cast());
for (slot, vector) in vectors.iter().take(count).enumerate() {
let ptr = base.add((group + slot) * row_len + offset);
_mm_storeu_si128(
ptr.cast(),
_mm_xor_si128(_mm_loadu_si128(ptr.cast()), scale_ssse3(source, vector)),
);
}
}
offset += 16;
}
for slot in 0..count {
let tail = unsafe {
core::slice::from_raw_parts_mut(
base.add((group + slot) * row_len + offset),
row_len - offset,
)
};
mul_add_scalar(tail, &coeffs.resolved(group + slot), &src[offset..]);
}
}
}
#[allow(clippy::used_underscore_binding)]
#[archmage::arcane(import_intrinsics)]
pub fn mul_add_gather_ssse3(
_token: archmage::X64V2Token,
dst: &mut [u8],
coeffs: &(impl Coeffs + ?Sized),
srcs: &[&[u8]],
) {
check_elements("gf16::mul_add_gather_ssse3", dst.len());
assert_eq!(
coeffs.count(),
srcs.len(),
"gf16::mul_add_gather_ssse3: coefficients is {} but sources is {}",
coeffs.count(),
srcs.len(),
);
for (index, &src) in srcs.iter().enumerate() {
assert_eq!(
dst.len(),
src.len(),
"gf16::mul_add_gather_ssse3: dst is {} bytes but source {index} is {} bytes",
dst.len(),
src.len(),
);
}
if dst.is_empty() || srcs.is_empty() {
return;
}
let tail_start = dst.len() & !15;
let (dst_lanes, dst_tail) = dst.as_chunks_mut::<16>();
for block in (0..coeffs.count()).step_by(TERM_TILE) {
let count = (coeffs.count() - block).min(TERM_TILE);
let vectors: [NibbleSsse3; TERM_TILE] =
core::array::from_fn(|i| nibble_ssse3(&coeffs.resolved(block + i.min(count - 1))));
for (w, dst_lane) in dst_lanes.iter_mut().enumerate() {
let mut acc = _mm_loadu_si128(&*dst_lane);
for i in 0..count {
let (src_lanes, _) = srcs[block + i].as_chunks::<16>();
let source = _mm_loadu_si128(&src_lanes[w]);
acc = _mm_xor_si128(acc, scale_ssse3(source, &vectors[i]));
}
_mm_storeu_si128(dst_lane, acc);
}
for i in 0..count {
mul_add_scalar(
dst_tail,
&coeffs.resolved(block + i),
&srcs[block + i][tail_start..],
);
}
}
}
#[archmage::arcane(import_intrinsics)]
pub fn mul_add_matrix_ssse3<E>(
token: archmage::X64V2Token,
rows: &mut [u8],
row_len: usize,
nrows: usize,
terms: &[(&[E], &[u8])],
) where
for<'a> [(&'a [E], &'a [u8])]: Matrix<TowerTables>,
{
check_elements("gf16::mul_add_matrix_ssse3", row_len);
let used = nrows
.checked_mul(row_len)
.expect("gf16::mul_add_matrix_ssse3: row geometry overflows");
assert!(
rows.len() >= used,
"gf16::mul_add_matrix_ssse3: rows is {} bytes but {nrows} rows of {row_len} bytes need {used}",
rows.len(),
);
check_terms("gf16::mul_add_matrix_ssse3", row_len, nrows, terms);
if row_len == 0 || nrows == 0 || terms.is_empty() {
return;
}
mul_add_matrix_ssse3_with(token, rows, row_len, nrows, terms);
}
#[allow(clippy::used_underscore_binding)]
#[allow(unsafe_code)]
#[archmage::arcane(import_intrinsics)]
pub fn mul_add_matrix_ssse3_with<M: Matrix<TowerTables> + ?Sized>(
_token: archmage::X64V2Token,
rows: &mut [u8],
row_len: usize,
nrows: usize,
terms: &M,
) {
check_elements("gf16::mul_add_matrix_ssse3_with", row_len);
if row_len == 0 || nrows == 0 || terms.len() == 0 {
return;
}
let used = nrows
.checked_mul(row_len)
.expect("gf16::mul_add_matrix_ssse3_with: row geometry overflows");
assert!(
rows.len() >= used,
"gf16::mul_add_matrix_ssse3_with: rows is {} bytes but {nrows} rows of {row_len} bytes need {used}",
rows.len(),
);
let nrows = nrows.min(rows.len() / row_len);
let mut span = row_len;
for term in 0..terms.len() {
span = span.min(terms.source(term).len());
}
if nrows == 0 || span == 0 {
return;
}
unsafe { matrix_rows(rows, row_len, span, nrows, terms) };
}
#[allow(unsafe_code)]
#[archmage::rite(v2)]
unsafe fn matrix_rows<M: Matrix<TowerTables> + ?Sized>(
rows: &mut [u8],
row_len: usize,
span: usize,
nrows: usize,
terms: &M,
) {
let vector_len = span & !15;
let base = rows.as_mut_ptr();
for group in (0..nrows).step_by(4) {
let row_count = (nrows - group).min(4);
for block in (0..terms.len()).step_by(TERM_TILE) {
let term_count = (terms.len() - block).min(TERM_TILE);
let vectors: [[NibbleSsse3; 4]; TERM_TILE] =
core::array::from_fn(|t| {
core::array::from_fn(|r| {
nibble_ssse3(&terms.coefficient(
block + t.min(term_count - 1),
group + r.min(row_count - 1),
))
})
});
let mut offset = 0;
while offset < vector_len {
unsafe {
let mut acc = [_mm_setzero_si128(); 4];
for (r, slot) in acc.iter_mut().take(row_count).enumerate() {
*slot = _mm_loadu_si128(base.add((group + r) * row_len + offset).cast());
}
for (t, vector) in vectors.iter().take(term_count).enumerate() {
let source =
_mm_loadu_si128(terms.source(block + t).as_ptr().add(offset).cast());
for (value, scale) in acc.iter_mut().zip(vector).take(row_count) {
*value = _mm_xor_si128(*value, scale_ssse3(source, scale));
}
}
for (r, &slot) in acc.iter().take(row_count).enumerate() {
_mm_storeu_si128(base.add((group + r) * row_len + offset).cast(), slot);
}
}
offset += 16;
}
for r in 0..row_count {
let tail = unsafe {
core::slice::from_raw_parts_mut(
base.add((group + r) * row_len + vector_len),
span - vector_len,
)
};
for term in block..block + term_count {
mul_add_scalar(
tail,
&terms.coefficient(term, group + r),
&terms.source(term)[vector_len..],
);
}
}
}
}
}
#[archmage::rite(v2)]
fn scale_delta(v: __m128i, lo: __m128i, hi: __m128i, nibble: __m128i) -> __m128i {
_mm_xor_si128(
_mm_shuffle_epi8(lo, _mm_and_si128(v, nibble)),
_mm_shuffle_epi8(hi, _mm_and_si128(_mm_srli_epi16::<4>(v), nibble)),
)
}
#[allow(clippy::used_underscore_binding)]
#[archmage::arcane(import_intrinsics)]
pub fn mul_elementwise_ssse3(
_token: archmage::X64V2Token,
dst: &mut [u8],
a: &[u8],
b: &[u8],
delta: &'static ScaleTable,
reduction: u8,
) {
check_elements("gf16::mul_elementwise_ssse3", dst.len());
assert_eq!(
dst.len(),
a.len(),
"gf16::mul_elementwise_ssse3: dst is {} bytes but a is {} bytes",
dst.len(),
a.len(),
);
assert_eq!(
dst.len(),
b.len(),
"gf16::mul_elementwise_ssse3: dst is {} bytes but b is {} bytes",
dst.len(),
b.len(),
);
let swap = swap_mask_ssse3();
let even = _mm_set1_epi16(0x00ff);
let nibble = _mm_set1_epi8(0x0f);
let (delta_lo, delta_hi) = (_mm_loadu_si128(&delta.lo), _mm_loadu_si128(&delta.hi));
let (a_lanes, a_rest) = a.as_chunks::<16>();
let (b_lanes, b_rest) = b.as_chunks::<16>();
let (dst_lanes, dst_rest) = dst.as_chunks_mut::<16>();
for ((dst_lane, x_lane), y_lane) in dst_lanes.iter_mut().zip(a_lanes).zip(b_lanes) {
let x = _mm_loadu_si128(x_lane);
let y = _mm_loadu_si128(y_lane);
let direct = gf8::multiply_vectors_ssse3(x, y, reduction);
let crossed = gf8::multiply_vectors_ssse3(x, _mm_shuffle_epi8(y, swap), reduction);
let delta_bd = scale_delta(_mm_shuffle_epi8(direct, swap), delta_lo, delta_hi, nibble);
let constant = _mm_xor_si128(direct, delta_bd);
let cross_sum = _mm_xor_si128(crossed, _mm_shuffle_epi8(crossed, swap));
let extension = _mm_xor_si128(cross_sum, direct);
let product = _mm_xor_si128(
_mm_and_si128(constant, even),
_mm_andnot_si128(even, extension),
);
_mm_storeu_si128(dst_lane, product);
}
for (d, (x, y)) in dst_rest
.chunks_exact_mut(2)
.zip(a_rest.chunks_exact(2).zip(b_rest.chunks_exact(2)))
{
d.copy_from_slice(&tail_product(
[x[0], x[1]],
[y[0], y[1]],
delta.coeff,
reduction,
));
}
}
#[allow(clippy::used_underscore_binding)]
#[archmage::arcane(import_intrinsics)]
pub fn mul_elementwise_assign_ssse3(
_token: archmage::X64V2Token,
dst: &mut [u8],
src: &[u8],
delta: &'static ScaleTable,
reduction: u8,
) {
check_elements("gf16::mul_elementwise_assign_ssse3", dst.len());
assert_eq!(
dst.len(),
src.len(),
"gf16::mul_elementwise_assign_ssse3: dst is {} bytes but src is {} bytes",
dst.len(),
src.len(),
);
let swap = swap_mask_ssse3();
let even = _mm_set1_epi16(0x00ff);
let nibble = _mm_set1_epi8(0x0f);
let (delta_lo, delta_hi) = (_mm_loadu_si128(&delta.lo), _mm_loadu_si128(&delta.hi));
let (src_lanes, src_rest) = src.as_chunks::<16>();
let (dst_lanes, dst_rest) = dst.as_chunks_mut::<16>();
for (dst_lane, y_lane) in dst_lanes.iter_mut().zip(src_lanes) {
let x = _mm_loadu_si128(&*dst_lane);
let y = _mm_loadu_si128(y_lane);
let direct = gf8::multiply_vectors_ssse3(x, y, reduction);
let crossed = gf8::multiply_vectors_ssse3(x, _mm_shuffle_epi8(y, swap), reduction);
let delta_bd = scale_delta(_mm_shuffle_epi8(direct, swap), delta_lo, delta_hi, nibble);
let constant = _mm_xor_si128(direct, delta_bd);
let cross_sum = _mm_xor_si128(crossed, _mm_shuffle_epi8(crossed, swap));
let extension = _mm_xor_si128(cross_sum, direct);
let product = _mm_xor_si128(
_mm_and_si128(constant, even),
_mm_andnot_si128(even, extension),
);
_mm_storeu_si128(dst_lane, product);
}
for (d, y) in dst_rest.chunks_exact_mut(2).zip(src_rest.chunks_exact(2)) {
let product = tail_product([d[0], d[1]], [y[0], y[1]], delta.coeff, reduction);
d.copy_from_slice(&product);
}
}