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 "ggml-metal-device.h"
2
3#include "ggml-metal-impl.h"
4#include "ggml-metal-tuning.h"
5
6#include "ggml-impl.h"
7
8#include <cassert>
9#include <memory>
10#include <string>
11#include <unordered_map>
12
13struct ggml_metal_device_deleter {
14    void operator()(ggml_metal_device_t ctx) {
15        ggml_metal_device_free(ctx);
16    }
17};
18
19typedef std::unique_ptr<ggml_metal_device, ggml_metal_device_deleter> ggml_metal_device_ptr;
20
21ggml_metal_device_t ggml_metal_device_get(int device, int n_devices) {
22    static std::vector<ggml_metal_device_ptr> devs;
23
24    devs.emplace_back(ggml_metal_device_init(device, n_devices));
25
26    return devs.back().get();
27}
28
29struct ggml_metal_pipelines {
30    std::unordered_map<std::string, ggml_metal_pipeline_t> data;
31};
32
33ggml_metal_pipelines_t ggml_metal_pipelines_init(void) {
34    ggml_metal_pipelines_t res = new ggml_metal_pipelines();
35
36    return res;
37}
38
39void ggml_metal_pipelines_free(ggml_metal_pipelines_t ppls) {
40    if (!ppls) {
41        return;
42    }
43
44    for (auto it = ppls->data.begin(); it != ppls->data.end(); ++it) {
45        ggml_metal_pipeline_free(it->second);
46    }
47
48    delete ppls;
49}
50
51void ggml_metal_pipelines_add(ggml_metal_pipelines_t ppls, const char * name, ggml_metal_pipeline_t pipeline) {
52    ppls->data[name] = pipeline;
53}
54
55ggml_metal_pipeline_t ggml_metal_pipelines_get(ggml_metal_pipelines_t ppls, const char * name) {
56    if (ppls->data.find(name) == ppls->data.end()) {
57        return nullptr;
58    }
59
60    return ppls->data[name];
61}
62
63struct ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_base(ggml_metal_library_t lib, ggml_op op) {
64    char base[256];
65    char name[256];
66
67    const char * op_str = "undefined";
68    switch (op) {
69        case GGML_OP_ADD_ID: op_str = "add_id"; break;
70        default: GGML_ABORT("fatal error");
71    };
72
73    snprintf(base, 256, "kernel_%s", op_str);
74    snprintf(name, 256, "%s", base);
75
76    ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name);
77    if (!res.pipeline) {
78        res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr);
79    }
80
81    return res;
82}
83
84ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_cpy(ggml_metal_library_t lib, ggml_type tsrc, ggml_type tdst) {
85    char base[256];
86    char name[256];
87
88    snprintf(base, 256, "kernel_cpy_%s_%s", ggml_type_name(tsrc), ggml_type_name(tdst));
89    snprintf(name, 256, "%s", base);
90
91    ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name);
92    if (!res.pipeline) {
93        res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr);
94    }
95
96    return res;
97}
98
99ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_pool_1d(ggml_metal_library_t lib, const ggml_tensor * op, ggml_op_pool op_pool) {
100    GGML_ASSERT(ggml_is_contiguous(op->src[0]));
101    GGML_ASSERT(op->src[0]->type == GGML_TYPE_F32 && op->src[0]->type == op->type);
102
103    const char * pool_str = "undefined";
104    switch (op_pool) {
105        case GGML_OP_POOL_AVG: pool_str = "avg"; break;
106        case GGML_OP_POOL_MAX: pool_str = "max"; break;
107        default: GGML_ASSERT(false && "not implemented");
108    };
109
110    char base[256];
111    char name[256];
112
113    snprintf(base, sizeof(base), "kernel_pool_1d_%s_%s", pool_str, ggml_type_name(op->src[0]->type));
114    snprintf(name, sizeof(name), "%s", base);
115
116    ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name);
117    if (!res.pipeline) {
118        res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr);
119    }
120
121    return res;
122}
123
124ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_pool_2d(ggml_metal_library_t lib, const ggml_tensor * op, ggml_op_pool op_pool) {
125    GGML_ASSERT(ggml_is_contiguous(op->src[0]));
126    GGML_ASSERT(op->src[0]->type == GGML_TYPE_F32 && op->src[0]->type == op->type);
127
128    const char * pool_str = "undefined";
129    switch (op_pool) {
130        case GGML_OP_POOL_AVG: pool_str = "avg"; break;
131        case GGML_OP_POOL_MAX: pool_str = "max"; break;
132        default: GGML_ASSERT(false && "not implemented");
133    };
134
135    char base[256];
136    char name[256];
137
138    snprintf(base, 256, "kernel_pool_2d_%s_%s", pool_str, ggml_type_name(op->src[0]->type));
139    snprintf(name, 256, "%s", base);
140
141    ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name);
142    if (!res.pipeline) {
143        res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr);
144    }
145
146    return res;
147}
148
149ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_get_rows(ggml_metal_library_t lib, ggml_type tsrc) {
150    char base[256];
151    char name[256];
152
153    snprintf(base, 256, "kernel_get_rows_%s", ggml_type_name(tsrc));
154    snprintf(name, 256, "%s", base);
155
156    ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name);
157    if (!res.pipeline) {
158        res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr);
159    }
160
161    return res;
162}
163
164ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_set_rows(ggml_metal_library_t lib, const ggml_tensor * op) {
165    char base[256];
166    char name[256];
167
168    const auto tsrc = op->src[0]->type;
169    const auto tidx = op->src[1]->type;
170    const auto tdst = op->type;
171
172    snprintf(base, 256, "kernel_set_rows_%s_%s_%s", ggml_type_name(tsrc), ggml_type_name(tidx), ggml_type_name(tdst));
173    snprintf(name, 256, "%s", base);
174
175    ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name);
176    if (!res.pipeline) {
177        res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr);
178    }
179
180    return res;
181}
182
183ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_diag(ggml_metal_library_t lib, const ggml_tensor * op) {
184    char base[256];
185    char name[256];
186
187    const int n = op->src[0]->ne[0];
188
189    snprintf(base, 256, "kernel_diag_%s", ggml_type_name(op->src[0]->type));
190    snprintf(name, 256, "%s_n=%d", base, n);
191
192    ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name);
193    if (!res.pipeline) {
194        res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr);
195    }
196
197    res.nsg  = 1;
198    res.smem = 0;
199
200    return res;
201}
202
203ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_repeat(ggml_metal_library_t lib, ggml_type tsrc) {
204    char base[256];
205    char name[256];
206
207    snprintf(base, 256, "kernel_repeat_%s", ggml_type_name(tsrc));
208    snprintf(name, 256, "%s", base);
209
210    ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name);
211    if (!res.pipeline) {
212        res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr);
213    }
214
215    return res;
216}
217
218ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_concat(ggml_metal_library_t lib, ggml_type tsrc) {
219    char base[256];
220    char name[256];
221
222    snprintf(base, 256, "kernel_concat_%s", ggml_type_name(tsrc));
223    snprintf(name, 256, "%s", base);
224
225    ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name);
226    if (!res.pipeline) {
227        res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr);
228    }
229
230    return res;
231}
232
233ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_unary(ggml_metal_library_t lib, const ggml_tensor * op) {
234    char base[256];
235    char name[256];
236
237    int op_num = -1;
238
239    switch (op->op) {
240        case GGML_OP_SCALE:      op_num = OP_UNARY_NUM_SCALE;      break;
241        case GGML_OP_FILL:       op_num = OP_UNARY_NUM_FILL;       break;
242        case GGML_OP_CLAMP:      op_num = OP_UNARY_NUM_CLAMP;      break;
243        case GGML_OP_SQR:        op_num = OP_UNARY_NUM_SQR;        break;
244        case GGML_OP_SQRT:       op_num = OP_UNARY_NUM_SQRT;       break;
245        case GGML_OP_SIN:        op_num = OP_UNARY_NUM_SIN;        break;
246        case GGML_OP_COS:        op_num = OP_UNARY_NUM_COS;        break;
247        case GGML_OP_LOG:        op_num = OP_UNARY_NUM_LOG;        break;
248        case GGML_OP_LEAKY_RELU: op_num = OP_UNARY_NUM_LEAKY_RELU; break;
249        case GGML_OP_UNARY:
250            switch (ggml_get_unary_op(op)) {
251                case GGML_UNARY_OP_TANH:        op_num = OP_UNARY_NUM_TANH;        break;
252                case GGML_UNARY_OP_RELU:        op_num = OP_UNARY_NUM_RELU;        break;
253                case GGML_UNARY_OP_SIGMOID:     op_num = OP_UNARY_NUM_SIGMOID;     break;
254                case GGML_UNARY_OP_GELU:        op_num = OP_UNARY_NUM_GELU;        break;
255                case GGML_UNARY_OP_GELU_ERF:    op_num = OP_UNARY_NUM_GELU_ERF;    break;
256                case GGML_UNARY_OP_GELU_QUICK:  op_num = OP_UNARY_NUM_GELU_QUICK;  break;
257                case GGML_UNARY_OP_SILU:        op_num = OP_UNARY_NUM_SILU;        break;
258                case GGML_UNARY_OP_ELU:         op_num = OP_UNARY_NUM_ELU;         break;
259                case GGML_UNARY_OP_NEG:         op_num = OP_UNARY_NUM_NEG;         break;
260                case GGML_UNARY_OP_ABS:         op_num = OP_UNARY_NUM_ABS;         break;
261                case GGML_UNARY_OP_SGN:         op_num = OP_UNARY_NUM_SGN;         break;
262                case GGML_UNARY_OP_STEP:        op_num = OP_UNARY_NUM_STEP;        break;
263                case GGML_UNARY_OP_HARDSWISH:   op_num = OP_UNARY_NUM_HARDSWISH;   break;
264                case GGML_UNARY_OP_HARDSIGMOID: op_num = OP_UNARY_NUM_HARDSIGMOID; break;
265                case GGML_UNARY_OP_EXP:         op_num = OP_UNARY_NUM_EXP;         break;
266                case GGML_UNARY_OP_SOFTPLUS:    op_num = OP_UNARY_NUM_SOFTPLUS;    break;
267                case GGML_UNARY_OP_EXPM1:       op_num = OP_UNARY_NUM_EXPM1;       break;
268                case GGML_UNARY_OP_FLOOR:       op_num = OP_UNARY_NUM_FLOOR;       break;
269                case GGML_UNARY_OP_CEIL:        op_num = OP_UNARY_NUM_CEIL;        break;
270                case GGML_UNARY_OP_ROUND:       op_num = OP_UNARY_NUM_ROUND;       break;
271                case GGML_UNARY_OP_TRUNC:       op_num = OP_UNARY_NUM_TRUNC;       break;
272                case GGML_UNARY_OP_XIELU:       op_num = OP_UNARY_NUM_XIELU;       break;
273                default: GGML_ABORT("fatal error");
274            } break;
275        default: GGML_ABORT("fatal error");
276    };
277
278    const char * t0_str = ggml_type_name(op->src[0]->type);
279    const char * t_str  = ggml_type_name(op->type);
280
281    const bool is_c4 = op->src[0]->ne[0] % 4 == 0;
282    const bool is_cnt = ggml_is_contiguous(op->src[0]) && ggml_nelements(op) < 32768;
283
284    snprintf(base, 256, "kernel_unary_%s_%s%s", t0_str, t_str, is_c4 ? "_4" : "");
285    snprintf(name, 256, "%s_op=%d_cnt=%d", base, op_num, is_cnt);
286
287    ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name);
288    if (!res.pipeline) {
289        ggml_metal_cv_t cv = ggml_metal_cv_init();
290
291        ggml_metal_cv_set_int16(cv, op_num, FC_UNARY + 0);
292        ggml_metal_cv_set_bool (cv, is_cnt, FC_UNARY + 1);
293
294        res = ggml_metal_library_compile_pipeline(lib, base, name, cv);
295
296        ggml_metal_cv_free(cv);
297    }
298
299    res.c4  = is_c4;
300    res.cnt = is_cnt;
301
302    return res;
303}
304
305ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_glu(ggml_metal_library_t lib, const ggml_tensor * op) {
306    GGML_ASSERT(ggml_is_contiguous_1(op->src[0]));
307
308    char base[256];
309    char name[256];
310
311    const char * op_str = "undefined";
312    switch (op->op) {
313        case GGML_OP_GLU:
314            switch (ggml_get_glu_op(op)) {
315                case GGML_GLU_OP_REGLU:        op_str = "reglu";        break;
316                case GGML_GLU_OP_GEGLU:        op_str = "geglu";        break;
317                case GGML_GLU_OP_SWIGLU:       op_str = "swiglu";       break;
318                case GGML_GLU_OP_SWIGLU_OAI:   op_str = "swiglu_oai";   break;
319                case GGML_GLU_OP_GEGLU_ERF:    op_str = "geglu_erf";    break;
320                case GGML_GLU_OP_GEGLU_QUICK:  op_str = "geglu_quick";  break;
321                case GGML_GLU_OP_SWIGLU_CLAMP: op_str = "swiglu_clamp"; break;
322                default: GGML_ABORT("fatal error");
323            } break;
324        default: GGML_ABORT("fatal error");
325    };
326
327    snprintf(base, 256, "kernel_%s_%s", op_str, ggml_type_name(op->src[0]->type));
328    snprintf(name, 256, "%s", base);
329
330    ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name);
331    if (!res.pipeline) {
332        res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr);
333    }
334
335    return res;
336}
337
338ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_sum(ggml_metal_library_t lib, const ggml_tensor * op) {
339    assert(op->op == GGML_OP_SUM);
340
341    char base[256];
342    char name[256];
343
344    snprintf(base, 256, "kernel_op_sum_%s", ggml_type_name(op->src[0]->type));
345    snprintf(name, 256, "%s", base);
346
347    ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name);
348    if (!res.pipeline) {
349        res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr);
350    }
351
352    return res;
353}
354
355ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_sum_rows(ggml_metal_library_t lib, const ggml_tensor * op) {
356    GGML_ASSERT(ggml_is_contiguous_rows(op->src[0]));
357
358    char base[256];
359    char name[256];
360
361    int op_num = -1;
362
363    switch (op->op) {
364        case GGML_OP_SUM_ROWS: op_num = OP_SUM_ROWS_NUM_SUM_ROWS; break;
365        case GGML_OP_MEAN:     op_num = OP_SUM_ROWS_NUM_MEAN;     break;
366        default: GGML_ABORT("fatal error");
367    };
368
369    const char * t0_str = ggml_type_name(op->src[0]->type);
370    const char * t_str  = ggml_type_name(op->type);
371
372    const bool is_c4 = op->src[0]->ne[0] % 4 == 0;
373
374    snprintf(base, 256, "kernel_sum_rows_%s_%s%s", t0_str, t_str, is_c4 ? "_4" : "");
375    snprintf(name, 256, "%s_op=%d", base, op_num);
376
377    ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name);
378    if (!res.pipeline) {
379        ggml_metal_cv_t cv = ggml_metal_cv_init();
380
381        ggml_metal_cv_set_int16(cv, op_num, FC_SUM_ROWS + 0);
382
383        res = ggml_metal_library_compile_pipeline(lib, base, name, cv);
384
385        ggml_metal_cv_free(cv);
386    }
387
388    res.smem = 32*sizeof(float);
389
390    if (is_c4) {
391        res.smem *= 4;
392    }
393
394    res.c4  = is_c4;
395
396    return res;
397}
398
399ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_cumsum_blk(ggml_metal_library_t lib, const ggml_tensor * op) {
400    GGML_ASSERT(op->op == GGML_OP_CUMSUM);
401
402    char base[256];
403    char name[256];
404
405    snprintf(base, 256, "kernel_cumsum_blk_%s", ggml_type_name(op->src[0]->type));
406    snprintf(name, 256, "%s", base);
407
408    ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name);
409    if (!res.pipeline) {
410        res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr);
411    }
412
413    return res;
414}
415
416ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_cumsum_add(ggml_metal_library_t lib, const ggml_tensor * op) {
417    GGML_ASSERT(op->op == GGML_OP_CUMSUM);
418
419    char base[256];
420    char name[256];
421
422    snprintf(base, 256, "kernel_cumsum_add_%s", ggml_type_name(op->src[0]->type));
423    snprintf(name, 256, "%s", base);
424
425    ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name);
426    if (!res.pipeline) {
427        res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr);
428    }
429
430    return res;
431}
432
433ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_tri(ggml_metal_library_t lib, const ggml_tensor * op) {
434    GGML_ASSERT(op->op == GGML_OP_TRI);
435    GGML_ASSERT(op->src[0]->nb[0] == ggml_type_size(op->src[0]->type));
436
437    char base[256];
438    char name[256];
439
440    const char * op_str = "tri";
441    const int ttype = op->op_params[0];
442
443    snprintf(base, 256, "kernel_%s_%s_%d", op_str, ggml_type_name(op->src[0]->type), ttype);
444
445    snprintf(name, 256, "%s", base);
446
447    ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name);
448    if (!res.pipeline) {
449        res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr);
450    }
451
452    return res;
453}
454
455ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_soft_max(ggml_metal_library_t lib, const ggml_tensor * op) {
456    GGML_ASSERT(!op->src[1] || op->src[1]->type == GGML_TYPE_F16 || op->src[1]->type == GGML_TYPE_F32);
457
458    char base[256];
459    char name[256];
460
461    const char * suffix = "";
462
463    if (op->src[0]->ne[0] % 4 == 0) {
464        suffix = "_4";
465    }
466
467    const ggml_type tsrc1 = op->src[1] ? op->src[1]->type : GGML_TYPE_F32;
468
469    snprintf(base, 256, "kernel_soft_max_%s%s", ggml_type_name(tsrc1), suffix);
470    snprintf(name, 256, "%s", base);
471
472    ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name);
473    if (!res.pipeline) {
474        res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr);
475    }
476
477    res.smem = 32*sizeof(float);
478
479    return res;
480}
481
482ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_lightning_indexer(
483        ggml_metal_library_t lib,
484        const ggml_tensor * op) {
485    GGML_ASSERT(op->op == GGML_OP_LIGHTNING_INDEXER);
486
487    char name[256];
488
489    snprintf(name, 256, "kernel_lightning_indexer_%s", ggml_type_name(op->src[1]->type));
490
491    ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name);
492    if (!res.pipeline) {
493        res = ggml_metal_library_compile_pipeline(lib, name, name, nullptr);
494    }
495
496    return res;
497}
498
499ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_dsv4_hc(ggml_metal_library_t lib, ggml_op op) {
500    const char * name = nullptr;
501
502    switch (op) {
503        case GGML_OP_DSV4_HC_COMB: name = "kernel_dsv4_hc_comb_f32"; break;
504        case GGML_OP_DSV4_HC_PRE:  name = "kernel_dsv4_hc_pre_f32";  break;
505        case GGML_OP_DSV4_HC_POST: name = "kernel_dsv4_hc_post_f32"; break;
506        default: GGML_ABORT("fatal error");
507    }
508
509    ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name);
510    if (!res.pipeline) {
511        res = ggml_metal_library_compile_pipeline(lib, name, name, nullptr);
512    }
513
514    return res;
515}
516
517ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_ssm_conv(ggml_metal_library_t lib, const ggml_tensor * op) {
518    GGML_ASSERT(op->src[0]->type == GGML_TYPE_F32);
519    GGML_ASSERT(op->src[1]->type == GGML_TYPE_F32);
520
521    GGML_ASSERT(ggml_is_contiguous(op->src[0]));
522    GGML_ASSERT(ggml_is_contiguous(op->src[1]));
523
524    char base[256];
525    char name[256];
526
527    const char * suffix = "";
528
529    if (op->src[1]->ne[0] % 4 == 0) {
530        suffix = "_4";
531    }
532
533    snprintf(base, 256, "kernel_ssm_conv_%s_%s%s", ggml_type_name(op->src[0]->type), ggml_type_name(op->src[1]->type), suffix);
534    snprintf(name, 256, "%s", base);
535
536    ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name);
537    if (!res.pipeline) {
538        res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr);
539    }
540
541    return res;
542}
543
544ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_ssm_conv_batched(ggml_metal_library_t lib, const ggml_tensor * op, int ssm_conv_bs) {
545    GGML_ASSERT(op->src[0]->type == GGML_TYPE_F32);
546    GGML_ASSERT(op->src[1]->type == GGML_TYPE_F32);
547
548    GGML_ASSERT(ggml_is_contiguous(op->src[0]));
549    GGML_ASSERT(ggml_is_contiguous(op->src[1]));
550
551    char base[256];
552    char name[256];
553
554    const char * suffix = "";
555    if (op->src[1]->ne[0] % 4 == 0) {
556        suffix = "_4";
557    }
558
559    snprintf(base, 256, "kernel_ssm_conv_%s_%s_batched%s", ggml_type_name(op->src[0]->type), ggml_type_name(op->src[1]->type), suffix);
560    snprintf(name, 256, "%s_ssm_conv_bs=%d", base, ssm_conv_bs);
561
562    ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name);
563    if (!res.pipeline) {
564        ggml_metal_cv_t cv = ggml_metal_cv_init();
565
566        ggml_metal_cv_set_int16(cv, ssm_conv_bs, FC_SSM_CONV + 0);
567
568        res = ggml_metal_library_compile_pipeline(lib, base, name, cv);
569
570        ggml_metal_cv_free(cv);
571    }
572
573    return res;
574}
575
576ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_ssm_scan(ggml_metal_library_t lib, const ggml_tensor * op, bool tail)  {
577    GGML_TENSOR_LOCALS( int32_t, ne0, op->src[0], ne);
578
579    char base[256];
580    char name[256];
581
582    const int nsg = (ne00 + 31)/32;
583
584    snprintf(base, 256, "kernel_ssm_scan_%s%s", ggml_type_name(op->src[0]->type), tail ? "_tail" : "");
585    snprintf(name, 256, "%s_nsg=%d", base, nsg);
586
587    ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name);
588    if (!res.pipeline) {
589        res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr);
590    }
591
592    // Shared memory layout:
593    // - sgptg * NW floats for partial sums (nsg * 32)
594    // - sgptg floats for shared_x_dt (nsg)
595    // - sgptg floats for shared_dA (nsg)
596    // Total: nsg * (32 + 2) floats
597    res.smem = GGML_PAD((32 + 2)*sizeof(float)*nsg, 16);
598
599    return res;
600}
601
602ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_ssm_scan_ssd_mma(ggml_metal_library_t lib, const ggml_tensor * op)  {
603    char base[256];
604    char name[256];
605
606    snprintf(base, 256, "kernel_ssm_scan_ssd_mma_%s", ggml_type_name(op->src[0]->type));
607    snprintf(name, 256, "%s", base);
608
609    ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name);
610    if (!res.pipeline) {
611        res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr);
612    }
613
614    // acs/exp(acs)/state-decay vectors + dtX + SAM rows + two 8x8 tiles per simdgroup
615    res.smem = (3*OP_SSM_SCAN_SSD_CS +
616                OP_SSM_SCAN_SSD_CS*OP_SSM_SCAN_SSD_HD +
617                OP_SSM_SCAN_SSD_NSG*8*OP_SSM_SCAN_SSD_CS +
618                OP_SSM_SCAN_SSD_NSG*2*8*8)*sizeof(float);
619
620    return res;
621}
622
623ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_rwkv(ggml_metal_library_t lib, const ggml_tensor * op) {
624    char base[256];
625    char name[256];
626
627    const int64_t C = op->ne[0];
628    const int64_t H = op->src[0]->ne[1];
629
630    switch (op->op) {
631        case GGML_OP_RWKV_WKV6:
632            {
633                GGML_ASSERT(op->src[5]->type == GGML_TYPE_F32);
634                GGML_ASSERT(C % H == 0);
635                GGML_ASSERT(C / H == 64);
636
637                snprintf(base, 256, "kernel_rwkv_wkv6_%s", ggml_type_name(op->src[0]->type));
638            } break;
639        case GGML_OP_RWKV_WKV7:
640            {
641                GGML_ASSERT(op->src[6]->type == GGML_TYPE_F32);
642                GGML_ASSERT(C % H == 0);
643                GGML_ASSERT(C / H == 64);
644
645                snprintf(base, 256, "kernel_rwkv_wkv7_%s", ggml_type_name(op->src[0]->type));
646            } break;
647        default:
648            GGML_ABORT("fatal error");
649    }
650
651    snprintf(name, 256, "%s", base);
652
653    ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name);
654    if (!res.pipeline) {
655        res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr);
656    }
657
658    return res;
659}
660
661ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_gated_delta_net(ggml_metal_library_t lib, const ggml_tensor * op) {
662    char base[256];
663    char name[256];
664
665    // v is src[2], dimensions: S_v = ne[0], H = ne[1]
666    const int ne20 = op->src[2]->ne[0]; // S_v
667    const int ne21 = op->src[2]->ne[1]; // H
668    const int ne30 = op->src[3]->ne[0]; // G
669    // state is src[5], 4D [S_v, S_v, H_v, n_seqs] (s0 only); K is op param 0.
670    const int K = ggml_get_op_params_i32(op, 0);
671
672    const int nsg = op->src[2]->ne[0]/32;
673
674    GGML_ASSERT(op->src[5]->type == GGML_TYPE_F32);
675    GGML_ASSERT(op->ne[0] == ne20 * ne21);
676    GGML_ASSERT(ne20 % 32 == 0);
677
678    snprintf(base, 256, "kernel_gated_delta_net_%s_%d", ggml_type_name(op->src[0]->type), nsg);
679    snprintf(name, 256, "%s_ne20=%d_ne30=%d_K=%d", base, ne20, ne30, K);
680
681    ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name);
682    if (!res.pipeline) {
683        ggml_metal_cv_t cv = ggml_metal_cv_init();
684
685        ggml_metal_cv_set_int16(cv, ne20, FC_GATED_DELTA_NET + 0);
686        ggml_metal_cv_set_int16(cv, ne30, FC_GATED_DELTA_NET + 1);
687        ggml_metal_cv_set_int16(cv, K,    FC_GATED_DELTA_NET + 2);
688
689        res = ggml_metal_library_compile_pipeline(lib, base, name, cv);
690
691        ggml_metal_cv_free(cv);
692    }
693
694    res.nsg = nsg;
695
696    return res;
697}
698
699ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_solve_tri(ggml_metal_library_t lib, const ggml_tensor * op) {
700    char base[256];
701    char name[256];
702
703    const int nsg = 8;
704    const int n   = op->src[1]->ne[1];
705    const int k   = op->src[1]->ne[0];
706
707    snprintf(base, 256, "kernel_solve_tri_%s", ggml_type_name(op->src[0]->type));
708    snprintf(name, 256, "%s_nsg=%d_n=%d_k=%d", base, nsg, n, k);
709
710    ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name);
711    if (!res.pipeline) {
712        ggml_metal_cv_t cv = ggml_metal_cv_init();
713
714        ggml_metal_cv_set_int16(cv, nsg, FC_SOLVE_TRI + 0);
715        ggml_metal_cv_set_int16(cv, n,   FC_SOLVE_TRI + 1);
716        ggml_metal_cv_set_int16(cv, k,   FC_SOLVE_TRI + 2);
717
718        res = ggml_metal_library_compile_pipeline(lib, base, name, cv);
719
720        ggml_metal_cv_free(cv);
721    }
722
723    res.nsg  = nsg;
724    res.smem = GGML_PAD(GGML_PAD(n, 32)*nsg*sizeof(float), 16);
725
726    return res;
727}
728
729ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_ext(ggml_metal_library_t lib, const ggml_tensor * op, int nsg, int nxpsg, int r1ptg) {
730    char base[256];
731    char name[256];
732
733    const ggml_type tsrc0 = op->src[0]->type;
734    const ggml_type tsrc1 = op->src[1]->type;
735    const int       ne12  = op->src[1]->ne[2];
736    const int       r2    = ne12 / op->src[0]->ne[2];
737    const int       r3    = op->src[1]->ne[3] / op->src[0]->ne[3];
738
739    GGML_ASSERT(ne12 <= INT16_MAX && r2 <= INT16_MAX && r3 <= INT16_MAX);
740
741    snprintf(base, 256, "kernel_mul_mv_ext_%s_%s_r1_%d", ggml_type_name(tsrc0), ggml_type_name(tsrc1), r1ptg);
742    snprintf(name, 256, "%s_nsg=%d_nxpsg=%d_ne12=%d_r2=%d_r3=%d", base, nsg, nxpsg, ne12, r2, r3);
743
744    ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name);
745    if (!res.pipeline) {
746        ggml_metal_cv_t cv = ggml_metal_cv_init();
747
748        ggml_metal_cv_set_int16(cv, nsg,            FC_MUL_MV + 0);
749        ggml_metal_cv_set_int16(cv, nxpsg,          FC_MUL_MV + 1);
750        ggml_metal_cv_set_int16(cv, (int16_t) ne12, FC_MUL_MV + 2);
751        ggml_metal_cv_set_int16(cv, (int16_t) r2,   FC_MUL_MV + 3);
752        ggml_metal_cv_set_int16(cv, (int16_t) r3,   FC_MUL_MV + 4);
753
754        res = ggml_metal_library_compile_pipeline(lib, base, name, cv);
755
756        ggml_metal_cv_free(cv);
757    }
758
759    return res;
760}
761
762ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mm(ggml_metal_library_t lib, const ggml_tensor * op) {
763    char base[256];
764    char name[256];
765
766    const ggml_type tsrc0 = op->src[0]->type;
767    const ggml_type tsrc1 = op->src[1]->type;
768
769    const bool bc_inp = op->src[0]->ne[0] % 32 != 0;
770
771    constexpr int NRA = SZ_SIMDGROUP * N_MM_BLOCK_Y * N_MM_SIMD_GROUP_Y;
772    constexpr int NRB = SZ_SIMDGROUP * N_MM_BLOCK_X * N_MM_SIMD_GROUP_X;
773
774    const bool has_tensor = ggml_metal_device_get_props(ggml_metal_library_get_device(lib))->has_tensor;
775
776    const bool bc_out = has_tensor
777        ? (op->ne[0] % NRA != 0 || op->ne[1] % NRB != 0)
778        : (op->ne[0] % 64  != 0 || op->ne[1] % 32  != 0);
779
780    GGML_ASSERT(op->src[1]->ne[2] <= INT16_MAX && op->src[1]->ne[3] <= INT16_MAX);
781    const int16_t ne12 = (int16_t) op->src[1]->ne[2];
782    const int16_t ne13 = (int16_t) op->src[1]->ne[3];
783    const int16_t r2   = (int16_t) (ne12 / op->src[0]->ne[2]);
784    const int16_t r3   = (int16_t) (ne13 / op->src[0]->ne[3]);
785
786    snprintf(base, 256, "kernel_mul_mm_%s_%s", ggml_type_name(tsrc0), ggml_type_name(tsrc1));
787    snprintf(name, 256, "%s_bci=%d_bco=%d_ne12=%d_ne13=%d_r2=%d_r3=%d",
788             base, bc_inp, bc_out, ne12, ne13, r2, r3);
789
790    ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name);
791    if (!res.pipeline) {
792        ggml_metal_cv_t cv = ggml_metal_cv_init();
793
794        ggml_metal_cv_set_bool(cv, bc_inp, FC_MUL_MM + 0);
795        ggml_metal_cv_set_bool(cv, bc_out, FC_MUL_MM + 1);
796        ggml_metal_cv_set_int16(cv, ne12,  FC_MUL_MM + 2);
797        ggml_metal_cv_set_int16(cv, ne13,  FC_MUL_MM + 3);
798        ggml_metal_cv_set_int16(cv, r2,    FC_MUL_MM + 4);
799        ggml_metal_cv_set_int16(cv, r3,    FC_MUL_MM + 5);
800
801        res = ggml_metal_library_compile_pipeline(lib, base, name, cv);
802
803        ggml_metal_cv_free(cv);
804    }
805
806    if (has_tensor) {
807        res.nr0 = NRA;
808        res.nr1 = NRB;
809
810        const size_t smem_a = NRA * N_MM_NK_TOTAL * sizeof(ggml_fp16_t);
811        res.smem = smem_a;
812    } else {
813        res.nr0 = 64;
814        res.nr1 = 32;
815
816        res.smem = bc_out ? 8192 : (4096 + 2048);
817    }
818
819    res.nsg = N_MM_SIMD_GROUP_X * N_MM_SIMD_GROUP_Y;
820
821    return res;
822}
823
824ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv(ggml_metal_library_t lib, const ggml_tensor * op) {
825    GGML_TENSOR_LOCALS( int32_t, ne0, op->src[0], ne);
826    GGML_TENSOR_LOCALS( int32_t, ne1, op->src[1], ne);
827
828    char base[256];
829    char name[256];
830
831    int nsg = 0; // number of simdgroups
832    int nr0 = 0; // number of src0 rows per simdgroup
833    int nr1 = 1; // number of src1 rows per threadgroup
834
835    size_t smem = 0; // shared memory
836
837    const ggml_type tsrc0 = op->src[0]->type;
838    const ggml_type tsrc1 = op->src[1]->type;
839
840    const char * suffix = "";
841
842    bool split = false;
843
844    // use custom matrix x vector kernel
845    switch (tsrc0) {
846        case GGML_TYPE_F32:
847        case GGML_TYPE_F16:
848        case GGML_TYPE_BF16:
849            {
850                if (ne00 < 32) {
851                    nsg = 1;
852                    nr0 = 32;
853                    nr1 = 1;
854                    suffix = "_short";
855                } else {
856                    nsg = std::min(4, (ne00 + 127) / 128);
857                    nr0 = 2;
858                    nr1 = 1;
859                    smem = 32*sizeof(float)*nr0;
860                    suffix = ne00 % 4 == 0 ? "_4" : "";
861                }
862            } break;
863        case GGML_TYPE_Q1_0:
864            {
865                nsg = N_SG_Q1_0;
866                nr0 = N_R0_Q1_0;
867            } break;
868        case GGML_TYPE_Q2_0:
869            {
870                nsg = N_SG_Q2_0;
871                nr0 = N_R0_Q2_0;
872            } break;
873        case GGML_TYPE_Q4_0:
874            {
875                nsg = N_SG_Q4_0;
876                nr0 = N_R0_Q4_0;
877            } break;
878        case GGML_TYPE_Q4_1:
879            {
880                nsg = N_SG_Q4_1;
881                nr0 = N_R0_Q4_1;
882            } break;
883        case GGML_TYPE_Q5_0:
884            {
885                nsg = N_SG_Q5_0;
886                nr0 = N_R0_Q5_0;
887            } break;
888        case GGML_TYPE_Q5_1:
889            {
890                nsg = N_SG_Q5_1;
891                nr0 = N_R0_Q5_1;
892            } break;
893        case GGML_TYPE_Q8_0:
894            {
895                nsg = N_SG_Q8_0;
896                nr0 = N_R0_Q8_0;
897                smem = 32*sizeof(float)*N_R0_Q8_0;
898            } break;
899        case GGML_TYPE_MXFP4:
900            {
901                nsg = N_SG_MXFP4;
902                nr0 = N_R0_MXFP4;
903                smem = 32*sizeof(float);
904            } break;
905        case GGML_TYPE_Q2_K:
906            {
907                nsg = N_SG_Q2_K;
908                nr0 = N_R0_Q2_K;
909            } break;
910        case GGML_TYPE_Q3_K:
911            {
912                nsg = N_SG_Q3_K;
913                nr0 = N_R0_Q3_K;
914            } break;
915        case GGML_TYPE_Q4_K:
916            {
917                nsg = N_SG_Q4_K;
918                nr0 = N_R0_Q4_K;
919            } break;
920        case GGML_TYPE_Q5_K:
921            {
922                nsg = N_SG_Q5_K;
923                nr0 = N_R0_Q5_K;
924            } break;
925        case GGML_TYPE_Q6_K:
926            {
927                nsg = N_SG_Q6_K;
928                nr0 = N_R0_Q6_K;
929            } break;
930        case GGML_TYPE_IQ2_XXS:
931            {
932                nsg = N_SG_IQ2_XXS;
933                nr0 = N_R0_IQ2_XXS;
934                smem = 256*8+128;
935
936                const int nb32 = ne00/32;
937                if (nb32 < 32 && (32 % nb32) == 0) {
938                    nr0 = N_R0_IQ2_XXS_SPLIT;
939                    split = true;
940                }
941            } break;
942        case GGML_TYPE_IQ2_XS:
943            {
944                nsg = N_SG_IQ2_XS;
945                nr0 = N_R0_IQ2_XS;
946                smem = 512*8+128;
947
948                const int nb32 = ne00/32;
949                if (nb32 < 32 && (32 % nb32) == 0) {
950                    nr0 = N_R0_IQ2_XS_SPLIT;
951                    split = true;
952                }
953            } break;
954        case GGML_TYPE_IQ3_XXS:
955            {
956                nsg = N_SG_IQ3_XXS;
957                nr0 = N_R0_IQ3_XXS;
958                smem = 256*4+128;
959
960                // split the rows across threads when there are fewer than 32 chunks per row
961                const int nb32 = ne00/32;
962                if (nb32 < 32 && (32 % nb32) == 0) {
963                    nr0 = N_R0_IQ3_XXS_SPLIT;
964                    split = true;
965                }
966            } break;
967        case GGML_TYPE_IQ3_S:
968            {
969                nsg = N_SG_IQ3_S;
970                nr0 = N_R0_IQ3_S;
971                smem = 512*4;
972
973                const int nb32 = ne00/32;
974                if (nb32 < 32 && (32 % nb32) == 0) {
975                    nr0 = N_R0_IQ3_S_SPLIT;
976                    split = true;
977                }
978            } break;
979        case GGML_TYPE_IQ2_S:
980            {
981                nsg = N_SG_IQ2_S;
982                nr0 = N_R0_IQ2_S;
983
984                const int nb32 = ne00/32;
985                if (nb32 < 32 && (32 % nb32) == 0) {
986                    nr0 = N_R0_IQ2_S_SPLIT;
987                    split = true;
988                }
989            } break;
990        case GGML_TYPE_IQ1_S:
991            {
992                nsg = N_SG_IQ1_S;
993                nr0 = N_R0_IQ1_S;
994
995                const int nb32 = ne00/32;
996                if (nb32 < 32 && (32 % nb32) == 0) {
997                    nr0 = N_R0_IQ1_S_SPLIT;
998                    split = true;
999                }
1000            } break;
1001        case GGML_TYPE_IQ1_M:
1002            {
1003                nsg = N_SG_IQ1_M;
1004                nr0 = N_R0_IQ1_M;
1005
1006                const int nb32 = ne00/32;
1007                if (nb32 < 32 && (32 % nb32) == 0) {
1008                    nr0 = N_R0_IQ1_M_SPLIT;
1009                    split = true;
1010                }
1011            } break;
1012        case GGML_TYPE_IQ4_NL:
1013            {
1014                nsg = N_SG_IQ4_NL;
1015                nr0 = N_R0_IQ4_NL;
1016                smem = 32*sizeof(float);
1017            } break;
1018        case GGML_TYPE_IQ4_XS:
1019            {
1020                nsg = N_SG_IQ4_XS;
1021                nr0 = N_R0_IQ4_XS;
1022                smem = 32*sizeof(float);
1023            } break;
1024        case GGML_TYPE_TQ2_0:
1025            {
1026                nsg = N_SG_TQ2_0;
1027                nr0 = N_R0_TQ2_0;
1028            } break;
1029        default:
1030            {
1031                GGML_LOG_ERROR("Asserting on type %d\n", (int) tsrc0);
1032                GGML_ABORT("not implemented");
1033            }
1034    };
1035
1036    GGML_ASSERT(ne12 <= INT16_MAX && ne13 <= INT16_MAX);
1037    const int16_t r2 = (int16_t) (ne12 / ne02);
1038    const int16_t r3 = (int16_t) (ne13 / ne03);
1039
1040    snprintf(base, 256, "kernel_mul_mv_%s_%s%s", ggml_type_name(tsrc0), ggml_type_name(tsrc1), suffix);
1041    snprintf(name, 256, "%s_nsg=%d_ne12=%d_r2=%d_r3=%d_split=%d", base, nsg, ne12, r2, r3, split);
1042
1043    ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name);
1044    if (!res.pipeline) {
1045        ggml_metal_cv_t cv = ggml_metal_cv_init();
1046
1047        ggml_metal_cv_set_int16(cv, nsg,            FC_MUL_MV + 0);
1048        ggml_metal_cv_set_int16(cv, (int16_t) ne12, FC_MUL_MV + 2);
1049        ggml_metal_cv_set_int16(cv, r2,             FC_MUL_MV + 3);
1050        ggml_metal_cv_set_int16(cv, r3,             FC_MUL_MV + 4);
1051        ggml_metal_cv_set_bool (cv, split,          FC_MUL_MV + 5);
1052
1053        res = ggml_metal_library_compile_pipeline(lib, base, name, cv);
1054
1055        ggml_metal_cv_free(cv);
1056    }
1057
1058    res.nr0  = nr0;
1059    res.nr1  = nr1;
1060    res.nsg  = nsg;
1061    res.smem = smem;
1062
1063    return res;
1064}
1065
1066ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mm_id_map0(ggml_metal_library_t lib, int ne02, int ne20) {
1067    char base[256];
1068    char name[256];
1069
1070    snprintf(base, 256, "kernel_mul_mm_id_map0_ne20_%d", ne20);
1071    snprintf(name, 256, "%s_ne02=%d", base, ne02);
1072
1073    ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name);
1074    if (!res.pipeline) {
1075        res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr);
1076    }
1077
1078    res.smem = (size_t) ne02*ne20*sizeof(uint16_t);
1079    res.smem = GGML_PAD(res.smem, 16);
1080
1081    return res;
1082}
1083
1084ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mm_id(ggml_metal_library_t lib, const ggml_tensor * op) {
1085    char base[256];
1086    char name[256];
1087
1088    const ggml_type tsrc0 = op->src[0]->type;
1089    const ggml_type tsrc1 = op->src[1]->type;
1090
1091    const bool bc_inp = op->src[0]->ne[0] % 32 != 0;
1092
1093    snprintf(base, 256, "kernel_mul_mm_id_%s_%s", ggml_type_name(tsrc0), ggml_type_name(tsrc1));
1094    snprintf(name, 256, "%s_bci=%d", base, bc_inp);
1095
1096    ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name);
1097    if (!res.pipeline) {
1098        ggml_metal_cv_t cv = ggml_metal_cv_init();
1099
1100        ggml_metal_cv_set_bool(cv, bc_inp, FC_MUL_MM + 0);
1101
1102        res = ggml_metal_library_compile_pipeline(lib, base, name, cv);
1103
1104        ggml_metal_cv_free(cv);
1105    }
1106
1107    res.smem = 8192;
1108
1109    return res;
1110}
1111
1112ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_mul_mv_id(ggml_metal_library_t lib, const ggml_tensor * op) {
1113    GGML_TENSOR_LOCALS( int32_t, ne0, op->src[0], ne);
1114    GGML_TENSOR_LOCALS( int32_t, ne1, op->src[1], ne);
1115
1116    char base[256];
1117    char name[256];
1118
1119    int nsg = 0; // number of simdgroups
1120    int nr0 = 0; // number of src0 rows per simdgroup
1121    int nr1 = 1; // number of src1 rows per threadgroup
1122
1123    size_t smem = 0; // shared memory
1124
1125    const ggml_type tsrc0 = op->src[0]->type;
1126    const ggml_type tsrc1 = op->src[1]->type;
1127
1128    const char * suffix = "";
1129
1130    bool split = false;
1131
1132        // use custom matrix x vector kernel
1133    switch (tsrc0) {
1134        case GGML_TYPE_F32:
1135        case GGML_TYPE_F16:
1136        case GGML_TYPE_BF16:
1137            {
1138                nsg = std::min(4, (ne00 + 127) / 128);
1139                nr0 = 2;
1140                nr1 = 1;
1141                smem = 32*sizeof(float)*nr0;
1142                suffix = ne00 % 4 == 0 ? "_4" : "";
1143            } break;
1144        case GGML_TYPE_Q1_0:
1145            {
1146                nsg = N_SG_Q1_0;
1147                nr0 = N_R0_Q1_0;
1148            } break;
1149        case GGML_TYPE_Q2_0:
1150            {
1151                nsg = N_SG_Q2_0;
1152                nr0 = N_R0_Q2_0;
1153            } break;
1154        case GGML_TYPE_Q4_0:
1155            {
1156                nsg = N_SG_Q4_0;
1157                nr0 = N_R0_Q4_0;
1158            } break;
1159        case GGML_TYPE_Q4_1:
1160            {
1161                nsg = N_SG_Q4_1;
1162                nr0 = N_R0_Q4_1;
1163            } break;
1164        case GGML_TYPE_Q5_0:
1165            {
1166                nsg = N_SG_Q5_0;
1167                nr0 = N_R0_Q5_0;
1168            } break;
1169        case GGML_TYPE_Q5_1:
1170            {
1171                nsg = N_SG_Q5_1;
1172                nr0 = N_R0_Q5_1;
1173            } break;
1174        case GGML_TYPE_Q8_0:
1175            {
1176                nsg = N_SG_Q8_0;
1177                nr0 = N_R0_Q8_0;
1178                smem = 32*sizeof(float)*N_R0_Q8_0;
1179            } break;
1180        case GGML_TYPE_MXFP4:
1181            {
1182                nsg = N_SG_MXFP4;
1183                nr0 = N_R0_MXFP4;
1184                smem = 32*sizeof(float);
1185            } break;
1186        case GGML_TYPE_Q2_K:
1187            {
1188                nsg = N_SG_Q2_K;
1189                nr0 = N_R0_Q2_K;
1190            } break;
1191        case GGML_TYPE_Q3_K:
1192            {
1193                nsg = N_SG_Q3_K;
1194                nr0 = N_R0_Q3_K;
1195            } break;
1196        case GGML_TYPE_Q4_K:
1197            {
1198                nsg = N_SG_Q4_K;
1199                nr0 = N_R0_Q4_K;
1200            } break;
1201        case GGML_TYPE_Q5_K:
1202            {
1203                nsg = N_SG_Q5_K;
1204                nr0 = N_R0_Q5_K;
1205            } break;
1206        case GGML_TYPE_Q6_K:
1207            {
1208                nsg = N_SG_Q6_K;
1209                nr0 = N_R0_Q6_K;
1210            } break;
1211        case GGML_TYPE_IQ2_XXS:
1212            {
1213                nsg = N_SG_IQ2_XXS;
1214                nr0 = N_R0_IQ2_XXS;
1215                smem = 256*8+128;
1216
1217                const int nb32 = ne00/32;
1218                if (nb32 < 32 && (32 % nb32) == 0) {
1219                    nr0 = N_R0_IQ2_XXS_SPLIT;
1220                    split = true;
1221                }
1222            } break;
1223        case GGML_TYPE_IQ2_XS:
1224            {
1225                nsg = N_SG_IQ2_XS;
1226                nr0 = N_R0_IQ2_XS;
1227                smem = 512*8+128;
1228
1229                const int nb32 = ne00/32;
1230                if (nb32 < 32 && (32 % nb32) == 0) {
1231                    nr0 = N_R0_IQ2_XS_SPLIT;
1232                    split = true;
1233                }
1234            } break;
1235        case GGML_TYPE_IQ3_XXS:
1236            {
1237                nsg = N_SG_IQ3_XXS;
1238                nr0 = N_R0_IQ3_XXS;
1239                smem = 256*4+128;
1240
1241                // split the rows across threads when there are fewer than 32 chunks per row
1242                const int nb32 = ne00/32;
1243                if (nb32 < 32 && (32 % nb32) == 0) {
1244                    nr0 = N_R0_IQ3_XXS_SPLIT;
1245                    split = true;
1246                }
1247            } break;
1248        case GGML_TYPE_IQ3_S:
1249            {
1250                nsg = N_SG_IQ3_S;
1251                nr0 = N_R0_IQ3_S;
1252                smem = 512*4;
1253
1254                const int nb32 = ne00/32;
1255                if (nb32 < 32 && (32 % nb32) == 0) {
1256                    nr0 = N_R0_IQ3_S_SPLIT;
1257                    split = true;
1258                }
1259            } break;
1260        case GGML_TYPE_IQ2_S:
1261            {
1262                nsg = N_SG_IQ2_S;
1263                nr0 = N_R0_IQ2_S;
1264
1265                const int nb32 = ne00/32;
1266                if (nb32 < 32 && (32 % nb32) == 0) {
1267                    nr0 = N_R0_IQ2_S_SPLIT;
1268                    split = true;
1269                }
1270            } break;
1271        case GGML_TYPE_IQ1_S:
1272            {
1273                nsg = N_SG_IQ1_S;
1274                nr0 = N_R0_IQ1_S;
1275
1276                const int nb32 = ne00/32;
1277                if (nb32 < 32 && (32 % nb32) == 0) {
1278                    nr0 = N_R0_IQ1_S_SPLIT;
1279                    split = true;
1280                }
1281            } break;
1282        case GGML_TYPE_IQ1_M:
1283            {
1284                nsg = N_SG_IQ1_M;
1285                nr0 = N_R0_IQ1_M;
1286
1287                const int nb32 = ne00/32;
1288                if (nb32 < 32 && (32 % nb32) == 0) {
1289                    nr0 = N_R0_IQ1_M_SPLIT;
1290                    split = true;
1291                }
1292            } break;
1293        case GGML_TYPE_IQ4_NL:
1294            {
1295                nsg = N_SG_IQ4_NL;
1296                nr0 = N_R0_IQ4_NL;
1297                smem = 32*sizeof(float);
1298            } break;
1299        case GGML_TYPE_IQ4_XS:
1300            {
1301                nsg = N_SG_IQ4_XS;
1302                nr0 = N_R0_IQ4_XS;
1303                smem = 32*sizeof(float);
1304            } break;
1305        case GGML_TYPE_TQ2_0:
1306            {
1307                nsg = N_SG_TQ2_0;
1308                nr0 = N_R0_TQ2_0;
1309            } break;
1310        default:
1311            {
1312                GGML_LOG_ERROR("Asserting on type %d\n", (int)op->src[2]->type);
1313                GGML_ABORT("not implemented");
1314            }
1315    };
1316
1317    snprintf(base, 256, "kernel_mul_mv_id_%s_%s%s", ggml_type_name(tsrc0), ggml_type_name(tsrc1), suffix);
1318    snprintf(name, 256, "%s_nsg=%d_split=%d", base, nsg, split);
1319
1320    ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name);
1321    if (!res.pipeline) {
1322        ggml_metal_cv_t cv = ggml_metal_cv_init();
1323
1324        ggml_metal_cv_set_int16(cv, nsg, FC_MUL_MV + 0);
1325        ggml_metal_cv_set_int16(cv, 1,   FC_MUL_MV + 2);
1326        ggml_metal_cv_set_int16(cv, 1,   FC_MUL_MV + 3);
1327        ggml_metal_cv_set_int16(cv, 1,   FC_MUL_MV + 4);
1328        ggml_metal_cv_set_bool (cv, split, FC_MUL_MV + 5);
1329
1330        res = ggml_metal_library_compile_pipeline(lib, base, name, cv);
1331
1332        ggml_metal_cv_free(cv);
1333    }
1334
1335    res.nr0  = nr0;
1336    res.nr1  = nr1;
1337    res.nsg  = nsg;
1338    res.smem = smem;
1339
1340    return res;
1341}
1342
1343ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_argmax(ggml_metal_library_t lib, const ggml_tensor * op) {
1344    GGML_ASSERT(op->src[0]->type == GGML_TYPE_F32);
1345    GGML_ASSERT(ggml_is_contiguous_1(op->src[0]));
1346    GGML_ASSERT(op->src[0]->nb[0] == ggml_type_size(op->src[0]->type));
1347
1348    char base[256];
1349    char name[256];
1350
1351    snprintf(base, 256, "kernel_argmax_%s", ggml_type_name(op->src[0]->type));
1352    snprintf(name, 256, "%s", base);
1353
1354    ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name);
1355    if (!res.pipeline) {
1356        res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr);
1357    }
1358
1359    res.smem = 32*(sizeof(float) + sizeof(int32_t));
1360
1361    return res;
1362}
1363
1364ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_argsort(ggml_metal_library_t lib, const ggml_tensor * op) {
1365    assert(op->op == GGML_OP_ARGSORT);
1366
1367    char base[256];
1368    char name[256];
1369
1370    ggml_sort_order order = (ggml_sort_order) op->op_params[0];
1371
1372    const char * order_str = "undefined";
1373    switch (order) {
1374        case GGML_SORT_ORDER_ASC:  order_str = "asc";  break;
1375        case GGML_SORT_ORDER_DESC: order_str = "desc"; break;
1376        default: GGML_ABORT("fatal error");
1377    };
1378
1379    snprintf(base, 256, "kernel_argsort_%s_%s_%s", ggml_type_name(op->src[0]->type), ggml_type_name(op->type), order_str);
1380    snprintf(name, 256, "%s", base);
1381
1382    ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name);
1383    if (!res.pipeline) {
1384        res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr);
1385    }
1386
1387    return res;
1388}
1389
1390ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_argsort_merge(ggml_metal_library_t lib, const ggml_tensor * op) {
1391    assert(op->op == GGML_OP_ARGSORT);
1392
1393    char base[256];
1394    char name[256];
1395
1396    ggml_sort_order order = (ggml_sort_order) op->op_params[0];
1397
1398    const char * order_str = "undefined";
1399    switch (order) {
1400        case GGML_SORT_ORDER_ASC:  order_str = "asc";  break;
1401        case GGML_SORT_ORDER_DESC: order_str = "desc"; break;
1402        default: GGML_ABORT("fatal error");
1403    };
1404
1405    snprintf(base, 256, "kernel_argsort_merge_%s_%s_%s", ggml_type_name(op->src[0]->type), ggml_type_name(op->type), order_str);
1406    snprintf(name, 256, "%s", base);
1407
1408    ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name);
1409    if (!res.pipeline) {
1410        res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr);
1411    }
1412
1413    return res;
1414}
1415
1416ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_fwht(ggml_metal_library_t lib, int n) {
1417    char base[256];
1418    char name[256];
1419
1420    snprintf(base, 256, "kernel_fwht_f32_%d", n);
1421    snprintf(name, 256, "%s", base);
1422
1423    ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name);
1424    if (!res.pipeline) {
1425        res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr);
1426    }
1427
1428    return res;
1429}
1430
1431// note: reuse the argsort kernel for the bitonic top_k fallback
1432ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_top_k(ggml_metal_library_t lib, const ggml_tensor * op) {
1433    assert(op->op == GGML_OP_TOP_K);
1434
1435    char base[256];
1436    char name[256];
1437
1438    // note: the top_k kernel is always descending order
1439    ggml_sort_order order = GGML_SORT_ORDER_DESC;
1440
1441    const char * order_str = "undefined";
1442    switch (order) {
1443        case GGML_SORT_ORDER_ASC:  order_str = "asc";  break;
1444        case GGML_SORT_ORDER_DESC: order_str = "desc"; break;
1445        default: GGML_ABORT("fatal error");
1446    };
1447
1448    snprintf(base, 256, "kernel_argsort_%s_%s_%s", ggml_type_name(op->src[0]->type), ggml_type_name(op->type), order_str);
1449    snprintf(name, 256, "%s", base);
1450
1451    ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name);
1452    if (!res.pipeline) {
1453        res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr);
1454    }
1455
1456    return res;
1457}
1458
1459ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_top_k_radix(ggml_metal_library_t lib, const ggml_tensor * op) {
1460    assert(op->op == GGML_OP_TOP_K);
1461
1462    char base[256];
1463    char name[256];
1464
1465    snprintf(base, 256, "kernel_top_k_%s_%s", ggml_type_name(op->src[0]->type), ggml_type_name(op->type));
1466    snprintf(name, 256, "%s", base);
1467
1468    ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name);
1469    if (!res.pipeline) {
1470        res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr);
1471    }
1472
1473    return res;
1474}
1475
1476ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_top_k_merge(ggml_metal_library_t lib, const ggml_tensor * op) {
1477    assert(op->op == GGML_OP_TOP_K);
1478
1479    char base[256];
1480    char name[256];
1481
1482    ggml_sort_order order = GGML_SORT_ORDER_DESC;
1483
1484    const char * order_str = "undefined";
1485    switch (order) {
1486        case GGML_SORT_ORDER_ASC:  order_str = "asc";  break;
1487        case GGML_SORT_ORDER_DESC: order_str = "desc"; break;
1488        default: GGML_ABORT("fatal error");
1489    };
1490
1491    snprintf(base, 256, "kernel_argsort_merge_%s_%s_%s", ggml_type_name(op->src[0]->type), ggml_type_name(op->type), order_str);
1492    snprintf(name, 256, "%s", base);
1493
1494    ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name);
1495    if (!res.pipeline) {
1496        res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr);
1497    }
1498
1499    return res;
1500}
1501
1502ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_attn_ext_pad(
1503        ggml_metal_library_t lib,
1504        const struct ggml_tensor * op,
1505        bool    has_mask,
1506        int32_t ncpsg) {
1507    assert(op->op == GGML_OP_FLASH_ATTN_EXT);
1508    GGML_UNUSED(op);
1509
1510    char base[256];
1511    char name[256];
1512
1513    snprintf(base, 256, "kernel_%s",
1514            "flash_attn_ext_pad");
1515
1516    snprintf(name, 256, "%s_mask=%d_ncpsg=%d",
1517            base,
1518            has_mask,
1519            ncpsg);
1520
1521    ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name);
1522    if (!res.pipeline) {
1523        ggml_metal_cv_t cv = ggml_metal_cv_init();
1524
1525        ggml_metal_cv_set_bool(cv, has_mask,  FC_FLASH_ATTN_EXT_PAD + 0);
1526        //ggml_metal_cv_set_bool(cv, has_sinks, FC_FLASH_ATTN_EXT_PAD + 1);
1527        //ggml_metal_cv_set_bool(cv, has_bias,  FC_FLASH_ATTN_EXT_PAD + 2);
1528        //ggml_metal_cv_set_bool(cv, has_scap,  FC_FLASH_ATTN_EXT_PAD + 3);
1529
1530        //ggml_metal_cv_set_int32(cv, ns10, FC_FLASH_ATTN_EXT_PAD + 20);
1531        //ggml_metal_cv_set_int32(cv, ns20, FC_FLASH_ATTN_EXT_PAD + 21);
1532        //ggml_metal_cv_set_int32(cv, nsg,  FC_FLASH_ATTN_EXT_PAD + 22);
1533        //ggml_metal_cv_set_int32(cv, nwg,  FC_FLASH_ATTN_EXT_PAD + 23);
1534        //ggml_metal_cv_set_int32(cv, nqptg, FC_FLASH_ATTN_EXT_PAD + 24);
1535        ggml_metal_cv_set_int32(cv, ncpsg, FC_FLASH_ATTN_EXT_PAD + 25);
1536
1537        res = ggml_metal_library_compile_pipeline(lib, base, name, cv);
1538
1539        ggml_metal_cv_free(cv);
1540    }
1541
1542    return res;
1543}
1544
1545ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_attn_ext_kv_f16(
1546        ggml_metal_library_t lib,
1547        const ggml_tensor * op) {
1548    assert(op->op == GGML_OP_FLASH_ATTN_EXT);
1549
1550    char base[256];
1551
1552    snprintf(base, 256, "kernel_flash_attn_ext_kv_%s_f16", ggml_type_name(op->src[1]->type));
1553
1554    ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, base);
1555    if (!res.pipeline) {
1556        res = ggml_metal_library_compile_pipeline(lib, base, base, nullptr);
1557    }
1558
1559    return res;
1560}
1561
1562ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_attn_ext_blk(
1563        ggml_metal_library_t lib,
1564        const struct ggml_tensor * op,
1565        int32_t nqptg,
1566        int32_t ncpsg) {
1567    assert(op->op == GGML_OP_FLASH_ATTN_EXT);
1568    GGML_UNUSED(op);
1569
1570    char base[256];
1571    char name[256];
1572
1573    snprintf(base, 256, "kernel_%s",
1574            "flash_attn_ext_blk");
1575
1576    snprintf(name, 256, "%s_nqptg=%d_ncpsg=%d",
1577            base,
1578            nqptg,
1579            ncpsg);
1580
1581    ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name);
1582    if (!res.pipeline) {
1583        ggml_metal_cv_t cv = ggml_metal_cv_init();
1584
1585        //ggml_metal_cv_set_bool(cv, has_mask,  FC_FLASH_ATTN_EXT_BLK + 0);
1586        //ggml_metal_cv_set_bool(cv, has_sinks, FC_FLASH_ATTN_EXT_BLK + 1);
1587        //ggml_metal_cv_set_bool(cv, has_bias,  FC_FLASH_ATTN_EXT_BLK + 2);
1588        //ggml_metal_cv_set_bool(cv, has_scap,  FC_FLASH_ATTN_EXT_BLK + 3);
1589
1590        //ggml_metal_cv_set_int32(cv, ns10, FC_FLASH_ATTN_EXT_BLK + 20);
1591        //ggml_metal_cv_set_int32(cv, ns20, FC_FLASH_ATTN_EXT_BLK + 21);
1592        //ggml_metal_cv_set_int32(cv, nsg,  FC_FLASH_ATTN_EXT_BLK + 22);
1593        //ggml_metal_cv_set_int32(cv, nwg,  FC_FLASH_ATTN_EXT_BLK + 23);
1594        ggml_metal_cv_set_int32(cv, nqptg, FC_FLASH_ATTN_EXT_BLK + 24);
1595        ggml_metal_cv_set_int32(cv, ncpsg, FC_FLASH_ATTN_EXT_BLK + 25);
1596
1597        res = ggml_metal_library_compile_pipeline(lib, base, name, cv);
1598
1599        ggml_metal_cv_free(cv);
1600    }
1601
1602    return res;
1603}
1604
1605ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_attn_ext(
1606        ggml_metal_library_t lib,
1607        const ggml_tensor * op,
1608        bool    has_mask,
1609        bool    has_sinks,
1610        bool    has_bias,
1611        bool    has_scap,
1612        bool    has_kvpad,
1613        int32_t nsg,
1614        bool    use_kv_f16,
1615        int32_t ns10,
1616        int32_t ns20) {
1617    assert(op->op == GGML_OP_FLASH_ATTN_EXT);
1618
1619    char base[256];
1620    char name[256];
1621
1622    const int32_t dk = (int32_t) op->src[1]->ne[0];
1623    const int32_t dv = (int32_t) op->src[2]->ne[0];
1624
1625    const char * type = use_kv_f16 ? "f16" : ggml_type_name(op->src[1]->type);
1626
1627    // do bounds checks for the mask?
1628    const bool bc_mask = op->src[3] && (op->src[3]->ne[1] % 8 != 0);
1629
1630    snprintf(base, 256, "kernel_%s_%s_dk%d_dv%d",
1631            "flash_attn_ext",
1632            type,
1633            dk,
1634            dv);
1635
1636    snprintf(name, 256, "%s_mask=%d_sinks=%d_bias=%d_scap=%d_kvpad=%d_bcm=%d_ns10=%d_ns20=%d_nsg=%d",
1637            base,
1638            has_mask,
1639            has_sinks,
1640            has_bias,
1641            has_scap,
1642            has_kvpad,
1643            bc_mask,
1644            ns10,
1645            ns20,
1646            nsg);
1647
1648    ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name);
1649    if (!res.pipeline) {
1650        ggml_metal_cv_t cv = ggml_metal_cv_init();
1651
1652        ggml_metal_cv_set_bool(cv, has_mask,  FC_FLASH_ATTN_EXT + 0);
1653        ggml_metal_cv_set_bool(cv, has_sinks, FC_FLASH_ATTN_EXT + 1);
1654        ggml_metal_cv_set_bool(cv, has_bias,  FC_FLASH_ATTN_EXT + 2);
1655        ggml_metal_cv_set_bool(cv, has_scap,  FC_FLASH_ATTN_EXT + 3);
1656        ggml_metal_cv_set_bool(cv, has_kvpad, FC_FLASH_ATTN_EXT + 4);
1657
1658        ggml_metal_cv_set_bool(cv, bc_mask, FC_FLASH_ATTN_EXT + 10);
1659
1660        ggml_metal_cv_set_int32(cv, ns10, FC_FLASH_ATTN_EXT + 20);
1661        ggml_metal_cv_set_int32(cv, ns20, FC_FLASH_ATTN_EXT + 21);
1662        ggml_metal_cv_set_int32(cv, nsg,  FC_FLASH_ATTN_EXT + 22);
1663
1664        res = ggml_metal_library_compile_pipeline(lib, base, name, cv);
1665
1666        ggml_metal_cv_free(cv);
1667    }
1668
1669    return res;
1670}
1671
1672ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_attn_ext_vec_idx(
1673        ggml_metal_library_t lib,
1674        const ggml_tensor * op) {
1675    assert(op->op == GGML_OP_FLASH_ATTN_EXT);
1676    assert(op->src[3]);
1677
1678    char name[256];
1679
1680    snprintf(name, 256, "kernel_flash_attn_ext_vec_idx");
1681
1682    ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name);
1683    if (!res.pipeline) {
1684        res = ggml_metal_library_compile_pipeline(lib, name, name, nullptr);
1685    }
1686
1687    GGML_UNUSED(op);
1688
1689    return res;
1690}
1691
1692ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_attn_ext_vec(
1693        ggml_metal_library_t lib,
1694        const ggml_tensor * op,
1695        bool    has_mask,
1696        bool    has_sinks,
1697        bool    has_bias,
1698        bool    has_scap,
1699        bool    has_kvpad,
1700        bool    has_sparse,
1701        int32_t nqpsg,
1702        int32_t ne,
1703        int32_t nsg,
1704        int32_t nwg,
1705        bool    use_kv_f16,
1706        int32_t ns10,
1707        int32_t ns20) {
1708    assert(op->op == GGML_OP_FLASH_ATTN_EXT);
1709
1710    char base[256];
1711    char name[256];
1712
1713    const int32_t dk = (int32_t) op->src[1]->ne[0];
1714    const int32_t dv = (int32_t) op->src[2]->ne[0];
1715
1716    const char * type = use_kv_f16 ? "f16" : ggml_type_name(op->src[1]->type);
1717
1718    char qne_suffix[16] = {0};
1719    if (!(nqpsg == 1 && ne == ggml_metal_tuning::fa_vec_baseline_ne(dk, dv))) {
1720        snprintf(qne_suffix, sizeof(qne_suffix), "_q%d_ne%d", nqpsg, ne);
1721    }
1722
1723    snprintf(base, 256, "kernel_%s_%s_dk%d_dv%d%s",
1724            "flash_attn_ext_vec",
1725            type,
1726            dk,
1727            dv,
1728            qne_suffix);
1729
1730    snprintf(name, 256, "%s_mask=%d_sink=%d_bias=%d_scap=%d_kvpad=%d_sparse=%d_ns10=%d_ns20=%d_nsg=%d_nwg=%d",
1731            base,
1732            has_mask,
1733            has_sinks,
1734            has_bias,
1735            has_scap,
1736            has_kvpad,
1737            has_sparse,
1738            ns10,
1739            ns20,
1740            nsg, nwg);
1741
1742    ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name);
1743    if (!res.pipeline) {
1744        ggml_metal_cv_t cv = ggml_metal_cv_init();
1745
1746        ggml_metal_cv_set_bool(cv, has_mask,  FC_FLASH_ATTN_EXT_VEC + 0);
1747        ggml_metal_cv_set_bool(cv, has_sinks, FC_FLASH_ATTN_EXT_VEC + 1);
1748        ggml_metal_cv_set_bool(cv, has_bias,  FC_FLASH_ATTN_EXT_VEC + 2);
1749        ggml_metal_cv_set_bool(cv, has_scap,  FC_FLASH_ATTN_EXT_VEC + 3);
1750        ggml_metal_cv_set_bool(cv, has_kvpad,  FC_FLASH_ATTN_EXT_VEC + 4);
1751        ggml_metal_cv_set_bool(cv, has_sparse, FC_FLASH_ATTN_EXT_VEC + 5);
1752
1753        ggml_metal_cv_set_int32(cv, ns10, FC_FLASH_ATTN_EXT_VEC + 20);
1754        ggml_metal_cv_set_int32(cv, ns20, FC_FLASH_ATTN_EXT_VEC + 21);
1755        ggml_metal_cv_set_int32(cv, nsg,  FC_FLASH_ATTN_EXT_VEC + 22);
1756        ggml_metal_cv_set_int32(cv, nwg,  FC_FLASH_ATTN_EXT_VEC + 23);
1757
1758        res = ggml_metal_library_compile_pipeline(lib, base, name, cv);
1759
1760        ggml_metal_cv_free(cv);
1761    }
1762
1763    return res;
1764}
1765
1766ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_flash_attn_ext_vec_reduce(
1767        ggml_metal_library_t lib,
1768        const ggml_tensor * op,
1769        int32_t dv,
1770        int32_t nwg) {
1771    assert(op->op == GGML_OP_FLASH_ATTN_EXT);
1772
1773    char base[256];
1774    char name[256];
1775
1776    snprintf(base, 256, "kernel_flash_attn_ext_vec_reduce");
1777    snprintf(name, 256, "%s_dv=%d_nwg=%d", base, dv, nwg);
1778
1779    ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name);
1780    if (!res.pipeline) {
1781        ggml_metal_cv_t cv = ggml_metal_cv_init();
1782
1783        ggml_metal_cv_set_int32(cv, dv,  FC_FLASH_ATTN_EXT_VEC_REDUCE + 0);
1784        ggml_metal_cv_set_int32(cv, nwg, FC_FLASH_ATTN_EXT_VEC_REDUCE + 1);
1785
1786        res = ggml_metal_library_compile_pipeline(lib, base, name, cv);
1787
1788        ggml_metal_cv_free(cv);
1789    }
1790
1791    return res;
1792
1793    GGML_UNUSED(op);
1794}
1795
1796ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_bin(ggml_metal_library_t lib, const ggml_tensor * op, int32_t n_fuse) {
1797    char base[256];
1798    char name[256];
1799
1800    int op_num = -1;
1801
1802    switch (op->op) {
1803        case GGML_OP_ADD: op_num = 0; break;
1804        case GGML_OP_SUB: op_num = 1; break;
1805        case GGML_OP_MUL: op_num = 2; break;
1806        case GGML_OP_DIV: op_num = 3; break;
1807        default: GGML_ABORT("fatal error");
1808    };
1809
1810    const char * t0_str = ggml_type_name(op->src[0]->type);
1811    const char * t1_str = ggml_type_name(op->src[1]->type);
1812    const char * t_str  = ggml_type_name(op->type);
1813
1814    const bool is_c4 = (op->src[0]->ne[0] % 4 == 0) && (op->src[1]->ne[0] % 4 == 0);
1815
1816    const bool is_cb = op->src[0]->ne[0] != op->src[1]->ne[0];
1817    const bool is_rb = ggml_is_contiguous(op->src[0]) && ggml_is_contiguous(op->src[1]) && (ggml_nrows(op->src[1]) == 1) && ggml_nelements(op) < 65536;
1818
1819    snprintf(base, 256, "kernel_bin_fuse_%s_%s_%s%s", t0_str, t1_str, t_str, is_c4 ? "_4" : "");
1820    snprintf(name, 256, "%s_op=%d_nf=%d_rb=%d_cb=%d", base, op_num, n_fuse, is_rb, is_cb);
1821
1822    ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name);
1823    if (!res.pipeline) {
1824        ggml_metal_cv_t cv = ggml_metal_cv_init();
1825
1826        ggml_metal_cv_set_int16(cv, op_num, FC_BIN + 0);
1827        ggml_metal_cv_set_int16(cv, n_fuse, FC_BIN + 1);
1828        ggml_metal_cv_set_bool (cv, is_rb,  FC_BIN + 2);
1829        ggml_metal_cv_set_bool (cv, is_cb,  FC_BIN + 3);
1830
1831        res = ggml_metal_library_compile_pipeline(lib, base, name, cv);
1832
1833        ggml_metal_cv_free(cv);
1834    }
1835
1836    res.c4  = is_c4;
1837    res.cnt = is_rb;
1838
1839    return res;
1840}
1841
1842ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_bin_one(ggml_metal_library_t lib, ggml_op op) {
1843    char base[256];
1844    char name[256];
1845
1846    int op_num = -1;
1847
1848    switch (op) {
1849        case GGML_OP_ADD: op_num = 0; break;
1850        case GGML_OP_SUB: op_num = 1; break;
1851        case GGML_OP_MUL: op_num = 2; break;
1852        case GGML_OP_DIV: op_num = 3; break;
1853        default: GGML_ABORT("fatal error");
1854    };
1855
1856    snprintf(base, 256, "kernel_bin_fuse_%s_%s_%s", "f32", "f32", "f32");
1857    snprintf(name, 256, "%s_op=%d_nf=%d", base, op_num, 1);
1858
1859    ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name);
1860    if (!res.pipeline) {
1861        ggml_metal_cv_t cv = ggml_metal_cv_init();
1862
1863        ggml_metal_cv_set_int16(cv, op_num, FC_BIN + 0);
1864        ggml_metal_cv_set_int16(cv, 1,      FC_BIN + 1);
1865        ggml_metal_cv_set_bool (cv, false,  FC_BIN + 2);
1866
1867        res = ggml_metal_library_compile_pipeline(lib, base, name, cv);
1868
1869        ggml_metal_cv_free(cv);
1870    }
1871
1872    return res;
1873}
1874
1875ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_l2_norm(ggml_metal_library_t lib, const ggml_tensor * op) {
1876    assert(op->op == GGML_OP_L2_NORM);
1877
1878    char base[256];
1879    char name[256];
1880
1881    const bool is_c4 = op->src[0]->ne[0] % 4 == 0;
1882
1883    const char * t0_str = ggml_type_name(op->src[0]->type);
1884    const char * t_str  = ggml_type_name(op->type);
1885
1886    snprintf(base, 256, "kernel_l2_norm_%s_%s%s", t0_str, t_str, is_c4 ? "_4" : "");
1887    snprintf(name, 256, "%s", base);
1888
1889    ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name);
1890    if (!res.pipeline) {
1891        res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr);
1892    }
1893
1894    res.c4   = is_c4;
1895    res.smem = 32*sizeof(float);
1896
1897    return res;
1898}
1899
1900ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_group_norm(ggml_metal_library_t lib, const ggml_tensor * op) {
1901    assert(op->op == GGML_OP_GROUP_NORM);
1902
1903    GGML_ASSERT(ggml_is_contiguous(op->src[0]));
1904
1905    char base[256];
1906    char name[256];
1907
1908    snprintf(base, 256, "kernel_group_norm_f32");
1909    snprintf(name, 256, "%s", base);
1910
1911    ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name);
1912    if (!res.pipeline) {
1913        res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr);
1914    }
1915
1916    res.smem = 32*sizeof(float);
1917
1918    return res;
1919}
1920
1921ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_norm(ggml_metal_library_t lib, const ggml_tensor * op, int n_fuse) {
1922    assert(op->op == GGML_OP_NORM || op->op == GGML_OP_RMS_NORM);
1923
1924    GGML_ASSERT(ggml_is_contiguous_rows(op->src[0]));
1925
1926    char base[256];
1927    char name[256];
1928
1929    const char * suffix = "";
1930    if (op->ne[0] % 4 == 0) {
1931        suffix = "_4";
1932    }
1933
1934    switch (op->op) {
1935        case GGML_OP_NORM:
1936            switch (n_fuse) {
1937                case 1: snprintf(base, 256, "kernel_norm_f32%s", suffix);         break;
1938                case 2: snprintf(base, 256, "kernel_norm_mul_f32%s", suffix);     break;
1939                case 3: snprintf(base, 256, "kernel_norm_mul_add_f32%s", suffix); break;
1940                default: GGML_ABORT("fatal error");
1941            } break;
1942        case GGML_OP_RMS_NORM:
1943            switch (n_fuse) {
1944                case 1: snprintf(base, 256, "kernel_rms_norm_f32%s", suffix);         break;
1945                case 2: snprintf(base, 256, "kernel_rms_norm_mul_f32%s", suffix);     break;
1946                case 3: snprintf(base, 256, "kernel_rms_norm_mul_add_f32%s", suffix); break;
1947                default: GGML_ABORT("fatal error");
1948            } break;
1949        default: GGML_ABORT("fatal error");
1950    }
1951
1952    snprintf(name, 256, "%s", base);
1953
1954    ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name);
1955    if (!res.pipeline) {
1956        res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr);
1957    }
1958
1959    res.smem = 32*sizeof(float);
1960
1961    return res;
1962}
1963
1964ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_rope(ggml_metal_library_t lib, const ggml_tensor * op) {
1965    assert(op->op == GGML_OP_ROPE || op->op == GGML_OP_ROPE_BACK);
1966
1967    const bool is_back = op->op == GGML_OP_ROPE_BACK;
1968
1969    char base[256];
1970    char name[256];
1971
1972    const int mode = ((const int32_t *) op->op_params)[2];
1973
1974    const bool is_neox   = mode & GGML_ROPE_TYPE_NEOX;
1975    const bool is_mrope  = mode & GGML_ROPE_TYPE_MROPE;
1976    const bool is_imrope = mode == GGML_ROPE_TYPE_IMROPE;
1977    const bool is_vision = mode == GGML_ROPE_TYPE_VISION;
1978
1979    if (is_neox) {
1980        snprintf(base, 256, "kernel_rope_neox_%s", ggml_type_name(op->src[0]->type));
1981    } else if ((is_mrope || is_imrope) && !is_vision) {
1982        GGML_ASSERT(op->src[1]->ne[0]*4 >= op->src[0]->ne[2]); // need at least 4 pos per token
1983        snprintf(base, 256, "kernel_rope_multi_%s", ggml_type_name(op->src[0]->type));
1984    } else if (is_vision) {
1985        GGML_ASSERT(op->src[1]->ne[0]*4 >= op->src[0]->ne[2]); // need at least 4 pos per token
1986        snprintf(base, 256, "kernel_rope_vision_%s", ggml_type_name(op->src[0]->type));
1987    } else {
1988        snprintf(base, 256, "kernel_rope_norm_%s", ggml_type_name(op->src[0]->type));
1989    }
1990
1991    snprintf(name, 256, "%s_imrope=%d_is_back=%d", base, is_imrope ? 1 : 0, is_back ? 1 : 0);
1992
1993    ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name);
1994    if (!res.pipeline) {
1995        ggml_metal_cv_t cv = ggml_metal_cv_init();
1996
1997        ggml_metal_cv_set_bool(cv, is_imrope, FC_ROPE + 0);
1998        ggml_metal_cv_set_bool(cv, is_back,   FC_ROPE + 1);
1999
2000        res = ggml_metal_library_compile_pipeline(lib, base, name, cv);
2001
2002        ggml_metal_cv_free(cv);
2003    }
2004
2005    return res;
2006}
2007
2008ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_im2col(ggml_metal_library_t lib, const ggml_tensor * op) {
2009    assert(op->op == GGML_OP_IM2COL);
2010
2011    GGML_TENSOR_LOCALS(int64_t, ne0, op->src[0], ne);
2012
2013    GGML_ASSERT(ggml_is_contiguous(op->src[1]));
2014    GGML_ASSERT(op->src[1]->type == GGML_TYPE_F32);
2015    GGML_ASSERT(op->type         == GGML_TYPE_F16 || op->type == GGML_TYPE_F32);
2016
2017    const bool is_2D = ((const int32_t *)(op->op_params))[6] == 1;
2018    const int64_t KH = is_2D ? ne01 : 1;
2019    const int64_t KW = ne00;
2020
2021    char base[256];
2022    char name[256];
2023
2024    if (KH*KW <= 1024) {
2025        snprintf(base, 256, "kernel_im2col_%s", ggml_type_name(op->type));
2026    } else {
2027        snprintf(base, 256, "kernel_im2col_ext_%s", ggml_type_name(op->type));
2028    }
2029    snprintf(name, 256, "%s", base);
2030
2031    ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name);
2032    if (!res.pipeline) {
2033        res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr);
2034    }
2035
2036    return res;
2037}
2038
2039ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_conv_transpose_1d(ggml_metal_library_t lib, const ggml_tensor * op) {
2040    assert(op->op == GGML_OP_CONV_TRANSPOSE_1D);
2041
2042    GGML_ASSERT(ggml_is_contiguous(op->src[0]));
2043    GGML_ASSERT(ggml_is_contiguous(op->src[1]));
2044    GGML_ASSERT(op->src[0]->type == GGML_TYPE_F16 || op->src[0]->type == GGML_TYPE_F32);
2045    GGML_ASSERT(op->src[1]->type == GGML_TYPE_F32);
2046    GGML_ASSERT(op->type         == GGML_TYPE_F32);
2047
2048    char base[256];
2049    char name[256];
2050
2051    snprintf(base, 256, "kernel_conv_transpose_1d_%s_%s", ggml_type_name(op->src[0]->type), ggml_type_name(op->src[1]->type));
2052    snprintf(name, 256, "%s", base);
2053
2054    ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name);
2055    if (!res.pipeline) {
2056        res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr);
2057    }
2058
2059    return res;
2060}
2061
2062ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_col2im_1d(ggml_metal_library_t lib, const ggml_tensor * op) {
2063    assert(op->op == GGML_OP_COL2IM_1D);
2064
2065    GGML_ASSERT(ggml_is_contiguous(op->src[0]));
2066    GGML_ASSERT(op->src[0]->type == GGML_TYPE_F32 || op->src[0]->type == GGML_TYPE_F16 || op->src[0]->type == GGML_TYPE_BF16);
2067
2068    char base[256];
2069    char name[256];
2070
2071    snprintf(base, 256, "kernel_col2im_1d_%s", ggml_type_name(op->src[0]->type));
2072    snprintf(name, 256, "%s", base);
2073
2074    ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name);
2075    if (!res.pipeline) {
2076        res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr);
2077    }
2078
2079    return res;
2080}
2081
2082ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_snake(ggml_metal_library_t lib, enum ggml_type type) {
2083    GGML_ASSERT(type == GGML_TYPE_F32 || type == GGML_TYPE_F16 || type == GGML_TYPE_BF16);
2084
2085    char base[256];
2086    char name[256];
2087
2088    snprintf(base, 256, "kernel_snake_%s", ggml_type_name(type));
2089    snprintf(name, 256, "%s", base);
2090
2091    ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name);
2092    if (!res.pipeline) {
2093        res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr);
2094    }
2095
2096    return res;
2097}
2098
2099ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_conv_transpose_2d(ggml_metal_library_t lib, const ggml_tensor * op) {
2100    assert(op->op == GGML_OP_CONV_TRANSPOSE_2D);
2101
2102    GGML_ASSERT(ggml_is_contiguous(op->src[0]));
2103    GGML_ASSERT(ggml_is_contiguous(op->src[1]));
2104    GGML_ASSERT(op->src[0]->type == GGML_TYPE_F16 || op->src[0]->type == GGML_TYPE_F32);
2105    GGML_ASSERT(op->src[1]->type == GGML_TYPE_F32);
2106    GGML_ASSERT(op->type         == GGML_TYPE_F32);
2107
2108    char base[256];
2109    char name[256];
2110
2111    snprintf(base, 256, "kernel_conv_transpose_2d_%s_%s", ggml_type_name(op->src[0]->type), ggml_type_name(op->src[1]->type));
2112    snprintf(name, 256, "%s", base);
2113
2114    ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name);
2115    if (!res.pipeline) {
2116        res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr);
2117    }
2118
2119    return res;
2120}
2121
2122ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_conv_2d(ggml_metal_library_t lib, const ggml_tensor * op) {
2123    assert(op->op == GGML_OP_CONV_2D);
2124
2125    GGML_ASSERT(ggml_is_contiguous(op->src[0]));
2126    GGML_ASSERT(op->src[0]->type == GGML_TYPE_F16 || op->src[0]->type == GGML_TYPE_F32);
2127    GGML_ASSERT(op->src[1]->type == GGML_TYPE_F32);
2128    GGML_ASSERT(op->type         == GGML_TYPE_F32);
2129
2130    char base[256];
2131    char name[256];
2132
2133    snprintf(base, 256, "kernel_conv_2d_%s_%s", ggml_type_name(op->src[0]->type), ggml_type_name(op->src[1]->type));
2134    snprintf(name, 256, "%s", base);
2135
2136    ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name);
2137    if (!res.pipeline) {
2138        res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr);
2139    }
2140
2141    return res;
2142}
2143
2144ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_conv_2d_dw(ggml_metal_library_t lib, const ggml_tensor * op, bool tiled) {
2145    assert(op->op == GGML_OP_CONV_2D_DW);
2146
2147    GGML_ASSERT(op->src[0]->type == GGML_TYPE_F16 || op->src[0]->type == GGML_TYPE_F32);
2148    GGML_ASSERT(op->src[1]->type == GGML_TYPE_F32);
2149    GGML_ASSERT(op->type         == GGML_TYPE_F32);
2150
2151    char base[256];
2152    char name[256];
2153
2154    snprintf(base, 256, "kernel_conv_2d_dw%s_%s_%s",
2155             tiled ? "_tiled" : "",
2156             ggml_type_name(op->src[0]->type), ggml_type_name(op->src[1]->type));
2157    snprintf(name, 256, "%s", base);
2158
2159    ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name);
2160    if (!res.pipeline) {
2161        res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr);
2162    }
2163
2164    return res;
2165}
2166
2167ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_conv_3d(ggml_metal_library_t lib, const ggml_tensor * op) {
2168    assert(op->op == GGML_OP_CONV_3D);
2169
2170    GGML_ASSERT(ggml_is_contiguous(op->src[0]));
2171    GGML_ASSERT(op->src[0]->type == GGML_TYPE_F16 || op->src[0]->type == GGML_TYPE_F32);
2172    GGML_ASSERT(op->src[1]->type == GGML_TYPE_F32);
2173    GGML_ASSERT(op->type         == GGML_TYPE_F32);
2174
2175    char base[256];
2176    char name[256];
2177
2178    snprintf(base, 256, "kernel_conv_3d_%s_%s", ggml_type_name(op->src[0]->type), ggml_type_name(op->src[1]->type));
2179    snprintf(name, 256, "%s", base);
2180
2181    ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name);
2182    if (!res.pipeline) {
2183        res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr);
2184    }
2185
2186    return res;
2187}
2188
2189ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_upscale(ggml_metal_library_t lib, const ggml_tensor * op) {
2190    assert(op->op == GGML_OP_UPSCALE);
2191
2192    char base[256];
2193    char name[256];
2194
2195    const int32_t mode_flags = ggml_get_op_params_i32(op, 0);
2196    const ggml_scale_mode mode = (ggml_scale_mode) (mode_flags & 0xFF);
2197
2198    const bool antialias = (mode_flags & GGML_SCALE_FLAG_ANTIALIAS);
2199
2200    if (mode == GGML_SCALE_MODE_BILINEAR) {
2201        snprintf(base, 256, "kernel_upscale_bilinear_%s", ggml_type_name(op->src[0]->type));
2202    } else if (mode == GGML_SCALE_MODE_BICUBIC) {
2203        snprintf(base, 256, "kernel_upscale_bicubic_%s", ggml_type_name(op->src[0]->type));
2204    } else {
2205        snprintf(base, 256, "kernel_upscale_nearest_%s", ggml_type_name(op->src[0]->type));
2206    }
2207    snprintf(name, 256, "%s_aa=%d", base, antialias);
2208
2209    ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name);
2210    if (!res.pipeline) {
2211        ggml_metal_cv_t cv = ggml_metal_cv_init();
2212
2213        ggml_metal_cv_set_bool(cv, antialias, FC_UPSCALE + 0);
2214
2215        res = ggml_metal_library_compile_pipeline(lib, base, name, cv);
2216
2217        ggml_metal_cv_free(cv);
2218    }
2219
2220    return res;
2221}
2222
2223ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_roll(ggml_metal_library_t lib, const ggml_tensor * op) {
2224    assert(op->op == GGML_OP_ROLL);
2225
2226    char base[256];
2227    char name[256];
2228
2229    snprintf(base, 256, "kernel_roll_%s", ggml_type_name(op->src[0]->type));
2230    snprintf(name, 256, "%s", base);
2231
2232    ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name);
2233    if (!res.pipeline) {
2234        res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr);
2235    }
2236
2237    return res;
2238}
2239
2240ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_pad(ggml_metal_library_t lib, const ggml_tensor * op) {
2241    assert(op->op == GGML_OP_PAD);
2242
2243    char base[256];
2244    char name[256];
2245
2246    // note: this is slower
2247    //const bool is_c4 = op->src[0]->ne[0] % 4 == 0 && op->ne[0] % 4 == 0;
2248    const bool is_c4 = false;
2249
2250    snprintf(base, 256, "kernel_pad_%s%s", ggml_type_name(op->src[0]->type), is_c4 ? "_4" : "");
2251    snprintf(name, 256, "%s", base);
2252
2253    ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name);
2254    if (res.pipeline) {
2255        return res;
2256    }
2257
2258    res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr);
2259
2260    res.c4 = is_c4;
2261
2262    return res;
2263}
2264
2265ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_pad_reflect_1d(ggml_metal_library_t lib, const ggml_tensor * op) {
2266    assert(op->op == GGML_OP_PAD_REFLECT_1D);
2267
2268    char base[256];
2269    char name[256];
2270
2271    snprintf(base, 256, "kernel_pad_reflect_1d_%s", ggml_type_name(op->src[0]->type));
2272    snprintf(name, 256, "%s", base);
2273
2274    ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name);
2275    if (!res.pipeline) {
2276        res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr);
2277    }
2278
2279    return res;
2280}
2281
2282ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_arange(ggml_metal_library_t lib, const ggml_tensor * op) {
2283    assert(op->op == GGML_OP_ARANGE);
2284
2285    char base[256];
2286    char name[256];
2287
2288    snprintf(base, 256, "kernel_arange_%s", ggml_type_name(op->type));
2289    snprintf(name, 256, "%s", base);
2290
2291    ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name);
2292    if (!res.pipeline) {
2293        res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr);
2294    }
2295
2296    return res;
2297}
2298
2299ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_timestep_embedding(ggml_metal_library_t lib, const ggml_tensor * op) {
2300    assert(op->op == GGML_OP_TIMESTEP_EMBEDDING);
2301
2302    char base[256];
2303    char name[256];
2304
2305    snprintf(base, 256, "kernel_timestep_embedding_%s", ggml_type_name(op->src[0]->type));
2306    snprintf(name, 256, "%s", base);
2307
2308    ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name);
2309    if (!res.pipeline) {
2310        res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr);
2311    }
2312
2313    return res;
2314}
2315
2316ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_opt_step_adamw(ggml_metal_library_t lib, const ggml_tensor * op) {
2317    assert(op->op == GGML_OP_OPT_STEP_ADAMW);
2318
2319    char base[256];
2320    char name[256];
2321
2322    snprintf(base, 256, "kernel_opt_step_adamw_%s", ggml_type_name(op->src[0]->type));
2323    snprintf(name, 256, "%s", base);
2324
2325    ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name);
2326    if (!res.pipeline) {
2327        res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr);
2328    }
2329
2330    return res;
2331}
2332
2333ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_opt_step_sgd(ggml_metal_library_t lib, const ggml_tensor * op) {
2334    assert(op->op == GGML_OP_OPT_STEP_SGD);
2335
2336    char base[256];
2337    char name[256];
2338
2339    snprintf(base, 256, "kernel_opt_step_sgd_%s", ggml_type_name(op->src[0]->type));
2340    snprintf(name, 256, "%s", base);
2341
2342    ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name);
2343    if (!res.pipeline) {
2344        res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr);
2345    }
2346
2347    return res;
2348}
2349
2350ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_silu_back(ggml_metal_library_t lib, const ggml_tensor * op) {
2351    assert(op->op == GGML_OP_SILU_BACK);
2352
2353    char base[256];
2354    char name[256];
2355
2356    snprintf(base, 256, "kernel_silu_back_%s", ggml_type_name(op->src[0]->type));
2357    snprintf(name, 256, "%s", base);
2358
2359    ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name);
2360    if (!res.pipeline) {
2361        res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr);
2362    }
2363
2364    return res;
2365}
2366
2367ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_memset(ggml_metal_library_t lib, const ggml_tensor *  op) {
2368    GGML_ASSERT(op->type == GGML_TYPE_I64);
2369
2370    char base[256];
2371    char name[256];
2372
2373    snprintf(base, 256, "kernel_memset_%s", ggml_type_name(op->type));
2374    snprintf(name, 256, "%s", base);
2375
2376    ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name);
2377    if (!res.pipeline) {
2378        res = ggml_metal_library_compile_pipeline(lib, base, name, nullptr);
2379    }
2380
2381    return res;
2382}
2383
2384ggml_metal_pipeline_with_params ggml_metal_library_get_pipeline_count_equal(ggml_metal_library_t lib, const ggml_tensor *  op) {
2385    assert(op->op == GGML_OP_COUNT_EQUAL);
2386
2387    GGML_TENSOR_LOCALS(int64_t, ne0, op->src[0], ne);
2388
2389    GGML_ASSERT(op->src[0]->type == op->src[1]->type);
2390    GGML_ASSERT(op->src[0]->type == GGML_TYPE_I32);
2391    GGML_ASSERT(op->type == GGML_TYPE_I64);
2392
2393    // note: the kernel only supports i32 output due to metal atomic add only supporting atomic_int
2394    GGML_ASSERT(ggml_nelements(op->src[0]) < (1LL << 31));
2395
2396    char base[256];
2397    char name[256];
2398
2399    int nsg = 1;
2400    while (32*nsg < ne00 && nsg < 32) {
2401        nsg *= 2;
2402    }
2403
2404    snprintf(base, 256, "kernel_count_equal_%s", ggml_type_name(op->src[0]->type));
2405    snprintf(name, 256, "%s_nsg=%d", base, nsg);
2406
2407    ggml_metal_pipeline_with_params res = ggml_metal_library_get_pipeline(lib, name);
2408    if (!res.pipeline) {
2409        ggml_metal_cv_t cv = ggml_metal_cv_init();
2410
2411        ggml_metal_cv_set_int16(cv, nsg, FC_COUNT_EQUAL + 0);
2412
2413        res = ggml_metal_library_compile_pipeline(lib, base, name, cv);
2414
2415        ggml_metal_cv_free(cv);
2416    }
2417
2418    res.smem = 32 * sizeof(int32_t);
2419    res.nsg  = nsg;
2420
2421    return res;
2422}