// References:
// [1] saffah_cc_v41_agg1, https://duck.ac/submission/110267 : copied
// the accepted no-A-pack Strassen/AVX2 engine, including the paired
// leaf kernel. No separate license notice was displayed.
// [2] saffah_codex_6s_agg2, https://duck.ac/submission/110143 : original
// paired leaf idea already cited in [1].
// Approach: Quantize each input at 40 levels and use the faster paired
// kernel on the two direct-A/sum-B products as well as the two established
// sum-A/direct-B products. Test for possible signed 16-bit quartet overflow
// and output error before relying on this precision tradeoff.
// Purpose: Experimental official speed and correctness validation for q=40.
#pragma GCC optimize("O3,unroll-loops,no-strict-aliasing,web,ira-loop-pressure")
#pragma GCC target("avx2,fma")
#include <immintrin.h>
#include <stdint.h>
#include <stddef.h>
static constexpr int MR=4, MC=208, NR=16, QA=40, QB=40;
alignas(4096) static int8_t aq[2*2048*2048], bq[2*2048*2048];
alignas(4096) static unsigned char workspace[184*1024*1024];
alignas(64) static uint8_t bp[1024*1024]; /* m2k_: the A panel `ap` is gone */
alignas(64) static int bsum[1024];
struct Mat { const int8_t *x; int ld; const int8_t *y=nullptr; int ly=0; int sign=0; };
// 取 16 个 int8；x±y 都在 int8 域完成（各源 |v|<=46 => 和 |v|<=92 <=127，不溢出）
static inline __m128i ld(Mat a,int i,int j) {
__m128i x=_mm_loadu_si128((const __m128i*)(a.x+i*a.ld+j));
if(a.sign) {__m128i y=_mm_loadu_si128((const __m128i*)(a.y+i*a.ly+j));
x=a.sign>0?_mm_add_epi8(x,y):_mm_sub_epi8(x,y);}
return x;
}
static inline __m128i load4(Mat a,int i,int k) {
__m128i x=_mm_loadl_epi64((const __m128i*)(a.x+i*a.ld+k));
if(a.sign) {
__m128i y=_mm_loadl_epi64((const __m128i*)(a.y+i*a.ly+k));
x=a.sign>0?_mm_add_epi16(x,y):_mm_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=28672;
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, %%ymm8\n\t"
"vpmaddubsw %%ymm12, %%ymm13, %%ymm9\n\t"
"vpbroadcastd 1024(%[a]), %%ymm13\n\t"
"vpmaddubsw %%ymm11, %%ymm13, %%ymm14\n\t"
"vpmaddubsw %%ymm12, %%ymm13, %%ymm15\n\t"
"vpmaddwd %%ymm10, %%ymm8, %%ymm8\n\t"
"vpmaddwd %%ymm10, %%ymm9, %%ymm9\n\t"
"vpmaddwd %%ymm10, %%ymm14, %%ymm14\n\t"
"vpmaddwd %%ymm10, %%ymm15, %%ymm15\n\t"
"vpaddd %%ymm8, %%ymm0, %%ymm0\n\t"
"vpaddd %%ymm9, %%ymm1, %%ymm1\n\t"
"vpaddd %%ymm14, %%ymm2, %%ymm2\n\t"
"vpaddd %%ymm15, %%ymm3, %%ymm3\n\t"
"vpbroadcastd 2048(%[a]), %%ymm13\n\t"
"vpmaddubsw %%ymm11, %%ymm13, %%ymm8\n\t"
"vpmaddubsw %%ymm12, %%ymm13, %%ymm9\n\t"
"vpbroadcastd 3072(%[a]), %%ymm13\n\t"
"vpmaddubsw %%ymm11, %%ymm13, %%ymm14\n\t"
"vpmaddubsw %%ymm12, %%ymm13, %%ymm15\n\t"
"vpmaddwd %%ymm10, %%ymm8, %%ymm8\n\t"
"vpmaddwd %%ymm10, %%ymm9, %%ymm9\n\t"
"vpmaddwd %%ymm10, %%ymm14, %%ymm14\n\t"
"vpmaddwd %%ymm10, %%ymm15, %%ymm15\n\t"
"vpaddd %%ymm8, %%ymm4, %%ymm4\n\t"
"vpaddd %%ymm9, %%ymm5, %%ymm5\n\t"
"vpaddd %%ymm14, %%ymm6, %%ymm6\n\t"
"vpaddd %%ymm15, %%ymm7, %%ymm7\n\t"
"vmovdqa 64(%[b]), %%ymm11\n\t"
"vmovdqa 96(%[b]), %%ymm12\n\t"
"vpbroadcastd 4(%[a]), %%ymm13\n\t"
"vpmaddubsw %%ymm11, %%ymm13, %%ymm8\n\t"
"vpmaddubsw %%ymm12, %%ymm13, %%ymm9\n\t"
"vpbroadcastd 1028(%[a]), %%ymm13\n\t"
"vpmaddubsw %%ymm11, %%ymm13, %%ymm14\n\t"
"vpmaddubsw %%ymm12, %%ymm13, %%ymm15\n\t"
"vpmaddwd %%ymm10, %%ymm8, %%ymm8\n\t"
"vpmaddwd %%ymm10, %%ymm9, %%ymm9\n\t"
"vpmaddwd %%ymm10, %%ymm14, %%ymm14\n\t"
"vpmaddwd %%ymm10, %%ymm15, %%ymm15\n\t"
"vpaddd %%ymm8, %%ymm0, %%ymm0\n\t"
"vpaddd %%ymm9, %%ymm1, %%ymm1\n\t"
"vpaddd %%ymm14, %%ymm2, %%ymm2\n\t"
"vpaddd %%ymm15, %%ymm3, %%ymm3\n\t"
"vpbroadcastd 2052(%[a]), %%ymm13\n\t"
"vpmaddubsw %%ymm11, %%ymm13, %%ymm8\n\t"
"vpmaddubsw %%ymm12, %%ymm13, %%ymm9\n\t"
"vpbroadcastd 3076(%[a]), %%ymm13\n\t"
"vpmaddubsw %%ymm11, %%ymm13, %%ymm14\n\t"
"vpmaddubsw %%ymm12, %%ymm13, %%ymm15\n\t"
"vpmaddwd %%ymm10, %%ymm8, %%ymm8\n\t"
"vpmaddwd %%ymm10, %%ymm9, %%ymm9\n\t"
"vpmaddwd %%ymm10, %%ymm14, %%ymm14\n\t"
"vpmaddwd %%ymm10, %%ymm15, %%ymm15\n\t"
"vpaddd %%ymm8, %%ymm4, %%ymm4\n\t"
"vpaddd %%ymm9, %%ymm5, %%ymm5\n\t"
"vpaddd %%ymm14, %%ymm6, %%ymm6\n\t"
"vpaddd %%ymm15, %%ymm7, %%ymm7\n\t"
"vmovdqa 128(%[b]), %%ymm11\n\t"
"vmovdqa 160(%[b]), %%ymm12\n\t"
"vpbroadcastd 8(%[a]), %%ymm13\n\t"
"vpmaddubsw %%ymm11, %%ymm13, %%ymm8\n\t"
"vpmaddubsw %%ymm12, %%ymm13, %%ymm9\n\t"
"vpbroadcastd 1032(%[a]), %%ymm13\n\t"
"vpmaddubsw %%ymm11, %%ymm13, %%ymm14\n\t"
"vpmaddubsw %%ymm12, %%ymm13, %%ymm15\n\t"
"vpmaddwd %%ymm10, %%ymm8, %%ymm8\n\t"
"vpmaddwd %%ymm10, %%ymm9, %%ymm9\n\t"
"vpmaddwd %%ymm10, %%ymm14, %%ymm14\n\t"
"vpmaddwd %%ymm10, %%ymm15, %%ymm15\n\t"
"vpaddd %%ymm8, %%ymm0, %%ymm0\n\t"
"vpaddd %%ymm9, %%ymm1, %%ymm1\n\t"
"vpaddd %%ymm14, %%ymm2, %%ymm2\n\t"
"vpaddd %%ymm15, %%ymm3, %%ymm3\n\t"
"vpbroadcastd 2056(%[a]), %%ymm13\n\t"
"vpmaddubsw %%ymm11, %%ymm13, %%ymm8\n\t"
"vpmaddubsw %%ymm12, %%ymm13, %%ymm9\n\t"
"vpbroadcastd 3080(%[a]), %%ymm13\n\t"
"vpmaddubsw %%ymm11, %%ymm13, %%ymm14\n\t"
"vpmaddubsw %%ymm12, %%ymm13, %%ymm15\n\t"
"vpmaddwd %%ymm10, %%ymm8, %%ymm8\n\t"
"vpmaddwd %%ymm10, %%ymm9, %%ymm9\n\t"
"vpmaddwd %%ymm10, %%ymm14, %%ymm14\n\t"
"vpmaddwd %%ymm10, %%ymm15, %%ymm15\n\t"
"vpaddd %%ymm8, %%ymm4, %%ymm4\n\t"
"vpaddd %%ymm9, %%ymm5, %%ymm5\n\t"
"vpaddd %%ymm14, %%ymm6, %%ymm6\n\t"
"vpaddd %%ymm15, %%ymm7, %%ymm7\n\t"
"vmovdqa 192(%[b]), %%ymm11\n\t"
"vmovdqa 224(%[b]), %%ymm12\n\t"
"vpbroadcastd 12(%[a]), %%ymm13\n\t"
"vpmaddubsw %%ymm11, %%ymm13, %%ymm8\n\t"
"vpmaddubsw %%ymm12, %%ymm13, %%ymm9\n\t"
"vpbroadcastd 1036(%[a]), %%ymm13\n\t"
"vpmaddubsw %%ymm11, %%ymm13, %%ymm14\n\t"
"vpmaddubsw %%ymm12, %%ymm13, %%ymm15\n\t"
"vpmaddwd %%ymm10, %%ymm8, %%ymm8\n\t"
"vpmaddwd %%ymm10, %%ymm9, %%ymm9\n\t"
"vpmaddwd %%ymm10, %%ymm14, %%ymm14\n\t"
"vpmaddwd %%ymm10, %%ymm15, %%ymm15\n\t"
"vpaddd %%ymm8, %%ymm0, %%ymm0\n\t"
"vpaddd %%ymm9, %%ymm1, %%ymm1\n\t"
"vpaddd %%ymm14, %%ymm2, %%ymm2\n\t"
"vpaddd %%ymm15, %%ymm3, %%ymm3\n\t"
"vpbroadcastd 2060(%[a]), %%ymm13\n\t"
"vpmaddubsw %%ymm11, %%ymm13, %%ymm8\n\t"
"vpmaddubsw %%ymm12, %%ymm13, %%ymm9\n\t"
"vpbroadcastd 3084(%[a]), %%ymm13\n\t"
"vpmaddubsw %%ymm11, %%ymm13, %%ymm14\n\t"
"vpmaddubsw %%ymm12, %%ymm13, %%ymm15\n\t"
"vpmaddwd %%ymm10, %%ymm8, %%ymm8\n\t"
"vpmaddwd %%ymm10, %%ymm9, %%ymm9\n\t"
"vpmaddwd %%ymm10, %%ymm14, %%ymm14\n\t"
"vpmaddwd %%ymm10, %%ymm15, %%ymm15\n\t"
"vpaddd %%ymm8, %%ymm4, %%ymm4\n\t"
"vpaddd %%ymm9, %%ymm5, %%ymm5\n\t"
"vpaddd %%ymm14, %%ymm6, %%ymm6\n\t"
"vpaddd %%ymm15, %%ymm7, %%ymm7\n\t"
"add $16, %[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 __attribute__((noinline)) void kernel_pair2(const uint8_t *a,const uint8_t *b,uint32_t *out,const int *sums) {
ptrdiff_t stride=28672;
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"
"vpcmpeqw %%ymm15, %%ymm15, %%ymm15\n\t"
"vpsrlw $15, %%ymm15, %%ymm15\n\t"
".p2align 5\n\t"
"1:\n\t"
"vmovdqa 0(%[b]), %%ymm11\n\t"
"vmovdqa 32(%[b]), %%ymm12\n\t"
"vmovdqa 64(%[b]), %%ymm13\n\t"
"vmovdqa 96(%[b]), %%ymm14\n\t"
"vpbroadcastd 0(%[a]), %%ymm10\n\t"
"vpmaddubsw %%ymm11, %%ymm10, %%ymm8\n\t"
"vpmaddubsw %%ymm12, %%ymm10, %%ymm9\n\t"
"vpbroadcastd 4(%[a]), %%ymm10\n\t"
"vpmaddubsw %%ymm13, %%ymm10, %%ymm4\n\t"
"vpmaddubsw %%ymm14, %%ymm10, %%ymm5\n\t"
"vpaddw %%ymm4, %%ymm8, %%ymm8\n\t"
"vpaddw %%ymm5, %%ymm9, %%ymm9\n\t"
"vpmaddwd %%ymm15, %%ymm8, %%ymm8\n\t"
"vpmaddwd %%ymm15, %%ymm9, %%ymm9\n\t"
"vpaddd %%ymm8, %%ymm0, %%ymm0\n\t"
"vpaddd %%ymm9, %%ymm1, %%ymm1\n\t"
"vpbroadcastd 1024(%[a]), %%ymm10\n\t"
"vpmaddubsw %%ymm11, %%ymm10, %%ymm8\n\t"
"vpmaddubsw %%ymm12, %%ymm10, %%ymm9\n\t"
"vpbroadcastd 1028(%[a]), %%ymm10\n\t"
"vpmaddubsw %%ymm13, %%ymm10, %%ymm4\n\t"
"vpmaddubsw %%ymm14, %%ymm10, %%ymm5\n\t"
"vpaddw %%ymm4, %%ymm8, %%ymm8\n\t"
"vpaddw %%ymm5, %%ymm9, %%ymm9\n\t"
"vpmaddwd %%ymm15, %%ymm8, %%ymm8\n\t"
"vpmaddwd %%ymm15, %%ymm9, %%ymm9\n\t"
"vpaddd %%ymm8, %%ymm2, %%ymm2\n\t"
"vpaddd %%ymm9, %%ymm3, %%ymm3\n\t"
"vmovdqa 128(%[b]), %%ymm11\n\t"
"vmovdqa 160(%[b]), %%ymm12\n\t"
"vmovdqa 192(%[b]), %%ymm13\n\t"
"vmovdqa 224(%[b]), %%ymm14\n\t"
"vpbroadcastd 8(%[a]), %%ymm10\n\t"
"vpmaddubsw %%ymm11, %%ymm10, %%ymm8\n\t"
"vpmaddubsw %%ymm12, %%ymm10, %%ymm9\n\t"
"vpbroadcastd 12(%[a]), %%ymm10\n\t"
"vpmaddubsw %%ymm13, %%ymm10, %%ymm4\n\t"
"vpmaddubsw %%ymm14, %%ymm10, %%ymm5\n\t"
"vpaddw %%ymm4, %%ymm8, %%ymm8\n\t"
"vpaddw %%ymm5, %%ymm9, %%ymm9\n\t"
"vpmaddwd %%ymm15, %%ymm8, %%ymm8\n\t"
"vpmaddwd %%ymm15, %%ymm9, %%ymm9\n\t"
"vpaddd %%ymm8, %%ymm0, %%ymm0\n\t"
"vpaddd %%ymm9, %%ymm1, %%ymm1\n\t"
"vpbroadcastd 1032(%[a]), %%ymm10\n\t"
"vpmaddubsw %%ymm11, %%ymm10, %%ymm8\n\t"
"vpmaddubsw %%ymm12, %%ymm10, %%ymm9\n\t"
"vpbroadcastd 1036(%[a]), %%ymm10\n\t"
"vpmaddubsw %%ymm13, %%ymm10, %%ymm4\n\t"
"vpmaddubsw %%ymm14, %%ymm10, %%ymm5\n\t"
"vpaddw %%ymm4, %%ymm8, %%ymm8\n\t"
"vpaddw %%ymm5, %%ymm9, %%ymm9\n\t"
"vpmaddwd %%ymm15, %%ymm8, %%ymm8\n\t"
"vpmaddwd %%ymm15, %%ymm9, %%ymm9\n\t"
"vpaddd %%ymm8, %%ymm2, %%ymm2\n\t"
"vpaddd %%ymm9, %%ymm3, %%ymm3\n\t"
"add $16, %[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"
: [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","ymm8","ymm9","ymm10","ymm11","ymm12","ymm13","ymm14","ymm15");
}
static void leaf(Mat a,Mat b,uint32_t *c,int lr) {
constexpr int n=1024,ldc=1024;
{
constexpr int pc=0,kc=1024;
const __m128i ones=_mm_set1_epi8(1);
const __m256i bscale=_mm256_set1_epi32(80);
for(int jb=0;jb<n;jb+=32) {
__m256i s00=_mm256_setzero_si256(),s01=_mm256_setzero_si256();
__m256i s10=_mm256_setzero_si256(),s11=_mm256_setzero_si256();
for(int k=0;k<kc;k+=4) {
{
int j=jb+0;
__m128i 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);
__m128i t0=_mm_unpacklo_epi8(v0,v1),t1=_mm_unpackhi_epi8(v0,v1);
__m128i t2=_mm_unpacklo_epi8(v2,v3),t3=_mm_unpackhi_epi8(v2,v3);
__m128i wa=_mm_add_epi16(_mm_maddubs_epi16(ones,t0),_mm_maddubs_epi16(ones,t2));
__m128i wb=_mm_add_epi16(_mm_maddubs_epi16(ones,t1),_mm_maddubs_epi16(ones,t3));
s00=_mm256_add_epi32(s00,_mm256_cvtepi16_epi32(wa));
s01=_mm256_add_epi32(s01,_mm256_cvtepi16_epi32(wb));
uint8_t *q=bp+j*kc+k*16;
_mm_store_si128((__m128i*)(q), _mm_unpacklo_epi16(t0,t2));
_mm_store_si128((__m128i*)(q+16),_mm_unpackhi_epi16(t0,t2));
_mm_store_si128((__m128i*)(q+32),_mm_unpacklo_epi16(t1,t3));
_mm_store_si128((__m128i*)(q+48),_mm_unpackhi_epi16(t1,t3));
}
{
int j=jb+16;
__m128i 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);
__m128i t0=_mm_unpacklo_epi8(v0,v1),t1=_mm_unpackhi_epi8(v0,v1);
__m128i t2=_mm_unpacklo_epi8(v2,v3),t3=_mm_unpackhi_epi8(v2,v3);
__m128i wa=_mm_add_epi16(_mm_maddubs_epi16(ones,t0),_mm_maddubs_epi16(ones,t2));
__m128i wb=_mm_add_epi16(_mm_maddubs_epi16(ones,t1),_mm_maddubs_epi16(ones,t3));
s10=_mm256_add_epi32(s10,_mm256_cvtepi16_epi32(wa));
s11=_mm256_add_epi32(s11,_mm256_cvtepi16_epi32(wb));
uint8_t *q=bp+j*kc+k*16;
_mm_store_si128((__m128i*)(q), _mm_unpacklo_epi16(t0,t2));
_mm_store_si128((__m128i*)(q+16),_mm_unpackhi_epi16(t0,t2));
_mm_store_si128((__m128i*)(q+32),_mm_unpacklo_epi16(t1,t3));
_mm_store_si128((__m128i*)(q+48),_mm_unpackhi_epi16(t1,t3));
}
}
_mm256_store_si256((__m256i*)(bsum+jb),_mm256_mullo_epi32(s00,bscale));
_mm256_store_si256((__m256i*)(bsum+jb+8),_mm256_mullo_epi32(s01,bscale));
_mm256_store_si256((__m256i*)(bsum+jb+16),_mm256_mullo_epi32(s10,bscale));
_mm256_store_si256((__m256i*)(bsum+jb+24),_mm256_mullo_epi32(s11,bscale));
}
_mm_sfence();
for(int ic=0;ic<n;ic+=MC) {
int mc=n-ic<MC?n-ic:MC;
/* m2k_: packA and the `ap` panel are gone. `aq` already holds the biased
bytes (quant_sums<false>), so each kernel reads its rows straight out of aq:
row i at +1024*i, k-group g at +4*g -- exactly the displacements above. */
for(int j=0;j<n;j+=NR) for(int r=0;r<mc;r+=MR) {
uint32_t *out=c+(((size_t)((ic+r)*64+(j>>4))*7+(size_t)lr)*16);
const uint8_t *pa=(const uint8_t*)(a.x+(size_t)(ic+r)*a.ld);
if(lr==1 || lr==2 || lr==3 || lr==4) {
kernel_pair2(pa,bp+j*kc,out,bsum+j);
kernel_pair2(pa+2048,bp+j*kc,out+2*7168,bsum+j);
} else kernel(pa,bp+j*kc,out,bsum+j);
}
}
}
}
static void combine(int h,uint32_t *c,int ldc,uint32_t *m) {
size_t z=(size_t)h*h+64;
for(int i=0;i<h;i++)for(int j=0;j<h;j+=8) {
size_t t=((size_t)(i*64+(j>>4))*7)*16+(size_t)(j&15);
__m256i m1=_mm256_load_si256((__m256i*)(m+t)),m2=_mm256_load_si256((__m256i*)(m+t+16)),
m3=_mm256_load_si256((__m256i*)(m+t+32)),m4=_mm256_load_si256((__m256i*)(m+t+48)),
m5=_mm256_load_si256((__m256i*)(m+t+64)),m6=_mm256_load_si256((__m256i*)(m+t+80)),m7=_mm256_load_si256((__m256i*)(m+t+96));
__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);
/* fused: convert and store as float here, so no separate int32->float pass is needed */
{
__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));
}
}
}
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+64;
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);
}
_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);
}
/* m2k_: +B on every quantised value, in the int16 domain (1 uop / 16 values; compiled away
when B==0). B=45 on the A call means the *sum* of two operands carries +90, i.e. exactly
the bias the mikro needs on its unsigned byte operand -- so aq can feed the kernel as-is. */
template<int B> static inline __m256i quant16b(const float *a,__m256 scale) {
__m256i v=quant16(a,scale);
if (B) v=_mm256_add_epi16(v,_mm256_set1_epi16((short)B));
return v;
}
template<bool ISB> static void quant_sums(const float *a,int8_t *dst,int8_t *diag,int q) {
constexpr int h=1024,n=2048;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+=32) {
__m256i x0a=quant16b<ISB?0:40>(a+i*n+j,scale),x0b=quant16b<ISB?0:40>(a+i*n+j+16,scale);
__m256i x1a=quant16b<ISB?0:40>(a+i*n+j+h,scale),x1b=quant16b<ISB?0:40>(a+i*n+j+h+16,scale);
__m256i x2a=quant16b<ISB?0:40>(a+(i+h)*n+j,scale),x2b=quant16b<ISB?0:40>(a+(i+h)*n+j+16,scale);
__m256i x3a=quant16b<ISB?0:40>(a+(i+h)*n+j+h,scale),x3b=quant16b<ISB?0:40>(a+(i+h)*n+j+h+16,scale);
size_t t=(size_t)i*h+j;
__m256i v0a=_mm256_add_epi16(x0a,x3a),v0b=_mm256_add_epi16(x0b,x3b);
__m256i v1a=ISB?_mm256_sub_epi16(x1a,x3a):_mm256_add_epi16(x2a,x3a);
__m256i v1b=ISB?_mm256_sub_epi16(x1b,x3b):_mm256_add_epi16(x2b,x3b);
__m256i v2a=ISB?_mm256_sub_epi16(x2a,x0a):_mm256_add_epi16(x0a,x1a);
__m256i v2b=ISB?_mm256_sub_epi16(x2b,x0b):_mm256_add_epi16(x0b,x1b);
__m256i v3a=ISB?_mm256_add_epi16(x0a,x1a):_mm256_sub_epi16(x2a,x0a);
__m256i v3b=ISB?_mm256_add_epi16(x0b,x1b):_mm256_sub_epi16(x2b,x0b);
__m256i v4a=ISB?_mm256_add_epi16(x2a,x3a):_mm256_sub_epi16(x1a,x3a);
__m256i v4b=ISB?_mm256_add_epi16(x2b,x3b):_mm256_sub_epi16(x1b,x3b);
/* m2k_: A side packs UNSIGNED and tops up where the +90 was lost: the two
subtractive combinations ((x2+45)-(x0+45)) and the two diagonal outputs (x+45,
which the mikro sees raw and therefore needs at +90). B side unchanged. */
const __m256i b90=_mm256_set1_epi16(80), b45=_mm256_set1_epi16(40);
#define PK(A,B) _mm256_permute4x64_epi64(ISB?_mm256_packs_epi16(A,B):_mm256_packus_epi16(A,B),0xd8)
#define PK9(A,B) _mm256_permute4x64_epi64(ISB?_mm256_packs_epi16(A,B):_mm256_packus_epi16(_mm256_add_epi16(A,b90),_mm256_add_epi16(B,b90)),0xd8)
#define PK45(A,B) _mm256_permute4x64_epi64(ISB?_mm256_packs_epi16(A,B):_mm256_packus_epi16(_mm256_add_epi16(A,b45),_mm256_add_epi16(B,b45)),0xd8)
_mm256_stream_si256((__m256i*)(dst+t), PK(v0a,v0b));
_mm256_stream_si256((__m256i*)(dst+z+t), PK(v1a,v1b));
_mm256_stream_si256((__m256i*)(dst+2*z+t), PK(v2a,v2b));
_mm256_stream_si256((__m256i*)(dst+3*z+t), PK9(v3a,v3b));
_mm256_stream_si256((__m256i*)(dst+4*z+t), PK9(v4a,v4b));
_mm256_stream_si256((__m256i*)(diag+t), PK45(x0a,x0b));
_mm256_stream_si256((__m256i*)(diag+z+t), PK45(x3a,x3b));
}
}
static void quant_all(const float *src,int8_t *dst,int scale) {
const __m256 q=_mm256_set1_ps((float)scale);
for(size_t i=0;i<(size_t)2048*2048;i+=32) {
__m256i x=quant16(src+i,q),y=quant16(src+i+16,q);
_mm256_store_si256((__m256i*)(dst+i),
_mm256_permute4x64_epi64(_mm256_packs_epi16(x,y),0xd8));
}
}
void matrix_multiply(int n,const float *a,const float *b,float *c) {
(void)n;
constexpr int h=1024;
constexpr size_t z=(size_t)h*h;
quant_sums<false>(a,aq,aq+5*z,QA);
quant_sums<true>(b,bq,bq+5*z,QB);
_mm_sfence();
static const int ai[7]={0,1,5,6,2,3,4};
static const int bi[7]={0,5,1,2,6,3,4};
uint32_t *m=(uint32_t*)workspace;
for(int r=0;r<7;r++) {
Mat x{aq+(size_t)ai[r]*z,h},y{bq+(size_t)bi[r]*z,h};
leaf(x,y,m,r);
}
_mm_sfence();
combine(h,(uint32_t*)c,2048,m);
_mm_sfence();
}
| Compilation | N/A | N/A | Compile OK | Score: N/A | 显示更多 |
| Testcase #1 | 77.911 ms | 59 MB + 16 KB | Accepted | Score: 100 | 显示更多 |