代码

// Floating-point Strassen multiplication with a classical AVX2/FMA leaf.
// The top level keeps sums as views, reducing temporary written memory.
#pragma GCC optimize("O3")
#pragma GCC target("avx2,fma")
#include <immintrin.h>
#include <stddef.h>
#include <stdlib.h>
#include <string.h>
namespace mat {
struct View { const float*x; int stride; const float*y=nullptr; int sign=0; };
static inline float get(View a,int i,int j) {
    const size_t t=(size_t)i*a.stride+j;
    return a.sign ? a.x[t]+float(a.sign)*a.y[t] : a.x[t];
}
static inline View block(View a,int i,int j){
    const size_t t=(size_t)i*a.stride+j;
    a.x+=t;if(a.y)a.y+=t;return a;
}
static inline View sum(View a,View b,int sign) {return {a.x,a.stride,b.x,sign};}
constexpr int MR=6,NR=16,KC=256,MAX_LEAF=512,ROWS=516;
alignas(4096) static float ap[ROWS*KC],bp[MAX_LEAF*MAX_LEAF],ct[ROWS*MAX_LEAF];
static inline void micro(const float*a,const float*b,float*c,int stride){
    __m256 c00=_mm256_load_ps(c+0*stride),c01=_mm256_load_ps(c+0*stride+8);
    __m256 c10=_mm256_load_ps(c+1*stride),c11=_mm256_load_ps(c+1*stride+8);
    __m256 c20=_mm256_load_ps(c+2*stride),c21=_mm256_load_ps(c+2*stride+8);
    __m256 c30=_mm256_load_ps(c+3*stride),c31=_mm256_load_ps(c+3*stride+8);
    __m256 c40=_mm256_load_ps(c+4*stride),c41=_mm256_load_ps(c+4*stride+8);
    __m256 c50=_mm256_load_ps(c+5*stride),c51=_mm256_load_ps(c+5*stride+8);
#pragma GCC unroll 4
    for(int k=0;k<KC;++k){
        __m256 b0=_mm256_load_ps(b),b1=_mm256_load_ps(b+8);
        __m256 a0=_mm256_broadcast_ss(a+0);c00=_mm256_fmadd_ps(a0,b0,c00);c01=_mm256_fmadd_ps(a0,b1,c01);
        __m256 a1=_mm256_broadcast_ss(a+1);c10=_mm256_fmadd_ps(a1,b0,c10);c11=_mm256_fmadd_ps(a1,b1,c11);
        __m256 a2=_mm256_broadcast_ss(a+2);c20=_mm256_fmadd_ps(a2,b0,c20);c21=_mm256_fmadd_ps(a2,b1,c21);
        __m256 a3=_mm256_broadcast_ss(a+3);c30=_mm256_fmadd_ps(a3,b0,c30);c31=_mm256_fmadd_ps(a3,b1,c31);
        __m256 a4=_mm256_broadcast_ss(a+4);c40=_mm256_fmadd_ps(a4,b0,c40);c41=_mm256_fmadd_ps(a4,b1,c41);
        __m256 a5=_mm256_broadcast_ss(a+5);c50=_mm256_fmadd_ps(a5,b0,c50);c51=_mm256_fmadd_ps(a5,b1,c51);
        a+=MR;b+=NR;
    }
    _mm256_store_ps(c+0*stride,c00);_mm256_store_ps(c+0*stride+8,c01);
    _mm256_store_ps(c+1*stride,c10);_mm256_store_ps(c+1*stride+8,c11);
    _mm256_store_ps(c+2*stride,c20);_mm256_store_ps(c+2*stride+8,c21);
    _mm256_store_ps(c+3*stride,c30);_mm256_store_ps(c+3*stride+8,c31);
    _mm256_store_ps(c+4*stride,c40);_mm256_store_ps(c+4*stride+8,c41);
    _mm256_store_ps(c+5*stride,c50);_mm256_store_ps(c+5*stride+8,c51);
}
static void classical(int n,View a,View b,float*c,int stride){
    if(n%KC||n>MAX_LEAF){
        for(int i=0;i<n;++i)for(int j=0;j<n;++j){
            float v=0;for(int k=0;k<n;++k)v+=get(a,i,k)*get(b,k,j);
            c[(size_t)i*stride+j]=v;
        }
        return;
    }
    const int rows=(n+MR-1)/MR*MR;
    memset(ct,0,(size_t)rows*n*sizeof(float));
    for(int pc=0;pc<n;pc+=KC)for(int q=0;q<n;q+=NR)
        for(int k=0;k<KC;++k)for(int j=0;j<NR;++j)
            bp[((pc/KC*(n/NR)+q/NR)*KC+k)*NR+j]=get(b,pc+k,q+j);
    for(int pc=0;pc<n;pc+=KC){
        for(int r=0;r<rows;r+=MR)for(int k=0;k<KC;++k)for(int i=0;i<MR;++i)
            ap[(r/MR*KC+k)*MR+i]=r+i<n?get(a,r+i,pc+k):0;
        for(int q=0;q<n;q+=NR){
            const float*packed=bp+((pc/KC)*(n/NR)+q/NR)*KC*NR;
            for(int r=0;r<rows;r+=MR)micro(ap+r*KC,packed,ct+(size_t)r*n+q,n);
        }
    }
    for(int i=0;i<n;++i)memcpy(c+(size_t)i*stride,ct+(size_t)i*n,(size_t)n*sizeof(float));
}
static void linear(int n,View a,View b,float*out,int sign){
    for(int i=0;i<n;++i)for(int j=0;j<n;++j)
        out[(size_t)i*n+j]=get(a,i,j)+float(sign)*get(b,i,j);
}
static void add(int n,float*c,int stride,const float*p,int sign){
    for(int i=0;i<n;++i)for(int j=0;j<n;++j)
        c[(size_t)i*stride+j]+=float(sign)*p[(size_t)i*n+j];
}
static void rec(int n,View a,View b,float*c,int stride,float*work,bool lazy=false){
    if(n<=512){classical(n,a,b,c,stride);return;}
    const int h=n/2;const size_t z=(size_t)h*h;
    float*s=work,*t=s+z,*p=lazy?work:t+z,*next=p+z;
    View aa[4]={block(a,0,0),block(a,0,h),block(a,h,0),block(a,h,h)};
    View bb[4]={block(b,0,0),block(b,0,h),block(b,h,0),block(b,h,h)};
    float*cc[4]={c,c+h,c+(size_t)h*stride,c+(size_t)h*stride+h};
    const int ai[7]={0,2,0,3,0,2,1},aj[7]={3,3,-1,-1,1,0,3},as[7]={1,1,0,0,1,-1,-1};
    const int bi[7]={0,0,1,2,3,0,2},bj[7]={3,-1,3,0,-1,1,3},bs[7]={1,0,-1,-1,0,1,1};
    for(int r=0;r<7;++r){
        View x=aa[ai[r]],y=bb[bi[r]];
        if(as[r]){
            if(lazy)x=sum(x,aa[aj[r]],as[r]);
            else {linear(h,x,aa[aj[r]],s,as[r]);x={s,h};}
        }
        if(bs[r]){
            if(lazy)y=sum(y,bb[bj[r]],bs[r]);
            else {linear(h,y,bb[bj[r]],t,bs[r]);y={t,h};}
        }
        rec(h,x,y,p,h,next);
        if(r==0){
            for(int i=0;i<h;++i){
                memcpy(cc[0]+(size_t)i*stride,p+(size_t)i*h,(size_t)h*sizeof(float));
                memcpy(cc[3]+(size_t)i*stride,p+(size_t)i*h,(size_t)h*sizeof(float));
                memset(cc[1]+(size_t)i*stride,0,(size_t)h*sizeof(float));
                memset(cc[2]+(size_t)i*stride,0,(size_t)h*sizeof(float));
            }
        }
        if(r==1){add(h,cc[2],stride,p,1);add(h,cc[3],stride,p,-1);}
        if(r==2){add(h,cc[1],stride,p,1);add(h,cc[3],stride,p,1);}
        if(r==3){add(h,cc[0],stride,p,1);add(h,cc[2],stride,p,1);}
        if(r==4){add(h,cc[0],stride,p,-1);add(h,cc[1],stride,p,1);}
        if(r==5)add(h,cc[3],stride,p,1);
        if(r==6)add(h,cc[0],stride,p,1);
    }
}
}
void matrix_multiply(int n,const float*A,const float*B,float*C){
    if(n<=512||(n&(n-1))){mat::classical(n,{A,n},{B,n},C,n);return;}
    float*work=nullptr;
    if(posix_memalign((void**)&work,4096,(size_t)n*n*sizeof(float)/2))abort();
    mat::rec(n,{A,n},{B,n},C,n,work,true);
    free(work);
}

评测结果