hey! this is a static git repo viewer. commit diffs, file history, older commits and other revisions won't load without javascript.

you are still able to view source files on branch/tag tips, but the experience is degraded.

char-fork/llama.cpp

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

Foad Abo Dahoodmetal : fix idle threads in the remaining iq mul_mv kernels for ne00 < 1024 (#28692)aac8102302026-09-11

1#include "common.h"
2#include "dequantize.h"
3// Q1_0 dot product: dot = d * (2 * Σ(yl[i] where bit=1) - sumy)
4inline float block_q_n_dot_y(device const block_q1_0 * qb_curr, float sumy, thread float * yl, int il) {
5    device const uint8_t * qs = qb_curr->qs + il / 8;
6    const uint8_t b0 = qs[0];
7    const uint8_t b1 = qs[1];
8
9    float acc = 0.0f;
10
11    acc += select(0.0f, yl[ 0], bool(b0 & 0x01));
12    acc += select(0.0f, yl[ 1], bool(b0 & 0x02));
13    acc += select(0.0f, yl[ 2], bool(b0 & 0x04));
14    acc += select(0.0f, yl[ 3], bool(b0 & 0x08));
15    acc += select(0.0f, yl[ 4], bool(b0 & 0x10));
16    acc += select(0.0f, yl[ 5], bool(b0 & 0x20));
17    acc += select(0.0f, yl[ 6], bool(b0 & 0x40));
18    acc += select(0.0f, yl[ 7], bool(b0 & 0x80));
19
20    acc += select(0.0f, yl[ 8], bool(b1 & 0x01));
21    acc += select(0.0f, yl[ 9], bool(b1 & 0x02));
22    acc += select(0.0f, yl[10], bool(b1 & 0x04));
23    acc += select(0.0f, yl[11], bool(b1 & 0x08));
24    acc += select(0.0f, yl[12], bool(b1 & 0x10));
25    acc += select(0.0f, yl[13], bool(b1 & 0x20));
26    acc += select(0.0f, yl[14], bool(b1 & 0x40));
27    acc += select(0.0f, yl[15], bool(b1 & 0x80));
28
29    return qb_curr->d * (2.0f * acc - sumy);
30}
31
32// Q2_0 dot: d * (sum_lo(y) + 2*sum_hi(y) - sumy) via per-bit conditional adds
33inline float block_q_n_dot_y(device const block_q2_0 * qb_curr, float sumy, thread float * yl, int il) {
34    device const uint8_t * qs = qb_curr->qs + (il / 4);
35    const uint8_t b0 = qs[0];
36    const uint8_t b1 = qs[1];
37    const uint8_t b2 = qs[2];
38    const uint8_t b3 = qs[3];
39
40    // Accumulate where low bit is set (bits 0,2,4,6 of each byte)
41    float acc_lo = 0.0f;
42    acc_lo += select(0.0f, yl[ 0], bool(b0 & 0x01));
43    acc_lo += select(0.0f, yl[ 1], bool(b0 & 0x04));
44    acc_lo += select(0.0f, yl[ 2], bool(b0 & 0x10));
45    acc_lo += select(0.0f, yl[ 3], bool(b0 & 0x40));
46    acc_lo += select(0.0f, yl[ 4], bool(b1 & 0x01));
47    acc_lo += select(0.0f, yl[ 5], bool(b1 & 0x04));
48    acc_lo += select(0.0f, yl[ 6], bool(b1 & 0x10));
49    acc_lo += select(0.0f, yl[ 7], bool(b1 & 0x40));
50    acc_lo += select(0.0f, yl[ 8], bool(b2 & 0x01));
51    acc_lo += select(0.0f, yl[ 9], bool(b2 & 0x04));
52    acc_lo += select(0.0f, yl[10], bool(b2 & 0x10));
53    acc_lo += select(0.0f, yl[11], bool(b2 & 0x40));
54    acc_lo += select(0.0f, yl[12], bool(b3 & 0x01));
55    acc_lo += select(0.0f, yl[13], bool(b3 & 0x04));
56    acc_lo += select(0.0f, yl[14], bool(b3 & 0x10));
57    acc_lo += select(0.0f, yl[15], bool(b3 & 0x40));
58
59    // Accumulate where high bit is set (bits 1,3,5,7 of each byte)
60    float acc_hi = 0.0f;
61    acc_hi += select(0.0f, yl[ 0], bool(b0 & 0x02));
62    acc_hi += select(0.0f, yl[ 1], bool(b0 & 0x08));
63    acc_hi += select(0.0f, yl[ 2], bool(b0 & 0x20));
64    acc_hi += select(0.0f, yl[ 3], bool(b0 & 0x80));
65    acc_hi += select(0.0f, yl[ 4], bool(b1 & 0x02));
66    acc_hi += select(0.0f, yl[ 5], bool(b1 & 0x08));
67    acc_hi += select(0.0f, yl[ 6], bool(b1 & 0x20));
68    acc_hi += select(0.0f, yl[ 7], bool(b1 & 0x80));
69    acc_hi += select(0.0f, yl[ 8], bool(b2 & 0x02));
70    acc_hi += select(0.0f, yl[ 9], bool(b2 & 0x08));
71    acc_hi += select(0.0f, yl[10], bool(b2 & 0x20));
72    acc_hi += select(0.0f, yl[11], bool(b2 & 0x80));
73    acc_hi += select(0.0f, yl[12], bool(b3 & 0x02));
74    acc_hi += select(0.0f, yl[13], bool(b3 & 0x08));
75    acc_hi += select(0.0f, yl[14], bool(b3 & 0x20));
76    acc_hi += select(0.0f, yl[15], bool(b3 & 0x80));
77
78    return qb_curr->d * (acc_lo + 2.0f * acc_hi - sumy);
79}
80
81// function for calculate inner product between half a q4_0 block and 16 floats (yl), sumy is SUM(yl[i])
82// il indicates where the q4 quants begin (0 or QK4_0/4)
83// we assume that the yl's have been multiplied with the appropriate scale factor
84// that corresponds to the missing bit shifts (1, 1/16, 1/256, 1/4096)
85inline float block_q_n_dot_y(device const block_q4_0 * qb_curr, float sumy, thread float * yl, int il) {
86    float d = qb_curr->d;
87
88    float acc[4] = { 0.0f, 0.0f, 0.0f, 0.0f };
89
90    device const uint16_t * qs = ((device const uint16_t *) qb_curr + 1 + il/2);
91
92    for (int i = 0; i < 8; i += 2) {
93        acc[0] += yl[i + 0] * (qs[i / 2] & 0x000F);
94        acc[1] += yl[i + 1] * (qs[i / 2] & 0x0F00);
95        acc[2] += yl[i + 8] * (qs[i / 2] & 0x00F0);
96        acc[3] += yl[i + 9] * (qs[i / 2] & 0xF000);
97    }
98
99    return d * (sumy * -8.f + acc[0] + acc[1] + acc[2] + acc[3]);
100}
101
102// function for calculate inner product between half a q4_1 block and 16 floats (yl), sumy is SUM(yl[i])
103// il indicates where the q4 quants begin (0 or QK4_0/4)
104// we assume that the yl's have been multiplied with the appropriate scale factor
105// that corresponds to the missing bit shifts (1, 1/16, 1/256, 1/4096)
106inline float block_q_n_dot_y(device const block_q4_1 * qb_curr, float sumy, thread float * yl, int il) {
107    float d = qb_curr->d;
108    float m = qb_curr->m;
109
110    float acc[4] = { 0.0f, 0.0f, 0.0f, 0.0f };
111
112    device const uint16_t * qs = ((device const uint16_t *) qb_curr + 2 + il/2);
113
114    for (int i = 0; i < 8; i+=2) {
115        acc[0] += yl[i + 0] * (qs[i / 2] & 0x000F);
116        acc[1] += yl[i + 1] * (qs[i / 2] & 0x0F00);
117        acc[2] += yl[i + 8] * (qs[i / 2] & 0x00F0);
118        acc[3] += yl[i + 9] * (qs[i / 2] & 0xF000);
119    }
120
121    return d * (acc[0] + acc[1] + acc[2] + acc[3]) + sumy * m;
122}
123
124// function for calculate inner product between half a q5_0 block and 16 floats (yl), sumy is SUM(yl[i])
125// il indicates where the q5 quants begin (0 or QK5_0/4)
126// we assume that the yl's have been multiplied with the appropriate scale factor
127// that corresponds to the missing bit shifts (1, 1/16, 1/256, 1/4096)
128inline float block_q_n_dot_y(device const block_q5_0 * qb_curr, float sumy, thread float * yl, int il) {
129    float d = qb_curr->d;
130
131    float acc[4] = { 0.0f, 0.0f, 0.0f, 0.0f };
132
133    device const uint16_t * qs =  ((device const uint16_t *)qb_curr + 3 + il/2);
134           const uint32_t   qh = *((device const uint32_t *)qb_curr->qh);
135
136    for (int i = 0; i < 8; i+=2) {
137        acc[0] += yl[i + 0] * ((qs[i / 2] & 0x000F) | ((qh >> (i+0+il        ) << 4 ) & 0x00010));
138        acc[1] += yl[i + 1] * ((qs[i / 2] & 0x0F00) | ((qh >> (i+1+il        ) << 12) & 0x01000));
139        acc[2] += yl[i + 8] * ((qs[i / 2] & 0x00F0) | ((qh >> (i+0+il+QK5_0/2) << 8 ) & 0x00100));
140        acc[3] += yl[i + 9] * ((qs[i / 2] & 0xF000) | ((qh >> (i+1+il+QK5_0/2) << 16) & 0x10000));
141    }
142
143    return d * (sumy * -16.f + acc[0] + acc[1] + acc[2] + acc[3]);
144}
145
146// function for calculate inner product between half a q5_1 block and 16 floats (yl), sumy is SUM(yl[i])
147// il indicates where the q5 quants begin (0 or QK5_1/4)
148// we assume that the yl's have been multiplied with the appropriate scale factor
149// that corresponds to the missing bit shifts (1, 1/16, 1/256, 1/4096)
150inline float block_q_n_dot_y(device const block_q5_1 * qb_curr, float sumy, thread float * yl, int il) {
151    float d = qb_curr->d;
152    float m = qb_curr->m;
153
154    float acc[4] = { 0.0f, 0.0f, 0.0f, 0.0f };
155
156    device const uint16_t * qs =  ((device const uint16_t *)qb_curr + 4 + il/2);
157           const uint32_t   qh = *((device const uint32_t *)qb_curr->qh);
158
159    for (int i = 0; i < 8; i+=2) {
160        acc[0] += yl[i + 0] * ((qs[i / 2] & 0x000F) | ((qh >> (i+0+il        ) << 4 ) & 0x00010));
161        acc[1] += yl[i + 1] * ((qs[i / 2] & 0x0F00) | ((qh >> (i+1+il        ) << 12) & 0x01000));
162        acc[2] += yl[i + 8] * ((qs[i / 2] & 0x00F0) | ((qh >> (i+0+il+QK5_0/2) << 8 ) & 0x00100));
163        acc[3] += yl[i + 9] * ((qs[i / 2] & 0xF000) | ((qh >> (i+1+il+QK5_0/2) << 16) & 0x10000));
164    }
165
166    return d * (acc[0] + acc[1] + acc[2] + acc[3]) + sumy * m;
167}
168
169template<short NR0>
170static inline void helper_mv_reduce_and_write(
171        device float * dst_f32,
172        float sumf[NR0],
173        const int r0,
174        const int ne01,
175        ushort tiisg,
176        ushort sgitg,
177        threadgroup char * shmem) {
178    constexpr short NW = N_SIMDWIDTH;
179
180    threadgroup float * shmem_f32[NR0];
181
182    for (short row = 0; row < NR0; ++row) {
183        shmem_f32[row] = (threadgroup float *) shmem + NW*row;
184
185        if (sgitg == 0) {
186            shmem_f32[row][tiisg] = 0.0f;
187        }
188
189        sumf[row] = simd_sum(sumf[row]);
190    }
191
192    threadgroup_barrier(mem_flags::mem_threadgroup);
193
194    for (short row = 0; row < NR0; ++row) {
195        if (tiisg == 0) {
196            shmem_f32[row][sgitg] = sumf[row];
197        }
198    }
199
200    threadgroup_barrier(mem_flags::mem_threadgroup);
201
202    for (short row = 0; row < NR0 && r0 + row < ne01; ++row) {
203        float tot = simd_sum(shmem_f32[row][tiisg]);
204
205        if (tiisg == 0 && sgitg == 0) {
206            dst_f32[r0 + row] = tot;
207        }
208    }
209}
210
211constant short FC_mul_mv_nsg   [[function_constant(FC_MUL_MV + 0)]];
212constant short FC_mul_mv_nxpsg [[function_constant(FC_MUL_MV + 1)]];
213constant short FC_mul_mv_ne12  [[function_constant(FC_MUL_MV + 2)]];
214constant short FC_mul_mv_r2    [[function_constant(FC_MUL_MV + 3)]];
215constant short FC_mul_mv_r3    [[function_constant(FC_MUL_MV + 4)]];
216constant bool  FC_mul_mv_split [[function_constant(FC_MUL_MV + 5)]];
217
218template<typename block_q_type, short NR0, typename args_t>
219void mul_vec_q_n_f32_impl(
220        args_t args,
221        device const char * src0,
222        device const char * src1,
223        device       char * dst,
224        threadgroup  char * shmem,
225        uint3  tgpig,
226        ushort tiisg,
227        ushort sgitg) {
228    const short NSG = FC_mul_mv_nsg;
229
230    constexpr short NW = N_SIMDWIDTH;
231    constexpr short NQ = 16;
232
233    const int nb = args.ne00/QK4_0;
234
235    const int r0 = (tgpig.x*NSG + sgitg)*NR0;
236  //const int r0 =  tgpig.x*NR0;
237    const int r1 =  tgpig.y;
238    const int im =  tgpig.z;
239
240    const uint i12 = im%FC_mul_mv_ne12;
241    const uint i13 = im/FC_mul_mv_ne12;
242
243  //const uint64_t offset0 = r0*args.nb01 + (i12/FC_mul_mv_r2)*args.nb02 + (i13/FC_mul_mv_r3)*args.nb03;
244    const uint64_t offset1 = r1*args.nb11 + (i12        )*args.nb12 + (i13        )*args.nb13;
245
246  //device const block_q_type * x = (device const block_q_type *) (src0 + offset0);
247    device const float        * y = (device const float        *) (src1 + offset1);
248
249    // pointers to src0 rows
250    device const block_q_type * ax[NR0];
251    FOR_UNROLL (int row = 0; row < NR0; ++row) {
252        const uint64_t offset0 = (r0 + row)*args.nb01 + (i12/FC_mul_mv_r2)*args.nb02 + (i13/FC_mul_mv_r3)*args.nb03;
253
254        ax[row] = (device const block_q_type *) ((device char *) src0 + offset0);
255    }
256
257    float sumf[NR0] = {0.f};
258
259    const short ix = (tiisg/(NW/NQ));
260    const short il = (tiisg%(NW/NQ))*8;
261
262    //const int ib0 = sgitg*NQ + ix;
263    const int ib0 = ix;
264
265    float yl[16]; // src1 vector cache
266
267    //device const float * yb = y + ix*QK4_0 + il;
268    device const float * yb = y + ib0*QK4_0 + il;
269
270    // each thread in a SIMD group deals with half a block.
271    //for (int ib = ib0; ib < nb; ib += NSG*NQ) {
272    for (int ib = ib0; ib < nb; ib += NQ) {
273        float sumy[2] = { 0.f, 0.f };
274
275        FOR_UNROLL (short i = 0; i < 8; i += 2) {
276            sumy[0]  += yb[i +  0] + yb[i +  1];
277            yl[i + 0] = yb[i +  0];
278            yl[i + 1] = yb[i +  1]/256.f;
279
280            sumy[1]  += yb[i + 16] + yb[i + 17];
281            yl[i + 8] = yb[i + 16]/16.f;
282            yl[i + 9] = yb[i + 17]/4096.f;
283        }
284
285        FOR_UNROLL (short row = 0; row < NR0; row++) {
286            sumf[row] += block_q_n_dot_y(ax[row] + ib, sumy[0] + sumy[1], yl, il);
287        }
288
289        yb += QK4_0 * 16;
290        //yb += NSG*NQ*QK4_0;
291    }
292
293    device float * dst_f32 = (device float *) dst + im*args.ne0*args.ne1 + r1*args.ne0;
294
295    //helper_mv_reduce_and_write<NR0>(dst_f32, sumf, r0, args.ne01, tiisg, sgitg, shmem);
296
297    for (int row = 0; row < NR0; ++row) {
298        const float tot = simd_sum(sumf[row]);
299
300        if (tiisg == 0 && r0 + row < args.ne01) {
301            dst_f32[r0 + row] = tot;
302        }
303    }
304}
305
306template<int nr0, typename args_t>
307void kernel_mul_mv_q1_0_f32_impl(
308        args_t args,
309        device const char * src0,
310        device const char * src1,
311        device       char * dst,
312        threadgroup  char * shmem,
313        uint3  tgpig,
314        ushort tiisg,
315        ushort sgitg) {
316    const short NSG = FC_mul_mv_nsg;
317
318    const int nb = args.ne00/QK1_0;
319
320    const int r0 = tgpig.x;
321    const int r1 = tgpig.y;
322    const int im = tgpig.z;
323
324    const int first_row = (r0 * NSG + sgitg) * nr0;
325
326    const uint i12 = im%FC_mul_mv_ne12;
327    const uint i13 = im/FC_mul_mv_ne12;
328
329    const uint64_t offset1 = r1*args.nb11 + (i12)*args.nb12 + (i13)*args.nb13;
330
331    device const float * y = (device const float *) (src1 + offset1);
332
333    device const block_q1_0 * ax[nr0];
334    for (int row = 0; row < nr0; ++row) {
335        const uint64_t offset0 = (first_row + row)*args.nb01 + (i12/FC_mul_mv_r2)*args.nb02 + (i13/FC_mul_mv_r3)*args.nb03;
336        ax[row] = (device const block_q1_0 *) ((device char *) src0 + offset0);
337    }
338
339    float yl[16];
340    float sumf[nr0] = {0.f};
341
342    const short ix = (tiisg/8);
343    const short il = (tiisg%8)*16;
344
345    device const float * yb = y + ix*QK1_0 + il;
346
347    for (int ib = ix; ib < nb; ib += N_SIMDWIDTH/8) {
348        float sumy = 0.f;
349
350        FOR_UNROLL (short i = 0; i < 16; i++) {
351            yl[i] = yb[i];
352            sumy += yb[i];
353        }
354
355        FOR_UNROLL (short row = 0; row < nr0; row++) {
356            sumf[row] += block_q_n_dot_y(ax[row] + ib, sumy, yl, il);
357        }
358
359        yb += QK1_0 * (N_SIMDWIDTH/8);
360    }
361
362    device float * dst_f32 = (device float *) dst + (uint64_t)im*args.ne0*args.ne1 + (uint64_t)r1*args.ne0;
363
364    for (int row = 0; row < nr0; ++row) {
365        const float tot = simd_sum(sumf[row]);
366
367        if (tiisg == 0 && first_row + row < args.ne01) {
368            dst_f32[first_row + row] = tot;
369        }
370    }
371}
372
373[[host_name("kernel_mul_mv_q1_0_f32")]]
374kernel void kernel_mul_mv_q1_0_f32(
375        constant ggml_metal_kargs_mul_mv & args,
376        device const char * src0,
377        device const char * src1,
378        device       char * dst,
379        uint3  tgpig[[threadgroup_position_in_grid]],
380        ushort tiisg[[thread_index_in_simdgroup]],
381        ushort sgitg[[simdgroup_index_in_threadgroup]]) {
382    kernel_mul_mv_q1_0_f32_impl<N_R0_Q1_0, constant ggml_metal_kargs_mul_mv &>(args, src0, src1, dst, nullptr, tgpig, tiisg, sgitg);
383}
384
385template<int nr0, typename args_t>
386void kernel_mul_mv_q2_0_f32_impl(
387        args_t args,
388        device const char * src0,
389        device const char * src1,
390        device       char * dst,
391        threadgroup  char * shmem,
392        uint3  tgpig,
393        ushort tiisg,
394        ushort sgitg) {
395    const short NSG = FC_mul_mv_nsg;
396
397    const int nb = args.ne00/QK2_0;
398
399    const int r0 = tgpig.x;
400    const int r1 = tgpig.y;
401    const int im = tgpig.z;
402
403    const int first_row = (r0 * NSG + sgitg) * nr0;
404
405    const uint i12 = im%FC_mul_mv_ne12;
406    const uint i13 = im/FC_mul_mv_ne12;
407
408    const uint64_t offset1 = r1*args.nb11 + (i12)*args.nb12 + (i13)*args.nb13;
409
410    device const float * y = (device const float *) (src1 + offset1);
411
412    device const block_q2_0 * ax[nr0];
413    for (int row = 0; row < nr0; ++row) {
414        const uint64_t offset0 = (first_row + row)*args.nb01 + (i12/FC_mul_mv_r2)*args.nb02 + (i13/FC_mul_mv_r3)*args.nb03;
415        ax[row] = (device const block_q2_0 *) ((device char *) src0 + offset0);
416    }
417
418    float yl[16];
419    float sumf[nr0] = {0.f};
420
421    // group 64: 4 sub-blocks of 16 weights per Q2_0 block
422    const short ix = (tiisg/4);
423    const short il = (tiisg%4)*16;
424
425    device const float * yb = y + ix*QK2_0 + il;
426
427    for (int ib = ix; ib < nb; ib += N_SIMDWIDTH/4) {
428        float sumy = 0.f;
429
430        FOR_UNROLL (short i = 0; i < 16; i++) {
431            yl[i] = yb[i];
432            sumy += yb[i];
433        }
434
435        FOR_UNROLL (short row = 0; row < nr0; row++) {
436            sumf[row] += block_q_n_dot_y(ax[row] + ib, sumy, yl, il);
437        }
438
439        yb += QK2_0 * (N_SIMDWIDTH/4);
440    }
441
442    device float * dst_f32 = (device float *) dst + (uint64_t)im*args.ne0*args.ne1 + (uint64_t)r1*args.ne0;
443
444    for (int row = 0; row < nr0; ++row) {
445        const float tot = simd_sum(sumf[row]);
446
447        if (tiisg == 0 && first_row + row < args.ne01) {
448            dst_f32[first_row + row] = tot;
449        }
450    }
451}
452
453[[host_name("kernel_mul_mv_q2_0_f32")]]
454kernel void kernel_mul_mv_q2_0_f32(
455        constant ggml_metal_kargs_mul_mv & args,
456        device const char * src0,
457        device const char * src1,
458        device       char * dst,
459        uint3  tgpig[[threadgroup_position_in_grid]],
460        ushort tiisg[[thread_index_in_simdgroup]],
461        ushort sgitg[[simdgroup_index_in_threadgroup]]) {
462    kernel_mul_mv_q2_0_f32_impl<N_R0_Q2_0, constant ggml_metal_kargs_mul_mv &>(args, src0, src1, dst, nullptr, tgpig, tiisg, sgitg);
463}
464
465kernel void kernel_mul_mv_q4_0_f32(
466        constant ggml_metal_kargs_mul_mv & args,
467        device const char * src0,
468        device const char * src1,
469        device       char * dst,
470        threadgroup  char * shmem [[threadgroup(0)]],
471        uint3  tgpig[[threadgroup_position_in_grid]],
472        ushort tiisg[[thread_index_in_simdgroup]],
473        ushort sgitg[[simdgroup_index_in_threadgroup]]) {
474    mul_vec_q_n_f32_impl<block_q4_0, N_R0_Q4_0, constant ggml_metal_kargs_mul_mv &>(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg);
475}
476
477kernel void kernel_mul_mv_q4_1_f32(
478        constant ggml_metal_kargs_mul_mv & args,
479        device const char * src0,
480        device const char * src1,
481        device       char * dst,
482        threadgroup  char * shmem [[threadgroup(0)]],
483        uint3  tgpig[[threadgroup_position_in_grid]],
484        ushort tiisg[[thread_index_in_simdgroup]],
485        ushort sgitg[[simdgroup_index_in_threadgroup]]) {
486     mul_vec_q_n_f32_impl<block_q4_1, N_R0_Q4_1, constant ggml_metal_kargs_mul_mv &>(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg);
487}
488
489kernel void kernel_mul_mv_q5_0_f32(
490        constant ggml_metal_kargs_mul_mv & args,
491        device const char * src0,
492        device const char * src1,
493        device       char * dst,
494        threadgroup  char * shmem [[threadgroup(0)]],
495        uint3  tgpig[[threadgroup_position_in_grid]],
496        ushort tiisg[[thread_index_in_simdgroup]],
497        ushort sgitg[[simdgroup_index_in_threadgroup]]) {
498    mul_vec_q_n_f32_impl<block_q5_0, N_R0_Q5_0, constant ggml_metal_kargs_mul_mv &>(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg);
499}
500
501kernel void kernel_mul_mv_q5_1_f32(
502        constant ggml_metal_kargs_mul_mv & args,
503        device const char * src0,
504        device const char * src1,
505        device       char * dst,
506        threadgroup  char * shmem [[threadgroup(0)]],
507        uint3  tgpig[[threadgroup_position_in_grid]],
508        ushort tiisg[[thread_index_in_simdgroup]],
509        ushort sgitg[[simdgroup_index_in_threadgroup]]) {
510    mul_vec_q_n_f32_impl<block_q5_1, N_R0_Q5_1, constant ggml_metal_kargs_mul_mv &>(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg);
511}
512
513template<short NR0, typename args_t>
514void kernel_mul_mv_q8_0_f32_impl(
515        args_t args,
516        device const char * src0,
517        device const char * src1,
518        device       char * dst,
519        threadgroup  char * shmem,
520        uint3  tgpig,
521        ushort tiisg,
522        ushort sgitg) {
523    const short NSG = FC_mul_mv_nsg;
524
525    constexpr short NW = N_SIMDWIDTH;
526    constexpr short NQ = 8;
527
528    const int nb = args.ne00/QK8_0;
529
530    const int r0 = tgpig.x*NR0;
531    const int r1 = tgpig.y;
532    const int im = tgpig.z;
533
534    const uint i12 = im%FC_mul_mv_ne12;
535    const uint i13 = im/FC_mul_mv_ne12;
536
537  //const uint64_t offset0 = r0*args.nb01 + (i12/FC_mul_mv_r2)*args.nb02 + (i13/FC_mul_mv_r3)*args.nb03;
538    const uint64_t offset1 = r1*args.nb11 + (i12        )*args.nb12 + (i13        )*args.nb13;
539
540  //device const block_q8_0 * x = (device const block_q8_0 *) (src0 + offset0);
541    device const float      * y = (device const float      *) (src1 + offset1);
542
543    // pointers to src0 rows
544    device const block_q8_0 * ax[NR0];
545    FOR_UNROLL (short row = 0; row < NR0; ++row) {
546        const uint64_t offset0 = (r0 + row)*args.nb01 + (i12/FC_mul_mv_r2)*args.nb02 + (i13/FC_mul_mv_r3)*args.nb03;
547
548        ax[row] = (device const block_q8_0 *) ((device char *) src0 + offset0);
549    }
550
551    float sumf[NR0] = { 0.f };
552
553    const short ix = tiisg/(NW/NQ);
554    const short il = tiisg%(NW/NQ);
555
556    const int ib0 = sgitg*NQ + ix;
557
558    float yl[NQ];
559
560    device const float * yb = y + ib0*QK8_0 + il*NQ;
561
562    // each thread in a SIMD group deals with NQ quants at a time
563    for (int ib = ib0; ib < nb; ib += NSG*NQ) {
564        for (short i = 0; i < NQ; ++i) {
565            yl[i] = yb[i];
566        }
567
568        for (short row = 0; row < NR0; row++) {
569            device const int8_t * qs = ax[row][ib].qs + il*NQ;
570
571            float sumq = 0.f;
572            FOR_UNROLL (short i = 0; i < NQ; ++i) {
573                sumq += qs[i] * yl[i];
574            }
575
576            sumf[row] += sumq*ax[row][ib].d;
577        }
578
579        yb += NSG*NQ*QK8_0;
580    }
581
582    device float * dst_f32 = (device float *) dst + (uint64_t)im*args.ne0*args.ne1 + (uint64_t)r1*args.ne0;
583
584    helper_mv_reduce_and_write<NR0>(dst_f32, sumf, r0, args.ne01, tiisg, sgitg, shmem);
585}
586
587[[host_name("kernel_mul_mv_q8_0_f32")]]
588kernel void kernel_mul_mv_q8_0_f32(
589        constant ggml_metal_kargs_mul_mv & args,
590        device const char * src0,
591        device const char * src1,
592        device       char * dst,
593        threadgroup  char * shmem [[threadgroup(0)]],
594        uint3  tgpig[[threadgroup_position_in_grid]],
595        ushort tiisg[[thread_index_in_simdgroup]],
596        ushort sgitg[[simdgroup_index_in_threadgroup]]) {
597    kernel_mul_mv_q8_0_f32_impl<N_R0_Q8_0, constant ggml_metal_kargs_mul_mv &>(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg);
598}
599
600// mat-vec kernel processing in chunks of float4
601// chpb - chunks per quantization block
602template<short r1ptg, typename q_t, short chpb, void (*deq_t4)(device const q_t *, short, thread float4 &) >
603void kernel_mul_mv_ext_q4_f32_impl(
604        constant ggml_metal_kargs_mul_mv_ext & args,
605        device const char * src0,
606        device const char * src1,
607        device       char * dst,
608        uint3   tgpig[[threadgroup_position_in_grid]],
609        ushort  tiisg[[thread_index_in_simdgroup]],
610        ushort  sgitg[[simdgroup_index_in_threadgroup]]) {
611    const short NSG   = FC_mul_mv_nsg;
612    const short nxpsg = FC_mul_mv_nxpsg;
613
614    const short chpt = 4; // chunks per thread
615
616  //const short nxpsg = (32);
617    const short nypsg = (32/nxpsg);
618
619    const short tx = tiisg%nxpsg;
620    const short ty = tiisg/nxpsg;
621
622    const int i01 = tgpig.x*(nypsg*NSG) + nypsg*sgitg + ty;
623    const int i11 = tgpig.y*r1ptg;
624    const int i1m = tgpig.z;
625
626    const int i12 = i1m%FC_mul_mv_ne12;
627    const int i13 = i1m/FC_mul_mv_ne12;
628
629    const uint64_t offset0 = i01*args.nb01 + (i12/FC_mul_mv_r2)*args.nb02 + (i13/FC_mul_mv_r3)*args.nb03;
630    const uint64_t offset1 = i11*args.nb11 + (i12        )*args.nb12 + (i13        )*args.nb13;
631
632    device const q_t * xq = (i01 < args.ne01) ? (device const q_t *) (src0 + offset0) + tx/chpb : (device const q_t *) src0;
633
634    device const float4 * y4[r1ptg];
635
636    for (int ir1 = 0; ir1 < r1ptg; ++ir1) {
637        y4[ir1] = (i11 + ir1 < args.ne11) ? (device const float4 *) (src1 + offset1 + ir1*args.nb11) + tx : (device const float4 *) src1;
638    }
639
640    float sumf[r1ptg] = { [ 0 ... r1ptg - 1 ] = 0.0f };
641
642    short cch = tx%chpb; // current chunk index
643
644    for (int ich = tx; 4*ich < args.ne00; ich += chpt*nxpsg) {
645        float4 lx[chpt];
646
647#pragma unroll(chpt)
648        for (short ch = 0; ch < chpt; ++ch) {
649            deq_t4(xq, cch, lx[ch]);
650
651            cch += nxpsg;
652            if (cch >= chpb) {
653                xq  += cch/chpb;
654                cch %= chpb;
655            }
656        }
657
658#pragma unroll(chpt)
659        for (short ch = 0; ch < chpt; ++ch) {
660#pragma unroll(r1ptg)
661            for (short ir1 = 0; ir1 < r1ptg; ++ir1) {
662                sumf[ir1] += dot(lx[ch], y4[ir1][ch*nxpsg]);
663            }
664        }
665
666#pragma unroll(r1ptg)
667        for (short ir1 = 0; ir1 < r1ptg; ++ir1) {
668            y4[ir1] += chpt*nxpsg;
669        }
670    }
671
672    // reduce only the threads in each row
673    for (short ir1 = 0; ir1 < r1ptg; ++ir1) {
674        if (nxpsg >= 32) {
675            sumf[ir1] += simd_shuffle_down(sumf[ir1], 16);
676        }
677        if (nxpsg >= 16) {
678            sumf[ir1] += simd_shuffle_down(sumf[ir1],  8);
679        }
680        if (nxpsg >= 8) {
681            sumf[ir1] += simd_shuffle_down(sumf[ir1],  4);
682        }
683        if (nxpsg >= 4) {
684            sumf[ir1] += simd_shuffle_down(sumf[ir1],  2);
685        }
686        if (nxpsg >= 2) {
687            sumf[ir1] += simd_shuffle_down(sumf[ir1],  1);
688        }
689
690        //sumf[ir1] = simd_sum(sumf[ir1]);
691    }
692
693    if (tx == 0) {
694        for (short ir1 = 0; ir1 < r1ptg && i11 + ir1 < args.ne11; ++ir1) {
695            device float * dst_f32 = (device float *) dst + (uint64_t)i1m*args.ne0*args.ne1 + (uint64_t)(i11 + ir1)*args.ne0;
696
697            if (i01 < args.ne01) {
698                dst_f32[i01] = sumf[ir1];
699            }
700        }
701    }
702}
703
704// mat-vec kernel processing in chunks of float4x4
705template<short r1ptg, typename q_t, short chpb, void (*deq_t4x4)(device const q_t *, short, thread float4x4 &) >
706void kernel_mul_mv_ext_q4x4_f32_impl(
707        constant ggml_metal_kargs_mul_mv_ext & args,
708        device const char * src0,
709        device const char * src1,
710        device       char * dst,
711        uint3   tgpig[[threadgroup_position_in_grid]],
712        ushort  tiisg[[thread_index_in_simdgroup]],
713        ushort  sgitg[[simdgroup_index_in_threadgroup]]) {
714    const short NSG   = FC_mul_mv_nsg;
715    const short nxpsg = FC_mul_mv_nxpsg;
716
717    const short chpt = 1;
718
719  //const short nxpsg = (32);
720    const short nypsg = (32/nxpsg);
721
722    const short tx = tiisg%nxpsg;
723    const short ty = tiisg/nxpsg;
724
725    const int i01 = tgpig.x*(nypsg*NSG) + nypsg*sgitg + ty;
726    const int i11 = tgpig.y*r1ptg;
727    const int i1m = tgpig.z;
728
729    const int i12 = i1m%FC_mul_mv_ne12;
730    const int i13 = i1m/FC_mul_mv_ne12;
731
732    const uint64_t offset0 = i01*args.nb01 + (i12/FC_mul_mv_r2)*args.nb02 + (i13/FC_mul_mv_r3)*args.nb03;
733    const uint64_t offset1 = i11*args.nb11 + (i12        )*args.nb12 + (i13        )*args.nb13;
734
735    device const q_t * xq = (i01 < args.ne01) ? (device const q_t *) (src0 + offset0) + tx/chpb : (device const q_t *) src0;
736
737    device const float4x4 * y4x4[r1ptg];
738
739    for (int ir1 = 0; ir1 < r1ptg; ++ir1) {
740        y4x4[ir1] = (i11 + ir1 < args.ne11) ? (device const float4x4 *) (src1 + offset1 + ir1*args.nb11) + tx : (device const float4x4 *) src1;
741    }
742
743    float sumf[r1ptg] = { [ 0 ... r1ptg - 1 ] = 0.0f };
744
745    short cch = tx%chpb;
746
747    for (int ich = tx; 16*ich < args.ne00; ich += chpt*nxpsg) {
748        float4x4 lx[chpt];
749
750#pragma unroll(chpt)
751        for (short ch = 0; ch < chpt; ++ch) {
752            deq_t4x4(xq, cch, lx[ch]);
753
754            cch += nxpsg;
755            if (cch >= chpb) {
756                xq  += cch/chpb;
757                cch %= chpb;
758            }
759        }
760
761#pragma unroll(chpt)
762        for (short ch = 0; ch < chpt; ++ch) {
763#pragma unroll(r1ptg)
764            for (short ir1 = 0; ir1 < r1ptg; ++ir1) {
765                sumf[ir1] +=
766                    dot(lx[ch][0], y4x4[ir1][ch*nxpsg][0]) +
767                    dot(lx[ch][1], y4x4[ir1][ch*nxpsg][1]) +
768                    dot(lx[ch][2], y4x4[ir1][ch*nxpsg][2]) +
769                    dot(lx[ch][3], y4x4[ir1][ch*nxpsg][3]);
770
771            }
772        }
773
774#pragma unroll(r1ptg)
775        for (short ir1 = 0; ir1 < r1ptg; ++ir1) {
776            y4x4[ir1] += chpt*nxpsg;
777        }
778    }
779
780    for (short ir1 = 0; ir1 < r1ptg; ++ir1) {
781        if (nxpsg >= 32) {
782            sumf[ir1] += simd_shuffle_down(sumf[ir1], 16);
783        }
784        if (nxpsg >= 16) {
785            sumf[ir1] += simd_shuffle_down(sumf[ir1],  8);
786        }
787        if (nxpsg >= 8) {
788            sumf[ir1] += simd_shuffle_down(sumf[ir1],  4);
789        }
790        if (nxpsg >= 4) {
791            sumf[ir1] += simd_shuffle_down(sumf[ir1],  2);
792        }
793        if (nxpsg >= 2) {
794            sumf[ir1] += simd_shuffle_down(sumf[ir1],  1);
795        }
796
797        //sumf[ir1] = simd_sum(sumf[ir1]);
798    }
799
800    if (tx == 0) {
801        for (short ir1 = 0; ir1 < r1ptg && i11 + ir1 < args.ne11; ++ir1) {
802            device float * dst_f32 = (device float *) dst + (uint64_t)i1m*args.ne0*args.ne1 + (uint64_t)(i11 + ir1)*args.ne0;
803
804            if (i01 < args.ne01) {
805                dst_f32[i01] = sumf[ir1];
806            }
807        }
808    }
809}
810
811// dispatchers needed for compile-time nxpsg
812// epb - elements per quantization block
813template<short r1ptg, typename q_t, short epb, void (*deq_t4)(device const q_t *, short, thread float4 &)>
814kernel void kernel_mul_mv_ext_q4_f32_disp(
815        constant ggml_metal_kargs_mul_mv_ext & args,
816        device const char * src0,
817        device const char * src1,
818        device       char * dst,
819        uint3   tgpig[[threadgroup_position_in_grid]],
820        ushort  tiisg[[thread_index_in_simdgroup]],
821        ushort  sgitg[[simdgroup_index_in_threadgroup]]) {
822    kernel_mul_mv_ext_q4_f32_impl<r1ptg, q_t, epb/4, deq_t4>(args, src0, src1, dst, tgpig, tiisg, sgitg);
823}
824
825template<short r1ptg, typename q_t, short epb, void (*deq_t4x4)(device const q_t *, short, thread float4x4 &)>
826kernel void kernel_mul_mv_ext_q4x4_f32_disp(
827        constant ggml_metal_kargs_mul_mv_ext & args,
828        device const char * src0,
829        device const char * src1,
830        device       char * dst,
831        uint3   tgpig[[threadgroup_position_in_grid]],
832        ushort  tiisg[[thread_index_in_simdgroup]],
833        ushort  sgitg[[simdgroup_index_in_threadgroup]]) {
834    kernel_mul_mv_ext_q4x4_f32_impl<r1ptg, q_t, epb/16, deq_t4x4>(args, src0, src1, dst, tgpig, tiisg, sgitg);
835}
836
837typedef decltype(kernel_mul_mv_ext_q4_f32_disp  <2, block_q8_0, 32,  dequantize_q8_0_t4>) mul_mv_ext_q4_f32_t;
838typedef decltype(kernel_mul_mv_ext_q4x4_f32_disp<2, block_q4_K, 256, dequantize_q4_K>)    mul_mv_ext_q4x4_f32_t;
839
840template [[host_name("kernel_mul_mv_ext_f32_f32_r1_2")]]    kernel mul_mv_ext_q4_f32_t kernel_mul_mv_ext_q4_f32_disp<2, float4,       4,  dequantize_f32_t4>;
841template [[host_name("kernel_mul_mv_ext_f32_f32_r1_3")]]    kernel mul_mv_ext_q4_f32_t kernel_mul_mv_ext_q4_f32_disp<3, float4,       4,  dequantize_f32_t4>;
842template [[host_name("kernel_mul_mv_ext_f32_f32_r1_4")]]    kernel mul_mv_ext_q4_f32_t kernel_mul_mv_ext_q4_f32_disp<4, float4,       4,  dequantize_f32_t4>;
843template [[host_name("kernel_mul_mv_ext_f32_f32_r1_5")]]    kernel mul_mv_ext_q4_f32_t kernel_mul_mv_ext_q4_f32_disp<5, float4,       4,  dequantize_f32_t4>;
844
845template [[host_name("kernel_mul_mv_ext_f16_f32_r1_2")]]    kernel mul_mv_ext_q4_f32_t kernel_mul_mv_ext_q4_f32_disp<2, half4,        4,  dequantize_f16_t4>;
846template [[host_name("kernel_mul_mv_ext_f16_f32_r1_3")]]    kernel mul_mv_ext_q4_f32_t kernel_mul_mv_ext_q4_f32_disp<3, half4,        4,  dequantize_f16_t4>;
847template [[host_name("kernel_mul_mv_ext_f16_f32_r1_4")]]    kernel mul_mv_ext_q4_f32_t kernel_mul_mv_ext_q4_f32_disp<4, half4,        4,  dequantize_f16_t4>;
848template [[host_name("kernel_mul_mv_ext_f16_f32_r1_5")]]    kernel mul_mv_ext_q4_f32_t kernel_mul_mv_ext_q4_f32_disp<5, half4,        4,  dequantize_f16_t4>;
849
850#if defined(GGML_METAL_HAS_BF16)
851template [[host_name("kernel_mul_mv_ext_bf16_f32_r1_2")]]   kernel mul_mv_ext_q4_f32_t kernel_mul_mv_ext_q4_f32_disp<2, bfloat4,      4,  dequantize_bf16_t4>;
852template [[host_name("kernel_mul_mv_ext_bf16_f32_r1_3")]]   kernel mul_mv_ext_q4_f32_t kernel_mul_mv_ext_q4_f32_disp<3, bfloat4,      4,  dequantize_bf16_t4>;
853template [[host_name("kernel_mul_mv_ext_bf16_f32_r1_4")]]   kernel mul_mv_ext_q4_f32_t kernel_mul_mv_ext_q4_f32_disp<4, bfloat4,      4,  dequantize_bf16_t4>;
854template [[host_name("kernel_mul_mv_ext_bf16_f32_r1_5")]]   kernel mul_mv_ext_q4_f32_t kernel_mul_mv_ext_q4_f32_disp<5, bfloat4,      4,  dequantize_bf16_t4>;
855#endif
856
857template [[host_name("kernel_mul_mv_ext_q1_0_f32_r1_2")]]   kernel mul_mv_ext_q4_f32_t kernel_mul_mv_ext_q4_f32_disp<2, block_q1_0,   128, dequantize_q1_0_t4>;
858template [[host_name("kernel_mul_mv_ext_q1_0_f32_r1_3")]]   kernel mul_mv_ext_q4_f32_t kernel_mul_mv_ext_q4_f32_disp<3, block_q1_0,   128, dequantize_q1_0_t4>;
859template [[host_name("kernel_mul_mv_ext_q1_0_f32_r1_4")]]   kernel mul_mv_ext_q4_f32_t kernel_mul_mv_ext_q4_f32_disp<4, block_q1_0,   128, dequantize_q1_0_t4>;
860template [[host_name("kernel_mul_mv_ext_q1_0_f32_r1_5")]]   kernel mul_mv_ext_q4_f32_t kernel_mul_mv_ext_q4_f32_disp<5, block_q1_0,   128, dequantize_q1_0_t4>;
861
862template [[host_name("kernel_mul_mv_ext_q2_0_f32_r1_2")]]   kernel mul_mv_ext_q4_f32_t kernel_mul_mv_ext_q4_f32_disp<2, block_q2_0,    64, dequantize_q2_0_t4>;
863template [[host_name("kernel_mul_mv_ext_q2_0_f32_r1_3")]]   kernel mul_mv_ext_q4_f32_t kernel_mul_mv_ext_q4_f32_disp<3, block_q2_0,    64, dequantize_q2_0_t4>;
864template [[host_name("kernel_mul_mv_ext_q2_0_f32_r1_4")]]   kernel mul_mv_ext_q4_f32_t kernel_mul_mv_ext_q4_f32_disp<4, block_q2_0,    64, dequantize_q2_0_t4>;
865template [[host_name("kernel_mul_mv_ext_q2_0_f32_r1_5")]]   kernel mul_mv_ext_q4_f32_t kernel_mul_mv_ext_q4_f32_disp<5, block_q2_0,    64, dequantize_q2_0_t4>;
866
867template [[host_name("kernel_mul_mv_ext_q4_0_f32_r1_2")]]   kernel mul_mv_ext_q4_f32_t kernel_mul_mv_ext_q4_f32_disp<2, block_q4_0,   32, dequantize_q4_0_t4>;
868template [[host_name("kernel_mul_mv_ext_q4_0_f32_r1_3")]]   kernel mul_mv_ext_q4_f32_t kernel_mul_mv_ext_q4_f32_disp<3, block_q4_0,   32, dequantize_q4_0_t4>;
869template [[host_name("kernel_mul_mv_ext_q4_0_f32_r1_4")]]   kernel mul_mv_ext_q4_f32_t kernel_mul_mv_ext_q4_f32_disp<4, block_q4_0,   32, dequantize_q4_0_t4>;
870template [[host_name("kernel_mul_mv_ext_q4_0_f32_r1_5")]]   kernel mul_mv_ext_q4_f32_t kernel_mul_mv_ext_q4_f32_disp<5, block_q4_0,   32, dequantize_q4_0_t4>;
871
872template [[host_name("kernel_mul_mv_ext_q4_1_f32_r1_2")]]   kernel mul_mv_ext_q4_f32_t kernel_mul_mv_ext_q4_f32_disp<2, block_q4_1,   32, dequantize_q4_1_t4>;
873template [[host_name("kernel_mul_mv_ext_q4_1_f32_r1_3")]]   kernel mul_mv_ext_q4_f32_t kernel_mul_mv_ext_q4_f32_disp<3, block_q4_1,   32, dequantize_q4_1_t4>;
874template [[host_name("kernel_mul_mv_ext_q4_1_f32_r1_4")]]   kernel mul_mv_ext_q4_f32_t kernel_mul_mv_ext_q4_f32_disp<4, block_q4_1,   32, dequantize_q4_1_t4>;
875template [[host_name("kernel_mul_mv_ext_q4_1_f32_r1_5")]]   kernel mul_mv_ext_q4_f32_t kernel_mul_mv_ext_q4_f32_disp<5, block_q4_1,   32, dequantize_q4_1_t4>;
876
877template [[host_name("kernel_mul_mv_ext_q5_0_f32_r1_2")]]   kernel mul_mv_ext_q4_f32_t kernel_mul_mv_ext_q4_f32_disp<2, block_q5_0,   32, dequantize_q5_0_t4>;
878template [[host_name("kernel_mul_mv_ext_q5_0_f32_r1_3")]]   kernel mul_mv_ext_q4_f32_t kernel_mul_mv_ext_q4_f32_disp<3, block_q5_0,   32, dequantize_q5_0_t4>;
879template [[host_name("kernel_mul_mv_ext_q5_0_f32_r1_4")]]   kernel mul_mv_ext_q4_f32_t kernel_mul_mv_ext_q4_f32_disp<4, block_q5_0,   32, dequantize_q5_0_t4>;
880template [[host_name("kernel_mul_mv_ext_q5_0_f32_r1_5")]]   kernel mul_mv_ext_q4_f32_t kernel_mul_mv_ext_q4_f32_disp<5, block_q5_0,   32, dequantize_q5_0_t4>;
881
882template [[host_name("kernel_mul_mv_ext_q5_1_f32_r1_2")]]   kernel mul_mv_ext_q4_f32_t kernel_mul_mv_ext_q4_f32_disp<2, block_q5_1,   32, dequantize_q5_1_t4>;
883template [[host_name("kernel_mul_mv_ext_q5_1_f32_r1_3")]]   kernel mul_mv_ext_q4_f32_t kernel_mul_mv_ext_q4_f32_disp<3, block_q5_1,   32, dequantize_q5_1_t4>;
884template [[host_name("kernel_mul_mv_ext_q5_1_f32_r1_4")]]   kernel mul_mv_ext_q4_f32_t kernel_mul_mv_ext_q4_f32_disp<4, block_q5_1,   32, dequantize_q5_1_t4>;
885template [[host_name("kernel_mul_mv_ext_q5_1_f32_r1_5")]]   kernel mul_mv_ext_q4_f32_t kernel_mul_mv_ext_q4_f32_disp<5, block_q5_1,   32, dequantize_q5_1_t4>;
886
887template [[host_name("kernel_mul_mv_ext_q8_0_f32_r1_2")]]   kernel mul_mv_ext_q4_f32_t kernel_mul_mv_ext_q4_f32_disp<2, block_q8_0,   32, dequantize_q8_0_t4>;
888template [[host_name("kernel_mul_mv_ext_q8_0_f32_r1_3")]]   kernel mul_mv_ext_q4_f32_t kernel_mul_mv_ext_q4_f32_disp<3, block_q8_0,   32, dequantize_q8_0_t4>;
889template [[host_name("kernel_mul_mv_ext_q8_0_f32_r1_4")]]   kernel mul_mv_ext_q4_f32_t kernel_mul_mv_ext_q4_f32_disp<4, block_q8_0,   32, dequantize_q8_0_t4>;
890template [[host_name("kernel_mul_mv_ext_q8_0_f32_r1_5")]]   kernel mul_mv_ext_q4_f32_t kernel_mul_mv_ext_q4_f32_disp<5, block_q8_0,   32, dequantize_q8_0_t4>;
891
892template [[host_name("kernel_mul_mv_ext_mxfp4_f32_r1_2")]]  kernel mul_mv_ext_q4_f32_t kernel_mul_mv_ext_q4_f32_disp<2, block_mxfp4,  32, dequantize_mxfp4_t4>;
893template [[host_name("kernel_mul_mv_ext_mxfp4_f32_r1_3")]]  kernel mul_mv_ext_q4_f32_t kernel_mul_mv_ext_q4_f32_disp<3, block_mxfp4,  32, dequantize_mxfp4_t4>;
894template [[host_name("kernel_mul_mv_ext_mxfp4_f32_r1_4")]]  kernel mul_mv_ext_q4_f32_t kernel_mul_mv_ext_q4_f32_disp<4, block_mxfp4,  32, dequantize_mxfp4_t4>;
895template [[host_name("kernel_mul_mv_ext_mxfp4_f32_r1_5")]]  kernel mul_mv_ext_q4_f32_t kernel_mul_mv_ext_q4_f32_disp<5, block_mxfp4,  32, dequantize_mxfp4_t4>;
896
897template [[host_name("kernel_mul_mv_ext_iq4_nl_f32_r1_2")]] kernel mul_mv_ext_q4_f32_t kernel_mul_mv_ext_q4_f32_disp<2, block_iq4_nl, 32, dequantize_iq4_nl_t4>;
898template [[host_name("kernel_mul_mv_ext_iq4_nl_f32_r1_3")]] kernel mul_mv_ext_q4_f32_t kernel_mul_mv_ext_q4_f32_disp<3, block_iq4_nl, 32, dequantize_iq4_nl_t4>;
899template [[host_name("kernel_mul_mv_ext_iq4_nl_f32_r1_4")]] kernel mul_mv_ext_q4_f32_t kernel_mul_mv_ext_q4_f32_disp<4, block_iq4_nl, 32, dequantize_iq4_nl_t4>;
900template [[host_name("kernel_mul_mv_ext_iq4_nl_f32_r1_5")]] kernel mul_mv_ext_q4_f32_t kernel_mul_mv_ext_q4_f32_disp<5, block_iq4_nl, 32, dequantize_iq4_nl_t4>;
901
902template [[host_name("kernel_mul_mv_ext_q4_K_f32_r1_2")]] kernel mul_mv_ext_q4x4_f32_t kernel_mul_mv_ext_q4x4_f32_disp<2, block_q4_K, 256, dequantize_q4_K>;
903template [[host_name("kernel_mul_mv_ext_q4_K_f32_r1_3")]] kernel mul_mv_ext_q4x4_f32_t kernel_mul_mv_ext_q4x4_f32_disp<3, block_q4_K, 256, dequantize_q4_K>;
904template [[host_name("kernel_mul_mv_ext_q4_K_f32_r1_4")]] kernel mul_mv_ext_q4x4_f32_t kernel_mul_mv_ext_q4x4_f32_disp<4, block_q4_K, 256, dequantize_q4_K>;
905template [[host_name("kernel_mul_mv_ext_q4_K_f32_r1_5")]] kernel mul_mv_ext_q4x4_f32_t kernel_mul_mv_ext_q4x4_f32_disp<5, block_q4_K, 256, dequantize_q4_K>;
906
907template [[host_name("kernel_mul_mv_ext_q5_K_f32_r1_2")]] kernel mul_mv_ext_q4x4_f32_t kernel_mul_mv_ext_q4x4_f32_disp<2, block_q5_K, 256, dequantize_q5_K>;
908template [[host_name("kernel_mul_mv_ext_q5_K_f32_r1_3")]] kernel mul_mv_ext_q4x4_f32_t kernel_mul_mv_ext_q4x4_f32_disp<3, block_q5_K, 256, dequantize_q5_K>;
909template [[host_name("kernel_mul_mv_ext_q5_K_f32_r1_4")]] kernel mul_mv_ext_q4x4_f32_t kernel_mul_mv_ext_q4x4_f32_disp<4, block_q5_K, 256, dequantize_q5_K>;
910template [[host_name("kernel_mul_mv_ext_q5_K_f32_r1_5")]] kernel mul_mv_ext_q4x4_f32_t kernel_mul_mv_ext_q4x4_f32_disp<5, block_q5_K, 256, dequantize_q5_K>;
911
912template [[host_name("kernel_mul_mv_ext_q6_K_f32_r1_2")]] kernel mul_mv_ext_q4x4_f32_t kernel_mul_mv_ext_q4x4_f32_disp<2, block_q6_K, 256, dequantize_q6_K>;
913template [[host_name("kernel_mul_mv_ext_q6_K_f32_r1_3")]] kernel mul_mv_ext_q4x4_f32_t kernel_mul_mv_ext_q4x4_f32_disp<3, block_q6_K, 256, dequantize_q6_K>;
914template [[host_name("kernel_mul_mv_ext_q6_K_f32_r1_4")]] kernel mul_mv_ext_q4x4_f32_t kernel_mul_mv_ext_q4x4_f32_disp<4, block_q6_K, 256, dequantize_q6_K>;
915template [[host_name("kernel_mul_mv_ext_q6_K_f32_r1_5")]] kernel mul_mv_ext_q4x4_f32_t kernel_mul_mv_ext_q4x4_f32_disp<5, block_q6_K, 256, dequantize_q6_K>;
916
917template [[host_name("kernel_mul_mv_ext_q2_K_f32_r1_2")]] kernel mul_mv_ext_q4x4_f32_t kernel_mul_mv_ext_q4x4_f32_disp<2, block_q2_K, 256, dequantize_q2_K>;
918template [[host_name("kernel_mul_mv_ext_q2_K_f32_r1_3")]] kernel mul_mv_ext_q4x4_f32_t kernel_mul_mv_ext_q4x4_f32_disp<3, block_q2_K, 256, dequantize_q2_K>;
919template [[host_name("kernel_mul_mv_ext_q2_K_f32_r1_4")]] kernel mul_mv_ext_q4x4_f32_t kernel_mul_mv_ext_q4x4_f32_disp<4, block_q2_K, 256, dequantize_q2_K>;
920template [[host_name("kernel_mul_mv_ext_q2_K_f32_r1_5")]] kernel mul_mv_ext_q4x4_f32_t kernel_mul_mv_ext_q4x4_f32_disp<5, block_q2_K, 256, dequantize_q2_K>;
921
922template [[host_name("kernel_mul_mv_ext_q3_K_f32_r1_2")]] kernel mul_mv_ext_q4x4_f32_t kernel_mul_mv_ext_q4x4_f32_disp<2, block_q3_K, 256, dequantize_q3_K>;
923template [[host_name("kernel_mul_mv_ext_q3_K_f32_r1_3")]] kernel mul_mv_ext_q4x4_f32_t kernel_mul_mv_ext_q4x4_f32_disp<3, block_q3_K, 256, dequantize_q3_K>;
924template [[host_name("kernel_mul_mv_ext_q3_K_f32_r1_4")]] kernel mul_mv_ext_q4x4_f32_t kernel_mul_mv_ext_q4x4_f32_disp<4, block_q3_K, 256, dequantize_q3_K>;
925template [[host_name("kernel_mul_mv_ext_q3_K_f32_r1_5")]] kernel mul_mv_ext_q4x4_f32_t kernel_mul_mv_ext_q4x4_f32_disp<5, block_q3_K, 256, dequantize_q3_K>;
926
927template<typename T0, typename T1, short NR0, typename args_t>
928void kernel_mul_mv_t_t_impl(
929        args_t args,
930        device const char * src0,
931        device const char * src1,
932        device       char * dst,
933        threadgroup  char * shmem,
934        uint3  tgpig,
935        ushort tiisg,
936        ushort sgitg) {
937    const short NSG = FC_mul_mv_nsg;
938
939    constexpr short NW = N_SIMDWIDTH;
940    constexpr short NB = 32;
941    constexpr short NF = 8;
942
943    const int nb = args.ne00/NB;
944
945    const int r0 = tgpig.x*NR0;
946    const int r1 = tgpig.y;
947    const int im = tgpig.z;
948
949    const uint i12 = im%FC_mul_mv_ne12;
950    const uint i13 = im/FC_mul_mv_ne12;
951
952  //const uint64_t offset0 = r0*args.nb01 + (i12/FC_mul_mv_r2)*args.nb02 + (i13/FC_mul_mv_r3)*args.nb03;
953    const uint64_t offset1 = r1*args.nb11 + (i12        )*args.nb12 + (i13        )*args.nb13;
954
955  //device const T0 * x = (device const T0 *) (src0 + offset0);
956    device const T1 * y = (device const T1 *) (src1 + offset1);
957
958    // pointers to src0 rows
959    device const T0 * ax [NR0];
960    FOR_UNROLL (short row = 0; row < NR0; ++row) {
961        const uint64_t offset0 = (r0 + row)*args.nb01 + (i12/FC_mul_mv_r2)*args.nb02 + (i13/FC_mul_mv_r3)*args.nb03;
962
963        ax[row] = (device const T0 *) ((device char *) src0 + offset0);
964    }
965
966    float sumf[NR0] = { 0.f };
967
968    const short ix = tiisg/(NW/NF);
969    const short il = tiisg%(NW/NF);
970
971    const int ib0 = sgitg*NF + ix;
972
973    T1 yl[NF];
974
975    device const T1 * yb = y + (ib0*NB + il*NF);
976
977    for (int ib = ib0; ib < nb; ib += NSG*NF) {
978        for (short i = 0; i < NF; ++i) {
979            yl[i] = yb[i];
980        }
981
982        for (short row = 0; row < NR0; row++) {
983            device const T0 * xb = ax[row] + (ib*NB + il*NF);
984
985            float sumq = 0.f;
986            FOR_UNROLL (short i = 0; i < NF; ++i) {
987                sumq += xb[i] * yl[i];
988            }
989
990            sumf[row] += sumq;
991        }
992
993        yb += NSG*NF*NW;
994    }
995
996    for (int i = nb*NB + sgitg*NW + tiisg; i < args.ne00; i += NW*NSG) {
997        for (short row = 0; row < NR0; row++) {
998            sumf[row] += ax[row][i] * y[i];
999        }
1000    }
1001
1002    device float * dst_f32 = (device float *) dst + (uint64_t)im*args.ne0*args.ne1 + (uint64_t)r1*args.ne0;
1003
1004    helper_mv_reduce_and_write<NR0>(dst_f32, sumf, r0, args.ne01, tiisg, sgitg, shmem);
1005}
1006
1007template<typename T0, typename T1, typename args_t>
1008void kernel_mul_mv_t_t_disp(
1009        args_t args,
1010        device const char * src0,
1011        device const char * src1,
1012        device       char * dst,
1013        threadgroup  char * shmem,
1014        uint3  tgpig,
1015        ushort tiisg,
1016        ushort sgitg) {
1017    switch (args.nr0) {
1018      //case 1: kernel_mul_mv_t_t_impl<T0, T1, 1, args_t>(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg); break;
1019        case 2: kernel_mul_mv_t_t_impl<T0, T1, 2, args_t>(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg); break;
1020      //case 3: kernel_mul_mv_t_t_impl<T0, T1, 3, args_t>(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg); break;
1021      //case 4: kernel_mul_mv_t_t_impl<T0, T1, 4, args_t>(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg); break;
1022    }
1023}
1024
1025template<typename T0, typename T1>
1026kernel void kernel_mul_mv_t_t(
1027        constant ggml_metal_kargs_mul_mv & args,
1028        device const char * src0,
1029        device const char * src1,
1030        device       char * dst,
1031        threadgroup  char * shmem [[threadgroup(0)]],
1032        uint3  tgpig[[threadgroup_position_in_grid]],
1033        ushort tiisg[[thread_index_in_simdgroup]],
1034        ushort sgitg[[simdgroup_index_in_threadgroup]]) {
1035    kernel_mul_mv_t_t_disp<T0, T1, constant ggml_metal_kargs_mul_mv &>(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg);
1036}
1037
1038typedef decltype(kernel_mul_mv_t_t<half, half>) mul_mv_t_t;
1039
1040template [[host_name("kernel_mul_mv_f32_f32")]]   kernel mul_mv_t_t kernel_mul_mv_t_t<float, float>;
1041template [[host_name("kernel_mul_mv_f16_f32")]]   kernel mul_mv_t_t kernel_mul_mv_t_t<half,  float>;
1042template [[host_name("kernel_mul_mv_f16_f16")]]   kernel mul_mv_t_t kernel_mul_mv_t_t<half,  half>;
1043#if defined(GGML_METAL_HAS_BF16)
1044template [[host_name("kernel_mul_mv_bf16_f32")]]  kernel mul_mv_t_t kernel_mul_mv_t_t<bfloat, float>;
1045template [[host_name("kernel_mul_mv_bf16_bf16")]] kernel mul_mv_t_t kernel_mul_mv_t_t<bfloat, bfloat>;
1046#endif
1047
1048template<typename T0, typename T04, typename T1, typename T14, short NR0, typename args_t>
1049void kernel_mul_mv_t_t_4_impl(
1050        args_t args,
1051        device const char * src0,
1052        device const char * src1,
1053        device       char * dst,
1054        threadgroup  char * shmem,
1055        uint3  tgpig,
1056        ushort tiisg,
1057        ushort sgitg) {
1058    const short NSG = FC_mul_mv_nsg;
1059
1060    constexpr short NW = N_SIMDWIDTH;
1061    constexpr short NB  = 32;
1062    constexpr short NF  = 16;
1063    constexpr short NF4 = NF/4;
1064
1065    const int nb = args.ne00/NB;
1066
1067    const int r0 = tgpig.x*NR0;
1068    const int r1 = tgpig.y;
1069    const int im = tgpig.z;
1070
1071    const uint i12 = im%FC_mul_mv_ne12;
1072    const uint i13 = im/FC_mul_mv_ne12;
1073
1074  //const uint64_t offset0 = r0*args.nb01 + (i12/FC_mul_mv_r2)*args.nb02 + (i13/FC_mul_mv_r3)*args.nb03;
1075    const uint64_t offset1 = r1*args.nb11 + (i12        )*args.nb12 + (i13        )*args.nb13;
1076
1077    device const T1  * y  = (device const T1  *) (src1 + offset1);
1078    device const T14 * y4 = (device const T14 *) (src1 + offset1);
1079
1080    // pointers to src0 rows
1081    device const T0  * ax [NR0];
1082    device const T04 * ax4[NR0];
1083    FOR_UNROLL (short row = 0; row < NR0; ++row) {
1084        const uint64_t offset0 = (r0 + row)*args.nb01 + (i12/FC_mul_mv_r2)*args.nb02 + (i13/FC_mul_mv_r3)*args.nb03;
1085
1086        ax [row] = (device const T0  *) ((device char *) src0 + offset0);
1087        ax4[row] = (device const T04 *) ((device char *) src0 + offset0);
1088    }
1089
1090    float sumf[NR0] = { 0.f };
1091
1092    const short ix = tiisg/(NW/NF);
1093    const short il = tiisg%(NW/NF);
1094
1095    const int ib0 = sgitg*NF + ix;
1096
1097    T14 yl4[NF4];
1098
1099    device const T14 * yb4 = y4 + (ib0*NB + il*NF)/4;
1100
1101    for (int ib = ib0; ib < nb; ib += NSG*NF) {
1102        for (short i = 0; i < NF4; ++i) {
1103            yl4[i] = yb4[i];
1104        }
1105
1106        for (short row = 0; row < NR0; row++) {
1107            device const T04 * xb4 = ax4[row] + (ib*NB + il*NF)/4;
1108
1109            float sumq = 0.f;
1110            FOR_UNROLL (short i = 0; i < NF4; ++i) {
1111                sumq += dot(float4(xb4[i]), float4(yl4[i]));
1112            }
1113
1114            sumf[row] += sumq;
1115        }
1116
1117        yb4 += NSG*NF*NW/4;
1118    }
1119
1120    for (int i = nb*NB + sgitg*NW + tiisg; i < args.ne00; i += NW*NSG) {
1121        for (short row = 0; row < NR0; row++) {
1122            sumf[row] += ax[row][i] * y[i];
1123        }
1124    }
1125
1126    device float * dst_f32 = (device float *) dst + (uint64_t)im*args.ne0*args.ne1 + (uint64_t)r1*args.ne0;
1127
1128    helper_mv_reduce_and_write<NR0>(dst_f32, sumf, r0, args.ne01, tiisg, sgitg, shmem);
1129}
1130
1131template<typename T0, typename T04, typename T1, typename T14, typename args_t>
1132void kernel_mul_mv_t_t_4_disp(
1133        args_t args,
1134        device const char * src0,
1135        device const char * src1,
1136        device       char * dst,
1137        threadgroup  char * shmem,
1138        uint3  tgpig,
1139        ushort tiisg,
1140        ushort sgitg) {
1141    switch (args.nr0) {
1142      //case 1: kernel_mul_mv_t_t_4_impl<T0, T04, T1, T14, 1, args_t>(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg); break;
1143        case 2: kernel_mul_mv_t_t_4_impl<T0, T04, T1, T14, 2, args_t>(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg); break;
1144      //case 3: kernel_mul_mv_t_t_4_impl<T0, T04, T1, T14, 3, args_t>(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg); break;
1145      //case 4: kernel_mul_mv_t_t_4_impl<T0, T04, T1, T14, 4, args_t>(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg); break;
1146    };
1147}
1148
1149template<typename T0, typename T04, typename T1, typename T14>
1150kernel void kernel_mul_mv_t_t_4(
1151        constant ggml_metal_kargs_mul_mv & args,
1152        device const char * src0,
1153        device const char * src1,
1154        device       char * dst,
1155        threadgroup  char * shmem [[threadgroup(0)]],
1156        uint3  tgpig[[threadgroup_position_in_grid]],
1157        ushort tiisg[[thread_index_in_simdgroup]],
1158        ushort sgitg[[simdgroup_index_in_threadgroup]]) {
1159    kernel_mul_mv_t_t_4_disp<T0, T04, T1, T14, constant ggml_metal_kargs_mul_mv &>(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg);
1160}
1161
1162typedef decltype(kernel_mul_mv_t_t_4<half, half4, half, half4>) mul_mv_t_t_4;
1163
1164template [[host_name("kernel_mul_mv_f32_f32_4")]]   kernel mul_mv_t_t_4 kernel_mul_mv_t_t_4<float, float4, float, float4>;
1165template [[host_name("kernel_mul_mv_f16_f32_4")]]   kernel mul_mv_t_t_4 kernel_mul_mv_t_t_4<half,  half4,  float, float4>;
1166template [[host_name("kernel_mul_mv_f16_f16_4")]]   kernel mul_mv_t_t_4 kernel_mul_mv_t_t_4<half,  half4,  half,  half4>;
1167#if defined(GGML_METAL_HAS_BF16)
1168template [[host_name("kernel_mul_mv_bf16_f32_4")]]  kernel mul_mv_t_t_4 kernel_mul_mv_t_t_4<bfloat, bfloat4, float,  float4>;
1169template [[host_name("kernel_mul_mv_bf16_bf16_4")]] kernel mul_mv_t_t_4 kernel_mul_mv_t_t_4<bfloat, bfloat4, bfloat, bfloat4>;
1170#endif
1171
1172template<typename T0, typename T1, typename args_t>
1173void kernel_mul_mv_t_t_short_impl(
1174        args_t args,
1175        device const char * src0,
1176        device const char * src1,
1177        device       char * dst,
1178        uint3  tgpig,
1179        ushort tiisg) {
1180    const int r0 = tgpig.x*32 + tiisg;
1181    const int r1 = tgpig.y;
1182    const int im = tgpig.z;
1183
1184    if (r0 >= args.ne01) {
1185        return;
1186    }
1187
1188    const uint i12 = im%FC_mul_mv_ne12;
1189    const uint i13 = im/FC_mul_mv_ne12;
1190
1191    const uint64_t offset0 = r0*args.nb01 + (i12/FC_mul_mv_r2)*args.nb02 + (i13/FC_mul_mv_r3)*args.nb03;
1192
1193    device const T0 * x = (device const T0 *) (src0 + offset0);
1194
1195    device float * dst_f32 = (device float *) dst + (uint64_t)im*args.ne0*args.ne1;
1196
1197    const uint64_t offset1 = r1*args.nb11 + (i12   )*args.nb12 + (i13   )*args.nb13;
1198
1199    device const T1 * y = (device const T1 *) (src1 + offset1);
1200
1201    float res = 0.0f;
1202
1203    for (int i = 0; i < args.ne00; ++i) {
1204        res += (float) x[i] * (float) y[i];
1205    }
1206
1207    dst_f32[(uint64_t)r1*args.ne0 + r0] = res;
1208}
1209
1210template<typename T0, typename T1>
1211kernel void kernel_mul_mv_t_t_short(
1212        constant ggml_metal_kargs_mul_mv & args,
1213        device const char * src0,
1214        device const char * src1,
1215        device       char * dst,
1216        uint3  tgpig[[threadgroup_position_in_grid]],
1217        ushort tiisg[[thread_index_in_simdgroup]]) {
1218    kernel_mul_mv_t_t_short_impl<T0, T1, constant ggml_metal_kargs_mul_mv &>(
1219        args,
1220        src0,
1221        src1,
1222        dst,
1223        tgpig,
1224        tiisg);
1225}
1226
1227typedef decltype(kernel_mul_mv_t_t_short<half, half>) mul_mv_t_t_short_t;
1228
1229template [[host_name("kernel_mul_mv_f32_f32_short")]]  kernel mul_mv_t_t_short_t kernel_mul_mv_t_t_short<float, float>;
1230template [[host_name("kernel_mul_mv_f16_f32_short")]]  kernel mul_mv_t_t_short_t kernel_mul_mv_t_t_short<half,  float>;
1231template [[host_name("kernel_mul_mv_f16_f16_short")]]  kernel mul_mv_t_t_short_t kernel_mul_mv_t_t_short<half,  half>;
1232#if defined(GGML_METAL_HAS_BF16)
1233template [[host_name("kernel_mul_mv_bf16_f32_short")]]  kernel mul_mv_t_t_short_t kernel_mul_mv_t_t_short<bfloat, float>;
1234template [[host_name("kernel_mul_mv_bf16_bf16_short")]] kernel mul_mv_t_t_short_t kernel_mul_mv_t_t_short<bfloat, bfloat>;
1235#endif
1236
1237template<int nr0, typename args_t>
1238void kernel_mul_mv_q2_K_f32_impl(
1239        args_t args,
1240        device const char * src0,
1241        device const char * src1,
1242        device       char * dst,
1243        threadgroup  char * shmem,
1244        uint3  tgpig,
1245        ushort tiisg,
1246        ushort sgitg) {
1247    const short NSG = FC_mul_mv_nsg;
1248
1249    const int nb = args.ne00/QK_K;
1250
1251    const int r0 = tgpig.x;
1252    const int r1 = tgpig.y;
1253    const int im = tgpig.z;
1254
1255    const int first_row = (r0 * NSG + sgitg) * nr0;
1256
1257    const uint i12 = im%FC_mul_mv_ne12;
1258    const uint i13 = im/FC_mul_mv_ne12;
1259
1260    const uint64_t offset0 = first_row*args.nb01 + (i12/FC_mul_mv_r2)*args.nb02 + (i13/FC_mul_mv_r3)*args.nb03;
1261    const uint64_t offset1 =        r1*args.nb11 + (i12        )*args.nb12 + (i13        )*args.nb13;
1262
1263    device const block_q2_K * x = (device const block_q2_K *) (src0 + offset0);
1264    device const float      * y = (device const float      *) (src1 + offset1);
1265
1266    float yl[32];
1267    float sumf[nr0]={0.f};
1268
1269    const short ix = tiisg/8;  // 0...3
1270    const short it = tiisg%8;  // 0...7
1271    const short iq = it/4;     // 0 or 1
1272    const short ir = it%4;     // 0...3
1273    const short is = (8*ir)/16;// 0 or 1
1274
1275    device const float * y4 = y + ix * QK_K + 128 * iq + 8 * ir;
1276
1277    for (int ib = ix; ib < nb; ib += 4) {
1278        float4 sumy = {0.f, 0.f, 0.f, 0.f};
1279        for (short i = 0; i < 8; ++i) {
1280            yl[i+ 0] = y4[i+ 0]; sumy[0] += yl[i+ 0];
1281            yl[i+ 8] = y4[i+32]; sumy[1] += yl[i+ 8];
1282            yl[i+16] = y4[i+64]; sumy[2] += yl[i+16];
1283            yl[i+24] = y4[i+96]; sumy[3] += yl[i+24];
1284        }
1285
1286        device const uint8_t  * sc = (device const uint8_t  *)x[ib].scales + 8*iq + is;
1287        device const uint16_t * qs = (device const uint16_t *)x[ib].qs + 16 * iq + 4 * ir;
1288        device const half     * dh = &x[ib].d;
1289
1290        for (short row = 0; row < nr0; row++) {
1291            float4 acc1 = {0.f, 0.f, 0.f, 0.f};
1292            float4 acc2 = {0.f, 0.f, 0.f, 0.f};
1293            for (int i = 0; i < 8; i += 2) {
1294                acc1[0] += yl[i+ 0] * (qs[i/2] & 0x0003);
1295                acc2[0] += yl[i+ 1] * (qs[i/2] & 0x0300);
1296                acc1[1] += yl[i+ 8] * (qs[i/2] & 0x000c);
1297                acc2[1] += yl[i+ 9] * (qs[i/2] & 0x0c00);
1298                acc1[2] += yl[i+16] * (qs[i/2] & 0x0030);
1299                acc2[2] += yl[i+17] * (qs[i/2] & 0x3000);
1300                acc1[3] += yl[i+24] * (qs[i/2] & 0x00c0);
1301                acc2[3] += yl[i+25] * (qs[i/2] & 0xc000);
1302            }
1303            float dall = dh[0];
1304            float dmin = dh[1] * 1.f/16.f;
1305            sumf[row] += dall * ((acc1[0] + 1.f/256.f * acc2[0]) * (sc[0] & 0xF) * 1.f/ 1.f +
1306                                 (acc1[1] + 1.f/256.f * acc2[1]) * (sc[2] & 0xF) * 1.f/ 4.f +
1307                                 (acc1[2] + 1.f/256.f * acc2[2]) * (sc[4] & 0xF) * 1.f/16.f +
1308                                 (acc1[3] + 1.f/256.f * acc2[3]) * (sc[6] & 0xF) * 1.f/64.f) -
1309                         dmin * (sumy[0] * (sc[0] & 0xF0) + sumy[1] * (sc[2] & 0xF0) + sumy[2] * (sc[4] & 0xF0) + sumy[3] * (sc[6] & 0xF0));
1310
1311            qs += args.nb01/2;
1312            sc += args.nb01;
1313            dh += args.nb01/2;
1314        }
1315
1316        y4 += 4 * QK_K;
1317    }
1318
1319    device float * dst_f32 = (device float *) dst + (uint64_t)im*args.ne0*args.ne1 + (uint64_t)r1*args.ne0;
1320
1321    for (int row = 0; row < nr0 && first_row + row < args.ne0; ++row) {
1322        float sum_all = simd_sum(sumf[row]);
1323        if (tiisg == 0) {
1324            dst_f32[first_row + row] = sum_all;
1325        }
1326    }
1327}
1328
1329[[host_name("kernel_mul_mv_q2_K_f32")]]
1330kernel void kernel_mul_mv_q2_K_f32(
1331        constant ggml_metal_kargs_mul_mv & args,
1332        device const char * src0,
1333        device const char * src1,
1334        device       char * dst,
1335        uint3  tgpig[[threadgroup_position_in_grid]],
1336        ushort tiisg[[thread_index_in_simdgroup]],
1337        ushort sgitg[[simdgroup_index_in_threadgroup]]) {
1338
1339    kernel_mul_mv_q2_K_f32_impl<N_R0_Q2_K, constant ggml_metal_kargs_mul_mv &>(args, src0, src1, dst, nullptr, tgpig, tiisg, sgitg);
1340}
1341
1342template<int nr0, typename args_t>
1343void kernel_mul_mv_q3_K_f32_impl(
1344        args_t args,
1345        device const char * src0,
1346        device const char * src1,
1347        device       char * dst,
1348        threadgroup  char * shmem,
1349        uint3  tgpig,
1350        ushort tiisg,
1351        ushort sgitg) {
1352    const short NSG = FC_mul_mv_nsg;
1353
1354    const int nb = args.ne00/QK_K;
1355
1356    const int r0 = tgpig.x;
1357    const int r1 = tgpig.y;
1358    const int im = tgpig.z;
1359
1360    const int first_row = (r0 * NSG + sgitg) * nr0;
1361
1362    const uint i12 = im%FC_mul_mv_ne12;
1363    const uint i13 = im/FC_mul_mv_ne12;
1364
1365    const uint64_t offset0 = first_row*args.nb01 + (i12/FC_mul_mv_r2)*args.nb02 + (i13/FC_mul_mv_r3)*args.nb03;
1366    const uint64_t offset1 =        r1*args.nb11 + (i12        )*args.nb12 + (i13        )*args.nb13;
1367
1368    device const block_q3_K * x = (device const block_q3_K *) (src0 + offset0);
1369    device const float     * yy = (device const float      *) (src1 + offset1);
1370
1371    float yl[32];
1372
1373    //const uint16_t kmask1 = 0x3030;
1374    //const uint16_t kmask2 = 0x0f0f;
1375
1376    const short tid = tiisg/4;
1377    const short ix  = tiisg%4;
1378    const short ip  = tid/4;          // 0 or 1
1379    const short il  = 2*((tid%4)/2);  // 0 or 2
1380    const short ir  = tid%2;
1381    const short l0  = 8*ir;
1382
1383    // One would think that the Metal compiler would figure out that ip and il can only have
1384    // 4 possible states, and optimize accordingly. Well, no. It needs help, and we do it
1385    // with these two tales.
1386    //
1387    // Possible masks for the high bit
1388    const ushort4 mm[4] = {{0x0001, 0x0100, 0x0002, 0x0200},  // ip = 0, il = 0
1389                           {0x0004, 0x0400, 0x0008, 0x0800},  // ip = 0, il = 2
1390                           {0x0010, 0x1000, 0x0020, 0x2000},  // ip = 1, il = 0
1391                           {0x0040, 0x4000, 0x0080, 0x8000}}; // ip = 1, il = 2
1392
1393    // Possible masks for the low 2 bits
1394    const int4 qm[2] = {{0x0003, 0x0300, 0x000c, 0x0c00}, {0x0030, 0x3000, 0x00c0, 0xc000}};
1395
1396    const ushort4 hm = mm[2*ip + il/2];
1397
1398    const short shift = 2*il;
1399
1400    const float v1 = il == 0 ? 4.f : 64.f;
1401    const float v2 = 4.f * v1;
1402
1403    const uint16_t s_shift1 = 4*ip;
1404    const uint16_t s_shift2 = s_shift1 + il;
1405
1406    const short q_offset = 32*ip + l0;
1407    const short y_offset = 128*ip + 32*il + l0;
1408
1409    device const float * y1 = yy + ix*QK_K + y_offset;
1410
1411    uint32_t scales32, aux32;
1412    thread uint16_t * scales16 = (thread uint16_t *)&scales32;
1413    thread const int8_t * scales = (thread const int8_t *)&scales32;
1414
1415    float sumf1[nr0] = {0.f};
1416    float sumf2[nr0] = {0.f};
1417
1418    for (int i = ix; i < nb; i += 4) {
1419        for (short l = 0; l < 8; ++l) {
1420            yl[l+ 0] = y1[l+ 0];
1421            yl[l+ 8] = y1[l+16];
1422            yl[l+16] = y1[l+32];
1423            yl[l+24] = y1[l+48];
1424        }
1425
1426        device const uint16_t * q = (device const uint16_t *)(x[i].qs + q_offset);
1427        device const uint16_t * h = (device const uint16_t *)(x[i].hmask + l0);
1428        device const uint16_t * a = (device const uint16_t *)(x[i].scales);
1429        device const half * dh = &x[i].d;
1430
1431        for (short row = 0; row < nr0; ++row) {
1432            const float d_all = (float)dh[0];
1433
1434            scales16[0] = a[4];
1435            scales16[1] = a[5];
1436            aux32 = ((scales32 >> s_shift2) << 4) & 0x30303030;
1437            scales16[0] = a[il+0];
1438            scales16[1] = a[il+1];
1439            scales32 = ((scales32 >> s_shift1) & 0x0f0f0f0f) | aux32;
1440
1441            float s1 = 0, s2 = 0, s3 = 0, s4 = 0, s5 = 0, s6 = 0;
1442            for (short l = 0; l < 8; l += 2) {
1443                const int32_t qs = q[l/2];
1444                s1 += yl[l+0] * (qs & qm[il/2][0]);
1445                s2 += yl[l+1] * (qs & qm[il/2][1]);
1446                s3 += ((h[l/2] & hm[0]) ? 0.f : yl[l+0]) + ((h[l/2] & hm[1]) ? 0.f : yl[l+1]);
1447                s4 += yl[l+16] * (qs & qm[il/2][2]);
1448                s5 += yl[l+17] * (qs & qm[il/2][3]);
1449                s6 += ((h[l/2] & hm[2]) ? 0.f : yl[l+16]) + ((h[l/2] & hm[3]) ? 0.f : yl[l+17]);
1450            }
1451            float d1 = d_all * (s1 + 1.f/256.f * s2 - s3*v1);
1452            float d2 = d_all * (s4 + 1.f/256.f * s5 - s6*v2);
1453            sumf1[row] += d1 * (scales[0] - 32);
1454            sumf2[row] += d2 * (scales[2] - 32);
1455
1456            s1 = s2 = s3 = s4 = s5 = s6 = 0;
1457            for (short l = 0; l < 8; l += 2) {
1458                const int32_t qs = q[l/2+8];
1459                s1 += yl[l+8] * (qs & qm[il/2][0]);
1460                s2 += yl[l+9] * (qs & qm[il/2][1]);
1461                s3 += ((h[l/2+8] & hm[0]) ? 0.f : yl[l+8]) + ((h[l/2+8] & hm[1]) ? 0.f : yl[l+9]);
1462                s4 += yl[l+24] * (qs & qm[il/2][2]);
1463                s5 += yl[l+25] * (qs & qm[il/2][3]);
1464                s6 += ((h[l/2+8] & hm[2]) ? 0.f : yl[l+24]) + ((h[l/2+8] & hm[3]) ? 0.f : yl[l+25]);
1465            }
1466            d1 = d_all * (s1 + 1.f/256.f * s2 - s3*v1);
1467            d2 = d_all * (s4 + 1.f/256.f * s5 - s6*v2);
1468            sumf1[row] += d1 * (scales[1] - 32);
1469            sumf2[row] += d2 * (scales[3] - 32);
1470
1471            q  += args.nb01/2;
1472            h  += args.nb01/2;
1473            a  += args.nb01/2;
1474            dh += args.nb01/2;
1475        }
1476
1477        y1 += 4 * QK_K;
1478    }
1479
1480    for (int row = 0; row < nr0; ++row) {
1481        const float sumf = (sumf1[row] + 0.25f * sumf2[row]) / (1 << shift);
1482        sumf1[row] = simd_sum(sumf);
1483    }
1484
1485    device float * dst_f32 = (device float *) dst + (uint64_t)im*args.ne0*args.ne1 + (uint64_t)r1*args.ne0;
1486
1487    if (tiisg == 0) {
1488        for (int row = 0; row < nr0 && first_row + row < args.ne0; ++row) {
1489            dst_f32[first_row + row] = sumf1[row];
1490        }
1491    }
1492}
1493
1494[[host_name("kernel_mul_mv_q3_K_f32")]]
1495kernel void kernel_mul_mv_q3_K_f32(
1496        constant ggml_metal_kargs_mul_mv & args,
1497        device const char * src0,
1498        device const char * src1,
1499        device       char * dst,
1500        uint3  tgpig[[threadgroup_position_in_grid]],
1501        ushort tiisg[[thread_index_in_simdgroup]],
1502        ushort sgitg[[simdgroup_index_in_threadgroup]]) {
1503
1504    kernel_mul_mv_q3_K_f32_impl<N_R0_Q3_K, constant ggml_metal_kargs_mul_mv &>(args, src0, src1, dst, nullptr, tgpig, tiisg, sgitg);
1505}
1506
1507template<int nr0, typename args_t>
1508void kernel_mul_mv_q4_K_f32_impl(
1509        args_t args,
1510        device const char * src0,
1511        device const char * src1,
1512        device       char * dst,
1513        threadgroup  char * shmem,
1514        uint3  tgpig,
1515        ushort tiisg,
1516        ushort sgitg) {
1517    const short NSG = FC_mul_mv_nsg;
1518
1519    constexpr uint16_t kmask1 = 0x3f3f;
1520    constexpr uint16_t kmask2 = 0x0f0f;
1521    constexpr uint16_t kmask3 = 0xc0c0;
1522
1523    const short ix = tiisg/8;  // 0...3
1524    const short it = tiisg%8;  // 0...7
1525    const short iq = it/4;     // 0 or 1
1526    const short ir = it%4;     // 0...3
1527
1528    const int nb = args.ne00/QK_K;
1529
1530    const int r0 = tgpig.x;
1531    const int r1 = tgpig.y;
1532    const int im = tgpig.z;
1533
1534    const int first_row = (r0 * NSG + sgitg) * nr0;
1535
1536    const uint i12 = im%FC_mul_mv_ne12;
1537    const uint i13 = im/FC_mul_mv_ne12;
1538
1539    const uint64_t offset0 = first_row*args.nb01 + (i12/FC_mul_mv_r2)*args.nb02 + (i13/FC_mul_mv_r3)*args.nb03;
1540    const uint64_t offset1 =        r1*args.nb11 + (i12        )*args.nb12 + (i13        )*args.nb13;
1541
1542    device const block_q4_K * x = (device const block_q4_K *) (src0 + offset0);
1543    device const float      * y = (device const float      *) (src1 + offset1);
1544
1545    float yl[16];
1546    float yh[16];
1547
1548    float sumf[nr0]={0.f};
1549
1550    device const float * y4 = y + ix * QK_K + 64 * iq + 8 * ir;
1551
1552    uint16_t sc16[4];
1553    thread const uint8_t * sc8 = (thread const uint8_t *)sc16;
1554
1555    for (int ib = ix; ib < nb; ib += 4) {
1556        float4 sumy = {0.f, 0.f, 0.f, 0.f};
1557
1558        for (short i = 0; i < 8; ++i) {
1559            yl[i+0] = y4[i+  0]; sumy[0] += yl[i+0];
1560            yl[i+8] = y4[i+ 32]; sumy[1] += yl[i+8];
1561            yh[i+0] = y4[i+128]; sumy[2] += yh[i+0];
1562            yh[i+8] = y4[i+160]; sumy[3] += yh[i+8];
1563        }
1564
1565        device const uint16_t * sc = (device const uint16_t *)x[ib].scales + iq;
1566        device const uint16_t * q1 = (device const uint16_t *)x[ib].qs + 16 * iq + 4 * ir;
1567        device const half     * dh = &x[ib].d;
1568
1569        for (short row = 0; row < nr0; row++) {
1570            sc16[0] = sc[0] & kmask1;
1571            sc16[1] = sc[2] & kmask1;
1572            sc16[2] = ((sc[4] >> 0) & kmask2) | ((sc[0] & kmask3) >> 2);
1573            sc16[3] = ((sc[4] >> 4) & kmask2) | ((sc[2] & kmask3) >> 2);
1574
1575            device const uint16_t * q2 = q1 + 32;
1576
1577            float4 acc1 = {0.f, 0.f, 0.f, 0.f};
1578            float4 acc2 = {0.f, 0.f, 0.f, 0.f};
1579
1580            FOR_UNROLL (short i = 0; i < 4; ++i) {
1581                acc1[0] += yl[2*i + 0] * (q1[i] & 0x000F);
1582                acc1[1] += yl[2*i + 1] * (q1[i] & 0x0F00);
1583                acc1[2] += yl[2*i + 8] * (q1[i] & 0x00F0);
1584                acc1[3] += yl[2*i + 9] * (q1[i] & 0xF000);
1585                acc2[0] += yh[2*i + 0] * (q2[i] & 0x000F);
1586                acc2[1] += yh[2*i + 1] * (q2[i] & 0x0F00);
1587                acc2[2] += yh[2*i + 8] * (q2[i] & 0x00F0);
1588                acc2[3] += yh[2*i + 9] * (q2[i] & 0xF000);
1589            }
1590
1591            sumf[row] += dh[0] * ((acc1[0] + 1.f/256.f * acc1[1]) * sc8[0] +
1592                                  (acc1[2] + 1.f/256.f * acc1[3]) * sc8[1] * 1.f/16.f +
1593                                  (acc2[0] + 1.f/256.f * acc2[1]) * sc8[4] +
1594                                  (acc2[2] + 1.f/256.f * acc2[3]) * sc8[5] * 1.f/16.f) -
1595                         dh[1] * (sumy[0] * sc8[2] + sumy[1] * sc8[3] + sumy[2] * sc8[6] + sumy[3] * sc8[7]);
1596
1597            q1 += args.nb01/2;
1598            sc += args.nb01/2;
1599            dh += args.nb01/2;
1600        }
1601
1602        y4 += 4 * QK_K;
1603    }
1604
1605    device float * dst_f32 = (device float *) dst + (int64_t)im*args.ne0*args.ne1 + (int64_t)r1*args.ne0;
1606
1607    for (int row = 0; row < nr0 && first_row + row < args.ne0; ++row) {
1608        float sum_all = simd_sum(sumf[row]);
1609        if (tiisg == 0) {
1610            dst_f32[first_row + row] = sum_all;
1611        }
1612    }
1613}
1614
1615[[host_name("kernel_mul_mv_q4_K_f32")]]
1616kernel void kernel_mul_mv_q4_K_f32(
1617        constant ggml_metal_kargs_mul_mv & args,
1618        device const char * src0,
1619        device const char * src1,
1620        device       char * dst,
1621        uint3  tgpig[[threadgroup_position_in_grid]],
1622        ushort tiisg[[thread_index_in_simdgroup]],
1623        ushort sgitg[[simdgroup_index_in_threadgroup]]) {
1624
1625    kernel_mul_mv_q4_K_f32_impl<N_R0_Q4_K, constant ggml_metal_kargs_mul_mv &>(args, src0, src1, dst, nullptr, tgpig, tiisg, sgitg);
1626}
1627
1628template<int nr0, typename args_t>
1629void kernel_mul_mv_q5_K_f32_impl(
1630        args_t args,
1631        device const char * src0,
1632        device const char * src1,
1633        device       char * dst,
1634        threadgroup  char * shmem,
1635        uint3  tgpig,
1636        ushort tiisg,
1637        ushort sgitg) {
1638    const short NSG = FC_mul_mv_nsg;
1639
1640    const int nb = args.ne00/QK_K;
1641
1642    const int r0 = tgpig.x;
1643    const int r1 = tgpig.y;
1644    const int im = tgpig.z;
1645
1646    const int first_row = (r0 * NSG + sgitg) * nr0;
1647
1648    const uint i12 = im%FC_mul_mv_ne12;
1649    const uint i13 = im/FC_mul_mv_ne12;
1650
1651    const uint64_t offset0 = first_row*args.nb01 + (i12/FC_mul_mv_r2)*args.nb02 + (i13/FC_mul_mv_r3)*args.nb03;
1652    const uint64_t offset1 =        r1*args.nb11 + (i12        )*args.nb12 + (i13        )*args.nb13;
1653
1654    device const block_q5_K * x = (device const block_q5_K *) (src0 + offset0);
1655    device const float     * yy = (device const float      *) (src1 + offset1);
1656
1657    float sumf[nr0]={0.f};
1658
1659    float yl[16], yh[16];
1660
1661    constexpr uint16_t kmask1 = 0x3f3f;
1662    constexpr uint16_t kmask2 = 0x0f0f;
1663    constexpr uint16_t kmask3 = 0xc0c0;
1664
1665    const short tid = tiisg/4;
1666    const short ix  = tiisg%4;
1667    const short iq  = tid/4;
1668    const short ir  = tid%4;
1669
1670    const short l0 = 8*ir;
1671    const short q_offset = 32*iq + l0;
1672    const short y_offset = 64*iq + l0;
1673
1674    const uint8_t hm1 = 1u << (2*iq);
1675    const uint8_t hm2 = hm1 << 1;
1676    const uint8_t hm3 = hm1 << 4;
1677    const uint8_t hm4 = hm2 << 4;
1678
1679    uint16_t sc16[4];
1680    thread const uint8_t * sc8 = (thread const uint8_t *)sc16;
1681
1682    device const float * y1 = yy + ix*QK_K + y_offset;
1683
1684    for (int i = ix; i < nb; i += 4) {
1685        device const uint8_t * q1 = x[i].qs + q_offset;
1686        device const uint8_t * qh = x[i].qh + l0;
1687        device const half * dh = &x[i].d;
1688        device const uint16_t * a = (device const uint16_t *)x[i].scales + iq;
1689
1690        device const float * y2 = y1 + 128;
1691        float4 sumy = {0.f, 0.f, 0.f, 0.f};
1692        for (short l = 0; l < 8; ++l) {
1693            yl[l+0] = y1[l+ 0]; sumy[0] += yl[l+0];
1694            yl[l+8] = y1[l+32]; sumy[1] += yl[l+8];
1695            yh[l+0] = y2[l+ 0]; sumy[2] += yh[l+0];
1696            yh[l+8] = y2[l+32]; sumy[3] += yh[l+8];
1697        }
1698
1699        for (short row = 0; row < nr0; ++row) {
1700            device const uint8_t * q2 = q1 + 64;
1701
1702            sc16[0] = a[0] & kmask1;
1703            sc16[1] = a[2] & kmask1;
1704            sc16[2] = ((a[4] >> 0) & kmask2) | ((a[0] & kmask3) >> 2);
1705            sc16[3] = ((a[4] >> 4) & kmask2) | ((a[2] & kmask3) >> 2);
1706
1707            float4 acc1 = {0.f};
1708            float4 acc2 = {0.f};
1709            FOR_UNROLL (short l = 0; l < 8; ++l) {
1710                uint8_t h = qh[l];
1711                acc1[0] += yl[l+0] * (q1[l] & 0x0F);
1712                acc1[1] += yl[l+8] * (q1[l] & 0xF0);
1713                acc1[2] += yh[l+0] * (q2[l] & 0x0F);
1714                acc1[3] += yh[l+8] * (q2[l] & 0xF0);
1715                acc2[0] += h & hm1 ? yl[l+0] : 0.f;
1716                acc2[1] += h & hm2 ? yl[l+8] : 0.f;
1717                acc2[2] += h & hm3 ? yh[l+0] : 0.f;
1718                acc2[3] += h & hm4 ? yh[l+8] : 0.f;
1719            }
1720
1721            sumf[row] += dh[0] * (sc8[0] * (acc1[0]      + 16.f*acc2[0]) +
1722                                  sc8[1] * (acc1[1]/16.f + 16.f*acc2[1]) +
1723                                  sc8[4] * (acc1[2]      + 16.f*acc2[2]) +
1724                                  sc8[5] * (acc1[3]/16.f + 16.f*acc2[3])) -
1725                         dh[1] * (sumy[0] * sc8[2] + sumy[1] * sc8[3] + sumy[2] * sc8[6] + sumy[3] * sc8[7]);
1726
1727            q1 += args.nb01;
1728            qh += args.nb01;
1729            dh += args.nb01/2;
1730            a  += args.nb01/2;
1731        }
1732
1733        y1 += 4 * QK_K;
1734    }
1735
1736    device float * dst_f32 = (device float *) dst + (uint64_t)im*args.ne0*args.ne1 + (uint64_t)r1*args.ne0;
1737
1738    for (int row = 0; row < nr0 && first_row + row < args.ne0; ++row) {
1739        const float tot = simd_sum(sumf[row]);
1740        if (tiisg == 0) {
1741            dst_f32[first_row + row] = tot;
1742        }
1743    }
1744}
1745
1746[[host_name("kernel_mul_mv_q5_K_f32")]]
1747kernel void kernel_mul_mv_q5_K_f32(
1748        constant ggml_metal_kargs_mul_mv & args,
1749        device const char * src0,
1750        device const char * src1,
1751        device       char * dst,
1752        uint3  tgpig[[threadgroup_position_in_grid]],
1753        ushort tiisg[[thread_index_in_simdgroup]],
1754        ushort sgitg[[simdgroup_index_in_threadgroup]]) {
1755
1756    kernel_mul_mv_q5_K_f32_impl<N_R0_Q5_K, constant ggml_metal_kargs_mul_mv &>(args, src0, src1, dst, nullptr, tgpig, tiisg, sgitg);
1757}
1758
1759template<int nr0, typename args_t>
1760void kernel_mul_mv_q6_K_f32_impl(
1761        args_t args,
1762        device const char * src0,
1763        device const char * src1,
1764        device       char * dst,
1765        threadgroup  char * shmem,
1766        uint3  tgpig,
1767        ushort tiisg,
1768        ushort sgitg) {
1769    const short NSG = FC_mul_mv_nsg;
1770
1771    constexpr uint8_t kmask1 = 0x03;
1772    constexpr uint8_t kmask2 = 0x0C;
1773    constexpr uint8_t kmask3 = 0x30;
1774    constexpr uint8_t kmask4 = 0xC0;
1775
1776    const int nb = args.ne00/QK_K;
1777
1778    const int r0 = tgpig.x;
1779    const int r1 = tgpig.y;
1780    const int im = tgpig.z;
1781
1782    const int first_row = (r0 * NSG + sgitg) * nr0;
1783
1784    const uint i12 = im%FC_mul_mv_ne12;
1785    const uint i13 = im/FC_mul_mv_ne12;
1786
1787    const uint64_t offset0 = first_row*args.nb01 + (i12/FC_mul_mv_r2)*args.nb02 + (i13/FC_mul_mv_r3)*args.nb03;
1788    const uint64_t offset1 =        r1*args.nb11 + (i12        )*args.nb12 + (i13        )*args.nb13;
1789
1790    device const block_q6_K * x = (device const block_q6_K *) (src0 + offset0);
1791    device const float     * yy = (device const float      *) (src1 + offset1);
1792
1793    float sumf[nr0] = { 0.f };
1794
1795    float yl[16];
1796
1797    const short tid = tiisg/2;
1798    const short ix  = tiisg%2;
1799    const short ip  = tid/8;         // 0 or 1
1800    const short il  = tid%8;
1801    const short l0  = 4*il;
1802    const short is  = 8*ip + l0/16;
1803
1804    const short y_offset   = 128*ip + l0;
1805    const short q_offset_l =  64*ip + l0;
1806    const short q_offset_h =  32*ip + l0;
1807
1808    for (int i = ix; i < nb; i += 2) {
1809        device const uint8_t * q1 = x[i].ql + q_offset_l;
1810        device const uint8_t * q2 = q1 + 32;
1811        device const uint8_t * qh = x[i].qh + q_offset_h;
1812        device const int8_t  * sc = x[i].scales + is;
1813        device const half    * dh = &x[i].d;
1814
1815        device const float * y = yy + i * QK_K + y_offset;
1816
1817        for (short l = 0; l < 4; ++l) {
1818            yl[4*l + 0] = y[l +  0];
1819            yl[4*l + 1] = y[l + 32];
1820            yl[4*l + 2] = y[l + 64];
1821            yl[4*l + 3] = y[l + 96];
1822        }
1823
1824        for (short row = 0; row < nr0; ++row) {
1825            float4 sums = {0.f, 0.f, 0.f, 0.f};
1826
1827            FOR_UNROLL (short l = 0; l < 4; ++l) {
1828                sums[0] += yl[4*l + 0] * ((int8_t)((q1[l] & 0xF) | ((qh[l] & kmask1) << 4)) - 32);
1829                sums[1] += yl[4*l + 1] * ((int8_t)((q2[l] & 0xF) | ((qh[l] & kmask2) << 2)) - 32);
1830                sums[2] += yl[4*l + 2] * ((int8_t)((q1[l]  >> 4) | ((qh[l] & kmask3) << 0)) - 32);
1831                sums[3] += yl[4*l + 3] * ((int8_t)((q2[l]  >> 4) | ((qh[l] & kmask4) >> 2)) - 32);
1832            }
1833
1834            sumf[row] += dh[0] * (sums[0] * sc[0] + sums[1] * sc[2] + sums[2] * sc[4] + sums[3] * sc[6]);
1835
1836            q1 += args.nb01;
1837            q2 += args.nb01;
1838            qh += args.nb01;
1839            sc += args.nb01;
1840            dh += args.nb01/2;
1841        }
1842    }
1843
1844    device float * dst_f32 = (device float *) dst + (uint64_t)im*args.ne0*args.ne1 + (uint64_t)r1*args.ne0;
1845
1846    for (int row = 0; row < nr0 && first_row + row < args.ne0; ++row) {
1847        float sum_all = simd_sum(sumf[row]);
1848        if (tiisg == 0) {
1849            dst_f32[first_row + row] = sum_all;
1850        }
1851    }
1852}
1853
1854[[host_name("kernel_mul_mv_q6_K_f32")]]
1855kernel void kernel_mul_mv_q6_K_f32(
1856        constant ggml_metal_kargs_mul_mv & args,
1857        device const char * src0,
1858        device const char * src1,
1859        device       char * dst,
1860        uint3  tgpig[[threadgroup_position_in_grid]],
1861        ushort tiisg[[thread_index_in_simdgroup]],
1862        ushort sgitg[[simdgroup_index_in_threadgroup]]) {
1863
1864    kernel_mul_mv_q6_K_f32_impl<N_R0_Q6_K, constant ggml_metal_kargs_mul_mv &>(args, src0, src1, dst, nullptr, tgpig, tiisg, sgitg);
1865}
1866
1867// ======================= "True" 2-bit
1868
1869template<int nr0, typename args_t>
1870void kernel_mul_mv_iq2_xxs_f32_impl(
1871        args_t args,
1872        device const char * src0,
1873        device const char * src1,
1874        device       char * dst,
1875        threadgroup  char * shmem,
1876        uint3  tgpig,
1877        ushort tiisg,
1878        ushort sgitg) {
1879    const short NSG = FC_mul_mv_nsg;
1880
1881    const int nb = args.ne00/QK_K;
1882
1883    const int r0 = tgpig.x;
1884    const int r1 = tgpig.y;
1885    const int im = tgpig.z;
1886
1887    const int first_row = (r0 * NSG + sgitg) * nr0;
1888
1889    const uint i12 = im%FC_mul_mv_ne12;
1890    const uint i13 = im/FC_mul_mv_ne12;
1891
1892    const int nb32 = nb * (QK_K / 32);
1893
1894    const short ntx  = FC_mul_mv_split ? nb32 : 32;
1895    const short nrep = 32 / ntx;
1896
1897    const short ix   = tiisg % ntx;
1898    const short irep = tiisg / ntx;
1899
1900    const short row0 = (nr0 * irep      ) / nrep;
1901    const short row1 = (nr0 * (irep + 1)) / nrep;
1902
1903    const uint64_t offset0 = (first_row + row0)*args.nb01 + (i12/FC_mul_mv_r2)*args.nb02 + (i13/FC_mul_mv_r3)*args.nb03;
1904    const uint64_t offset1 =                r1*args.nb11 + (i12        )*args.nb12 + (i13        )*args.nb13;
1905
1906    device const block_iq2_xxs * x = (device const block_iq2_xxs *) (src0 + offset0);
1907    device const float         * y = (device const float         *) (src1 + offset1);
1908
1909    float yl[32];
1910    float sumf[nr0]={0.f};
1911
1912    threadgroup uint64_t * svalues = (threadgroup uint64_t *)(shmem);
1913    threadgroup uint8_t  * ssigns  = (threadgroup uint8_t  *)(svalues + 256);
1914    {
1915        int nval = 4;
1916        int pos  = (32*sgitg + tiisg)*nval;
1917        for (int i = 0; i < nval; ++i) svalues[pos + i] = iq2xxs_grid[pos + i];
1918        nval = 2;
1919        pos  = (32*sgitg + tiisg)*nval;
1920        for (int i = 0; i < nval; ++i) ssigns[pos+i] = ksigns_iq2xs[pos+i];
1921        threadgroup_barrier(mem_flags::mem_threadgroup);
1922    }
1923
1924    device const float * y4 = y + 32 * ix;
1925
1926    for (int ib32 = ix; ib32 < nb32; ib32 += ntx) {
1927        for (short i = 0; i < 32; ++i) {
1928            yl[i] = y4[i];
1929        }
1930
1931        const int ibl = ib32 / (QK_K / 32);
1932        const int ib  = ib32 % (QK_K / 32);
1933
1934        device const block_iq2_xxs * xr = x + ibl;
1935        device const uint16_t * q2 = xr->qs + 4 * ib;
1936        device const half * dh = &xr->d;
1937
1938        for (short row = row0; row < row1; row++) {
1939            const float db = dh[0];
1940            device const uint8_t * aux8 = (device const uint8_t *)q2;
1941            const uint32_t aux32 = q2[2] | (q2[3] << 16);
1942            const float d = db * (0.5f + (aux32 >> 28));
1943
1944            float sum = 0;
1945            for (short l = 0; l < 4; ++l) {
1946                const threadgroup uint8_t * grid = (const threadgroup uint8_t *)(svalues + aux8[l]);
1947                const uint8_t signs = ssigns[(aux32 >> 7*l) & 127];
1948                for (short j = 0; j < 8; ++j) {
1949                    sum += yl[8*l + j] * grid[j] * (signs & kmask_iq2xs[j] ? -1.f : 1.f);
1950                }
1951            }
1952            sumf[row] += d * sum;
1953
1954            dh += args.nb01/2;
1955            q2 += args.nb01/2;
1956        }
1957
1958        y4 += 32 * ntx;
1959    }
1960
1961    device float * dst_f32 = (device float *) dst + (uint64_t)im*args.ne0*args.ne1 + (uint64_t)r1*args.ne0;
1962
1963    for (int row = 0; row < nr0 && first_row + row < args.ne0; ++row) {
1964        float sum_all = simd_sum(sumf[row]);
1965        if (tiisg == 0) {
1966            dst_f32[first_row + row] = sum_all * 0.25f;
1967        }
1968    }
1969}
1970
1971template<typename args_t>
1972void kernel_mul_mv_iq2_xxs_f32_disp(
1973        args_t args,
1974        device const char * src0,
1975        device const char * src1,
1976        device       char * dst,
1977        threadgroup  char * shmem,
1978        uint3  tgpig,
1979        ushort tiisg,
1980        ushort sgitg) {
1981    if (FC_mul_mv_split) {
1982        kernel_mul_mv_iq2_xxs_f32_impl<N_R0_IQ2_XXS_SPLIT, args_t>(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg);
1983    } else {
1984        kernel_mul_mv_iq2_xxs_f32_impl<N_R0_IQ2_XXS, args_t>(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg);
1985    }
1986}
1987
1988[[host_name("kernel_mul_mv_iq2_xxs_f32")]]
1989kernel void kernel_mul_mv_iq2_xxs_f32(
1990        constant ggml_metal_kargs_mul_mv & args,
1991        device const char * src0,
1992        device const char * src1,
1993        device       char * dst,
1994        threadgroup  char * shmem [[threadgroup(0)]],
1995        uint3  tgpig[[threadgroup_position_in_grid]],
1996        ushort tiisg[[thread_index_in_simdgroup]],
1997        ushort sgitg[[simdgroup_index_in_threadgroup]]) {
1998    kernel_mul_mv_iq2_xxs_f32_disp<constant ggml_metal_kargs_mul_mv &>(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg);
1999}
2000
2001template<int nr0, typename args_t>
2002void kernel_mul_mv_iq2_xs_f32_impl(
2003        args_t args,
2004        device const char * src0,
2005        device const char * src1,
2006        device       char * dst,
2007        threadgroup  char * shmem,
2008        uint3  tgpig,
2009        ushort tiisg,
2010        ushort sgitg) {
2011    const short NSG = FC_mul_mv_nsg;
2012
2013    const int nb = args.ne00/QK_K;
2014
2015    const int r0 = tgpig.x;
2016    const int r1 = tgpig.y;
2017    const int im = tgpig.z;
2018
2019    const int first_row = (r0 * NSG + sgitg) * nr0;
2020
2021    const uint i12 = im%FC_mul_mv_ne12;
2022    const uint i13 = im/FC_mul_mv_ne12;
2023
2024    const int nb32 = nb * (QK_K / 32);
2025
2026    const short ntx  = FC_mul_mv_split ? nb32 : 32;
2027    const short nrep = 32 / ntx;
2028
2029    const short ix   = tiisg % ntx;
2030    const short irep = tiisg / ntx;
2031
2032    const short row0 = (nr0 * irep      ) / nrep;
2033    const short row1 = (nr0 * (irep + 1)) / nrep;
2034
2035    const uint64_t offset0 = (first_row + row0)*args.nb01 + (i12/FC_mul_mv_r2)*args.nb02 + (i13/FC_mul_mv_r3)*args.nb03;
2036    const uint64_t offset1 =                r1*args.nb11 + (i12        )*args.nb12 + (i13        )*args.nb13;
2037
2038    device const block_iq2_xs * x = (device const block_iq2_xs *) (src0 + offset0);
2039    device const float        * y = (device const float        *) (src1 + offset1);
2040
2041    float yl[32];
2042    float sumf[nr0]={0.f};
2043
2044    threadgroup uint64_t * svalues = (threadgroup uint64_t *)(shmem);
2045    threadgroup uint8_t  * ssigns  = (threadgroup uint8_t  *)(svalues + 512);
2046    {
2047        int nval = 8;
2048        int pos  = (32*sgitg + tiisg)*nval;
2049        for (int i = 0; i < nval; ++i) svalues[pos + i] = iq2xs_grid[pos + i];
2050        nval = 2;
2051        pos  = (32*sgitg + tiisg)*nval;
2052        for (int i = 0; i < nval; ++i) ssigns[pos+i] = ksigns_iq2xs[pos+i];
2053        threadgroup_barrier(mem_flags::mem_threadgroup);
2054    }
2055
2056    device const float * y4 = y + 32 * ix;
2057
2058    for (int ib32 = ix; ib32 < nb32; ib32 += ntx) {
2059        for (short i = 0; i < 32; ++i) {
2060            yl[i] = y4[i];
2061        }
2062
2063        const int ibl = ib32 / (QK_K / 32);
2064        const int ib  = ib32 % (QK_K / 32);
2065
2066        device const block_iq2_xs * xr = x + ibl;
2067        device const uint16_t * q2 = xr->qs + 4 * ib;
2068        device const uint8_t  * sc = xr->scales + ib;
2069        device const half * dh = &xr->d;
2070
2071        for (short row = row0; row < row1; row++) {
2072            const float db = dh[0];
2073            const uint8_t ls1 = sc[0] & 0xf;
2074            const uint8_t ls2 = sc[0] >>  4;
2075            const float d1 = db * (0.5f + ls1);
2076            const float d2 = db * (0.5f + ls2);
2077
2078            float sum1 = 0, sum2 = 0;
2079            for (short l = 0; l < 2; ++l) {
2080                const threadgroup uint8_t * grid = (const threadgroup uint8_t *)(svalues + (q2[l] & 511));
2081                const uint8_t signs = ssigns[(q2[l] >> 9)];
2082                for (short j = 0; j < 8; ++j) {
2083                    sum1 += yl[8*l + j] * grid[j] * (signs & kmask_iq2xs[j] ? -1.f : 1.f);
2084                }
2085            }
2086            for (short l = 2; l < 4; ++l) {
2087                const threadgroup uint8_t * grid = (const threadgroup uint8_t *)(svalues + (q2[l] & 511));
2088                const uint8_t signs = ssigns[(q2[l] >> 9)];
2089                for (short j = 0; j < 8; ++j) {
2090                    sum2 += yl[8*l + j] * grid[j] * (signs & kmask_iq2xs[j] ? -1.f : 1.f);
2091                }
2092            }
2093            sumf[row] += d1 * sum1 + d2 * sum2;
2094
2095            dh += args.nb01/2;
2096            q2 += args.nb01/2;
2097            sc += args.nb01;
2098        }
2099
2100        y4 += 32 * ntx;
2101    }
2102
2103    device float * dst_f32 = (device float *) dst + (uint64_t)im*args.ne0*args.ne1 + (uint64_t)r1*args.ne0;
2104
2105    for (int row = 0; row < nr0 && first_row + row < args.ne0; ++row) {
2106        float sum_all = simd_sum(sumf[row]);
2107        if (tiisg == 0) {
2108            dst_f32[first_row + row] = sum_all * 0.25f;
2109        }
2110    }
2111}
2112
2113template<typename args_t>
2114void kernel_mul_mv_iq2_xs_f32_disp(
2115        args_t args,
2116        device const char * src0,
2117        device const char * src1,
2118        device       char * dst,
2119        threadgroup  char * shmem,
2120        uint3  tgpig,
2121        ushort tiisg,
2122        ushort sgitg) {
2123    if (FC_mul_mv_split) {
2124        kernel_mul_mv_iq2_xs_f32_impl<N_R0_IQ2_XS_SPLIT, args_t>(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg);
2125    } else {
2126        kernel_mul_mv_iq2_xs_f32_impl<N_R0_IQ2_XS, args_t>(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg);
2127    }
2128}
2129
2130[[host_name("kernel_mul_mv_iq2_xs_f32")]]
2131kernel void kernel_mul_mv_iq2_xs_f32(
2132        constant ggml_metal_kargs_mul_mv & args,
2133        device const char * src0,
2134        device const char * src1,
2135        device       char * dst,
2136        threadgroup  char * shmem [[threadgroup(0)]],
2137        uint3  tgpig[[threadgroup_position_in_grid]],
2138        ushort tiisg[[thread_index_in_simdgroup]],
2139        ushort sgitg[[simdgroup_index_in_threadgroup]]) {
2140
2141    kernel_mul_mv_iq2_xs_f32_disp<constant ggml_metal_kargs_mul_mv &>(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg);
2142}
2143
2144// FC_mul_mv_split: for nb32 < 32 (nb32 divides 32), 32/nb32 threads share each chunk and each takes a slice of the rows
2145template<int nr0, typename args_t>
2146void kernel_mul_mv_iq3_xxs_f32_impl(
2147        args_t args,
2148        device const char * src0,
2149        device const char * src1,
2150        device       char * dst,
2151        threadgroup  char * shmem,
2152        uint3  tgpig,
2153        ushort tiisg,
2154        ushort sgitg) {
2155    const short NSG = FC_mul_mv_nsg;
2156
2157    const int nb = args.ne00/QK_K;
2158
2159    const int r0 = tgpig.x;
2160    const int r1 = tgpig.y;
2161    const int im = tgpig.z;
2162
2163    const int first_row = (r0 * NSG + sgitg) * nr0;
2164
2165    const uint i12 = im%FC_mul_mv_ne12;
2166    const uint i13 = im/FC_mul_mv_ne12;
2167
2168    const int nb32 = nb * (QK_K / 32);
2169
2170    const short ntx  = FC_mul_mv_split ? nb32 : 32;
2171    const short nrep = 32 / ntx;
2172
2173    const short ix   = tiisg % ntx;
2174    const short irep = tiisg / ntx;
2175
2176    const short row0 = (nr0 * irep      ) / nrep;
2177    const short row1 = (nr0 * (irep + 1)) / nrep;
2178
2179    const uint64_t offset0 = (first_row + row0)*args.nb01 + (i12/FC_mul_mv_r2)*args.nb02 + (i13/FC_mul_mv_r3)*args.nb03;
2180    const uint64_t offset1 =                r1*args.nb11 + (i12        )*args.nb12 + (i13        )*args.nb13;
2181
2182    device const block_iq3_xxs * x = (device const block_iq3_xxs *) (src0 + offset0);
2183    device const float         * y = (device const float         *) (src1 + offset1);
2184
2185    float yl[32];
2186    float sumf[nr0]={0.f};
2187
2188    threadgroup uint32_t * svalues = (threadgroup uint32_t *)(shmem);
2189    threadgroup uint8_t  * ssigns  = (threadgroup uint8_t  *)(svalues + 256);
2190    {
2191        int nval = 4;
2192        int pos  = (32*sgitg + tiisg)*nval;
2193        for (int i = 0; i < nval; ++i) svalues[pos + i] = iq3xxs_grid[pos + i];
2194        nval = 2;
2195        pos  = (32*sgitg + tiisg)*nval;
2196        for (int i = 0; i < nval; ++i) ssigns[pos+i] = ksigns_iq2xs[pos+i];
2197        threadgroup_barrier(mem_flags::mem_threadgroup);
2198    }
2199
2200    device const float * y4 = y + 32 * ix;
2201
2202    for (int ib32 = ix; ib32 < nb32; ib32 += ntx) {
2203        for (short i = 0; i < 32; ++i) {
2204            yl[i] = y4[i];
2205        }
2206
2207        const int ibl = ib32 / (QK_K / 32);
2208        const int ib  = ib32 % (QK_K / 32);
2209
2210        device const block_iq3_xxs * xr = x + ibl;
2211        device const uint8_t  * q3 = xr->qs + 8 * ib;
2212        device const uint16_t * gas = (device const uint16_t *)(xr->qs + QK_K/4) + 2 * ib;
2213        device const half * dh = &xr->d;
2214
2215        for (short row = row0; row < row1; row++) {
2216            const float db = dh[0];
2217            const uint32_t aux32 = gas[0] | (gas[1] << 16);
2218            const float d = db * (0.5f + (aux32 >> 28));
2219
2220            float2 sum = {0};
2221            for (short l = 0; l < 4; ++l) {
2222                const threadgroup uint8_t * grid1 = (const threadgroup uint8_t *)(svalues + q3[2*l+0]);
2223                const threadgroup uint8_t * grid2 = (const threadgroup uint8_t *)(svalues + q3[2*l+1]);
2224                const uint8_t signs = ssigns[(aux32 >> 7*l) & 127];
2225                for (short j = 0; j < 4; ++j) {
2226                    sum[0] += yl[8*l + j + 0] * grid1[j] * (signs & kmask_iq2xs[j+0] ? -1.f : 1.f);
2227                    sum[1] += yl[8*l + j + 4] * grid2[j] * (signs & kmask_iq2xs[j+4] ? -1.f : 1.f);
2228                }
2229            }
2230            sumf[row] += d * (sum[0] + sum[1]);
2231
2232            dh  += args.nb01/2;
2233            q3  += args.nb01;
2234            gas += args.nb01/2;
2235        }
2236
2237        y4 += 32 * ntx;
2238    }
2239
2240    device float * dst_f32 = (device float *) dst + (uint64_t)im*args.ne0*args.ne1 + (uint64_t)r1*args.ne0;
2241
2242    for (int row = 0; row < nr0 && first_row + row < args.ne0; ++row) {
2243        float sum_all = simd_sum(sumf[row]);
2244        if (tiisg == 0) {
2245            dst_f32[first_row + row] = sum_all * 0.5f;
2246        }
2247    }
2248}
2249
2250template<typename args_t>
2251void kernel_mul_mv_iq3_xxs_f32_disp(
2252        args_t args,
2253        device const char * src0,
2254        device const char * src1,
2255        device       char * dst,
2256        threadgroup  char * shmem,
2257        uint3  tgpig,
2258        ushort tiisg,
2259        ushort sgitg) {
2260    if (FC_mul_mv_split) {
2261        kernel_mul_mv_iq3_xxs_f32_impl<N_R0_IQ3_XXS_SPLIT, args_t>(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg);
2262    } else {
2263        kernel_mul_mv_iq3_xxs_f32_impl<N_R0_IQ3_XXS, args_t>(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg);
2264    }
2265}
2266
2267[[host_name("kernel_mul_mv_iq3_xxs_f32")]]
2268kernel void kernel_mul_mv_iq3_xxs_f32(
2269        constant ggml_metal_kargs_mul_mv & args,
2270        device const char * src0,
2271        device const char * src1,
2272        device       char * dst,
2273        threadgroup  char * shmem [[threadgroup(0)]],
2274        uint3  tgpig[[threadgroup_position_in_grid]],
2275        ushort tiisg[[thread_index_in_simdgroup]],
2276        ushort sgitg[[simdgroup_index_in_threadgroup]]) {
2277
2278    kernel_mul_mv_iq3_xxs_f32_disp<constant ggml_metal_kargs_mul_mv &>(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg);
2279}
2280
2281template<int nr0, typename args_t>
2282void kernel_mul_mv_iq3_s_f32_impl(
2283        args_t args,
2284        device const char * src0,
2285        device const char * src1,
2286        device       char * dst,
2287        threadgroup  char * shmem,
2288        uint3  tgpig,
2289        ushort tiisg,
2290        ushort sgitg) {
2291    const short NSG = FC_mul_mv_nsg;
2292
2293    const int nb = args.ne00/QK_K;
2294
2295    const int r0 = tgpig.x;
2296    const int r1 = tgpig.y;
2297    const int im = tgpig.z;
2298
2299    const int first_row = (r0 * NSG + sgitg) * nr0;
2300
2301    const uint i12 = im%FC_mul_mv_ne12;
2302    const uint i13 = im/FC_mul_mv_ne12;
2303
2304    const int nb32 = nb * (QK_K / 32);
2305
2306    const short ntx  = FC_mul_mv_split ? nb32 : 32;
2307    const short nrep = 32 / ntx;
2308
2309    const short ix   = tiisg % ntx;
2310    const short irep = tiisg / ntx;
2311
2312    const short row0 = (nr0 * irep      ) / nrep;
2313    const short row1 = (nr0 * (irep + 1)) / nrep;
2314
2315    const uint64_t offset0 = (first_row + row0)*args.nb01 + (i12/FC_mul_mv_r2)*args.nb02 + (i13/FC_mul_mv_r3)*args.nb03;
2316    const uint64_t offset1 =                r1*args.nb11 + (i12        )*args.nb12 + (i13        )*args.nb13;
2317
2318    device const block_iq3_s * x = (device const block_iq3_s *) (src0 + offset0);
2319    device const float       * y = (device const float       *) (src1 + offset1);
2320
2321    float yl[32];
2322    float sumf[nr0]={0.f};
2323
2324    threadgroup uint32_t * svalues = (threadgroup uint32_t *) shmem;
2325    {
2326        int nval = 8;
2327        int pos  = (32*sgitg + tiisg)*nval;
2328        for (int i = 0; i < nval; ++i) svalues[pos + i] = iq3s_grid[pos + i];
2329        threadgroup_barrier(mem_flags::mem_threadgroup);
2330    }
2331
2332    device const float * y4 = y + 32 * ix;
2333
2334    for (int ib32 = ix; ib32 < nb32; ib32 += ntx) {
2335        for (short i = 0; i < 32; ++i) {
2336            yl[i] = y4[i];
2337        }
2338
2339        const int ibl = ib32 / (QK_K / 32);
2340        const int ib  = ib32 % (QK_K / 32);
2341
2342        device const block_iq3_s * xr = x + ibl;
2343        device const uint8_t * qs = xr->qs + 8 * ib;
2344        device const uint8_t * qh = xr->qh + ib;
2345        device const uint8_t * sc = xr->scales + (ib/2);
2346        device const uint8_t * signs = xr->signs + 4 * ib;
2347        device const half * dh = &xr->d;
2348
2349        for (short row = row0; row < row1; row++) {
2350            const float db = dh[0];
2351            const float d = db * (1 + 2*((sc[0] >> 4*(ib%2)) & 0xf));
2352
2353            float2 sum = {0};
2354            for (short l = 0; l < 4; ++l) {
2355                const threadgroup uint32_t * table1 = qh[0] & kmask_iq2xs[2*l+0] ? svalues + 256 : svalues;
2356                const threadgroup uint32_t * table2 = qh[0] & kmask_iq2xs[2*l+1] ? svalues + 256 : svalues;
2357                const threadgroup uint8_t * grid1 = (const threadgroup uint8_t *)(table1 + qs[2*l+0]);
2358                const threadgroup uint8_t * grid2 = (const threadgroup uint8_t *)(table2 + qs[2*l+1]);
2359                for (short j = 0; j < 4; ++j) {
2360                    sum[0] += yl[8*l + j + 0] * grid1[j] * select(1, -1, signs[l] & kmask_iq2xs[j+0]);
2361                    sum[1] += yl[8*l + j + 4] * grid2[j] * select(1, -1, signs[l] & kmask_iq2xs[j+4]);
2362                }
2363            }
2364            sumf[row] += d * (sum[0] + sum[1]);
2365
2366            dh    += args.nb01/2;
2367            qs    += args.nb01;
2368            qh    += args.nb01;
2369            sc    += args.nb01;
2370            signs += args.nb01;
2371        }
2372
2373        y4 += 32 * ntx;
2374    }
2375
2376    device float * dst_f32 = (device float *) dst + (uint64_t)im*args.ne0*args.ne1 + (uint64_t)r1*args.ne0;
2377
2378    for (int row = 0; row < nr0 && first_row + row < args.ne0; ++row) {
2379        float sum_all = simd_sum(sumf[row]);
2380        if (tiisg == 0) {
2381            dst_f32[first_row + row] = sum_all;
2382        }
2383    }
2384}
2385
2386template<typename args_t>
2387void kernel_mul_mv_iq3_s_f32_disp(
2388        args_t args,
2389        device const char * src0,
2390        device const char * src1,
2391        device       char * dst,
2392        threadgroup  char * shmem,
2393        uint3  tgpig,
2394        ushort tiisg,
2395        ushort sgitg) {
2396    if (FC_mul_mv_split) {
2397        kernel_mul_mv_iq3_s_f32_impl<N_R0_IQ3_S_SPLIT, args_t>(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg);
2398    } else {
2399        kernel_mul_mv_iq3_s_f32_impl<N_R0_IQ3_S, args_t>(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg);
2400    }
2401}
2402
2403[[host_name("kernel_mul_mv_iq3_s_f32")]]
2404kernel void kernel_mul_mv_iq3_s_f32(
2405        constant ggml_metal_kargs_mul_mv & args,
2406        device const char * src0,
2407        device const char * src1,
2408        device       char * dst,
2409        threadgroup  char * shmem [[threadgroup(0)]],
2410        uint3  tgpig[[threadgroup_position_in_grid]],
2411        ushort tiisg[[thread_index_in_simdgroup]],
2412        ushort sgitg[[simdgroup_index_in_threadgroup]]) {
2413
2414    kernel_mul_mv_iq3_s_f32_disp<constant ggml_metal_kargs_mul_mv &>(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg);
2415}
2416
2417template<int nr0, typename args_t>
2418void kernel_mul_mv_iq2_s_f32_impl(
2419        args_t args,
2420        device const char * src0,
2421        device const char * src1,
2422        device       char * dst,
2423        threadgroup  char * shmem,
2424        uint3  tgpig,
2425        ushort tiisg,
2426        ushort sgitg) {
2427    const short NSG = FC_mul_mv_nsg;
2428
2429    const int nb = args.ne00/QK_K;
2430
2431    const int r0 = tgpig.x;
2432    const int r1 = tgpig.y;
2433    const int im = tgpig.z;
2434
2435    const int first_row = (r0 * NSG + sgitg) * nr0;
2436
2437    const uint i12 = im%FC_mul_mv_ne12;
2438    const uint i13 = im/FC_mul_mv_ne12;
2439
2440    const int nb32 = nb * (QK_K / 32);
2441
2442    const short ntx  = FC_mul_mv_split ? nb32 : 32;
2443    const short nrep = 32 / ntx;
2444
2445    const short ix   = tiisg % ntx;
2446    const short irep = tiisg / ntx;
2447
2448    const short row0 = (nr0 * irep      ) / nrep;
2449    const short row1 = (nr0 * (irep + 1)) / nrep;
2450
2451    const uint64_t offset0 = (first_row + row0)*args.nb01 + (i12/FC_mul_mv_r2)*args.nb02 + (i13/FC_mul_mv_r3)*args.nb03;
2452    const uint64_t offset1 =                r1*args.nb11 + (i12        )*args.nb12 + (i13        )*args.nb13;
2453
2454    device const block_iq2_s * x = (device const block_iq2_s *) (src0 + offset0);
2455    device const float       * y = (device const float       *) (src1 + offset1);
2456
2457    float yl[32];
2458    float sumf[nr0]={0.f};
2459
2460    //threadgroup uint64_t * svalues = (threadgroup uint64_t *) shmem;
2461    //{
2462    //    int nval = 32;
2463    //    int pos  = (32*sgitg + tiisg)*nval;
2464    //    for (int i = 0; i < nval; ++i) svalues[pos + i] = iq2s_grid[pos + i];
2465    //    threadgroup_barrier(mem_flags::mem_threadgroup);
2466    //}
2467
2468    device const float * y4 = y + 32 * ix;
2469
2470    for (int ib32 = ix; ib32 < nb32; ib32 += ntx) {
2471        for (short i = 0; i < 32; ++i) {
2472            yl[i] = y4[i];
2473        }
2474
2475        const int ibl = ib32 / (QK_K / 32);
2476        const int ib  = ib32 % (QK_K / 32);
2477
2478        device const block_iq2_s * xr = x + ibl;
2479        device const uint8_t * qs = xr->qs + 4 * ib;
2480        device const uint8_t * qh = xr->qh + ib;
2481        device const uint8_t * sc = xr->scales + ib;
2482        device const uint8_t * signs = qs + QK_K/8;
2483        device const half * dh = &xr->d;
2484
2485        for (short row = row0; row < row1; row++) {
2486            const float db = dh[0];
2487            const float d1 = db * (0.5f + (sc[0] & 0xf));
2488            const float d2 = db * (0.5f + (sc[0] >>  4));
2489
2490            float2 sum = {0};
2491            for (short l = 0; l < 2; ++l) {
2492                //const threadgroup uint8_t * grid1 = (const threadgroup uint8_t *)(svalues + (qs[l+0] | ((qh[0] << (8-2*l)) & 0x300)));
2493                //const threadgroup uint8_t * grid2 = (const threadgroup uint8_t *)(svalues + (qs[l+2] | ((qh[0] << (4-2*l)) & 0x300)));
2494                constant uint8_t * grid1 = (constant uint8_t *)(iq2s_grid + (qs[l+0] | ((qh[0] << (8-2*l)) & 0x300)));
2495                constant uint8_t * grid2 = (constant uint8_t *)(iq2s_grid + (qs[l+2] | ((qh[0] << (4-2*l)) & 0x300)));
2496                for (short j = 0; j < 8; ++j) {
2497                    sum[0] += yl[8*l + j +  0] * grid1[j] * select(1, -1, signs[l+0] & kmask_iq2xs[j]);
2498                    sum[1] += yl[8*l + j + 16] * grid2[j] * select(1, -1, signs[l+2] & kmask_iq2xs[j]);
2499                }
2500            }
2501            sumf[row] += d1 * sum[0] + d2 * sum[1];
2502
2503            dh    += args.nb01/2;
2504            qs    += args.nb01;
2505            qh    += args.nb01;
2506            sc    += args.nb01;
2507            signs += args.nb01;
2508        }
2509
2510        y4 += 32 * ntx;
2511    }
2512
2513    device float * dst_f32 = (device float *) dst + (uint64_t)im*args.ne0*args.ne1 + (uint64_t)r1*args.ne0;
2514
2515    for (int row = 0; row < nr0 && first_row + row < args.ne0; ++row) {
2516        float sum_all = simd_sum(sumf[row]);
2517        if (tiisg == 0) {
2518            dst_f32[first_row + row] = sum_all * 0.25f;
2519        }
2520    }
2521}
2522
2523template<typename args_t>
2524void kernel_mul_mv_iq2_s_f32_disp(
2525        args_t args,
2526        device const char * src0,
2527        device const char * src1,
2528        device       char * dst,
2529        threadgroup  char * shmem,
2530        uint3  tgpig,
2531        ushort tiisg,
2532        ushort sgitg) {
2533    if (FC_mul_mv_split) {
2534        kernel_mul_mv_iq2_s_f32_impl<N_R0_IQ2_S_SPLIT, args_t>(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg);
2535    } else {
2536        kernel_mul_mv_iq2_s_f32_impl<N_R0_IQ2_S, args_t>(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg);
2537    }
2538}
2539
2540[[host_name("kernel_mul_mv_iq2_s_f32")]]
2541kernel void kernel_mul_mv_iq2_s_f32(
2542        constant ggml_metal_kargs_mul_mv & args,
2543        device const char * src0,
2544        device const char * src1,
2545        device       char * dst,
2546        threadgroup  char * shmem [[threadgroup(0)]],
2547        uint3  tgpig[[threadgroup_position_in_grid]],
2548        ushort tiisg[[thread_index_in_simdgroup]],
2549        ushort sgitg[[simdgroup_index_in_threadgroup]]) {
2550
2551    kernel_mul_mv_iq2_s_f32_disp<constant ggml_metal_kargs_mul_mv &>(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg);
2552}
2553
2554template<int nr0, typename args_t>
2555void kernel_mul_mv_iq1_s_f32_impl(
2556        args_t args,
2557        device const char * src0,
2558        device const char * src1,
2559        device       char * dst,
2560        threadgroup  char * shmem,
2561        uint3  tgpig,
2562        ushort tiisg,
2563        ushort sgitg) {
2564    const short NSG = FC_mul_mv_nsg;
2565
2566    const int nb = args.ne00/QK_K;
2567
2568    const int r0 = tgpig.x;
2569    const int r1 = tgpig.y;
2570    const int im = tgpig.z;
2571
2572    const int first_row = (r0 * NSG + sgitg) * nr0;
2573
2574    const uint i12 = im%FC_mul_mv_ne12;
2575    const uint i13 = im/FC_mul_mv_ne12;
2576
2577    const int nb32 = nb * (QK_K / 32);
2578
2579    const short ntx  = FC_mul_mv_split ? nb32 : 32;
2580    const short nrep = 32 / ntx;
2581
2582    const short ix   = tiisg % ntx;
2583    const short irep = tiisg / ntx;
2584
2585    const short row0 = (nr0 * irep      ) / nrep;
2586    const short row1 = (nr0 * (irep + 1)) / nrep;
2587
2588    const uint64_t offset0 = (first_row + row0)*args.nb01 + (i12/FC_mul_mv_r2)*args.nb02 + (i13/FC_mul_mv_r3)*args.nb03;
2589    const uint64_t offset1 =                r1*args.nb11 + (i12        )*args.nb12 + (i13        )*args.nb13;
2590
2591    device const block_iq1_s * x = (device const block_iq1_s *) (src0 + offset0);
2592    device const float       * y = (device const float       *) (src1 + offset1);
2593
2594    float yl[32];
2595    float sumf[nr0]={0.f};
2596
2597    device const float * y4 = y + 32 * ix;
2598
2599    for (int ib32 = ix; ib32 < nb32; ib32 += ntx) {
2600        float sumy = 0;
2601        for (short i = 0; i < 32; ++i) {
2602            yl[i] = y4[i];
2603            sumy += yl[i];
2604        }
2605
2606        const int ibl = ib32 / (QK_K / 32);
2607        const int ib  = ib32 % (QK_K / 32);
2608
2609        device const block_iq1_s * xr = x + ibl;
2610        device const uint8_t  * qs = xr->qs + 4 * ib;
2611        device const uint16_t * qh = xr->qh + ib;
2612        device const half     * dh = &xr->d;
2613
2614        for (short row = row0; row < row1; row++) {
2615            constant uint8_t * grid1 = (constant uint8_t *)(iq1s_grid_gpu + (qs[0] | ((qh[0] << 8) & 0x700)));
2616            constant uint8_t * grid2 = (constant uint8_t *)(iq1s_grid_gpu + (qs[1] | ((qh[0] << 5) & 0x700)));
2617            constant uint8_t * grid3 = (constant uint8_t *)(iq1s_grid_gpu + (qs[2] | ((qh[0] << 2) & 0x700)));
2618            constant uint8_t * grid4 = (constant uint8_t *)(iq1s_grid_gpu + (qs[3] | ((qh[0] >> 1) & 0x700)));
2619
2620            float sum = 0;
2621            for (short j = 0; j < 4; ++j) {
2622                sum += yl[j+ 0] * (grid1[j] & 0xf) + yl[j+ 4] * (grid1[j] >> 4)
2623                     + yl[j+ 8] * (grid2[j] & 0xf) + yl[j+12] * (grid2[j] >> 4)
2624                     + yl[j+16] * (grid3[j] & 0xf) + yl[j+20] * (grid3[j] >> 4)
2625                     + yl[j+24] * (grid4[j] & 0xf) + yl[j+28] * (grid4[j] >> 4);
2626            }
2627            sumf[row] += (float)dh[0] * (sum + sumy * (qh[0] & 0x8000 ? -1 - IQ1S_DELTA : -1 + IQ1S_DELTA)) * (2*((qh[0] >> 12) & 7) + 1);
2628
2629            dh += args.nb01/2;
2630            qs += args.nb01;
2631            qh += args.nb01/2;
2632        }
2633
2634        y4 += 32 * ntx;
2635    }
2636
2637    device float * dst_f32 = (device float *) dst + (uint64_t)im*args.ne0*args.ne1 + (uint64_t)r1*args.ne0;
2638
2639    for (int row = 0; row < nr0 && first_row + row < args.ne0; ++row) {
2640        float sum_all = simd_sum(sumf[row]);
2641        if (tiisg == 0) {
2642            dst_f32[first_row + row] = sum_all;
2643        }
2644    }
2645}
2646
2647template<typename args_t>
2648void kernel_mul_mv_iq1_s_f32_disp(
2649        args_t args,
2650        device const char * src0,
2651        device const char * src1,
2652        device       char * dst,
2653        threadgroup  char * shmem,
2654        uint3  tgpig,
2655        ushort tiisg,
2656        ushort sgitg) {
2657    if (FC_mul_mv_split) {
2658        kernel_mul_mv_iq1_s_f32_impl<N_R0_IQ1_S_SPLIT, args_t>(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg);
2659    } else {
2660        kernel_mul_mv_iq1_s_f32_impl<N_R0_IQ1_S, args_t>(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg);
2661    }
2662}
2663
2664[[host_name("kernel_mul_mv_iq1_s_f32")]]
2665kernel void kernel_mul_mv_iq1_s_f32(
2666        constant ggml_metal_kargs_mul_mv & args,
2667        device const char * src0,
2668        device const char * src1,
2669        device       char * dst,
2670        uint3  tgpig[[threadgroup_position_in_grid]],
2671        ushort tiisg[[thread_index_in_simdgroup]],
2672        ushort sgitg[[simdgroup_index_in_threadgroup]]) {
2673
2674    kernel_mul_mv_iq1_s_f32_disp<constant ggml_metal_kargs_mul_mv &>(args, src0, src1, dst, nullptr, tgpig, tiisg, sgitg);
2675}
2676
2677template<int nr0, typename args_t>
2678void kernel_mul_mv_iq1_m_f32_impl(
2679        args_t args,
2680        device const char * src0,
2681        device const char * src1,
2682        device       char * dst,
2683        threadgroup  char * shmem,
2684        uint3  tgpig,
2685        ushort tiisg,
2686        ushort sgitg) {
2687    const short NSG = FC_mul_mv_nsg;
2688
2689    const int nb = args.ne00/QK_K;
2690
2691    const int r0 = tgpig.x;
2692    const int r1 = tgpig.y;
2693    const int im = tgpig.z;
2694
2695    const int first_row = (r0 * NSG + sgitg) * nr0;
2696
2697    const uint i12 = im%FC_mul_mv_ne12;
2698    const uint i13 = im/FC_mul_mv_ne12;
2699
2700    const int nb32 = nb * (QK_K / 32);
2701
2702    const short ntx  = FC_mul_mv_split ? nb32 : 32;
2703    const short nrep = 32 / ntx;
2704
2705    const short ix   = tiisg % ntx;
2706    const short irep = tiisg / ntx;
2707
2708    const short row0 = (nr0 * irep      ) / nrep;
2709    const short row1 = (nr0 * (irep + 1)) / nrep;
2710
2711    const uint64_t offset0 = (first_row + row0)*args.nb01 + (i12/FC_mul_mv_r2)*args.nb02 + (i13/FC_mul_mv_r3)*args.nb03;
2712    const uint64_t offset1 =                r1*args.nb11 + (i12        )*args.nb12 + (i13        )*args.nb13;
2713
2714    device const block_iq1_m * x = (device const block_iq1_m *) (src0 + offset0);
2715    device const float       * y = (device const float       *) (src1 + offset1);
2716
2717    float yl[32];
2718    float sumf[nr0]={0.f};
2719
2720    device const float * y4 = y + 32 * ix;
2721
2722    iq1m_scale_t scale;
2723
2724    for (int ib32 = ix; ib32 < nb32; ib32 += ntx) {
2725        float4 sumy = {0.f};
2726        for (short i = 0; i < 8; ++i) {
2727            yl[i+ 0] = y4[i+ 0]; sumy[0] += yl[i+ 0];
2728            yl[i+ 8] = y4[i+ 8]; sumy[1] += yl[i+ 8];
2729            yl[i+16] = y4[i+16]; sumy[2] += yl[i+16];
2730            yl[i+24] = y4[i+24]; sumy[3] += yl[i+24];
2731        }
2732
2733        const int ibl = ib32 / (QK_K / 32);
2734        const int ib  = ib32 % (QK_K / 32);
2735
2736        device const block_iq1_m * xr = x + ibl;
2737        device const uint8_t  * qs = xr->qs + 4 * ib;
2738        device const uint8_t  * qh = xr->qh + 2 * ib;
2739        device const uint16_t * sc = (device const uint16_t *)xr->scales;
2740
2741        for (short row = row0; row < row1; row++) {
2742            scale.u16 = (sc[0] >> 12) | ((sc[1] >> 8) & 0x00f0) | ((sc[2] >> 4) & 0x0f00) | (sc[3] & 0xf000);
2743
2744            constant uint8_t * grid1 = (constant uint8_t *)(iq1s_grid_gpu + (qs[0] | ((qh[0] << 8) & 0x700)));
2745            constant uint8_t * grid2 = (constant uint8_t *)(iq1s_grid_gpu + (qs[1] | ((qh[0] << 4) & 0x700)));
2746            constant uint8_t * grid3 = (constant uint8_t *)(iq1s_grid_gpu + (qs[2] | ((qh[1] << 8) & 0x700)));
2747            constant uint8_t * grid4 = (constant uint8_t *)(iq1s_grid_gpu + (qs[3] | ((qh[1] << 4) & 0x700)));
2748
2749            float2 sum = {0.f};
2750            for (short j = 0; j < 4; ++j) {
2751                sum[0] += yl[j+ 0] * (grid1[j] & 0xf) + yl[j+ 4] * (grid1[j] >> 4)
2752                        + yl[j+ 8] * (grid2[j] & 0xf) + yl[j+12] * (grid2[j] >> 4);
2753                sum[1] += yl[j+16] * (grid3[j] & 0xf) + yl[j+20] * (grid3[j] >> 4)
2754                        + yl[j+24] * (grid4[j] & 0xf) + yl[j+28] * (grid4[j] >> 4);
2755            }
2756            const float delta1 = sumy[0] * (qh[0] & 0x08 ? -1 - IQ1M_DELTA : -1 + IQ1M_DELTA) + sumy[1] * (qh[0] & 0x80 ? -1 - IQ1M_DELTA : -1 + IQ1M_DELTA);
2757            const float delta2 = sumy[2] * (qh[1] & 0x08 ? -1 - IQ1M_DELTA : -1 + IQ1M_DELTA) + sumy[3] * (qh[1] & 0x80 ? -1 - IQ1M_DELTA : -1 + IQ1M_DELTA);
2758
2759            sumf[row] += (float)scale.f16 * ((sum[0] + delta1) * (2*((sc[ib/2] >> (6*(ib%2)+0)) & 7) + 1) +
2760                                             (sum[1] + delta2) * (2*((sc[ib/2] >> (6*(ib%2)+3)) & 7) + 1));
2761
2762            sc += args.nb01/2;
2763            qs += args.nb01;
2764            qh += args.nb01;
2765        }
2766
2767        y4 += 32 * ntx;
2768    }
2769
2770    device float * dst_f32 = (device float *) dst + (uint64_t)im*args.ne0*args.ne1 + (uint64_t)r1*args.ne0;
2771
2772    for (int row = 0; row < nr0 && first_row + row < args.ne0; ++row) {
2773        float sum_all = simd_sum(sumf[row]);
2774        if (tiisg == 0) {
2775            dst_f32[first_row + row] = sum_all;
2776        }
2777    }
2778}
2779
2780template<typename args_t>
2781void kernel_mul_mv_iq1_m_f32_disp(
2782        args_t args,
2783        device const char * src0,
2784        device const char * src1,
2785        device       char * dst,
2786        threadgroup  char * shmem,
2787        uint3  tgpig,
2788        ushort tiisg,
2789        ushort sgitg) {
2790    if (FC_mul_mv_split) {
2791        kernel_mul_mv_iq1_m_f32_impl<N_R0_IQ1_M_SPLIT, args_t>(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg);
2792    } else {
2793        kernel_mul_mv_iq1_m_f32_impl<N_R0_IQ1_M, args_t>(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg);
2794    }
2795}
2796
2797[[host_name("kernel_mul_mv_iq1_m_f32")]]
2798kernel void kernel_mul_mv_iq1_m_f32(
2799        constant ggml_metal_kargs_mul_mv & args,
2800        device const char * src0,
2801        device const char * src1,
2802        device       char * dst,
2803        uint3  tgpig[[threadgroup_position_in_grid]],
2804        ushort tiisg[[thread_index_in_simdgroup]],
2805        ushort sgitg[[simdgroup_index_in_threadgroup]]) {
2806
2807    kernel_mul_mv_iq1_m_f32_disp<constant ggml_metal_kargs_mul_mv &>(args, src0, src1, dst, nullptr, tgpig, tiisg, sgitg);
2808}
2809
2810template<int NR0, typename args_t>
2811void kernel_mul_mv_iq4_nl_f32_impl(
2812        args_t args,
2813        device const char * src0,
2814        device const char * src1,
2815        device       char * dst,
2816        threadgroup  char * shmem,
2817        uint3  tgpig,
2818        ushort tiisg,
2819        ushort sgitg) {
2820    const short NSG = FC_mul_mv_nsg;
2821
2822    threadgroup float * shmem_f32 = (threadgroup float *) shmem;
2823
2824    const int r0 = tgpig.x;
2825    const int r1 = tgpig.y;
2826    const int im = tgpig.z;
2827
2828    const int first_row = (r0 * NSG + sgitg) * NR0;
2829
2830    const uint i12 = im%FC_mul_mv_ne12;
2831    const uint i13 = im/FC_mul_mv_ne12;
2832
2833    const uint64_t offset0 = first_row*args.nb01 + (i12/FC_mul_mv_r2)*args.nb02 + (i13/FC_mul_mv_r3)*args.nb03;
2834    const uint64_t offset1 =        r1*args.nb11 + (i12        )*args.nb12 + (i13        )*args.nb13;
2835
2836    device const block_iq4_nl * x = (device const block_iq4_nl *) (src0 + offset0);
2837    device const float        * y = (device const float        *) (src1 + offset1);
2838
2839    const int nb   = args.ne00/QK4_NL;
2840    const int ns01 = args.nb01/args.nb00;
2841
2842    const short ix = tiisg/2;  // 0...15
2843    const short it = tiisg%2;  // 0 or 1
2844
2845    shmem_f32[tiisg] = kvalues_iq4nl_f[tiisg%16];
2846    threadgroup_barrier(mem_flags::mem_threadgroup);
2847
2848    float4 yl[4];
2849    float sumf[NR0]={0.f};
2850
2851    device const float * yb = y + ix*QK4_NL + it*8;
2852
2853    uint32_t aux32[2];
2854    thread const uint8_t * q8 = (thread const uint8_t *)aux32;
2855
2856    float4 qf1, qf2;
2857
2858    // [TAG_MUL_MV_WEIRD]
2859    for (int ib = ix; ib < nb && ib < ns01; ib += 16) {
2860        device const float4 * y4 = (device const float4 *)yb;
2861        yl[0] = y4[0];
2862        yl[1] = y4[4];
2863        yl[2] = y4[1];
2864        yl[3] = y4[5];
2865
2866        for (short row = 0; row < NR0; row++) {
2867            device const block_iq4_nl & xb = x[row*ns01 + ib];
2868            device const uint16_t * q4 = (device const uint16_t *)(xb.qs + 8*it);
2869
2870            float4 acc1 = {0.f}, acc2 = {0.f};
2871
2872            aux32[0] = q4[0] | (q4[1] << 16);
2873            aux32[1] = (aux32[0] >> 4) & 0x0f0f0f0f;
2874            aux32[0] &= 0x0f0f0f0f;
2875            qf1 = {shmem_f32[q8[0]], shmem_f32[q8[1]], shmem_f32[q8[2]], shmem_f32[q8[3]]};
2876            qf2 = {shmem_f32[q8[4]], shmem_f32[q8[5]], shmem_f32[q8[6]], shmem_f32[q8[7]]};
2877            acc1 += yl[0] * qf1;
2878            acc2 += yl[1] * qf2;
2879
2880            aux32[0] = q4[2] | (q4[3] << 16);
2881            aux32[1] = (aux32[0] >> 4) & 0x0f0f0f0f;
2882            aux32[0] &= 0x0f0f0f0f;
2883            qf1 = {shmem_f32[q8[0]], shmem_f32[q8[1]], shmem_f32[q8[2]], shmem_f32[q8[3]]};
2884            qf2 = {shmem_f32[q8[4]], shmem_f32[q8[5]], shmem_f32[q8[6]], shmem_f32[q8[7]]};
2885            acc1 += yl[2] * qf1;
2886            acc2 += yl[3] * qf2;
2887
2888            acc1 += acc2;
2889
2890            sumf[row] += (float)xb.d * (acc1[0] + acc1[1] + acc1[2] + acc1[3]);
2891        }
2892
2893        yb += 16 * QK4_NL;
2894    }
2895
2896    device float * dst_f32 = (device float *) dst + (uint64_t)im*args.ne0*args.ne1 + (uint64_t)r1*args.ne0;
2897
2898    for (int row = 0; row < NR0 && first_row + row < args.ne0; ++row) {
2899        float sum_all = simd_sum(sumf[row]);
2900        if (tiisg == 0) {
2901            dst_f32[first_row + row] = sum_all;
2902        }
2903    }
2904}
2905
2906[[host_name("kernel_mul_mv_iq4_nl_f32")]]
2907kernel void kernel_mul_mv_iq4_nl_f32(
2908        constant ggml_metal_kargs_mul_mv & args,
2909        device const char * src0,
2910        device const char * src1,
2911        device       char * dst,
2912        threadgroup  char * shmem [[threadgroup(0)]],
2913        uint3  tgpig[[threadgroup_position_in_grid]],
2914        ushort tiisg[[thread_index_in_simdgroup]],
2915        ushort sgitg[[simdgroup_index_in_threadgroup]]) {
2916
2917    kernel_mul_mv_iq4_nl_f32_impl<N_R0_IQ4_NL, constant ggml_metal_kargs_mul_mv &>(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg);
2918}
2919
2920template<int NR0, typename args_t>
2921void kernel_mul_mv_iq4_xs_f32_impl(
2922        args_t args,
2923        device const char * src0,
2924        device const char * src1,
2925        device       char * dst,
2926        threadgroup  char * shmem,
2927        uint3  tgpig,
2928        ushort tiisg,
2929        ushort sgitg) {
2930    const short NSG = FC_mul_mv_nsg;
2931
2932    threadgroup float * shmem_f32 = (threadgroup float *) shmem;
2933
2934    const int r0 = tgpig.x;
2935    const int r1 = tgpig.y;
2936    const int im = tgpig.z;
2937    const int first_row = (r0 * NSG + sgitg) * NR0;
2938
2939    const uint i12 = im%FC_mul_mv_ne12;
2940    const uint i13 = im/FC_mul_mv_ne12;
2941
2942    const uint64_t offset0 = first_row*args.nb01 + (i12/FC_mul_mv_r2)*args.nb02 + (i13/FC_mul_mv_r3)*args.nb03;
2943    const uint64_t offset1 =        r1*args.nb11 + (i12        )*args.nb12 + (i13        )*args.nb13;
2944
2945    device const block_iq4_xs * x = (device const block_iq4_xs *) (src0 + offset0);
2946    device const float        * y = (device const float        *) (src1 + offset1);
2947
2948    const int nb   = args.ne00/QK_K;
2949    const int ns01 = args.nb01/args.nb00;
2950
2951    const short ix = tiisg/16;  // 0 or 1
2952    const short it = tiisg%16;  // 0...15
2953    const short ib = it/2;
2954    const short il = it%2;
2955
2956    shmem_f32[tiisg] = kvalues_iq4nl_f[tiisg%16];
2957    threadgroup_barrier(mem_flags::mem_threadgroup);
2958
2959    float4 yl[4];
2960    float sumf[NR0]={0.f};
2961
2962    device const float * yb = y + ix * QK_K + ib * 32 + il * 8;
2963
2964    uint32_t aux32[2];
2965    thread const uint8_t * q8 = (thread const uint8_t *)aux32;
2966
2967    float4 qf1, qf2;
2968
2969    // [TAG_MUL_MV_WEIRD]
2970    for (int ibl = ix; ibl < nb && ibl < ns01; ibl += 2) {
2971        device const float4 * y4 = (device const float4 *)yb;
2972        yl[0] = y4[0];
2973        yl[1] = y4[4];
2974        yl[2] = y4[1];
2975        yl[3] = y4[5];
2976
2977        for (short row = 0; row < NR0; ++row) {
2978            device const block_iq4_xs & xb = x[row*ns01 + ibl];
2979            device const uint32_t * q4 = (device const uint32_t *)(xb.qs + 16*ib + 8*il);
2980
2981            float4 acc1 = {0.f}, acc2 = {0.f};
2982
2983            aux32[0] = (q4[0]     ) & 0x0f0f0f0f;
2984            aux32[1] = (q4[0] >> 4) & 0x0f0f0f0f;
2985            qf1 = {shmem_f32[q8[0]], shmem_f32[q8[1]], shmem_f32[q8[2]], shmem_f32[q8[3]]};
2986            qf2 = {shmem_f32[q8[4]], shmem_f32[q8[5]], shmem_f32[q8[6]], shmem_f32[q8[7]]};
2987            acc1 += yl[0] * qf1;
2988            acc2 += yl[1] * qf2;
2989
2990            aux32[0] = (q4[1]     ) & 0x0f0f0f0f;
2991            aux32[1] = (q4[1] >> 4) & 0x0f0f0f0f;
2992            qf1 = {shmem_f32[q8[0]], shmem_f32[q8[1]], shmem_f32[q8[2]], shmem_f32[q8[3]]};
2993            qf2 = {shmem_f32[q8[4]], shmem_f32[q8[5]], shmem_f32[q8[6]], shmem_f32[q8[7]]};
2994            acc1 += yl[2] * qf1;
2995            acc2 += yl[3] * qf2;
2996
2997            acc1 += acc2;
2998
2999            const int ls = (((xb.scales_l[ib/2] >> 4*(ib%2)) & 0xf) | (((xb.scales_h >> 2*ib) & 3) << 4)) - 32;
3000            sumf[row] += (float)xb.d * ls * (acc1[0] + acc1[1] + acc1[2] + acc1[3]);
3001        }
3002
3003        yb += 2 * QK_K;
3004    }
3005
3006    device float * dst_f32 = (device float *) dst + (uint64_t)im*args.ne0*args.ne1 + (uint64_t)r1*args.ne0;
3007
3008    for (int row = 0; row < NR0 && first_row + row < args.ne0; ++row) {
3009        float sum_all = simd_sum(sumf[row]);
3010        if (tiisg == 0) {
3011            dst_f32[first_row + row] = sum_all;
3012        }
3013    }
3014}
3015
3016[[host_name("kernel_mul_mv_iq4_xs_f32")]]
3017kernel void kernel_mul_mv_iq4_xs_f32(
3018        constant ggml_metal_kargs_mul_mv & args,
3019        device const char * src0,
3020        device const char * src1,
3021        device       char * dst,
3022        threadgroup  char * shmem [[threadgroup(0)]],
3023        uint3  tgpig[[threadgroup_position_in_grid]],
3024        ushort tiisg[[thread_index_in_simdgroup]],
3025        ushort sgitg[[simdgroup_index_in_threadgroup]]) {
3026
3027    kernel_mul_mv_iq4_xs_f32_impl<N_R0_IQ4_XS, constant ggml_metal_kargs_mul_mv &>(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg);
3028}
3029
3030template<int NR0, typename args_t>
3031void kernel_mul_mv_mxfp4_f32_impl(
3032        args_t args,
3033        device const char * src0,
3034        device const char * src1,
3035        device       char * dst,
3036        threadgroup  char * shmem,
3037        uint3  tgpig,
3038        ushort tiisg,
3039        ushort sgitg) {
3040    const short NSG = FC_mul_mv_nsg;
3041
3042    threadgroup float * shmem_f32 = (threadgroup float *) shmem;
3043
3044    const int r0 = tgpig.x;
3045    const int r1 = tgpig.y;
3046    const int im = tgpig.z;
3047
3048    const int first_row = (r0 * NSG + sgitg) * NR0;
3049
3050    const uint i12 = im%FC_mul_mv_ne12;
3051    const uint i13 = im/FC_mul_mv_ne12;
3052
3053    const uint64_t offset0 = first_row*args.nb01 + (i12/FC_mul_mv_r2)*args.nb02 + (i13/FC_mul_mv_r3)*args.nb03;
3054    const uint64_t offset1 =        r1*args.nb11 + (i12        )*args.nb12 + (i13        )*args.nb13;
3055
3056    device const block_mxfp4 * x = (device const block_mxfp4 *) (src0 + offset0);
3057    device const float       * y = (device const float       *) (src1 + offset1);
3058
3059    const int nb   = args.ne00/QK_MXFP4;
3060    const int ns01 = args.nb01/args.nb00; // this can be larger than nb for permuted src0 tensors
3061
3062    const short ix = tiisg/2;  // 0...15
3063    const short it = tiisg%2;  // 0 or 1
3064
3065    shmem_f32[tiisg] = kvalues_mxfp4_f[tiisg%16];
3066    threadgroup_barrier(mem_flags::mem_threadgroup);
3067
3068    float4 yl[4];
3069    float sumf[NR0]={0.f};
3070
3071    device const float * yb = y + ix*QK_MXFP4 + it*8;
3072
3073    // note: just the check `ib < nb` is enough, but adding the redundant `&& ib < ns01` check makes the kernel a bit faster
3074    //       no idea why that is - needs some deeper investigation [TAG_MUL_MV_WEIRD]
3075    for (int ib = ix; ib < nb && ib < ns01; ib += 16) {
3076        device const float4 * y4 = (device const float4 *) yb;
3077
3078        yl[0] = y4[0];
3079        yl[1] = y4[4];
3080        yl[2] = y4[1];
3081        yl[3] = y4[5];
3082
3083        FOR_UNROLL (short row = 0; row < NR0; row++) {
3084            device const block_mxfp4 & xb = x[row*ns01 + ib];
3085            device const uint8_t     * q2 = (device const uint8_t *)(xb.qs + 8*it);
3086
3087            float4 acc1 = yl[0]*float4(shmem_f32[q2[0] &  0x0F], shmem_f32[q2[1] &  0x0F], shmem_f32[q2[2] &  0x0F], shmem_f32[q2[3] &  0x0F]);
3088            float4 acc2 = yl[1]*float4(shmem_f32[q2[0] >> 4   ], shmem_f32[q2[1] >> 4   ], shmem_f32[q2[2] >> 4   ], shmem_f32[q2[3] >> 4   ]);
3089            float4 acc3 = yl[2]*float4(shmem_f32[q2[4] &  0x0F], shmem_f32[q2[5] &  0x0F], shmem_f32[q2[6] &  0x0F], shmem_f32[q2[7] &  0x0F]);
3090            float4 acc4 = yl[3]*float4(shmem_f32[q2[4] >> 4   ], shmem_f32[q2[5] >> 4   ], shmem_f32[q2[6] >> 4   ], shmem_f32[q2[7] >> 4   ]);
3091
3092            acc1 = (acc1 + acc3) + (acc2 + acc4);
3093
3094            sumf[row] += e8m0_to_fp32(xb.e) * ((acc1[0] + acc1[1]) + (acc1[2] + acc1[3]));
3095        }
3096
3097        yb += 16 * QK_MXFP4;
3098    }
3099
3100    device float * dst_f32 = (device float *) dst + (uint64_t)im*args.ne0*args.ne1 + (uint64_t)r1*args.ne0;
3101
3102    for (int row = 0; row < NR0 && first_row + row < args.ne0; ++row) {
3103        float sum_all = simd_sum(sumf[row]);
3104        if (tiisg == 0) {
3105            dst_f32[first_row + row] = sum_all;
3106        }
3107    }
3108}
3109
3110[[host_name("kernel_mul_mv_mxfp4_f32")]]
3111kernel void kernel_mul_mv_mxfp4_f32(
3112        constant ggml_metal_kargs_mul_mv & args,
3113        device const char * src0,
3114        device const char * src1,
3115        device       char * dst,
3116        threadgroup  char * shmem [[threadgroup(0)]],
3117        uint3  tgpig[[threadgroup_position_in_grid]],
3118        ushort tiisg[[thread_index_in_simdgroup]],
3119        ushort sgitg[[simdgroup_index_in_threadgroup]]) {
3120
3121    kernel_mul_mv_mxfp4_f32_impl<N_R0_MXFP4, constant ggml_metal_kargs_mul_mv &>(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg);
3122}
3123
3124template<int nr0, typename args_t>
3125void kernel_mul_mv_tq2_0_f32_impl(
3126        args_t args,
3127        device const char * src0,
3128        device const char * src1,
3129        device       char * dst,
3130        threadgroup  char * shmem,
3131        uint3  tgpig,
3132        ushort tiisg,
3133        ushort sgitg) {
3134    const short NSG = FC_mul_mv_nsg;
3135
3136    const int nb = args.ne00/QK_K;
3137
3138    const int r0 = tgpig.x;
3139    const int r1 = tgpig.y;
3140    const int im = tgpig.z;
3141
3142    const int first_row = (r0 * NSG + sgitg) * nr0;
3143
3144    const uint i12 = im%FC_mul_mv_ne12;
3145    const uint i13 = im/FC_mul_mv_ne12;
3146
3147    const uint64_t offset1 =        r1*args.nb11 + (i12        )*args.nb12 + (i13        )*args.nb13;
3148
3149    device const float * y = (device const float *) (src1 + offset1);
3150
3151    device const block_tq2_0 * ax[nr0];
3152    for (int row = 0; row < nr0; ++row) {
3153        const uint64_t offset0 = (first_row + row)*args.nb01 + (i12/FC_mul_mv_r2)*args.nb02 + (i13/FC_mul_mv_r3)*args.nb03;
3154        ax[row] = (device const block_tq2_0 *) ((device char *) src0 + offset0);
3155    }
3156
3157    float sumf[nr0] = {0.f};
3158
3159    // 8 threads per block, NBLOCK blocks per pass, 2 halves per block per pass
3160    constexpr short NBLOCK = 4;
3161
3162    constexpr short NB = N_SIMDWIDTH/NBLOCK; // threads per block
3163
3164    const short blk = tiisg / NB;    // 0..NBLOCK-1, block handled by this thread
3165    const short htg = tiisg % NB;    // 0..NB-1, thread within block (0..7)
3166
3167    // byte and y base offsets within the block (32 elements per thread, 4 per byte)
3168    device const float4 * yb4 = (device const float4 *)(y + 4*htg + blk*QK_K);
3169
3170    // hoisted per-byte coefficients (from y) and total y-sum, shared across rows
3171    // ref: https://github.com/ggml-org/llama.cpp/pull/26980
3172    float4 coef[4];
3173
3174    for (int ib = blk; ib < nb; ib += NBLOCK) {
3175        FOR_UNROLL (short h0 = 0; h0 < 2; ++h0) {
3176            const float4 y0 = yb4[ 0 + 32*h0];
3177            const float4 y1 = yb4[ 8 + 32*h0];
3178            const float4 y2 = yb4[16 + 32*h0];
3179            const float4 y3 = yb4[24 + 32*h0];
3180
3181            float sumy = 0.f;
3182            FOR_UNROLL (short j = 0; j < 4; ++j) {
3183                coef[j] = float4(
3184                        y0[j],
3185                        y1[j] - 4.0f*y0[j],
3186                        y2[j] - 4.0f*y1[j],
3187                        y3[j] - 4.0f*y2[j]);
3188
3189                sumy += (y0[j] + y1[j]) + (y2[j] + y3[j]);
3190            }
3191
3192            FOR_UNROLL (short row = 0; row < nr0; ++row) {
3193                device const block_tq2_0 & xb = ax[row][ib];
3194                device const uchar * qs = xb.qs + 4*htg + 32*h0;
3195
3196                float sum = -sumy;
3197                FOR_UNROLL (short j = 0; j < 4; ++j) {
3198                    // express the 2-bit field shifts (v>>2, v>>4, v>>6) as float floor ops
3199                    const float v = (float)qs[j];
3200
3201                    const float f0 = v;
3202                    const float f1 = floor(v*0.25f);    // v>>2
3203                    const float f2 = floor(v*0.0625);   // v>>4
3204                    const float f3 = floor(v*0.015625); // v>>6
3205
3206                    sum += coef[j][0]*f0 + coef[j][1]*f1 + coef[j][2]*f2 + coef[j][3]*f3;
3207                }
3208
3209                sumf[row] += xb.d * sum;
3210            }
3211        }
3212
3213        yb4 += QK_K * NBLOCK / 4;
3214    }
3215
3216    device float * dst_f32 = (device float *) dst + (uint64_t)im*args.ne0*args.ne1 + (uint64_t)r1*args.ne0;
3217
3218    for (int row = 0; row < nr0; ++row) {
3219        const float tot = simd_sum(sumf[row]);
3220        if (tiisg == 0 && first_row + row < args.ne01) {
3221            dst_f32[first_row + row] = tot;
3222        }
3223    }
3224}
3225
3226[[host_name("kernel_mul_mv_tq2_0_f32")]]
3227kernel void kernel_mul_mv_tq2_0_f32(
3228        constant ggml_metal_kargs_mul_mv & args,
3229        device const char * src0,
3230        device const char * src1,
3231        device       char * dst,
3232        uint3  tgpig[[threadgroup_position_in_grid]],
3233        ushort tiisg[[thread_index_in_simdgroup]],
3234        ushort sgitg[[simdgroup_index_in_threadgroup]]) {
3235
3236    kernel_mul_mv_tq2_0_f32_impl<N_R0_TQ2_0, constant ggml_metal_kargs_mul_mv &>(args, src0, src1, dst, nullptr, tgpig, tiisg, sgitg);
3237}
3238
3239//
3240// matrix-vector multiplication
3241//
3242
3243typedef void (kernel_mul_mv_disp_t)(
3244        ggml_metal_kargs_mul_mv args,
3245        device const char * src0,
3246        device const char * src1,
3247        device       char * dst,
3248        uint3  tgpig,
3249        ushort tiisg);
3250
3251typedef void (kernel_mul_mv2_disp_t)(
3252        ggml_metal_kargs_mul_mv args,
3253        device const char * src0,
3254        device const char * src1,
3255        device       char * dst,
3256        threadgroup  char * shmem,
3257        uint3  tgpig,
3258        ushort tiisg,
3259        ushort sgitg);
3260
3261template<kernel_mul_mv_disp_t disp_fn>
3262void mmv_fn(
3263        ggml_metal_kargs_mul_mv args,
3264        device const char * src0,
3265        device const char * src1,
3266        device       char * dst,
3267        threadgroup  char * shmem,
3268        uint3  tgpig,
3269        ushort tiitg,
3270        ushort tiisg,
3271        ushort sgitg) {
3272    disp_fn(args, src0, src1, dst, tgpig, tiisg);
3273}
3274
3275template<kernel_mul_mv2_disp_t disp_fn>
3276void mmv_fn(
3277        ggml_metal_kargs_mul_mv args,
3278        device const char * src0,
3279        device const char * src1,
3280        device       char * dst,
3281        threadgroup  char * shmem,
3282        uint3  tgpig,
3283        ushort tiitg,
3284        ushort tiisg,
3285        ushort sgitg) {
3286    disp_fn(args, src0, src1, dst, shmem, tgpig, tiisg, sgitg);
3287}
3288
3289typedef decltype(mmv_fn<kernel_mul_mv_t_t_disp<half, half, ggml_metal_kargs_mul_mv>>) mul_mv_disp_fn_t;
3290
3291template<mul_mv_disp_fn_t disp_fn>
3292kernel void kernel_mul_mv_id(
3293        constant ggml_metal_kargs_mul_mv_id & args,
3294        device const char * src0s,
3295        device const char * src1,
3296        device       char * dst,
3297        device const char * ids,
3298        threadgroup  char * shmem [[threadgroup(0)]],
3299        uint3  tgpig[[threadgroup_position_in_grid]],
3300        ushort tiitg[[thread_index_in_threadgroup]],
3301        ushort tiisg[[thread_index_in_simdgroup]],
3302        ushort sgitg[[simdgroup_index_in_threadgroup]]) {
3303    const int iid1 = tgpig.z/args.nei0;
3304    const int idx  = tgpig.z%args.nei0;
3305
3306    tgpig.z = 0;
3307
3308    const int32_t i02 = ((device const int32_t *) (ids + iid1*args.nbi1))[idx];
3309
3310    const int64_t i11 = idx % args.ne11;
3311    const int64_t i12 = iid1;
3312
3313    const int64_t i1 = idx;
3314    const int64_t i2 = i12;
3315
3316    device const char * src0_cur = src0s + i02*args.nb02;
3317    device const char * src1_cur = src1  + i11*args.nb11 + i12*args.nb12;
3318
3319    device char * dst_cur = dst + (i1*args.ne0 + i2*args.ne1*args.ne0)*sizeof(float);
3320
3321    ggml_metal_kargs_mul_mv args0 = {
3322        /*.ne00 =*/ args.ne00,
3323        /*.ne01 =*/ args.ne01,
3324        /*.ne02 =*/ 1, // args.ne02,
3325        /*.nb00 =*/ args.nb00,
3326        /*.nb01 =*/ args.nb01,
3327        /*.nb02 =*/ args.nb02,
3328        /*.nb03 =*/ args.nb02, // args.ne02 == 1
3329        /*.ne10 =*/ args.ne10,
3330        /*.ne11 =*/ 1, // args.ne11,
3331        /*.ne12 =*/ 1, // args.ne12,
3332        /*.nb10 =*/ args.nb10,
3333        /*.nb11 =*/ args.nb11,
3334        /*.nb12 =*/ args.nb12,
3335        /*.nb13 =*/ args.nb12, // ne12 == 1
3336        /*.ne0  =*/ args.ne0,
3337        /*.ne1  =*/ 1, // args.ne1,
3338        /*.nr0  =*/ args.nr0,
3339        /*.r2   =*/ 1,
3340        /*.r3   =*/ 1,
3341    };
3342
3343    disp_fn(
3344        args0,
3345        /* src0 */ src0_cur,
3346        /* src1 */ src1_cur,
3347        /* dst  */ dst_cur,
3348        shmem,
3349        tgpig,
3350        tiitg,
3351        tiisg,
3352        sgitg);
3353}
3354
3355typedef decltype(kernel_mul_mv_id<mmv_fn<kernel_mul_mv_t_t_disp<float, float>>>) kernel_mul_mv_id_t;
3356
3357typedef decltype(kernel_mul_mv_id<mmv_fn<kernel_mul_mv_t_t_4_disp<float, float4, float, float4>>>) kernel_mul_mv_id_4_t;
3358
3359template [[host_name("kernel_mul_mv_id_f32_f32")]]     kernel kernel_mul_mv_id_t kernel_mul_mv_id<mmv_fn<kernel_mul_mv_t_t_disp<float, float>>>;
3360template [[host_name("kernel_mul_mv_id_f16_f32")]]     kernel kernel_mul_mv_id_t kernel_mul_mv_id<mmv_fn<kernel_mul_mv_t_t_disp<half,  float>>>;
3361#if defined(GGML_METAL_HAS_BF16)
3362template [[host_name("kernel_mul_mv_id_bf16_f32")]]    kernel kernel_mul_mv_id_t kernel_mul_mv_id<mmv_fn<kernel_mul_mv_t_t_disp<bfloat, float>>>;
3363#endif
3364template [[host_name("kernel_mul_mv_id_f32_f32_4")]]   kernel kernel_mul_mv_id_4_t kernel_mul_mv_id<mmv_fn<kernel_mul_mv_t_t_4_disp<float, float4, float, float4>>>;
3365template [[host_name("kernel_mul_mv_id_f16_f32_4")]]   kernel kernel_mul_mv_id_4_t kernel_mul_mv_id<mmv_fn<kernel_mul_mv_t_t_4_disp<half,  half4,  float, float4>>>;
3366#if defined(GGML_METAL_HAS_BF16)
3367template [[host_name("kernel_mul_mv_id_bf16_f32_4")]]  kernel kernel_mul_mv_id_4_t kernel_mul_mv_id<mmv_fn<kernel_mul_mv_t_t_4_disp<bfloat, bfloat4, float, float4>>>;
3368#endif
3369
3370template [[host_name("kernel_mul_mv_id_q8_0_f32")]]    kernel kernel_mul_mv_id_t kernel_mul_mv_id<mmv_fn<kernel_mul_mv_q8_0_f32_impl<N_R0_Q8_0>>>;
3371
3372template [[host_name("kernel_mul_mv_id_q1_0_f32")]]    kernel kernel_mul_mv_id_t kernel_mul_mv_id<mmv_fn<kernel_mul_mv_q1_0_f32_impl<N_R0_Q1_0>>>;
3373template [[host_name("kernel_mul_mv_id_q2_0_f32")]]    kernel kernel_mul_mv_id_t kernel_mul_mv_id<mmv_fn<kernel_mul_mv_q2_0_f32_impl<N_R0_Q2_0>>>;
3374template [[host_name("kernel_mul_mv_id_q4_0_f32")]]    kernel kernel_mul_mv_id_t kernel_mul_mv_id<mmv_fn<mul_vec_q_n_f32_impl<block_q4_0, N_R0_Q4_0>>>;
3375template [[host_name("kernel_mul_mv_id_q4_1_f32")]]    kernel kernel_mul_mv_id_t kernel_mul_mv_id<mmv_fn<mul_vec_q_n_f32_impl<block_q4_1, N_R0_Q4_1>>>;
3376template [[host_name("kernel_mul_mv_id_q5_0_f32")]]    kernel kernel_mul_mv_id_t kernel_mul_mv_id<mmv_fn<mul_vec_q_n_f32_impl<block_q5_0, N_R0_Q5_0>>>;
3377template [[host_name("kernel_mul_mv_id_q5_1_f32")]]    kernel kernel_mul_mv_id_t kernel_mul_mv_id<mmv_fn<mul_vec_q_n_f32_impl<block_q5_1, N_R0_Q5_1>>>;
3378
3379template [[host_name("kernel_mul_mv_id_mxfp4_f32")]]   kernel kernel_mul_mv_id_t kernel_mul_mv_id<mmv_fn<kernel_mul_mv_mxfp4_f32_impl<N_R0_MXFP4>>>;
3380
3381template [[host_name("kernel_mul_mv_id_q2_K_f32")]]    kernel kernel_mul_mv_id_t kernel_mul_mv_id<mmv_fn<kernel_mul_mv_q2_K_f32_impl   <N_R0_Q2_K>>>;
3382template [[host_name("kernel_mul_mv_id_q3_K_f32")]]    kernel kernel_mul_mv_id_t kernel_mul_mv_id<mmv_fn<kernel_mul_mv_q3_K_f32_impl   <N_R0_Q3_K>>>;
3383template [[host_name("kernel_mul_mv_id_q4_K_f32")]]    kernel kernel_mul_mv_id_t kernel_mul_mv_id<mmv_fn<kernel_mul_mv_q4_K_f32_impl   <N_R0_Q4_K>>>;
3384template [[host_name("kernel_mul_mv_id_q5_K_f32")]]    kernel kernel_mul_mv_id_t kernel_mul_mv_id<mmv_fn<kernel_mul_mv_q5_K_f32_impl   <N_R0_Q5_K>>>;
3385template [[host_name("kernel_mul_mv_id_q6_K_f32")]]    kernel kernel_mul_mv_id_t kernel_mul_mv_id<mmv_fn<kernel_mul_mv_q6_K_f32_impl   <N_R0_Q6_K>>>;
3386template [[host_name("kernel_mul_mv_id_iq1_s_f32")]]   kernel kernel_mul_mv_id_t kernel_mul_mv_id<mmv_fn<kernel_mul_mv_iq1_s_f32_disp<ggml_metal_kargs_mul_mv>>>;
3387template [[host_name("kernel_mul_mv_id_iq1_m_f32")]]   kernel kernel_mul_mv_id_t kernel_mul_mv_id<mmv_fn<kernel_mul_mv_iq1_m_f32_disp<ggml_metal_kargs_mul_mv>>>;
3388template [[host_name("kernel_mul_mv_id_iq2_xxs_f32")]] kernel kernel_mul_mv_id_t kernel_mul_mv_id<mmv_fn<kernel_mul_mv_iq2_xxs_f32_disp<ggml_metal_kargs_mul_mv>>>;
3389template [[host_name("kernel_mul_mv_id_iq2_xs_f32")]]  kernel kernel_mul_mv_id_t kernel_mul_mv_id<mmv_fn<kernel_mul_mv_iq2_xs_f32_disp<ggml_metal_kargs_mul_mv>>>;
3390template [[host_name("kernel_mul_mv_id_iq3_xxs_f32")]] kernel kernel_mul_mv_id_t kernel_mul_mv_id<mmv_fn<kernel_mul_mv_iq3_xxs_f32_disp<ggml_metal_kargs_mul_mv>>>;
3391template [[host_name("kernel_mul_mv_id_iq3_s_f32")]]   kernel kernel_mul_mv_id_t kernel_mul_mv_id<mmv_fn<kernel_mul_mv_iq3_s_f32_disp<ggml_metal_kargs_mul_mv>>>;
3392template [[host_name("kernel_mul_mv_id_iq2_s_f32")]]   kernel kernel_mul_mv_id_t kernel_mul_mv_id<mmv_fn<kernel_mul_mv_iq2_s_f32_disp<ggml_metal_kargs_mul_mv>>>;
3393template [[host_name("kernel_mul_mv_id_iq4_nl_f32")]]  kernel kernel_mul_mv_id_t kernel_mul_mv_id<mmv_fn<kernel_mul_mv_iq4_nl_f32_impl <N_R0_IQ4_NL>>>;
3394template [[host_name("kernel_mul_mv_id_iq4_xs_f32")]]  kernel kernel_mul_mv_id_t kernel_mul_mv_id<mmv_fn<kernel_mul_mv_iq4_xs_f32_impl <N_R0_IQ4_XS>>>;
3395template [[host_name("kernel_mul_mv_id_tq2_0_f32")]]   kernel kernel_mul_mv_id_t kernel_mul_mv_id<mmv_fn<kernel_mul_mv_tq2_0_f32_impl  <N_R0_TQ2_0>>>;