ggml-cpu: add AVX2 vec_dot for Q2_0 x Q8_0 on x86

ggml/src/ggml-cpu/arch-fallback.h

@@ -83,7 +83,6 @@
 #define ggml_gemm_q2_K_8x8_q8_K_generic ggml_gemm_q2_K_8x8_q8_K
 #elif defined(__x86_64__) || defined(__i386__) || defined(_M_IX86) || defined(_M_X64)
 // quants.c
-#define ggml_vec_dot_q2_0_q8_0_generic ggml_vec_dot_q2_0_q8_0
 // repack.cpp
 #define ggml_quantize_mat_q8_0_4x4_generic ggml_quantize_mat_q8_0_4x4
 #define ggml_quantize_mat_q8_K_4x4_generic ggml_quantize_mat_q8_K_4x4

ggml/src/ggml-cpu/arch/x86/quants.c

@@ -698,6 +698,64 @@ void ggml_vec_dot_q1_0_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const voi
 #endif
 }
 
+void ggml_vec_dot_q2_0_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) {
+    const int qk = QK2_0;
+    const int nb = n / qk;
+
+    assert(n % qk == 0);
+    assert(nrc == 1);
+    UNUSED(nrc);
+    UNUSED(bx);
+    UNUSED(by);
+    UNUSED(bs);
+
+    const block_q2_0 * GGML_RESTRICT x = vx;
+    const block_q8_0 * GGML_RESTRICT y = vy;
+
+#if defined(__AVX2__)
+    // weight (4*b + j) lives in bits [2j, 2j+1] of byte b: broadcast the 8 bytes to 4 lanes and shift lane j by 2j,
+    // which yields order (8*j + b); the q8 block is permuted to the same order, so the dot product is unchanged
+    const __m256i shifts = _mm256_set_epi64x(6, 4, 2, 0);
+    const __m256i m3     = _mm256_set1_epi8(0x03);
+    const __m256i one    = _mm256_set1_epi8(1);
+    // per-128-bit-lane byte shuffle: lane bytes b*4+j -> j*4+b
+    const __m256i shuf = _mm256_set_epi8(
+        15, 11, 7, 3,  14, 10, 6, 2,  13, 9, 5, 1,  12, 8, 4, 0,
+        15, 11, 7, 3,  14, 10, 6, 2,  13, 9, 5, 1,  12, 8, 4, 0);
+    // interleave the 32-bit chunks of the two lanes: (lo0, hi0, lo1, hi1, ...)
+    const __m256i perm = _mm256_set_epi32(7, 3, 6, 2, 5, 1, 4, 0);
+
+    __m256 acc = _mm256_setzero_ps();
+
+    for (int i = 0; i < nb; i++) {
+        const float d0 = GGML_CPU_FP16_TO_FP32(x[i].d);
+
+        // one Q2_0 block (64 weights) maps to two Q8_0 blocks (2 * 32)
+        for (int k = 0; k < 2; k++) {
+            const block_q8_0 * GGML_RESTRICT yb = &y[i * 2 + k];
+            const __m256 d = _mm256_set1_ps(d0 * GGML_CPU_FP16_TO_FP32(yb->d));
+
+            int64_t raw;
+            memcpy(&raw, &x[i].qs[k * 8], sizeof(raw));
+            __m256i qx = _mm256_srlv_epi64(_mm256_set1_epi64x(raw), shifts);
+            qx = _mm256_and_si256(qx, m3);
+            qx = _mm256_sub_epi8(qx, one);           // {0,1,2,3} -> {-1,0,1,2}
+
+            __m256i qy = _mm256_loadu_si256((const __m256i *) yb->qs);
+            qy = _mm256_shuffle_epi8(qy, shuf);
+            qy = _mm256_permutevar8x32_epi32(qy, perm);
+
+            const __m256 q = mul_sum_i8_pairs_float(qx, qy);
+            acc = _mm256_fmadd_ps(d, q, acc);
+        }
+    }
+
+    *s = hsum_float_8(acc);
+#else
+    ggml_vec_dot_q2_0_q8_0_generic(n, s, bs, vx, bx, vy, by, nrc);
+#endif
+}
+
 void ggml_vec_dot_q4_0_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) {
     const int qk = QK8_0;
     const int nb = n / qk;