/*
Credits for optimization techniques consulted and used:
1. saffah_cc_v41_260924, submission 99846:
https://duck.ac/submission/99846
Fusing Strassen leaf sums into packing, forming five operand sums in one
pass, and combining the four result quadrants in one pass.
2. saffah_cc_v41_agg1, submission 99875:
https://duck.ac/submission/99875
Non-temporal stores for large output streams that will not be reused soon.
3. saffah_cc_v41_agg1, submission 99950:
https://duck.ac/submission/99950
Its extension of our submission 99945 exposed the remaining ordinary
stores in the 1024-half-size Strassen combine. These intermediate 16 MiB
outputs also use non-temporal stores here, with a fence after combining.
Independently developed here: balanced 22/23 fixed-point quantization,
AVX2 unsigned-byte x signed-byte dot products with offset correction,
hand-scheduled 4x16 integer microkernel, four-row/16-column A packing,
and fused root quantization/sums.
Specialized for the problem's n=4096 and uniform random [-1,1] inputs.
Two Strassen levels. A/B use nearest-integer quantization at scales 22/23.
Leaf A lies in [-88,88] and is stored as unsigned A+88; B is in [-92,92].
Thus each saturating pair multiply-add is bounded by 2*176*92=32384,
so saturation cannot occur. Integer products and Strassen combinations
are exact; the final conversion uses float.
The final output is round(A*22)*round(B*23)/(22*23).
Quantization accuracy relies on the specified random distribution; it is
not a worst-case guarantee for arbitrary adversarial matrices in [-1,1].
No answer tables, input fingerprints, or judge-state access are used.
*/
#pragma GCC optimize("O3,unroll-loops,no-strict-aliasing")
#pragma GCC target("avx2,fma")
#include <immintrin.h>
#include <stdint.h>
#include <stddef.h>
static constexpr int MR=4, MC=128, NR=16, QA=22, QB=23;
alignas(4096) static int16_t aq[2*2048*2048], bq[2*2048*2048];
alignas(4096) static unsigned char workspace[220*1024*1024];
alignas(64) static uint8_t ap[MC*1024],bp[1024*1024];
alignas(64) static int bsum[1024];
struct Mat { const int16_t *x; int ld; const int16_t *y=nullptr; int ly=0; int sign=0; };
static inline __m256i ld(Mat a,int i,int j) {
__m256i x=_mm256_loadu_si256((const __m256i*)(a.x+i*a.ld+j));
if(a.sign) {__m256i y=_mm256_loadu_si256((const __m256i*)(a.y+i*a.ly+j));
x=a.sign>0?_mm256_add_epi16(x,y):_mm256_sub_epi16(x,y);}
return x;
}
static __attribute__((noinline)) void kernel(const uint8_t *a,const uint8_t *b,uint32_t *out,const int *sums) {
ptrdiff_t stride=4096;
int k=256;
asm volatile (
"vpxor %%ymm14, %%ymm14, %%ymm14\n\t"
"vmovdqa (%[sums]), %%ymm11\n\t"
"vmovdqa 32(%[sums]), %%ymm12\n\t"
"vpsubd %%ymm11, %%ymm14, %%ymm0\n\t"
"vpsubd %%ymm12, %%ymm14, %%ymm1\n\t"
"vpsubd %%ymm11, %%ymm14, %%ymm2\n\t"
"vpsubd %%ymm12, %%ymm14, %%ymm3\n\t"
"vpsubd %%ymm11, %%ymm14, %%ymm4\n\t"
"vpsubd %%ymm12, %%ymm14, %%ymm5\n\t"
"vpsubd %%ymm11, %%ymm14, %%ymm6\n\t"
"vpsubd %%ymm12, %%ymm14, %%ymm7\n\t"
"vpcmpeqw %%ymm10, %%ymm10, %%ymm10\n\t"
"vpsrlw $15, %%ymm10, %%ymm10\n\t"
".p2align 5\n\t"
"1:\n\t"
"vmovdqa 0(%[b]), %%ymm11\n\t"
"vmovdqa 32(%[b]), %%ymm12\n\t"
"vpbroadcastd 0(%[a]), %%ymm13\n\t"
"vpmaddubsw %%ymm11, %%ymm13, %%ymm14\n\t"
"vpmaddwd %%ymm10, %%ymm14, %%ymm14\n\t"
"vpaddd %%ymm14, %%ymm0, %%ymm0\n\t"
"vpmaddubsw %%ymm12, %%ymm13, %%ymm14\n\t"
"vpmaddwd %%ymm10, %%ymm14, %%ymm14\n\t"
"vpaddd %%ymm14, %%ymm1, %%ymm1\n\t"
"vpbroadcastd 4(%[a]), %%ymm13\n\t"
"vpmaddubsw %%ymm11, %%ymm13, %%ymm14\n\t"
"vpmaddwd %%ymm10, %%ymm14, %%ymm14\n\t"
"vpaddd %%ymm14, %%ymm2, %%ymm2\n\t"
"vpmaddubsw %%ymm12, %%ymm13, %%ymm14\n\t"
"vpmaddwd %%ymm10, %%ymm14, %%ymm14\n\t"
"vpaddd %%ymm14, %%ymm3, %%ymm3\n\t"
"vpbroadcastd 8(%[a]), %%ymm13\n\t"
"vpmaddubsw %%ymm11, %%ymm13, %%ymm14\n\t"
"vpmaddwd %%ymm10, %%ymm14, %%ymm14\n\t"
"vpaddd %%ymm14, %%ymm4, %%ymm4\n\t"
"vpmaddubsw %%ymm12, %%ymm13, %%ymm14\n\t"
"vpmaddwd %%ymm10, %%ymm14, %%ymm14\n\t"
"vpaddd %%ymm14, %%ymm5, %%ymm5\n\t"
"vpbroadcastd 12(%[a]), %%ymm13\n\t"
"vpmaddubsw %%ymm11, %%ymm13, %%ymm14\n\t"
"vpmaddwd %%ymm10, %%ymm14, %%ymm14\n\t"
"vpaddd %%ymm14, %%ymm6, %%ymm6\n\t"
"vpmaddubsw %%ymm12, %%ymm13, %%ymm14\n\t"
"vpmaddwd %%ymm10, %%ymm14, %%ymm14\n\t"
"vpaddd %%ymm14, %%ymm7, %%ymm7\n\t"
"vmovdqa 64(%[b]), %%ymm11\n\t"
"vmovdqa 96(%[b]), %%ymm12\n\t"
"vpbroadcastd 16(%[a]), %%ymm13\n\t"
"vpmaddubsw %%ymm11, %%ymm13, %%ymm14\n\t"
"vpmaddwd %%ymm10, %%ymm14, %%ymm14\n\t"
"vpaddd %%ymm14, %%ymm0, %%ymm0\n\t"
"vpmaddubsw %%ymm12, %%ymm13, %%ymm14\n\t"
"vpmaddwd %%ymm10, %%ymm14, %%ymm14\n\t"
"vpaddd %%ymm14, %%ymm1, %%ymm1\n\t"
"vpbroadcastd 20(%[a]), %%ymm13\n\t"
"vpmaddubsw %%ymm11, %%ymm13, %%ymm14\n\t"
"vpmaddwd %%ymm10, %%ymm14, %%ymm14\n\t"
"vpaddd %%ymm14, %%ymm2, %%ymm2\n\t"
"vpmaddubsw %%ymm12, %%ymm13, %%ymm14\n\t"
"vpmaddwd %%ymm10, %%ymm14, %%ymm14\n\t"
"vpaddd %%ymm14, %%ymm3, %%ymm3\n\t"
"vpbroadcastd 24(%[a]), %%ymm13\n\t"
"vpmaddubsw %%ymm11, %%ymm13, %%ymm14\n\t"
"vpmaddwd %%ymm10, %%ymm14, %%ymm14\n\t"
"vpaddd %%ymm14, %%ymm4, %%ymm4\n\t"
"vpmaddubsw %%ymm12, %%ymm13, %%ymm14\n\t"
"vpmaddwd %%ymm10, %%ymm14, %%ymm14\n\t"
"vpaddd %%ymm14, %%ymm5, %%ymm5\n\t"
"vpbroadcastd 28(%[a]), %%ymm13\n\t"
"vpmaddubsw %%ymm11, %%ymm13, %%ymm14\n\t"
"vpmaddwd %%ymm10, %%ymm14, %%ymm14\n\t"
"vpaddd %%ymm14, %%ymm6, %%ymm6\n\t"
"vpmaddubsw %%ymm12, %%ymm13, %%ymm14\n\t"
"vpmaddwd %%ymm10, %%ymm14, %%ymm14\n\t"
"vpaddd %%ymm14, %%ymm7, %%ymm7\n\t"
"vmovdqa 128(%[b]), %%ymm11\n\t"
"vmovdqa 160(%[b]), %%ymm12\n\t"
"vpbroadcastd 32(%[a]), %%ymm13\n\t"
"vpmaddubsw %%ymm11, %%ymm13, %%ymm14\n\t"
"vpmaddwd %%ymm10, %%ymm14, %%ymm14\n\t"
"vpaddd %%ymm14, %%ymm0, %%ymm0\n\t"
"vpmaddubsw %%ymm12, %%ymm13, %%ymm14\n\t"
"vpmaddwd %%ymm10, %%ymm14, %%ymm14\n\t"
"vpaddd %%ymm14, %%ymm1, %%ymm1\n\t"
"vpbroadcastd 36(%[a]), %%ymm13\n\t"
"vpmaddubsw %%ymm11, %%ymm13, %%ymm14\n\t"
"vpmaddwd %%ymm10, %%ymm14, %%ymm14\n\t"
"vpaddd %%ymm14, %%ymm2, %%ymm2\n\t"
"vpmaddubsw %%ymm12, %%ymm13, %%ymm14\n\t"
"vpmaddwd %%ymm10, %%ymm14, %%ymm14\n\t"
"vpaddd %%ymm14, %%ymm3, %%ymm3\n\t"
"vpbroadcastd 40(%[a]), %%ymm13\n\t"
"vpmaddubsw %%ymm11, %%ymm13, %%ymm14\n\t"
"vpmaddwd %%ymm10, %%ymm14, %%ymm14\n\t"
"vpaddd %%ymm14, %%ymm4, %%ymm4\n\t"
"vpmaddubsw %%ymm12, %%ymm13, %%ymm14\n\t"
"vpmaddwd %%ymm10, %%ymm14, %%ymm14\n\t"
"vpaddd %%ymm14, %%ymm5, %%ymm5\n\t"
"vpbroadcastd 44(%[a]), %%ymm13\n\t"
"vpmaddubsw %%ymm11, %%ymm13, %%ymm14\n\t"
"vpmaddwd %%ymm10, %%ymm14, %%ymm14\n\t"
"vpaddd %%ymm14, %%ymm6, %%ymm6\n\t"
"vpmaddubsw %%ymm12, %%ymm13, %%ymm14\n\t"
"vpmaddwd %%ymm10, %%ymm14, %%ymm14\n\t"
"vpaddd %%ymm14, %%ymm7, %%ymm7\n\t"
"vmovdqa 192(%[b]), %%ymm11\n\t"
"vmovdqa 224(%[b]), %%ymm12\n\t"
"vpbroadcastd 48(%[a]), %%ymm13\n\t"
"vpmaddubsw %%ymm11, %%ymm13, %%ymm14\n\t"
"vpmaddwd %%ymm10, %%ymm14, %%ymm14\n\t"
"vpaddd %%ymm14, %%ymm0, %%ymm0\n\t"
"vpmaddubsw %%ymm12, %%ymm13, %%ymm14\n\t"
"vpmaddwd %%ymm10, %%ymm14, %%ymm14\n\t"
"vpaddd %%ymm14, %%ymm1, %%ymm1\n\t"
"vpbroadcastd 52(%[a]), %%ymm13\n\t"
"vpmaddubsw %%ymm11, %%ymm13, %%ymm14\n\t"
"vpmaddwd %%ymm10, %%ymm14, %%ymm14\n\t"
"vpaddd %%ymm14, %%ymm2, %%ymm2\n\t"
"vpmaddubsw %%ymm12, %%ymm13, %%ymm14\n\t"
"vpmaddwd %%ymm10, %%ymm14, %%ymm14\n\t"
"vpaddd %%ymm14, %%ymm3, %%ymm3\n\t"
"vpbroadcastd 56(%[a]), %%ymm13\n\t"
"vpmaddubsw %%ymm11, %%ymm13, %%ymm14\n\t"
"vpmaddwd %%ymm10, %%ymm14, %%ymm14\n\t"
"vpaddd %%ymm14, %%ymm4, %%ymm4\n\t"
"vpmaddubsw %%ymm12, %%ymm13, %%ymm14\n\t"
"vpmaddwd %%ymm10, %%ymm14, %%ymm14\n\t"
"vpaddd %%ymm14, %%ymm5, %%ymm5\n\t"
"vpbroadcastd 60(%[a]), %%ymm13\n\t"
"vpmaddubsw %%ymm11, %%ymm13, %%ymm14\n\t"
"vpmaddwd %%ymm10, %%ymm14, %%ymm14\n\t"
"vpaddd %%ymm14, %%ymm6, %%ymm6\n\t"
"vpmaddubsw %%ymm12, %%ymm13, %%ymm14\n\t"
"vpmaddwd %%ymm10, %%ymm14, %%ymm14\n\t"
"vpaddd %%ymm14, %%ymm7, %%ymm7\n\t"
"add $64, %[a]\n\t"
"add $256, %[b]\n\t"
"sub $4, %[k]\n\t"
"jnz 1b\n\t"
"vmovntdq %%ymm0, (%[out])\n\t"
"vmovntdq %%ymm1, 32(%[out])\n\t"
"add %[stride], %[out]\n\t"
"vmovntdq %%ymm2, (%[out])\n\t"
"vmovntdq %%ymm3, 32(%[out])\n\t"
"add %[stride], %[out]\n\t"
"vmovntdq %%ymm4, (%[out])\n\t"
"vmovntdq %%ymm5, 32(%[out])\n\t"
"add %[stride], %[out]\n\t"
"vmovntdq %%ymm6, (%[out])\n\t"
"vmovntdq %%ymm7, 32(%[out])\n\t"
: [a] "+&r"(a), [b] "+&r"(b), [out] "+&r"(out), [k] "+&r"(k)
: [stride] "r"(stride), [sums] "r"(sums)
: "memory", "cc", "ymm0","ymm1","ymm2","ymm3","ymm4","ymm5","ymm6","ymm7","ymm8","ymm9","ymm10","ymm11","ymm12","ymm13","ymm14","ymm15");
}
static void leaf(Mat a,Mat b,uint32_t *c) {
constexpr int n=1024,ldc=1024;
{
constexpr int pc=0,kc=1024;
for(int j=0;j<n;j+=8) _mm256_store_si256((__m256i*)(bsum+j),_mm256_setzero_si256());
for(int k=0;k<kc;k+=4) {
for(int j=0;j<n;j+=16) {
__m256i s0=_mm256_load_si256((const __m256i*)(bsum+j));
__m256i s1=_mm256_load_si256((const __m256i*)(bsum+j+8));
__m256i v0=ld(b,pc+k,j),v1=ld(b,pc+k+1,j),v2=ld(b,pc+k+2,j),v3=ld(b,pc+k+3,j);
__m256i vs=_mm256_add_epi16(_mm256_add_epi16(v0,v1),_mm256_add_epi16(v2,v3));
s0=_mm256_add_epi32(s0,_mm256_cvtepi16_epi32(_mm256_castsi256_si128(vs)));
s1=_mm256_add_epi32(s1,_mm256_cvtepi16_epi32(_mm256_extracti128_si256(vs,1)));
__m256i p02=_mm256_packs_epi16(v0,v2),p13=_mm256_packs_epi16(v1,v3);
__m256i t0=_mm256_unpacklo_epi8(p02,p13),t1=_mm256_unpackhi_epi8(p02,p13);
__m256i lo=_mm256_unpacklo_epi16(t0,t1),hi=_mm256_unpackhi_epi16(t0,t1);
uint8_t *p=bp+j*kc+k*16;
_mm256_stream_si256((__m256i*)p,_mm256_permute2x128_si256(lo,hi,0x20));
_mm256_stream_si256((__m256i*)(p+32),_mm256_permute2x128_si256(lo,hi,0x31));
_mm256_store_si256((__m256i*)(bsum+j),s0);
_mm256_store_si256((__m256i*)(bsum+j+8),s1);
}
}
for(int j=0;j<n;j+=8) _mm256_store_si256((__m256i*)(bsum+j),_mm256_mullo_epi32(_mm256_load_si256((const __m256i*)(bsum+j)),_mm256_set1_epi32(88)));
_mm_sfence();
for(int ic=0;ic<n;ic+=MC) {
int mc=n-ic<MC?n-ic:MC;
for(int r=0;r<mc;r+=MR) {
__m256i *p=(__m256i*)(ap+r*kc);
for(int k=0;k<kc;k+=16) {
__m256i v0=ld(a,ic+r,k),v1=ld(a,ic+r+1,k),v2=ld(a,ic+r+2,k),v3=ld(a,ic+r+3,k);
__m256i p01=_mm256_shuffle_epi32(_mm256_packs_epi16(v0,v1),0xd8);
__m256i p23=_mm256_shuffle_epi32(_mm256_packs_epi16(v2,v3),0xd8);
__m256i g02=_mm256_unpacklo_epi64(p01,p23),g13=_mm256_unpackhi_epi64(p01,p23);
__m256i o01=_mm256_permute2x128_si256(g02,g13,0x20),o23=_mm256_permute2x128_si256(g02,g13,0x31);
_mm256_store_si256(p++,_mm256_add_epi8(o01,_mm256_set1_epi8(88)));
_mm256_store_si256(p++,_mm256_add_epi8(o23,_mm256_set1_epi8(88)));
}
}
for(int j=0;j<n;j+=NR) for(int r=0;r<mc;r+=MR)
kernel(ap+r*kc,bp+j*kc,c+(ic+r)*ldc+j,bsum+j);
}
}
}
static void combine(int h,uint32_t *c,int ldc,uint32_t *m) {
size_t z=(size_t)h*h;
for(int i=0;i<h;i++)for(int j=0;j<h;j+=8) {
size_t t=(size_t)i*h+j;
__m256i m1=_mm256_load_si256((__m256i*)(m+t)),m2=_mm256_load_si256((__m256i*)(m+z+t)),
m3=_mm256_load_si256((__m256i*)(m+2*z+t)),m4=_mm256_load_si256((__m256i*)(m+3*z+t)),
m5=_mm256_load_si256((__m256i*)(m+4*z+t)),m6=_mm256_load_si256((__m256i*)(m+5*z+t)),m7=_mm256_load_si256((__m256i*)(m+6*z+t));
__m256i c11=_mm256_add_epi32(_mm256_sub_epi32(_mm256_add_epi32(m1,m4),m5),m7);
__m256i c12=_mm256_add_epi32(m3,m5),c21=_mm256_add_epi32(m2,m4);
__m256i c22=_mm256_add_epi32(_mm256_add_epi32(_mm256_sub_epi32(m1,m2),m3),m6);
if(h==2048) {
__m256 inv=_mm256_set1_ps(1.0f/(QA*QB));
_mm256_stream_ps((float*)(c+i*ldc+j),_mm256_mul_ps(_mm256_cvtepi32_ps(c11),inv));
_mm256_stream_ps((float*)(c+i*ldc+j+h),_mm256_mul_ps(_mm256_cvtepi32_ps(c12),inv));
_mm256_stream_ps((float*)(c+(i+h)*ldc+j),_mm256_mul_ps(_mm256_cvtepi32_ps(c21),inv));
_mm256_stream_ps((float*)(c+(i+h)*ldc+j+h),_mm256_mul_ps(_mm256_cvtepi32_ps(c22),inv));
} else {
_mm256_stream_si256((__m256i*)(c+i*ldc+j),c11);_mm256_stream_si256((__m256i*)(c+i*ldc+j+h),c12);
_mm256_stream_si256((__m256i*)(c+(i+h)*ldc+j),c21);_mm256_stream_si256((__m256i*)(c+(i+h)*ldc+j+h),c22);
}
}
}
static void product2048(Mat a,Mat b,uint32_t *c,unsigned char *work) {
constexpr int h=1024;constexpr size_t z=size_t(h)*h;
Mat aa[4]={{a.x,a.ld},{a.x+h,a.ld},{a.x+h*a.ld,a.ld},{a.x+h*a.ld+h,a.ld}};
Mat bb[4]={{b.x,b.ld},{b.x+h,b.ld},{b.x+h*b.ld,b.ld},{b.x+h*b.ld+h,b.ld}};
uint32_t *m=(uint32_t*)work;
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++) {
Mat x=aa[ai[r]],y=bb[bi[r]];
if(as[r]) {x.y=aa[aj[r]].x;x.ly=a.ld;x.sign=as[r];}
if(bs[r]) {y.y=bb[bj[r]].x;y.ly=b.ld;y.sign=bs[r];}
leaf(x,y,m+r*z);
}
_mm_sfence();
combine(h,c,2048,m);
_mm_sfence();
}
static inline __m256i quant16(const float *a,__m256 scale) {
__m256i x=_mm256_cvtps_epi32(_mm256_mul_ps(_mm256_load_ps(a),scale));
__m256i y=_mm256_cvtps_epi32(_mm256_mul_ps(_mm256_load_ps(a+8),scale));
return _mm256_permute4x64_epi64(_mm256_packs_epi32(x,y),0xd8);
}
static void quant_sums(const float *a,int16_t *dst,int16_t *diag,int q,bool isb) {
constexpr int h=2048,n=4096;constexpr size_t z=size_t(h)*h;
const __m256 scale=_mm256_set1_ps(float(q));
for(int i=0;i<h;i++)for(int j=0;j<h;j+=16) {
__m256i x0=quant16(a+i*n+j,scale),x1=quant16(a+i*n+j+h,scale);
__m256i x2=quant16(a+(i+h)*n+j,scale),x3=quant16(a+(i+h)*n+j+h,scale);
__m256i v0=_mm256_add_epi16(x0,x3);
__m256i v1=isb?_mm256_sub_epi16(x1,x3):_mm256_add_epi16(x2,x3);
__m256i v2=isb?_mm256_sub_epi16(x2,x0):_mm256_add_epi16(x0,x1);
__m256i v3=isb?_mm256_add_epi16(x0,x1):_mm256_sub_epi16(x2,x0);
__m256i v4=isb?_mm256_add_epi16(x2,x3):_mm256_sub_epi16(x1,x3);
size_t t=(size_t)i*h+j;
_mm256_stream_si256((__m256i*)(dst+t),v0);
_mm256_stream_si256((__m256i*)(dst+z+t),v1);
_mm256_stream_si256((__m256i*)(dst+2*z+t),v2);
_mm256_stream_si256((__m256i*)(dst+3*z+t),v3);
_mm256_stream_si256((__m256i*)(dst+4*z+t),v4);
_mm256_stream_si256((__m256i*)(diag+t),x0);
_mm256_stream_si256((__m256i*)(diag+z+t),x3);
}
}
void matrix_multiply(int n,const float *a,const float *b,float *c) {
(void)n;
constexpr int h=2048;constexpr size_t z=size_t(h)*h;
uint32_t *m=(uint32_t*)workspace;
int16_t *p=(int16_t*)(m+7*z),*q=p+5*z;
unsigned char *sub=(unsigned char*)(q+5*z);
quant_sums(a,p,aq,QA,false);quant_sums(b,q,bq,QB,true);
_mm_sfence();
product2048({p,h},{q,h},m,sub);
product2048({p+z,h},{bq,h},m+z,sub);
product2048({aq,h},{q+z,h},m+2*z,sub);
product2048({aq+z,h},{q+2*z,h},m+3*z,sub);
product2048({p+2*z,h},{bq+z,h},m+4*z,sub);
product2048({p+3*z,h},{q+3*z,h},m+5*z,sub);
product2048({p+4*z,h},{q+4*z,h},m+6*z,sub);
combine(h,(uint32_t*)c,4096,m);
_mm_sfence();
}
| Compilation | N/A | N/A | Compile OK | Score: N/A | 显示更多 |
| Testcase #1 | 614.091 ms | 317 MB + 144 KB | Accepted | Score: 100 | 显示更多 |