iq2_k: AVX2

1 files changed, 56 insertions, 15 deletions

diff --git a/ggml/src/iqk/iqk_mul_mat.cpp b/ggml/src/iqk/iqk_mul_mat.cpp
index ad09d341..08c29cc5 100644
--- a/ggml/src/iqk/iqk_mul_mat.cpp
+++ b/ggml/src/iqk/iqk_mul_mat.cpp
@@ -1105,24 +1105,68 @@ struct DequantizerIQ4XS final : public BaseDequantizer<block_iq4_xs> {
     const __m256i values;
 };
 
-struct DequantizerIQ4K final : public BaseDequantizer<block_iq4_k> {
-    DequantizerIQ4K(const void * vx, size_t bx) : BaseDequantizer(vx, bx), values(load_iq4nl_values()) {}
+struct IQXKScales {
+    IQXKScales(int8_t shift, int8_t min_val) : min(_mm256_set1_epi16(min_val)), eshift(_mm_set1_epi8(shift)) {}
     template <typename Q8>
-    inline void new_block(int i, const Q8& q8, __m256 * accm, __m256i * scales) {
-        d = GGML_FP16_TO_FP32(x[i].d);
-        auto scales8 = make_scales(x[i].scales_l, (const uint16_t *)x[i].scales_h);
+    inline void process(int i, float d, uint16_t extra, __m128i scales8, const Q8& q8, __m256 * accm, __m256i * scales) const {
         auto scales16 = _mm256_cvtepi8_epi16(_mm_shuffle_epi8(scales8, hshuff));
-        auto extra128 = _mm_set1_epi16(x[i].extra);
+        auto extra128 = _mm_set1_epi16(extra);
         extra128 = _mm_cmpeq_epi8(_mm_and_si128(extra128, emask), emask);
-        extra128 = _mm_and_si128(extra128, e4);
+        extra128 = _mm_and_si128(extra128, eshift);
         extra128 = _mm_shuffle_epi8(extra128, eshuffle);
-        auto scales_s = _mm256_mullo_epi16(scales16, _mm256_add_epi16(_mm256_set1_epi16(-128), _mm256_cvtepi8_epi16(extra128)));
+        auto scales_s = _mm256_mullo_epi16(scales16, _mm256_add_epi16(min, _mm256_cvtepi8_epi16(extra128)));
         for (int iy = 0; iy < Q8::nrc_y; ++iy) {
             const __m256i prod  = _mm256_madd_epi16(scales_s, q8.load_bsums(iy, i));
             accm[iy] = _mm256_fmadd_ps(_mm256_set1_ps(d * q8.scale(iy, i)), _mm256_cvtepi32_ps(prod), accm[iy]);
         }
         prepare_scales_16(scales16, scales);
     }
+    const __m256i min;
+    const __m128i eshift;
+    const __m128i hshuff   = _mm_set_epi32(0x0f070e06, 0x0d050c04, 0x0b030a02, 0x09010800);
+    const __m128i emask    = _mm_set_epi32(0x80804040, 0x20201010, 0x08080404, 0x02020101);
+    const __m128i eshuffle = _mm_set_epi32(0x0f0d0b09, 0x07050301, 0x0e0c0a08, 0x06040200);
+};
+
+struct DequantizerIQ2K final : public BaseDequantizer<block_iq2_k> {
+    DequantizerIQ2K(const void * vx, size_t bx) : BaseDequantizer(vx, bx), iqxk(5, -32), values(load_values()) {}
+    template <typename Q8>
+    inline void new_block(int i, const Q8& q8, __m256 * accm, __m256i * scales) {
+        d = GGML_FP16_TO_FP32(x[i].d);
+        iqxk.process(i, d, x[i].extra, make_scales(x[i].scales), q8, accm, scales);
+    }
+    inline void prepare(int i, int j) {
+        bits.prepare(x[i].qs, j);
+        bits.values[0] = _mm256_shuffle_epi8(values, bits.values[0]);
+        bits.values[1] = _mm256_shuffle_epi8(values, bits.values[1]);
+        bits.values[2] = _mm256_shuffle_epi8(values, bits.values[2]);
+        bits.values[3] = _mm256_shuffle_epi8(values, bits.values[3]);
+    }
+    static inline __m256i load_values() {
+        static const uint8_t kvalues_iq2nl[16] = {1, 19, 33, 49, 0, 0, 0, 0,  6, 24, 38, 54, 0, 0, 0, 0};
+        auto val128 = _mm_loadu_si128((const __m128i *)kvalues_iq2nl);
+        return MM256_SET_M128I(val128, val128);
+    }
+    inline __m128i make_scales(const uint8_t * scales_l) const {
+        uint64_t aux64; std::memcpy(&aux64, scales_l, 8);
+        auto scl = _mm_and_si128(_mm_set_epi64x(aux64 >> 4, aux64), maskl);
+        return _mm_add_epi8(_mm_slli_epi16(scl, 1), m15);
+    }
+
+    Q2Bits bits;
+    const IQXKScales iqxk;
+    const __m256i values;
+    const __m128i m15      = _mm_set1_epi8(-15);
+    const __m128i maskl    = _mm_set1_epi8(0xf);
+};
+
+struct DequantizerIQ4K final : public BaseDequantizer<block_iq4_k> {
+    DequantizerIQ4K(const void * vx, size_t bx) : BaseDequantizer(vx, bx), iqxk(4, -128), values(load_iq4nl_values()) {}
+    template <typename Q8>
+    inline void new_block(int i, const Q8& q8, __m256 * accm, __m256i * scales) {
+        d = GGML_FP16_TO_FP32(x[i].d);
+        iqxk.process(i, d, x[i].extra, make_scales(x[i].scales_l, (const uint16_t *)x[i].scales_h), q8, accm, scales);
+    }
     inline void prepare(int i, int j) {
         bits.prepare16(x[i].qs, j);
         bits.values[0] = _mm256_shuffle_epi8(values, bits.values[0]);
@@ -1136,20 +1180,16 @@ struct DequantizerIQ4K final : public BaseDequantizer<block_iq4_k> {
         auto scl = _mm_and_si128(_mm_set_epi64x(aux64 >> 4, aux64), maskl);
         const uint32_t aux32 = scales_h[0] | (scales_h[1] << 16);
         auto aux = _mm_and_si128(_mm_set_epi32(aux32 >> 2, aux32, aux32 << 2, aux32 << 4), maskh);
-        auto sch = _mm_shuffle_epi8(aux, hshuff);
+        auto sch = _mm_shuffle_epi8(aux, iqxk.hshuff);
         return _mm_add_epi8(_mm_or_si128(scl, sch), m32);
     }
 
     Q4Bits bits;
+    const IQXKScales iqxk;
     const __m256i values;
     const __m128i maskl    = _mm_set1_epi8(0xf);
     const __m128i maskh    = _mm_set1_epi8(0x30);
-    const __m128i hshuff   = _mm_set_epi32(0x0f070e06, 0x0d050c04, 0x0b030a02, 0x09010800);
     const __m128i m32      = _mm_set1_epi8(-32);
-    const __m128i emask    = _mm_set_epi32(0x80804040, 0x20201010, 0x08080404, 0x02020101);
-    const __m128i e4       = _mm_set1_epi8(4);
-    const __m128i eshuffle = _mm_set_epi32(0x0f0d0b09, 0x07050301, 0x0e0c0a08, 0x06040200);
-
 };
 
 struct DequantizerQ5K final : public BaseDequantizer<block_q5_K> {
@@ -2887,7 +2927,8 @@ template <typename Dequantizer> void MulMat::set_functions(MulMat& m) {
             if constexpr (std::is_same_v<Dequantizer, DequantizerQ2K> ||
                           std::is_same_v<Dequantizer, DequantizerQ3K> ||
                           std::is_same_v<Dequantizer, DequantizerQ6K> ||
-                          std::is_same_v<Dequantizer, DequantizerIQ4K>) {
+                          std::is_same_v<Dequantizer, DequantizerIQ4K>||
+                          std::is_same_v<Dequantizer, DequantizerIQ2K>) {
                 m.funcs[0] = mul_mat_qY_K_q8_K_T<Dequantizer, 1>;
                 m.funcs[1] = mul_mat_qY_K_q8_K_T<Dequantizer, 2>;
                 m.funcs[2] = mul_mat_qY_K_q8_K_T<Dequantizer, 3>;