use core::arch::aarch64::*;
use crate::kernel::Matrix;
use crate::kernel::gf16::{
CoeffKind, Coeffs, mul_add_scalar, mul_assign_scalar, mul_into_scalar, tail_product,
};
use crate::kernel::proven_checks::{check_equal, check_row_span, check_terms};
use crate::kernel::tables::TowerTables;
const TERM_BLOCK: usize = 8;
#[derive(Clone, Copy)]
struct Factors {
lo: [uint8x16_t; 4],
hi: [uint8x16_t; 4],
}
#[inline]
#[archmage::rite(neon)]
fn empty_factors() -> Factors {
let zero = vdupq_n_u8(0);
Factors {
lo: [zero; 4],
hi: [zero; 4],
}
}
#[inline]
#[archmage::rite(neon, import_intrinsics)]
fn load_factors(tables: &TowerTables) -> Factors {
let f = &tables.factors;
Factors {
lo: [
vld1q_u8(&f[0].lo),
vld1q_u8(&f[1].lo),
vld1q_u8(&f[2].lo),
vld1q_u8(&f[3].lo),
],
hi: [
vld1q_u8(&f[0].hi),
vld1q_u8(&f[1].hi),
vld1q_u8(&f[2].hi),
vld1q_u8(&f[3].hi),
],
}
}
#[derive(Clone, Copy)]
struct Nibbles {
source_lo: uint8x16_t,
source_hi: uint8x16_t,
swapped_lo: uint8x16_t,
swapped_hi: uint8x16_t,
}
#[inline]
#[archmage::rite(neon)]
fn split(source: uint8x16_t) -> Nibbles {
let mask = vdupq_n_u8(0x0f);
let swapped = vrev16q_u8(source);
Nibbles {
source_lo: vandq_u8(source, mask),
source_hi: vshrq_n_u8(source, 4),
swapped_lo: vandq_u8(swapped, mask),
swapped_hi: vshrq_n_u8(swapped, 4),
}
}
#[inline]
#[archmage::rite(neon)]
fn scaled(nibbles: Nibbles, factors: &Factors) -> uint8x16_t {
let Nibbles {
source_lo,
source_hi,
swapped_lo,
swapped_hi,
} = nibbles;
let direct_even = veorq_u8(
vqtbl1q_u8(factors.lo[0], source_lo),
vqtbl1q_u8(factors.hi[0], source_hi),
);
let direct_odd = veorq_u8(
vqtbl1q_u8(factors.lo[1], source_lo),
vqtbl1q_u8(factors.hi[1], source_hi),
);
let cross_even = veorq_u8(
vqtbl1q_u8(factors.lo[2], swapped_lo),
vqtbl1q_u8(factors.hi[2], swapped_hi),
);
let cross_odd = veorq_u8(
vqtbl1q_u8(factors.lo[3], swapped_lo),
vqtbl1q_u8(factors.hi[3], swapped_hi),
);
let even_lanes = vreinterpretq_u8_u16(vdupq_n_u16(0x00ff));
vbslq_u8(
even_lanes,
veorq_u8(direct_even, cross_even),
veorq_u8(direct_odd, cross_odd),
)
}
#[inline]
#[archmage::rite(neon)]
fn scaled_vector(source: uint8x16_t, factors: &Factors) -> uint8x16_t {
scaled(split(source), factors)
}
#[allow(clippy::used_underscore_binding)]
#[archmage::arcane]
pub fn mul_add_neon(_token: archmage::NeonToken, dst: &mut [u8], tables: &TowerTables, src: &[u8]) {
check_equal("gf16::mul_add_neon", "dst", dst.len(), "src", src.len());
mul_add_impl(dst, tables, src);
}
#[archmage::rite(neon, import_intrinsics)]
fn mul_add_impl(dst: &mut [u8], tables: &TowerTables, src: &[u8]) {
let span = dst.len().min(src.len());
let vector_len = span & !15;
let pair_len = vector_len & !31;
let factors = load_factors(tables);
let (dst_tiles, _) = dst[..pair_len].as_chunks_mut::<32>();
let (src_tiles, _) = src[..pair_len].as_chunks::<32>();
for (d_tile, s_tile) in dst_tiles.iter_mut().zip(src_tiles) {
let (d0, d1) = d_tile.split_at_mut(16);
let (d0, d1): (&mut [u8; 16], &mut [u8; 16]) =
(d0.try_into().unwrap(), d1.try_into().unwrap());
let (s0, s1): (&[u8; 16], &[u8; 16]) = {
let (s0, s1) = s_tile.split_at(16);
(s0.try_into().unwrap(), s1.try_into().unwrap())
};
let p0 = scaled_vector(vld1q_u8(s0), &factors);
let p1 = scaled_vector(vld1q_u8(s1), &factors);
let d0v = vld1q_u8(&*d0);
let d1v = vld1q_u8(&*d1);
vst1q_u8(d0, veorq_u8(d0v, p0));
vst1q_u8(d1, veorq_u8(d1v, p1));
}
let (dst_lanes, _) = dst[pair_len..vector_len].as_chunks_mut::<16>();
let (src_lanes, _) = src[pair_len..vector_len].as_chunks::<16>();
for (d, s) in dst_lanes.iter_mut().zip(src_lanes) {
let p = scaled_vector(vld1q_u8(s), &factors);
let cur = vld1q_u8(&*d);
vst1q_u8(d, veorq_u8(cur, p));
}
mul_add_scalar(&mut dst[vector_len..span], tables, &src[vector_len..span]);
}
#[allow(clippy::used_underscore_binding)]
#[archmage::arcane]
pub fn mul_assign_neon(_token: archmage::NeonToken, dst: &mut [u8], tables: &TowerTables) {
mul_assign_impl(dst, tables);
}
#[archmage::rite(neon, import_intrinsics)]
fn mul_assign_impl(dst: &mut [u8], tables: &TowerTables) {
let vector_len = dst.len() & !15;
let pair_len = vector_len & !31;
let factors = load_factors(tables);
let (dst_tiles, _) = dst[..pair_len].as_chunks_mut::<32>();
for d_tile in dst_tiles {
let (d0, d1) = d_tile.split_at_mut(16);
let (d0, d1): (&mut [u8; 16], &mut [u8; 16]) =
(d0.try_into().unwrap(), d1.try_into().unwrap());
let p0 = scaled_vector(vld1q_u8(&*d0), &factors);
let p1 = scaled_vector(vld1q_u8(&*d1), &factors);
vst1q_u8(d0, p0);
vst1q_u8(d1, p1);
}
let (dst_lanes, _) = dst[pair_len..vector_len].as_chunks_mut::<16>();
for d in dst_lanes {
vst1q_u8(d, scaled_vector(vld1q_u8(&*d), &factors));
}
mul_assign_scalar(&mut dst[vector_len..], tables);
}
#[allow(clippy::used_underscore_binding)]
#[archmage::arcane]
pub fn mul_into_neon(
_token: archmage::NeonToken,
dst: &mut [u8],
tables: &TowerTables,
src: &[u8],
) {
check_equal("gf16::mul_into_neon", "dst", dst.len(), "src", src.len());
mul_into_impl(dst, tables, src);
}
#[archmage::rite(neon, import_intrinsics)]
fn mul_into_impl(dst: &mut [u8], tables: &TowerTables, src: &[u8]) {
let span = dst.len().min(src.len());
let vector_len = span & !15;
let pair_len = vector_len & !31;
let factors = load_factors(tables);
let (dst_tiles, _) = dst[..pair_len].as_chunks_mut::<32>();
let (src_tiles, _) = src[..pair_len].as_chunks::<32>();
for (d_tile, s_tile) in dst_tiles.iter_mut().zip(src_tiles) {
let (d0, d1) = d_tile.split_at_mut(16);
let (d0, d1): (&mut [u8; 16], &mut [u8; 16]) =
(d0.try_into().unwrap(), d1.try_into().unwrap());
let (s0, s1): (&[u8; 16], &[u8; 16]) = {
let (s0, s1) = s_tile.split_at(16);
(s0.try_into().unwrap(), s1.try_into().unwrap())
};
vst1q_u8(d0, scaled_vector(vld1q_u8(s0), &factors));
vst1q_u8(d1, scaled_vector(vld1q_u8(s1), &factors));
}
let (dst_lanes, _) = dst[pair_len..vector_len].as_chunks_mut::<16>();
let (src_lanes, _) = src[pair_len..vector_len].as_chunks::<16>();
for (d, s) in dst_lanes.iter_mut().zip(src_lanes) {
vst1q_u8(d, scaled_vector(vld1q_u8(s), &factors));
}
mul_into_scalar(&mut dst[vector_len..span], tables, &src[vector_len..span]);
}
#[allow(unsafe_code)]
#[archmage::rite(neon)]
unsafe fn xor_row(row: *mut u8, span: usize, src: &[u8]) {
let len = span & !15;
let src_ptr = src.as_ptr();
let mut offset = 0;
while offset < len {
unsafe {
let dst_ptr = row.add(offset);
let current = vld1q_u8(dst_ptr);
let source = vld1q_u8(src_ptr.add(offset));
vst1q_u8(dst_ptr, veorq_u8(current, source));
}
offset += 16;
}
let tail = unsafe { core::slice::from_raw_parts_mut(row.add(len), span - len) };
for (d, &s) in tail.iter_mut().zip(&src[len..span]) {
*d ^= s;
}
}
#[allow(unsafe_code)]
#[archmage::rite(neon)]
unsafe fn scatter_group<const N: usize>(
base: *mut u8,
span: usize,
starts: [usize; N],
group: [TowerTables; N],
src: &[u8],
) {
let len = span & !15;
let mut factors = [empty_factors(); N];
for (factor, tables) in factors.iter_mut().zip(&group) {
*factor = load_factors(tables);
}
let src_ptr = src.as_ptr();
let mut offset = 0;
while offset < len {
unsafe {
let nibbles = split(vld1q_u8(src_ptr.add(offset)));
for (&start, factor) in starts.iter().zip(&factors) {
let row = base.add(start + offset);
vst1q_u8(row, veorq_u8(vld1q_u8(row), scaled(nibbles, factor)));
}
}
offset += 16;
}
for (&start, tables) in starts.iter().zip(&group) {
let tail = unsafe { core::slice::from_raw_parts_mut(base.add(start + len), span - len) };
mul_add_scalar(tail, tables, &src[len..span]);
}
}
#[allow(unsafe_code)]
#[allow(clippy::used_underscore_binding)]
#[archmage::arcane]
pub fn mul_add_scatter_neon(
_token: archmage::NeonToken,
rows: &mut [u8],
row_len: usize,
coeffs: &(impl Coeffs + ?Sized),
src: &[u8],
) {
check_equal(
"gf16::mul_add_scatter_neon",
"row_len",
row_len,
"src",
src.len(),
);
check_row_span(
"gf16::mul_add_scatter_neon",
rows.len(),
row_len,
coeffs.count(),
);
if row_len == 0 || coeffs.count() == 0 || src.is_empty() {
return;
}
unsafe { mul_add_scatter_impl(rows, row_len, coeffs, src) };
}
#[allow(unsafe_code)]
#[archmage::rite(neon, import_intrinsics)]
unsafe fn mul_add_scatter_impl(
rows: &mut [u8],
row_len: usize,
coeffs: &(impl Coeffs + ?Sized),
src: &[u8],
) {
let nrows = coeffs.count().min(rows.len() / row_len);
let span = row_len.min(src.len());
let base = rows.as_mut_ptr();
let mut starts = [0usize; 4];
let mut group = [TowerTables::zero(); 4];
let mut count = 0usize;
for j in 0..nrows {
if coeffs.kind(j) == CoeffKind::Zero {
continue;
}
let start = j * row_len;
if coeffs.kind(j) == CoeffKind::One {
unsafe { xor_row(base.add(start), span, src) };
continue;
}
let tables = coeffs.resolved(j);
starts[count] = start;
group[count] = tables;
count += 1;
if count == 4 {
unsafe { scatter_group::<4>(base, span, starts, group, src) };
count = 0;
}
}
match count {
1 => unsafe {
scatter_group::<1>(base, span, [starts[0]], [group[0]], src);
},
2 => unsafe {
scatter_group::<2>(
base,
span,
[starts[0], starts[1]],
[group[0], group[1]],
src,
);
},
3 => unsafe {
scatter_group::<3>(
base,
span,
[starts[0], starts[1], starts[2]],
[group[0], group[1], group[2]],
src,
);
},
_ => {}
}
}
#[derive(Clone, Copy)]
enum Mode {
Skip,
Xor,
Mul,
}
#[allow(unsafe_code)]
#[archmage::rite(neon)]
unsafe fn matrix_group<const N: usize, E>(
base: *mut u8,
span: usize,
starts: [usize; N],
g: usize,
terms: &[(&[E], &[u8])],
) where
for<'a> [(&'a [E], &'a [u8])]: Matrix<TowerTables>,
{
let len = span & !15;
for (block_index, block) in terms.chunks(TERM_BLOCK).enumerate() {
let block_start = block_index * TERM_BLOCK;
let mut cache = [[empty_factors(); N]; TERM_BLOCK];
let mut modes = [[Mode::Skip; N]; TERM_BLOCK];
for t in 0..block.len() {
for k in 0..N {
let tables = terms.coefficient(block_start + t, g + k);
modes[t][k] = match CoeffKind::of(&tables) {
CoeffKind::Zero => Mode::Skip,
CoeffKind::One => Mode::Xor,
CoeffKind::General => {
cache[t][k] = load_factors(&tables);
Mode::Mul
}
};
}
}
let mut offset = 0;
while offset < len {
let mut tile = [vdupq_n_u8(0); N];
unsafe {
for (acc, &start) in tile.iter_mut().zip(&starts) {
*acc = vld1q_u8(base.add(start + offset));
}
}
for ((&(_, src), slots), kinds) in block.iter().zip(&cache).zip(&modes) {
let source = unsafe { vld1q_u8(src.as_ptr().add(offset)) };
let nibbles = split(source);
for ((acc, factor), kind) in tile.iter_mut().zip(slots).zip(kinds) {
match kind {
Mode::Skip => {}
Mode::Xor => *acc = veorq_u8(*acc, source),
Mode::Mul => *acc = veorq_u8(*acc, scaled(nibbles, factor)),
}
}
}
unsafe {
for (acc, &start) in tile.iter().zip(&starts) {
vst1q_u8(base.add(start + offset), *acc);
}
}
offset += 16;
}
if len < span {
for (k, &start) in starts.iter().enumerate() {
let tail =
unsafe { core::slice::from_raw_parts_mut(base.add(start + len), span - len) };
for t in 0..block.len() {
let tables = terms.coefficient(block_start + t, g + k);
mul_add_scalar(tail, &tables, &block[t].1[len..span]);
}
}
}
}
}
#[allow(unsafe_code)]
#[allow(clippy::used_underscore_binding)]
#[archmage::arcane]
pub fn mul_add_matrix_neon<E>(
_token: archmage::NeonToken,
rows: &mut [u8],
row_len: usize,
nrows: usize,
terms: &[(&[E], &[u8])],
) where
for<'a> [(&'a [E], &'a [u8])]: Matrix<TowerTables>,
{
check_row_span("gf16::mul_add_matrix_neon", rows.len(), row_len, nrows);
check_terms("gf16::mul_add_matrix_neon", row_len, nrows, terms);
if row_len == 0 || nrows == 0 || terms.is_empty() {
return;
}
unsafe { mul_add_matrix_impl(rows, row_len, nrows, terms) };
}
#[allow(unsafe_code)]
#[archmage::rite(neon, import_intrinsics)]
unsafe fn mul_add_matrix_impl<E>(
rows: &mut [u8],
row_len: usize,
nrows: usize,
terms: &[(&[E], &[u8])],
) where
for<'a> [(&'a [E], &'a [u8])]: Matrix<TowerTables>,
{
let count = nrows.min(rows.len() / row_len);
let mut span = row_len;
for &(_, src) in terms {
span = span.min(src.len());
}
if count == 0 || span == 0 {
return;
}
let base = rows.as_mut_ptr();
let mut g = 0;
while g + 4 <= count {
let starts = [
g * row_len,
(g + 1) * row_len,
(g + 2) * row_len,
(g + 3) * row_len,
];
unsafe { matrix_group::<4, _>(base, span, starts, g, terms) };
g += 4;
}
match count - g {
1 => unsafe { matrix_group::<1, _>(base, span, [g * row_len], g, terms) },
2 => unsafe {
matrix_group::<2, _>(base, span, [g * row_len, (g + 1) * row_len], g, terms);
},
3 => unsafe {
matrix_group::<3, _>(
base,
span,
[g * row_len, (g + 1) * row_len, (g + 2) * row_len],
g,
terms,
);
},
_ => {}
}
}
#[allow(clippy::used_underscore_binding)]
#[archmage::arcane]
pub fn mul_add_gather_neon(
_token: archmage::NeonToken,
dst: &mut [u8],
coeffs: &(impl Coeffs + ?Sized),
srcs: &[&[u8]],
) {
check_equal(
"gf16::mul_add_gather_neon",
"coefficients",
coeffs.count(),
"sources",
srcs.len(),
);
for (index, &src) in srcs.iter().enumerate() {
check_equal(
"gf16::mul_add_gather_neon",
"dst",
dst.len(),
format_args!("source {index}"),
src.len(),
);
}
if dst.is_empty() || srcs.is_empty() {
return;
}
mul_add_gather_impl(dst, coeffs, srcs);
}
#[archmage::rite(neon, import_intrinsics)]
fn mul_add_gather_impl(dst: &mut [u8], coeffs: &(impl Coeffs + ?Sized), srcs: &[&[u8]]) {
let vector_len = dst.len() & !15;
let (dst_lanes, dst_tail) = dst.as_chunks_mut::<16>();
for block in (0..coeffs.count()).step_by(TERM_BLOCK) {
let count = (coeffs.count() - block).min(TERM_BLOCK);
let mut factors = [empty_factors(); TERM_BLOCK];
for i in 0..count {
factors[i] = load_factors(&coeffs.resolved(block + i));
}
for (t, d) in dst_lanes.iter_mut().enumerate() {
let mut acc = vld1q_u8(&*d);
for i in 0..count {
let source: &[u8; 16] = srcs[block + i][t * 16..t * 16 + 16].try_into().unwrap();
acc = veorq_u8(acc, scaled_vector(vld1q_u8(source), &factors[i]));
}
vst1q_u8(d, acc);
}
for i in 0..count {
mul_add_scalar(
dst_tail,
&coeffs.resolved(block + i),
&srcs[block + i][vector_len..],
);
}
}
}
#[inline]
#[archmage::rite(neon)]
fn multiply_base_vectors(
mut a: uint8x16_t,
mut b: uint8x16_t,
reduction: uint8x16_t,
) -> uint8x16_t {
let one = vdupq_n_u8(1);
let high_bit = vdupq_n_u8(0x80);
let mut product = vdupq_n_u8(0);
for _ in 0..8 {
let active = vceqq_u8(vandq_u8(b, one), one);
product = veorq_u8(product, vandq_u8(a, active));
let reduce = vceqq_u8(vandq_u8(a, high_bit), high_bit);
a = veorq_u8(vshlq_n_u8(a, 1), vandq_u8(reduction, reduce));
b = vshrq_n_u8(b, 1);
}
product
}
#[allow(clippy::used_underscore_binding)]
#[archmage::arcane]
pub fn mul_elementwise_neon(
_token: archmage::NeonToken,
dst: &mut [u8],
a: &[u8],
b: &[u8],
b_raw: u8,
reduction: u8,
) {
check_equal("gf16::mul_elementwise_neon", "dst", dst.len(), "a", a.len());
check_equal("gf16::mul_elementwise_neon", "dst", dst.len(), "b", b.len());
elementwise_impl(dst, a, b, b_raw, reduction);
}
#[archmage::rite(neon, import_intrinsics)]
fn elementwise_impl(dst: &mut [u8], a: &[u8], b: &[u8], b_raw: u8, reduction: u8) {
let even = vreinterpretq_u8_u16(vdupq_n_u16(0x00ff));
let delta_even = vreinterpretq_u8_u16(vdupq_n_u16(u16::from_le_bytes([b_raw, 0])));
let reduction_byte = reduction;
let reduction = vdupq_n_u8(reduction);
let (dst_lanes, dst_tail) = dst.as_chunks_mut::<16>();
let (a_lanes, a_tail) = a.as_chunks::<16>();
let (b_lanes, b_tail) = b.as_chunks::<16>();
for ((d, x), y) in dst_lanes.iter_mut().zip(a_lanes).zip(b_lanes) {
let xv = vld1q_u8(x);
let yv = vld1q_u8(y);
let direct = multiply_base_vectors(xv, yv, reduction);
let crossed = multiply_base_vectors(xv, vrev16q_u8(yv), reduction);
let delta_bd = multiply_base_vectors(vrev16q_u8(direct), delta_even, reduction);
let constant = veorq_u8(direct, delta_bd);
let extension = veorq_u8(veorq_u8(crossed, vrev16q_u8(crossed)), direct);
vst1q_u8(d, vbslq_u8(even, constant, extension));
}
for (d, (x, y)) in dst_tail
.chunks_exact_mut(2)
.zip(a_tail.chunks_exact(2).zip(b_tail.chunks_exact(2)))
{
let product = tail_product([x[0], x[1]], [y[0], y[1]], b_raw, reduction_byte);
d.copy_from_slice(&product);
}
}