// 1002: four-step (cache-blocked) NTT mod 998244353, Montgomery arithmetic, AVX2.
// n = m = 1e6; coefficients < 10 so result coefficients (< 8.1e7) are exact in Z_p.
//
// Pipeline (N = N1*N2):
// L : natural x[n1*N2+n2] -> T1[n2*N1+n1] (tiled transpose + Montgomery convert)
// D1 : DIF_N1 on every row of T1 (natural in, digit-reversed out)
// P : T1[n2][p] -> T2[PERM[p]][n2] (tiled transpose; row permutation is free)
// W : T2[k1][n2] *= w_N^{k1*n2} (per-row geometric progression)
// D2 : DIF_N2 on every row of T2
// spectrum: T2[k1*N2+p] = X[k1 + N1*perm2(p)]
// inverse: DIT_N2 -> untwiddle -> P^-1 -> DIT_N1 -> store natural (with 1/N scale).
#include <immintrin.h>
#include <cstdint>
#include <cstdio>
#include <cstring>
#include <cstdlib>
typedef uint32_t u32;
typedef uint64_t u64;
static const u32 P = 998244353u;
static const u32 RMOD = 301989884u; // 2^32 mod P = mont(1)
static const u32 P1M = P - RMOD; // mont(-1)
static const u32 R2MOD = 932051910u; // 2^64 mod P
static const u32 NPM = 998244351u; // -P^{-1} mod 2^32
static const u32 IVAL = 911660635u; // g^((P-1)/4) (plain)
static const u32 IVM = 691295370u; // mont(I)
static const u32 IIVM = 306948983u; // mont(-I)
#define AVX2 __attribute__((target("avx2")))
static inline u32 addm(u32 a, u32 b) { u32 s = a + b; return s >= P ? s - P : s; }
static inline u32 subm(u32 a, u32 b) { return a >= b ? a - b : a + P - b; }
static inline u32 mmul1(u32 a, u32 b) {
u64 t = (u64)a * b;
u32 m = (u32)((u32)t * NPM);
u32 r = (u32)((t + (u64)m * P) >> 32);
return r >= P ? r - P : r;
}
static u32 fpow(u32 a, u64 e) { u64 r = 1, b = a; while (e) { if (e & 1) r = r * b % P; b = b * b % P; e >>= 1; } return (u32)r; }
AVX2 static inline __m256i vaddp(__m256i a, __m256i b) {
const __m256i pv = _mm256_set1_epi32((int)P);
__m256i s = _mm256_add_epi32(a, b);
return _mm256_min_epu32(s, _mm256_sub_epi32(s, pv));
}
AVX2 static inline __m256i vsubp(__m256i a, __m256i b) {
const __m256i pv = _mm256_set1_epi32((int)P);
__m256i d = _mm256_sub_epi32(a, b);
return _mm256_min_epu32(d, _mm256_add_epi32(d, pv));
}
AVX2 static inline __m256i vmul(__m256i a, __m256i b) {
const __m256i np = _mm256_set1_epi32((int)NPM);
const __m256i pv = _mm256_set1_epi32((int)P);
__m256i lo = _mm256_mul_epu32(a, b);
__m256i hi = _mm256_mul_epu32(_mm256_srli_epi64(a, 32), _mm256_srli_epi64(b, 32));
__m256i ml = _mm256_mul_epu32(lo, np);
__m256i mh = _mm256_mul_epu32(hi, np);
__m256i sl = _mm256_add_epi64(lo, _mm256_mul_epu32(ml, pv));
__m256i sh = _mm256_add_epi64(hi, _mm256_mul_epu32(mh, pv));
__m256i r = _mm256_blend_epi32(_mm256_srli_epi64(sl, 32), sh, 0xAA);
return _mm256_min_epu32(r, _mm256_sub_epi32(r, pv));
}
AVX2 static inline __m256i vload(const u32 *p) { return _mm256_loadu_si256((const __m256i *)p); }
AVX2 static inline void vstore(u32 *p, __m256i v) { _mm256_storeu_si256((__m256i *)p, v); }
AVX2 static inline void tr8x8(__m256i *r) {
__m256i t0 = _mm256_unpacklo_epi32(r[0], r[1]), t1 = _mm256_unpackhi_epi32(r[0], r[1]);
__m256i t2 = _mm256_unpacklo_epi32(r[2], r[3]), t3 = _mm256_unpackhi_epi32(r[2], r[3]);
__m256i t4 = _mm256_unpacklo_epi32(r[4], r[5]), t5 = _mm256_unpackhi_epi32(r[4], r[5]);
__m256i t6 = _mm256_unpacklo_epi32(r[6], r[7]), t7 = _mm256_unpackhi_epi32(r[6], r[7]);
__m256i u0 = _mm256_unpacklo_epi64(t0, t2), u1 = _mm256_unpackhi_epi64(t0, t2);
__m256i u2 = _mm256_unpacklo_epi64(t1, t3), u3 = _mm256_unpackhi_epi64(t1, t3);
__m256i u4 = _mm256_unpacklo_epi64(t4, t6), u5 = _mm256_unpackhi_epi64(t4, t6);
__m256i u6 = _mm256_unpacklo_epi64(t5, t7), u7 = _mm256_unpackhi_epi64(t5, t7);
r[0] = _mm256_permute2x128_si256(u0, u4, 0x20);
r[1] = _mm256_permute2x128_si256(u1, u5, 0x20);
r[2] = _mm256_permute2x128_si256(u2, u6, 0x20);
r[3] = _mm256_permute2x128_si256(u3, u7, 0x20);
r[4] = _mm256_permute2x128_si256(u0, u4, 0x31);
r[5] = _mm256_permute2x128_si256(u1, u5, 0x31);
r[6] = _mm256_permute2x128_si256(u2, u6, 0x31);
r[7] = _mm256_permute2x128_si256(u3, u7, 0x31);
}
// ---- scalar reference row transforms (small sizes / fallback) ----
static void dif4s(u32 *a, int n, int M, const u32 *T1, const u32 *T2, const u32 *T3) {
int L = M << 2;
for (int base = 0; base < n; base += L) {
u32 *p = a + base;
for (int j = 0; j < M; j++) {
u32 x0 = p[j], x1 = p[j + M], x2 = p[j + 2 * M], x3 = p[j + 3 * M];
u32 t0 = addm(x0, x2), t1 = subm(x0, x2), t2 = addm(x1, x3), t3 = mmul1(subm(x1, x3), IVM);
p[j] = addm(t0, t2);
p[j + M] = mmul1(addm(t1, t3), T1[j]);
p[j + 2 * M] = mmul1(subm(t0, t2), T2[j]);
p[j + 3 * M] = mmul1(subm(t1, t3), T3[j]);
}
}
}
static void dit4s(u32 *a, int n, int M, const u32 *T1, const u32 *T2, const u32 *T3) {
int L = M << 2;
for (int base = 0; base < n; base += L) {
u32 *p = a + base;
for (int j = 0; j < M; j++) {
u32 c0 = p[j], c1 = mmul1(p[j + M], T1[j]), c2 = mmul1(p[j + 2 * M], T2[j]), c3 = mmul1(p[j + 3 * M], T3[j]);
u32 t0 = addm(c0, c2), t1 = subm(c0, c2), t2 = addm(c1, c3), t3 = mmul1(subm(c1, c3), IIVM);
p[j] = addm(t0, t2);
p[j + M] = addm(t1, t3);
p[j + 2 * M] = subm(t0, t2);
p[j + 3 * M] = subm(t1, t3);
}
}
}
AVX2 static void dif4v(u32 *a, int n, int M, const u32 *T1, const u32 *T2, const u32 *T3) {
const __m256i iv = _mm256_set1_epi32((int)IVM);
const int L = M << 2;
for (int base = 0; base < n; base += L) {
u32 *p = a + base;
int j = 0;
for (; j + 16 <= M; j += 16) {
__m256i A0 = vload(p + j), A1 = vload(p + j + M), A2 = vload(p + j + 2 * M), A3 = vload(p + j + 3 * M);
__m256i B0 = vload(p + j + 8), B1 = vload(p + j + M + 8), B2 = vload(p + j + 2 * M + 8), B3 = vload(p + j + 3 * M + 8);
__m256i ta0 = vaddp(A0, A2), ta1 = vsubp(A0, A2), ta2 = vaddp(A1, A3), ta3 = vmul(vsubp(A1, A3), iv);
__m256i tb0 = vaddp(B0, B2), tb1 = vsubp(B0, B2), tb2 = vaddp(B1, B3), tb3 = vmul(vsubp(B1, B3), iv);
__m256i oa1 = vaddp(ta1, ta3), oa3 = vsubp(ta1, ta3), ob1 = vaddp(tb1, tb3), ob3 = vsubp(tb1, tb3);
__m256i oa0 = vaddp(ta0, ta2), oa2 = vsubp(ta0, ta2), ob0 = vaddp(tb0, tb2), ob2 = vsubp(tb0, tb2);
vstore(p + j, oa0);
vstore(p + j + 8, ob0);
vstore(p + j + M, vmul(oa1, vload(T1 + j)));
vstore(p + j + M + 8, vmul(ob1, vload(T1 + j + 8)));
vstore(p + j + 2 * M, vmul(oa2, vload(T2 + j)));
vstore(p + j + 2 * M + 8, vmul(ob2, vload(T2 + j + 8)));
vstore(p + j + 3 * M, vmul(oa3, vload(T3 + j)));
vstore(p + j + 3 * M + 8, vmul(ob3, vload(T3 + j + 8)));
}
for (; j < M; j += 8) {
__m256i x0 = vload(p + j), x1 = vload(p + j + M), x2 = vload(p + j + 2 * M), x3 = vload(p + j + 3 * M);
__m256i t0 = vaddp(x0, x2), t1 = vsubp(x0, x2), t2 = vaddp(x1, x3), t3 = vmul(vsubp(x1, x3), iv);
vstore(p + j, vaddp(t0, t2));
vstore(p + j + M, vmul(vaddp(t1, t3), vload(T1 + j)));
vstore(p + j + 2 * M, vmul(vsubp(t0, t2), vload(T2 + j)));
vstore(p + j + 3 * M, vmul(vsubp(t1, t3), vload(T3 + j)));
}
}
}
AVX2 static void dit4v(u32 *a, int n, int M, const u32 *T1, const u32 *T2, const u32 *T3) {
const __m256i iv = _mm256_set1_epi32((int)IIVM);
const int L = M << 2;
for (int base = 0; base < n; base += L) {
u32 *p = a + base;
int j = 0;
for (; j + 16 <= M; j += 16) {
__m256i c0 = vload(p + j), d0 = vload(p + j + 8);
__m256i c1 = vmul(vload(p + j + M), vload(T1 + j));
__m256i d1 = vmul(vload(p + j + M + 8), vload(T1 + j + 8));
__m256i c2 = vmul(vload(p + j + 2 * M), vload(T2 + j));
__m256i d2 = vmul(vload(p + j + 2 * M + 8), vload(T2 + j + 8));
__m256i c3 = vmul(vload(p + j + 3 * M), vload(T3 + j));
__m256i d3 = vmul(vload(p + j + 3 * M + 8), vload(T3 + j + 8));
__m256i ta0 = vaddp(c0, c2), ta1 = vsubp(c0, c2), ta2 = vaddp(c1, c3), ta3 = vmul(vsubp(c1, c3), iv);
__m256i tb0 = vaddp(d0, d2), tb1 = vsubp(d0, d2), tb2 = vaddp(d1, d3), tb3 = vmul(vsubp(d1, d3), iv);
vstore(p + j, vaddp(ta0, ta2));
vstore(p + j + 8, vaddp(tb0, tb2));
vstore(p + j + M, vaddp(ta1, ta3));
vstore(p + j + M + 8, vaddp(tb1, tb3));
vstore(p + j + 2 * M, vsubp(ta0, ta2));
vstore(p + j + 2 * M + 8, vsubp(tb0, tb2));
vstore(p + j + 3 * M, vsubp(ta1, ta3));
vstore(p + j + 3 * M + 8, vsubp(tb1, tb3));
}
for (; j < M; j += 8) {
__m256i c0 = vload(p + j);
__m256i c1 = vmul(vload(p + j + M), vload(T1 + j));
__m256i c2 = vmul(vload(p + j + 2 * M), vload(T2 + j));
__m256i c3 = vmul(vload(p + j + 3 * M), vload(T3 + j));
__m256i t0 = vaddp(c0, c2), t1 = vsubp(c0, c2), t2 = vaddp(c1, c3), t3 = vmul(vsubp(c1, c3), iv);
vstore(p + j, vaddp(t0, t2));
vstore(p + j + M, vaddp(t1, t3));
vstore(p + j + 2 * M, vsubp(t0, t2));
vstore(p + j + 3 * M, vsubp(t1, t3));
}
}
}
AVX2 static void stage2(u32 *a, int n) {
for (int i = 0; i < n; i += 8) {
__m256i v = vload(a + i);
__m256i t = _mm256_shuffle_epi32(v, 0xB1);
__m256i s = vaddp(v, t);
__m256i d = vsubp(t, v);
vstore(a + i, _mm256_blend_epi32(s, d, 0xAA));
}
}
// L=8 radix-4 stage (M=2): 8 blocks (64 elements) per iteration via 8x8 transposes
AVX2 static void tail8(u32 *a, int n, const u32 *T1, const u32 *T2, const u32 *T3, int dir) {
const __m256i iv = _mm256_set1_epi32((int)(dir ? IIVM : IVM));
const __m256i t1 = _mm256_set1_epi32((int)T1[1]), t2 = _mm256_set1_epi32((int)T2[1]), t3 = _mm256_set1_epi32((int)T3[1]);
for (int g = 0; g < n; g += 64) {
__m256i r[8];
for (int i = 0; i < 8; i++) r[i] = vload(a + g + 8 * i);
tr8x8(r); // r[j] = position j across the 8 blocks
for (int j = 0; j < 2; j++) {
__m256i x0 = r[j], x1 = r[j + 2], x2 = r[j + 4], x3 = r[j + 6];
if (dir && j) { x1 = vmul(x1, t1); x2 = vmul(x2, t2); x3 = vmul(x3, t3); }
__m256i a0 = vaddp(x0, x2), a1 = vsubp(x0, x2), a2 = vaddp(x1, x3), a3 = vmul(vsubp(x1, x3), iv);
__m256i b0 = vaddp(a0, a2), b1 = vaddp(a1, a3), b2 = vsubp(a0, a2), b3 = vsubp(a1, a3);
if (!dir && j) { b1 = vmul(b1, t1); b2 = vmul(b2, t2); b3 = vmul(b3, t3); }
r[j] = b0; r[j + 2] = b1; r[j + 4] = b2; r[j + 6] = b3;
}
tr8x8(r);
for (int i = 0; i < 8; i++) vstore(a + g + 8 * i, r[i]);
}
}
// L=16 radix-4 stage (M=4): 2 blocks (32 elements) per iteration
AVX2 static void tail16(u32 *a, int n, const u32 *T1, const u32 *T2, const u32 *T3, int dir) {
const __m256i iv = _mm256_set1_epi32((int)(dir ? IIVM : IVM));
const __m256i t1 = _mm256_broadcastsi128_si256(_mm_loadu_si128((const __m128i *)T1));
const __m256i t2 = _mm256_broadcastsi128_si256(_mm_loadu_si128((const __m128i *)T2));
const __m256i t3 = _mm256_broadcastsi128_si256(_mm_loadu_si128((const __m128i *)T3));
for (int base = 0; base < n; base += 32) {
__m256i A = vload(a + base), B = vload(a + base + 8), C = vload(a + base + 16), D = vload(a + base + 24);
__m256i x0 = _mm256_permute2x128_si256(A, C, 0x20);
__m256i x1 = _mm256_permute2x128_si256(A, C, 0x31);
__m256i x2 = _mm256_permute2x128_si256(B, D, 0x20);
__m256i x3 = _mm256_permute2x128_si256(B, D, 0x31);
if (dir) { x1 = vmul(x1, t1); x2 = vmul(x2, t2); x3 = vmul(x3, t3); }
__m256i r0 = vaddp(x0, x2), r1 = vsubp(x0, x2), r2 = vaddp(x1, x3), r3 = vmul(vsubp(x1, x3), iv);
__m256i b0 = vaddp(r0, r2), b1 = vaddp(r1, r3), b2 = vsubp(r0, r2), b3 = vsubp(r1, r3);
if (!dir) { b1 = vmul(b1, t1); b2 = vmul(b2, t2); b3 = vmul(b3, t3); }
vstore(a + base, _mm256_permute2x128_si256(b0, b1, 0x20));
vstore(a + base + 8, _mm256_permute2x128_si256(b2, b3, 0x20));
vstore(a + base + 16, _mm256_permute2x128_si256(b0, b1, 0x31));
vstore(a + base + 24, _mm256_permute2x128_si256(b2, b3, 0x31));
}
}
// L=4 radix-4 stage (M=1, no twiddles): 2 groups (8 elements) per iteration
AVX2 static void tail4(u32 *a, int n, int dir) {
const __m256i K = _mm256_setr_epi32((int)RMOD, (int)(dir ? IIVM : IVM), (int)P1M, (int)(dir ? IVM : IIVM),
(int)RMOD, (int)(dir ? IIVM : IVM), (int)P1M, (int)(dir ? IVM : IIVM));
for (int i = 0; i < n; i += 8) {
__m256i v = vload(a + i);
__m256i w = _mm256_shuffle_epi32(v, 0x4E);
__m256i s = vaddp(v, w);
__m256i d = vsubp(v, w);
__m256i E = _mm256_unpacklo_epi32(s, d);
__m256i G = _mm256_shuffle_epi32(E, 0xD8);
vstore(a + i, vaddp(_mm256_shuffle_epi32(G, 0x88), vmul(_mm256_shuffle_epi32(G, 0xDD), K)));
}
}
// ---- plans ----
struct StageTab { int M; const u32 *w1, *w2, *w3; };
struct RowPlan { int n; int ns; int r2; StageTab st[8]; };
static u32 *buildPlan(RowPlan &pl, int n, int inv, u32 *pool) {
int Ls[8], c = 0, L = n;
while (L >= 4) { Ls[c++] = L; L >>= 2; }
pl.n = n; pl.r2 = (L == 2); pl.ns = c;
u32 *p = pool;
for (int k = 0; k < c; k++) {
int Li = inv ? Ls[c - 1 - k] : Ls[k];
int M = Li >> 2;
u32 wl = fpow(3, (P - 1) / Li);
if (inv) wl = fpow(wl, P - 2);
u32 *t1 = p, *t2 = p + M, *t3 = p + 2 * M; p += 3 * M;
pl.st[k].M = M; pl.st[k].w1 = t1; pl.st[k].w2 = t2; pl.st[k].w3 = t3;
u64 cur = 1;
for (int j = 0; j < M; j++) {
u32 v1 = (u32)cur, v2 = (u32)((u64)v1 * v1 % P), v3 = (u32)((u64)v2 * v1 % P);
t1[j] = (u32)((u64)v1 * RMOD % P);
t2[j] = (u32)((u64)v2 * RMOD % P);
t3[j] = (u32)((u64)v3 * RMOD % P);
cur = cur * wl % P;
}
}
return p;
}
AVX2 static void rowFwd(u32 *a, const RowPlan &pl) {
int n = pl.n;
for (int s = 0; s < pl.ns; s++) {
int M = pl.st[s].M;
if (M >= 8) dif4v(a, n, M, pl.st[s].w1, pl.st[s].w2, pl.st[s].w3);
else if (M == 4 && n >= 32) tail16(a, n, pl.st[s].w1, pl.st[s].w2, pl.st[s].w3, 0);
else if (M == 2 && n >= 64) tail8(a, n, pl.st[s].w1, pl.st[s].w2, pl.st[s].w3, 0);
else if (M == 1 && n >= 8) tail4(a, n, 0);
else dif4s(a, n, M, pl.st[s].w1, pl.st[s].w2, pl.st[s].w3);
}
if (pl.r2) stage2(a, n);
}
AVX2 static void rowInv(u32 *a, const RowPlan &pl) {
int n = pl.n;
if (pl.r2) stage2(a, n);
for (int s = 0; s < pl.ns; s++) {
int M = pl.st[s].M;
if (M >= 8) dit4v(a, n, M, pl.st[s].w1, pl.st[s].w2, pl.st[s].w3);
else if (M == 4 && n >= 32) tail16(a, n, pl.st[s].w1, pl.st[s].w2, pl.st[s].w3, 1);
else if (M == 2 && n >= 64) tail8(a, n, pl.st[s].w1, pl.st[s].w2, pl.st[s].w3, 1);
else if (M == 1 && n >= 8) tail4(a, n, 1);
else dit4s(a, n, M, pl.st[s].w1, pl.st[s].w2, pl.st[s].w3);
}
}
// ---- global config ----
static int N1 = 2048, N2 = 1024;
static int S1 = 2048 + 16, S2 = 1024 + 16; // padded row strides (avoid L1 set conflicts)
static u32 *BUF0, *BUF1, *BUF2;
static u32 *PERM;
static u32 *FROOT, *IROOT;
static u32 SCLN; // mont(1/N)
static RowPlan fwd1, inv1, fwd2, inv2;
static void buildPerm(int n, u32 *perm) {
int Ls[8], rs[8], c = 0, L = n;
while (L >= 4) { Ls[c] = L; rs[c] = 4; c++; L >>= 2; }
if (L == 2) { Ls[c] = 2; rs[c] = 2; c++; }
for (int p = 0; p < n; p++) {
long k = 0, R = 1;
for (int i = 0; i < c; i++) { k += (long)((p % Ls[i]) / (Ls[i] / rs[i])) * R; R *= rs[i]; }
perm[p] = (u32)k;
}
}
// natural x -> T1[n2*N1+n1] in Montgomery form, zero padded (16x16 macro tiles)
AVX2 static void loadNat(const u32 *x, int len) {
const __m256i r2 = _mm256_set1_epi32((int)R2MOD);
const __m256i z = _mm256_setzero_si256();
u32 tmp[8];
int full = (len / N2) & ~7; // n1 rows entirely below len (multiple of 8)
if (full > N1) full = N1;
for (int n1 = 0; n1 < N1; n1 += 16) {
int mode = (n1 + 16 <= full) ? 0 : (n1 >= full + 16 ? 2 : 1); // 0 fast, 1 boundary, 2 zero
for (int n2 = 0; n2 < N2; n2 += 16)
for (int rh = 0; rh < 2; rh++)
for (int h = 0; h < 2; h++) {
__m256i tile[8];
int r0 = n1 + rh * 8, c0 = n2 + h * 8;
if (mode == 0) {
for (int i = 0; i < 8; i++) tile[i] = vload(x + (size_t)(r0 + i) * N2 + c0);
} else if (mode == 2) {
for (int i = 0; i < 8; i++) tile[i] = z;
} else {
for (int i = 0; i < 8; i++) {
int base = (r0 + i) * N2 + c0;
if (base + 8 <= len) tile[i] = vload(x + base);
else { for (int t = 0; t < 8; t++) tmp[t] = (base + t < len) ? x[base + t] : 0; tile[i] = vload(tmp); }
}
}
tr8x8(tile);
if (mode == 2) for (int i = 0; i < 8; i++) vstore(BUF0 + (size_t)(c0 + i) * S1 + r0, z);
else for (int i = 0; i < 8; i++) vstore(BUF0 + (size_t)(c0 + i) * S1 + r0, vmul(tile[i], r2));
}
}
}
AVX2 static void storeNat(u32 *out, int outlen) {
const __m256i sc = _mm256_set1_epi32((int)SCLN);
u32 tmp[8];
for (int n1 = 0; n1 < N1; n1 += 16)
for (int n2 = 0; n2 < N2; n2 += 16)
for (int rh = 0; rh < 2; rh++)
for (int h = 0; h < 2; h++) {
__m256i tile[8];
int r0 = n1 + rh * 8, c0 = n2 + h * 8;
for (int i = 0; i < 8; i++) tile[i] = vload(BUF0 + (size_t)(c0 + i) * S1 + r0);
tr8x8(tile);
for (int i = 0; i < 8; i++) {
int base = (r0 + i) * N2 + c0;
if (base + 8 <= outlen) vstore(out + base, vmul(tile[i], sc));
else { vstore(tmp, vmul(tile[i], sc)); for (int t = 0; t < 8 && base + t < outlen; t++) out[base + t] = tmp[t]; }
}
}
}
AVX2 static void transpF(const u32 *src, u32 *dst) {
for (int r0 = 0; r0 < N2; r0 += 16)
for (int c0 = 0; c0 < N1; c0 += 16)
for (int h = 0; h < 2; h++)
for (int rh = 0; rh < 2; rh++) {
__m256i r[8];
for (int i = 0; i < 8; i++) r[i] = vload(src + (size_t)(r0 + rh * 8 + i) * S1 + c0 + h * 8);
tr8x8(r);
for (int i = 0; i < 8; i++) vstore(dst + (size_t)PERM[c0 + h * 8 + i] * S2 + r0 + rh * 8, r[i]);
}
}
AVX2 static void transpI(const u32 *src, u32 *dst) {
for (int r0 = 0; r0 < N2; r0 += 16)
for (int c0 = 0; c0 < N1; c0 += 16)
for (int h = 0; h < 2; h++)
for (int rh = 0; rh < 2; rh++) {
__m256i r[8];
for (int i = 0; i < 8; i++) r[i] = vload(src + (size_t)PERM[c0 + h * 8 + i] * S2 + r0 + rh * 8);
tr8x8(r);
for (int i = 0; i < 8; i++) vstore(dst + (size_t)(r0 + rh * 8 + i) * S1 + c0 + h * 8, r[i]);
}
}
// multiply row (length len = 32*k) by the progression u^0..u^{len-1}
AVX2 static void twiddleRow(u32 *row, int len, u32 u) {
u32 R[40], Q[40];
R[0] = RMOD;
for (int d = 1; d <= 8; d++) R[d] = mmul1(R[d - 1], u);
const __m256i u8 = _mm256_set1_epi32((int)R[8]);
for (int d = 8; d < 32; d += 8) vstore(R + d, vmul(vload(R + d - 8), u8));
u32 u32v = mmul1(R[24], R[8]);
Q[0] = RMOD;
for (int c = 1; c <= 8; c++) Q[c] = mmul1(Q[c - 1], u32v);
const __m256i uq = _mm256_set1_epi32((int)Q[8]);
for (int c = 8; c < 32; c += 8) vstore(Q + c, vmul(vload(Q + c - 8), uq));
for (int c = 0; c < len / 32; c++) {
const __m256i q = _mm256_set1_epi32((int)Q[c]);
for (int d = 0; d < 32; d += 8) {
__m256i w = vmul(q, vload(R + d));
u32 *rp = row + c * 32 + d;
vstore(rp, vmul(vload(rp), w));
}
}
}
AVX2 static void conv(const u32 *a, int la, const u32 *b, int lb, u32 *c, int lc) {
loadNat(a, la);
for (int n2 = 0; n2 < N2; n2++) rowFwd(BUF0 + (size_t)n2 * S1, fwd1);
transpF(BUF0, BUF1);
for (int k1 = 0; k1 < N1; k1++) {
twiddleRow(BUF1 + (size_t)k1 * S2, N2, FROOT[k1]);
rowFwd(BUF1 + (size_t)k1 * S2, fwd2);
}
loadNat(b, lb);
for (int n2 = 0; n2 < N2; n2++) rowFwd(BUF0 + (size_t)n2 * S1, fwd1);
transpF(BUF0, BUF2);
for (int k1 = 0; k1 < N1; k1++) {
twiddleRow(BUF2 + (size_t)k1 * S2, N2, FROOT[k1]);
rowFwd(BUF2 + (size_t)k1 * S2, fwd2);
}
for (int k1 = 0; k1 < N1; k1++)
for (int i = 0; i < N2; i += 8) vstore(BUF1 + (size_t)k1 * S2 + i, vmul(vload(BUF1 + (size_t)k1 * S2 + i), vload(BUF2 + (size_t)k1 * S2 + i)));
for (int k1 = 0; k1 < N1; k1++) {
rowInv(BUF1 + (size_t)k1 * S2, inv2);
twiddleRow(BUF1 + (size_t)k1 * S2, N2, IROOT[k1]);
}
transpI(BUF1, BUF0);
for (int n2 = 0; n2 < N2; n2++) rowInv(BUF0 + (size_t)n2 * S1, inv1);
storeNat(c, lc);
}
static void setup(int n1, int n2, u32 *b0, u32 *b1, u32 *b2, u32 *perm, u32 *fr, u32 *ir, u32 *pool) {
N1 = n1; N2 = n2; BUF0 = b0; BUF1 = b1; BUF2 = b2; PERM = perm; FROOT = fr; IROOT = ir;
S1 = n1 + 16; S2 = n2 + 16;
buildPerm(n1, perm);
u32 *p = pool;
p = buildPlan(fwd1, n1, 0, p);
p = buildPlan(inv1, n1, 1, p);
p = buildPlan(fwd2, n2, 0, p);
p = buildPlan(inv2, n2, 1, p);
SCLN = fpow((u32)(n1 * n2), P - 2); // plain 1/N (buffer holds Montgomery form)
u32 w = fpow(3, (P - 1) / (n1 * n2)), wi = fpow(w, P - 2);
u64 cur = 1, curi = 1;
for (int k = 0; k < n1; k++) {
fr[k] = (u32)(cur * RMOD % P);
ir[k] = (u32)(curi * RMOD % P);
cur = cur * w % P;
curi = curi * wi % P;
}
}
// ---------------- judge entry ----------------
#define MAXN (2048 * (1024 + 16))
static u32 g_buf0[MAXN], g_buf1[MAXN], g_buf2[MAXN];
static u32 g_perm[4096], g_froot[4096], g_iroot[4096];
static u32 g_pool[80000];
static int g_ready = 0;
static void ensure(void) {
if (g_ready) return;
g_ready = 1;
setup(2048, 1024, g_buf0, g_buf1, g_buf2, g_perm, g_froot, g_iroot, g_pool);
}
static void convScalar(const u32 *a, int la, const u32 *b, int lb, u32 *c) {
int lc = la + lb - 1;
for (int i = 0; i < lc; i++) c[i] = 0;
for (int i = 0; i < la; i++) {
u64 ai = a[i];
if (!ai) continue;
for (int j = 0; j < lb; j++) c[i + j] = (u32)((c[i + j] + ai * b[j]) % P);
}
}
void poly_multiply(unsigned *a, int n, unsigned *b, int m, unsigned *c) {
int la = n + 1, lb = m + 1, lc = la + lb - 1;
if (lc <= 1024 || (u64)la * lb <= 200000) { convScalar(a, la, b, lb, c); return; }
ensure();
conv(a, la, b, lb, c, lc);
}
| Compilation | N/A | N/A | Compile OK | Score: N/A | 显示更多 |
| Testcase #1 | 47.441 ms | 32 MB + 20 KB | Accepted | Score: 100 | 显示更多 |