代码

// 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=4,NR=16,MC=128,NC=128,KC=256;
alignas(4096) static float ap[MC*KC],bp[NC*KC];
static inline void micro(const float*a,const float*b,float*c,int stride){
    __m256 c00=_mm256_loadu_ps(c),c01=_mm256_loadu_ps(c+8);
    __m256 c10=_mm256_loadu_ps(c+stride),c11=_mm256_loadu_ps(c+stride+8);
    __m256 c20=_mm256_loadu_ps(c+2*stride),c21=_mm256_loadu_ps(c+2*stride+8);
    __m256 c30=_mm256_loadu_ps(c+3*stride),c31=_mm256_loadu_ps(c+3*stride+8);
#pragma GCC unroll 8
    for(int k=0;k<KC;++k){
        __m256 b0=_mm256_loadu_ps(b),b1=_mm256_loadu_ps(b+8);
        __m256 a0=_mm256_broadcast_ss(a);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);
        a+=MR;b+=NR;
    }
    _mm256_storeu_ps(c,c00);_mm256_storeu_ps(c+8,c01);
    _mm256_storeu_ps(c+stride,c10);_mm256_storeu_ps(c+stride+8,c11);
    _mm256_storeu_ps(c+2*stride,c20);_mm256_storeu_ps(c+2*stride+8,c21);
    _mm256_storeu_ps(c+3*stride,c30);_mm256_storeu_ps(c+3*stride+8,c31);
}
static void classical(int n,View a,View b,float*c,int stride){
    for(int i=0;i<n;++i)memset(c+(size_t)i*stride,0,(size_t)n*sizeof(float));
    if(n%MC){
        for(int i=0;i<n;++i)for(int k=0;k<n;++k)for(int j=0;j<n;++j)
            c[(size_t)i*stride+j]+=get(a,i,k)*get(b,k,j);
        return;
    }
    for(int ic=0;ic<n;ic+=MC)for(int pc=0;pc<n;pc+=KC){
        for(int r=0;r<MC;r+=MR)for(int k=0;k<KC;++k)for(int i=0;i<MR;++i)
            ap[(r/MR*KC+k)*MR+i]=get(a,ic+r+i,pc+k);
        for(int jc=0;jc<n;jc+=NC){
            for(int q=0;q<NC;q+=NR)for(int k=0;k<KC;++k)for(int j=0;j<NR;++j)
                bp[(q/NR*KC+k)*NR+j]=get(b,pc+k,jc+q+j);
            for(int r=0;r<MC;r+=MR)for(int q=0;q<NC;q+=NR)
                micro(ap+r*KC,bp+q*KC,c+(size_t)(ic+r)*stride+jc+q,stride);
        }
    }
}
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);
}

评测结果