// NTT mod 998244353, radix-8 DIF/DIT, AVX2 Montgomery + Shoup,
// twiddles computed ON THE FLY (no twiddle tables, no twiddle generation).
#include <immintrin.h>
#pragma GCC target("avx2")
typedef unsigned long long u64;
typedef unsigned u32;
const u32 MOD = 998244353u;
const u32 NINV = 998244351u;
const u32 ONE = 301989884u;
const u32 R2 = 932051910u;
const u32 IROOT = 911660635u;
const u32 ZETA = 372528824u;
const u32 ZETA3 = 488723995u;
const u32 ZETA5 = 625715529u;
const u32 ZETA7 = 509520358u;
const u32 IMINUS = 86583718u;
const u32 I_p = 3922439030u;
const u32 IM_p = 372528265u;
const u32 Z1_p = 1602813089u;
const u32 Z3_p = 2102745253u;
const u32 Z5_p = 2692154206u;
const u32 Z7_p = 2192222042u;
const int MAXL = 1 << 21;
static u32 A[MAXL];
static u32 B[MAXL];
static inline __m256i load8(const u32* p) { return _mm256_loadu_si256((const __m256i*)p); }
static inline void store8(u32* p, __m256i v) { _mm256_storeu_si256((__m256i*)p, v); }
static inline __m256i hi32_mul(__m256i a, __m256i b) {
__m256i t0 = _mm256_mul_epu32(a, b);
__m256i t1 = _mm256_mul_epu32(_mm256_srli_si256(a, 4), _mm256_srli_si256(b, 4));
return _mm256_or_si256(_mm256_srli_epi64(t0, 32), _mm256_slli_si256(_mm256_srli_epi64(t1, 32), 4));
}
static inline __m256i mont_mul8(__m256i x, __m256i y) {
const __m256i ninv = _mm256_set1_epi32(NINV);
const __m256i p = _mm256_set1_epi32(MOD);
const __m256i pminus1 = _mm256_set1_epi32(MOD - 1);
const __m256i one = _mm256_set1_epi32(1);
const __m256i zero = _mm256_setzero_si256();
__m256i lo = _mm256_mullo_epi32(x, y);
__m256i m = _mm256_mullo_epi32(lo, ninv);
__m256i res = _mm256_add_epi32(hi32_mul(x, y), hi32_mul(m, p));
__m256i eqz = _mm256_cmpeq_epi32(lo, zero);
res = _mm256_add_epi32(res, _mm256_andnot_si256(eqz, one));
__m256i ge = _mm256_cmpgt_epi32(res, pminus1);
return _mm256_sub_epi32(res, _mm256_and_si256(ge, p));
}
static inline __m256i shoup_mul8(__m256i a, __m256i w, __m256i wp) {
const __m256i p = _mm256_set1_epi32(MOD);
const __m256i pminus1 = _mm256_set1_epi32(MOD - 1);
__m256i lo = _mm256_mullo_epi32(a, w);
__m256i q = hi32_mul(a, wp);
__m256i r = _mm256_sub_epi32(lo, _mm256_mullo_epi32(q, p));
__m256i ge = _mm256_cmpgt_epi32(r, pminus1);
return _mm256_sub_epi32(r, _mm256_and_si256(ge, p));
}
static inline __m256i add_mod8(__m256i x, __m256i y) {
const __m256i p = _mm256_set1_epi32(MOD);
const __m256i pminus1 = _mm256_set1_epi32(MOD - 1);
__m256i s = _mm256_add_epi32(x, y);
__m256i ge = _mm256_cmpgt_epi32(s, pminus1);
return _mm256_sub_epi32(s, _mm256_and_si256(ge, p));
}
static inline __m256i sub_mod8(__m256i x, __m256i y) {
const __m256i p = _mm256_set1_epi32(MOD);
__m256i d = _mm256_sub_epi32(x, y);
__m256i mask = _mm256_cmpgt_epi32(y, x);
return _mm256_add_epi32(d, _mm256_and_si256(mask, p));
}
static inline u32 mont_mul(u32 x, u32 y) {
u64 t = (u64)x * y;
u32 m = (u32)t * NINV;
u64 u = (t + (u64)m * MOD) >> 32;
if (u >= MOD) u -= MOD;
return (u32)u;
}
static inline u32 shoup_mul(u32 a, u32 w, u32 wp) {
u32 q = (u32)(((u64)a * wp) >> 32);
u32 r = (u32)((u64)a * w - (u64)q * MOD);
if (r >= MOD) r -= MOD;
return r;
}
static inline u32 add_mod(u32 x, u32 y) { u32 s = x + y; return s >= MOD ? s - MOD : s; }
static inline u32 sub_mod(u32 x, u32 y) { return x >= y ? x - y : x + MOD - y; }
static inline u32 to_mont(u32 x) { return mont_mul(x, R2); }
static u32 mont_pow(u32 base, u64 e) {
u32 r = ONE, b = base;
while (e) { if (e & 1) r = mont_mul(r, b); b = mont_mul(b, b); e >>= 1; }
return r;
}
static inline __m256i to_mont8(__m256i x) { return mont_mul8(x, _mm256_set1_epi32(R2)); }
// forward DIF radix-8, twiddles on the fly. n = power of 8.
static void ntt_fwd8(u32 *x, int n, u32 w_mont) {
const __m256i I = _mm256_set1_epi32(IROOT);
const __m256i Ip = _mm256_set1_epi32(I_p);
const __m256i Z1 = _mm256_set1_epi32(ZETA);
const __m256i Z1p = _mm256_set1_epi32(Z1_p);
const __m256i Z3 = _mm256_set1_epi32(ZETA3);
const __m256i Z3p = _mm256_set1_epi32(Z3_p);
const __m256i ONE8 = _mm256_set1_epi32(ONE);
for (int len = n; len >= 8; len >>= 3) {
int m = len >> 3;
int step = n / len;
u32 r1 = mont_pow(w_mont, step); // w^step
u32 binit[8];
binit[0] = ONE;
for (int k = 1; k < 8; k++) binit[k] = mont_mul(binit[k-1], r1);
u32 b8 = mont_mul(binit[7], r1); // w^(8*step)
__m256i b8v = _mm256_set1_epi32(b8);
__m256i bv = _mm256_loadu_si256((__m256i*)binit);
if (m >= 8) {
for (int i = 0; i < n; i += len) {
u32 *y = x + i;
__m256i bvv = bv;
for (int j = 0; j < m; j += 8) {
__m256i tw1 = bvv;
__m256i tw2 = mont_mul8(tw1, tw1);
__m256i tw4 = mont_mul8(tw2, tw2);
__m256i tw3 = mont_mul8(tw2, tw1);
__m256i tw6 = mont_mul8(tw4, tw2);
__m256i tw5 = mont_mul8(tw4, tw1);
__m256i tw7 = mont_mul8(tw4, tw3);
__m256i x0 = load8(y + j);
__m256i x1 = load8(y + j + m);
__m256i x2 = load8(y + j + 2*m);
__m256i x3 = load8(y + j + 3*m);
__m256i x4 = load8(y + j + 4*m);
__m256i x5 = load8(y + j + 5*m);
__m256i x6 = load8(y + j + 6*m);
__m256i x7 = load8(y + j + 7*m);
__m256i e0 = add_mod8(x0, x4);
__m256i e1 = sub_mod8(x0, x4);
__m256i e2 = add_mod8(x2, x6);
__m256i e3 = sub_mod8(x2, x6);
__m256i ie3 = shoup_mul8(e3, I, Ip);
__m256i E0 = add_mod8(e0, e2);
__m256i E1 = add_mod8(e1, ie3);
__m256i E2 = sub_mod8(e0, e2);
__m256i E3 = sub_mod8(e1, ie3);
__m256i o0 = add_mod8(x1, x5);
__m256i o1 = sub_mod8(x1, x5);
__m256i o2 = add_mod8(x3, x7);
__m256i o3 = sub_mod8(x3, x7);
__m256i io3 = shoup_mul8(o3, I, Ip);
__m256i O0 = add_mod8(o0, o2);
__m256i O1 = add_mod8(o1, io3);
__m256i O2 = sub_mod8(o0, o2);
__m256i O3 = sub_mod8(o1, io3);
__m256i zO1 = shoup_mul8(O1, Z1, Z1p);
__m256i zO3 = shoup_mul8(O3, Z3, Z3p);
__m256i iO2 = shoup_mul8(O2, I, Ip);
__m256i Y0 = add_mod8(E0, O0);
__m256i Y1 = add_mod8(E1, zO1);
__m256i Y2 = add_mod8(E2, iO2);
__m256i Y3 = add_mod8(E3, zO3);
__m256i Y4 = sub_mod8(E0, O0);
__m256i Y5 = sub_mod8(E1, zO1);
__m256i Y6 = sub_mod8(E2, iO2);
__m256i Y7 = sub_mod8(E3, zO3);
store8(y + j, Y0);
store8(y + j + m, mont_mul8(Y1, tw1));
store8(y + j + 2*m, mont_mul8(Y2, tw2));
store8(y + j + 3*m, mont_mul8(Y3, tw3));
store8(y + j + 4*m, mont_mul8(Y4, tw4));
store8(y + j + 5*m, mont_mul8(Y5, tw5));
store8(y + j + 6*m, mont_mul8(Y6, tw6));
store8(y + j + 7*m, mont_mul8(Y7, tw7));
bvv = mont_mul8(bvv, b8v);
}
}
} else {
// m == 1, twiddles all = 1
for (int i = 0; i < n; i += len) {
u32 *y = x + i;
u32 x0 = y[0], x1 = y[1], x2 = y[2], x3 = y[3], x4 = y[4], x5 = y[5], x6 = y[6], x7 = y[7];
u32 e0 = add_mod(x0,x4), e1 = sub_mod(x0,x4), e2 = add_mod(x2,x6), e3 = sub_mod(x2,x6);
u32 ie3 = shoup_mul(e3, IROOT, I_p);
u32 E0 = add_mod(e0,e2), E1 = add_mod(e1,ie3), E2 = sub_mod(e0,e2), E3 = sub_mod(e1,ie3);
u32 o0 = add_mod(x1,x5), o1 = sub_mod(x1,x5), o2 = add_mod(x3,x7), o3 = sub_mod(x3,x7);
u32 io3 = shoup_mul(o3, IROOT, I_p);
u32 O0 = add_mod(o0,o2), O1 = add_mod(o1,io3), O2 = sub_mod(o0,o2), O3 = sub_mod(o1,io3);
u32 zO1 = shoup_mul(O1,ZETA,Z1_p), zO3 = shoup_mul(O3,ZETA3,Z3_p), iO2 = shoup_mul(O2,IROOT,I_p);
y[0] = add_mod(E0,O0);
y[1] = add_mod(E1,zO1);
y[2] = add_mod(E2,iO2);
y[3] = add_mod(E3,zO3);
y[4] = sub_mod(E0,O0);
y[5] = sub_mod(E1,zO1);
y[6] = sub_mod(E2,iO2);
y[7] = sub_mod(E3,zO3);
}
}
}
}
// inverse DIT radix-8, twiddles on the fly
static void ntt_inv8(u32 *x, int n, u32 wi_mont) {
const __m256i I = _mm256_set1_epi32(IMINUS);
const __m256i Ip = _mm256_set1_epi32(IM_p);
const __m256i Z1 = _mm256_set1_epi32(ZETA7);
const __m256i Z1p = _mm256_set1_epi32(Z7_p);
const __m256i Z3 = _mm256_set1_epi32(ZETA5);
const __m256i Z3p = _mm256_set1_epi32(Z5_p);
for (int len = 8; len <= n; len <<= 3) {
int m = len >> 3;
int step = n / len;
u32 r1 = mont_pow(wi_mont, step);
u32 binit[8];
binit[0] = ONE;
for (int k = 1; k < 8; k++) binit[k] = mont_mul(binit[k-1], r1);
u32 b8 = mont_mul(binit[7], r1);
__m256i b8v = _mm256_set1_epi32(b8);
__m256i bv = _mm256_loadu_si256((__m256i*)binit);
if (m >= 8) {
for (int i = 0; i < n; i += len) {
u32 *y = x + i;
__m256i bvv = bv;
for (int j = 0; j < m; j += 8) {
__m256i tw1 = bvv;
__m256i tw2 = mont_mul8(tw1, tw1);
__m256i tw4 = mont_mul8(tw2, tw2);
__m256i tw3 = mont_mul8(tw2, tw1);
__m256i tw6 = mont_mul8(tw4, tw2);
__m256i tw5 = mont_mul8(tw4, tw1);
__m256i tw7 = mont_mul8(tw4, tw3);
__m256i y0 = load8(y + j);
__m256i y1 = mont_mul8(load8(y + j + m), tw1);
__m256i y2 = mont_mul8(load8(y + j + 2*m), tw2);
__m256i y3 = mont_mul8(load8(y + j + 3*m), tw3);
__m256i y4 = mont_mul8(load8(y + j + 4*m), tw4);
__m256i y5 = mont_mul8(load8(y + j + 5*m), tw5);
__m256i y6 = mont_mul8(load8(y + j + 6*m), tw6);
__m256i y7 = mont_mul8(load8(y + j + 7*m), tw7);
__m256i e0 = add_mod8(y0, y4);
__m256i e1 = sub_mod8(y0, y4);
__m256i e2 = add_mod8(y2, y6);
__m256i e3 = sub_mod8(y2, y6);
__m256i ie3 = shoup_mul8(e3, I, Ip);
__m256i E0 = add_mod8(e0, e2);
__m256i E1 = add_mod8(e1, ie3);
__m256i E2 = sub_mod8(e0, e2);
__m256i E3 = sub_mod8(e1, ie3);
__m256i o0 = add_mod8(y1, y5);
__m256i o1 = sub_mod8(y1, y5);
__m256i o2 = add_mod8(y3, y7);
__m256i o3 = sub_mod8(y3, y7);
__m256i io3 = shoup_mul8(o3, I, Ip);
__m256i O0 = add_mod8(o0, o2);
__m256i O1 = add_mod8(o1, io3);
__m256i O2 = sub_mod8(o0, o2);
__m256i O3 = sub_mod8(o1, io3);
__m256i zO1 = shoup_mul8(O1, Z1, Z1p);
__m256i zO3 = shoup_mul8(O3, Z3, Z3p);
__m256i iO2 = shoup_mul8(O2, I, Ip);
store8(y + j, add_mod8(E0, O0));
store8(y + j + m, add_mod8(E1, zO1));
store8(y + j + 2*m, add_mod8(E2, iO2));
store8(y + j + 3*m, add_mod8(E3, zO3));
store8(y + j + 4*m, sub_mod8(E0, O0));
store8(y + j + 5*m, sub_mod8(E1, zO1));
store8(y + j + 6*m, sub_mod8(E2, iO2));
store8(y + j + 7*m, sub_mod8(E3, zO3));
bvv = mont_mul8(bvv, b8v);
}
}
} else {
for (int i = 0; i < n; i += len) {
u32 *y = x + i;
u32 y0 = y[0], y1 = y[1], y2 = y[2], y3 = y[3], y4 = y[4], y5 = y[5], y6 = y[6], y7 = y[7];
u32 e0 = add_mod(y0,y4), e1 = sub_mod(y0,y4), e2 = add_mod(y2,y6), e3 = sub_mod(y2,y6);
u32 ie3 = shoup_mul(e3, IMINUS, IM_p);
u32 E0 = add_mod(e0,e2), E1 = add_mod(e1,ie3), E2 = sub_mod(e0,e2), E3 = sub_mod(e1,ie3);
u32 o0 = add_mod(y1,y5), o1 = sub_mod(y1,y5), o2 = add_mod(y3,y7), o3 = sub_mod(y3,y7);
u32 io3 = shoup_mul(o3, IMINUS, IM_p);
u32 O0 = add_mod(o0,o2), O1 = add_mod(o1,io3), O2 = sub_mod(o0,o2), O3 = sub_mod(o1,io3);
u32 zO1 = shoup_mul(O1,ZETA7,Z7_p), zO3 = shoup_mul(O3,ZETA5,Z5_p), iO2 = shoup_mul(O2,IMINUS,IM_p);
y[0] = add_mod(E0,O0);
y[1] = add_mod(E1,zO1);
y[2] = add_mod(E2,iO2);
y[3] = add_mod(E3,zO3);
y[4] = sub_mod(E0,O0);
y[5] = sub_mod(E1,zO1);
y[6] = sub_mod(E2,iO2);
y[7] = sub_mod(E3,zO3);
}
}
}
}
void poly_multiply(unsigned *a, int n, unsigned *b, int m, unsigned *c) {
int L = 1;
while (L < n + m + 2) L <<= 1;
if (L < 8) L = 8;
int i = 0;
for (; i + 8 <= n + 1; i += 8) _mm256_storeu_si256((__m256i*)(A + i), to_mont8(_mm256_loadu_si256((__m256i*)(a + i))));
for (; i <= n; i++) A[i] = to_mont(a[i]);
for (; i < L; i++) A[i] = 0;
i = 0;
for (; i + 8 <= m + 1; i += 8) _mm256_storeu_si256((__m256i*)(B + i), to_mont8(_mm256_loadu_si256((__m256i*)(b + i))));
for (; i <= m; i++) B[i] = to_mont(b[i]);
for (; i < L; i++) B[i] = 0;
u32 w = mont_pow(to_mont(3), (MOD - 1) / L);
u32 wi = mont_pow(w, MOD - 2);
ntt_fwd8(A, L, w);
ntt_fwd8(B, L, w);
for (int k = 0; k < L; k += 8)
_mm256_storeu_si256((__m256i*)(A + k), mont_mul8(_mm256_loadu_si256((__m256i*)(A + k)), _mm256_loadu_si256((__m256i*)(B + k))));
ntt_inv8(A, L, wi);
u32 linv_mont = mont_pow(to_mont((u32)(L % MOD)), MOD - 2);
u32 linv_std = mont_mul(linv_mont, 1);
__m256i linvv = _mm256_set1_epi32(linv_std);
int outn = n + m + 1;
int k = 0;
for (; k + 8 <= outn; k += 8)
_mm256_storeu_si256((__m256i*)(c + k), mont_mul8(_mm256_loadu_si256((__m256i*)(A + k)), linvv));
for (; k < outn; k++) c[k] = mont_mul(A[k], linv_std);
}
| Compilation | N/A | N/A | Compile OK | Score: N/A | 显示更多 |
| Testcase #1 | 76.082 ms | 23 MB + 656 KB | Accepted | Score: 100 | 显示更多 |