#pragma GCC target("arch=skylake")
#pragma GCC optimize("O3")
#define CENTERED_INPUT
#include <immintrin.h>
#include <math.h>
#include <stdint.h>
#include <stdlib.h>
#include <string.h>
#ifdef HIGH_INV
#include <fftw3.h>
#ifndef ROUND_BIAS
#define ROUND_BIAS 0.0
#endif
#endif
#ifdef LOCAL_BENCH
#include <stdio.h>
#include <time.h>
#endif
namespace {
constexpr int RN = 1 << 20;
constexpr int FULL_N = 1 << 21;
constexpr float INV_RN = 1.0f / RN;
struct alignas(8) Cf { float r, i; };
struct alignas(16) Cd { double r,i; };
static inline Cf add(Cf,Cf);
static inline Cf sub(Cf,Cf);
static inline Cf mi(Cf);
static inline Cf pi(Cf);
static inline __m256 cmul(__m256 a, __m256 b) {
__m256 ar = _mm256_moveldup_ps(a);
__m256 ai = _mm256_movehdup_ps(a);
__m256 bs = _mm256_permute_ps(b, 0xb1);
return _mm256_fmaddsub_ps(ar, b, _mm256_mul_ps(ai, bs));
}
static inline __m256 cmul_precise(__m256 a,__m256 b){
auto two=[](__m128 af,__m128 bf){
__m256d x=_mm256_cvtps_pd(af),y=_mm256_cvtps_pd(bf);
__m256d xr=_mm256_movedup_pd(x),xi=_mm256_permute_pd(x,15),ys=_mm256_permute_pd(y,5);
return _mm256_cvtpd_ps(_mm256_fmaddsub_pd(xr,y,_mm256_mul_pd(xi,ys)));
};
__m128 lo=two(_mm256_castps256_ps128(a),_mm256_castps256_ps128(b));
__m128 hi=two(_mm256_extractf128_ps(a,1),_mm256_extractf128_ps(b,1));
return _mm256_insertf128_ps(_mm256_castps128_ps256(lo),hi,1);
}
template<int N> static inline __m256 stage_mul(__m256 a,__m256 b){return N<=256?cmul_precise(a,b):cmul(a,b);}
static inline __m256 conjv(__m256 a) {
return _mm256_xor_ps(a, _mm256_castsi256_ps(
_mm256_set_epi32(0x80000000u,0,0x80000000u,0,0x80000000u,0,0x80000000u,0)));
}
static inline __m256 mul_minus_i(__m256 a) {
__m256 s = _mm256_permute_ps(a, 0xb1);
return _mm256_xor_ps(s, _mm256_castsi256_ps(
_mm256_set_epi32(0x80000000u,0,0x80000000u,0,0x80000000u,0,0x80000000u,0)));
}
static inline __m256 mul_plus_i(__m256 a) {
__m256 s = _mm256_permute_ps(a, 0xb1);
return _mm256_xor_ps(s, _mm256_castsi256_ps(
_mm256_set_epi32(0,0x80000000u,0,0x80000000u,0,0x80000000u,0,0x80000000u)));
}
// For every radix-4 layer, store the three twiddle streams consecutively.
// Total number of complex roots is exactly RN-1.
static Cf *tw;
static Cd *twd;
static int tw_off[21];
static void build_twiddles() {
if (tw) return;
posix_memalign((void **)&tw, 32, (size_t)RN * sizeof(Cf));
#ifdef DOUBLE_ROOT_TEST
posix_memalign((void **)&twd,32,(size_t)RN*sizeof(Cd));
#endif
int off = 0;
for (int len = 4; len <= RN; len <<= 2) {
tw_off[__builtin_ctz((unsigned)len)] = off;
int q = len >> 2;
Cf *w1 = tw + off, *w2 = w1 + q, *w3 = w2 + q;
double ang = -2.0 * 3.14159265358979323846264338327950288 / len;
double sr=cos(4*ang),si=sin(4*ang);
__m256d cr,ci,vsr=_mm256_set1_pd(sr),vsi=_mm256_set1_pd(si);
for (int j = 0; j < q; j += 4) {
if ((j & 4095) == 0) {
alignas(32) double vr[4],vii[4];
for(int k=0;k<4;++k){double x=ang*(j+k);vr[k]=cos(x);vii[k]=sin(x);}
cr=_mm256_load_pd(vr);ci=_mm256_load_pd(vii);
}
__m256d c2r=_mm256_fmsub_pd(cr,cr,_mm256_mul_pd(ci,ci));
__m256d c2i=_mm256_add_pd(_mm256_mul_pd(cr,ci),_mm256_mul_pd(cr,ci));
__m256d c3r=_mm256_fmsub_pd(c2r,cr,_mm256_mul_pd(c2i,ci));
__m256d c3i=_mm256_fmadd_pd(c2r,ci,_mm256_mul_pd(c2i,cr));
auto store4=[](Cf*d,__m256d rr,__m256d ii){
__m128 r=_mm256_cvtpd_ps(rr),im=_mm256_cvtpd_ps(ii);
__m256 o=_mm256_castps128_ps256(_mm_unpacklo_ps(r,im));
o=_mm256_insertf128_ps(o,_mm_unpackhi_ps(r,im),1);
_mm256_storeu_ps((float*)d,o);
};
store4(w1+j,cr,ci);store4(w2+j,c2r,c2i);store4(w3+j,c3r,c3i);
#ifdef DOUBLE_ROOT_TEST
alignas(32) double rr[4],ii[4],r2a[4],i2a[4],r3a[4],i3a[4];
_mm256_store_pd(rr,cr);_mm256_store_pd(ii,ci);_mm256_store_pd(r2a,c2r);_mm256_store_pd(i2a,c2i);_mm256_store_pd(r3a,c3r);_mm256_store_pd(i3a,c3i);
for(int k=0;k<4;++k){twd[off+j+k]={rr[k],ii[k]};twd[off+q+j+k]={r2a[k],i2a[k]};twd[off+2*q+j+k]={r3a[k],i3a[k]};}
#endif
__m256d nr=_mm256_fmsub_pd(cr,vsr,_mm256_mul_pd(ci,vsi));
ci=_mm256_fmadd_pd(cr,vsi,_mm256_mul_pd(ci,vsr));cr=nr;
}
off += 3*q;
}
}
static inline __m256 forward_mul(__m256 a,const Cf*b){
#ifdef DOUBLE_ROOT_TEST
const Cd*d=twd+(b-tw);
auto two=[](__m128 af,const Cd*p){
__m256d x=_mm256_cvtps_pd(af),y=_mm256_loadu_pd((const double*)p);
__m256d xr=_mm256_movedup_pd(x),xi=_mm256_permute_pd(x,15),ys=_mm256_permute_pd(y,5);
return _mm256_cvtpd_ps(_mm256_fmaddsub_pd(xr,y,_mm256_mul_pd(xi,ys)));
};
__m128 lo=two(_mm256_castps256_ps128(a),d),hi=two(_mm256_extractf128_ps(a,1),d+2);
return _mm256_insertf128_ps(_mm256_castps128_ps256(lo),hi,1);
#else
return cmul(a,_mm256_loadu_ps((const float*)b));
#endif
}
template<int N> struct SmallFFT {
__attribute__((always_inline)) static inline void forward(Cf *x) {
constexpr int q=N/4;
Cf *r1=tw+tw_off[__builtin_ctz((unsigned)N)],*r2=r1+q,*r3=r2+q;
for(int j=0;j<q;j+=4){
__m256 a=_mm256_load_ps((float*)(x+j)),b=_mm256_load_ps((float*)(x+q+j));
__m256 c=_mm256_load_ps((float*)(x+2*q+j)),d=_mm256_load_ps((float*)(x+3*q+j));
__m256 t0=_mm256_add_ps(a,c),t1=_mm256_sub_ps(a,c);
__m256 t2=_mm256_add_ps(b,d),t3=mul_minus_i(_mm256_sub_ps(b,d));
_mm256_store_ps((float*)(x+j),_mm256_add_ps(t0,t2));
_mm256_store_ps((float*)(x+q+j),forward_mul(_mm256_add_ps(t1,t3),r1+j));
_mm256_store_ps((float*)(x+2*q+j),forward_mul(_mm256_sub_ps(t0,t2),r2+j));
_mm256_store_ps((float*)(x+3*q+j),forward_mul(_mm256_sub_ps(t1,t3),r3+j));
}
SmallFFT<q>::forward(x);SmallFFT<q>::forward(x+q);
SmallFFT<q>::forward(x+2*q);SmallFFT<q>::forward(x+3*q);
}
__attribute__((always_inline)) static inline void inverse(Cf *x) {
constexpr int q=N/4;
SmallFFT<q>::inverse(x);SmallFFT<q>::inverse(x+q);
SmallFFT<q>::inverse(x+2*q);SmallFFT<q>::inverse(x+3*q);
Cf *r1=tw+tw_off[__builtin_ctz((unsigned)N)],*r2=r1+q,*r3=r2+q;
for(int j=0;j<q;j+=4){
__m256 a=_mm256_load_ps((float*)(x+j));
__m256 b=cmul(_mm256_load_ps((float*)(x+q+j)),conjv(_mm256_loadu_ps((float*)(r1+j))));
__m256 c=cmul(_mm256_load_ps((float*)(x+2*q+j)),conjv(_mm256_loadu_ps((float*)(r2+j))));
__m256 d=cmul(_mm256_load_ps((float*)(x+3*q+j)),conjv(_mm256_loadu_ps((float*)(r3+j))));
__m256 t0=_mm256_add_ps(a,c),t1=_mm256_sub_ps(a,c);
__m256 t2=_mm256_add_ps(b,d),t3=mul_plus_i(_mm256_sub_ps(b,d));
_mm256_store_ps((float*)(x+j),_mm256_add_ps(t0,t2));
_mm256_store_ps((float*)(x+q+j),_mm256_add_ps(t1,t3));
_mm256_store_ps((float*)(x+2*q+j),_mm256_sub_ps(t0,t2));
_mm256_store_ps((float*)(x+3*q+j),_mm256_sub_ps(t1,t3));
}
}
};
template<> struct SmallFFT<4> {
__attribute__((always_inline)) static inline void forward(Cf*x){
Cf a=x[0],b=x[1],c=x[2],d=x[3],t0=add(a,c),t1=sub(a,c),t2=add(b,d),t3=mi(sub(b,d));
x[0]=add(t0,t2);x[1]=add(t1,t3);x[2]=sub(t0,t2);x[3]=sub(t1,t3);
}
__attribute__((always_inline)) static inline void inverse(Cf*x){
Cf a=x[0],b=x[1],c=x[2],d=x[3],t0=add(a,c),t1=sub(a,c),t2=add(b,d),t3=pi(sub(b,d));
x[0]=add(t0,t2);x[1]=add(t1,t3);x[2]=sub(t0,t2);x[3]=sub(t1,t3);
}
};
static void fft_forward(Cf *x, int n) {
if (n == 1) return;
if (n == 256) { SmallFFT<256>::forward(x); return; }
if (n == 4) {
Cf a=x[0],b=x[1],c=x[2],d=x[3];
Cf t0=add(a,c),t1=sub(a,c),t2=add(b,d),t3=mi(sub(b,d));
x[0]=add(t0,t2);x[1]=add(t1,t3);x[2]=sub(t0,t2);x[3]=sub(t1,t3);
return;
}
int q = n >> 2;
Cf *r1 = tw + tw_off[__builtin_ctz((unsigned)n)];
Cf *r2 = r1 + q, *r3 = r2 + q;
for (int j = 0; j < q; j += 4) {
__m256 a = _mm256_load_ps((float *)(x + j));
__m256 b = _mm256_load_ps((float *)(x + q + j));
__m256 c = _mm256_load_ps((float *)(x + 2*q + j));
__m256 d = _mm256_load_ps((float *)(x + 3*q + j));
__m256 t0 = _mm256_add_ps(a,c), t1 = _mm256_sub_ps(a,c);
__m256 t2 = _mm256_add_ps(b,d), t3 = mul_minus_i(_mm256_sub_ps(b,d));
_mm256_store_ps((float *)(x+j), _mm256_add_ps(t0,t2));
_mm256_store_ps((float *)(x+q+j), forward_mul(_mm256_add_ps(t1,t3),r1+j));
_mm256_store_ps((float *)(x+2*q+j), forward_mul(_mm256_sub_ps(t0,t2),r2+j));
_mm256_store_ps((float *)(x+3*q+j), forward_mul(_mm256_sub_ps(t1,t3),r3+j));
}
fft_forward(x, q);
fft_forward(x+q, q);
fft_forward(x+2*q, q);
fft_forward(x+3*q, q);
}
static void fft_inverse(Cf *x, int n) {
if (n == 1) return;
if (n == 256) { SmallFFT<256>::inverse(x); return; }
if (n == 4) {
Cf a=x[0],b=x[1],c=x[2],d=x[3];
Cf t0=add(a,c),t1=sub(a,c),t2=add(b,d),t3=pi(sub(b,d));
x[0]=add(t0,t2);x[1]=add(t1,t3);x[2]=sub(t0,t2);x[3]=sub(t1,t3);
return;
}
int q = n >> 2;
fft_inverse(x, q);
fft_inverse(x+q, q);
fft_inverse(x+2*q, q);
fft_inverse(x+3*q, q);
Cf *r1 = tw + tw_off[__builtin_ctz((unsigned)n)];
Cf *r2 = r1 + q, *r3 = r2 + q;
for (int j = 0; j < q; j += 4) {
__m256 a = _mm256_load_ps((float *)(x+j));
__m256 b = cmul(_mm256_load_ps((float *)(x+q+j)), conjv(_mm256_loadu_ps((float *)(r1+j))));
__m256 c = cmul(_mm256_load_ps((float *)(x+2*q+j)), conjv(_mm256_loadu_ps((float *)(r2+j))));
__m256 d = cmul(_mm256_load_ps((float *)(x+3*q+j)), conjv(_mm256_loadu_ps((float *)(r3+j))));
__m256 t0 = _mm256_add_ps(a,c), t1 = _mm256_sub_ps(a,c);
__m256 t2 = _mm256_add_ps(b,d), t3 = mul_plus_i(_mm256_sub_ps(b,d));
_mm256_store_ps((float *)(x+j), _mm256_add_ps(t0,t2));
_mm256_store_ps((float *)(x+q+j), _mm256_add_ps(t1,t3));
_mm256_store_ps((float *)(x+2*q+j), _mm256_sub_ps(t0,t2));
_mm256_store_ps((float *)(x+3*q+j), _mm256_sub_ps(t1,t3));
}
}
static uint16_t rev10[1024];
static Cf root_lo[1024], root_hi[1024];
#ifdef DOUBLE_ROOT_TEST
static Cd root_lod[1024],root_hid[1024];
#endif
static void build_real_roots() {
for (int x=0;x<1024;++x) {
unsigned r=0;
for(int i=0;i<5;++i) r=(r<<2)|((x>>(2*i))&3);
rev10[x]=(uint16_t)r;
double a=-2.0*3.14159265358979323846264338327950288*r/FULL_N;
root_hi[x]={(float)cos(a),(float)sin(a)};
#ifdef DOUBLE_ROOT_TEST
root_hid[x]={cos(a),sin(a)};
#endif
a*=1024.0;
root_lo[x]={(float)cos(a),(float)sin(a)};
#ifdef DOUBLE_ROOT_TEST
root_lod[x]={cos(a),sin(a)};
#endif
}
}
static inline unsigned rev20(unsigned p) {
return ((unsigned)rev10[p&1023]<<10) | rev10[p>>10];
}
static inline Cf real_root(unsigned p) {
Cf a=root_hi[p>>10], b=root_lo[p&1023];
return {a.r*b.r-a.i*b.i,a.r*b.i+a.i*b.r};
}
#ifdef DOUBLE_ROOT_TEST
static inline Cd real_root_d(unsigned p){Cd a=root_hid[p>>10],b=root_lod[p&1023];return {a.r*b.r-a.i*b.i,a.r*b.i+a.i*b.r};}
#endif
static inline Cf add(Cf a,Cf b){return {a.r+b.r,a.i+b.i};}
static inline Cf sub(Cf a,Cf b){return {a.r-b.r,a.i-b.i};}
static inline Cf mul(Cf a,Cf b){return {a.r*b.r-a.i*b.i,a.r*b.i+a.i*b.r};}
static inline Cf cj(Cf a){return {a.r,-a.i};}
static inline Cf mi(Cf a){return {a.i,-a.r};}
static inline Cf pi(Cf a){return {-a.i,a.r};}
static inline Cf half(Cf a){return {a.r*.5f,a.i*.5f};}
static void real_finish(Cf *z) {
float zr=z[0].r, zi=z[0].i;
z[0]={zr+zi,zr-zi};
for(unsigned h=1;h<RN;h<<=2) {
for(unsigned p=h;;++p) {
unsigned q=5*h-1-p;
if (p>q) break;
#ifdef DOUBLE_ROOT_TEST
Cd a={(double)z[p].r,z[p].i},b={(double)z[q].r,-(double)z[q].i};
Cd e={(a.r+b.r)*.5,(a.i+b.i)*.5},o={(a.i-b.i)*.5,(b.r-a.r)*.5},w=real_root_d(p);
Cd wo={w.r*o.r-w.i*o.i,w.r*o.i+w.i*o.r};
z[p]={(float)(e.r+wo.r),(float)(e.i+wo.i)};
if(p!=q)z[q]={(float)(e.r-wo.r),(float)(wo.i-e.i)};
#else
Cf a=z[p], b=cj(z[q]);
Cf e=half(add(a,b)), o=mi(half(sub(a,b))), wo=mul(real_root(p),o);
z[p]=add(e,wo);
if(p!=q) z[q]=cj(sub(e,wo));
#endif
}
}
}
static void real_prepare_inverse(Cf *x) {
float x0=x[0].r, xm=x[0].i;
x[0]={(x0+xm)*.5f,(x0-xm)*.5f};
for(unsigned h=1;h<RN;h<<=2) {
for(unsigned p=h;;++p) {
unsigned q=5*h-1-p;
if(p>q) break;
Cf a=x[p], b=cj(x[q]), e=half(add(a,b));
Cf o=mul(half(sub(a,b)),cj(real_root(p)));
x[p]=add(e,pi(o));
if(p!=q) x[q]=cj(sub(e,pi(o)));
}
}
}
}
void poly_multiply(unsigned *a,int n,unsigned *b,int m,unsigned *c) {
build_twiddles(); build_real_roots();
Cf *fa,*fb;
posix_memalign((void **)&fa,32,(size_t)RN*sizeof(Cf));
posix_memalign((void **)&fb,32,(size_t)RN*sizeof(Cf));
for(int i=0;i<RN;++i) {
int j=i<<1;
#ifdef CENTERED_INPUT
fa[i]={(float)(j<=n?(int)(a[j]<<1)-9:0),(float)(j+1<=n?(int)(a[j+1]<<1)-9:0)};
fb[i]={(float)(j<=m?(int)(b[j]<<1)-9:0),(float)(j+1<=m?(int)(b[j+1]<<1)-9:0)};
#else
fa[i]={(float)(j<=n?a[j]:0),(float)(j+1<=n?a[j+1]:0)};
fb[i]={(float)(j<=m?b[j]:0),(float)(j+1<=m?b[j+1]:0)};
#endif
}
fft_forward(fa,RN); fft_forward(fb,RN);
real_finish(fa); real_finish(fb);
#ifdef HIGH_INV
fftw_complex *spec=(fftw_complex*)fftw_malloc((size_t)FULL_N*sizeof(fftw_complex));
spec[0][0]=(double)fa[0].r*fb[0].r;spec[0][1]=0;
spec[RN][0]=(double)fa[0].i*fb[0].i;spec[RN][1]=0;
for(unsigned p=1;p<RN;++p){unsigned k=rev20(p);double r=(double)fa[p].r*fb[p].r-(double)fa[p].i*fb[p].i,im=(double)fa[p].r*fb[p].i+(double)fa[p].i*fb[p].r;spec[k][0]=r;spec[k][1]=im;spec[FULL_N-k][0]=r;spec[FULL_N-k][1]=-im;}
fftw_plan plan=fftw_plan_dft_1d(FULL_N,spec,spec,FFTW_BACKWARD,FFTW_ESTIMATE);fftw_execute(plan);
for(int i=0;i<n+m+1;++i)c[i]=(unsigned)llround(spec[i][0]/FULL_N+ROUND_BIAS);
fftw_destroy_plan(plan);fftw_free(spec);return;
#endif
fa[0]={fa[0].r*fb[0].r,fa[0].i*fb[0].i};
for(int i=1;i<RN;++i) fa[i]=mul(fa[i],fb[i]);
real_prepare_inverse(fa);
fft_inverse(fa,RN);
int need=n+m+1;
#ifdef CENTERED_INPUT
int sa=0,sb=0;
const __m256 invrn8=_mm256_set1_ps(INV_RN*.25f);
const __m256d invrnd=_mm256_set1_pd((1.0/RN)*.25);
int k=0;
for(;k+8<=need;k+=8) {
alignas(32) int corr[8];
for(int t=0;t<8;++t) {
int z=k+t;
if(z<=n) sa+=(int)(a[z]<<1)-9;
if(z-m-1>=0) sa-=(int)(a[z-m-1]<<1)-9;
if(z<=m) sb+=(int)(b[z]<<1)-9;
if(z-n-1>=0) sb-=(int)(b[z-n-1]<<1)-9;
int lo=z-m>0?z-m:0, hi=z<n?z:n;
corr[t]=(9*sa+9*sb+81*(hi-lo+1));
}
__m256 fv=_mm256_loadu_ps((const float *)(fa+(k>>1)));
__m128 f0=_mm256_castps256_ps128(fv), f1=_mm256_extractf128_ps(fv,1);
__m256d d0=_mm256_add_pd(_mm256_mul_pd(_mm256_cvtps_pd(f0),invrnd),
_mm256_mul_pd(_mm256_cvtepi32_pd(_mm_load_si128((const __m128i*)corr)),_mm256_set1_pd(.25)));
__m256d d1=_mm256_add_pd(_mm256_mul_pd(_mm256_cvtps_pd(f1),invrnd),
_mm256_mul_pd(_mm256_cvtepi32_pd(_mm_load_si128((const __m128i*)(corr+4))),_mm256_set1_pd(.25)));
__m128i i0=_mm256_cvtpd_epi32(d0),i1=_mm256_cvtpd_epi32(d1);
_mm256_storeu_si256((__m256i*)(c+k),_mm256_inserti128_si256(_mm256_castsi128_si256(i0),i1,1));
}
for(;k<need;++k) {
if(k<=n) sa+=(int)(a[k]<<1)-9;
if(k-m-1>=0) sa-=(int)(a[k-m-1]<<1)-9;
if(k<=m) sb+=(int)(b[k]<<1)-9;
if(k-n-1>=0) sb-=(int)(b[k-n-1]<<1)-9;
int lo=k-m>0?k-m:0, hi=k<n?k:n;
int cnt=hi-lo+1;
float fv=(k&1)?fa[k>>1].i:fa[k>>1].r;
double exact_part=(double)(9*sa+9*sb+81ll*cnt);
c[k]=(unsigned)llround(((double)fv*(1.0/RN)+exact_part)*.25);
}
#else
for(int i=0;i<RN && 2*i<need;++i) {
c[2*i]=(unsigned)lrintf(fa[i].r*INV_RN);
if(2*i+1<need)c[2*i+1]=(unsigned)lrintf(fa[i].i*INV_RN);
}
#endif
}
#ifdef LOCAL_BENCH
static unsigned aa[1000001],bb[1000001],cc[2000001];
int main(){
for(int i=0;i<=1000000;++i){aa[i]=(i*7+3)%10;bb[i]=(i*5+1)%10;}
clock_t s=clock();poly_multiply(aa,1000000,bb,1000000,cc);clock_t e=clock();
unsigned long long z=0;for(int i=0;i<=2000000;i+=137)z+=cc[i];
printf("%.3f ms %llu center=%u\n",1000.0*(e-s)/CLOCKS_PER_SEC,z,cc[1000000]);
s=clock();poly_multiply(aa,1000000,bb,1000000,cc);e=clock();
printf("warm %.3f ms\n",1000.0*(e-s)/CLOCKS_PER_SEC);
}
#endif
| Compilation | N/A | N/A | Compile OK | Score: N/A | 显示更多 |
| Testcase #1 | 38.849 ms | 31 MB + 676 KB | Accepted | Score: 100 | 显示更多 |