代码

#pragma GCC optimize("O3,unroll-loops,omit-frame-pointer")
#pragma GCC target("arch=haswell")
#include <immintrin.h>

enum { N=1024,H=512,MAIN=504,PANELS=42 };
static double apack[H*H] __attribute__((aligned(4096)));
static double bpanel[H*12] __attribute__((aligned(4096)));

static void pack_a(const double*p,const double*q,int sub){
    double*d=apack;
    for(int i=0;i<H;++i,p+=N,d+=H){
        if(!q)for(int j=0;j<H;j+=4)_mm256_store_pd(d+j,_mm256_load_pd(p+j));
        else if(sub){for(int j=0;j<H;j+=4)_mm256_store_pd(d+j,_mm256_sub_pd(_mm256_load_pd(p+j),_mm256_load_pd(q+j)));q+=N;}
        else{for(int j=0;j<H;j+=4)_mm256_store_pd(d+j,_mm256_add_pd(_mm256_load_pd(p+j),_mm256_load_pd(q+j)));q+=N;}
    }
}
static void pack12(const double*p,const double*q,int sub,int col){
    double*d=bpanel;p+=col;q=q?q+col:0;
    for(int k=0;k<H;++k,d+=12,p+=N){
        __m256d x0=_mm256_load_pd(p),x1=_mm256_load_pd(p+4),x2=_mm256_load_pd(p+8);
        if(q){__m256d y0=_mm256_load_pd(q),y1=_mm256_load_pd(q+4),y2=_mm256_load_pd(q+8);
            x0=sub?_mm256_sub_pd(x0,y0):_mm256_add_pd(x0,y0);x1=sub?_mm256_sub_pd(x1,y1):_mm256_add_pd(x1,y1);x2=sub?_mm256_sub_pd(x2,y2):_mm256_add_pd(x2,y2);q+=N;}
        _mm256_store_pd(d,x0);_mm256_store_pd(d+4,x1);_mm256_store_pd(d+8,x2);
    }
}
static void pack8(const double*p,const double*q,int sub){
    double*d=bpanel;p+=MAIN;q=q?q+MAIN:0;
    for(int k=0;k<H;++k,d+=8,p+=N){__m256d x0=_mm256_load_pd(p),x1=_mm256_load_pd(p+4);
        if(q){__m256d y0=_mm256_load_pd(q),y1=_mm256_load_pd(q+4);x0=sub?_mm256_sub_pd(x0,y0):_mm256_add_pd(x0,y0);x1=sub?_mm256_sub_pd(x1,y1):_mm256_add_pd(x1,y1);q+=N;}
        _mm256_store_pd(d,x0);_mm256_store_pd(d+4,x1);}
}
static inline void put3(double*d,__m256d a,__m256d b,__m256d c,int op){
    if(!op){_mm256_store_pd(d,a);_mm256_store_pd(d+4,b);_mm256_store_pd(d+8,c);}
    else if(op>0){_mm256_store_pd(d,_mm256_add_pd(_mm256_load_pd(d),a));_mm256_store_pd(d+4,_mm256_add_pd(_mm256_load_pd(d+4),b));_mm256_store_pd(d+8,_mm256_add_pd(_mm256_load_pd(d+8),c));}
    else{_mm256_store_pd(d,_mm256_sub_pd(_mm256_load_pd(d),a));_mm256_store_pd(d+4,_mm256_sub_pd(_mm256_load_pd(d+4),b));_mm256_store_pd(d+8,_mm256_sub_pd(_mm256_load_pd(d+8),c));}
}
static inline void put2(double*d,__m256d a,__m256d b,int op){
    if(!op){_mm256_store_pd(d,a);_mm256_store_pd(d+4,b);}
    else if(op>0){_mm256_store_pd(d,_mm256_add_pd(_mm256_load_pd(d),a));_mm256_store_pd(d+4,_mm256_add_pd(_mm256_load_pd(d+4),b));}
    else{_mm256_store_pd(d,_mm256_sub_pd(_mm256_load_pd(d),a));_mm256_store_pd(d+4,_mm256_sub_pd(_mm256_load_pd(d+4),b));}
}
static inline void kernel12(const double*aa,double*C1,double*C2,int op1,int op2){
    __m256d z=_mm256_setzero_pd(),c00=z,c01=z,c02=z,c10=z,c11=z,c12=z,c20=z,c21=z,c22=z,c30=z,c31=z,c32=z;
    const double*a0=aa,*a1=aa+H,*a2=aa+2*H,*a3=aa+3*H,*bp=bpanel;
    for(int k=0;k<H;++k,bp+=12){__m256d b0=_mm256_load_pd(bp),b1=_mm256_load_pd(bp+4),b2=_mm256_load_pd(bp+8),a=_mm256_broadcast_sd(a0+k);
        c00=_mm256_fmadd_pd(a,b0,c00);c01=_mm256_fmadd_pd(a,b1,c01);c02=_mm256_fmadd_pd(a,b2,c02);a=_mm256_broadcast_sd(a1+k);
        c10=_mm256_fmadd_pd(a,b0,c10);c11=_mm256_fmadd_pd(a,b1,c11);c12=_mm256_fmadd_pd(a,b2,c12);a=_mm256_broadcast_sd(a2+k);
        c20=_mm256_fmadd_pd(a,b0,c20);c21=_mm256_fmadd_pd(a,b1,c21);c22=_mm256_fmadd_pd(a,b2,c22);a=_mm256_broadcast_sd(a3+k);
        c30=_mm256_fmadd_pd(a,b0,c30);c31=_mm256_fmadd_pd(a,b1,c31);c32=_mm256_fmadd_pd(a,b2,c32);}
    put3(C1,c00,c01,c02,op1);put3(C1+N,c10,c11,c12,op1);put3(C1+2*N,c20,c21,c22,op1);put3(C1+3*N,c30,c31,c32,op1);
    if(C2){put3(C2,c00,c01,c02,op2);put3(C2+N,c10,c11,c12,op2);put3(C2+2*N,c20,c21,c22,op2);put3(C2+3*N,c30,c31,c32,op2);}
}
static inline void kernel8(const double*aa,double*C1,double*C2,int op1,int op2){
    __m256d z=_mm256_setzero_pd(),c00=z,c01=z,c10=z,c11=z,c20=z,c21=z,c30=z,c31=z;const double*a0=aa,*a1=aa+H,*a2=aa+2*H,*a3=aa+3*H,*bp=bpanel;
    for(int k=0;k<H;++k,bp+=8){__m256d b0=_mm256_load_pd(bp),b1=_mm256_load_pd(bp+4),a=_mm256_broadcast_sd(a0+k);c00=_mm256_fmadd_pd(a,b0,c00);c01=_mm256_fmadd_pd(a,b1,c01);a=_mm256_broadcast_sd(a1+k);c10=_mm256_fmadd_pd(a,b0,c10);c11=_mm256_fmadd_pd(a,b1,c11);a=_mm256_broadcast_sd(a2+k);c20=_mm256_fmadd_pd(a,b0,c20);c21=_mm256_fmadd_pd(a,b1,c21);a=_mm256_broadcast_sd(a3+k);c30=_mm256_fmadd_pd(a,b0,c30);c31=_mm256_fmadd_pd(a,b1,c31);}
    put2(C1,c00,c01,op1);put2(C1+N,c10,c11,op1);put2(C1+2*N,c20,c21,op1);put2(C1+3*N,c30,c31,op1);if(C2){put2(C2,c00,c01,op2);put2(C2+N,c10,c11,op2);put2(C2+2*N,c20,c21,op2);put2(C2+3*N,c30,c31,op2);}
}
static void product(const double*B,const double*BQ,int sub,double*C1,int op1,double*C2,int op2){
    for(int panel=0;panel<PANELS;++panel){int col=panel*12;pack12(B,BQ,sub,col);for(int i=0;i<H;i+=4)kernel12(apack+i*H,C1+i*N+col,C2?C2+i*N+col:0,op1,op2);}
    pack8(B,BQ,sub);for(int i=0;i<H;i+=4)kernel8(apack+i*H,C1+i*N+MAIN,C2?C2+i*N+MAIN:0,op1,op2);
}
void matrix_multiply(int n,const double*A,const double*B,double*C){(void)n;const double*A00=A,*A01=A+H,*A10=A+H*N,*A11=A10+H,*B00=B,*B01=B+H,*B10=B+H*N,*B11=B10+H;double*C00=C,*C01=C+H,*C10=C+H*N,*C11=C10+H;
    pack_a(A00,A11,0);product(B00,B11,0,C00,0,C11,0);pack_a(A10,A11,0);product(B00,0,0,C10,0,C11,-1);pack_a(A00,0,0);product(B01,B11,1,C01,0,C11,1);pack_a(A11,0,0);product(B10,B00,1,C00,1,C10,1);pack_a(A00,A01,0);product(B11,0,0,C00,-1,C01,1);pack_a(A10,A00,1);product(B00,B01,0,C11,1,0,0);pack_a(A01,A11,1);product(B10,B11,0,C00,1,0,0);_mm256_zeroupper();}

评测结果