// NTT mod 998244353, radix-8 DIF/DIT, all-Shoup AVX2, INTERLEAVED (w,w') twiddles.
// Shoup (2 vpmuludq/mul) instead of Montgomery (4 vpmuludq/mul), with the twiddle
// quotient w' stored interleaved alongside w so the 14 loads are ONE contiguous stream.
#pragma GCC target("avx2")
#pragma GCC optimize("O3,unroll-loops")
#include <immintrin.h>
typedef unsigned long long u64;
typedef unsigned int u32;
const u32 MOD = 998244353u;
const u32 NINV = 998244351u;
const u32 R2 = 932051910u;
const u32 ONE = 301989884u; // 2^32 mod p
const u32 G = 3u;
const u32 IROOT = 911660635u; // i
const u32 ZETA = 372528824u; // zeta (primitive 8th root)
const u32 ZETA3 = 488723995u;
const u32 ZETA5 = 625715529u;
const u32 ZETA7 = 509520358u;
const u32 IMINUS = 86583718u; // -i
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 u64 MAGIC = 18479187002ULL; // floor(2^64/p)
const int MAXL = 1 << 21;
const int MAXTW = 2 * (1 << 21) + 128;
alignas(32) static u32 A[MAXL];
alignas(32) static u32 B[MAXL];
alignas(32) static u32 TW_fwd[MAXTW];
alignas(32) static u32 TW_inv[MAXTW];
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& modv, const __m256i& shuf) {
__m256i lo = _mm256_mullo_epi32(a, w);
__m256i qe = _mm256_mul_epu32(a, wp);
__m256i qo = _mm256_mul_epu32(_mm256_srli_epi64(a, 32), _mm256_srli_epi64(wp, 32));
__m256i qe2 = _mm256_srli_epi64(qe, 32);
__m256i qo2 = _mm256_srli_epi64(qo, 32);
__m256i e = _mm256_shuffle_epi8(qe2, shuf);
__m256i o = _mm256_shuffle_epi8(qo2, shuf);
__m256i q = _mm256_unpacklo_epi32(e, o);
__m256i r = _mm256_sub_epi32(lo, _mm256_mullo_epi32(q, modv));
return _mm256_min_epu32(r, _mm256_sub_epi32(r, modv));
}
static inline __m256i add_mod8(__m256i x, __m256i y, const __m256i& p) {
__m256i s = _mm256_add_epi32(x, y);
return _mm256_min_epu32(s, _mm256_sub_epi32(s, p));
}
static inline __m256i sub_mod8(__m256i x, __m256i y, const __m256i& p) {
__m256i d = _mm256_sub_epi32(x, y);
return _mm256_min_epu32(d, _mm256_add_epi32(d, 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 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 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 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 u32 pow_mod(u32 base, u64 e) {
u64 r = 1, b = base;
while (e) { if (e & 1) r = r * b % MOD; b = b * b % MOD; e >>= 1; }
return (u32)r;
}
static u32 shp_quot(u32 w) { return (u32)(((u64)w * MAGIC) >> 32); }
// Generate roots rts[i] = w^i (normal form) for i = 0..n-1.
static void gen_roots_normal(u32* rts, int n, u32 w) {
u32 wp = shp_quot(w);
u32 p8[8];
p8[0] = 1;
for (int k = 1; k < 8; k++) p8[k] = shoup_mul(p8[k-1], w, wp);
u32 w8 = shoup_mul(p8[7], w, wp);
u32 w8p = shp_quot(w8);
__m256i pv = load8(p8);
u32 base = 1;
for (int i = 0; i < n; i += 8) {
u32 basep = shp_quot(base);
store8(rts + i, shoup_mul8(pv, _mm256_set1_epi32(base), _mm256_set1_epi32(basep), _mm256_set1_epi32(MOD), _mm256_setr_epi8(0,1,2,3,8,9,10,11,0,0,0,0,0,0,0,0, 0,1,2,3,8,9,10,11,0,0,0,0,0,0,0,0)));
base = shoup_mul(base, w8, w8p);
}
}
// Build INTERLEAVED twiddles: for each 8-position block store
// [w1[8], w1'[8], w2[8], w2'[8], ..., w7[8], w7'[8]] (14*8 = 112 u32 per block).
static void build_twiddles8_fwd(const u32* rts, u32* tw, int n) {
int off = 0;
for (int len = n; len >= 8; len >>= 3) {
int m = len >> 3;
int step = n / len;
if (m >= 8) {
for (int j = 0; j < m; j += 8) {
u32* wb = tw + off + (j >> 3) * 112;
for (int s = 1; s <= 7; s++) {
u32* base = wb + 2 * (s - 1) * 8;
for (int k = 0; k < 8; k++) {
u32 val = rts[s * (j + k) * step];
base[k] = val;
base[8 + k] = shp_quot(val);
}
}
}
} else {
u32* wb = tw + off;
for (int s = 1; s <= 7; s++) {
u32 val = rts[s * 0 * step];
wb[2*(s-1)] = val;
wb[2*(s-1)+1] = shp_quot(val);
}
}
off += 14 * m;
}
}
static void build_twiddles8_inv(const u32* rts, u32* tw, int n) {
int off = 0;
for (int len = 8; len <= n; len <<= 3) {
int m = len >> 3;
int step = n / len;
if (m >= 8) {
for (int j = 0; j < m; j += 8) {
u32* wb = tw + off + (j >> 3) * 112;
for (int s = 1; s <= 7; s++) {
u32* base = wb + 2 * (s - 1) * 8;
for (int k = 0; k < 8; k++) {
u32 val = rts[s * (j + k) * step];
base[k] = val;
base[8 + k] = shp_quot(val);
}
}
}
} else {
u32* wb = tw + off;
for (int s = 1; s <= 7; s++) {
u32 val = rts[s * 0 * step];
wb[2*(s-1)] = val;
wb[2*(s-1)+1] = shp_quot(val);
}
}
off += 14 * m;
}
}
static void ntt_fwd8(u32* x, int n, const u32* tw) {
const __m256i modv = _mm256_set1_epi32(MOD);
const __m256i shuf = _mm256_setr_epi8(0,1,2,3,8,9,10,11,0,0,0,0,0,0,0,0, 0,1,2,3,8,9,10,11,0,0,0,0,0,0,0,0);
const __m256i iv = _mm256_set1_epi32(IROOT), ipv = _mm256_set1_epi32(I_p);
const __m256i z1v = _mm256_set1_epi32(ZETA), z1pv = _mm256_set1_epi32(Z1_p);
const __m256i z3v = _mm256_set1_epi32(ZETA3), z3pv = _mm256_set1_epi32(Z3_p);
int off = 0;
for (int len = n; len >= 8; len >>= 3) {
int m = len >> 3;
if (m >= 8) {
for (int i = 0; i < n; i += len) {
u32* y = x + i;
for (int j = 0; j < m; j += 8) {
const u32* wb = tw + off + (j >> 3) * 112;
__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,modv);
__m256i e1 = sub_mod8(x0, x4,modv);
__m256i e2 = add_mod8(x2, x6,modv);
__m256i e3 = sub_mod8(x2, x6,modv);
__m256i ie3 = shoup_mul8(e3, iv, ipv, modv, shuf);
__m256i E0 = add_mod8(e0, e2,modv);
__m256i E1 = add_mod8(e1, ie3,modv);
__m256i E2 = sub_mod8(e0, e2,modv);
__m256i E3 = sub_mod8(e1, ie3,modv);
__m256i o0 = add_mod8(x1, x5,modv);
__m256i o1 = sub_mod8(x1, x5,modv);
__m256i o2 = add_mod8(x3, x7,modv);
__m256i o3 = sub_mod8(x3, x7,modv);
__m256i io3 = shoup_mul8(o3, iv, ipv, modv, shuf);
__m256i O0 = add_mod8(o0, o2,modv);
__m256i O1 = add_mod8(o1, io3,modv);
__m256i O2 = sub_mod8(o0, o2,modv);
__m256i O3 = sub_mod8(o1, io3,modv);
__m256i zO1 = shoup_mul8(O1, z1v, z1pv, modv, shuf);
__m256i zO3 = shoup_mul8(O3, z3v, z3pv, modv, shuf);
__m256i iO2 = shoup_mul8(O2, iv, ipv, modv, shuf);
__m256i Y0 = add_mod8(E0, O0,modv);
__m256i Y1 = add_mod8(E1, zO1,modv);
__m256i Y2 = add_mod8(E2, iO2,modv);
__m256i Y3 = add_mod8(E3, zO3,modv);
__m256i Y4 = sub_mod8(E0, O0,modv);
__m256i Y5 = sub_mod8(E1, zO1,modv);
__m256i Y6 = sub_mod8(E2, iO2,modv);
__m256i Y7 = sub_mod8(E3, zO3,modv);
store8(y + j, Y0);
store8(y + j + m, shoup_mul8(Y1, load8(wb + 0), load8(wb + 8), modv, shuf));
store8(y + j + 2*m, shoup_mul8(Y2, load8(wb + 16), load8(wb + 24), modv, shuf));
store8(y + j + 3*m, shoup_mul8(Y3, load8(wb + 32), load8(wb + 40), modv, shuf));
store8(y + j + 4*m, shoup_mul8(Y4, load8(wb + 48), load8(wb + 56), modv, shuf));
store8(y + j + 5*m, shoup_mul8(Y5, load8(wb + 64), load8(wb + 72), modv, shuf));
store8(y + j + 6*m, shoup_mul8(Y6, load8(wb + 80), load8(wb + 88), modv, shuf));
store8(y + j + 7*m, shoup_mul8(Y7, load8(wb + 96), load8(wb + 104), modv, shuf));
}
}
} else {
for (int i = 0; i < n; i += len) {
u32* y = x + i;
const u32* wb = tw + off;
u32 x0 = y[0], x1 = y[m], x2 = y[2*m], x3 = y[3*m], x4 = y[4*m], x5 = y[5*m], x6 = y[6*m], x7 = y[7*m];
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[m] = shoup_mul(add_mod(E1,zO1), wb[0], wb[1]);
y[2*m] = shoup_mul(add_mod(E2,iO2), wb[2], wb[3]);
y[3*m] = shoup_mul(add_mod(E3,zO3), wb[4], wb[5]);
y[4*m] = shoup_mul(sub_mod(E0,O0), wb[6], wb[7]);
y[5*m] = shoup_mul(sub_mod(E1,zO1), wb[8], wb[9]);
y[6*m] = shoup_mul(sub_mod(E2,iO2), wb[10], wb[11]);
y[7*m] = shoup_mul(sub_mod(E3,zO3), wb[12], wb[13]);
}
}
off += 14 * m;
}
}
static void ntt_inv8(u32* x, int n, const u32* tw) {
const __m256i modv = _mm256_set1_epi32(MOD);
const __m256i shuf = _mm256_setr_epi8(0,1,2,3,8,9,10,11,0,0,0,0,0,0,0,0, 0,1,2,3,8,9,10,11,0,0,0,0,0,0,0,0);
const __m256i iv = _mm256_set1_epi32(IMINUS), ipv = _mm256_set1_epi32(IM_p);
const __m256i z1v = _mm256_set1_epi32(ZETA7), z1pv = _mm256_set1_epi32(Z7_p);
const __m256i z3v = _mm256_set1_epi32(ZETA5), z3pv = _mm256_set1_epi32(Z5_p);
for (int i = 0; i < n; i += 8) {
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, 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(x1,x5), o1 = sub_mod(x1,x5), o2 = add_mod(x3,x7), o3 = sub_mod(x3,x7);
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);
}
int off = 14;
for (int len = 64; len <= n; len <<= 3) {
int m = len >> 3;
if (m >= 8) {
for (int i = 0; i < n; i += len) {
u32* y = x + i;
for (int j = 0; j < m; j += 8) {
const u32* wb = tw + off + (j >> 3) * 112;
__m256i y0 = load8(y + j);
__m256i y1 = shoup_mul8(load8(y + j + m), load8(wb + 0), load8(wb + 8), modv, shuf);
__m256i y2 = shoup_mul8(load8(y + j + 2*m), load8(wb + 16), load8(wb + 24), modv, shuf);
__m256i y3 = shoup_mul8(load8(y + j + 3*m), load8(wb + 32), load8(wb + 40), modv, shuf);
__m256i y4 = shoup_mul8(load8(y + j + 4*m), load8(wb + 48), load8(wb + 56), modv, shuf);
__m256i y5 = shoup_mul8(load8(y + j + 5*m), load8(wb + 64), load8(wb + 72), modv, shuf);
__m256i y6 = shoup_mul8(load8(y + j + 6*m), load8(wb + 80), load8(wb + 88), modv, shuf);
__m256i y7 = shoup_mul8(load8(y + j + 7*m), load8(wb + 96), load8(wb + 104), modv, shuf);
__m256i e0 = add_mod8(y0, y4,modv);
__m256i e1 = sub_mod8(y0, y4,modv);
__m256i e2 = add_mod8(y2, y6,modv);
__m256i e3 = sub_mod8(y2, y6,modv);
__m256i ie3 = shoup_mul8(e3, iv, ipv, modv, shuf);
__m256i E0 = add_mod8(e0, e2,modv);
__m256i E1 = add_mod8(e1, ie3,modv);
__m256i E2 = sub_mod8(e0, e2,modv);
__m256i E3 = sub_mod8(e1, ie3,modv);
__m256i o0 = add_mod8(y1, y5,modv);
__m256i o1 = sub_mod8(y1, y5,modv);
__m256i o2 = add_mod8(y3, y7,modv);
__m256i o3 = sub_mod8(y3, y7,modv);
__m256i io3 = shoup_mul8(o3, iv, ipv, modv, shuf);
__m256i O0 = add_mod8(o0, o2,modv);
__m256i O1 = add_mod8(o1, io3,modv);
__m256i O2 = sub_mod8(o0, o2,modv);
__m256i O3 = sub_mod8(o1, io3,modv);
__m256i zO1 = shoup_mul8(O1, z1v, z1pv, modv, shuf);
__m256i zO3 = shoup_mul8(O3, z3v, z3pv, modv, shuf);
__m256i iO2 = shoup_mul8(O2, iv, ipv, modv, shuf);
store8(y + j, add_mod8(E0, O0,modv));
store8(y + j + m, add_mod8(E1, zO1,modv));
store8(y + j + 2*m, add_mod8(E2, iO2,modv));
store8(y + j + 3*m, add_mod8(E3, zO3,modv));
store8(y + j + 4*m, sub_mod8(E0, O0,modv));
store8(y + j + 5*m, sub_mod8(E1, zO1,modv));
store8(y + j + 6*m, sub_mod8(E2, iO2,modv));
store8(y + j + 7*m, sub_mod8(E3, zO3,modv));
}
}
} else {
for (int i = 0; i < n; i += len) {
u32* y = x + i;
const u32* wb = tw + off;
u32 y0 = y[0];
u32 y1 = shoup_mul(y[m], wb[0], wb[1]);
u32 y2 = shoup_mul(y[2*m], wb[2], wb[3]);
u32 y3 = shoup_mul(y[3*m], wb[4], wb[5]);
u32 y4 = shoup_mul(y[4*m], wb[6], wb[7]);
u32 y5 = shoup_mul(y[5*m], wb[8], wb[9]);
u32 y6 = shoup_mul(y[6*m], wb[10], wb[11]);
u32 y7 = shoup_mul(y[7*m], wb[12], wb[13]);
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[m] = add_mod(E1,zO1);
y[2*m] = add_mod(E2,iO2);
y[3*m] = add_mod(E3,zO3);
y[4*m] = sub_mod(E0,O0);
y[5*m] = sub_mod(E1,zO1);
y[6*m] = sub_mod(E2,iO2);
y[7*m] = sub_mod(E3,zO3);
}
}
off += 14 * m;
}
}
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) store8(A + i, mont_mul8(load8((const u32*)(a + i)), _mm256_set1_epi32(R2)));
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) store8(B + i, mont_mul8(load8((const u32*)(b + i)), _mm256_set1_epi32(R2)));
for (; i <= m; i++) B[i] = to_mont(b[i]);
for (; i < L; i++) B[i] = 0;
u32 w = pow_mod(G, (MOD - 1) / L);
u32 wi = pow_mod(w, MOD - 2);
static u32 roots[MAXL], roots_inv[MAXL];
gen_roots_normal(roots, L, w);
gen_roots_normal(roots_inv, L, wi);
build_twiddles8_fwd(roots, TW_fwd, L);
build_twiddles8_inv(roots_inv, TW_inv, L);
ntt_fwd8(A, L, TW_fwd);
ntt_fwd8(B, L, TW_fwd);
for (int k = 0; k < L; k += 8)
store8(A + k, mont_mul8(load8(A + k), load8(B + k)));
ntt_inv8(A, L, TW_inv);
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)
store8(c + k, mont_mul8(load8(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 | 60.694 ms | 71 MB + 656 KB | Accepted | Score: 100 | 显示更多 |