iq3_kt: AVX2 GEMM

Author: Iwan Kawrakow

ggml/src/ggml.c

@@ -1612,11 +1612,6 @@ static const ggml_type_traits_t type_traits[GGML_TYPE_COUNT] = {
 #else
         .vec_dot_type             = GGML_TYPE_Q8_0_X4,
 #endif
-//#ifdef __ARM_NEON
-//        .vec_dot_type             = GGML_TYPE_F16,
-//#else
-//        .vec_dot_type             = GGML_TYPE_F32,
-//#endif
         .nrows                    = 1,
         .row_meta_size            = 4,
     },
@@ -1629,11 +1624,16 @@ static const ggml_type_traits_t type_traits[GGML_TYPE_COUNT] = {
         .from_float               = quantize_row_iq3_kt,
         .from_float_ref           = (ggml_from_float_t)quantize_row_iq3_kt_ref,
         .vec_dot                  = vec_dot_iq3_kt_q8_k,
-#ifdef __ARM_NEON
-        .vec_dot_type             = GGML_TYPE_F16,
+#if defined __AVX2__
+        .vec_dot_type             = GGML_TYPE_Q8_2_X4,
 #else
-        .vec_dot_type             = GGML_TYPE_F32,
+        .vec_dot_type             = GGML_TYPE_Q8_0_X4,
 #endif
+//#ifdef __ARM_NEON
+//        .vec_dot_type             = GGML_TYPE_F16,
+//#else
+//        .vec_dot_type             = GGML_TYPE_F32,
+//#endif
         .nrows                    = 1,
         .row_meta_size            = 4,
     },

ggml/src/iqk/iqk_gemm_ktquants.cpp

@@ -98,7 +98,7 @@ struct Trellis2 {
 };
 
 
-template <bool is_8 = false>
+template <bool is_8 = false, bool is_abs = false>
 struct Trellis3 {
     constexpr static uint32_t ka = 0xCBAC1FED;
     constexpr static uint32_t ka1 = ka*ka;
@@ -127,7 +127,11 @@ struct Trellis3 {
         auto dot = _mm256_maddubs_epi16(v8, _mm256_set1_epi32(0x01010101));
         auto i8  = _mm256_add_epi32(_mm256_set1_epi32(-126), _mm256_madd_epi16(dot, _mm256_set1_epi16(1)));
 #endif
-        return _mm256_cvtepi32_ps(i8);
+        if constexpr (is_abs) {
+            return _mm256_cvtepi32_ps(_mm256_sign_epi32(i8, i8));
+        } else {
+            return _mm256_cvtepi32_ps(i8);
+        }
     }
     inline __m256 gen8(uint32_t val) const {
         auto v8 = _mm256_and_si256(next8(val), _mm256_set1_epi32(0x3f3f3f3f));
@@ -137,11 +141,14 @@ struct Trellis3 {
         auto dot = _mm256_maddubs_epi16(v8, _mm256_set1_epi32(0x01010101));
         auto i8  = _mm256_add_epi32(_mm256_set1_epi32(-126), _mm256_madd_epi16(dot, _mm256_set1_epi16(1)));
 #endif
-        return _mm256_cvtepi32_ps(i8);
+        if constexpr (is_abs) {
+            return _mm256_cvtepi32_ps(_mm256_sign_epi32(i8, i8));
+        } else {
+            return _mm256_cvtepi32_ps(i8);
+        }
     }
-    template <bool is_unsigned = false>
     inline __m256i next32(const uint32_t * val) const {
-        const __m256i offset = is_unsigned ? _mm256_setzero_si256() : _mm256_set1_epi32(-126);
+        const __m256i offset = _mm256_set1_epi32(-126);
         __m256i aux[4];
         for (int i = 0; i < 4; ++i) {
             auto i8 = _mm256_and_si256(next8(val[2*i+0], val[2*i+1]), _mm256_set1_epi32(0x3f3f3f3f));
@@ -156,11 +163,15 @@ struct Trellis3 {
         aux[2] = _mm256_packs_epi32(aux[2], aux[3]); // 16, 17, 18, 19, 24, 25, 26, 27, 20, 21, 22, 23, 28, 29, 30, 31
         aux[0] = _mm256_packs_epi16(aux[0], aux[2]); //  0,  1,  2,  3,  8,  9, 10, 11, 16, 17, 18, 19, 24, 25, 26, 27
                                                      //  4,  5,  6,  7, 12, 13, 14, 15, 20, 21, 22, 23, 28, 29, 30, 31
-        return _mm256_permutevar8x32_epi32(aux[0], shuffle);
+        if constexpr (is_abs) {
+            auto result = _mm256_permutevar8x32_epi32(aux[0], shuffle);
+            return _mm256_sign_epi8(result, result);
+        } else {
+            return _mm256_permutevar8x32_epi32(aux[0], shuffle);
+        }
     }
-    template <bool is_unsigned = false>
     inline __m256i next32(const uint16_t * val, uint32_t v0) const {
-        const __m256i offset = is_unsigned ? _mm256_setzero_si256() : _mm256_set1_epi32(-126);
+        const __m256i offset = _mm256_set1_epi32(-126);
         __m256i aux[4];
         for (int i = 0; i < 4; ++i) {
             auto i8 = _mm256_and_si256(next8(v0 + val[i]), _mm256_set1_epi32(0x3f3f3f3f));
@@ -175,11 +186,15 @@ struct Trellis3 {
         aux[2] = _mm256_packs_epi32(aux[2], aux[3]); // 16, 17, 18, 19, 24, 25, 26, 27, 20, 21, 22, 23, 28, 29, 30, 31
         aux[0] = _mm256_packs_epi16(aux[0], aux[2]); //  0,  1,  2,  3,  8,  9, 10, 11, 16, 17, 18, 19, 24, 25, 26, 27
                                                      //  4,  5,  6,  7, 12, 13, 14, 15, 20, 21, 22, 23, 28, 29, 30, 31
-        return _mm256_permutevar8x32_epi32(aux[0], shuffle);
+        if constexpr (is_abs) {
+            auto result = _mm256_permutevar8x32_epi32(aux[0], shuffle);
+            return _mm256_sign_epi8(result, result);
+        } else {
+            return _mm256_permutevar8x32_epi32(aux[0], shuffle);
+        }
     }
-    template <bool is_unsigned = false>
     inline void next64(const uint32_t * val, __m256i * result) const {
-        const __m256i offset = is_unsigned ? _mm256_setzero_si256() : _mm256_set1_epi32(-126);
+        const __m256i offset = _mm256_set1_epi32(-126);
         auto vka3 = _mm256_set1_epi32(ka3);
         __m256i aux[8];
         for (int i = 0; i < 4; ++i) {
@@ -203,6 +218,9 @@ struct Trellis3 {
             aux[4*k+0] = _mm256_packs_epi16(aux[4*k+0], aux[4*k+2]); //  0,  1,  2,  3,  8,  9, 10, 11, 16, 17, 18, 19, 24, 25, 26, 27
                                                                      //  4,  5,  6,  7, 12, 13, 14, 15, 20, 21, 22, 23, 28, 29, 30, 31
             result[k] = _mm256_permutevar8x32_epi32(aux[4*k+0], shuffle);
+            if constexpr (is_abs) {
+                result[k] = _mm256_sign_epi8(result[k], result[k]);
+            }
         }
     }
 };
@@ -449,6 +467,70 @@ void mul_mat_iq2_kt_q8_2_x4_T(int n, const void * vx, size_t bx, const DataInfo&
     }
 }
 
+void iqk_dequantize_iq3_kt_q80_r8(int n, const void * vx, size_t bx, void * vy, int nrc_x) {
+    GGML_ASSERT(n%QK_K == 0);
+    GGML_ASSERT(nrc_x%8 == 0);
+    const int nb = n/QK_K;
+
+    Trellis3<false, true> trellis;
+
+    auto shifts = _mm_set_epi32(0, 0, 4, 0);
+
+    block_q8_0_r8 * y = (block_q8_0_r8 *)vy;
+
+    const block_iq3_kt * x8[8];
+    float dkt[8];
+    float ls[8];
+    float ls_all[64];
+    uint32_t idx[8];
+    uint32_t sign_bits[16];
+
+    for (int ix = 0; ix < nrc_x; ix += 8) {
+        for (int k = 0; k < 8; ++k) {
+            const float * dptr = (const float *)((const char*)vx + (ix+k)*bx);
+            dkt[k] = dptr[0];
+            x8[k] = (const block_iq3_kt *)(dptr + 1);
+        }
+        auto vd = _mm256_mul_ps(_mm256_set1_ps(1.01f), _mm256_loadu_ps(dkt));
+
+        for (int i = 0; i < nb; ++i) {
+            for (int k = 0; k < 8; ++k) {
+                auto s8 = _mm_set1_epi32(*(const uint32_t *)x8[k][i].scales);
+                s8 = _mm_and_si128(_mm_srlv_epi32(s8, shifts), _mm_set1_epi8(0xf));
+                auto s32 = _mm256_cvtepi8_epi32(s8);
+                _mm256_storeu_ps(ls_all + 8*k, _mm256_cvtepi32_ps(s32));
+            }
+            auto mask = _mm256_set1_epi8(1);
+            for (int ib = 0; ib < QK_K/32; ++ib) {
+                for (int k = 0; k < 8; ++k) ls[k] = ls_all[8*k+ib];
+                auto scales = _mm256_mul_ps(vd, _mm256_loadu_ps(ls));
+                _mm_storeu_si128((__m128i *)y[ib].d, _mm256_cvtps_ph(scales, _MM_FROUND_TO_NEAREST_INT));
+                for (int j = 0; j < 4; ++j) {
+                    for (int k = 0; k < 8; ++k) {
+                        const uint16_t * ql = (const uint16_t *)x8[k][i].ql;
+                        idx[k] = ql[4*ib+j] + 4096;
+                        auto qh = (const uint32_t *)x8[k][i].qh;
+                        sign_bits[k+0] = qh[2*j+0];
+                        sign_bits[k+8] = qh[2*j+1];
+                    }
+                    __m256i packed[2];
+                    trellis.next64(idx, packed);
+                    auto signs1 = _mm256_loadu_si256((const __m256i *)sign_bits+0);
+                    auto signs2 = _mm256_loadu_si256((const __m256i *)sign_bits+1);
+                    signs1 = _mm256_or_si256(_mm256_cmpeq_epi8(_mm256_and_si256(signs1, mask), mask), _mm256_set1_epi8(1));
+                    signs2 = _mm256_or_si256(_mm256_cmpeq_epi8(_mm256_and_si256(signs2, mask), mask), _mm256_set1_epi8(1));
+                    packed[0] = _mm256_sign_epi8(packed[0], signs1);
+                    packed[1] = _mm256_sign_epi8(packed[1], signs2);
+                    _mm256_storeu_si256((__m256i *)y[ib].qs+2*j+0, packed[0]);
+                    _mm256_storeu_si256((__m256i *)y[ib].qs+2*j+1, packed[1]);
+                }
+                mask = _mm256_slli_epi16(mask, 1);
+            }
+            y += 8; // = QK_K/32;
+        }
+    }
+}
+
 inline __m256 abs_ps(__m256 vals) {
     // Clear sign-bit of all the 32-bit floats in vals
     __m256 sign_bit = _mm256_set1_ps(-0.0f);
@@ -887,10 +969,10 @@ bool iqk_set_kernels_ktquants(int ne00, int typeA, int typeB, std::array<mul_mat
 
 }
 
-bool iqk_dequantize_ktquants(int type, int n, const void * vx, size_t bx, void * y, size_t stride_y, int nrc_x) {
+bool iqk_dequantize_ktquants(int type, int n, const void * vx, size_t bx, void * y, [[maybe_unused]] size_t stride_y, int nrc_x) {
     switch (type) {
         case GGML_TYPE_IQ2_KT: iqk_dequantize_iq2_kt_q80_r8(n, vx, bx, y, nrc_x); break;
-        case GGML_TYPE_IQ3_KT: iqk_dequantize_iq3_kt(n, vx, bx, (float *)y, stride_y, nrc_x); break;
+        case GGML_TYPE_IQ3_KT: iqk_dequantize_iq3_kt_q80_r8(n, vx, bx, y, nrc_x); break;
         case GGML_TYPE_IQ4_KT: iqk_dequantize_iq4_kt_q80_r8(n, vx, bx, y, nrc_x); break;
         default: return false;
     }

ggml/src/iqk/iqk_mul_mat.cpp

@@ -243,7 +243,7 @@ struct MulMat {
             case GGML_TYPE_Q4_K   : return nrc_y >= 32 ? GGML_TYPE_Q8_1    : type;
             case GGML_TYPE_Q5_K   : return nrc_y >= 32 ? GGML_TYPE_Q8_1    : type;
             case GGML_TYPE_IQ2_KT : return nrc_y >= 32 ? GGML_TYPE_Q8_0_R8 : type;
-            case GGML_TYPE_IQ3_KT : return nrc_y >= 32 ? GGML_TYPE_F32 : type;
+            case GGML_TYPE_IQ3_KT : return nrc_y >= 32 ? GGML_TYPE_Q8_0_R8 : type;
             case GGML_TYPE_IQ4_KT : return nrc_y >= 32 ? GGML_TYPE_Q8_0_R8 : type;
             default: break;
         }