pub const Q: i32 = 3329;
pub const N: usize = 256;
const fn pow_mod(mut b: i64, mut e: u32, m: i64) -> i64 {
let mut r = 1i64;
b %= m;
while e > 0 {
if e & 1 == 1 {
r = r * b % m;
}
b = b * b % m;
e >>= 1;
}
r
}
const fn brv7(x: u32) -> u32 {
let mut r = 0;
let mut i = 0;
while i < 7 {
r |= ((x >> i) & 1) << (6 - i);
i += 1;
}
r
}
pub const ZETAS: [i32; 128] = {
let mut z = [0i32; 128];
let mut k = 0;
while k < 128 {
z[k] = pow_mod(17, brv7(k as u32), Q as i64) as i32;
k += 1;
}
z
};
pub const ZETAS_INV: [i32; 128] = {
let mut z = [0i32; 128];
let mut k = 0;
while k < 128 {
z[k] = pow_mod(17, 128 - brv7(k as u32), Q as i64) as i32;
k += 1;
}
z
};
pub const GAMMAS: [i32; 128] = {
let mut z = [0i32; 128];
let mut k = 0;
while k < 128 {
z[k] = pow_mod(17, 2 * brv7(k as u32) + 1, Q as i64) as i32;
k += 1;
}
z
};
pub struct Field {
pub p: i64,
pub zetas: [i64; 128],
pub zetas_inv: [i64; 128],
pub gammas: [i64; 128],
pub n_inv: i64,
}
impl Field {
pub const fn with_root(p: i64, zeta: i64) -> Self {
let (mut zetas, mut zetas_inv, mut gammas) = ([0i64; 128], [0i64; 128], [0i64; 128]);
let mut k = 0;
while k < 128 {
let b = brv7(k as u32);
zetas[k] = pow_mod(zeta, b, p);
zetas_inv[k] = pow_mod(zeta, 128 - b, p);
gammas[k] = pow_mod(zeta, 2 * b + 1, p);
k += 1;
}
Self { p, zetas, zetas_inv, gammas, n_inv: pow_mod(128, (p - 2) as u32, p) }
}
pub const fn new(p: i64) -> Self {
let mut g = 2;
loop {
let z = pow_mod(g, ((p - 1) / 256) as u32, p);
if pow_mod(z, 128, p) == p - 1 {
return Self::with_root(p, z);
}
g += 1;
}
}
pub fn ntt(&self, a: &mut [i64; N]) {
let p = self.p;
let mut k = 1;
let mut len = 128;
while len >= 2 {
for start in (0..N).step_by(2 * len) {
let z = self.zetas[k];
k += 1;
for j in start..start + len {
let t = (z * a[j + len]).rem_euclid(p);
a[j + len] = (a[j] - t).rem_euclid(p);
a[j] = (a[j] + t).rem_euclid(p);
}
}
len >>= 1;
}
}
pub fn intt(&self, a: &mut [i64; N]) {
let p = self.p;
let mut len = 2;
while len <= 128 {
for (b, start) in (0..N).step_by(2 * len).enumerate() {
let z = self.zetas_inv[128 / len + b];
for j in start..start + len {
let t = a[j];
a[j] = (t + a[j + len]).rem_euclid(p);
a[j + len] = (z * (a[j + len] - t)).rem_euclid(p);
}
}
len <<= 1;
}
for v in a.iter_mut() {
*v = (*v * self.n_inv).rem_euclid(p);
}
}
pub fn basemul(&self, a: &[i64; N], b: &[i64; N]) -> [i64; N] {
let p = self.p;
let mut c = [0i64; N];
for i in 0..128 {
let (a0, a1, b0, b1) = (a[2 * i], a[2 * i + 1], b[2 * i], b[2 * i + 1]);
c[2 * i] = (a0 * b0 + (a1 * b1).rem_euclid(p) * self.gammas[i]).rem_euclid(p);
c[2 * i + 1] = (a0 * b1 + a1 * b0).rem_euclid(p);
}
c
}
pub fn negacyclic(&self, a: &[i64; N], b: &[i64; N]) -> [i64; N] {
let (mut x, mut y) = (*a, *b);
for v in x.iter_mut().chain(y.iter_mut()) {
*v = v.rem_euclid(self.p);
}
self.ntt(&mut x);
self.ntt(&mut y);
let mut c = self.basemul(&x, &y);
self.intt(&mut c);
c
}
}
pub const F3329: Field = Field::with_root(3329, 17);
pub const F7681: Field = Field::new(7681);
pub const F12289: Field = Field::new(12289);
pub fn schoolbook(a: &[i64; N], b: &[i64; N]) -> [i64; N] {
let mut c = [0i64; N];
for i in 0..N {
for j in 0..N {
let v = a[i] * b[j];
if i + j < N {
c[i + j] += v;
} else {
c[i + j - N] -= v;
}
}
}
c
}
pub fn crt(r: &[i64], fields: &[&Field]) -> i128 {
let mut x = 0i128;
let mut m = 1i128;
for (&ri, f) in r.iter().zip(fields) {
let p = i128::from(f.p);
let inv = i128::from(pow_mod((m % p) as i64, (f.p - 2) as u32, f.p));
let t = ((i128::from(ri) - x) % p + p) % p * inv % p;
x += t * m;
m *= p;
}
if x > m / 2 {
x - m
} else {
x
}
}
pub fn reverse_negacyclic(a: &[i16; N]) -> [i16; N] {
core::array::from_fn(|i| if i == 0 { a[0] } else { -a[N - i] })
}
const QINV: i32 = {
let mut x: i64 = 1;
let mut i = 0;
while i < 5 {
x = x * (2 - Q as i64 * x) % 65536;
i += 1;
}
x.rem_euclid(65536) as i32
};
const fn centered(x: i64) -> i16 {
let r = ((x % Q as i64) + Q as i64) % Q as i64;
(if r > (Q as i64) / 2 { r - Q as i64 } else { r }) as i16
}
const fn mont(z: i32) -> i16 {
centered(z as i64 * 65536)
}
#[cfg(target_arch = "x86_64")]
const fn mont_qinv(z: i32) -> i16 {
(mont(z) as i32).wrapping_mul(QINV) as i16
}
pub const MONT: i64 = pow_mod(2, 16, Q as i64);
pub const MONT_INV: i64 = pow_mod(MONT, (Q - 2) as u32, Q as i64);
pub fn ntt_ref(a: &[i16; N]) -> [i32; N] {
let mut r: [i64; N] = core::array::from_fn(|i| i64::from(a[i]));
let mut k = 1;
let mut len = 128;
while len >= 2 {
let mut start = 0;
while start < N {
let z = i64::from(ZETAS[k]);
k += 1;
for j in start..start + len {
let t = (z * r[j + len]).rem_euclid(Q as i64);
r[j + len] = (r[j] - t).rem_euclid(Q as i64);
r[j] = (r[j] + t).rem_euclid(Q as i64);
}
start += 2 * len;
}
len >>= 1;
}
core::array::from_fn(|i| r[i] as i32)
}
pub fn ntt_scalar(a: &mut [i16; N]) {
let fq = |b: i16, z: i32| -> i16 {
let p = i32::from(b) * i32::from(mont(z));
let u = (p as i16 as i32).wrapping_mul(QINV) as i16;
((p - i32::from(u) * Q) >> 16) as i16
};
let mut k = 1;
let mut len = 128;
while len >= 2 {
let mut start = 0;
while start < N {
let z = ZETAS[k];
k += 1;
for j in start..start + len {
let t = fq(a[j + len], z);
a[j + len] = a[j] - t;
a[j] += t;
}
start += 2 * len;
}
len >>= 1;
}
}
pub fn ntt(a: &mut [i16; N]) {
#[cfg(target_arch = "x86_64")]
if slain_rt::cpu::avx2() {
unsafe { x86::ntt(a) };
return;
}
ntt_scalar(a);
}
pub fn intt(a: &mut [i16; N]) {
#[cfg(target_arch = "x86_64")]
if slain_rt::cpu::avx2() {
unsafe { x86::intt(a) };
return;
}
intt_scalar(a);
}
pub fn intt_scalar(a: &mut [i16; N]) {
let mut v: [i64; N] = core::array::from_fn(|i| i64::from(a[i]).rem_euclid(Q as i64));
F3329.intt(&mut v);
*a = core::array::from_fn(|i| (v[i] * MONT % Q as i64) as i16);
}
pub fn basemul(a: &[i16; N], b: &[i16; N], c: &mut [i16; N]) {
#[cfg(target_arch = "x86_64")]
if slain_rt::cpu::avx2() {
unsafe { x86::basemul(a, b, c) };
return;
}
basemul_scalar(a, b, c);
}
pub fn basemul_scalar(a: &[i16; N], b: &[i16; N], c: &mut [i16; N]) {
let x: [i64; N] = core::array::from_fn(|i| i64::from(a[i]).rem_euclid(Q as i64));
let y: [i64; N] = core::array::from_fn(|i| i64::from(b[i]).rem_euclid(Q as i64));
let v = F3329.basemul(&x, &y);
*c = core::array::from_fn(|i| (v[i] * MONT_INV % Q as i64) as i16);
}
pub enum Coeff {}
pub enum Ntt {}
pub enum Product {}
pub struct Poly<D> {
a: [i16; N],
domain: core::marker::PhantomData<fn() -> D>,
}
impl<D> Clone for Poly<D> {
fn clone(&self) -> Self {
*self
}
}
impl<D> Copy for Poly<D> {}
impl<D> Poly<D> {
const fn wrap(a: [i16; N]) -> Self {
Self { a, domain: core::marker::PhantomData }
}
#[must_use]
pub const fn as_array(&self) -> &[i16; N] {
&self.a
}
#[must_use]
pub const fn into_array(self) -> [i16; N] {
self.a
}
}
impl Poly<Coeff> {
#[must_use]
pub fn new(a: [i16; N]) -> Option<Self> {
a.iter().all(|&x| i32::from(x).abs() < Q).then(|| Self::wrap(a))
}
#[must_use]
pub fn from_u8(a: &[u8; N]) -> Self {
Self::wrap(a.map(i16::from))
}
#[must_use]
pub fn reduce(a: [i16; N]) -> Self {
Self::wrap(a.map(|x| centered(i64::from(x))))
}
#[must_use]
pub fn ntt(self) -> Poly<Ntt> {
let mut a = self.a;
ntt(&mut a);
Poly::wrap(a)
}
}
impl Poly<Ntt> {
#[must_use]
pub fn mul(&self, b: &Self) -> Poly<Product> {
let mut c = [0i16; N];
#[cfg(target_arch = "x86_64")]
if slain_rt::cpu::avx2() {
unsafe { x86::basemul_bounded(&self.a, &b.a, &mut c) };
return Poly::wrap(c);
}
basemul_scalar(&self.a, &b.a, &mut c);
Poly::wrap(c)
}
}
impl Poly<Product> {
#[must_use]
pub fn intt(self) -> Poly<Coeff> {
let mut a = self.a;
intt(&mut a);
Poly::wrap(a)
}
}
#[cfg(target_arch = "x86_64")]
pub mod x86 {
use super::{mont, mont_qinv, GAMMAS, N, Q, QINV, ZETAS, ZETAS_INV};
use core::arch::x86_64::*;
#[repr(C, align(32))]
struct Lanes([i16; 16]);
struct Tables {
l8: [[Lanes; 2]; 8],
l4: [[Lanes; 2]; 8],
l2: [[Lanes; 2]; 8],
}
const fn pair(zs: &[i32; 128], f: [usize; 16]) -> [Lanes; 2] {
let mut z = [0i16; 16];
let mut zq = [0i16; 16];
let mut i = 0;
while i < 16 {
z[i] = mont(zs[f[i]]);
zq[i] = mont_qinv(zs[f[i]]);
i += 1;
}
[Lanes(z), Lanes(zq)]
}
const fn build(zs: &[i32; 128], layer: u32) -> [[Lanes; 2]; 8] {
let mut out = [pair(zs, [0; 16]), pair(zs, [0; 16]), pair(zs, [0; 16]), pair(zs, [0; 16]), pair(zs, [0; 16]), pair(zs, [0; 16]), pair(zs, [0; 16]), pair(zs, [0; 16])];
let mut p = 0;
while p < 8 {
let g = 2 * p;
let mut f = [0usize; 16];
let mut i = 0;
while i < 16 {
f[i] = match layer {
8 => 16 + g + i / 8,
4 => 32 + 2 * g + i / 4,
_ => 64 + 4 * g + i / 2,
};
i += 1;
}
out[p] = pair(zs, f);
p += 1;
}
out
}
const TABLES: Tables = Tables { l8: build(&ZETAS, 8), l4: build(&ZETAS, 4), l2: build(&ZETAS, 2) };
const TABLES_INV: Tables = Tables { l8: build(&ZETAS_INV, 8), l4: build(&ZETAS_INV, 4), l2: build(&ZETAS_INV, 2) };
const GAMMA_LANES: [[Lanes; 2]; 16] = {
let mut out = [const { [Lanes([0; 16]), Lanes([0; 16])] }; 16];
let mut v = 0;
while v < 16 {
let mut m = 0;
while m < 8 {
out[v][0].0[2 * m] = mont(GAMMAS[8 * v + m]);
out[v][1].0[2 * m] = mont_qinv(GAMMAS[8 * v + m]);
m += 1;
}
v += 1;
}
out
};
#[derive(Clone, Copy)]
enum Row {}
#[derive(Clone, Copy)]
enum Dist8 {}
#[derive(Clone, Copy)]
enum Dist4 {}
#[derive(Clone, Copy)]
enum Dist2 {}
#[derive(Clone, Copy)]
#[repr(transparent)]
struct V<L>(__m256i, core::marker::PhantomData<L>);
type Row16 = V<Row>;
type Span8 = V<Dist8>;
type Span4 = V<Dist4>;
type Span2 = V<Dist2>;
#[derive(Clone, Copy)]
struct Tw<L>(__m256i, __m256i, core::marker::PhantomData<L>);
trait Perm128 {
type To;
}
impl Perm128 for Row {
type To = Dist8;
}
impl Perm128 for Dist8 {
type To = Row;
}
trait Unpack64 {
type To;
}
impl Unpack64 for Dist8 {
type To = Dist4;
}
impl Unpack64 for Dist4 {
type To = Dist8;
}
trait Blend32 {
type To;
}
impl Blend32 for Dist4 {
type To = Dist2;
}
impl Blend32 for Dist2 {
type To = Dist4;
}
trait Layer {
fn of(t: &Tables) -> &[[Lanes; 2]; 8];
}
impl Layer for Dist8 {
fn of(t: &Tables) -> &[[Lanes; 2]; 8] {
&t.l8
}
}
impl Layer for Dist4 {
fn of(t: &Tables) -> &[[Lanes; 2]; 8] {
&t.l4
}
}
impl Layer for Dist2 {
fn of(t: &Tables) -> &[[Lanes; 2]; 8] {
&t.l2
}
}
impl<L> V<L> {
#[inline(always)]
const fn of(v: __m256i) -> Self {
Self(v, core::marker::PhantomData)
}
}
impl Row16 {
#[inline(always)]
unsafe fn take(r: &__m256i) -> Self {
Self::of(*r)
}
#[inline(always)]
unsafe fn put(self, r: &mut __m256i) {
*r = self.0;
}
}
#[inline(always)]
unsafe fn tw_of<L: Layer>(t: &Tables, i: usize) -> Tw<L> {
let [z, zq] = &L::of(t)[i];
Tw(ld(z), ld(zq), core::marker::PhantomData)
}
#[inline(always)]
unsafe fn perm<L: Perm128>(a: V<L>, b: V<L>) -> (V<L::To>, V<L::To>) {
let (x, y) = perm_raw(a.0, b.0);
(V::of(x), V::of(y))
}
#[inline(always)]
unsafe fn unpack<L: Unpack64>(a: V<L>, b: V<L>) -> (V<L::To>, V<L::To>) {
let (x, y) = unp(a.0, b.0);
(V::of(x), V::of(y))
}
#[inline(always)]
unsafe fn blend<L: Blend32>(a: V<L>, b: V<L>) -> (V<L::To>, V<L::To>) {
let (x, y) = bl(a.0, b.0);
(V::of(x), V::of(y))
}
#[inline(always)]
unsafe fn fqv<L>(b: V<L>, w: Tw<L>, q: __m256i) -> V<L> {
V::of(fq(b.0, w.0, w.1, q))
}
#[inline(always)]
unsafe fn addv<L>(a: V<L>, b: V<L>) -> V<L> {
V::of(_mm256_add_epi16(a.0, b.0))
}
#[inline(always)]
unsafe fn subv<L>(a: V<L>, b: V<L>) -> V<L> {
V::of(_mm256_sub_epi16(a.0, b.0))
}
#[inline(always)]
unsafe fn bfv<L>(a: &mut V<L>, b: &mut V<L>, w: Tw<L>, q: __m256i) {
bf(&mut a.0, &mut b.0, w.0, w.1, q);
}
#[inline(always)]
unsafe fn gsv<L>(a: &mut V<L>, b: &mut V<L>, w: Tw<L>, q: __m256i) {
gs(&mut a.0, &mut b.0, w.0, w.1, q);
}
#[inline(always)]
unsafe fn fq(b: __m256i, z: __m256i, zq: __m256i, q: __m256i) -> __m256i {
_mm256_sub_epi16(_mm256_mulhi_epi16(b, z), _mm256_mulhi_epi16(_mm256_mullo_epi16(b, zq), q))
}
#[inline(always)]
unsafe fn bf(a: &mut __m256i, b: &mut __m256i, z: __m256i, zq: __m256i, q: __m256i) {
let t = fq(*b, z, zq, q);
*b = _mm256_sub_epi16(*a, t);
*a = _mm256_add_epi16(*a, t);
}
#[inline(always)]
unsafe fn ld(t: &Lanes) -> __m256i {
_mm256_load_si256(t.0.as_ptr().cast())
}
#[inline(always)]
unsafe fn tw(k: usize) -> (__m256i, __m256i) {
(_mm256_set1_epi16(mont(ZETAS[k])), _mm256_set1_epi16(mont_qinv(ZETAS[k])))
}
#[target_feature(enable = "avx2")]
pub unsafe fn ntt(a: &mut [i16; N]) {
let p = a.as_mut_ptr().cast::<__m256i>();
let q = _mm256_set1_epi16(Q as i16);
let (z, zq) = tw(1);
for k in 0..8 {
let mut x = _mm256_loadu_si256(p.add(k));
let mut y = _mm256_loadu_si256(p.add(k + 8));
bf(&mut x, &mut y, z, zq, q);
_mm256_storeu_si256(p.add(k), x);
_mm256_storeu_si256(p.add(k + 8), y);
}
half::<0, true>(p, q);
half::<1, true>(p, q);
}
#[cfg(test)]
#[target_feature(enable = "avx2")]
pub unsafe fn ntt_prev(a: &mut [i16; N]) {
let p = a.as_mut_ptr().cast::<__m256i>();
let q = _mm256_set1_epi16(Q as i16);
let (z, zq) = tw(1);
for k in 0..8 {
let mut x = _mm256_loadu_si256(p.add(k));
let mut y = _mm256_loadu_si256(p.add(k + 8));
bf(&mut x, &mut y, z, zq, q);
_mm256_storeu_si256(p.add(k), x);
_mm256_storeu_si256(p.add(k + 8), y);
}
half::<0, false>(p, q);
half::<1, false>(p, q);
}
#[inline(always)]
unsafe fn half<const H: usize, const PAIRED: bool>(p: *mut __m256i, q: __m256i) {
let b = p.add(8 * H);
let (mut r0, mut r1, mut r2, mut r3) = (_mm256_loadu_si256(b), _mm256_loadu_si256(b.add(1)), _mm256_loadu_si256(b.add(2)), _mm256_loadu_si256(b.add(3)));
let (mut r4, mut r5, mut r6, mut r7) = (_mm256_loadu_si256(b.add(4)), _mm256_loadu_si256(b.add(5)), _mm256_loadu_si256(b.add(6)), _mm256_loadu_si256(b.add(7)));
let (z, zq) = tw(2 + H);
bf(&mut r0, &mut r4, z, zq, q);
bf(&mut r1, &mut r5, z, zq, q);
bf(&mut r2, &mut r6, z, zq, q);
bf(&mut r3, &mut r7, z, zq, q);
let (z, zq) = tw(4 + 2 * H);
bf(&mut r0, &mut r2, z, zq, q);
bf(&mut r1, &mut r3, z, zq, q);
let (z, zq) = tw(5 + 2 * H);
bf(&mut r4, &mut r6, z, zq, q);
bf(&mut r5, &mut r7, z, zq, q);
let (z, zq) = tw(8 + 4 * H);
bf(&mut r0, &mut r1, z, zq, q);
let (z, zq) = tw(9 + 4 * H);
bf(&mut r2, &mut r3, z, zq, q);
let (z, zq) = tw(10 + 4 * H);
bf(&mut r4, &mut r5, z, zq, q);
let (z, zq) = tw(11 + 4 * H);
bf(&mut r6, &mut r7, z, zq, q);
if PAIRED {
tail2(&mut r0, &mut r1, 4 * H, &mut r2, &mut r3, 4 * H + 1, q);
tail2(&mut r4, &mut r5, 4 * H + 2, &mut r6, &mut r7, 4 * H + 3, q);
} else {
tail(&mut r0, &mut r1, 4 * H, q);
tail(&mut r2, &mut r3, 4 * H + 1, q);
tail(&mut r4, &mut r5, 4 * H + 2, q);
tail(&mut r6, &mut r7, 4 * H + 3, q);
}
_mm256_storeu_si256(b, r0);
_mm256_storeu_si256(b.add(1), r1);
_mm256_storeu_si256(b.add(2), r2);
_mm256_storeu_si256(b.add(3), r3);
_mm256_storeu_si256(b.add(4), r4);
_mm256_storeu_si256(b.add(5), r5);
_mm256_storeu_si256(b.add(6), r6);
_mm256_storeu_si256(b.add(7), r7);
}
#[inline(always)]
unsafe fn twi(k: usize) -> (__m256i, __m256i) {
(_mm256_set1_epi16(mont(ZETAS_INV[k])), _mm256_set1_epi16(mont_qinv(ZETAS_INV[k])))
}
#[inline(always)]
unsafe fn gs(a: &mut __m256i, b: &mut __m256i, z: __m256i, zq: __m256i, q: __m256i) {
let t = *a;
*a = _mm256_add_epi16(t, *b);
*b = fq(_mm256_sub_epi16(*b, t), z, zq, q);
}
#[inline(always)]
unsafe fn barrett(a: __m256i, q: __m256i) -> __m256i {
let t = _mm256_mulhrs_epi16(_mm256_mulhi_epi16(a, _mm256_set1_epi16(20159)), _mm256_set1_epi16(32));
_mm256_sub_epi16(a, _mm256_mullo_epi16(t, q))
}
#[inline(always)]
unsafe fn montmul(a: __m256i, b: __m256i, q: __m256i) -> __m256i {
let u = _mm256_mullo_epi16(_mm256_mullo_epi16(a, b), _mm256_set1_epi16(QINV as i16));
_mm256_sub_epi16(_mm256_mulhi_epi16(a, b), _mm256_mulhi_epi16(u, q))
}
#[inline(always)]
unsafe fn perm_raw(a: __m256i, b: __m256i) -> (__m256i, __m256i) {
(_mm256_permute2x128_si256::<0x20>(a, b), _mm256_permute2x128_si256::<0x31>(a, b))
}
#[inline(always)]
unsafe fn unp(a: __m256i, b: __m256i) -> (__m256i, __m256i) {
(_mm256_unpacklo_epi64(a, b), _mm256_unpackhi_epi64(a, b))
}
#[inline(always)]
unsafe fn bl(a: __m256i, b: __m256i) -> (__m256i, __m256i) {
(_mm256_blend_epi32::<0xaa>(a, _mm256_slli_epi64::<32>(b)), _mm256_blend_epi32::<0xaa>(_mm256_srli_epi64::<32>(a), b))
}
#[inline(always)]
unsafe fn itail(r0: &mut __m256i, r1: &mut __m256i, t: usize, q: __m256i) {
let (a8, b8): (Span8, Span8) = perm(Row16::take(r0), Row16::take(r1));
let (x4, y4): (Span4, Span4) = unpack(a8, b8);
let (mut p2, mut q2): (Span2, Span2) = blend(x4, y4);
gsv(&mut p2, &mut q2, tw_of(&TABLES_INV, t), q);
let (mut x4, mut y4): (Span4, Span4) = blend(p2, q2);
gsv(&mut x4, &mut y4, tw_of(&TABLES_INV, t), q);
let (mut a8, mut b8): (Span8, Span8) = unpack(x4, y4);
gsv(&mut a8, &mut b8, tw_of(&TABLES_INV, t), q);
let (o0, o1): (Row16, Row16) = perm(a8, b8);
o0.put(r0);
o1.put(r1);
}
#[inline(always)]
unsafe fn ihalf<const H: usize>(p: *mut __m256i, q: __m256i) {
let b = p.add(8 * H);
let (mut r0, mut r1, mut r2, mut r3) = (barrett(_mm256_loadu_si256(b), q), barrett(_mm256_loadu_si256(b.add(1)), q), barrett(_mm256_loadu_si256(b.add(2)), q), barrett(_mm256_loadu_si256(b.add(3)), q));
let (mut r4, mut r5, mut r6, mut r7) = (barrett(_mm256_loadu_si256(b.add(4)), q), barrett(_mm256_loadu_si256(b.add(5)), q), barrett(_mm256_loadu_si256(b.add(6)), q), barrett(_mm256_loadu_si256(b.add(7)), q));
itail(&mut r0, &mut r1, 4 * H, q);
itail(&mut r2, &mut r3, 4 * H + 1, q);
itail(&mut r4, &mut r5, 4 * H + 2, q);
itail(&mut r6, &mut r7, 4 * H + 3, q);
(r0, r1, r2, r3) = (barrett(r0, q), barrett(r1, q), barrett(r2, q), barrett(r3, q));
(r4, r5, r6, r7) = (barrett(r4, q), barrett(r5, q), barrett(r6, q), barrett(r7, q));
let (z, zq) = twi(8 + 4 * H);
gs(&mut r0, &mut r1, z, zq, q);
let (z, zq) = twi(9 + 4 * H);
gs(&mut r2, &mut r3, z, zq, q);
let (z, zq) = twi(10 + 4 * H);
gs(&mut r4, &mut r5, z, zq, q);
let (z, zq) = twi(11 + 4 * H);
gs(&mut r6, &mut r7, z, zq, q);
let (z, zq) = twi(4 + 2 * H);
gs(&mut r0, &mut r2, z, zq, q);
gs(&mut r1, &mut r3, z, zq, q);
let (z, zq) = twi(5 + 2 * H);
gs(&mut r4, &mut r6, z, zq, q);
gs(&mut r5, &mut r7, z, zq, q);
let (z, zq) = twi(2 + H);
gs(&mut r0, &mut r4, z, zq, q);
gs(&mut r1, &mut r5, z, zq, q);
gs(&mut r2, &mut r6, z, zq, q);
gs(&mut r3, &mut r7, z, zq, q);
for (i, v) in [r0, r1, r2, r3, r4, r5, r6, r7].into_iter().enumerate() {
_mm256_storeu_si256(b.add(i), barrett(v, q));
}
}
#[target_feature(enable = "avx2")]
pub unsafe fn intt(a: &mut [i16; N]) {
let p = a.as_mut_ptr().cast::<__m256i>();
let q = _mm256_set1_epi16(Q as i16);
ihalf::<0>(p, q);
ihalf::<1>(p, q);
const ZS: i32 = ((ZETAS_INV[1] as i64 * 512).rem_euclid(Q as i64)) as i32;
let (zs, zsq) = (_mm256_set1_epi16(mont(ZS)), _mm256_set1_epi16(mont_qinv(ZS)));
let (s, sq) = (_mm256_set1_epi16(mont(512)), _mm256_set1_epi16(mont_qinv(512)));
for k in 0..8 {
let x = _mm256_loadu_si256(p.add(k));
let y = _mm256_loadu_si256(p.add(k + 8));
_mm256_storeu_si256(p.add(k), fq(_mm256_add_epi16(x, y), s, sq, q));
_mm256_storeu_si256(p.add(k + 8), fq(_mm256_sub_epi16(y, x), zs, zsq, q));
}
}
#[cfg(test)]
#[target_feature(enable = "avx2")]
pub unsafe fn intt_prev(a: &mut [i16; N]) {
let p = a.as_mut_ptr().cast::<__m256i>();
let q = _mm256_set1_epi16(Q as i16);
ihalf::<0>(p, q);
ihalf::<1>(p, q);
let (z, zq) = twi(1);
let (s, sq) = (_mm256_set1_epi16(mont(512)), _mm256_set1_epi16(mont_qinv(512)));
for k in 0..8 {
let mut x = _mm256_loadu_si256(p.add(k));
let mut y = _mm256_loadu_si256(p.add(k + 8));
gs(&mut x, &mut y, z, zq, q);
_mm256_storeu_si256(p.add(k), fq(x, s, sq, q));
_mm256_storeu_si256(p.add(k + 8), fq(y, s, sq, q));
}
}
#[target_feature(enable = "avx2")]
pub unsafe fn basemul(a: &[i16; N], b: &[i16; N], c: &mut [i16; N]) {
basemul_in::<true>(a, b, c);
}
#[target_feature(enable = "avx2")]
pub(super) unsafe fn basemul_bounded(a: &[i16; N], b: &[i16; N], c: &mut [i16; N]) {
basemul_in::<false>(a, b, c);
}
#[inline(always)]
unsafe fn basemul_in<const REDUCE: bool>(a: &[i16; N], b: &[i16; N], c: &mut [i16; N]) {
let q = _mm256_set1_epi16(Q as i16);
let swap = _mm256_setr_epi8(2, 3, 0, 1, 6, 7, 4, 5, 10, 11, 8, 9, 14, 15, 12, 13, 2, 3, 0, 1, 6, 7, 4, 5, 10, 11, 8, 9, 14, 15, 12, 13);
let (pa, pb, pc) = (a.as_ptr().cast::<__m256i>(), b.as_ptr().cast::<__m256i>(), c.as_mut_ptr().cast::<__m256i>());
for v in 0..16 {
let (x, y) = (_mm256_loadu_si256(pa.add(v)), _mm256_loadu_si256(pb.add(v)));
let (x, y) = if REDUCE { (barrett(x, q), barrett(y, q)) } else { (x, y) };
let pr = montmul(x, y, q);
let sm = montmul(x, _mm256_shuffle_epi8(y, swap), q);
let even = _mm256_add_epi16(pr, fq(_mm256_shuffle_epi8(pr, swap), ld(&GAMMA_LANES[v][0]), ld(&GAMMA_LANES[v][1]), q));
let odd = _mm256_add_epi16(sm, _mm256_shuffle_epi8(sm, swap));
_mm256_storeu_si256(pc.add(v), _mm256_blend_epi16::<0xaa>(even, odd));
}
}
#[inline(always)]
unsafe fn tail(r0: &mut __m256i, r1: &mut __m256i, t: usize, q: __m256i) {
let (mut a8, mut b8): (Span8, Span8) = perm(Row16::take(r0), Row16::take(r1));
bfv(&mut a8, &mut b8, tw_of(&TABLES, t), q);
let (mut x4, mut y4): (Span4, Span4) = unpack(a8, b8);
bfv(&mut x4, &mut y4, tw_of(&TABLES, t), q);
let (mut p2, mut q2): (Span2, Span2) = blend(x4, y4);
bfv(&mut p2, &mut q2, tw_of(&TABLES, t), q);
let (x4, y4): (Span4, Span4) = blend(p2, q2);
let (a8, b8): (Span8, Span8) = unpack(x4, y4);
let (o0, o1): (Row16, Row16) = perm(a8, b8);
o0.put(r0);
o1.put(r1);
}
#[inline(always)]
unsafe fn tail2(r0: &mut __m256i, r1: &mut __m256i, t: usize, s0: &mut __m256i, s1: &mut __m256i, u: usize, q: __m256i) {
let (mut a8, mut b8): (Span8, Span8) = perm(Row16::take(r0), Row16::take(r1));
let (mut c8, mut d8): (Span8, Span8) = perm(Row16::take(s0), Row16::take(s1));
let ta = fqv(b8, tw_of(&TABLES, t), q);
let tc = fqv(d8, tw_of(&TABLES, u), q);
b8 = subv(a8, ta);
d8 = subv(c8, tc);
a8 = addv(a8, ta);
c8 = addv(c8, tc);
let (mut x4, mut y4): (Span4, Span4) = unpack(a8, b8);
let (mut v4, mut w4): (Span4, Span4) = unpack(c8, d8);
let ta = fqv(y4, tw_of(&TABLES, t), q);
let tc = fqv(w4, tw_of(&TABLES, u), q);
y4 = subv(x4, ta);
w4 = subv(v4, tc);
x4 = addv(x4, ta);
v4 = addv(v4, tc);
let (mut p2, mut q2): (Span2, Span2) = blend(x4, y4);
let (mut m2, mut n2): (Span2, Span2) = blend(v4, w4);
let ta = fqv(q2, tw_of(&TABLES, t), q);
let tc = fqv(n2, tw_of(&TABLES, u), q);
q2 = subv(p2, ta);
n2 = subv(m2, tc);
p2 = addv(p2, ta);
m2 = addv(m2, tc);
let (x4, y4): (Span4, Span4) = blend(p2, q2);
let (v4, w4): (Span4, Span4) = blend(m2, n2);
let (a8, b8): (Span8, Span8) = unpack(x4, y4);
let (c8, d8): (Span8, Span8) = unpack(v4, w4);
let (o0, o1): (Row16, Row16) = perm(a8, b8);
let (o2, o3): (Row16, Row16) = perm(c8, d8);
o0.put(r0);
o2.put(s0);
o1.put(r1);
o3.put(s1);
}
}
#[cfg(test)]
mod tests {
use super::*;
fn rng(seed: u64) -> impl FnMut() -> u64 {
let mut s = seed | 1;
move || {
s ^= s << 13;
s ^= s >> 7;
s ^= s << 17;
s
}
}
fn canon(v: &[i16; N]) -> [i32; N] {
core::array::from_fn(|i| i32::from(v[i]).rem_euclid(Q))
}
#[test]
fn zetas_and_constants() {
assert_eq!(ZETAS[0], 1);
assert_eq!(ZETAS[1], 1729);
assert_eq!((QINV as i64 * Q as i64) % 65536, 1);
assert_eq!(pow_mod(17, 128, Q as i64), Q as i64 - 1);
}
#[test]
fn scalar_and_avx2_equal_reference_mod_q() {
let mut r = rng(42);
let mut inputs: Vec<[i16; N]> = (0..200).map(|_| core::array::from_fn(|_| (r() % Q as u64) as i16)).collect();
inputs.push([0; N]);
inputs.push([(Q - 1) as i16; N]);
inputs.push(core::array::from_fn(|i| if i % 2 == 0 { (Q - 1) as i16 } else { 0 }));
inputs.push(core::array::from_fn(|i| if i < 128 { (Q - 1) as i16 } else { 1 }));
for a in &inputs {
let want = ntt_ref(a);
let mut s = *a;
ntt_scalar(&mut s);
assert_eq!(canon(&s), want, "scalar");
#[cfg(target_arch = "x86_64")]
if slain_rt::cpu::avx2() {
let mut v = *a;
unsafe { x86::ntt(&mut v) };
assert_eq!(canon(&v), want, "avx2");
assert!(v.iter().all(|&x| i32::from(x).abs() < 8 * Q), "bound");
let mut v2 = *a;
unsafe { x86::ntt_prev(&mut v2) };
assert_eq!(v2, v, "ntt == previous ntt");
}
}
}
fn wide(a: &[i16; N]) -> [i64; N] {
core::array::from_fn(|i| i64::from(a[i]))
}
fn inputs(seed: u64, hi: u64) -> Vec<[i16; N]> {
let mut r = rng(seed);
let mut v: Vec<[i16; N]> = (0..100).map(|_| core::array::from_fn(|_| (r() % hi) as i16)).collect();
v.push([0; N]);
v.push([(hi - 1) as i16; N]);
v.push(core::array::from_fn(|i| if i % 2 == 0 { (hi - 1) as i16 } else { 0 }));
v
}
#[test]
fn field_3329_is_the_ml_kem_ntt() {
for a in inputs(5, Q as u64) {
let mut v = wide(&a);
F3329.ntt(&mut v);
assert_eq!(v.map(|x| x as i32), ntt_ref(&a));
}
}
#[test]
fn every_field_inverts_and_convolves_exactly() {
for f in [&F3329, &F7681, &F12289] {
assert_eq!(pow_mod(f.zetas[1], 2, f.p), f.p - 1, "zeta^128 = -1 for p {}", f.p);
let xs = inputs(11 + f.p as u64, 256);
for (a, b) in xs.iter().zip(xs.iter().rev()) {
let mut v = wide(a);
f.ntt(&mut v);
f.intt(&mut v);
assert_eq!(v, wide(a), "intt(ntt(a)) p {}", f.p);
let want = schoolbook(&wide(a), &wide(b)).map(|x| x.rem_euclid(f.p));
assert_eq!(f.negacyclic(&wide(a), &wide(b)), want, "negacyclic p {}", f.p);
}
}
}
#[test]
fn avx2_inverse_and_basemul_equal_the_scalar_field() {
let xs = inputs(23, Q as u64);
let negs = inputs(29, 511).into_iter().map(|a| a.map(|x| x - 255));
for (a, b) in xs.iter().copied().zip(negs) {
let (mut fa, mut fb) = (a, b);
ntt(&mut fa);
ntt(&mut fb);
let mut c = [0i16; N];
basemul(&fa, &fb, &mut c);
let mut cs = [0i16; N];
basemul_scalar(&fa, &fb, &mut cs);
assert_eq!(canon(&cs), canon(&c), "basemul scalar");
let mut is = cs;
intt_scalar(&mut is);
let want = F3329.basemul(&canon(&fa).map(i64::from), &canon(&fb).map(i64::from)).map(|x| x * MONT_INV % Q as i64);
assert_eq!(canon(&c).map(i64::from), want, "basemul");
assert!(c.iter().all(|&x| i32::from(x).abs() < 2 * Q), "basemul bound");
let mut ci = c;
intt(&mut ci);
let mut wi = want;
F3329.intt(&mut wi);
let wi = wi.map(|x| x * MONT % Q as i64);
assert_eq!(canon(&ci).map(i64::from), wi, "intt");
assert!(ci.iter().all(|&x| i32::from(x).abs() < Q), "intt bound");
assert_eq!(canon(&is), canon(&ci), "intt scalar");
let conv = schoolbook(&wide(&a), &wide(&b)).map(|x| x.rem_euclid(Q as i64));
assert_eq!(canon(&ci).map(i64::from), conv, "ntt -> basemul -> intt = negacyclic product");
}
for v in [[(Q - 1) as i16; N], [-(Q - 1) as i16; N], core::array::from_fn(|i| if i % 3 == 0 { 8 * Q as i16 - 1 } else { -8 * Q as i16 + 1 })] {
let mut ci = v;
intt(&mut ci);
let mut wi = canon(&v).map(i64::from);
F3329.intt(&mut wi);
let wi = wi.map(|x| x * MONT % Q as i64);
assert_eq!(canon(&ci).map(i64::from), wi, "intt extreme input");
}
}
#[test]
fn typed_ntt_mul_intt_is_the_negacyclic_product() {
let half = (Q / 2) as i16;
let mut xs: Vec<[i16; N]> = inputs(41, Q as u64).into_iter().map(|a| a.map(|x| x - half)).collect();
xs.push([(Q - 1) as i16; N]);
xs.push([-(Q - 1) as i16; N]);
xs.push(core::array::from_fn(|i| if i % 2 == 0 { (Q - 1) as i16 } else { -(Q - 1) as i16 }));
assert!(Poly::<Coeff>::new([Q as i16; N]).is_none());
assert_eq!(Poly::<Coeff>::reduce([Q as i16 + 5; N]).into_array(), [5; N]);
assert_eq!(Poly::<Coeff>::from_u8(&[255; N]).into_array(), [255; N]);
for (a, b) in xs.iter().zip(xs.iter().rev()) {
let (pa, pb) = (Poly::<Coeff>::new(*a).unwrap().ntt(), Poly::<Coeff>::new(*b).unwrap().ntt());
let prod = pa.mul(&pb);
assert!(prod.as_array().iter().all(|&x| i32::from(x).abs() < 8 * Q), "mul bound");
let c = prod.intt();
assert!(c.as_array().iter().all(|&x| i32::from(x).abs() < Q), "intt bound keeps the Coeff invariant");
let conv = schoolbook(&wide(a), &wide(b)).map(|x| x.rem_euclid(Q as i64));
assert_eq!(canon(c.as_array()).map(i64::from), conv);
assert_eq!(canon(c.ntt().as_array()), canon(Poly::<Coeff>::reduce(c.into_array()).ntt().as_array()));
}
}
#[cfg(target_arch = "x86_64")]
#[test]
fn bounded_basemul_is_exact_over_the_whole_ntt_output_range() {
if !slain_rt::cpu::avx2() {
return;
}
let top = (8 * Q - 1) as i16;
let mut r = rng(53);
let mut cases: Vec<[i16; N]> = (0..200).map(|_| core::array::from_fn(|_| ((r() % (16 * Q as u64 - 1)) as i32 - i32::from(top)) as i16)).collect();
cases.push([top; N]);
cases.push([-top; N]);
cases.push(core::array::from_fn(|i| if i % 2 == 0 { top } else { -top }));
cases.push(core::array::from_fn(|i| if i % 4 < 2 { top } else { -top }));
for (a, b) in cases.iter().zip(cases.iter().rev()) {
let (mut want, mut got) = ([0i16; N], [0i16; N]);
unsafe {
x86::basemul(a, b, &mut want);
x86::basemul_bounded(a, b, &mut got);
}
assert_eq!(canon(&got), canon(&want), "bounded basemul mod q");
assert!(got.iter().all(|&x| i32::from(x).abs() < 8 * Q), "bounded basemul output < 8q");
let (mut wi, mut gi) = (want, got);
intt(&mut wi);
intt(&mut gi);
assert_eq!(canon(&gi), canon(&wi), "intt of bounded product");
assert!(gi.iter().all(|&x| i32::from(x).abs() < Q), "intt output < q");
}
}
#[test]
fn crt_correlation_equals_the_direct_sum_at_every_offset() {
let mut r = rng(31);
let mut cases: Vec<([u8; 64], [u8; N])> = (0..60).map(|_| (core::array::from_fn(|_| r() as u8), core::array::from_fn(|_| r() as u8))).collect();
cases.push(([255; 64], [255; N]));
cases.push(([255; 64], [0; N]));
for (blk, tile) in cases {
let mut a = [0i16; N];
for k in 0..64 {
a[(k / 8) * 16 + k % 8] = i16::from(blk[k]);
}
let ar = reverse_negacyclic(&a);
let b: [i16; N] = tile.map(i16::from);
let (mut fa, mut fb) = (ar, b);
ntt(&mut fa);
ntt(&mut fb);
let mut c = [0i16; N];
basemul(&fa, &fb, &mut c);
intt(&mut c);
let r7681 = F7681.negacyclic(&wide(&ar), &wide(&b));
let r12289 = F12289.negacyclic(&wide(&ar), &wide(&b));
for dy in 0..9 {
for dx in 0..9 {
let d = dy * 16 + dx;
let direct: i64 = (0..64).map(|k| i64::from(blk[k]) * i64::from(tile[(k / 8 + dy) * 16 + k % 8 + dx])).sum();
let r3329 = i64::from(c[d]).rem_euclid(Q as i64);
assert_eq!(crt(&[r3329, r7681[d]], &[&F3329, &F7681]), i128::from(direct), "two primes, offset {dy},{dx}");
assert_eq!(crt(&[r3329, r7681[d], r12289[d]], &[&F3329, &F7681, &F12289]), i128::from(direct), "three primes");
}
}
}
}
#[cfg(target_arch = "x86_64")]
#[test]
#[ignore]
fn bench_correlation_per_second() {
use core::hint::black_box as slain;
if !slain_rt::cpu::avx2() {
return;
}
let mut r = rng(9);
let b: [i16; N] = core::array::from_fn(|_| (r() % 256) as i16);
let bu8: [u8; N] = b.map(|v| v as u8);
let fa_coeff: [i16; N] = core::array::from_fn(|_| (r() % 511) as i16 - 255);
let mut fa = fa_coeff;
ntt(&mut fa);
let reps = 200_000u32;
let names = ["ntt", "intt", "basemul", "piece", "typed_piece"];
let typed_fa = Poly::<Coeff>::new(core::array::from_fn(|i| fa_coeff[i])).unwrap().ntt();
let mut per = [[0f64; 15]; 5];
for round in 0..15 {
for (arm, out) in per.iter_mut().enumerate() {
let t = std::time::Instant::now();
for _ in 0..reps {
let mut x = slain(b);
let mut c = [0i16; N];
unsafe {
match arm {
0 => x86::ntt(&mut x),
1 => x86::intt(&mut x),
2 => x86::basemul(&fa, &x, &mut c),
3 => {
x86::ntt(&mut x);
x86::basemul(&fa, &x, &mut c);
x86::intt(&mut c);
}
_ => c = typed_fa.mul(&Poly::<Coeff>::from_u8(&slain(bu8)).ntt()).intt().into_array(),
}
}
slain((x, c));
}
out[round] = f64::from(reps) / t.elapsed().as_secs_f64();
}
}
for (name, v) in names.iter().zip(per.iter_mut()) {
v.sort_by(f64::total_cmp);
println!("NTT256 {name} {:.9} calls/s {:.9} ns per call (range {:.9}-{:.9})", v[7], 1e9 / v[7], v[0], v[14]);
}
}
#[cfg(target_arch = "x86_64")]
#[test]
fn intt_equals_intt_prev() {
if !slain_rt::cpu::avx2() {
return;
}
let mut r = rng(21);
for _ in 0..2000 {
let a: [i16; N] = core::array::from_fn(|_| r() as i16);
let (mut x, mut y) = (a, a);
unsafe {
x86::intt(&mut x);
x86::intt_prev(&mut y);
}
assert_eq!(canon(&x), canon(&y));
assert!(x.iter().all(|&v| i32::from(v).abs() < Q));
}
}
#[cfg(target_arch = "x86_64")]
#[test]
#[ignore]
fn bench_intt_paired() {
use core::hint::black_box as slain;
if !slain_rt::cpu::avx2() {
return;
}
let mut r = rng(5);
let a: [i16; N] = core::array::from_fn(|_| (r() % Q as u64) as i16);
let reps = 400_000u32;
let mut ratio = Vec::new();
for round in 0..15 {
let mut t = [0f64; 2];
for arm in [round % 2, 1 - round % 2] {
let s = std::time::Instant::now();
for _ in 0..reps {
let mut x = slain(a);
unsafe {
if arm == 0 {
x86::intt_prev(&mut x);
} else {
x86::intt(&mut x);
}
}
slain(x);
}
t[arm] = s.elapsed().as_secs_f64() * 1e9 / f64::from(reps);
}
ratio.push(t[0] / t[1]);
println!("round {round}: prev {:.9} ns new {:.9} ns prev/new {:.9}", t[0], t[1], t[0] / t[1]);
}
ratio.sort_by(f64::total_cmp);
println!("intt paired prev/new median {:.9} range {:.9}-{:.9}", ratio[7], ratio[0], ratio[14]);
}
#[cfg(target_arch = "x86_64")]
#[test]
#[ignore = "floor: just floor ntt_piece slain-num"]
fn floor_ntt_piece() {
use core::hint::black_box as slain;
if !slain_rt::cpu::avx2() {
return;
}
let fp0 = 3 * 56 + (3 * 48 + 3 * 56 + 3 * 16) + 17 * 16;
let mut r = rng(13);
let b: [i16; N] = core::array::from_fn(|_| (r() % 256) as i16);
let mut fa: [i16; N] = core::array::from_fn(|_| (r() % 511) as i16 - 255);
ntt(&mut fa);
let calls = 2_000_000u32;
let mut best = 0f64;
for _ in 0..3 {
let t = std::time::Instant::now();
for _ in 0..calls {
let mut x = slain(b);
let mut c = [0i16; N];
unsafe {
x86::ntt(&mut x);
x86::basemul(&fa, &x, &mut c);
x86::intt(&mut c);
}
slain(c);
}
best = best.max(f64::from(calls) / t.elapsed().as_secs_f64());
}
println!("FLOOR_FP0 {fp0}");
println!("FLOOR_PX 81");
println!("FLOOR_CALLS_PER_S {best:.9}");
println!("FLOOR_CORE_HZ {:.9}", slain_rt::cpu::core_hz());
}
#[cfg(target_arch = "x86_64")]
#[test]
#[ignore]
fn bench_ntt_per_second() {
use core::hint::black_box as slain;
if !slain_rt::cpu::avx2() {
return;
}
let mut r = rng(7);
let a: [i16; N] = core::array::from_fn(|_| (r() % Q as u64) as i16);
let reps = 200_000;
let (mut s, mut v) = (Vec::new(), Vec::new());
for round in 0..15 {
for arm in [round % 2, 1 - round % 2] {
let t = std::time::Instant::now();
for _ in 0..reps {
let mut x = slain(a);
if arm == 0 {
ntt_scalar(&mut x);
} else {
unsafe { x86::ntt(&mut x) };
}
slain(x);
}
let per_s = reps as f64 / t.elapsed().as_secs_f64();
if arm == 0 {
s.push(per_s)
} else {
v.push(per_s)
}
}
}
s.sort_by(f64::total_cmp);
v.sort_by(f64::total_cmp);
println!("NTT-256 q=3329: scalar {:.0}/s register-mapped AVX2 {:.0}/s {:.2}x (spread scalar {:.0}-{:.0}, avx2 {:.0}-{:.0})", s[7], v[7], v[7] / s[7], s[0], s[14], v[0], v[14]);
}
#[cfg(target_arch = "x86_64")]
#[test]
#[ignore]
fn bench_typed_mul_per_second() {
use core::hint::black_box as slain;
if !slain_rt::cpu::avx2() {
return;
}
let mut r = rng(9);
let a = Poly::<Coeff>::new(core::array::from_fn(|_| (r() % Q as u64) as i16)).unwrap().ntt();
let b = Poly::<Coeff>::new(core::array::from_fn(|_| (r() % Q as u64) as i16)).unwrap().ntt();
let reps = 1_000_000;
let (mut raw, mut typed) = (Vec::new(), Vec::new());
for round in 0..15 {
for arm in [round % 2, 1 - round % 2] {
let t = std::time::Instant::now();
for _ in 0..reps {
let (x, y) = (slain(a), slain(b));
if arm == 0 {
let mut c = [0i16; N];
unsafe { x86::basemul(x.as_array(), y.as_array(), &mut c) };
slain(c);
} else {
slain(x.mul(&y));
}
}
let per_s = reps as f64 / t.elapsed().as_secs_f64();
if arm == 0 {
raw.push(per_s)
} else {
typed.push(per_s)
}
}
}
raw.sort_by(f64::total_cmp);
typed.sort_by(f64::total_cmp);
println!(
"basemul-256 q=3329: raw (Barrett on both inputs) {:.3}/s typed Poly<Ntt>::mul (bound in the type, no Barrett) {:.3}/s {:.9}x (spread raw {:.3}-{:.3}, typed {:.3}-{:.3})",
raw[7],
typed[7],
typed[7] / raw[7],
raw[0],
raw[14],
typed[0],
typed[14]
);
}
}