代码

// Duck.ac 1002: exact convolution for coefficients 0..9, degrees <= 1,000,000.
// AVX2 / GCC 9.3, adapted from Qwerty1232: https://duck.ac/submission/28087
// Cleaned from https://duck.ac/submission/48181 (20.505208 ms).
// p = 39 * 2^21 + 1 > 81 * 1,000,001, so one modulus gives exact integers.
// The fast path fuses radix-4 NTT stages in 512-element blocks and reuses c.
#include <immintrin.h>

#include <algorithm>
#include <array>
#include <cassert>
#include <cstdint>
#include <cstring>
#include <vector>

#pragma GCC target("avx2,bmi")

using u32 = uint32_t;
using u64 = uint64_t;

struct Montgomery {
    u32 mod;   // mod
    u32 mod2;  // 2 * mod
    u32 n_inv; // n_inv * mod == -1 (mod 2^32)
    u32 r;     // 2^32 % mod
    u32 r2;    // (2^32)^2 % mod

    Montgomery() = default;
    Montgomery(u32 mod) : mod(mod) {
        assert(mod % 2 == 1);
        assert(mod < (1 << 30));
        mod2 = 2 * mod;
        n_inv = 1;
        for (int i = 0; i < 5; i++) {
            n_inv *= 2 + n_inv * mod;
        }
        r = (u64(1) << 32) % mod;
        r2 = u64(r) * r % mod;
    }

    u32 shrink(u32 val) const { return std::min(val, val - mod); }
    u32 shrink2(u32 val) const { return std::min(val, val - mod2); }

    template <bool strict = true> u32 reduce(u64 val) const {
        u32 res = (val + u32(val) * n_inv * u64(mod)) >> 32;
        if (strict) res = shrink(res);
        return res;
    }

    template <bool strict = true> u32 mul(u32 a, u32 b) const { return reduce<strict>(u64(a) * b); }

    template <bool input_in_space = false, bool output_in_space = false> u32 power(u32 b, u32 e) const {
        if (!input_in_space) b = mul<false>(b, r2);
        u32 r = output_in_space ? this->r : 1;
        for (; e > 0; e >>= 1) {
            if (e & 1) r = mul<false>(r, b);
            b = mul<false>(b, b);
        }
        return shrink(r);
    }
};

using i256 = __m256i;
using u32x8 = u32 __attribute__((vector_size(32)));
using u64x4 = u64 __attribute__((vector_size(32)));

u32x8 load_u32x8(const u32 *ptr) {
    return (u32x8)_mm256_load_si256((const i256 *)ptr);
}
void store_u32x8(u32 *ptr, u32x8 vec) {
    _mm256_store_si256((i256 *)ptr, (i256)vec);
}

struct MontgomeryAVX2 {
    static constexpr u32x8 mod = {81788929, 81788929, 81788929, 81788929, 81788929, 81788929, 81788929, 81788929};
    static constexpr u32x8 mod2 = {163577858, 163577858, 163577858, 163577858,
                                   163577858, 163577858, 163577858, 163577858};
    static constexpr u32x8 n_inv = {81788927, 81788927, 81788927, 81788927, 81788927, 81788927, 81788927, 81788927};
    static constexpr u32x8 r = {41942988, 41942988, 41942988, 41942988, 41942988, 41942988, 41942988, 41942988};
    static constexpr u32x8 r2 = {56088131, 56088131, 56088131, 56088131, 56088131, 56088131, 56088131, 56088131};
    MontgomeryAVX2() = default;
    explicit MontgomeryAVX2(u32 p) { assert(p == 81788929); }

    u32x8 shrink(u32x8 vec) const { return (u32x8)_mm256_min_epu32((i256)vec, _mm256_sub_epi32((i256)vec, (i256)mod)); }
    template <int Low = 0> u32x8 canonical_wide(u32x8 v) const {
        for (int shift = 5; shift >= Low; shift--) {
            u32x8 p = mod << shift;
            v = (u32x8)_mm256_min_epu32((i256)v, (i256)(v - p));
        }
        return v;
    }
    u32x8 shrink2(u32x8 vec) const {
        return (u32x8)_mm256_min_epu32((i256)vec, _mm256_sub_epi32((i256)vec, (i256)mod2));
    }
    u32x8 shrink2_n(u32x8 vec) const {
        return (u32x8)_mm256_min_epu32((i256)vec, _mm256_add_epi32((i256)vec, (i256)mod2));
    }

    template <bool strict = true> u32x8 reduce(u64x4 x0246, u64x4 x1357) const {
        u64x4 x0246_ninv = (u64x4)_mm256_mul_epu32((i256)x0246, (i256)n_inv);
        u64x4 x1357_ninv = (u64x4)_mm256_mul_epu32((i256)x1357, (i256)n_inv);
        u64x4 x0246_res = (u64x4)_mm256_add_epi64((i256)x0246, _mm256_mul_epu32((i256)x0246_ninv, (i256)mod));
        u64x4 x1357_res = (u64x4)_mm256_add_epi64((i256)x1357, _mm256_mul_epu32((i256)x1357_ninv, (i256)mod));
        u32x8 res = (u32x8)_mm256_or_si256(_mm256_bsrli_epi128((i256)x0246_res, 4), (i256)x1357_res);
        if (strict) res = shrink(res);
        return res;
    }

    template <bool strict = true, bool b_use_only_even = false> u32x8 mul_u32x8(u32x8 a, u32x8 b) const {
        u32x8 a_sh = (u32x8)_mm256_bsrli_epi128((i256)a, 4);
        u32x8 b_sh = b_use_only_even ? b : (u32x8)_mm256_bsrli_epi128((i256)b, 4);
        u64x4 x0246 = (u64x4)_mm256_mul_epu32((i256)a, (i256)b);
        u64x4 x1357 = (u64x4)_mm256_mul_epu32((i256)a_sh, (i256)b_sh);
        return reduce<strict>(x0246, x1357);
    }

    template <bool strict = true> u64x4 mul_u64x4(u64x4 a, u64x4 b) const {
        u64x4 pr = (u64x4)_mm256_mul_epu32((i256)a, (i256)b);
        u64x4 pr2 = (u64x4)_mm256_mul_epu32(_mm256_mul_epu32((i256)pr, (i256)n_inv), (i256)mod);
        u64x4 res = (u64x4)_mm256_bsrli_epi128(_mm256_add_epi64((i256)pr, (i256)pr2), 4);
        if (strict) res = (u64x4)shrink((u32x8)res);
        return res;
    }
};

// Assemble read-only roots at build time. C++ constexpr expansion exceeds the
// judge compiler's memory limit. n1/n2/n3 are w1/w2/w3 * n_inv modulo 2^32.
namespace fixed_roots {
struct Twiddle {
    u32 w1, w2, w3, n1, n2, n3;
};
struct Table {
    Twiddle data[65536];
};
extern const Table forward asm("poly_roots_forward");
extern const Table inverse asm("poly_roots_inverse");
struct DotTable {
    u32 data[32768][4];
};
extern const DotTable dot asm("poly_roots_dot");
} // namespace fixed_roots
asm(R"asm(
.pushsection .rodata
// Select by the number of trailing one bits in the table index.
.macro next_factor dest, mask, value, rest:vararg
.if ((_i & \mask) == 0)
.set \dest,\value
.else
next_factor \dest,(\mask*2),\rest
.endif
.endm
.p2align 6
.globl poly_roots_forward
.type poly_roots_forward,@object
poly_roots_forward:
.set _i,0
.set _w1,41942988
.set _w2,41942988
.rept 65536
.set _w12,(((_w1*_w2)%81788929)*1557504)%81788929
.long _w1,_w2,_w12,((_w1*81788927)&0xffffffff),((_w2*81788927)&0xffffffff),((_w12*81788927)&0xffffffff)
next_factor _fac, 1, 1977387,49739338,76551861,57685215,25318722,22305379,75160758,77449485,49050524,58847824,69356575,69052175,45043381,68811137,48691376,28111944,26577652
next_factor _fac2, 1, 57807995,1883838,27152551,62819432,22367481,25489457,54748607,18371892,60074596,43336831,16579980,78708963,26101542,51041304,60500196,40232015,28323882
.set _w1,(_w1*_fac2)%81788929
.set _w2,(_w2*_fac)%81788929
.set _i,_i+1
.endr
.size poly_roots_forward,.-poly_roots_forward
.p2align 6
.globl poly_roots_inverse
.type poly_roots_inverse,@object
poly_roots_inverse:
.set _i,0
.set _w1,41942988
.set _w2,41942988
.rept 65536
.set _w12,(((_w1*_w2)%81788929)*1557504)%81788929
.long _w1,_w2,_w12,((_w1*81788927)&0xffffffff),((_w2*81788927)&0xffffffff),((_w12*81788927)&0xffffffff)
next_factor _fac, 1, 1883838,34192649,65864533,47472559,32202660,46455854,18299665,51166265,46148164,40005067,42538512,22507185,19881487,13191717,67317322,1064838,16759432
next_factor _fac2, 1, 23980934,1977387,58967103,56026958,74557765,58488554,3169619,20142414,28119438,26733415,74787290,67900511,63391377,74641937,67976842,40043517,2457972
.set _w1,(_w1*_fac2)%81788929
.set _w2,(_w2*_fac)%81788929
.set _i,_i+1
.endr
.size poly_roots_inverse,.-poly_roots_inverse
.p2align 6
.globl poly_roots_dot
.type poly_roots_dot,@object
poly_roots_dot:
.set _i,0
.set _d0,41942988
.set _d1,42958308
.set _d2,28282409
.set _d3,36011086
.rept 32768
.long _d0,_d1,_d2,_d3
next_factor _fac, 1, 34192649,76852948,45870503,27153147,50722843,53125215,43544278,13378268,50576854,50366248,18491959,52344447,58962465,12062499,19859762,9337084
.set _d0,(_d0*_fac)%81788929
.set _d1,(_d1*_fac)%81788929
.set _d2,(_d2*_fac)%81788929
.set _d3,(_d3*_fac)%81788929
.set _i,_i+1
.endr
.size poly_roots_dot,.-poly_roots_dot
.purgem next_factor
.popsection
)asm");

class NTT {
    // Global coefficient offset of the quarter currently in local scra

评测结果