代码

// 1004: radix-4 (2 x 4^7 x 8) DIF/DIT, hand-written asm inner loops, 3-mod NTT.
#include <sys/auxv.h>
#include <stdint.h>
#include <string.h>
#include <immintrin.h>
#pragma GCC target("avx2")
#pragma GCC optimize("O3,unroll-loops")

typedef uint64_t u64; typedef uint32_t u32; typedef __uint128_t u128;

struct DuckInfo {
  uint64_t abi_version; const char *stdin_ptr; uint64_t stdin_size;
  char *stdout_ptr; uint64_t stdout_limit; uint64_t stdout_size;
  char *stderr_ptr; uint64_t stderr_limit; uint64_t stderr_size;
  const char *IB_ptr; uint64_t IB_limit; char *OB_ptr; uint64_t OB_limit;
  uint64_t tsc_frequency;
} __attribute__((packed));

static const u32 MODS[3] = {998244353u, 1004535809u, 469762049u};
static const u32 NINVS[3] = {998244351u, 1004535807u, 469762047u};
static const u32 R2S[3]   = {932051910u, 542374313u, 460175152u};
static const u32 ZETA[3]  = {372528824u, 395918948u, 129701348u};
static const u32 ZETA3[3] = {488723995u, 823254837u, 443138433u};
static const u32 ZETA5[3] = {625715529u, 608616861u, 340060701u};
static const u32 ZETA7[3] = {509520358u, 181280972u, 26623616u};
static const u32 IROOT[3] = {911660635u, 483363861u, 450151958u};
static const u32 IMINUS[3]= {86583718u, 521171948u, 19610091u};

#define SIZE (1<<18)
static u32 A[SIZE], B[SIZE];
static u32 TW[SIZE+64], ITW[SIZE+64];
static u32 tmp[SIZE/4+8];
static u32 limbsA[111112], limbsB[111112];
static u32 outlimbs[222224];
static u32 r0[SIZE], r1[SIZE], r2[SIZE];
static char tab3[1000][3];

__m256i modvec, ninvvec, shufvec;
__m256i z1v, z3v, ivv;
static u32 curmod, curninv;

static u32 powmod(u64 a, u64 e, u32 mod) { u64 r=1,b=a%mod; while(e){ if(e&1) r=r*b%mod; b=b*b%mod; e>>=1; } return (u32)r; }
static inline u32 mont_s(u32 a, u32 b, u32 mod, u32 ninv){ u64 t=(u64)a*b; u32 m=(u32)t*ninv; u32 u=(u32)((t+(u64)m*mod)>>32); if(u>=mod)u-=mod; return u; }
static inline u32 sadd(u32 a, u32 b, u32 mod){ u32 s=a+b; if(s>=mod)s-=mod; return s; }
static inline u32 ssub(u32 a, u32 b, u32 mod){ u32 d=a-b; if(d>mod)d+=mod; return d; }

static inline __m256i mont8(__m256i x, __m256i y) {
    __m256i ao=_mm256_srli_epi64(x,32);
    __m256i bo=_mm256_srli_epi64(y,32);
    __m256i te=_mm256_mul_epu32(x,y);
    __m256i to=_mm256_mul_epu32(ao,bo);
    __m256i me=_mm256_mul_epu32(te,ninvvec);
    __m256i mo=_mm256_mul_epu32(to,ninvvec);
    __m256i ue=_mm256_srli_epi64(_mm256_add_epi64(te,_mm256_mul_epu32(me,modvec)),32);
    __m256i uo=_mm256_srli_epi64(_mm256_add_epi64(to,_mm256_mul_epu32(mo,modvec)),32);
    __m256i u=_mm256_or_si256(ue,_mm256_slli_epi64(uo,32));
    return _mm256_min_epu32(u,_mm256_sub_epi32(u,modvec));
}
static inline __m256i addm(__m256i a, __m256i b){ __m256i s=_mm256_add_epi32(a,b); return _mm256_min_epu32(s,_mm256_sub_epi32(s,modvec)); }
static inline __m256i subm(__m256i a, __m256i b){ __m256i d=_mm256_sub_epi32(a,b); return _mm256_min_epu32(d,_mm256_add_epi32(d,modvec)); }
static inline __m256i loadu(const u32* p){ return _mm256_loadu_si256((const __m256i*)p); }
static inline void storeu(u32* p, __m256i v){ _mm256_storeu_si256((__m256i*)p, v); }

static void gen_powers(u32* tmp, int s, u32 wm, u32 mod, u32 ninv, u32 r2c) {
    tmp[0]=mont_s(1,r2c,mod,ninv);
    u32 P[8];
    P[0]=wm;
    for(int i=1;i<8;i++) P[i]=mont_s(P[i-1],wm,mod,ninv);
    __m256i Pv=_mm256_loadu_si256((__m256i*)P);
    __m256i carry=_mm256_set1_epi32((int)tmp[0]);
    int p=1;
    for(;p+8<=s;p+=8){
        __m256i chunk=mont8(carry,Pv);
        _mm256_storeu_si256((__m256i*)(tmp+p),chunk);
        carry=_mm256_set1_epi32((int)tmp[p+7]);
    }
    for(;p<s;p++) tmp[p]=mont_s(tmp[p-1],wm,mod,ninv);
}

static void gen_tw4(u32* tw, u32 w, u32 mod, u32 ninv, u32 r2c) {
    gen_powers(tw, SIZE/2, mont_s(w,r2c,mod,ninv), mod, ninv, r2c);
    int off=SIZE/2;
    for(int len=SIZE/2; len>=32; len>>=2){
        int s=len>>2; int step=SIZE/len;
        u32 ws=powmod(w,step,mod); u32 wm=mont_s(ws,r2c,mod,ninv);
        gen_powers(tmp,s,wm,mod,ninv,r2c);
        for(int p=0;p<s;p+=8){
            __m256i t1=loadu(tmp+p);
            __m256i t2=mont8(t1,t1);
            __m256i t3=mont8(t2,t1);
            u32* wb=tw+off+p*3;
            storeu(wb+0,t1); storeu(wb+8,t2); storeu(wb+16,t3);
        }
        off+=3*s;
    }
}

static void gen_itw4(u32* tw, u32 iw, u32 mod, u32 ninv, u32 r2c) {
    int off=0;
    for(int len=32; len<=SIZE/2; len<<=2){
        int s=len>>2; int step=SIZE/len;
        u32 ws=powmod(iw,step,mod); u32 wm=mont_s(ws,r2c,mod,ninv);
        gen_powers(tmp,s,wm,mod,ninv,r2c);
        for(int p=0;p<s;p+=8){
            __m256i t1=loadu(tmp+p);
            __m256i t2=mont8(t1,t1);
            __m256i t3=mont8(t2,t1);
            u32* wb=tw+off+p*3;
            storeu(wb+0,t1); storeu(wb+8,t2); storeu(wb+16,t3);
        }
        off+=3*s;
    }
    gen_powers(tw+off, SIZE/2, mont_s(iw,r2c,mod,ninv), mod, ninv, r2c);
}

// Vectorized len=8 DFT (base stage): 8 groups x 8 points per iter.
static void ntt8_vec(u32* x, __m256i z1, __m256i z3, __m256i iv) {
    for(int i=0;i<SIZE;i+=64){
        __m256i a0=loadu(x+i+0), a1=loadu(x+i+8), a2=loadu(x+i+16), a3=loadu(x+i+24);
        __m256i a4=loadu(x+i+32), a5=loadu(x+i+40), a6=loadu(x+i+48), a7=loadu(x+i+56);
        __m256i b0=_mm256_unpacklo_epi32(a0,a1), b1=_mm256_unpackhi_epi32(a0,a1);
        __m256i b2=_mm256_unpacklo_epi32(a2,a3), b3=_mm256_unpackhi_epi32(a2,a3);
        __m256i b4=_mm256_unpacklo_epi32(a4,a5), b5=_mm256_unpackhi_epi32(a4,a5);
        __m256i b6=_mm256_unpacklo_epi32(a6,a7), b7=_mm256_unpackhi_epi32(a6,a7);
        __m256i c0=_mm256_unpacklo_epi64(b0,b2), c1=_mm256_unpackhi_epi64(b0,b2);
        __m256i c2=_mm256_unpacklo_epi64(b1,b3), c3=_mm256_unpackhi_epi64(b1,b3);
        __m256i c4=_mm256_unpacklo_epi64(b4,b6), c5=_mm256_unpackhi_epi64(b4,b6);
        __m256i c6=_mm256_unpacklo_epi64(b5,b7), c7=_mm256_unpackhi_epi64(b5,b7);
        __m256i t0=_mm256_permute2f128_si256(c0,c4,0x20), t1=_mm256_permute2f128_si256(c1,c5,0x20);
        __m256i t2=_mm256_permute2f128_si256(c2,c6,0x20), t3=_mm256_permute2f128_si256(c3,c7,0x20);
        __m256i t4=_mm256_permute2f128_si256(c0,c4,0x31), t5=_mm256_permute2f128_si256(c1,c5,0x31);
        __m256i t6=_mm256_permute2f128_si256(c2,c6,0x31), t7=_mm256_permute2f128_si256(c3,c7,0x31);
        __m256i d0=addm(t0,t4), d1=addm(t1,t5), d2=addm(t2,t6), d3=addm(t3,t7);
        __m256i e0=subm(t0,t4);
        __m256i e1=mont8(subm(t1,t5),z1v);
        __m256i e2=mont8(subm(t2,t6),ivv);
        __m256i e3=mont8(subm(t3,t7),z3v);
        __m256i s02=addm(d0,d2), t02=subm(d0,d2);
        __m256i s13=addm(d1,d3), t13=mont8(subm(d1,d3),ivv);
        __m256i D0=addm(s02,s13), D2=subm(s02,s13);
        __m256i D1=addm(t02,t13), D3=subm(t02,t13);
        __m256i u02=addm(e0,e2), v02=subm(e0,e2);
        __m256i u13=addm(e1,e3), v13=mont8(subm(e1,e3),ivv);
        __m256i E0=addm(u02,u13), E2=subm(u02,u13);
        __m256i E1=addm(v02,v13), E3=subm(v02,v13);
        __m256i rb0=_mm256_unpacklo_epi32(D0,E0), rb1=_mm256_unpackhi_epi32(D0,E0);
        __m256i rb2=_mm256_unpacklo_epi32(D1,E1), rb3=_mm256_unpackhi_epi32(D1,E1);
        __m256i rb4=_mm256_unpacklo_epi32(D2,E2), rb5=_mm256_unpackhi_epi32(D2,E2);
        __m256i rb6=_mm256_unpacklo_epi32(D3,E3), rb7=_mm256_unpackhi_epi32(D3,E3);
        __m256i rc0=_mm256_unpacklo_epi64(rb0,rb2), rc1=_mm256_unpackhi_epi64(rb0,rb2);
        __m256i rc2=_mm256_unpacklo_epi64(rb1,rb3), rc3=_mm256_unpackhi_epi64(rb1,rb3);
        __m256i rc4=_mm256_unpacklo_epi64(rb4,rb6), rc5=_mm256_unpackhi_epi64(rb4,rb6);
        __m256i rc6=_mm256_unpacklo_epi64(rb5,rb7), rc7=_mm256_unpackhi_epi64(rb5,rb7);
        __m256i r0=_mm256_permute2f128_si256(rc0,rc4,0x20), r1=_mm256_permute2f128_si256(rc1,rc5,0x20);
        __m256i r2=_mm256_permute2f128_si256(rc2,rc6,0x20), r3=_mm256_permute2f128_si256(rc3,rc7,0x20);
        __m256i r4=_mm256_permute2f128_si256(rc0,rc4,0x31), r5=_mm256_permute2f128_si256(rc1,rc5,0x31);
        __m256i r6=_mm256_permute2f128_si256(rc2,rc6,0x31), r7=_mm256_permute2f128_si256(rc3,rc7,0x31);
        storeu(x+i+0,r0); storeu(x+i+8,r1); storeu(x+i+16,r2); storeu(x+i+24,r3);
        storeu(x+i+32,r4); storeu(x+i+40,r5); storeu(x+i+48,r6); storeu(x+i+56,r7);
    }
}

// ---- Hand-written asm inner loops ----
u32* g_y4; int g_s4; const u32* g_wb4; int g_cnt4;

__attribute__((noinline)) static void dif4_inner(void) {
  __asm__ volatile(
    ".intel_syntax noprefix\n"
    ".macro MONT8 dst, x, w, t1, t2\n"
    "  vpsrlq \\t1, \\x, 32\n"
    "  vpsrlq \\t2, \\w, 32\n"
    "  vpmuludq \\dst, \\x, \\w\n"
    "  vpmuludq \\t1, \\t1, \\t2\n"
    "  vpmuludq \\t2, \\dst, ymm9\n"
    "  vpmuludq \\t2, \\t2, ymm8\n"
    "  vpaddq \\dst, \\dst, \\t2\n"
    "  vpsrlq \\dst, \\dst, 32\n"
    "  vpmuludq \\t2, \\t1, ymm9\n"
    "  vpmuludq \\t2, \\t2, ymm8\n"
    "  vpaddq \\t1, \\t1, \\t2\n"
    "  vpsrlq \\t1, \\t1, 32\n"
    "  vpsllq \\t1, \\t1, 32\n"
    "  vpor \\dst, \\dst, \\t1\n"
    "  vpsubd \\t1, \\dst, ymm8\n"
    "  vpminud \\dst, \\dst, \\t1\n"
    ".endm\n"
    ".macro ADSUB sreg, dreg, t1, t2\n"
    "  vpsubd \\t1, \\sreg, \\dreg\n"
    "  vpaddd \\sreg, \\sreg, \\dreg\n"
    "  vpsubd \\t2, \\sreg, ymm8\n"
    "  vpminud \\sreg, \\sreg, \\t2\n"
    "  vpaddd \\t2, \\t1, ymm8\n"
    "  vpminud \\dreg, \\t1, \\t2\n"
    ".endm\n"
    "mov rdi, g_y4[rip]\n"
    "movsxd rsi, dword ptr g_s4[rip]\n"
    "mov r13, g_wb4[rip]\n"
    "movsxd r14, dword ptr g_cnt4[rip]\n"
    "vmovdqa ymm8, modvec[rip]\n"
    "vmovdqa ymm9, ninvvec[rip]\n"
    "vmovdqa ymm10, ivv[rip]\n"
    "lea rax, [rdi + rsi*4]\n"
    "lea rbx, [rdi + rsi*8]\n"
    "lea r8, [rdi + rsi*8]\n"
    "lea r8, [r8 + rsi*4]\n"
    ".p2align 4\n"
    "1:\n"
    "vmovdqu ymm0, [rdi]\n"
    "vmovdqu ymm1, [rax]\n"
    "vmovdqu ymm2, [rbx]\n"
    "vmovdqu ymm3, [r8]\n"
    "vmovdqu ymm4, [r13]\n"
    "vmovdqu ymm5, [r13+32]\n"
    "vmovdqu ymm6, [r13+64]\n"
    "ADSUB ymm0, ymm2, ymm13, ymm14\n"         // t0=ymm0, t1=ymm2
    "ADSUB ymm1, ymm3, ymm13, ymm14\n"         // t2=ymm1, d13=ymm3
    "MONT8 ymm7, ymm3, ymm10, ymm11, ymm12\n"  // t3=ymm7 = d13*i  (temps 11,12)
    "ADSUB ymm0, ymm1, ymm13, ymm14\n"         // D0=ymm0, D2=ymm1  (temps 13,14)
    "vmovdqu [rdi], ymm0\n"                    // store D0
    "MONT8 ymm5, ymm1, ymm5, ymm15, ymm3\n"    // Y2 = D2*w2  (dst ymm5, temps 15,3) [overlaps i-mul]
    "ADSUB ymm2, ymm7, ymm13, ymm14\n"         // D1=ymm2, D3=ymm7  (temps 13,14)
    "MONT8 ymm4, ymm2, ymm4, ymm11, ymm12\n"   // Y1 = D1*w1  (dst ymm4, temps 11,12) [reuse i-mul temps]
    "MONT8 ymm6, ymm7, ymm6, ymm13, ymm14\n"   // Y3 = D3*w3  (dst ymm6, temps 13,14) [reuse ADSUB temps]
    "vmovdqu [rax], ymm4\n"
    "vmovdqu [rbx], ymm5\n"
    "vmovdqu [r8], ymm6\n"
    "add rdi, 32\n"
    "add rax, 32\n"
    "add rbx, 32\n"
    "add r8, 32\n"
    "add r13, 96\n"
    "dec r14\n"
    "jnz 1b\n"
    ".att_syntax prefix\n"
    :
    :
    : "memory", "rax","rbx","rcx","rdx","rsi","rdi","r8","r9","r10","r11","r12","r13","r14","r15","rbp",
      "ymm0","ymm1","ymm2","ymm3","ymm4","ymm5","ymm6","ymm7","ymm8","ymm9","ymm10","ymm11","ymm12","ymm13","ymm14","ymm15"
  );
}

__attribute__((noinline)) static void dit4_inner(void) {
  __asm__ volatile(
    ".intel_syntax noprefix\n"
    "mov rdi, g_y4[rip]\n"
    "movsxd rsi, dword ptr g_s4[rip]\n"
    "mov r13, g_wb4[rip]\n"
    "movsxd r14, dword ptr g_cnt4[rip]\n"
    "vmovdqa ymm8, modvec[rip]\n"
    "vmovdqa ymm9, ninvvec[rip]\n"
    "vmovdqa ymm10, ivv[rip]\n"
    "lea rax, [rdi + rsi*4]\n"
    "lea rbx, [rdi + rsi*8]\n"
    "lea r8, [rdi + rsi*8]\n"
    "lea r8, [r8 + rsi*4]\n"
    ".p2align 4\n"
    "1:\n"
    "vmovdqu ymm0, [rdi]\n"
    "vmovdqu ymm1, [rax]\n"
    "vmovdqu ymm2, [rbx]\n"
    "vmovdqu ymm3, [r8]\n"
    "vmovdqu ymm4, [r13]\n"
    "vmovdqu ymm5, [r13+32]\n"
    "vmovdqu ymm6, [r13+64]\n"
    "MONT8 ymm1, ymm1, ymm4, ymm11, ymm12\n"   // x1 = X1*w1 (temps 11,12)
    "MONT8 ymm2, ymm2, ymm5, ymm13, ymm14\n"   // x2 = X2*w2 (temps 13,14)
    "MONT8 ymm3, ymm3, ymm6, ymm15, ymm7\n"    // x3 = X3*w3 (temps 15,7)
    "ADSUB ymm0, ymm2, ymm11, ymm12\n"          // t0=ymm0, t1=ymm2
    "ADSUB ymm1, ymm3, ymm13, ymm14\n"          // t2=ymm1, d13=ymm3
    "MONT8 ymm7, ymm3, ymm10, ymm15, ymm14\n"   // t3=ymm7 = d13*i (temps 15,14)
    "ADSUB ymm0, ymm1, ymm11, ymm12\n"          // o0=ymm0, o2=ymm1
    "ADSUB ymm2, ymm7, ymm13, ymm14\n"          // o1=ymm2, o3=ymm7
    "vmovdqu [rdi], ymm0\n"
    "vmovdqu [rax], ymm2\n"
    "vmovdqu [rbx], ymm1\n"
    "vmovdqu [r8], ymm7\n"
    "add rdi, 32\n"
    "add rax, 32\n"
    "add rbx, 32\n"
    "add r8, 32\n"
    "add r13, 96\n"
    "dec r14\n"
    "jnz 1b\n"
    ".att_syntax prefix\n"
    :
    :
    : "memory", "rax","rbx","rcx","rdx","rsi","rdi","r8","r9","r10","r11","r12","r13","r14","r15","rbp",
      "ymm0","ymm1","ymm2","ymm3","ymm4","ymm5","ymm6","ymm7","ymm8","ymm9","ymm10","ymm11","ymm12","ymm13","ymm14","ymm15"
  );
}

static void ntt_fwd(u32* x) {
    // radix-2 DIF (intrinsics)
    {
        int s=SIZE>>1;
        for(int p=0;p<s;p+=8){
            __m256i av=loadu(x+p);
            __m256i bv=loadu(x+p+s);
            __m256i wv=loadu(TW+p);
            storeu(x+p,addm(av,bv));
            storeu(x+p+s,mont8(subm(av,bv),wv));
        }
    }
    int off=SIZE>>1;
    for(int len=SIZE>>1; len>=32; len>>=2){
        int s=len>>2;
        for(int i=0;i<SIZE;i+=len){
            u32* y=x+i;
            g_y4=y; g_s4=s; g_wb4=TW+off; g_cnt4=s>>3;
            dif4_inner();
        }
        off+=3*s;
    }
    ntt8_vec(x, z1v, z3v, ivv);
}

static void ntt_inv(u32* x) {
    ntt8_vec(x, z1v, z3v, ivv);
    int off=0;
    for(int len=32; len<=SIZE/2; len<<=2){
        int s=len>>2;
        for(int i=0;i<SIZE;i+=len){
            u32* y=x+i;
            g_y4=y; g_s4=s; g_wb4=ITW+off; g_cnt4=s>>3;
            dit4_inner();
        }
        off+=3*s;
    }
    {
        int s=SIZE>>1;
        const u32* wb=ITW+off;
        for(int p=0;p<s;p+=8){
            __m256i av=loadu(x+p);
            __m256i bv=mont8(loadu(x+p+s),loadu(wb+p));
            storeu(x+p,addm(av,bv));
            storeu(x+p+s,subm(av,bv));
        }
    }
}

static void build_tab3(void){ for(int i=0;i<1000;i++){ int v=i; tab3[i][2]='0'+v%10; v/=10; tab3[i][1]='0'+v%10; v/=10; tab3[i][0]='0'+v%10; } }

#ifdef LOCAL_TEST
extern uintptr_t jd_getauxval(uintptr_t);
int jd_main(){
    DuckInfo* di=(DuckInfo*)jd_getauxval(0x6b637564ull);
#else
int main(){
    DuckInfo* di=(DuckInfo*)getauxval(0x6b637564ull);
#endif
    const char* in=di->stdin_ptr; u64 inlen=di->stdin_size;
    const char* p=in; const char* inend=in+inlen;
    while(p<inend&&(*p=='\n'||*p=='\r'||*p==' '||*p=='\t'))p++;
    const char* a_start=p; while(p<inend&&*p>='0'&&*p<='9')p++; const char* a_end=p;
    while(p<inend&&(*p=='\n'||*p=='\r'||*p==' '||*p=='\t'))p++;
    const char* b_start=p; while(p<inend&&*p>='0'&&*p<='9')p++; const char* b_end=p;
    const char* sa=a_start; while(sa<a_end-1&&*sa=='0')sa++;
    const char* sb=b_start; while(sb<b_end-1&&*sb=='0')sb++;
    int na=0,nb=0;
    {
        const char* pos=a_end;
        while(pos-9>=sa){const char* q=pos-9;
            u32 v=(u32)(q[0]-'0'); v=v*10+(u32)(q[1]-'0'); v=v*10+(u32)(q[2]-'0'); v=v*10+(u32)(q[3]-'0');
            v=v*10+(u32)(q[4]-'0'); v=v*10+(u32)(q[5]-'0'); v=v*10+(u32)(q[6]-'0'); v=v*10+(u32)(q[7]-'0'); v=v*10+(u32)(q[8]-'0');
            limbsA[na++]=v; pos=q;
        }
        if(pos>sa){u32 v=0; for(const char* q=sa;q<pos;q++) v=v*10+(u32)(*q-'0'); limbsA[na++]=v;}
    }
    {
        const char* pos=b_end;
        while(pos-9>=sb){const char* q=pos-9;
            u32 v=(u32)(q[0]-'0'); v=v*10+(u32)(q[1]-'0'); v=v*10+(u32)(q[2]-'0'); v=v*10+(u32)(q[3]-'0');
            v=v*10+(u32)(q[4]-'0'); v=v*10+(u32)(q[5]-'0'); v=v*10+(u32)(q[6]-'0'); v=v*10+(u32)(q[7]-'0'); v=v*10+(u32)(q[8]-'0');
            limbsB[nb++]=v; pos=q;
        }
        if(pos>sb){u32 v=0; for(const char* q=sb;q<pos;q++) v=v*10+(u32)(*q-'0'); limbsB[nb++]=v;}
    }

    shufvec=_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);

    for(int mi=0;mi<3;mi++){
        u32 mod=MODS[mi],ninv=NINVS[mi],r2c=R2S[mi];
        curmod=mod; curninv=ninv;
        modvec=_mm256_set1_epi32((int)mod); ninvvec=_mm256_set1_epi32((int)ninv);
        u32 w=powmod(3u,(mod-1)/SIZE,mod);
        u32 iw=powmod(w,mod-2,mod);
        u32 Iv=mont_s(IROOT[mi],r2c,mod,ninv);
        u32 Ivi=mont_s(IMINUS[mi],r2c,mod,ninv);
        u32 Z1=mont_s(ZETA[mi],r2c,mod,ninv), Z3=mont_s(ZETA3[mi],r2c,mod,ninv);
        u32 Z1i=mont_s(ZETA7[mi],r2c,mod,ninv), Z3i=mont_s(ZETA5[mi],r2c,mod,ninv);
        gen_tw4(TW,w,mod,ninv,r2c);
        gen_itw4(ITW,iw,mod,ninv,r2c);

        __m256i r2b=_mm256_set1_epi32((int)r2c);
        {int vlim=(na+7)>>3;for(int v=0;v<vlim;v++){__m256i x=loadu(limbsA+v*8);x=_mm256_min_epu32(x,_mm256_sub_epi32(x,modvec));x=_mm256_min_epu32(x,_mm256_sub_epi32(x,modvec));storeu(A+v*8,mont8(x,r2b));}memset(A+vlim*8,0,(SIZE-vlim*8)*4);}
        {int vlim=(nb+7)>>3;for(int v=0;v<vlim;v++){__m256i x=loadu(limbsB+v*8);x=_mm256_min_epu32(x,_mm256_sub_epi32(x,modvec));x=_mm256_min_epu32(x,_mm256_sub_epi32(x,modvec));storeu(B+v*8,mont8(x,r2b));}memset(B+vlim*8,0,(SIZE-vlim*8)*4);}

        z1v=_mm256_set1_epi32((int)Z1); z3v=_mm256_set1_epi32((int)Z3); ivv=_mm256_set1_epi32((int)Iv);
        ntt_fwd(A);
        ntt_fwd(B);
        for(int v=0;v<SIZE;v+=8) storeu(A+v,mont8(loadu(A+v),loadu(B+v)));
        z1v=_mm256_set1_epi32((int)Z1i); z3v=_mm256_set1_epi32((int)Z3i); ivv=_mm256_set1_epi32((int)Ivi);
        ntt_inv(A);
        u32 ninvn=powmod(SIZE,mod-2,mod);
        __m256i ninvb=_mm256_set1_epi32((int)ninvn);
        u32* dst=(mi==0)?r0:(mi==1)?r1:r2;
        for(int v=0;v<SIZE;v+=8) storeu(dst+v,mont8(loadu(A+v),ninvb));
    }

    {
        u32 m0=MODS[0],m1=MODS[1],m2=MODS[2];
        u64 M0=m0,M1=m1;
        u64 INV01=powmod(M0%m1,m1-2,m1);
        u128 M01_128=(u128)M0*M1; u64 M01=(u64)M01_128;
        u64 INV012=powmod(M01%m2,m2-2,m2);
        u64 M0m2=M0%m2;
        u64 M01_q=M01/1000000000u; u64 M01_r=M01%1000000000u;
        int outlen=na+nb-1; u64 carry=0; u64 m2x2=2ull*m2;
        for(int i=0;i<outlen;i++){
            u64 a0=r0[i],a1=r1[i],a2=r2[i];
            u64 t1=a1-a0+m1; if(t1>=m1) t1-=m1;
            t1=t1*INV01%m1;
            u64 a0m2=a0; if(a0m2>=m2x2) a0m2-=m2x2; if(a0m2>=m2) a0m2-=m2;
            u64 x1m2=(a0m2+t1*M0m2)%m2;
            u64 t2=a2-x1m2+m2; if(t2>=m2) t2-=m2;
            t2=t2*INV012%m2;
            u64 d=a0+t1*M0+carry;
            u64 s=d+t2*M01_r;
            u64 carr=t2*M01_q+s/1000000000u;
            outlimbs[i]=(u32)(s%1000000000u); carry=carr;
        }
        while(carry){ outlimbs[outlen++]=(u32)(carry%1000000000u); carry/=1000000000u; }
        int hi=outlen-1; while(hi>0&&outlimbs[hi]==0) hi--;
        build_tab3();
        char* out=di->stdout_ptr; char* o=out;
        {u32 v=outlimbs[hi]; char tmp[10]; int t=0; do{ tmp[t++]='0'+(char)(v%10); v/=10; }while(v); while(t>0)*o++=tmp[--t]; }
        for(int i=hi-1;i>=0;i--){ u32 v=outlimbs[i]; u32 g2=v/1000000u; u32 g1=(v/1000u)%1000u; u32 g0=v%1000u; const char* q;
            q=tab3[g2]; *o++=q[0];*o++=q[1];*o++=q[2];
            q=tab3[g1]; *o++=q[0];*o++=q[1];*o++=q[2];
            q=tab3[g0]; *o++=q[0];*o++=q[1];*o++=q[2]; }
        *o++='\n'; di->stdout_size=(u64)(o-out);
    }
#ifdef LOCAL_TEST
    return 0;
#else
    asm volatile("mov $60, %%eax; xor %%edi, %%edi; syscall" ::: "rax","rdi","rcx","r11","memory");
    __builtin_unreachable();
#endif
}

评测结果