arm: NEON vec_dot for PQ2_0

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

@@ -82,10 +82,10 @@
 #define ggml_gemm_q1_0_4x8_q8_0_generic ggml_gemm_q1_0_4x8_q8_0
 #define ggml_gemm_pq2_0_4x8_q8_0_generic ggml_gemm_pq2_0_4x8_q8_0
 #elif defined(__aarch64__) || defined(__arm__) || defined(_M_ARM) || defined(_M_ARM64)
-// PTQ1_0 currently has only the generic vec_dot; alias it here until a SIMD version lands
+// PTQ1_0's NEON vec_dot is unverified on real ARM hardware and measured slower
+// than the generic loop on two independent devices (Apple M5 Pro, Snapdragon 7
+// Gen 4); alias it back to generic until a validated replacement lands
 #define ggml_vec_dot_ptq1_0_q8_0_generic ggml_vec_dot_ptq1_0_q8_0
-// PQ2_0 x Q8_K has only the generic vec_dot outside x86; alias it until a SIMD version lands
-#define ggml_vec_dot_pq2_0_q8_K_generic ggml_vec_dot_pq2_0_q8_K
 // repack.cpp
 #define ggml_quantize_mat_q8_K_4x4_generic ggml_quantize_mat_q8_K_4x4
 #define ggml_quantize_mat_q8_K_4x8_generic ggml_quantize_mat_q8_K_4x8

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

@@ -336,12 +336,12 @@ void ggml_vec_dot_pq2_0_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const vo
             const uint8x8_t raw = vld1_u8(&x[i].qs[k * 8]);
             const uint8x16_t raw16 = vcombine_u8(raw, raw);
 
-            uint8x16_t bytes0 = vqtbl1q_u8(raw16, idx_lo);
+            uint8x16_t bytes0 = ggml_vqtbl1q_u8(raw16, idx_lo);
             int8x16_t qv0 = vsubq_s8(
                 vreinterpretq_s8_u8(vandq_u8(vshlq_u8(bytes0, shifts), mask2)),
                 one);
 
-            uint8x16_t bytes1 = vqtbl1q_u8(raw16, idx_hi);
+            uint8x16_t bytes1 = ggml_vqtbl1q_u8(raw16, idx_hi);
             int8x16_t qv1 = vsubq_s8(
                 vreinterpretq_s8_u8(vandq_u8(vshlq_u8(bytes1, shifts), mask2)),
                 one);
@@ -365,6 +365,68 @@ void ggml_vec_dot_pq2_0_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const vo
     *s = sumf;
 }
 
+void ggml_vec_dot_pq2_0_q8_K(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) {
+    assert(n % QK_K == 0);
+    assert(nrc == 1);
+    UNUSED(nrc);
+    UNUSED(bx);
+    UNUSED(by);
+    UNUSED(bs);
+
+    const block_pq2_0 * GGML_RESTRICT x = vx;
+    const block_q8_K  * GGML_RESTRICT y = vy;
+    const int nb = n / QK_PQ2_0;
+
+    float sumf = 0.0f;
+
+#if defined(__ARM_NEON)
+    // Same 2-bit codec as the Q8_0 path above: byte b of a 32-weight sub-block holds
+    // weights 4b..4b+3, LSB pair first, so replicating each byte four times and shifting
+    // by {0,2,4,6} lands code 4b+j at lane 4b+j. With one activation scale per 256 the
+    // four sub-block dots accumulate in int32 and the scale is applied once per block,
+    // instead of a float multiply-add per sub-block.
+    static const uint8_t tbl_idx_lo[16] = {0,0,0,0, 1,1,1,1, 2,2,2,2, 3,3,3,3};
+    static const uint8_t tbl_idx_hi[16] = {4,4,4,4, 5,5,5,5, 6,6,6,6, 7,7,7,7};
+    static const int8_t  shift_vals[16] = {0,-2,-4,-6, 0,-2,-4,-6, 0,-2,-4,-6, 0,-2,-4,-6};
+
+    const uint8x16_t idx_lo = vld1q_u8(tbl_idx_lo);
+    const uint8x16_t idx_hi = vld1q_u8(tbl_idx_hi);
+    const int8x16_t  shifts = vld1q_s8(shift_vals);
+    const uint8x16_t mask2  = vdupq_n_u8(0x03);
+    const int8x16_t  one    = vdupq_n_s8(1);
+
+    for (int i = 0; i < nb; i++) {
+        const block_q8_K * GGML_RESTRICT yb = &y[i >> 1];
+        const int8_t     * GGML_RESTRICT q8 = yb->qs + 128 * (i & 1);
+
+        int32x4_t acc = vdupq_n_s32(0);
+
+        for (int k = 0; k < 4; k++) {
+            const uint8x8_t  raw   = vld1_u8(&x[i].qs[8 * k]);
+            const uint8x16_t raw16 = vcombine_u8(raw, raw);
+
+            const int8x16_t qv0 = vsubq_s8(
+                vreinterpretq_s8_u8(vandq_u8(vshlq_u8(ggml_vqtbl1q_u8(raw16, idx_lo), shifts), mask2)), one);
+            const int8x16_t qv1 = vsubq_s8(
+                vreinterpretq_s8_u8(vandq_u8(vshlq_u8(ggml_vqtbl1q_u8(raw16, idx_hi), shifts), mask2)), one);
+
+            const int8x16_t y0 = vld1q_s8(q8 + 32 * k);
+            const int8x16_t y1 = vld1q_s8(q8 + 32 * k + 16);
+
+            acc = ggml_vdotq_s32(acc, qv0, y0);
+            acc = ggml_vdotq_s32(acc, qv1, y1);
+        }
+
+        sumf += (GGML_CPU_FP16_TO_FP32(x[i].d) * yb->d) * (float) vaddvq_s32(acc);
+    }
+#else
+    ggml_vec_dot_pq2_0_q8_K_generic(n, s, bs, vx, bx, vy, by, nrc);
+    return;
+#endif
+
+    *s = sumf;
+}
+
 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;

tests/test-quantize-fns.cpp

@@ -2,6 +2,7 @@
 
 #include "ggml.h"
 #include "ggml-cpu.h"
+#include "../ggml/src/ggml-quants.h"
 
 #undef NDEBUG
 #include <assert.h>
@@ -202,6 +203,84 @@ static int test_vec_dot_q(bool verbose) {
     return num_failed;
 }
 
+static int test_vec_dot_ternary(bool verbose) {
+    int num_failed = 0;
+    for (ggml_type type : {GGML_TYPE_PQ2_0, GGML_TYPE_PTQ1_0}) {
+        const auto * traits = ggml_get_type_traits(type);
+        const auto * cpu = ggml_get_type_traits_cpu(type);
+        // PQ2_0 dots against Q8_K (one float scale per 256), PTQ1_0 against Q8_0 (one fp16
+        // scale per 32), so the activation side is built per format. PQ2_0 needs whole Q8_K
+        // blocks, i.e. an even number of 128-weight blocks.
+        const bool  q8k      = type == GGML_TYPE_PQ2_0;
+        const ggml_type ytype = q8k ? GGML_TYPE_Q8_K : GGML_TYPE_Q8_0;
+        for (int nb : q8k ? std::vector<int>{2, 4} : std::vector<int>{1, 3}) {
+            const int n = nb * 128;
+            std::vector<block_pq2_0> pq(nb);
+            std::vector<block_ptq1_0> ptq(nb);
+            std::vector<block_q8_0> q8(nb * 4);
+            std::vector<block_q8_K> q8k_blocks(nb / 2 + 1);
+            std::vector<float> x(n), y(n);
+            const void * weights = type == GGML_TYPE_PQ2_0 ? (const void *) pq.data() : (const void *) ptq.data();
+            for (int pattern = 0; pattern < 256; ++pattern) {
+                for (int i = 0; i < nb; ++i) {
+                    pq[i].d = ptq[i].d = ggml_fp32_to_fp16(0.25f * (i + 1));
+                    for (size_t j = 0; j < sizeof(pq[i].qs); ++j) {
+                        pq[i].qs[j] = (uint8_t) (pattern + 17*j + i);
+                    }
+                    for (size_t j = 0; j < sizeof(ptq[i].qs); ++j) {
+                        ptq[i].qs[j] = (uint8_t) (pattern + 17*j + i);
+                    }
+                    for (size_t j = 0; j < sizeof(ptq[i].qh); ++j) {
+                        ptq[i].qh[j] = (uint8_t) (pattern + 37*j + i);
+                    }
+                }
+                for (int i = 0; i < nb * 4; ++i) {
+                    q8[i].d = ggml_fp32_to_fp16(0.125f * (i % 4 + 1));
+                    for (int j = 0; j < QK8_0; ++j) {
+                        q8[i].qs[j] = (int8_t) ((pattern + 13*j + i) % 256 - 128);
+                    }
+                }
+                for (size_t i = 0; i < q8k_blocks.size(); ++i) {
+                    q8k_blocks[i].d = 0.125f * (i % 4 + 1);
+                    for (int j = 0; j < QK_K; ++j) {
+                        q8k_blocks[i].qs[j] = (int8_t) ((pattern + 13*j + i) % 256 - 128);
+                    }
+                    // bsums is unused by the PQ2_0 dot but keep it consistent.
+                    for (int j = 0; j < QK_K/16; ++j) {
+                        int16_t s = 0;
+                        for (int t = 0; t < 16; ++t) s += q8k_blocks[i].qs[j*16 + t];
+                        q8k_blocks[i].bsums[j] = s;
+                    }
+                }
+                const void * acts = q8k ? (const void *) q8k_blocks.data() : (const void *) q8.data();
+                traits->to_float(weights, x.data(), n);
+                if (q8k) {
+                    // Q8_K is an activation-only type and has no to_float, so expand it here.
+                    for (int j = 0; j < n; ++j) {
+                        const block_q8_K & b = q8k_blocks[j / QK_K];
+                        y[j] = b.d * (float) b.qs[j % QK_K];
+                    }
+                } else {
+                    ggml_get_type_traits(ytype)->to_float(acts, y.data(), n);
+                }
+                const float ref = dot_product(x.data(), y.data(), n);
+                float result = INFINITY;
+                cpu->vec_dot(n, &result, 0, weights, 0, acts, 0, 1);
+                // Power-of-two scales keep this comparison exact.
+                const bool failed = result != ref;
+                num_failed += failed;
+                if (failed) {
+                    printf("%5s packed dot nb=%d pattern=%d: FAILED (ref=%f got=%f)\n", ggml_type_name(type), nb, pattern, ref, result);
+                }
+            }
+        }
+    }
+    if (num_failed || verbose) {
+        printf("ternary packed dot products: %s (%d failures)\n", RESULT_STR[num_failed != 0], num_failed);
+    }
+    return num_failed;
+}
+
 int main(int argc, char * argv[]) {
     bool verbose = false;
 
@@ -223,6 +302,7 @@ int main(int argc, char * argv[]) {
 
     num_failed += test_vec_dot_f32(verbose);
     num_failed += test_vec_dot_q(verbose);
+    num_failed += test_vec_dot_ternary(verbose);
 
     if (num_failed || verbose) {
         printf("%d tests failed\n", num_failed);