Skip to main content

slain_num/
ntt.rs

1pub const Q: i32 = 3329;
2pub const N: usize = 256;
3
4const fn pow_mod(mut b: i64, mut e: u32, m: i64) -> i64 {
5    let mut r = 1i64;
6    b %= m;
7    while e > 0 {
8        if e & 1 == 1 {
9            r = r * b % m;
10        }
11        b = b * b % m;
12        e >>= 1;
13    }
14    r
15}
16
17const fn brv7(x: u32) -> u32 {
18    let mut r = 0;
19    let mut i = 0;
20    while i < 7 {
21        r |= ((x >> i) & 1) << (6 - i);
22        i += 1;
23    }
24    r
25}
26
27pub const ZETAS: [i32; 128] = {
28    let mut z = [0i32; 128];
29    let mut k = 0;
30    while k < 128 {
31        z[k] = pow_mod(17, brv7(k as u32), Q as i64) as i32;
32        k += 1;
33    }
34    z
35};
36
37pub const ZETAS_INV: [i32; 128] = {
38    let mut z = [0i32; 128];
39    let mut k = 0;
40    while k < 128 {
41        z[k] = pow_mod(17, 128 - brv7(k as u32), Q as i64) as i32;
42        k += 1;
43    }
44    z
45};
46
47pub const GAMMAS: [i32; 128] = {
48    let mut z = [0i32; 128];
49    let mut k = 0;
50    while k < 128 {
51        z[k] = pow_mod(17, 2 * brv7(k as u32) + 1, Q as i64) as i32;
52        k += 1;
53    }
54    z
55};
56
57pub struct Field {
58    pub p: i64,
59    pub zetas: [i64; 128],
60    pub zetas_inv: [i64; 128],
61    pub gammas: [i64; 128],
62    pub n_inv: i64,
63}
64
65impl Field {
66    pub const fn with_root(p: i64, zeta: i64) -> Self {
67        let (mut zetas, mut zetas_inv, mut gammas) = ([0i64; 128], [0i64; 128], [0i64; 128]);
68        let mut k = 0;
69        while k < 128 {
70            let b = brv7(k as u32);
71            zetas[k] = pow_mod(zeta, b, p);
72            zetas_inv[k] = pow_mod(zeta, 128 - b, p);
73            gammas[k] = pow_mod(zeta, 2 * b + 1, p);
74            k += 1;
75        }
76        Self { p, zetas, zetas_inv, gammas, n_inv: pow_mod(128, (p - 2) as u32, p) }
77    }
78
79    pub const fn new(p: i64) -> Self {
80        let mut g = 2;
81        loop {
82            let z = pow_mod(g, ((p - 1) / 256) as u32, p);
83            if pow_mod(z, 128, p) == p - 1 {
84                return Self::with_root(p, z);
85            }
86            g += 1;
87        }
88    }
89
90    pub fn ntt(&self, a: &mut [i64; N]) {
91        let p = self.p;
92        let mut k = 1;
93        let mut len = 128;
94        while len >= 2 {
95            for start in (0..N).step_by(2 * len) {
96                let z = self.zetas[k];
97                k += 1;
98                for j in start..start + len {
99                    let t = (z * a[j + len]).rem_euclid(p);
100                    a[j + len] = (a[j] - t).rem_euclid(p);
101                    a[j] = (a[j] + t).rem_euclid(p);
102                }
103            }
104            len >>= 1;
105        }
106    }
107
108    pub fn intt(&self, a: &mut [i64; N]) {
109        let p = self.p;
110        let mut len = 2;
111        while len <= 128 {
112            for (b, start) in (0..N).step_by(2 * len).enumerate() {
113                let z = self.zetas_inv[128 / len + b];
114                for j in start..start + len {
115                    let t = a[j];
116                    a[j] = (t + a[j + len]).rem_euclid(p);
117                    a[j + len] = (z * (a[j + len] - t)).rem_euclid(p);
118                }
119            }
120            len <<= 1;
121        }
122        for v in a.iter_mut() {
123            *v = (*v * self.n_inv).rem_euclid(p);
124        }
125    }
126
127    pub fn basemul(&self, a: &[i64; N], b: &[i64; N]) -> [i64; N] {
128        let p = self.p;
129        let mut c = [0i64; N];
130        for i in 0..128 {
131            let (a0, a1, b0, b1) = (a[2 * i], a[2 * i + 1], b[2 * i], b[2 * i + 1]);
132            c[2 * i] = (a0 * b0 + (a1 * b1).rem_euclid(p) * self.gammas[i]).rem_euclid(p);
133            c[2 * i + 1] = (a0 * b1 + a1 * b0).rem_euclid(p);
134        }
135        c
136    }
137
138    pub fn negacyclic(&self, a: &[i64; N], b: &[i64; N]) -> [i64; N] {
139        let (mut x, mut y) = (*a, *b);
140        for v in x.iter_mut().chain(y.iter_mut()) {
141            *v = v.rem_euclid(self.p);
142        }
143        self.ntt(&mut x);
144        self.ntt(&mut y);
145        let mut c = self.basemul(&x, &y);
146        self.intt(&mut c);
147        c
148    }
149}
150
151pub const F3329: Field = Field::with_root(3329, 17);
152pub const F7681: Field = Field::new(7681);
153pub const F12289: Field = Field::new(12289);
154
155pub fn schoolbook(a: &[i64; N], b: &[i64; N]) -> [i64; N] {
156    let mut c = [0i64; N];
157    for i in 0..N {
158        for j in 0..N {
159            let v = a[i] * b[j];
160            if i + j < N {
161                c[i + j] += v;
162            } else {
163                c[i + j - N] -= v;
164            }
165        }
166    }
167    c
168}
169
170pub fn crt(r: &[i64], fields: &[&Field]) -> i128 {
171    let mut x = 0i128;
172    let mut m = 1i128;
173    for (&ri, f) in r.iter().zip(fields) {
174        let p = i128::from(f.p);
175        let inv = i128::from(pow_mod((m % p) as i64, (f.p - 2) as u32, f.p));
176        let t = ((i128::from(ri) - x) % p + p) % p * inv % p;
177        x += t * m;
178        m *= p;
179    }
180    if x > m / 2 {
181        x - m
182    } else {
183        x
184    }
185}
186
187pub fn reverse_negacyclic(a: &[i16; N]) -> [i16; N] {
188    core::array::from_fn(|i| if i == 0 { a[0] } else { -a[N - i] })
189}
190
191const QINV: i32 = {
192    let mut x: i64 = 1;
193    let mut i = 0;
194    while i < 5 {
195        x = x * (2 - Q as i64 * x) % 65536;
196        i += 1;
197    }
198    x.rem_euclid(65536) as i32
199};
200
201const fn centered(x: i64) -> i16 {
202    let r = ((x % Q as i64) + Q as i64) % Q as i64;
203    (if r > (Q as i64) / 2 { r - Q as i64 } else { r }) as i16
204}
205
206const fn mont(z: i32) -> i16 {
207    centered(z as i64 * 65536)
208}
209
210#[cfg(target_arch = "x86_64")]
211const fn mont_qinv(z: i32) -> i16 {
212    (mont(z) as i32).wrapping_mul(QINV) as i16
213}
214
215pub const MONT: i64 = pow_mod(2, 16, Q as i64);
216pub const MONT_INV: i64 = pow_mod(MONT, (Q - 2) as u32, Q as i64);
217
218pub fn ntt_ref(a: &[i16; N]) -> [i32; N] {
219    let mut r: [i64; N] = core::array::from_fn(|i| i64::from(a[i]));
220    let mut k = 1;
221    let mut len = 128;
222    while len >= 2 {
223        let mut start = 0;
224        while start < N {
225            let z = i64::from(ZETAS[k]);
226            k += 1;
227            for j in start..start + len {
228                let t = (z * r[j + len]).rem_euclid(Q as i64);
229                r[j + len] = (r[j] - t).rem_euclid(Q as i64);
230                r[j] = (r[j] + t).rem_euclid(Q as i64);
231            }
232            start += 2 * len;
233        }
234        len >>= 1;
235    }
236    core::array::from_fn(|i| r[i] as i32)
237}
238
239pub fn ntt_scalar(a: &mut [i16; N]) {
240    let fq = |b: i16, z: i32| -> i16 {
241        let p = i32::from(b) * i32::from(mont(z));
242        let u = (p as i16 as i32).wrapping_mul(QINV) as i16;
243        ((p - i32::from(u) * Q) >> 16) as i16
244    };
245    let mut k = 1;
246    let mut len = 128;
247    while len >= 2 {
248        let mut start = 0;
249        while start < N {
250            let z = ZETAS[k];
251            k += 1;
252            for j in start..start + len {
253                let t = fq(a[j + len], z);
254                a[j + len] = a[j] - t;
255                a[j] += t;
256            }
257            start += 2 * len;
258        }
259        len >>= 1;
260    }
261}
262
263pub fn ntt(a: &mut [i16; N]) {
264    #[cfg(target_arch = "x86_64")]
265    if slain_rt::cpu::avx2() {
266        unsafe { x86::ntt(a) };
267        return;
268    }
269    ntt_scalar(a);
270}
271
272pub fn intt(a: &mut [i16; N]) {
273    #[cfg(target_arch = "x86_64")]
274    if slain_rt::cpu::avx2() {
275        unsafe { x86::intt(a) };
276        return;
277    }
278    intt_scalar(a);
279}
280
281pub fn intt_scalar(a: &mut [i16; N]) {
282    let mut v: [i64; N] = core::array::from_fn(|i| i64::from(a[i]).rem_euclid(Q as i64));
283    F3329.intt(&mut v);
284    *a = core::array::from_fn(|i| (v[i] * MONT % Q as i64) as i16);
285}
286
287pub fn basemul(a: &[i16; N], b: &[i16; N], c: &mut [i16; N]) {
288    #[cfg(target_arch = "x86_64")]
289    if slain_rt::cpu::avx2() {
290        unsafe { x86::basemul(a, b, c) };
291        return;
292    }
293    basemul_scalar(a, b, c);
294}
295
296pub fn basemul_scalar(a: &[i16; N], b: &[i16; N], c: &mut [i16; N]) {
297    let x: [i64; N] = core::array::from_fn(|i| i64::from(a[i]).rem_euclid(Q as i64));
298    let y: [i64; N] = core::array::from_fn(|i| i64::from(b[i]).rem_euclid(Q as i64));
299    let v = F3329.basemul(&x, &y);
300    *c = core::array::from_fn(|i| (v[i] * MONT_INV % Q as i64) as i16);
301}
302
303pub enum Coeff {}
304pub enum Ntt {}
305pub enum Product {}
306
307pub struct Poly<D> {
308    a: [i16; N],
309    domain: core::marker::PhantomData<fn() -> D>,
310}
311
312impl<D> Clone for Poly<D> {
313    fn clone(&self) -> Self {
314        *self
315    }
316}
317
318impl<D> Copy for Poly<D> {}
319
320impl<D> Poly<D> {
321    const fn wrap(a: [i16; N]) -> Self {
322        Self { a, domain: core::marker::PhantomData }
323    }
324
325    #[must_use]
326    pub const fn as_array(&self) -> &[i16; N] {
327        &self.a
328    }
329
330    #[must_use]
331    pub const fn into_array(self) -> [i16; N] {
332        self.a
333    }
334}
335
336impl Poly<Coeff> {
337    #[must_use]
338    pub fn new(a: [i16; N]) -> Option<Self> {
339        a.iter().all(|&x| i32::from(x).abs() < Q).then(|| Self::wrap(a))
340    }
341
342    #[must_use]
343    pub fn from_u8(a: &[u8; N]) -> Self {
344        Self::wrap(a.map(i16::from))
345    }
346
347    #[must_use]
348    pub fn reduce(a: [i16; N]) -> Self {
349        Self::wrap(a.map(|x| centered(i64::from(x))))
350    }
351
352    #[must_use]
353    pub fn ntt(self) -> Poly<Ntt> {
354        let mut a = self.a;
355        ntt(&mut a);
356        Poly::wrap(a)
357    }
358}
359
360impl Poly<Ntt> {
361    #[must_use]
362    pub fn mul(&self, b: &Self) -> Poly<Product> {
363        let mut c = [0i16; N];
364        #[cfg(target_arch = "x86_64")]
365        if slain_rt::cpu::avx2() {
366            unsafe { x86::basemul_bounded(&self.a, &b.a, &mut c) };
367            return Poly::wrap(c);
368        }
369        basemul_scalar(&self.a, &b.a, &mut c);
370        Poly::wrap(c)
371    }
372}
373
374impl Poly<Product> {
375    #[must_use]
376    pub fn intt(self) -> Poly<Coeff> {
377        let mut a = self.a;
378        intt(&mut a);
379        Poly::wrap(a)
380    }
381}
382
383#[cfg(target_arch = "x86_64")]
384pub mod x86 {
385    use super::{mont, mont_qinv, GAMMAS, N, Q, QINV, ZETAS, ZETAS_INV};
386    use core::arch::x86_64::*;
387
388    #[repr(C, align(32))]
389    struct Lanes([i16; 16]);
390
391    struct Tables {
392        l8: [[Lanes; 2]; 8],
393        l4: [[Lanes; 2]; 8],
394        l2: [[Lanes; 2]; 8],
395    }
396
397    const fn pair(zs: &[i32; 128], f: [usize; 16]) -> [Lanes; 2] {
398        let mut z = [0i16; 16];
399        let mut zq = [0i16; 16];
400        let mut i = 0;
401        while i < 16 {
402            z[i] = mont(zs[f[i]]);
403            zq[i] = mont_qinv(zs[f[i]]);
404            i += 1;
405        }
406        [Lanes(z), Lanes(zq)]
407    }
408
409    const fn build(zs: &[i32; 128], layer: u32) -> [[Lanes; 2]; 8] {
410        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])];
411        let mut p = 0;
412        while p < 8 {
413            let g = 2 * p;
414            let mut f = [0usize; 16];
415            let mut i = 0;
416            while i < 16 {
417                f[i] = match layer {
418                    8 => 16 + g + i / 8,
419                    4 => 32 + 2 * g + i / 4,
420                    _ => 64 + 4 * g + i / 2,
421                };
422                i += 1;
423            }
424            out[p] = pair(zs, f);
425            p += 1;
426        }
427        out
428    }
429
430    const TABLES: Tables = Tables { l8: build(&ZETAS, 8), l4: build(&ZETAS, 4), l2: build(&ZETAS, 2) };
431    const TABLES_INV: Tables = Tables { l8: build(&ZETAS_INV, 8), l4: build(&ZETAS_INV, 4), l2: build(&ZETAS_INV, 2) };
432
433    const GAMMA_LANES: [[Lanes; 2]; 16] = {
434        let mut out = [const { [Lanes([0; 16]), Lanes([0; 16])] }; 16];
435        let mut v = 0;
436        while v < 16 {
437            let mut m = 0;
438            while m < 8 {
439                out[v][0].0[2 * m] = mont(GAMMAS[8 * v + m]);
440                out[v][1].0[2 * m] = mont_qinv(GAMMAS[8 * v + m]);
441                m += 1;
442            }
443            v += 1;
444        }
445        out
446    };
447
448    #[derive(Clone, Copy)]
449    enum Row {}
450    #[derive(Clone, Copy)]
451    enum Dist8 {}
452    #[derive(Clone, Copy)]
453    enum Dist4 {}
454    #[derive(Clone, Copy)]
455    enum Dist2 {}
456
457    #[derive(Clone, Copy)]
458    #[repr(transparent)]
459    struct V<L>(__m256i, core::marker::PhantomData<L>);
460
461    type Row16 = V<Row>;
462    type Span8 = V<Dist8>;
463    type Span4 = V<Dist4>;
464    type Span2 = V<Dist2>;
465
466    #[derive(Clone, Copy)]
467    struct Tw<L>(__m256i, __m256i, core::marker::PhantomData<L>);
468
469    trait Perm128 {
470        type To;
471    }
472    impl Perm128 for Row {
473        type To = Dist8;
474    }
475    impl Perm128 for Dist8 {
476        type To = Row;
477    }
478
479    trait Unpack64 {
480        type To;
481    }
482    impl Unpack64 for Dist8 {
483        type To = Dist4;
484    }
485    impl Unpack64 for Dist4 {
486        type To = Dist8;
487    }
488
489    trait Blend32 {
490        type To;
491    }
492    impl Blend32 for Dist4 {
493        type To = Dist2;
494    }
495    impl Blend32 for Dist2 {
496        type To = Dist4;
497    }
498
499    trait Layer {
500        fn of(t: &Tables) -> &[[Lanes; 2]; 8];
501    }
502    impl Layer for Dist8 {
503        fn of(t: &Tables) -> &[[Lanes; 2]; 8] {
504            &t.l8
505        }
506    }
507    impl Layer for Dist4 {
508        fn of(t: &Tables) -> &[[Lanes; 2]; 8] {
509            &t.l4
510        }
511    }
512    impl Layer for Dist2 {
513        fn of(t: &Tables) -> &[[Lanes; 2]; 8] {
514            &t.l2
515        }
516    }
517
518    impl<L> V<L> {
519        #[inline(always)]
520        const fn of(v: __m256i) -> Self {
521            Self(v, core::marker::PhantomData)
522        }
523    }
524
525    impl Row16 {
526        #[inline(always)]
527        unsafe fn take(r: &__m256i) -> Self {
528            Self::of(*r)
529        }
530
531        #[inline(always)]
532        unsafe fn put(self, r: &mut __m256i) {
533            *r = self.0;
534        }
535    }
536
537    #[inline(always)]
538    unsafe fn tw_of<L: Layer>(t: &Tables, i: usize) -> Tw<L> {
539        let [z, zq] = &L::of(t)[i];
540        Tw(ld(z), ld(zq), core::marker::PhantomData)
541    }
542
543    #[inline(always)]
544    unsafe fn perm<L: Perm128>(a: V<L>, b: V<L>) -> (V<L::To>, V<L::To>) {
545        let (x, y) = perm_raw(a.0, b.0);
546        (V::of(x), V::of(y))
547    }
548
549    #[inline(always)]
550    unsafe fn unpack<L: Unpack64>(a: V<L>, b: V<L>) -> (V<L::To>, V<L::To>) {
551        let (x, y) = unp(a.0, b.0);
552        (V::of(x), V::of(y))
553    }
554
555    #[inline(always)]
556    unsafe fn blend<L: Blend32>(a: V<L>, b: V<L>) -> (V<L::To>, V<L::To>) {
557        let (x, y) = bl(a.0, b.0);
558        (V::of(x), V::of(y))
559    }
560
561    #[inline(always)]
562    unsafe fn fqv<L>(b: V<L>, w: Tw<L>, q: __m256i) -> V<L> {
563        V::of(fq(b.0, w.0, w.1, q))
564    }
565
566    #[inline(always)]
567    unsafe fn addv<L>(a: V<L>, b: V<L>) -> V<L> {
568        V::of(_mm256_add_epi16(a.0, b.0))
569    }
570
571    #[inline(always)]
572    unsafe fn subv<L>(a: V<L>, b: V<L>) -> V<L> {
573        V::of(_mm256_sub_epi16(a.0, b.0))
574    }
575
576    #[inline(always)]
577    unsafe fn bfv<L>(a: &mut V<L>, b: &mut V<L>, w: Tw<L>, q: __m256i) {
578        bf(&mut a.0, &mut b.0, w.0, w.1, q);
579    }
580
581    #[inline(always)]
582    unsafe fn gsv<L>(a: &mut V<L>, b: &mut V<L>, w: Tw<L>, q: __m256i) {
583        gs(&mut a.0, &mut b.0, w.0, w.1, q);
584    }
585
586    #[inline(always)]
587    unsafe fn fq(b: __m256i, z: __m256i, zq: __m256i, q: __m256i) -> __m256i {
588        _mm256_sub_epi16(_mm256_mulhi_epi16(b, z), _mm256_mulhi_epi16(_mm256_mullo_epi16(b, zq), q))
589    }
590
591    #[inline(always)]
592    unsafe fn bf(a: &mut __m256i, b: &mut __m256i, z: __m256i, zq: __m256i, q: __m256i) {
593        let t = fq(*b, z, zq, q);
594        *b = _mm256_sub_epi16(*a, t);
595        *a = _mm256_add_epi16(*a, t);
596    }
597
598    #[inline(always)]
599    unsafe fn ld(t: &Lanes) -> __m256i {
600        _mm256_load_si256(t.0.as_ptr().cast())
601    }
602
603    #[inline(always)]
604    unsafe fn tw(k: usize) -> (__m256i, __m256i) {
605        (_mm256_set1_epi16(mont(ZETAS[k])), _mm256_set1_epi16(mont_qinv(ZETAS[k])))
606    }
607
608    #[target_feature(enable = "avx2")]
609    pub unsafe fn ntt(a: &mut [i16; N]) {
610        let p = a.as_mut_ptr().cast::<__m256i>();
611        let q = _mm256_set1_epi16(Q as i16);
612        let (z, zq) = tw(1);
613        for k in 0..8 {
614            let mut x = _mm256_loadu_si256(p.add(k));
615            let mut y = _mm256_loadu_si256(p.add(k + 8));
616            bf(&mut x, &mut y, z, zq, q);
617            _mm256_storeu_si256(p.add(k), x);
618            _mm256_storeu_si256(p.add(k + 8), y);
619        }
620        half::<0, true>(p, q);
621        half::<1, true>(p, q);
622    }
623
624    #[cfg(test)]
625    #[target_feature(enable = "avx2")]
626    pub unsafe fn ntt_prev(a: &mut [i16; N]) {
627        let p = a.as_mut_ptr().cast::<__m256i>();
628        let q = _mm256_set1_epi16(Q as i16);
629        let (z, zq) = tw(1);
630        for k in 0..8 {
631            let mut x = _mm256_loadu_si256(p.add(k));
632            let mut y = _mm256_loadu_si256(p.add(k + 8));
633            bf(&mut x, &mut y, z, zq, q);
634            _mm256_storeu_si256(p.add(k), x);
635            _mm256_storeu_si256(p.add(k + 8), y);
636        }
637        half::<0, false>(p, q);
638        half::<1, false>(p, q);
639    }
640
641    #[inline(always)]
642    unsafe fn half<const H: usize, const PAIRED: bool>(p: *mut __m256i, q: __m256i) {
643        let b = p.add(8 * H);
644        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)));
645        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)));
646        let (z, zq) = tw(2 + H);
647        bf(&mut r0, &mut r4, z, zq, q);
648        bf(&mut r1, &mut r5, z, zq, q);
649        bf(&mut r2, &mut r6, z, zq, q);
650        bf(&mut r3, &mut r7, z, zq, q);
651        let (z, zq) = tw(4 + 2 * H);
652        bf(&mut r0, &mut r2, z, zq, q);
653        bf(&mut r1, &mut r3, z, zq, q);
654        let (z, zq) = tw(5 + 2 * H);
655        bf(&mut r4, &mut r6, z, zq, q);
656        bf(&mut r5, &mut r7, z, zq, q);
657        let (z, zq) = tw(8 + 4 * H);
658        bf(&mut r0, &mut r1, z, zq, q);
659        let (z, zq) = tw(9 + 4 * H);
660        bf(&mut r2, &mut r3, z, zq, q);
661        let (z, zq) = tw(10 + 4 * H);
662        bf(&mut r4, &mut r5, z, zq, q);
663        let (z, zq) = tw(11 + 4 * H);
664        bf(&mut r6, &mut r7, z, zq, q);
665        if PAIRED {
666            tail2(&mut r0, &mut r1, 4 * H, &mut r2, &mut r3, 4 * H + 1, q);
667            tail2(&mut r4, &mut r5, 4 * H + 2, &mut r6, &mut r7, 4 * H + 3, q);
668        } else {
669            tail(&mut r0, &mut r1, 4 * H, q);
670            tail(&mut r2, &mut r3, 4 * H + 1, q);
671            tail(&mut r4, &mut r5, 4 * H + 2, q);
672            tail(&mut r6, &mut r7, 4 * H + 3, q);
673        }
674        _mm256_storeu_si256(b, r0);
675        _mm256_storeu_si256(b.add(1), r1);
676        _mm256_storeu_si256(b.add(2), r2);
677        _mm256_storeu_si256(b.add(3), r3);
678        _mm256_storeu_si256(b.add(4), r4);
679        _mm256_storeu_si256(b.add(5), r5);
680        _mm256_storeu_si256(b.add(6), r6);
681        _mm256_storeu_si256(b.add(7), r7);
682    }
683
684    #[inline(always)]
685    unsafe fn twi(k: usize) -> (__m256i, __m256i) {
686        (_mm256_set1_epi16(mont(ZETAS_INV[k])), _mm256_set1_epi16(mont_qinv(ZETAS_INV[k])))
687    }
688
689    #[inline(always)]
690    unsafe fn gs(a: &mut __m256i, b: &mut __m256i, z: __m256i, zq: __m256i, q: __m256i) {
691        let t = *a;
692        *a = _mm256_add_epi16(t, *b);
693        *b = fq(_mm256_sub_epi16(*b, t), z, zq, q);
694    }
695
696    #[inline(always)]
697    unsafe fn barrett(a: __m256i, q: __m256i) -> __m256i {
698        let t = _mm256_mulhrs_epi16(_mm256_mulhi_epi16(a, _mm256_set1_epi16(20159)), _mm256_set1_epi16(32));
699        _mm256_sub_epi16(a, _mm256_mullo_epi16(t, q))
700    }
701
702    #[inline(always)]
703    unsafe fn montmul(a: __m256i, b: __m256i, q: __m256i) -> __m256i {
704        let u = _mm256_mullo_epi16(_mm256_mullo_epi16(a, b), _mm256_set1_epi16(QINV as i16));
705        _mm256_sub_epi16(_mm256_mulhi_epi16(a, b), _mm256_mulhi_epi16(u, q))
706    }
707
708    #[inline(always)]
709    unsafe fn perm_raw(a: __m256i, b: __m256i) -> (__m256i, __m256i) {
710        (_mm256_permute2x128_si256::<0x20>(a, b), _mm256_permute2x128_si256::<0x31>(a, b))
711    }
712
713    #[inline(always)]
714    unsafe fn unp(a: __m256i, b: __m256i) -> (__m256i, __m256i) {
715        (_mm256_unpacklo_epi64(a, b), _mm256_unpackhi_epi64(a, b))
716    }
717
718    #[inline(always)]
719    unsafe fn bl(a: __m256i, b: __m256i) -> (__m256i, __m256i) {
720        (_mm256_blend_epi32::<0xaa>(a, _mm256_slli_epi64::<32>(b)), _mm256_blend_epi32::<0xaa>(_mm256_srli_epi64::<32>(a), b))
721    }
722
723    #[inline(always)]
724    unsafe fn itail(r0: &mut __m256i, r1: &mut __m256i, t: usize, q: __m256i) {
725        let (a8, b8): (Span8, Span8) = perm(Row16::take(r0), Row16::take(r1));
726        let (x4, y4): (Span4, Span4) = unpack(a8, b8);
727        let (mut p2, mut q2): (Span2, Span2) = blend(x4, y4);
728        gsv(&mut p2, &mut q2, tw_of(&TABLES_INV, t), q);
729        let (mut x4, mut y4): (Span4, Span4) = blend(p2, q2);
730        gsv(&mut x4, &mut y4, tw_of(&TABLES_INV, t), q);
731        let (mut a8, mut b8): (Span8, Span8) = unpack(x4, y4);
732        gsv(&mut a8, &mut b8, tw_of(&TABLES_INV, t), q);
733        let (o0, o1): (Row16, Row16) = perm(a8, b8);
734        o0.put(r0);
735        o1.put(r1);
736    }
737
738    #[inline(always)]
739    unsafe fn ihalf<const H: usize>(p: *mut __m256i, q: __m256i) {
740        let b = p.add(8 * H);
741        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));
742        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));
743        itail(&mut r0, &mut r1, 4 * H, q);
744        itail(&mut r2, &mut r3, 4 * H + 1, q);
745        itail(&mut r4, &mut r5, 4 * H + 2, q);
746        itail(&mut r6, &mut r7, 4 * H + 3, q);
747        (r0, r1, r2, r3) = (barrett(r0, q), barrett(r1, q), barrett(r2, q), barrett(r3, q));
748        (r4, r5, r6, r7) = (barrett(r4, q), barrett(r5, q), barrett(r6, q), barrett(r7, q));
749        let (z, zq) = twi(8 + 4 * H);
750        gs(&mut r0, &mut r1, z, zq, q);
751        let (z, zq) = twi(9 + 4 * H);
752        gs(&mut r2, &mut r3, z, zq, q);
753        let (z, zq) = twi(10 + 4 * H);
754        gs(&mut r4, &mut r5, z, zq, q);
755        let (z, zq) = twi(11 + 4 * H);
756        gs(&mut r6, &mut r7, z, zq, q);
757        let (z, zq) = twi(4 + 2 * H);
758        gs(&mut r0, &mut r2, z, zq, q);
759        gs(&mut r1, &mut r3, z, zq, q);
760        let (z, zq) = twi(5 + 2 * H);
761        gs(&mut r4, &mut r6, z, zq, q);
762        gs(&mut r5, &mut r7, z, zq, q);
763        let (z, zq) = twi(2 + H);
764        gs(&mut r0, &mut r4, z, zq, q);
765        gs(&mut r1, &mut r5, z, zq, q);
766        gs(&mut r2, &mut r6, z, zq, q);
767        gs(&mut r3, &mut r7, z, zq, q);
768        for (i, v) in [r0, r1, r2, r3, r4, r5, r6, r7].into_iter().enumerate() {
769            _mm256_storeu_si256(b.add(i), barrett(v, q));
770        }
771    }
772
773    #[target_feature(enable = "avx2")]
774    pub unsafe fn intt(a: &mut [i16; N]) {
775        let p = a.as_mut_ptr().cast::<__m256i>();
776        let q = _mm256_set1_epi16(Q as i16);
777        ihalf::<0>(p, q);
778        ihalf::<1>(p, q);
779        const ZS: i32 = ((ZETAS_INV[1] as i64 * 512).rem_euclid(Q as i64)) as i32;
780        let (zs, zsq) = (_mm256_set1_epi16(mont(ZS)), _mm256_set1_epi16(mont_qinv(ZS)));
781        let (s, sq) = (_mm256_set1_epi16(mont(512)), _mm256_set1_epi16(mont_qinv(512)));
782        for k in 0..8 {
783            let x = _mm256_loadu_si256(p.add(k));
784            let y = _mm256_loadu_si256(p.add(k + 8));
785            _mm256_storeu_si256(p.add(k), fq(_mm256_add_epi16(x, y), s, sq, q));
786            _mm256_storeu_si256(p.add(k + 8), fq(_mm256_sub_epi16(y, x), zs, zsq, q));
787        }
788    }
789
790    #[cfg(test)]
791    #[target_feature(enable = "avx2")]
792    pub unsafe fn intt_prev(a: &mut [i16; N]) {
793        let p = a.as_mut_ptr().cast::<__m256i>();
794        let q = _mm256_set1_epi16(Q as i16);
795        ihalf::<0>(p, q);
796        ihalf::<1>(p, q);
797        let (z, zq) = twi(1);
798        let (s, sq) = (_mm256_set1_epi16(mont(512)), _mm256_set1_epi16(mont_qinv(512)));
799        for k in 0..8 {
800            let mut x = _mm256_loadu_si256(p.add(k));
801            let mut y = _mm256_loadu_si256(p.add(k + 8));
802            gs(&mut x, &mut y, z, zq, q);
803            _mm256_storeu_si256(p.add(k), fq(x, s, sq, q));
804            _mm256_storeu_si256(p.add(k + 8), fq(y, s, sq, q));
805        }
806    }
807
808    #[target_feature(enable = "avx2")]
809    pub unsafe fn basemul(a: &[i16; N], b: &[i16; N], c: &mut [i16; N]) {
810        basemul_in::<true>(a, b, c);
811    }
812
813    #[target_feature(enable = "avx2")]
814    pub(super) unsafe fn basemul_bounded(a: &[i16; N], b: &[i16; N], c: &mut [i16; N]) {
815        basemul_in::<false>(a, b, c);
816    }
817
818    #[inline(always)]
819    unsafe fn basemul_in<const REDUCE: bool>(a: &[i16; N], b: &[i16; N], c: &mut [i16; N]) {
820        let q = _mm256_set1_epi16(Q as i16);
821        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);
822        let (pa, pb, pc) = (a.as_ptr().cast::<__m256i>(), b.as_ptr().cast::<__m256i>(), c.as_mut_ptr().cast::<__m256i>());
823        for v in 0..16 {
824            let (x, y) = (_mm256_loadu_si256(pa.add(v)), _mm256_loadu_si256(pb.add(v)));
825            let (x, y) = if REDUCE { (barrett(x, q), barrett(y, q)) } else { (x, y) };
826            let pr = montmul(x, y, q);
827            let sm = montmul(x, _mm256_shuffle_epi8(y, swap), q);
828            let even = _mm256_add_epi16(pr, fq(_mm256_shuffle_epi8(pr, swap), ld(&GAMMA_LANES[v][0]), ld(&GAMMA_LANES[v][1]), q));
829            let odd = _mm256_add_epi16(sm, _mm256_shuffle_epi8(sm, swap));
830            _mm256_storeu_si256(pc.add(v), _mm256_blend_epi16::<0xaa>(even, odd));
831        }
832    }
833
834    #[inline(always)]
835    unsafe fn tail(r0: &mut __m256i, r1: &mut __m256i, t: usize, q: __m256i) {
836        let (mut a8, mut b8): (Span8, Span8) = perm(Row16::take(r0), Row16::take(r1));
837        bfv(&mut a8, &mut b8, tw_of(&TABLES, t), q);
838        let (mut x4, mut y4): (Span4, Span4) = unpack(a8, b8);
839        bfv(&mut x4, &mut y4, tw_of(&TABLES, t), q);
840        let (mut p2, mut q2): (Span2, Span2) = blend(x4, y4);
841        bfv(&mut p2, &mut q2, tw_of(&TABLES, t), q);
842        let (x4, y4): (Span4, Span4) = blend(p2, q2);
843        let (a8, b8): (Span8, Span8) = unpack(x4, y4);
844        let (o0, o1): (Row16, Row16) = perm(a8, b8);
845        o0.put(r0);
846        o1.put(r1);
847    }
848
849    #[inline(always)]
850    unsafe fn tail2(r0: &mut __m256i, r1: &mut __m256i, t: usize, s0: &mut __m256i, s1: &mut __m256i, u: usize, q: __m256i) {
851        let (mut a8, mut b8): (Span8, Span8) = perm(Row16::take(r0), Row16::take(r1));
852        let (mut c8, mut d8): (Span8, Span8) = perm(Row16::take(s0), Row16::take(s1));
853        let ta = fqv(b8, tw_of(&TABLES, t), q);
854        let tc = fqv(d8, tw_of(&TABLES, u), q);
855        b8 = subv(a8, ta);
856        d8 = subv(c8, tc);
857        a8 = addv(a8, ta);
858        c8 = addv(c8, tc);
859        let (mut x4, mut y4): (Span4, Span4) = unpack(a8, b8);
860        let (mut v4, mut w4): (Span4, Span4) = unpack(c8, d8);
861        let ta = fqv(y4, tw_of(&TABLES, t), q);
862        let tc = fqv(w4, tw_of(&TABLES, u), q);
863        y4 = subv(x4, ta);
864        w4 = subv(v4, tc);
865        x4 = addv(x4, ta);
866        v4 = addv(v4, tc);
867        let (mut p2, mut q2): (Span2, Span2) = blend(x4, y4);
868        let (mut m2, mut n2): (Span2, Span2) = blend(v4, w4);
869        let ta = fqv(q2, tw_of(&TABLES, t), q);
870        let tc = fqv(n2, tw_of(&TABLES, u), q);
871        q2 = subv(p2, ta);
872        n2 = subv(m2, tc);
873        p2 = addv(p2, ta);
874        m2 = addv(m2, tc);
875        let (x4, y4): (Span4, Span4) = blend(p2, q2);
876        let (v4, w4): (Span4, Span4) = blend(m2, n2);
877        let (a8, b8): (Span8, Span8) = unpack(x4, y4);
878        let (c8, d8): (Span8, Span8) = unpack(v4, w4);
879        let (o0, o1): (Row16, Row16) = perm(a8, b8);
880        let (o2, o3): (Row16, Row16) = perm(c8, d8);
881        o0.put(r0);
882        o2.put(s0);
883        o1.put(r1);
884        o3.put(s1);
885    }
886}
887
888#[cfg(test)]
889mod tests {
890    use super::*;
891
892    fn rng(seed: u64) -> impl FnMut() -> u64 {
893        let mut s = seed | 1;
894        move || {
895            s ^= s << 13;
896            s ^= s >> 7;
897            s ^= s << 17;
898            s
899        }
900    }
901
902    fn canon(v: &[i16; N]) -> [i32; N] {
903        core::array::from_fn(|i| i32::from(v[i]).rem_euclid(Q))
904    }
905
906    #[test]
907    fn zetas_and_constants() {
908        assert_eq!(ZETAS[0], 1);
909        assert_eq!(ZETAS[1], 1729);
910        assert_eq!((QINV as i64 * Q as i64) % 65536, 1);
911        assert_eq!(pow_mod(17, 128, Q as i64), Q as i64 - 1);
912    }
913
914    #[test]
915    fn scalar_and_avx2_equal_reference_mod_q() {
916        let mut r = rng(42);
917        let mut inputs: Vec<[i16; N]> = (0..200).map(|_| core::array::from_fn(|_| (r() % Q as u64) as i16)).collect();
918        inputs.push([0; N]);
919        inputs.push([(Q - 1) as i16; N]);
920        inputs.push(core::array::from_fn(|i| if i % 2 == 0 { (Q - 1) as i16 } else { 0 }));
921        inputs.push(core::array::from_fn(|i| if i < 128 { (Q - 1) as i16 } else { 1 }));
922        for a in &inputs {
923            let want = ntt_ref(a);
924            let mut s = *a;
925            ntt_scalar(&mut s);
926            assert_eq!(canon(&s), want, "scalar");
927            #[cfg(target_arch = "x86_64")]
928            if slain_rt::cpu::avx2() {
929                let mut v = *a;
930                unsafe { x86::ntt(&mut v) };
931                assert_eq!(canon(&v), want, "avx2");
932                assert!(v.iter().all(|&x| i32::from(x).abs() < 8 * Q), "bound");
933                let mut v2 = *a;
934                unsafe { x86::ntt_prev(&mut v2) };
935                assert_eq!(v2, v, "ntt == previous ntt");
936            }
937        }
938    }
939
940    fn wide(a: &[i16; N]) -> [i64; N] {
941        core::array::from_fn(|i| i64::from(a[i]))
942    }
943
944    fn inputs(seed: u64, hi: u64) -> Vec<[i16; N]> {
945        let mut r = rng(seed);
946        let mut v: Vec<[i16; N]> = (0..100).map(|_| core::array::from_fn(|_| (r() % hi) as i16)).collect();
947        v.push([0; N]);
948        v.push([(hi - 1) as i16; N]);
949        v.push(core::array::from_fn(|i| if i % 2 == 0 { (hi - 1) as i16 } else { 0 }));
950        v
951    }
952
953    #[test]
954    fn field_3329_is_the_ml_kem_ntt() {
955        for a in inputs(5, Q as u64) {
956            let mut v = wide(&a);
957            F3329.ntt(&mut v);
958            assert_eq!(v.map(|x| x as i32), ntt_ref(&a));
959        }
960    }
961
962    #[test]
963    fn every_field_inverts_and_convolves_exactly() {
964        for f in [&F3329, &F7681, &F12289] {
965            assert_eq!(pow_mod(f.zetas[1], 2, f.p), f.p - 1, "zeta^128 = -1 for p {}", f.p);
966            let xs = inputs(11 + f.p as u64, 256);
967            for (a, b) in xs.iter().zip(xs.iter().rev()) {
968                let mut v = wide(a);
969                f.ntt(&mut v);
970                f.intt(&mut v);
971                assert_eq!(v, wide(a), "intt(ntt(a)) p {}", f.p);
972                let want = schoolbook(&wide(a), &wide(b)).map(|x| x.rem_euclid(f.p));
973                assert_eq!(f.negacyclic(&wide(a), &wide(b)), want, "negacyclic p {}", f.p);
974            }
975        }
976    }
977
978    #[test]
979    fn avx2_inverse_and_basemul_equal_the_scalar_field() {
980        let xs = inputs(23, Q as u64);
981        let negs = inputs(29, 511).into_iter().map(|a| a.map(|x| x - 255));
982        for (a, b) in xs.iter().copied().zip(negs) {
983            let (mut fa, mut fb) = (a, b);
984            ntt(&mut fa);
985            ntt(&mut fb);
986            let mut c = [0i16; N];
987            basemul(&fa, &fb, &mut c);
988            let mut cs = [0i16; N];
989            basemul_scalar(&fa, &fb, &mut cs);
990            assert_eq!(canon(&cs), canon(&c), "basemul scalar");
991            let mut is = cs;
992            intt_scalar(&mut is);
993            let want = F3329.basemul(&canon(&fa).map(i64::from), &canon(&fb).map(i64::from)).map(|x| x * MONT_INV % Q as i64);
994            assert_eq!(canon(&c).map(i64::from), want, "basemul");
995            assert!(c.iter().all(|&x| i32::from(x).abs() < 2 * Q), "basemul bound");
996            let mut ci = c;
997            intt(&mut ci);
998            let mut wi = want;
999            F3329.intt(&mut wi);
1000            let wi = wi.map(|x| x * MONT % Q as i64);
1001            assert_eq!(canon(&ci).map(i64::from), wi, "intt");
1002            assert!(ci.iter().all(|&x| i32::from(x).abs() < Q), "intt bound");
1003            assert_eq!(canon(&is), canon(&ci), "intt scalar");
1004            let conv = schoolbook(&wide(&a), &wide(&b)).map(|x| x.rem_euclid(Q as i64));
1005            assert_eq!(canon(&ci).map(i64::from), conv, "ntt -> basemul -> intt = negacyclic product");
1006        }
1007        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 })] {
1008            let mut ci = v;
1009            intt(&mut ci);
1010            let mut wi = canon(&v).map(i64::from);
1011            F3329.intt(&mut wi);
1012            let wi = wi.map(|x| x * MONT % Q as i64);
1013            assert_eq!(canon(&ci).map(i64::from), wi, "intt extreme input");
1014        }
1015    }
1016
1017    #[test]
1018    fn typed_ntt_mul_intt_is_the_negacyclic_product() {
1019        let half = (Q / 2) as i16;
1020        let mut xs: Vec<[i16; N]> = inputs(41, Q as u64).into_iter().map(|a| a.map(|x| x - half)).collect();
1021        xs.push([(Q - 1) as i16; N]);
1022        xs.push([-(Q - 1) as i16; N]);
1023        xs.push(core::array::from_fn(|i| if i % 2 == 0 { (Q - 1) as i16 } else { -(Q - 1) as i16 }));
1024        assert!(Poly::<Coeff>::new([Q as i16; N]).is_none());
1025        assert_eq!(Poly::<Coeff>::reduce([Q as i16 + 5; N]).into_array(), [5; N]);
1026        assert_eq!(Poly::<Coeff>::from_u8(&[255; N]).into_array(), [255; N]);
1027        for (a, b) in xs.iter().zip(xs.iter().rev()) {
1028            let (pa, pb) = (Poly::<Coeff>::new(*a).unwrap().ntt(), Poly::<Coeff>::new(*b).unwrap().ntt());
1029            let prod = pa.mul(&pb);
1030            assert!(prod.as_array().iter().all(|&x| i32::from(x).abs() < 8 * Q), "mul bound");
1031            let c = prod.intt();
1032            assert!(c.as_array().iter().all(|&x| i32::from(x).abs() < Q), "intt bound keeps the Coeff invariant");
1033            let conv = schoolbook(&wide(a), &wide(b)).map(|x| x.rem_euclid(Q as i64));
1034            assert_eq!(canon(c.as_array()).map(i64::from), conv);
1035            assert_eq!(canon(c.ntt().as_array()), canon(Poly::<Coeff>::reduce(c.into_array()).ntt().as_array()));
1036        }
1037    }
1038
1039    #[cfg(target_arch = "x86_64")]
1040    #[test]
1041    fn bounded_basemul_is_exact_over_the_whole_ntt_output_range() {
1042        if !slain_rt::cpu::avx2() {
1043            return;
1044        }
1045        let top = (8 * Q - 1) as i16;
1046        let mut r = rng(53);
1047        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();
1048        cases.push([top; N]);
1049        cases.push([-top; N]);
1050        cases.push(core::array::from_fn(|i| if i % 2 == 0 { top } else { -top }));
1051        cases.push(core::array::from_fn(|i| if i % 4 < 2 { top } else { -top }));
1052        for (a, b) in cases.iter().zip(cases.iter().rev()) {
1053            let (mut want, mut got) = ([0i16; N], [0i16; N]);
1054            unsafe {
1055                x86::basemul(a, b, &mut want);
1056                x86::basemul_bounded(a, b, &mut got);
1057            }
1058            assert_eq!(canon(&got), canon(&want), "bounded basemul mod q");
1059            assert!(got.iter().all(|&x| i32::from(x).abs() < 8 * Q), "bounded basemul output < 8q");
1060            let (mut wi, mut gi) = (want, got);
1061            intt(&mut wi);
1062            intt(&mut gi);
1063            assert_eq!(canon(&gi), canon(&wi), "intt of bounded product");
1064            assert!(gi.iter().all(|&x| i32::from(x).abs() < Q), "intt output < q");
1065        }
1066    }
1067
1068    #[test]
1069    fn crt_correlation_equals_the_direct_sum_at_every_offset() {
1070        let mut r = rng(31);
1071        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();
1072        cases.push(([255; 64], [255; N]));
1073        cases.push(([255; 64], [0; N]));
1074        for (blk, tile) in cases {
1075            let mut a = [0i16; N];
1076            for k in 0..64 {
1077                a[(k / 8) * 16 + k % 8] = i16::from(blk[k]);
1078            }
1079            let ar = reverse_negacyclic(&a);
1080            let b: [i16; N] = tile.map(i16::from);
1081            let (mut fa, mut fb) = (ar, b);
1082            ntt(&mut fa);
1083            ntt(&mut fb);
1084            let mut c = [0i16; N];
1085            basemul(&fa, &fb, &mut c);
1086            intt(&mut c);
1087            let r7681 = F7681.negacyclic(&wide(&ar), &wide(&b));
1088            let r12289 = F12289.negacyclic(&wide(&ar), &wide(&b));
1089            for dy in 0..9 {
1090                for dx in 0..9 {
1091                    let d = dy * 16 + dx;
1092                    let direct: i64 = (0..64).map(|k| i64::from(blk[k]) * i64::from(tile[(k / 8 + dy) * 16 + k % 8 + dx])).sum();
1093                    let r3329 = i64::from(c[d]).rem_euclid(Q as i64);
1094                    assert_eq!(crt(&[r3329, r7681[d]], &[&F3329, &F7681]), i128::from(direct), "two primes, offset {dy},{dx}");
1095                    assert_eq!(crt(&[r3329, r7681[d], r12289[d]], &[&F3329, &F7681, &F12289]), i128::from(direct), "three primes");
1096                }
1097            }
1098        }
1099    }
1100
1101    #[cfg(target_arch = "x86_64")]
1102    #[test]
1103    #[ignore]
1104    fn bench_correlation_per_second() {
1105        use core::hint::black_box as slain;
1106        if !slain_rt::cpu::avx2() {
1107            return;
1108        }
1109        let mut r = rng(9);
1110        let b: [i16; N] = core::array::from_fn(|_| (r() % 256) as i16);
1111        let bu8: [u8; N] = b.map(|v| v as u8);
1112        let fa_coeff: [i16; N] = core::array::from_fn(|_| (r() % 511) as i16 - 255);
1113        let mut fa = fa_coeff;
1114        ntt(&mut fa);
1115        let reps = 200_000u32;
1116        let names = ["ntt", "intt", "basemul", "piece", "typed_piece"];
1117        let typed_fa = Poly::<Coeff>::new(core::array::from_fn(|i| fa_coeff[i])).unwrap().ntt();
1118        let mut per = [[0f64; 15]; 5];
1119        for round in 0..15 {
1120            for (arm, out) in per.iter_mut().enumerate() {
1121                let t = std::time::Instant::now();
1122                for _ in 0..reps {
1123                    let mut x = slain(b);
1124                    let mut c = [0i16; N];
1125                    unsafe {
1126                        match arm {
1127                            0 => x86::ntt(&mut x),
1128                            1 => x86::intt(&mut x),
1129                            2 => x86::basemul(&fa, &x, &mut c),
1130                            3 => {
1131                                x86::ntt(&mut x);
1132                                x86::basemul(&fa, &x, &mut c);
1133                                x86::intt(&mut c);
1134                            }
1135                            _ => c = typed_fa.mul(&Poly::<Coeff>::from_u8(&slain(bu8)).ntt()).intt().into_array(),
1136                        }
1137                    }
1138                    slain((x, c));
1139                }
1140                out[round] = f64::from(reps) / t.elapsed().as_secs_f64();
1141            }
1142        }
1143        for (name, v) in names.iter().zip(per.iter_mut()) {
1144            v.sort_by(f64::total_cmp);
1145            println!("NTT256 {name} {:.9} calls/s  {:.9} ns per call  (range {:.9}-{:.9})", v[7], 1e9 / v[7], v[0], v[14]);
1146        }
1147    }
1148
1149    #[cfg(target_arch = "x86_64")]
1150    #[test]
1151    fn intt_equals_intt_prev() {
1152        if !slain_rt::cpu::avx2() {
1153            return;
1154        }
1155        let mut r = rng(21);
1156        for _ in 0..2000 {
1157            let a: [i16; N] = core::array::from_fn(|_| r() as i16);
1158            let (mut x, mut y) = (a, a);
1159            unsafe {
1160                x86::intt(&mut x);
1161                x86::intt_prev(&mut y);
1162            }
1163            assert_eq!(canon(&x), canon(&y));
1164            assert!(x.iter().all(|&v| i32::from(v).abs() < Q));
1165        }
1166    }
1167
1168    #[cfg(target_arch = "x86_64")]
1169    #[test]
1170    #[ignore]
1171    fn bench_intt_paired() {
1172        use core::hint::black_box as slain;
1173        if !slain_rt::cpu::avx2() {
1174            return;
1175        }
1176        let mut r = rng(5);
1177        let a: [i16; N] = core::array::from_fn(|_| (r() % Q as u64) as i16);
1178        let reps = 400_000u32;
1179        let mut ratio = Vec::new();
1180        for round in 0..15 {
1181            let mut t = [0f64; 2];
1182            for arm in [round % 2, 1 - round % 2] {
1183                let s = std::time::Instant::now();
1184                for _ in 0..reps {
1185                    let mut x = slain(a);
1186                    unsafe {
1187                        if arm == 0 {
1188                            x86::intt_prev(&mut x);
1189                        } else {
1190                            x86::intt(&mut x);
1191                        }
1192                    }
1193                    slain(x);
1194                }
1195                t[arm] = s.elapsed().as_secs_f64() * 1e9 / f64::from(reps);
1196            }
1197            ratio.push(t[0] / t[1]);
1198            println!("round {round}: prev {:.9} ns  new {:.9} ns  prev/new {:.9}", t[0], t[1], t[0] / t[1]);
1199        }
1200        ratio.sort_by(f64::total_cmp);
1201        println!("intt paired prev/new median {:.9}  range {:.9}-{:.9}", ratio[7], ratio[0], ratio[14]);
1202    }
1203
1204    #[cfg(target_arch = "x86_64")]
1205    #[test]
1206    #[ignore = "floor: just floor ntt_piece slain-num"]
1207    fn floor_ntt_piece() {
1208        use core::hint::black_box as slain;
1209        if !slain_rt::cpu::avx2() {
1210            return;
1211        }
1212        let fp0 = 3 * 56 + (3 * 48 + 3 * 56 + 3 * 16) + 17 * 16;
1213        let mut r = rng(13);
1214        let b: [i16; N] = core::array::from_fn(|_| (r() % 256) as i16);
1215        let mut fa: [i16; N] = core::array::from_fn(|_| (r() % 511) as i16 - 255);
1216        ntt(&mut fa);
1217        let calls = 2_000_000u32;
1218        let mut best = 0f64;
1219        for _ in 0..3 {
1220            let t = std::time::Instant::now();
1221            for _ in 0..calls {
1222                let mut x = slain(b);
1223                let mut c = [0i16; N];
1224                unsafe {
1225                    x86::ntt(&mut x);
1226                    x86::basemul(&fa, &x, &mut c);
1227                    x86::intt(&mut c);
1228                }
1229                slain(c);
1230            }
1231            best = best.max(f64::from(calls) / t.elapsed().as_secs_f64());
1232        }
1233        println!("FLOOR_FP0 {fp0}");
1234        println!("FLOOR_PX 81");
1235        println!("FLOOR_CALLS_PER_S {best:.9}");
1236        println!("FLOOR_CORE_HZ {:.9}", slain_rt::cpu::core_hz());
1237    }
1238
1239    #[cfg(target_arch = "x86_64")]
1240    #[test]
1241    #[ignore]
1242    fn bench_ntt_per_second() {
1243        use core::hint::black_box as slain;
1244        if !slain_rt::cpu::avx2() {
1245            return;
1246        }
1247        let mut r = rng(7);
1248        let a: [i16; N] = core::array::from_fn(|_| (r() % Q as u64) as i16);
1249        let reps = 200_000;
1250        let (mut s, mut v) = (Vec::new(), Vec::new());
1251        for round in 0..15 {
1252            for arm in [round % 2, 1 - round % 2] {
1253                let t = std::time::Instant::now();
1254                for _ in 0..reps {
1255                    let mut x = slain(a);
1256                    if arm == 0 {
1257                        ntt_scalar(&mut x);
1258                    } else {
1259                        unsafe { x86::ntt(&mut x) };
1260                    }
1261                    slain(x);
1262                }
1263                let per_s = reps as f64 / t.elapsed().as_secs_f64();
1264                if arm == 0 {
1265                    s.push(per_s)
1266                } else {
1267                    v.push(per_s)
1268                }
1269            }
1270        }
1271        s.sort_by(f64::total_cmp);
1272        v.sort_by(f64::total_cmp);
1273        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]);
1274    }
1275
1276    #[cfg(target_arch = "x86_64")]
1277    #[test]
1278    #[ignore]
1279    fn bench_typed_mul_per_second() {
1280        use core::hint::black_box as slain;
1281        if !slain_rt::cpu::avx2() {
1282            return;
1283        }
1284        let mut r = rng(9);
1285        let a = Poly::<Coeff>::new(core::array::from_fn(|_| (r() % Q as u64) as i16)).unwrap().ntt();
1286        let b = Poly::<Coeff>::new(core::array::from_fn(|_| (r() % Q as u64) as i16)).unwrap().ntt();
1287        let reps = 1_000_000;
1288        let (mut raw, mut typed) = (Vec::new(), Vec::new());
1289        for round in 0..15 {
1290            for arm in [round % 2, 1 - round % 2] {
1291                let t = std::time::Instant::now();
1292                for _ in 0..reps {
1293                    let (x, y) = (slain(a), slain(b));
1294                    if arm == 0 {
1295                        let mut c = [0i16; N];
1296                        unsafe { x86::basemul(x.as_array(), y.as_array(), &mut c) };
1297                        slain(c);
1298                    } else {
1299                        slain(x.mul(&y));
1300                    }
1301                }
1302                let per_s = reps as f64 / t.elapsed().as_secs_f64();
1303                if arm == 0 {
1304                    raw.push(per_s)
1305                } else {
1306                    typed.push(per_s)
1307                }
1308            }
1309        }
1310        raw.sort_by(f64::total_cmp);
1311        typed.sort_by(f64::total_cmp);
1312        println!(
1313            "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})",
1314            raw[7],
1315            typed[7],
1316            typed[7] / raw[7],
1317            raw[0],
1318            raw[14],
1319            typed[0],
1320            typed[14]
1321        );
1322    }
1323}