代码

// NTT mod 998244353, radix-4 DIF/DIT, AVX2 Montgomery, contiguous twiddle tables
#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 int MAXL = 1 << 21;
const int MAXTW = (1 << 20) * 3 + 64;

static u32 A[MAXL];
static u32 B[MAXL];
static u32 roots[MAXL];
static u32 roots_inv[MAXL];
static u32 TW_fwd[MAXTW];
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 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 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 void gen_roots(u32* rts, int n, u32 w_mont) {
    u32 p[8];
    p[0] = ONE;
    for (int k = 1; k < 8; k++) p[k] = mont_mul(p[k-1], w_mont);
    u32 w8 = mont_mul(p[7], w_mont);
    __m256i pv = _mm256_loadu_si256((__m256i*)p);
    u32 base = ONE;
    for (int i = 0; i < n; i += 8) {
        _mm256_storeu_si256((__m256i*)(rts + i), mont_mul8(pv, _mm256_set1_epi32(base)));
        base = mont_mul(base, w8);
    }
}
static inline __m256i to_mont8(__m256i x) { return mont_mul8(x, _mm256_set1_epi32(R2)); }


// Build per-stage contiguous twiddle tables: [w1[0..m), w2[0..m), w3[0..m)] for each radix-4 stage.
static void build_twiddles_fwd(const u32* rts, u32* tw, int n) {
    int off = 0;
    int start_len = (n % 3 == 2) ? (n >> 1) : n;
    for (int len = start_len; len >= 4; len >>= 2) {
        int m = len >> 2;
        int step = n / len;
        u32* w1 = tw + off;
        u32* w2 = tw + off + m;
        u32* w3 = tw + off + 2 * m;
        for (int j = 0; j < m; j++) w1[j] = rts[j * step];
        for (int j = 0; j < m; j++) w2[j] = rts[2 * j * step];
        for (int j = 0; j < m; j++) w3[j] = rts[3 * j * step];
        off += 3 * m;
    }
}
static void build_twiddles_inv(const u32* rts, u32* tw, int n) {
    int off = 0;
    int end_len = (n % 3 == 2) ? (n >> 1) : n;
    for (int len = 4; len <= end_len; len <<= 2) {
        int m = len >> 2;
        int step = n / len;
        u32* w1 = tw + off;
        u32* w2 = tw + off + m;
        u32* w3 = tw + off + 2 * m;
        for (int j = 0; j < m; j++) w1[j] = rts[j * step];
        for (int j = 0; j < m; j++) w2[j] = rts[2 * j * step];
        for (int j = 0; j < m; j++) w3[j] = rts[3 * j * step];
        off += 3 * m;
    }
}

static void ntt_fwd(u32 *x, int n, const u32 *rts, const u32 *tw, u32 iroot) {
    __m256i irootv = _mm256_set1_epi32(iroot);
    int start_len;
    if (n % 3 == 2) {
        int half = n >> 1;
        for (int j = 0; j < half; j += 8) {
            __m256i u = load8(x + j);
            __m256i v = load8(x + j + half);
            __m256i w = load8(rts + j);
            store8(x + j, add_mod8(u, v));
            store8(x + j + half, mont_mul8(sub_mod8(u, v), w));
        }
        start_len = n >> 1;
    } else {
        start_len = n;
    }
    int off = 0;
    for (int len = start_len; len >= 4; len >>= 2) {
        int m = len >> 2;
        if (m >= 16) {
            const u32* w1 = tw + off;
            const u32* w2 = tw + off + m;
            const u32* w3 = tw + off + 2 * m;
            for (int i = 0; i < n; i += len) {
                u32 *y = x + i;
                for (int j = 0; j < m; j += 8) {
                    __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 a0 = add_mod8(x0, x2);
                    __m256i a1 = sub_mod8(x0, x2);
                    __m256i a2 = add_mod8(x1, x3);
                    __m256i a3 = sub_mod8(x1, x3);
                    __m256i ia3 = mont_mul8(a3, irootv);
                    __m256i tw1 = load8(w1 + j);
                    __m256i tw2 = load8(w2 + j);
                    __m256i tw3 = load8(w3 + j);
                    __m256i t01 = add_mod8(a1, ia3);
                    __m256i t03 = sub_mod8(a1, ia3);
                    store8(y + j,       add_mod8(a0, a2));
                    store8(y + j + m,   mont_mul8(t01, tw1));
                    store8(y + j + 2*m, mont_mul8(sub_mod8(a0, a2), tw2));
                    store8(y + j + 3*m, mont_mul8(t03, tw3));
                }
            }
        } else {
            const u32* w1 = tw + off;
            const u32* w2 = tw + off + m;
            const u32* w3 = tw + off + 2 * m;
            for (int i = 0; i < n; i += len) {
                u32 *y = x + i;
                for (int j = 0; j < m; j++) {
                    u32 x0 = y[j], x1 = y[j + m], x2 = y[j + 2*m], x3 = y[j + 3*m];
                    u32 a0 = add_mod(x0, x2);
                    u32 a1 = sub_mod(x0, x2);
                    u32 a2 = add_mod(x1, x3);
                    u32 a3 = sub_mod(x1, x3);
                    u32 ia3 = mont_mul(a3, iroot);
                    y[j]       = add_mod(a0, a2);
                    y[j + m]   = mont_mul(add_mod(a1, ia3), w1[j]);
                    y[j + 2*m] = mont_mul(sub_mod(a0, a2), w2[j]);
                    y[j + 3*m] = mont_mul(sub_mod(a1, ia3), w3[j]);
                }
            }
        }
        off += 3 * m;
    }
}

static void ntt_inv(u32 *x, int n, const u32 *rts, const u32 *tw, u32 iroot) {
    __m256i irootv = _mm256_set1_epi32(iroot);
    int end_len = (n % 3 == 2) ? (n >> 1) : n;
    int off = 0;
    for (int len = 4; len <= end_len; len <<= 2) {
        int m = len >> 2;
        if (m >= 16) {
            const u32* w1 = tw + off;
            const u32* w2 = tw + off + m;
            const u32* w3 = tw + off + 2 * m;
            for (int i = 0; i < n; i += len) {
                u32 *y = x + i;
                for (int j = 0; j < m; j += 8) {
                    __m256i tw1 = load8(w1 + j);
                    __m256i tw2 = load8(w2 + j);
                    __m256i tw3 = load8(w3 + j);
                    __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 b0 = add_mod8(y0, y2);
                    __m256i b1 = sub_mod8(y0, y2);
                    __m256i b2 = add_mod8(y1, y3);
                    __m256i b3 = sub_mod8(y1, y3);
                    __m256i ib3 = mont_mul8(b3, irootv);
                    store8(y + j,       add_mod8(b0, b2));
                    store8(y + j + m,   add_mod8(b1, ib3));
                    store8(y + j + 2*m, sub_mod8(b0, b2));
                    store8(y + j + 3*m, sub_mod8(b1, ib3));
                }
            }
        } else {
            const u32* w1 = tw + off;
            const u32* w2 = tw + off + m;
            const u32* w3 = tw + off + 2 * m;
            for (int i = 0; i < n; i += len) {
                u32 *y = x + i;
                for (int j = 0; j < m; j++) {
                    u32 y0 = y[j];
                    u32 y1 = mont_mul(y[j + m], w1[j]);
                    u32 y2 = mont_mul(y[j + 2*m], w2[j]);
                    u32 y3 = mont_mul(y[j + 3*m], w3[j]);
                    u32 b0 = add_mod(y0, y2);
                    u32 b1 = sub_mod(y0, y2);
                    u32 b2 = add_mod(y1, y3);
                    u32 b3 = sub_mod(y1, y3);
                    u32 ib3 = mont_mul(b3, iroot);
                    y[j]       = add_mod(b0, b2);
                    y[j + m]   = add_mod(b1, ib3);
                    y[j + 2*m] = sub_mod(b0, b2);
                    y[j + 3*m] = sub_mod(b1, ib3);
                }
            }
        }
        off += 3 * m;
    }
    if (n % 3 == 2) {
        int half = n >> 1;
        for (int j = 0; j < half; j += 8) {
            __m256i u = load8(x + j);
            __m256i v = mont_mul8(load8(x + j + half), load8(rts + j));
            store8(x + j, add_mod8(u, v));
            store8(x + j + half, sub_mod8(u, v));
        }
    }
}


void poly_multiply(unsigned *a, int n, unsigned *b, int m, unsigned *c) {
    int L = 1;
    while (L < n + m + 2) L <<= 1;

    // vectorized fill + zero pad
    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);
    gen_roots(roots, L, w);
    gen_roots(roots_inv, L, wi);

    build_twiddles_fwd(roots, TW_fwd, L);
    build_twiddles_inv(roots_inv, TW_inv, L);

    u32 iplus = to_mont(IROOT);
    u32 iminus = mont_mul(iplus, to_mont(MOD - 1));

    ntt_fwd(A, L, roots, TW_fwd, iplus);
    ntt_fwd(B, L, roots, TW_fwd, iplus);

    // vectorized pointwise
    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_inv(A, L, roots_inv, TW_inv, iminus);

    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);

    // vectorized scale to c
    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);
}

评测结果