char-fork/llama.cpp

git clone https://git.t4t.associates/char-fork/llama.cpp

Kartik Sirohiggml: vectorize ggml_vec_dot_q4_1_q8_1 with WASM SIMD128 (#22209)4c5130961

char/aliceai-t5-moe
45.6 KiB1292 linesraw
1#define GGML_COMMON_IMPL_C
2#include "ggml-common.h"
3#include "ggml-quants.h"
4#include "ggml-impl.h"
5#include "ggml-cpu.h"
6#include "simd-mappings.h"
7
8#include "../../quants.h"
9#include "../../ggml-cpu-impl.h"
10
11#include <math.h>
12#include <string.h>
13#include <assert.h>
14#include <float.h>
15#include <stdlib.h> // for qsort
16#include <stdio.h>  // for GGML_ASSERT
17
18#define GROUP_MAX_EPS 1e-15f
19#define GROUP_MAX_EPS_IQ3_XXS 1e-8f
20#define GROUP_MAX_EPS_IQ2_S 1e-8f
21#define GROUP_MAX_EPS_IQ1_M 1e-7f
22#define GROUP_MAX_EPS_IQ1_S 1e-12f
23
24#define UNUSED GGML_UNUSED
25
26#if defined(__wasm_simd128__)
27#define B1(c,s,n)  0x ## n ## c ,  0x ## n ## s
28#define B2(c,s,n) B1(c,s,n ## c), B1(c,s,n ## s)
29#define B3(c,s,n) B2(c,s,n ## c), B2(c,s,n ## s)
30#define B4(c,s,n) B3(c,s,n ## c), B3(c,s,n ## s)
31#define B5(c,s,n) B4(c,s,n ## c), B4(c,s,n ## s)
32#define B6(c,s,n) B5(c,s,n ## c), B5(c,s,n ## s)
33#define B7(c,s,n) B6(c,s,n ## c), B6(c,s,n ## s)
34#define B8(c,s  ) B7(c,s,     c), B7(c,s,     s)
35
36// precomputed tables for expanding 8bits to 8 bytes:
37static const uint64_t table_b2b_0[1 << 8] = { B8(00, 10) }; // ( b) << 4
38static const uint64_t table_b2b_1[1 << 8] = { B8(10, 00) }; // (!b) << 4
39#endif
40
41void quantize_row_q8_0(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t k) {
42    assert(QK8_0 == 32);
43    assert(k % QK8_0 == 0);
44    const int nb = k / QK8_0;
45
46    block_q8_0 * GGML_RESTRICT y = vy;
47
48#if defined __wasm_simd128__
49    for (int i = 0; i < nb; i++) {
50        v128_t srcv [8];
51        v128_t asrcv[8];
52        v128_t amaxv[8];
53
54        for (int j = 0; j < 8; j++) srcv[j]  = wasm_v128_load(x + i*32 + 4*j);
55        for (int j = 0; j < 8; j++) asrcv[j] = wasm_f32x4_abs(srcv[j]);
56
57        for (int j = 0; j < 4; j++) amaxv[2*j] = wasm_f32x4_max(asrcv[2*j], asrcv[2*j+1]);
58        for (int j = 0; j < 2; j++) amaxv[4*j] = wasm_f32x4_max(amaxv[4*j], amaxv[4*j+2]);
59        for (int j = 0; j < 1; j++) amaxv[8*j] = wasm_f32x4_max(amaxv[8*j], amaxv[8*j+4]);
60
61        const float amax = MAX(MAX(wasm_f32x4_extract_lane(amaxv[0], 0),
62                                   wasm_f32x4_extract_lane(amaxv[0], 1)),
63                               MAX(wasm_f32x4_extract_lane(amaxv[0], 2),
64                                   wasm_f32x4_extract_lane(amaxv[0], 3)));
65
66        const float d = amax / ((1 << 7) - 1);
67        const float id = d ? 1.0f/d : 0.0f;
68
69        y[i].d = GGML_CPU_FP32_TO_FP16(d);
70
71        for (int j = 0; j < 8; j++) {
72            const v128_t v  = wasm_f32x4_mul(srcv[j], wasm_f32x4_splat(id));
73            const v128_t vi = wasm_i32x4_trunc_sat_f32x4(v);
74
75            y[i].qs[4*j + 0] = wasm_i32x4_extract_lane(vi, 0);
76            y[i].qs[4*j + 1] = wasm_i32x4_extract_lane(vi, 1);
77            y[i].qs[4*j + 2] = wasm_i32x4_extract_lane(vi, 2);
78            y[i].qs[4*j + 3] = wasm_i32x4_extract_lane(vi, 3);
79        }
80    }
81#else
82    GGML_UNUSED(nb);
83    // scalar
84    quantize_row_q8_0_ref(x, y, k);
85#endif
86}
87
88void quantize_row_q8_1(const float * GGML_RESTRICT x, void * GGML_RESTRICT vy, int64_t k) {
89    assert(k % QK8_1 == 0);
90    const int nb = k / QK8_1;
91
92    block_q8_1 * GGML_RESTRICT y = vy;
93#if defined __wasm_simd128__
94    for (int i = 0; i < nb; i++) {
95        v128_t srcv [8];
96        v128_t asrcv[8];
97        v128_t amaxv[8];
98
99        for (int j = 0; j < 8; j++) srcv[j]  = wasm_v128_load(x + i*32 + 4*j);
100        for (int j = 0; j < 8; j++) asrcv[j] = wasm_f32x4_abs(srcv[j]);
101
102        for (int j = 0; j < 4; j++) amaxv[2*j] = wasm_f32x4_max(asrcv[2*j], asrcv[2*j+1]);
103        for (int j = 0; j < 2; j++) amaxv[4*j] = wasm_f32x4_max(amaxv[4*j], amaxv[4*j+2]);
104        for (int j = 0; j < 1; j++) amaxv[8*j] = wasm_f32x4_max(amaxv[8*j], amaxv[8*j+4]);
105
106        const float amax = MAX(MAX(wasm_f32x4_extract_lane(amaxv[0], 0),
107                                   wasm_f32x4_extract_lane(amaxv[0], 1)),
108                               MAX(wasm_f32x4_extract_lane(amaxv[0], 2),
109                                   wasm_f32x4_extract_lane(amaxv[0], 3)));
110
111        const float d = amax / ((1 << 7) - 1);
112        const float id = d ? 1.0f/d : 0.0f;
113
114        y[i].d = GGML_CPU_FP32_TO_FP16(d);
115
116        v128_t accv = wasm_i32x4_splat(0);
117
118        for (int j = 0; j < 8; j++) {
119            const v128_t v  = wasm_f32x4_mul(srcv[j], wasm_f32x4_splat(id));
120            const v128_t vi = wasm_i32x4_trunc_sat_f32x4(v);
121
122            y[i].qs[4*j + 0] = wasm_i32x4_extract_lane(vi, 0);
123            y[i].qs[4*j + 1] = wasm_i32x4_extract_lane(vi, 1);
124            y[i].qs[4*j + 2] = wasm_i32x4_extract_lane(vi, 2);
125            y[i].qs[4*j + 3] = wasm_i32x4_extract_lane(vi, 3);
126
127            accv = wasm_i32x4_add(accv, vi);
128        }
129
130        y[i].s = GGML_CPU_FP32_TO_FP16(
131                d * (wasm_i32x4_extract_lane(accv, 0) +
132                     wasm_i32x4_extract_lane(accv, 1) +
133                     wasm_i32x4_extract_lane(accv, 2) +
134                     wasm_i32x4_extract_lane(accv, 3)));
135    }
136#else
137    GGML_UNUSED(nb);
138    // scalar
139    quantize_row_q8_1_ref(x, y, k);
140#endif
141}
142
143//===================================== Q8_K ==============================================
144
145void quantize_row_q8_K(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k) {
146#ifdef __wasm_simd128__
147    assert(k % QK_K == 0);
148    const int64_t nb = k / QK_K;
149    block_q8_K * GGML_RESTRICT yc = y; // Cast to proper type
150
151    for (int i = 0; i < nb; i++) {
152        const float * x_block = x + i * QK_K;
153
154        v128_t min_vec = wasm_v128_load(x_block);
155        v128_t max_vec = min_vec;
156
157        for (int j = 4; j < QK_K; j += 4) {
158            v128_t x_vec = wasm_v128_load(x_block + j);
159            max_vec = wasm_f32x4_pmax(max_vec, x_vec);
160            min_vec = wasm_f32x4_pmin(min_vec, x_vec);
161        }
162        max_vec = wasm_f32x4_pmax(max_vec, wasm_i32x4_shuffle(max_vec, max_vec, 2, 3, 0, 1));
163        max_vec = wasm_f32x4_pmax(max_vec, wasm_i32x4_shuffle(max_vec, max_vec, 1, 0, 3, 2));
164        min_vec = wasm_f32x4_pmin(min_vec, wasm_i32x4_shuffle(min_vec, min_vec, 2, 3, 0, 1));
165        min_vec = wasm_f32x4_pmin(min_vec, wasm_i32x4_shuffle(min_vec, min_vec, 1, 0, 3, 2));
166        float max = wasm_f32x4_extract_lane(max_vec, 0);
167        float min = wasm_f32x4_extract_lane(min_vec, 0);
168        float amax = -min > max ? min : max;
169
170        if (amax == 0.0f) {
171            yc[i].d = 0.0f;
172            const v128_t zero = wasm_i8x16_splat(0);
173            for (int j = 0; j < QK_K; j += 16) {
174                wasm_v128_store(yc[i].qs + j, zero);
175            }
176            continue;
177        }
178
179        const float iscale = -127.0f / amax;
180        const v128_t scale_vec = wasm_f32x4_splat(iscale);
181
182        // Process 16 elements per iteration
183        for (int j = 0, jb = 0; j < QK_K; j += 16, jb++) {
184            // Load and quantize 16 floats
185            v128_t x0 = wasm_v128_load(x_block + j);
186            v128_t x1 = wasm_v128_load(x_block + j + 4);
187            v128_t x2 = wasm_v128_load(x_block + j + 8);
188            v128_t x3 = wasm_v128_load(x_block + j + 12);
189
190            v128_t q0 = wasm_f32x4_nearest(wasm_f32x4_mul(x0, scale_vec));
191            v128_t q1 = wasm_f32x4_nearest(wasm_f32x4_mul(x1, scale_vec));
192            v128_t q2 = wasm_f32x4_nearest(wasm_f32x4_mul(x2, scale_vec));
193            v128_t q3 = wasm_f32x4_nearest(wasm_f32x4_mul(x3, scale_vec));
194
195            // Convert to i32 with saturation
196            v128_t i0 = wasm_i32x4_trunc_sat_f32x4(q0);
197            v128_t i1 = wasm_i32x4_trunc_sat_f32x4(q1);
198            v128_t i2 = wasm_i32x4_trunc_sat_f32x4(q2);
199            v128_t i3 = wasm_i32x4_trunc_sat_f32x4(q3);
200
201            // Pack into 16 i8 values
202            v128_t i8 = wasm_i8x16_narrow_i16x8(
203                wasm_i16x8_narrow_i32x4(i0, i1),
204                wasm_i16x8_narrow_i32x4(i2, i3)
205            );
206            wasm_v128_store(yc[i].qs + j, i8);
207
208            // Calculate bsums using SIMD
209            v128_t sum16 = wasm_i16x8_add(
210                wasm_i16x8_extend_low_i8x16(i8),
211                wasm_i16x8_extend_high_i8x16(i8)
212            );
213            v128_t sum32 = wasm_i32x4_add(
214                wasm_i32x4_extend_low_i16x8(sum16),
215                wasm_i32x4_extend_high_i16x8(sum16)
216            );
217            sum32 = wasm_i32x4_add(sum32, wasm_i32x4_shuffle(sum32, sum32, 2, 3, 0, 1));
218            sum32 = wasm_i32x4_add(sum32, wasm_i32x4_shuffle(sum32, sum32, 1, 0, 3, 2));
219            yc[i].bsums[jb] = wasm_i32x4_extract_lane(sum32, 0);
220        }
221
222        yc[i].d = 1.0f / iscale;
223    }
224#else
225    quantize_row_q8_K_ref(x, y, k);
226#endif
227}
228
229
230//===================================== Dot products =================================
231
232void 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) {
233    const int qk = QK8_0;
234    const int nb = n / qk;
235
236    assert(n % qk == 0);
237    assert(nrc == 1);
238    UNUSED(nrc);
239    UNUSED(bx);
240    UNUSED(by);
241    UNUSED(bs);
242
243    const block_q4_0 * GGML_RESTRICT x = vx;
244    const block_q8_0 * GGML_RESTRICT y = vy;
245
246    int ib = 0;
247    float sumf = 0;
248
249#if defined __wasm_simd128__
250    v128_t sumv = wasm_f32x4_splat(0.0f);
251
252    const v128_t m4b = wasm_i8x16_splat(0x0F);
253    const v128_t s8b = wasm_i8x16_splat(0x8);
254
255    for (; ib + 1 < nb; ib += 2) {
256        const block_q4_0 * GGML_RESTRICT x0 = &x[ib];
257        const block_q4_0 * GGML_RESTRICT x1 = &x[ib + 1];
258        const block_q8_0 * GGML_RESTRICT y0 = &y[ib];
259        const block_q8_0 * GGML_RESTRICT y1 = &y[ib + 1];
260
261        // Load and process x0
262        v128_t v0_0 = wasm_v128_load(x0->qs);
263        v128_t v0_0l = wasm_v128_and(v0_0, m4b);
264        v128_t v0_0h = wasm_u8x16_shr(v0_0, 4);
265        v128_t v0_0ls = wasm_i8x16_sub(v0_0l, s8b);
266        v128_t v0_0hs = wasm_i8x16_sub(v0_0h, s8b);
267
268        // Load y0 vectors
269        v128_t y0_l = wasm_v128_load(y0->qs);
270        v128_t y0_h = wasm_v128_load(y0->qs + 16);
271
272        // Extend to i16x8 and compute dot products
273        v128_t dx0l = wasm_i16x8_extend_low_i8x16(v0_0ls);
274        v128_t dx0h = wasm_i16x8_extend_high_i8x16(v0_0ls);
275        v128_t dx0hl = wasm_i16x8_extend_low_i8x16(v0_0hs);
276        v128_t dx0hh = wasm_i16x8_extend_high_i8x16(v0_0hs);
277
278        v128_t dy0ll = wasm_i16x8_extend_low_i8x16(y0_l);
279        v128_t dy0lh = wasm_i16x8_extend_high_i8x16(y0_l);
280        v128_t dy0hl = wasm_i16x8_extend_low_i8x16(y0_h);
281        v128_t dy0hh = wasm_i16x8_extend_high_i8x16(y0_h);
282
283        v128_t dp0 = wasm_i32x4_add(
284            wasm_i32x4_add(
285                wasm_i32x4_dot_i16x8(dx0l, dy0ll),
286                wasm_i32x4_dot_i16x8(dx0h, dy0lh)
287            ),
288            wasm_i32x4_add(
289                wasm_i32x4_dot_i16x8(dx0hl, dy0hl),
290                wasm_i32x4_dot_i16x8(dx0hh, dy0hh)
291            )
292        );
293
294        // Load and process x1
295        v128_t v0_1 = wasm_v128_load(x1->qs);
296        v128_t v0_1l = wasm_v128_and(v0_1, m4b);
297        v128_t v0_1h = wasm_u8x16_shr(v0_1, 4);
298        v128_t v0_1ls = wasm_i8x16_sub(v0_1l, s8b);
299        v128_t v0_1hs = wasm_i8x16_sub(v0_1h, s8b);
300
301        // Load y1 vectors
302        v128_t y1_l = wasm_v128_load(y1->qs);
303        v128_t y1_h = wasm_v128_load(y1->qs + 16);
304
305        // Extend to i16x8 and compute dot products
306        v128_t dx1l = wasm_i16x8_extend_low_i8x16(v0_1ls);
307        v128_t dx1h = wasm_i16x8_extend_high_i8x16(v0_1ls);
308        v128_t dx1hl = wasm_i16x8_extend_low_i8x16(v0_1hs);
309        v128_t dx1hh = wasm_i16x8_extend_high_i8x16(v0_1hs);
310
311        v128_t dy1ll = wasm_i16x8_extend_low_i8x16(y1_l);
312        v128_t dy1lh = wasm_i16x8_extend_high_i8x16(y1_l);
313        v128_t dy1hl = wasm_i16x8_extend_low_i8x16(y1_h);
314        v128_t dy1hh = wasm_i16x8_extend_high_i8x16(y1_h);
315
316        v128_t dp1 = wasm_i32x4_add(
317            wasm_i32x4_add(
318                wasm_i32x4_dot_i16x8(dx1l, dy1ll),
319                wasm_i32x4_dot_i16x8(dx1h, dy1lh)
320            ),
321            wasm_i32x4_add(
322                wasm_i32x4_dot_i16x8(dx1hl, dy1hl),
323                wasm_i32x4_dot_i16x8(dx1hh, dy1hh)
324            )
325        );
326
327        // Accumulate results with scaling
328        float scale0 = GGML_CPU_FP16_TO_FP32(x0->d) * GGML_CPU_FP16_TO_FP32(y0->d);
329        float scale1 = GGML_CPU_FP16_TO_FP32(x1->d) * GGML_CPU_FP16_TO_FP32(y1->d);
330
331        sumv = wasm_f32x4_add(sumv, wasm_f32x4_mul(wasm_f32x4_convert_i32x4(dp0), wasm_f32x4_splat(scale0)));
332        sumv = wasm_f32x4_add(sumv, wasm_f32x4_mul(wasm_f32x4_convert_i32x4(dp1), wasm_f32x4_splat(scale1)));
333    }
334
335    sumf = wasm_f32x4_extract_lane(sumv, 0) + wasm_f32x4_extract_lane(sumv, 1) +
336           wasm_f32x4_extract_lane(sumv, 2) + wasm_f32x4_extract_lane(sumv, 3);
337
338#endif
339    for (; ib < nb; ++ib) {
340        int sumi0 = 0;
341        int sumi1 = 0;
342
343        for (int j = 0; j < qk/2; ++j) {
344            const int v0 = (x[ib].qs[j] & 0x0F) - 8;
345            const int v1 = (x[ib].qs[j] >>   4) - 8;
346
347            sumi0 += (v0 * y[ib].qs[j]);
348            sumi1 += (v1 * y[ib].qs[j + qk/2]);
349        }
350
351        int sumi = sumi0 + sumi1;
352        sumf += sumi*GGML_CPU_FP16_TO_FP32(x[ib].d)*GGML_CPU_FP16_TO_FP32(y[ib].d);
353    }
354
355    *s = sumf;
356}
357
358void ggml_vec_dot_q4_1_q8_1(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) {
359    const int qk = QK8_1;
360    const int nb = n / qk;
361
362    assert(n % qk == 0);
363    assert(nrc == 1);
364    UNUSED(nrc);
365    UNUSED(bx);
366    UNUSED(by);
367    UNUSED(bs);
368
369    const block_q4_1 * GGML_RESTRICT x = vx;
370    const block_q8_1 * GGML_RESTRICT y = vy;
371
372    float sumf = 0;
373
374#if defined __wasm_simd128__
375    v128_t sumv = wasm_f32x4_splat(0.0f);
376    float summs = 0.0f;
377
378    for (int ib = 0; ib < nb; ++ib) {
379        const block_q4_1 * GGML_RESTRICT x0 = &x[ib];
380        const block_q8_1 * GGML_RESTRICT y0 = &y[ib];
381
382        summs += GGML_CPU_FP16_TO_FP32(x0->m) * GGML_CPU_FP16_TO_FP32(y0->s);
383
384        const v128_t raw  = wasm_v128_load(x0->qs);
385        const v128_t v0s  = wasm_v128_and(raw, wasm_i8x16_splat(0x0F));
386        const v128_t v1s  = wasm_u8x16_shr(raw, 4);
387
388        const v128_t ys_lo = wasm_v128_load(y0->qs);
389        const v128_t ys_hi = wasm_v128_load(y0->qs + 16);
390
391        const v128_t v0s_l = wasm_u16x8_extend_low_u8x16(v0s);
392        const v128_t v0s_h = wasm_u16x8_extend_high_u8x16(v0s);
393        const v128_t ylo_l = wasm_i16x8_extend_low_i8x16(ys_lo);
394        const v128_t ylo_h = wasm_i16x8_extend_high_i8x16(ys_lo);
395        const v128_t v1s_l = wasm_u16x8_extend_low_u8x16(v1s);
396        const v128_t v1s_h = wasm_u16x8_extend_high_u8x16(v1s);
397        const v128_t yhi_l = wasm_i16x8_extend_low_i8x16(ys_hi);
398        const v128_t yhi_h = wasm_i16x8_extend_high_i8x16(ys_hi);
399
400        const v128_t acc = wasm_i32x4_add(
401            wasm_i32x4_add(
402                wasm_i32x4_dot_i16x8(v0s_l, ylo_l),
403                wasm_i32x4_dot_i16x8(v0s_h, ylo_h)),
404            wasm_i32x4_add(
405                wasm_i32x4_dot_i16x8(v1s_l, yhi_l),
406                wasm_i32x4_dot_i16x8(v1s_h, yhi_h)));
407
408        sumv = wasm_f32x4_add(sumv,
409            wasm_f32x4_mul(
410                wasm_f32x4_convert_i32x4(acc),
411                wasm_f32x4_splat(GGML_CPU_FP16_TO_FP32(x0->d) * GGML_CPU_FP16_TO_FP32(y0->d))));
412    }
413
414    sumf = wasm_f32x4_extract_lane(sumv, 0) + wasm_f32x4_extract_lane(sumv, 1) +
415           wasm_f32x4_extract_lane(sumv, 2) + wasm_f32x4_extract_lane(sumv, 3) + summs;
416
417    *s = sumf;
418
419#else
420    UNUSED(nb);
421    UNUSED(x);
422    UNUSED(y);
423    UNUSED(sumf);
424
425    ggml_vec_dot_q4_1_q8_1_generic(
426        n, s, bs, vx, bx, vy, by, nrc);
427#endif
428}
429
430void ggml_vec_dot_q5_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) {
431    const int qk = QK8_0;
432    const int nb = n / qk;
433
434    int ib = 0;
435    float sumf = 0;
436
437    assert(n % qk == 0);
438    assert(qk == QK5_0);
439    assert(nrc == 1);
440    UNUSED(nrc);
441    UNUSED(bx);
442    UNUSED(by);
443    UNUSED(bs);
444
445    const block_q5_0 * GGML_RESTRICT x = vx;
446    const block_q8_0 * GGML_RESTRICT y = vy;
447
448#if defined __wasm_simd128__
449    v128_t sumv = wasm_f32x4_splat(0.0f);
450
451    uint32_t qh_;
452    uint64_t tmp[4];
453
454    // TODO: check if unrolling this is better
455    for (; ib < nb; ++ib) {
456        const block_q5_0 * GGML_RESTRICT x0 = &x[ib];
457        const block_q8_0 * GGML_RESTRICT y0 = &y[ib];
458
459        const v128_t m4b  = wasm_i8x16_splat(0x0F);
460
461        // extract the 5th bit
462        memcpy(&qh_, x0->qh, sizeof(qh_));
463
464        tmp[0] = table_b2b_1[(qh_ >>  0) & 0xFF];
465        tmp[1] = table_b2b_1[(qh_ >>  8) & 0xFF];
466        tmp[2] = table_b2b_1[(qh_ >> 16) & 0xFF];
467        tmp[3] = table_b2b_1[(qh_ >> 24)       ];
468
469        const v128_t qhl = wasm_v128_load(tmp + 0);
470        const v128_t qhh = wasm_v128_load(tmp + 2);
471
472        const v128_t v0 = wasm_v128_load(x0->qs);
473
474        // 4-bit -> 8-bit
475        const v128_t v0l = wasm_v128_and (v0, m4b);
476        const v128_t v0h = wasm_u8x16_shr(v0, 4);
477
478        // add high bit and sub 16 (equivalent to sub 0x10 when bit is zero)
479        const v128_t v0lf = wasm_i8x16_sub(v0l, qhl);
480        const v128_t v0hf = wasm_i8x16_sub(v0h, qhh);
481
482        // load y
483        const v128_t v1l = wasm_v128_load(y0->qs);
484        const v128_t v1h = wasm_v128_load(y0->qs + 16);
485
486        // int8x16 -> int16x8
487        const v128_t v0lfl = wasm_i16x8_extend_low_i8x16 (v0lf);
488        const v128_t v0lfh = wasm_i16x8_extend_high_i8x16(v0lf);
489        const v128_t v0hfl = wasm_i16x8_extend_low_i8x16 (v0hf);
490        const v128_t v0hfh = wasm_i16x8_extend_high_i8x16(v0hf);
491
492        const v128_t v1ll = wasm_i16x8_extend_low_i8x16 (v1l);
493        const v128_t v1lh = wasm_i16x8_extend_high_i8x16(v1l);
494        const v128_t v1hl = wasm_i16x8_extend_low_i8x16 (v1h);
495        const v128_t v1hh = wasm_i16x8_extend_high_i8x16(v1h);
496
497        // dot product
498        sumv = wasm_f32x4_add(sumv, wasm_f32x4_mul(wasm_f32x4_convert_i32x4(
499                        wasm_i32x4_add(
500                            wasm_i32x4_add(wasm_i32x4_dot_i16x8(v0lfl, v1ll),
501                                           wasm_i32x4_dot_i16x8(v0lfh, v1lh)),
502                            wasm_i32x4_add(wasm_i32x4_dot_i16x8(v0hfl, v1hl),
503                                           wasm_i32x4_dot_i16x8(v0hfh, v1hh)))),
504                    wasm_f32x4_splat(GGML_CPU_FP16_TO_FP32(x0->d) * GGML_CPU_FP16_TO_FP32(y0->d))));
505    }
506
507    sumf = wasm_f32x4_extract_lane(sumv, 0) + wasm_f32x4_extract_lane(sumv, 1) +
508           wasm_f32x4_extract_lane(sumv, 2) + wasm_f32x4_extract_lane(sumv, 3);
509
510    *s = sumf;
511#else
512    UNUSED(nb);
513    UNUSED(ib);
514    UNUSED(sumf);
515    UNUSED(x);
516    UNUSED(y);
517    ggml_vec_dot_q5_0_q8_0_generic(n, s, bs, vx, bx, vy, by, nrc);
518#endif
519}
520
521void ggml_vec_dot_q5_1_q8_1(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) {
522    const int qk = QK8_1;
523    const int nb = n / qk;
524
525    int ib = 0;
526    float sumf = 0;
527
528    assert(n % qk == 0);
529    assert(qk == QK5_1);
530    assert(nrc == 1);
531    UNUSED(nrc);
532    UNUSED(bx);
533    UNUSED(by);
534    UNUSED(bs);
535
536    const block_q5_1 * GGML_RESTRICT x = vx;
537    const block_q8_1 * GGML_RESTRICT y = vy;
538
539#if defined __wasm_simd128__
540    v128_t sumv = wasm_f32x4_splat(0.0f);
541
542    float summs = 0.0f;
543
544    uint32_t qh_;
545    uint64_t tmp[4];
546
547    // TODO: check if unrolling this is better
548    for (; ib < nb; ++ib) {
549        const block_q5_1 * GGML_RESTRICT x0 = &x[ib];
550        const block_q8_1 * GGML_RESTRICT y0 = &y[ib];
551
552        summs += GGML_CPU_FP16_TO_FP32(x0->m) * GGML_CPU_FP16_TO_FP32(y0->s);
553
554        const v128_t m4b = wasm_i8x16_splat(0x0F);
555
556        // extract the 5th bit
557        memcpy(&qh_, x0->qh, sizeof(qh_));
558
559        tmp[0] = table_b2b_0[(qh_ >>  0) & 0xFF];
560        tmp[1] = table_b2b_0[(qh_ >>  8) & 0xFF];
561        tmp[2] = table_b2b_0[(qh_ >> 16) & 0xFF];
562        tmp[3] = table_b2b_0[(qh_ >> 24)       ];
563
564        const v128_t qhl = wasm_v128_load(tmp + 0);
565        const v128_t qhh = wasm_v128_load(tmp + 2);
566
567        const v128_t v0 = wasm_v128_load(x0->qs);
568
569        // 4-bit -> 8-bit
570        const v128_t v0l = wasm_v128_and (v0, m4b);
571        const v128_t v0h = wasm_u8x16_shr(v0, 4);
572
573        // add high bit
574        const v128_t v0lf = wasm_v128_or(v0l, qhl);
575        const v128_t v0hf = wasm_v128_or(v0h, qhh);
576
577        // load y
578        const v128_t v1l = wasm_v128_load(y0->qs);
579        const v128_t v1h = wasm_v128_load(y0->qs + 16);
580
581        // int8x16 -> int16x8
582        const v128_t v0lfl = wasm_i16x8_extend_low_i8x16 (v0lf);
583        const v128_t v0lfh = wasm_i16x8_extend_high_i8x16(v0lf);
584        const v128_t v0hfl = wasm_i16x8_extend_low_i8x16 (v0hf);
585        const v128_t v0hfh = wasm_i16x8_extend_high_i8x16(v0hf);
586
587        const v128_t v1ll = wasm_i16x8_extend_low_i8x16 (v1l);
588        const v128_t v1lh = wasm_i16x8_extend_high_i8x16(v1l);
589        const v128_t v1hl = wasm_i16x8_extend_low_i8x16 (v1h);
590        const v128_t v1hh = wasm_i16x8_extend_high_i8x16(v1h);
591
592        // dot product
593        sumv = wasm_f32x4_add(sumv,
594                wasm_f32x4_mul(wasm_f32x4_convert_i32x4(wasm_i32x4_add(
595                            wasm_i32x4_add(wasm_i32x4_dot_i16x8(v0lfl, v1ll),
596                                           wasm_i32x4_dot_i16x8(v0lfh, v1lh)),
597                            wasm_i32x4_add(wasm_i32x4_dot_i16x8(v0hfl, v1hl),
598                                           wasm_i32x4_dot_i16x8(v0hfh, v1hh)))),
599                    wasm_f32x4_splat(GGML_CPU_FP16_TO_FP32(x0->d) * GGML_CPU_FP16_TO_FP32(y0->d))));
600    }
601
602    sumf = wasm_f32x4_extract_lane(sumv, 0) + wasm_f32x4_extract_lane(sumv, 1) +
603           wasm_f32x4_extract_lane(sumv, 2) + wasm_f32x4_extract_lane(sumv, 3) + summs;
604
605    *s = sumf;
606#else
607    UNUSED(nb);
608    UNUSED(ib);
609    UNUSED(sumf);
610    UNUSED(x);
611    UNUSED(y);
612    ggml_vec_dot_q5_1_q8_1_generic(n, s, bs, vx, bx, vy, by, nrc);
613#endif
614}
615
616void ggml_vec_dot_q8_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) {
617    const int qk = QK8_0;
618    const int nb = n / qk;
619
620    assert(n % qk == 0);
621    assert(nrc == 1);
622    UNUSED(nrc);
623    UNUSED(bx);
624    UNUSED(by);
625    UNUSED(bs);
626
627    const block_q8_0 * GGML_RESTRICT x = vx;
628    const block_q8_0 * GGML_RESTRICT y = vy;
629
630    int ib = 0;
631    float sumf = 0;
632
633#if defined __wasm_simd128__
634    v128_t sumv = wasm_f32x4_splat(0.0f);
635
636    for (; ib < nb; ++ib) {
637        const block_q8_0 * GGML_RESTRICT x0 = &x[ib];
638        const block_q8_0 * GGML_RESTRICT y0 = &y[ib];
639
640        const v128_t x0_0 = wasm_v128_load(x0->qs);
641        const v128_t x0_1 = wasm_v128_load(x0->qs + 16);
642        const v128_t y0_0 = wasm_v128_load(y0->qs);
643        const v128_t y0_1 = wasm_v128_load(y0->qs + 16);
644
645        // Extend 8-bit to 16-bit
646        const v128_t x0_0l = wasm_i16x8_extend_low_i8x16(x0_0);
647        const v128_t x0_0h = wasm_i16x8_extend_high_i8x16(x0_0);
648        const v128_t x0_1l = wasm_i16x8_extend_low_i8x16(x0_1);
649        const v128_t x0_1h = wasm_i16x8_extend_high_i8x16(x0_1);
650
651        const v128_t y0_0l = wasm_i16x8_extend_low_i8x16(y0_0);
652        const v128_t y0_0h = wasm_i16x8_extend_high_i8x16(y0_0);
653        const v128_t y0_1l = wasm_i16x8_extend_low_i8x16(y0_1);
654        const v128_t y0_1h = wasm_i16x8_extend_high_i8x16(y0_1);
655
656        // Compute dot products
657        const v128_t dx0_0 = wasm_i32x4_dot_i16x8(x0_0l, y0_0l);
658        const v128_t dx0_1 = wasm_i32x4_dot_i16x8(x0_0h, y0_0h);
659        const v128_t dx1_0 = wasm_i32x4_dot_i16x8(x0_1l, y0_1l);
660        const v128_t dx1_1 = wasm_i32x4_dot_i16x8(x0_1h, y0_1h);
661
662        // Sum all dot products
663        const v128_t sum_dots = wasm_i32x4_add(wasm_i32x4_add(dx0_0, dx0_1), wasm_i32x4_add(dx1_0, dx1_1));
664
665        // Convert to float and accumulate
666        const float scale = GGML_CPU_FP16_TO_FP32(x0->d) * GGML_CPU_FP16_TO_FP32(y0->d);
667        sumv = wasm_f32x4_add(sumv, wasm_f32x4_mul(wasm_f32x4_convert_i32x4(sum_dots), wasm_f32x4_splat(scale)));
668    }
669
670    sumf = wasm_f32x4_extract_lane(sumv, 0) + wasm_f32x4_extract_lane(sumv, 1) +
671           wasm_f32x4_extract_lane(sumv, 2) + wasm_f32x4_extract_lane(sumv, 3);
672
673    *s = sumf;
674#else
675    UNUSED(nb);
676    UNUSED(x);
677    UNUSED(y);
678    UNUSED(ib);
679    UNUSED(sumf);
680    ggml_vec_dot_q8_0_q8_0_generic(n, s, bs, vx, bx, vy, by, nrc);
681#endif
682}
683
684void ggml_vec_dot_q2_K_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) {
685    assert(nrc == 1);
686    UNUSED(nrc);
687    UNUSED(bx);
688    UNUSED(by);
689    UNUSED(bs);
690
691    const block_q2_K * GGML_RESTRICT x = vx;
692    const block_q8_K * GGML_RESTRICT y = vy;
693
694    const int nb = n / QK_K;
695
696#if defined __wasm_simd128__
697    float sumf = 0;
698
699    for (int i = 0; i < nb; ++i) {
700        const uint8_t * q2 = x[i].qs;
701        const int8_t * q8 = y[i].qs;
702        const uint8_t * sc = x[i].scales;
703
704        // Vectorized summs calculation
705        v128_t summs_vec = wasm_i32x4_splat(0);
706        {
707            v128_t sc_vec = wasm_v128_load(sc);
708            v128_t sc_upper = wasm_u8x16_shr(sc_vec, 4);
709
710            v128_t sc_low = wasm_u16x8_extend_low_u8x16(sc_upper);
711            v128_t sc_high = wasm_u16x8_extend_high_u8x16(sc_upper);
712
713            v128_t bsums1 = wasm_v128_load(&y[i].bsums[0]);
714            v128_t bsums2 = wasm_v128_load(&y[i].bsums[8]);
715
716            summs_vec = wasm_i32x4_add(
717                wasm_i32x4_add(wasm_i32x4_dot_i16x8(sc_low, bsums1),
718                               wasm_i32x4_dot_i16x8(sc_high, bsums2)),
719                summs_vec
720            );
721
722            summs_vec = wasm_i32x4_add(summs_vec, wasm_i32x4_shuffle(summs_vec, summs_vec, 2, 3, 0, 1));
723            summs_vec = wasm_i32x4_add(summs_vec, wasm_i32x4_shuffle(summs_vec, summs_vec, 1, 0, 3, 2));
724        }
725        int32_t summs = wasm_i32x4_extract_lane(summs_vec, 0);
726
727        // Vectorized isum calculation
728        int32_t isum = 0;
729        const uint8_t * sc_ptr = sc;
730        const int k_iters = QK_K/128;
731
732        for (int k = 0; k < k_iters; ++k) {
733            v128_t isum_vec = wasm_i32x4_splat(0);
734            int shift = 0;
735
736            for (int j = 0; j < 4; ++j) {
737                const int d0 = (sc_ptr[0] & 0xF);
738                const int d1 = (sc_ptr[1] & 0xF);
739                sc_ptr += 2;
740
741                // Process first 16 elements
742                v128_t q2_0 = wasm_v128_load(q2);
743                v128_t q8_0 = wasm_v128_load(q8);
744                v128_t q2_shift_0 = wasm_u8x16_shr(q2_0, shift);
745                v128_t q2_bits_0 = wasm_v128_and(q2_shift_0, wasm_i8x16_splat(0x03));
746
747                // Process next 16 elements
748                v128_t q2_1 = wasm_v128_load(q2 + 16);
749                v128_t q8_1 = wasm_v128_load(q8 + 16);
750                v128_t q2_shift_1 = wasm_u8x16_shr(q2_1, shift);
751                v128_t q2_bits_1 = wasm_v128_and(q2_shift_1, wasm_i8x16_splat(0x03));
752
753                // Calculate dot products
754                v128_t p0 = wasm_i32x4_dot_i16x8(
755                    wasm_i16x8_extend_low_i8x16(q8_0),
756                    wasm_i16x8_extend_low_i8x16(q2_bits_0)
757                );
758                v128_t p1 = wasm_i32x4_dot_i16x8(
759                    wasm_i16x8_extend_high_i8x16(q8_0),
760                    wasm_i16x8_extend_high_i8x16(q2_bits_0)
761                );
762                v128_t p2 = wasm_i32x4_dot_i16x8(
763                    wasm_i16x8_extend_low_i8x16(q8_1),
764                    wasm_i16x8_extend_low_i8x16(q2_bits_1)
765                );
766                v128_t p3 = wasm_i32x4_dot_i16x8(
767                    wasm_i16x8_extend_high_i8x16(q8_1),
768                    wasm_i16x8_extend_high_i8x16(q2_bits_1)
769                );
770
771                // Accumulate scaled results
772                v128_t scaled = wasm_i32x4_add(
773                    wasm_i32x4_mul(wasm_i32x4_add(p0, p1), wasm_i32x4_splat(d0)),
774                    wasm_i32x4_mul(wasm_i32x4_add(p2, p3), wasm_i32x4_splat(d1))
775                );
776
777                isum_vec = wasm_i32x4_add(isum_vec, scaled);
778                q8 += 32;
779                shift += 2;
780            }
781            q2 += 32;
782
783            // Horizontal sum of isum_vec
784            isum_vec = wasm_i32x4_add(isum_vec, wasm_i32x4_shuffle(isum_vec, isum_vec, 2, 3, 0, 1));
785            isum_vec = wasm_i32x4_add(isum_vec, wasm_i32x4_shuffle(isum_vec, isum_vec, 1, 0, 3, 2));
786            isum += wasm_i32x4_extract_lane(isum_vec, 0);
787        }
788
789        const float dall = GGML_CPU_FP16_TO_FP32(x[i].d) * y[i].d;
790        const float dmin = GGML_CPU_FP16_TO_FP32(x[i].dmin) * y[i].d;
791        sumf += dall * isum - dmin * summs;
792    }
793
794    *s = sumf;
795
796#else
797    UNUSED(x);
798    UNUSED(y);
799    UNUSED(nb);
800    ggml_vec_dot_q2_K_q8_K_generic(n, s, bs, vx, bx, vy, by, nrc);
801#endif
802}
803
804void ggml_vec_dot_q3_K_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) {
805    assert(n % QK_K == 0);
806    assert(nrc == 1);
807    UNUSED(nrc);
808    UNUSED(bx);
809    UNUSED(by);
810    UNUSED(bs);
811
812    const uint32_t kmask1 = 0x03030303;
813    const uint32_t kmask2 = 0x0f0f0f0f;
814
815    const block_q3_K * GGML_RESTRICT x = vx;
816    const block_q8_K * GGML_RESTRICT y = vy;
817
818    const int nb = n / QK_K;
819
820#if defined __wasm_simd128__
821    int8_t  aux8[QK_K];
822    float   sums[8] = {0};
823    uint32_t auxs[4];
824
825    float sumf = 0;
826    for (int i = 0; i < nb; ++i) {
827        const uint8_t * GGML_RESTRICT q3 = x[i].qs;
828        const uint8_t * GGML_RESTRICT hm = x[i].hmask;
829        const int8_t  * GGML_RESTRICT q8 = y[i].qs;
830
831        // Process blocks with SIMD
832        int8_t * a = aux8;
833        uint8_t m = 1;
834        for (int j = 0; j < QK_K; j += 128) {
835            for (int shift = 0; shift <= 6; shift += 2) {
836                v128_t v_m = wasm_i8x16_splat(m);
837                for (int l = 0; l < 32; l += 16) {
838                    v128_t v_q3 = wasm_v128_load(q3 + l);
839                    v128_t v_shift = wasm_i8x16_shr(v_q3, shift);
840                    v128_t v_low2 = wasm_v128_and(v_shift, wasm_i8x16_splat(0x03));
841
842                    v128_t v_hm = wasm_v128_load(hm + l);
843                    v128_t v_mask = wasm_v128_and(v_hm, v_m);
844                    v_mask = wasm_i8x16_ne(v_mask, wasm_i8x16_splat(0));
845
846                    v_low2 = wasm_i8x16_sub(v_low2, wasm_v128_and(wasm_i8x16_splat(4), wasm_v128_not(v_mask)));
847                    wasm_v128_store(a + l, v_low2);
848                }
849                a += 32;
850                m <<= 1;
851            }
852            q3 += 32;
853        }
854
855        // Extract scales
856        memcpy(auxs, x[i].scales, 12);
857        uint32_t tmp = auxs[2];
858        auxs[2] = ((auxs[0] >> 4) & kmask2) | (((tmp >> 4) & kmask1) << 4);
859        auxs[3] = ((auxs[1] >> 4) & kmask2) | (((tmp >> 6) & kmask1) << 4);
860        auxs[0] = (auxs[0] & kmask2) | (((tmp >> 0) & kmask1) << 4);
861        auxs[1] = (auxs[1] & kmask2) | (((tmp >> 2) & kmask1) << 4);
862        const int8_t * scales = (const int8_t *)auxs;
863
864        // SIMD dot product with register accumulators
865        v128_t v_acc0 = wasm_i32x4_splat(0);
866        v128_t v_acc1 = wasm_i32x4_splat(0);
867        a = aux8;
868        for (int j = 0; j < QK_K/16; ++j) {
869            const v128_t v_scale = wasm_i16x8_splat(scales[j] - 32);
870
871            // Process 16 elements per iteration
872            for (int k = 0; k < 2; ++k) {
873                const v128_t v_q8 = wasm_i16x8_load8x8(q8);
874                const v128_t v_a = wasm_i16x8_load8x8(a);
875
876                v128_t v_prod = wasm_i16x8_mul(v_q8, v_a);
877                v_prod = wasm_i16x8_mul(v_prod, v_scale);
878
879                v_acc0 = wasm_i32x4_add(v_acc0, wasm_i32x4_extend_low_i16x8(v_prod));
880                v_acc1 = wasm_i32x4_add(v_acc1, wasm_i32x4_extend_high_i16x8(v_prod));
881
882                q8 += 8;
883                a += 8;
884            }
885        }
886
887        // Accumulate results
888        const float d = GGML_CPU_FP16_TO_FP32(x[i].d) * y[i].d;
889        const v128_t v_d = wasm_f32x4_splat(d);
890        v128_t v_sum = wasm_f32x4_add(
891            wasm_f32x4_mul(wasm_f32x4_convert_i32x4(v_acc0), v_d),
892            wasm_f32x4_mul(wasm_f32x4_convert_i32x4(v_acc1), v_d)
893        );
894
895        // Accumulate into sums vector
896        wasm_v128_store(sums, wasm_f32x4_add(wasm_v128_load(sums), v_sum));
897    }
898
899    // Horizontal sum
900    v128_t v_sum = wasm_f32x4_add(wasm_v128_load(sums), wasm_v128_load(sums + 4));
901    sumf = wasm_f32x4_extract_lane(v_sum, 0) +
902           wasm_f32x4_extract_lane(v_sum, 1) +
903           wasm_f32x4_extract_lane(v_sum, 2) +
904           wasm_f32x4_extract_lane(v_sum, 3);
905
906    *s = sumf;
907
908#else
909    UNUSED(kmask1);
910    UNUSED(kmask2);
911    UNUSED(x);
912    UNUSED(y);
913    UNUSED(nb);
914    ggml_vec_dot_q3_K_q8_K_generic(n, s, bs, vx, bx, vy, by, nrc);
915#endif
916
917}
918
919void ggml_vec_dot_q4_K_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) {
920    assert(n % QK_K == 0);
921    assert(nrc == 1);
922    UNUSED(nrc);
923    UNUSED(bx);
924    UNUSED(by);
925    UNUSED(bs);
926
927    const block_q4_K * GGML_RESTRICT x = vx;
928    const block_q8_K * GGML_RESTRICT y = vy;
929
930    const int nb = n / QK_K;
931
932    static const uint32_t kmask1 = 0x3f3f3f3f;
933    static const uint32_t kmask2 = 0x0f0f0f0f;
934    static const uint32_t kmask3 = 0x03030303;
935
936    uint32_t utmp[4];
937
938#if defined __wasm_simd128__
939    const uint8_t * scales = (const uint8_t*)&utmp[0];
940    float sumf = 0;
941
942    for (int i = 0; i < nb; ++i) {
943        const float d = y[i].d * GGML_CPU_FP16_TO_FP32(x[i].d);
944        const float dmin = y[i].d * GGML_CPU_FP16_TO_FP32(x[i].dmin); // Corrected sign
945
946        const uint8_t * GGML_RESTRICT q4 = x[i].qs;
947        const int8_t  * GGML_RESTRICT q8 = y[i].qs;
948
949        // Process scales and mins
950        memcpy(utmp, x[i].scales, 12);
951        utmp[3] = ((utmp[2] >> 4) & kmask2) | (((utmp[1] >> 6) & kmask3) << 4);
952        const uint32_t uaux = utmp[1] & kmask1;
953        utmp[1] = (utmp[2] & kmask2) | (((utmp[0] >> 6) & kmask3) << 4);
954        utmp[2] = uaux;
955        utmp[0] &= kmask1;
956
957        // Sum mins * q8sums
958        int32_t sumi = 0;
959        const int16_t * GGML_RESTRICT q8sums = y[i].bsums;
960        const uint8_t * m = (const uint8_t *)&utmp[2];
961        for (int j = 0; j < 16; j += 2) {
962            sumi += (q8sums[j] + q8sums[j+1]) * m[j/2];
963        }
964        sumf -= dmin * sumi;
965
966        int32_t sumi1 = 0;
967        int32_t sumi2 = 0;
968
969        for (int j = 0; j < QK_K/64; ++j) {
970            // Load 64 4-bit weights (32 bytes)
971            const v128_t q4x0 = wasm_v128_load(q4);
972            const v128_t q4x1 = wasm_v128_load(q4 + 16);
973            q4 += 32;
974
975            // Split into low/high nibbles
976            const v128_t q4l0 = wasm_v128_and(q4x0, wasm_i8x16_splat(0x0F));
977            const v128_t q4h0 = wasm_u8x16_shr(q4x0, 4);
978            const v128_t q4l1 = wasm_v128_and(q4x1, wasm_i8x16_splat(0x0F));
979            const v128_t q4h1 = wasm_u8x16_shr(q4x1, 4);
980
981            // Load 64 8-bit values (64 bytes)
982            const v128_t q8x0 = wasm_v128_load(q8);
983            const v128_t q8x1 = wasm_v128_load(q8 + 16);
984            const v128_t q8x2 = wasm_v128_load(q8 + 32);
985            const v128_t q8x3 = wasm_v128_load(q8 + 48);
986            q8 += 64;
987
988            // Low nibble products
989            v128_t vacc1 = wasm_i32x4_dot_i16x8(
990                wasm_i16x8_extend_low_i8x16(q4l0),
991                wasm_i16x8_extend_low_i8x16(q8x0)
992            );
993            vacc1 = wasm_i32x4_add(vacc1, wasm_i32x4_dot_i16x8(
994                wasm_i16x8_extend_high_i8x16(q4l0),
995                wasm_i16x8_extend_high_i8x16(q8x0)
996            ));
997            vacc1 = wasm_i32x4_add(vacc1, wasm_i32x4_dot_i16x8(
998                wasm_i16x8_extend_low_i8x16(q4l1),
999                wasm_i16x8_extend_low_i8x16(q8x1)
1000            ));
1001            vacc1 = wasm_i32x4_add(vacc1, wasm_i32x4_dot_i16x8(
1002                wasm_i16x8_extend_high_i8x16(q4l1),
1003                wasm_i16x8_extend_high_i8x16(q8x1)
1004            ));
1005
1006            // High nibble products
1007            v128_t vacc2 = wasm_i32x4_dot_i16x8(
1008                wasm_i16x8_extend_low_i8x16(q4h0),
1009                wasm_i16x8_extend_low_i8x16(q8x2)
1010            );
1011            vacc2 = wasm_i32x4_add(vacc2, wasm_i32x4_dot_i16x8(
1012                wasm_i16x8_extend_high_i8x16(q4h0),
1013                wasm_i16x8_extend_high_i8x16(q8x2)
1014            ));
1015            vacc2 = wasm_i32x4_add(vacc2, wasm_i32x4_dot_i16x8(
1016                wasm_i16x8_extend_low_i8x16(q4h1),
1017                wasm_i16x8_extend_low_i8x16(q8x3)
1018            ));
1019            vacc2 = wasm_i32x4_add(vacc2, wasm_i32x4_dot_i16x8(
1020                wasm_i16x8_extend_high_i8x16(q4h1),
1021                wasm_i16x8_extend_high_i8x16(q8x3)
1022            ));
1023
1024            // Accumulate scaled results
1025            int32_t vacc1_sum = wasm_i32x4_extract_lane(vacc1, 0) + wasm_i32x4_extract_lane(vacc1, 1) +
1026                                wasm_i32x4_extract_lane(vacc1, 2) + wasm_i32x4_extract_lane(vacc1, 3);
1027            sumi1 += vacc1_sum * scales[2*j];
1028
1029            int32_t vacc2_sum = wasm_i32x4_extract_lane(vacc2, 0) + wasm_i32x4_extract_lane(vacc2, 1) +
1030                                wasm_i32x4_extract_lane(vacc2, 2) + wasm_i32x4_extract_lane(vacc2, 3);
1031            sumi2 += vacc2_sum * scales[2*j+1];
1032        }
1033
1034        sumf += d * (sumi1 + sumi2);
1035    }
1036
1037    *s = sumf;
1038
1039#else
1040    UNUSED(x);
1041    UNUSED(y);
1042    UNUSED(nb);
1043    UNUSED(kmask1);
1044    UNUSED(kmask2);
1045    UNUSED(kmask3);
1046    UNUSED(utmp);
1047    ggml_vec_dot_q4_K_q8_K_generic(n, s, bs, vx, bx, vy, by, nrc);
1048#endif
1049}
1050
1051void ggml_vec_dot_q5_K_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) {
1052    assert(n % QK_K == 0);
1053    assert(nrc == 1);
1054    UNUSED(nrc);
1055    UNUSED(bx);
1056    UNUSED(by);
1057    UNUSED(bs);
1058
1059    const block_q5_K * GGML_RESTRICT x = vx;
1060    const block_q8_K * GGML_RESTRICT y = vy;
1061
1062    const int nb = n / QK_K;
1063
1064    static const uint32_t kmask1 = 0x3f3f3f3f;
1065    static const uint32_t kmask2 = 0x0f0f0f0f;
1066    static const uint32_t kmask3 = 0x03030303;
1067
1068    uint32_t utmp[4];
1069
1070#if defined __wasm_simd128__
1071    //const uint8_t * scales = (const uint8_t*)&utmp[0];
1072    float sumf = 0;
1073
1074    for (int i = 0; i < nb; ++i) {
1075        const float d = y[i].d * GGML_CPU_FP16_TO_FP32(x[i].d);
1076        const float dmin = y[i].d * GGML_CPU_FP16_TO_FP32(x[i].dmin); // Fixed sign
1077
1078        const uint8_t * GGML_RESTRICT q5 = x[i].qs;
1079        const uint8_t * GGML_RESTRICT qh = x[i].qh;
1080        const int8_t  * GGML_RESTRICT q8 = y[i].qs;
1081
1082        // Process scales and mins
1083        memcpy(utmp, x[i].scales, 12);
1084        utmp[3] = ((utmp[2] >> 4) & kmask2) | (((utmp[1] >> 6) & kmask3) << 4);
1085        const uint32_t uaux = utmp[1] & kmask1;
1086        utmp[1] = (utmp[2] & kmask2) | (((utmp[0] >> 6) & kmask3) << 4);
1087        utmp[2] = uaux;
1088        utmp[0] &= kmask1;
1089
1090        // Sum mins * q8sums
1091        int32_t sumi_mins = 0;
1092        const int16_t * GGML_RESTRICT q8sums = y[i].bsums;
1093        const uint8_t * m = (const uint8_t *)&utmp[2];
1094        for (int j = 0; j < 16; j += 2) {
1095            sumi_mins += (q8sums[j] + q8sums[j+1]) * m[j/2];
1096        }
1097        sumf -= dmin * sumi_mins; // Correct subtraction
1098
1099        v128_t qh0 = wasm_v128_load(qh);
1100        v128_t qh1 = wasm_v128_load(qh + 16);
1101        const uint8_t * sc = (const uint8_t *)utmp;
1102
1103        int32_t sumi = 0;
1104
1105        for (int j = 0; j < QK_K/64; ++j) {
1106            const int shift = j * 2;
1107            v128_t qh_shift0 = wasm_u8x16_shr(qh0, shift);
1108            v128_t qh_shift1 = wasm_u8x16_shr(qh1, shift);
1109
1110            v128_t qh_low0 = wasm_i8x16_shl(wasm_v128_and(qh_shift0, wasm_i8x16_splat(0x01)), 4);
1111            v128_t qh_high0 = wasm_i8x16_shl(wasm_v128_and(qh_shift0, wasm_i8x16_splat(0x02)), 3);
1112            v128_t qh_low1 = wasm_i8x16_shl(wasm_v128_and(qh_shift1, wasm_i8x16_splat(0x01)), 4);
1113            v128_t qh_high1 = wasm_i8x16_shl(wasm_v128_and(qh_shift1, wasm_i8x16_splat(0x02)), 3);
1114
1115            v128_t q5_0 = wasm_v128_load(q5);
1116            v128_t q5_1 = wasm_v128_load(q5 + 16);
1117            q5 += 32;
1118
1119            v128_t q5l_0 = wasm_v128_or(wasm_v128_and(q5_0, wasm_i8x16_splat(0x0F)), qh_low0);
1120            v128_t q5h_0 = wasm_v128_or(wasm_u8x16_shr(q5_0, 4), qh_high0);
1121            v128_t q5l_1 = wasm_v128_or(wasm_v128_and(q5_1, wasm_i8x16_splat(0x0F)), qh_low1);
1122            v128_t q5h_1 = wasm_v128_or(wasm_u8x16_shr(q5_1, 4), qh_high1);
1123
1124            v128_t q8_0 = wasm_v128_load(q8);
1125            v128_t q8_1 = wasm_v128_load(q8 + 16);
1126            v128_t q8_2 = wasm_v128_load(q8 + 32);
1127            v128_t q8_3 = wasm_v128_load(q8 + 48);
1128            q8 += 64;
1129
1130            // Process low quants
1131            v128_t pl0 = wasm_i32x4_dot_i16x8(
1132                wasm_i16x8_extend_low_i8x16(q5l_0),
1133                wasm_i16x8_extend_low_i8x16(q8_0)
1134            );
1135            pl0 = wasm_i32x4_add(pl0, wasm_i32x4_dot_i16x8(
1136                wasm_i16x8_extend_high_i8x16(q5l_0),
1137                wasm_i16x8_extend_high_i8x16(q8_0)
1138            ));
1139            v128_t pl1 = wasm_i32x4_dot_i16x8(
1140                wasm_i16x8_extend_low_i8x16(q5l_1),
1141                wasm_i16x8_extend_low_i8x16(q8_1)
1142            );
1143            pl1 = wasm_i32x4_add(pl1, wasm_i32x4_dot_i16x8(
1144                wasm_i16x8_extend_high_i8x16(q5l_1),
1145                wasm_i16x8_extend_high_i8x16(q8_1)
1146            ));
1147            v128_t sum_low = wasm_i32x4_add(pl0, pl1);
1148
1149            // Process high quants
1150            v128_t ph0 = wasm_i32x4_dot_i16x8(
1151                wasm_i16x8_extend_low_i8x16(q5h_0),
1152                wasm_i16x8_extend_low_i8x16(q8_2)
1153            );
1154            ph0 = wasm_i32x4_add(ph0, wasm_i32x4_dot_i16x8(
1155                wasm_i16x8_extend_high_i8x16(q5h_0),
1156                wasm_i16x8_extend_high_i8x16(q8_2)
1157            ));
1158            v128_t ph1 = wasm_i32x4_dot_i16x8(
1159                wasm_i16x8_extend_low_i8x16(q5h_1),
1160                wasm_i16x8_extend_low_i8x16(q8_3)
1161            );
1162            ph1 = wasm_i32x4_add(ph1, wasm_i32x4_dot_i16x8(
1163                wasm_i16x8_extend_high_i8x16(q5h_1),
1164                wasm_i16x8_extend_high_i8x16(q8_3)
1165            ));
1166            v128_t sum_high = wasm_i32x4_add(ph0, ph1);
1167
1168            // Accumulate with scale factors
1169            int32_t sl = wasm_i32x4_extract_lane(sum_low, 0) + wasm_i32x4_extract_lane(sum_low, 1) +
1170                        wasm_i32x4_extract_lane(sum_low, 2) + wasm_i32x4_extract_lane(sum_low, 3);
1171            int32_t sh = wasm_i32x4_extract_lane(sum_high, 0) + wasm_i32x4_extract_lane(sum_high, 1) +
1172                        wasm_i32x4_extract_lane(sum_high, 2) + wasm_i32x4_extract_lane(sum_high, 3);
1173
1174            sumi += sl * sc[2*j] + sh * sc[2*j+1];
1175        }
1176
1177        sumf += d * sumi;
1178    }
1179
1180    *s = sumf;
1181
1182#else
1183    UNUSED(x);
1184    UNUSED(y);
1185    UNUSED(nb);
1186    UNUSED(kmask1);
1187    UNUSED(kmask2);
1188    UNUSED(kmask3);
1189    UNUSED(utmp);
1190    ggml_vec_dot_q5_K_q8_K_generic(n, s, bs, vx, bx, vy, by, nrc);
1191#endif
1192}
1193
1194void ggml_vec_dot_q6_K_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) {
1195    assert(n % QK_K == 0);
1196    assert(nrc == 1);
1197    UNUSED(nrc);
1198    UNUSED(bx);
1199    UNUSED(by);
1200    UNUSED(bs);
1201
1202    const block_q6_K * GGML_RESTRICT x = vx;
1203    const block_q8_K * GGML_RESTRICT y = vy;
1204
1205    const int nb = n / QK_K;
1206
1207#if defined __wasm_simd128__
1208    int8_t aux8[QK_K] __attribute__((aligned(16)));
1209    int32_t aux32[8] __attribute__((aligned(16))) = {0};
1210    float sums[8] __attribute__((aligned(16))) = {0};
1211
1212    for (int i = 0; i < nb; ++i) {
1213        // Unpack 6-bit quantized data into aux8 (unchanged)
1214        const uint8_t * GGML_RESTRICT q4 = x[i].ql;
1215        const uint8_t * GGML_RESTRICT qh = x[i].qh;
1216        int8_t * a = aux8;
1217        for (int j = 0; j < QK_K; j += 128) {
1218            for (int l = 0; l < 32; ++l) {
1219                a[l +  0] = (int8_t)((q4[l +  0] & 0xF) | (((qh[l] >> 0) & 3) << 4)) - 32;
1220                a[l + 32] = (int8_t)((q4[l + 32] & 0xF) | (((qh[l] >> 2) & 3) << 4)) - 32;
1221                a[l + 64] = (int8_t)((q4[l +  0] >>  4) | (((qh[l] >> 4) & 3) << 4)) - 32;
1222                a[l + 96] = (int8_t)((q4[l + 32] >>  4) | (((qh[l] >> 6) & 3) << 4)) - 32;
1223            }
1224            a += 128;
1225            q4 += 64;
1226            qh += 32;
1227        }
1228
1229        const int8_t * GGML_RESTRICT a_ptr = aux8;
1230        const int8_t * GGML_RESTRICT q8 = y[i].qs;
1231        v128_t acc0 = wasm_i32x4_splat(0);
1232        v128_t acc1 = wasm_i32x4_splat(0);
1233
1234        for (int j = 0; j < QK_K/16; ++j) {
1235            const int scale = x[i].scales[j];
1236            const v128_t vscale = wasm_i32x4_splat(scale);
1237
1238            // Load 16 elements from a and q8
1239            const v128_t a_vec = wasm_v128_load(a_ptr);
1240            const v128_t q8_vec = wasm_v128_load(q8);
1241
1242            // Process low 8 elements
1243            v128_t a_low = wasm_i16x8_extend_low_i8x16(a_vec);
1244            v128_t q8_low = wasm_i16x8_extend_low_i8x16(q8_vec);
1245            v128_t prod_low = wasm_i16x8_mul(a_low, q8_low);
1246            v128_t prod_lo_lo = wasm_i32x4_extend_low_i16x8(prod_low);
1247            v128_t prod_lo_hi = wasm_i32x4_extend_high_i16x8(prod_low);
1248
1249            // Process high 8 elements
1250            v128_t a_high = wasm_i16x8_extend_high_i8x16(a_vec);
1251            v128_t q8_high = wasm_i16x8_extend_high_i8x16(q8_vec);
1252            v128_t prod_high = wasm_i16x8_mul(a_high, q8_high);
1253            v128_t prod_hi_lo = wasm_i32x4_extend_low_i16x8(prod_high);
1254            v128_t prod_hi_hi = wasm_i32x4_extend_high_i16x8(prod_high);
1255
1256            // Scale and accumulate
1257            prod_lo_lo = wasm_i32x4_mul(prod_lo_lo, vscale);
1258            prod_lo_hi = wasm_i32x4_mul(prod_lo_hi, vscale);
1259            prod_hi_lo = wasm_i32x4_mul(prod_hi_lo, vscale);
1260            prod_hi_hi = wasm_i32x4_mul(prod_hi_hi, vscale);
1261
1262            acc0 = wasm_i32x4_add(acc0, wasm_i32x4_add(prod_lo_lo, prod_hi_lo));
1263            acc1 = wasm_i32x4_add(acc1, wasm_i32x4_add(prod_lo_hi, prod_hi_hi));
1264
1265            a_ptr += 16;
1266            q8 += 16;
1267        }
1268
1269        // Store accumulated results
1270        wasm_v128_store(&aux32[0], acc0);
1271        wasm_v128_store(&aux32[4], acc1);
1272
1273        const float d = GGML_CPU_FP16_TO_FP32(x[i].d) * y[i].d;
1274        for (int l = 0; l < 8; ++l) {
1275            sums[l] += d * aux32[l];
1276        }
1277    }
1278
1279    // Sum final results
1280    float sumf = 0;
1281    for (int l = 0; l < 8; ++l) {
1282        sumf += sums[l];
1283    }
1284    *s = sumf;
1285
1286#else
1287    UNUSED(x);
1288    UNUSED(y);
1289    UNUSED(nb);
1290    ggml_vec_dot_q6_K_q8_K_generic(n, s, bs, vx, bx, vy, by, nrc);
1291#endif
1292}