代码

// 1002: poly_multiply(a,n,b,m,c) with n=m=1e6, coefficients < 10.
// AVX2 double-precision complex FFT: pack a+i*b, transform, unpack, multiply, inverse.
#include <cmath>
#include <cstring>
#include <immintrin.h>

typedef unsigned u32;
typedef unsigned long long u64;

static const int MAXN = 1 << 21;
static double WR[MAXN], WI[MAXN];
static double IR[MAXN], II[MAXN];
static double ZR[MAXN], ZI[MAXN];
static double CR[MAXN], CI[MAXN];
static int REV[MAXN], PART[MAXN];

static void build_rev(int n, int bits) {
    REV[0] = 0;
    for (int i = 1; i < n; i++) REV[i] = (REV[i >> 1] >> 1) | ((i & 1) << (bits - 1));
    for (int i = 0; i < n; i++) {
        int k = REV[i];
        PART[i] = REV[(n - k) & (n - 1)];
    }
}

static void build_tables(int n) {
    const double PI = 3.14159265358979323846;
    double *wr = WR, *wi = WI;
    for (int len = n; len >= 2; len >>= 1) {
        int half = len >> 1;
        double ang = -2.0 * PI / len;
        double c = cos(ang), s = sin(ang), cr = 1.0, ci = 0.0;
        for (int j = 0; j < half; j++) {
            wr[j] = cr; wi[j] = ci;
            double nr = cr * c - ci * s, ni = cr * s + ci * c;
            cr = nr; ci = ni;
        }
        wr += half; wi += half;
    }
    double *vr = IR, *vi = II;
    for (int len = 2; len <= n; len <<= 1) {
        int half = len >> 1;
        double ang = 2.0 * PI / len;
        double c = cos(ang), s = sin(ang), cr = 1.0, ci = 0.0;
        for (int j = 0; j < half; j++) {
            vr[j] = cr; vi[j] = ci;
            double nr = cr * c - ci * s, ni = cr * s + ci * c;
            cr = nr; ci = ni;
        }
        vr += half; vi += half;
    }
}

__attribute__((target("avx2,fma")))
static void fft_dif(int n) {
    const double *wr = WR, *wi = WI;
    for (int len = n; len >= 2; len >>= 1) {
        int half = len >> 1;
        for (int i = 0; i < n; i += len) {
            double *ar = ZR + i, *ai = ZI + i;
            double *br = ar + half, *bi = ai + half;
            int j = 0;
            for (; j + 4 <= half; j += 4) {
                __m256d ur = _mm256_loadu_pd(ar + j), ui = _mm256_loadu_pd(ai + j);
                __m256d vr = _mm256_loadu_pd(br + j), vi = _mm256_loadu_pd(bi + j);
                __m256d sr = _mm256_add_pd(ur, vr), si = _mm256_add_pd(ui, vi);
                __m256d dr = _mm256_sub_pd(ur, vr), di = _mm256_sub_pd(ui, vi);
                __m256d twr = _mm256_loadu_pd(wr + j), twi = _mm256_loadu_pd(wi + j);
                __m256d nr = _mm256_fmsub_pd(dr, twr, _mm256_mul_pd(di, twi));
                __m256d ni = _mm256_fmadd_pd(dr, twi, _mm256_mul_pd(di, twr));
                _mm256_storeu_pd(ar + j, sr); _mm256_storeu_pd(ai + j, si);
                _mm256_storeu_pd(br + j, nr); _mm256_storeu_pd(bi + j, ni);
            }
            for (; j < half; j++) {
                double ur = ar[j], ui = ai[j], vr = br[j], vi = bi[j];
                ar[j] = ur + vr; ai[j] = ui + vi;
                double dr = ur - vr, di = ui - vi;
                br[j] = dr * wr[j] - di * wi[j];
                bi[j] = dr * wi[j] + di * wr[j];
            }
        }
        wr += half; wi += half;
    }
}

__attribute__((target("avx2,fma")))
static void fft_dit_inv(int n) {
    const double *wr = IR, *wi = II;
    for (int len = 2; len <= n; len <<= 1) {
        int half = len >> 1;
        for (int i = 0; i < n; i += len) {
            double *ar = CR + i, *ai = CI + i;
            double *br = ar + half, *bi = ai + half;
            int j = 0;
            for (; j + 4 <= half; j += 4) {
                __m256d ur = _mm256_loadu_pd(ar + j), ui = _mm256_loadu_pd(ai + j);
                __m256d vr = _mm256_loadu_pd(br + j), vi = _mm256_loadu_pd(bi + j);
                __m256d twr = _mm256_loadu_pd(wr + j), twi = _mm256_loadu_pd(wi + j);
                __m256d nr = _mm256_fmsub_pd(vr, twr, _mm256_mul_pd(vi, twi));
                __m256d ni = _mm256_fmadd_pd(vr, twi, _mm256_mul_pd(vi, twr));
                _mm256_storeu_pd(ar + j, _mm256_add_pd(ur, nr));
                _mm256_storeu_pd(ai + j, _mm256_add_pd(ui, ni));
                _mm256_storeu_pd(br + j, _mm256_sub_pd(ur, nr));
                _mm256_storeu_pd(bi + j, _mm256_sub_pd(ui, ni));
            }
            for (; j < half; j++) {
                double ur = ar[j], ui = ai[j];
                double vr = br[j], vi = bi[j];
                double nr = vr * wr[j] - vi * wi[j];
                double ni = vr * wi[j] + vi * wr[j];
                ar[j] = ur + nr; ai[j] = ui + ni;
                br[j] = ur - nr; bi[j] = ui - ni;
            }
        }
        wr += half; wi += half;
    }
}

extern "C" void poly_multiply(unsigned *a, int n, unsigned *b, int m, unsigned *c) {
    int N = 1, bits = 0;
    while (N < n + m + 1) { N <<= 1; bits++; }
    build_tables(N);
    build_rev(N, bits);
    for (int i = 0; i < N; i++) { ZR[i] = 0.0; ZI[i] = 0.0; }
    for (int i = 0; i <= n; i++) ZR[i] = (double)a[i];
    for (int i = 0; i <= m; i++) ZI[i] = (double)b[i];
    fft_dif(N);
    for (int i = 0; i < N; i++) {
        int q = PART[i];
        double zr = ZR[i], zi = ZI[i];
        double pr = ZR[q], pi = -ZI[q];
        double ar = (zr + pr) * 0.5, ai = (zi + pi) * 0.5;
        double br = (zi - pi) * 0.5, bi = -(zr - pr) * 0.5;
        CR[i] = ar * br - ai * bi;
        CI[i] = ar * bi + ai * br;
    }
    fft_dit_inv(N);
    double inv = 1.0 / N;
    for (int i = 0; i <= n + m; i++) {
        double v = CR[i] * inv;
        c[i] = (unsigned)(long long)(v + 0.5);
    }
}

评测结果