url/colibri
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Merge pull request #269 from michael-denyer/arm-i8mm-smmla

Browse Source 
ARM i8mm: SMMLA fast path for the int8/int4 IDOT drivers (opt-in via ARCH=native)
This commit is contained in:
Vincenzo
2026-07-15 16:20:32 +02:00
committed by GitHub
parent 52550ff571 d0c9941a67
commit 704b4d7eeb
3 changed files with 189 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
c/tests/test_idot.c
+55
View File
@@ -22,6 +22,50 @@ static int32_t ref_i4i8(const uint8_t *w4, const int8_t *x, int I){
return (int32_t)s;
}
/* Driver-level exactness: matmul_qt_ex on the IDOT path (allow_idot=1) must match
* a plain-C reference bit-for-bit. This exercises the SMMLA 2x2-tile drivers on the
* ARCH=native (i8mm) build and the SDOT drivers on the default build. The integer
* dot is exact and qrow_i8 is the shared quantizer, so any float bit mismatch is a
* driver bug (lane map, tiling, tail, or scale order), not rounding. */
static void fill_qt(QT *w, int fmt, int O, int I){
memset(w,0,sizeof *w);
w->fmt=fmt; w->O=O; w->I=I;
w->s=malloc((size_t)O*sizeof(float));
for(int o=0;o<O;o++) w->s[o]=0.001f+(float)(xr()%1000)*1e-6f;
if(fmt==1){
w->q8=malloc((size_t)O*I);
for(int64_t i=0;i<(int64_t)O*I;i++) w->q8[i]=(int8_t)((int)(xr()%256)-128);
}else{
size_t nb=(size_t)O*((I+1)/2);
w->q4=malloc(nb);
for(size_t i=0;i<nb;i++) w->q4[i]=(uint8_t)(xr()&0xFF);
}
}
static int check_driver(int fmt,int O,int I,int S){
QT w; fill_qt(&w,fmt,O,I); int rb=(I+1)/2;
float *x=malloc((size_t)S*I*sizeof(float));
float *y=malloc((size_t)S*O*sizeof(float));
float *yref=malloc((size_t)S*O*sizeof(float));
int8_t *xqr=malloc((size_t)S*I);
float *sxr=malloc((size_t)S*sizeof(float));
for(int64_t i=0;i<(int64_t)S*I;i++) x[i]=((float)(xr()%4001)-2000.f)/500.f;
matmul_qt_ex(y,x,&w,S,1);
for(int s=0;s<S;s++) sxr[s]=qrow_i8(x+(int64_t)s*I, xqr+(int64_t)s*I, I);
for(int o=0;o<O;o++) for(int s=0;s<S;s++){
int32_t d=fmt==1 ? ref_i8i8(w.q8+(int64_t)o*I, xqr+(int64_t)s*I, I)
: ref_i4i8(w.q4+(int64_t)o*rb, xqr+(int64_t)s*I, I);
yref[(int64_t)s*O+o]=(float)d*w.s[o]*sxr[s];
}
int rc=0;
for(int64_t i=0;i<(int64_t)S*O;i++)
if(memcmp(&y[i],&yref[i],sizeof(float))!=0){
fprintf(stderr,"FAIL driver fmt=%d O=%d I=%d S=%d idx=%lld: %.9g != %.9g\n",
fmt,O,I,S,(long long)i,(double)y[i],(double)yref[i]); rc=1; break;
}
free(w.s); free(w.q8); free(w.q4); free(x); free(y); free(yref); free(xqr); free(sxr);
return rc;
}
int main(void){
static const int sizes[]={1,2,15,16,17,31,32,33,63,64,65,100,127,128,1408,4096,4097};
static int8_t w[8192], x[8192]; static uint8_t w4[4096];
@@ -40,5 +84,16 @@ int main(void){
}
}
printf("idot kernel exactness (%s): ok\n", IDOT_KERNEL);
static const int Os[]={1,2,3,64,65};
static const int Is[]={16,17,100,1408};
static const int Ss[]={2,3,4,5,8};
for(int rep=0;rep<4;rep++)
for(unsigned a=0;a<sizeof Os/sizeof Os[0];a++)
for(unsigned b=0;b<sizeof Is/sizeof Is[0];b++)
for(unsigned c=0;c<sizeof Ss/sizeof Ss[0];c++)
for(int fmt=1;fmt<=2;fmt++)
if(check_driver(fmt,Os[a],Is[b],Ss[c])) return 1;
printf("idot driver exactness (%s): ok\n", IDOT_KERNEL);
return 0;
}