quant: AVX-512 arm for matmul_i3 (fmt=5 int3-g64)

c/colibri.c

@@ -6688,6 +6688,11 @@ int main(int argc, char **argv){
         if(!i4_acc512_selftest()) return 1;
         puts("AVX512 i4 selftest: ok"); return 0;
     }
+    if(getenv("I3_AVX512")) g_i3_avx512=atoi(getenv("I3_AVX512"))!=0;
+    if(getenv("I3_AVX512_TEST")){
+        if(!i3_avx512_selftest()) return 1;
+        puts("AVX512 i3 selftest: ok"); return 0;
+    }
 #endif
     const char *snap=getenv("SNAP"); if(!snap){fprintf(stderr,"SNAP=<dir>\n");return 1;}
     g_nopack = getenv("NOPACK")?1:0;

c/quant.h

@@ -295,10 +295,62 @@ static void matmul_i2(float *y, const float *x, const uint8_t *q2, const float *
 static inline int64_t i3_groups(int I){ return ((int64_t)I + I3_GROUP - 1) / I3_GROUP; }
 static inline int64_t i3_rowbytes(int I){ return i3_groups(I) * I3_GBYTES; }
 
+#if defined(__AVX512F__) && defined(__AVX512BW__)
+static int g_i3_avx512=1;
+/* one full 64-value group -> f32 partial. Relies on immintrin.h arriving via the
+ * __AVX2__-gated include above (AVX512F implies AVX2 on clang/gcc/MSVC), same as
+ * dot_i4f_avx512. */
+static inline float dot_i3g64_avx512(const uint8_t *lo, const uint8_t *hi, const float *x){
+    const __m128i m2=_mm_set1_epi8(3); const __m512i c4=_mm512_set1_epi8(4);
+    __m128i by=_mm_loadu_si128((const __m128i*)lo);
+    __m128i p0=_mm_and_si128(by,m2),                   p1=_mm_and_si128(_mm_srli_epi16(by,2),m2);
+    __m128i p2=_mm_and_si128(_mm_srli_epi16(by,4),m2), p3=_mm_and_si128(_mm_srli_epi16(by,6),m2);
+    __m128i l01=_mm_unpacklo_epi8(p0,p1), h01=_mm_unpackhi_epi8(p0,p1);
+    __m128i l23=_mm_unpacklo_epi8(p2,p3), h23=_mm_unpackhi_epi8(p2,p3);
+    __m512i lov=_mm512_inserti32x4(_mm512_inserti32x4(_mm512_inserti32x4(
+        _mm512_castsi128_si512(_mm_unpacklo_epi16(l01,l23)),
+        _mm_unpackhi_epi16(l01,l23),1),
+        _mm_unpacklo_epi16(h01,h23),2),
+        _mm_unpackhi_epi16(h01,h23),3);            /* byte k = low 2 bits of value k */
+    uint64_t hb; memcpy(&hb,hi,8);                 /* mask bit k = high bit of value k */
+    __m512i wq=_mm512_sub_epi8(_mm512_mask_add_epi8(lov,(__mmask64)hb,lov,c4),c4); /* [-4,3] in order */
+    __m512 ac0=_mm512_setzero_ps(), ac1=_mm512_setzero_ps();
+    ac0=_mm512_fmadd_ps(_mm512_loadu_ps(x),    _mm512_cvtepi32_ps(_mm512_cvtepi8_epi32(_mm512_castsi512_si128(wq))),      ac0);
+    ac1=_mm512_fmadd_ps(_mm512_loadu_ps(x+16), _mm512_cvtepi32_ps(_mm512_cvtepi8_epi32(_mm512_extracti32x4_epi32(wq,1))), ac1);
+    ac0=_mm512_fmadd_ps(_mm512_loadu_ps(x+32), _mm512_cvtepi32_ps(_mm512_cvtepi8_epi32(_mm512_extracti32x4_epi32(wq,2))), ac0);
+    ac1=_mm512_fmadd_ps(_mm512_loadu_ps(x+48), _mm512_cvtepi32_ps(_mm512_cvtepi8_epi32(_mm512_extracti32x4_epi32(wq,3))), ac1);
+    return _mm512_reduce_add_ps(_mm512_add_ps(ac0,ac1));
+}
+static int i3_avx512_selftest(void){
+    /* fixed group, asymmetric in every lane: pseudo-random 3-bit values with
+     * distinct nonzero integer activations. All terms and partials are small
+     * integers (exact in f32 under ANY summation order), so the compare is
+     * exact — any lane permutation, bias error or plane mix-up shifts the sum. */
+    uint8_t lo[16]={0}, hi[8]={0}; float x[I3_GROUP]; double ref=0;
+    uint64_t r=0x9E3779B97F4A7C15ull;
+    for(int k=0;k<I3_GROUP;k++){
+        r^=r<<13; r^=r>>7; r^=r<<17;
+        unsigned u=(unsigned)(r&7);
+        lo[k>>2]|=(uint8_t)((u&3)<<((k&3)*2));
+        hi[k>>3]|=(uint8_t)((u>>2)<<(k&7));
+        x[k]=(k&1)?-(float)(k+1):(float)(k+1);
+        ref+=(double)x[k]*((int)u-4);
+    }
+    float got=dot_i3g64_avx512(lo,hi,x);
+    if(got!=(float)ref){ fprintf(stderr,"AVX512 i3 selftest: %.9g != %.9g\n",got,ref); return 0; }
+    return 1;
+}
+#endif
+
 /* Dequant-on-use with PER-GROUP scale. Exact f32 path only (no IDOT in v1: int8
  * activations don't compose with per-group accumulation without a kernel
  * restructure — follow-up). NEON: low plane = matmul_i2's unpack, high plane
- * expanded via vtst on bit masks; x86 stays scalar for now (follow-up). */
+ * expanded via vtst on bit masks. AVX-512(F+BW): same unpack at 128-bit, high
+ * plane loaded as a __mmask64 (bit k = value k) driving a masked +4; one full
+ * group per iteration (dot_i3g64_avx512; I3_AVX512=0 falls back to scalar).
+ * Other x86 stays scalar (follow-up). Both vector arms reorder fma WITHIN a
+ * group only; the per-group partial is scaled by scale[g] and added to the
+ * row accumulator in scalar order, exactly like the scalar loop. */
 static void matmul_i3(float *y, const float *x, const uint8_t *q3, const float *scale, int S, int I, int O){
     int64_t ng=i3_groups(I), rb=i3_rowbytes(I);
     #pragma omp parallel for schedule(static)
@@ -312,7 +364,9 @@ static void matmul_i3(float *y, const float *x, const uint8_t *q3, const float *
                 const uint8_t *lo=wrow+g*I3_GBYTES, *hi=lo+16;
                 int base=(int)(g*I3_GROUP), n = I-base < I3_GROUP ? I-base : I3_GROUP;
                 float a=0; int k=0;
-#if defined(__ARM_NEON)
+#if defined(__AVX512F__) && defined(__AVX512BW__)
+                if(g_i3_avx512 && n==I3_GROUP){ a=dot_i3g64_avx512(lo,hi,xs+base); k=I3_GROUP; }
+#elif defined(__ARM_NEON)
                 if(n==I3_GROUP){
                     const uint8x8_t m2v=vdup_n_u8(3); const int8x16_t b4q=vdupq_n_s8(4);
                     const uint8x16_t bitm={1,2,4,8,16,32,64,128,1,2,4,8,16,32,64,128};

c/tests/test_int3.c

@@ -1,5 +1,5 @@
 /* int3-g64 (fmt=5) tests: pack layout, dequant round-trip vs plain-C reference,
- * matmul_i3 (NEON + scalar tail) vs reference dequant-matmul, per-row helpers,
+ * matmul_i3 (NEON/AVX-512 + scalar tail) vs reference dequant-matmul, per-row helpers,
  * the .qs-size format tag, and the quality claim in miniature (per-group int3
  * beats per-row int4 on rows with outliers — the #132 result this format ships). */
 #define main coli_glm_main_unused
@@ -11,7 +11,8 @@
 #include <math.h>
 
 static int fails = 0;
-#define CHECK(c) do{ if(!(c)){ printf("FAIL %s:%d: %s\n", __FILE__, __LINE__, #c); fails++; } }while(0)
+static int cur_I = 0, cur_S = 0;               /* shape under test, for failure triage */
+#define CHECK(c) do{ if(!(c)){ printf("FAIL %s:%d: %s (I=%d S=%d O=%d)\n", __FILE__, __LINE__, #c, cur_I, cur_S, (int)O); fails++; } }while(0)
 
 static uint64_t rng = 0x9E3779B97F4A7C15ull;
 static float rndf(void){ rng ^= rng << 13; rng ^= rng >> 7; rng ^= rng << 17;
@@ -44,15 +45,15 @@ static void unpack_i3(const uint8_t *q3, const float *s, float *dq, int O, int I
 }
 
 int main(void){
-    const int Is[]={64,128,192,100,65,7168};   /* incl. short tail groups and one real GLM dim */
+    const int Is[]={64,128,192,100,65,33,7,7168}; /* incl. short tail groups, I<64 (scalar-only group), one real GLM dim */
     enum { O=7, MAXI=7168 };
     static float w[(int64_t)O*MAXI], dq_ref[(int64_t)O*MAXI], dq_pk[(int64_t)O*MAXI];
     static float x[4*MAXI], y_ref[4*O], y_ker[4*O];
     static uint8_t q3[(int64_t)O*(MAXI/64+1)*24];
     static float   sc[(int64_t)O*(MAXI/64+1)];
 
     for(unsigned c=0;c<sizeof Is/sizeof *Is;c++){
-        int I=Is[c];
+        int I=Is[c]; cur_I=I; cur_S=0;
         for(int64_t i=0;i<(int64_t)O*I;i++) w[i]=rndf()*0.05f;
         w[3]=1.7f; w[(int64_t)2*I+5]=-2.2f;                 /* outliers */
 
@@ -65,8 +66,8 @@ int main(void){
         CHECK(bad==0);
 
         /* 2. matmul_i3 == matmul over the dequantized reference (fp tolerance:
-         *    NEON fma order differs from the scalar reference loop) */
-        for(int S=1;S<=4;S+=3){
+         *    NEON/AVX-512 fma order differs from the scalar reference loop) */
+        for(int S=1;S<=4;S+=3){ cur_S=S;
             for(int64_t i=0;i<(int64_t)S*I;i++) x[i]=rndf();
             matmul_i3(y_ker, x, q3, sc, S, I, O);
             for(int s=0;s<S;s++) for(int o=0;o<O;o++){
@@ -80,6 +81,7 @@ int main(void){
         }
 
         /* 3. QT plumbing: qt_alloc(bits=3) -> qt_fill -> matmul_qt & qt_bytes & helpers */
+        cur_S=1;
         QT t; qt_alloc(&t, O, I, 3);
         CHECK(t.fmt==5);
         qt_fill(&t, w, 3);
@@ -123,7 +125,7 @@ int main(void){
     /* 5. quality in miniature: on rows with outliers, per-group int3 must beat
      *    per-row int4 on reconstruction RMS (the #132 finding this format ships). */
     {
-        int I=1024;
+        int I=1024; cur_I=I; cur_S=1;
         for(int64_t i=0;i<(int64_t)O*I;i++) w[i]=rndf()*0.02f;
         for(int o=0;o<O;o++) w[(int64_t)o*I+(o*37)%I]=1.5f;    /* one outlier per row */
         ref_i3_dequant(w, dq_ref, O, I);