llama_cpp 0.1.4 → 0.2.0

Sign up to get free protection for your applications and to get access to all the features.
@@ -0,0 +1,783 @@
1
+ #import "ggml-metal.h"
2
+
3
+ #import "ggml.h"
4
+
5
+ #import <Foundation/Foundation.h>
6
+
7
+ #import <Metal/Metal.h>
8
+ #import <MetalPerformanceShaders/MetalPerformanceShaders.h>
9
+
10
+ #ifdef GGML_METAL_NDEBUG
11
+ #define metal_printf(...)
12
+ #else
13
+ #define metal_printf(...) fprintf(stderr, __VA_ARGS__)
14
+ #endif
15
+
16
+ #define UNUSED(x) (void)(x)
17
+
18
+ struct ggml_metal_buffer {
19
+ const char * name;
20
+
21
+ void * data;
22
+ size_t size;
23
+
24
+ id<MTLBuffer> metal;
25
+ };
26
+
27
+ struct ggml_metal_context {
28
+ float * logits;
29
+
30
+ id<MTLDevice> device;
31
+ id<MTLCommandQueue> queue;
32
+ id<MTLLibrary> library;
33
+
34
+ int n_buffers;
35
+ struct ggml_metal_buffer buffers[GGML_METAL_MAX_BUFFERS];
36
+
37
+ // custom kernels
38
+ #define GGML_METAL_DECL_KERNEL(name) \
39
+ id<MTLFunction> function_##name; \
40
+ id<MTLComputePipelineState> pipeline_##name
41
+
42
+ GGML_METAL_DECL_KERNEL(add);
43
+ GGML_METAL_DECL_KERNEL(mul);
44
+ GGML_METAL_DECL_KERNEL(mul_row); // TODO: avoid this extra kernel, instead extend the "mul" kernel to support broadcast
45
+ GGML_METAL_DECL_KERNEL(scale);
46
+ GGML_METAL_DECL_KERNEL(silu);
47
+ GGML_METAL_DECL_KERNEL(relu);
48
+ GGML_METAL_DECL_KERNEL(gelu);
49
+ GGML_METAL_DECL_KERNEL(soft_max);
50
+ GGML_METAL_DECL_KERNEL(diag_mask_inf);
51
+ GGML_METAL_DECL_KERNEL(get_rows_f16);
52
+ GGML_METAL_DECL_KERNEL(get_rows_q4_0);
53
+ GGML_METAL_DECL_KERNEL(get_rows_q4_1);
54
+ GGML_METAL_DECL_KERNEL(get_rows_q2_k);
55
+ GGML_METAL_DECL_KERNEL(get_rows_q4_k);
56
+ GGML_METAL_DECL_KERNEL(get_rows_q6_k);
57
+ GGML_METAL_DECL_KERNEL(rms_norm);
58
+ GGML_METAL_DECL_KERNEL(mul_mat_f16_f32);
59
+ GGML_METAL_DECL_KERNEL(mul_mat_q4_0_f32);
60
+ GGML_METAL_DECL_KERNEL(mul_mat_q4_1_f32);
61
+ GGML_METAL_DECL_KERNEL(mul_mat_q2_k_f32);
62
+ GGML_METAL_DECL_KERNEL(mul_mat_q4_k_f32);
63
+ GGML_METAL_DECL_KERNEL(mul_mat_q6_k_f32);
64
+ GGML_METAL_DECL_KERNEL(rope);
65
+ GGML_METAL_DECL_KERNEL(cpy_f32_f16);
66
+ GGML_METAL_DECL_KERNEL(cpy_f32_f32);
67
+
68
+ #undef GGML_METAL_DECL_KERNEL
69
+ };
70
+
71
+ // MSL code
72
+ // TODO: move the contents here when ready
73
+ // for now it is easier to work in a separate file
74
+ static NSString * const msl_library_source = @"see metal.metal";
75
+
76
+ // Here to assist with NSBundle Path Hack
77
+ @interface GGMLMetalClass : NSObject
78
+ @end
79
+ @implementation GGMLMetalClass
80
+ @end
81
+
82
+ struct ggml_metal_context * ggml_metal_init(void) {
83
+ fprintf(stderr, "%s: allocating\n", __func__);
84
+
85
+ struct ggml_metal_context * ctx = malloc(sizeof(struct ggml_metal_context));
86
+
87
+ ctx->device = MTLCreateSystemDefaultDevice();
88
+ ctx->queue = [ctx->device newCommandQueue];
89
+
90
+ // determine if we can use MPS
91
+ if (MPSSupportsMTLDevice(ctx->device)) {
92
+ fprintf(stderr, "%s: using MPS\n", __func__);
93
+ } else {
94
+ fprintf(stderr, "%s: not using MPS\n", __func__);
95
+ GGML_ASSERT(false && "MPS not supported");
96
+ }
97
+
98
+ #if 0
99
+ // compile from source string and show compile log
100
+ {
101
+ NSError * error = nil;
102
+
103
+ ctx->library = [ctx->device newLibraryWithSource:msl_library_source options:nil error:&error];
104
+ if (error) {
105
+ fprintf(stderr, "%s: error: %s\n", __func__, [[error description] UTF8String]);
106
+ exit(1);
107
+ }
108
+ }
109
+ #else
110
+ UNUSED(msl_library_source);
111
+
112
+ // read the source from "ggml-metal.metal" into a string and use newLibraryWithSource
113
+ {
114
+ NSError * error = nil;
115
+
116
+ //NSString * path = [[NSBundle mainBundle] pathForResource:@"../../examples/metal/metal" ofType:@"metal"];
117
+ NSBundle * bundle = [NSBundle bundleForClass:[GGMLMetalClass class]];
118
+ NSString * path = [bundle pathForResource:@"ggml-metal" ofType:@"metal"];
119
+ fprintf(stderr, "%s: loading '%s'\n", __func__, [path UTF8String]);
120
+
121
+ NSString * src = [NSString stringWithContentsOfFile:path encoding:NSUTF8StringEncoding error:&error];
122
+ if (error) {
123
+ fprintf(stderr, "%s: error: %s\n", __func__, [[error description] UTF8String]);
124
+ exit(1);
125
+ }
126
+
127
+ ctx->library = [ctx->device newLibraryWithSource:src options:nil error:&error];
128
+ if (error) {
129
+ fprintf(stderr, "%s: error: %s\n", __func__, [[error description] UTF8String]);
130
+ exit(1);
131
+ }
132
+ }
133
+ #endif
134
+
135
+ // load kernels
136
+ {
137
+ #define GGML_METAL_ADD_KERNEL(name) \
138
+ ctx->function_##name = [ctx->library newFunctionWithName:@"kernel_"#name]; \
139
+ ctx->pipeline_##name = [ctx->device newComputePipelineStateWithFunction:ctx->function_##name error:nil]; \
140
+ fprintf(stderr, "%s: loaded %-32s %16p\n", __func__, "kernel_"#name, (void *) ctx->pipeline_##name);
141
+
142
+ GGML_METAL_ADD_KERNEL(add);
143
+ GGML_METAL_ADD_KERNEL(mul);
144
+ GGML_METAL_ADD_KERNEL(mul_row);
145
+ GGML_METAL_ADD_KERNEL(scale);
146
+ GGML_METAL_ADD_KERNEL(silu);
147
+ GGML_METAL_ADD_KERNEL(relu);
148
+ GGML_METAL_ADD_KERNEL(gelu);
149
+ GGML_METAL_ADD_KERNEL(soft_max);
150
+ GGML_METAL_ADD_KERNEL(diag_mask_inf);
151
+ GGML_METAL_ADD_KERNEL(get_rows_f16);
152
+ GGML_METAL_ADD_KERNEL(get_rows_q4_0);
153
+ GGML_METAL_ADD_KERNEL(get_rows_q4_1);
154
+ GGML_METAL_ADD_KERNEL(get_rows_q2_k);
155
+ GGML_METAL_ADD_KERNEL(get_rows_q4_k);
156
+ GGML_METAL_ADD_KERNEL(get_rows_q6_k);
157
+ GGML_METAL_ADD_KERNEL(rms_norm);
158
+ GGML_METAL_ADD_KERNEL(mul_mat_f16_f32);
159
+ GGML_METAL_ADD_KERNEL(mul_mat_q4_0_f32);
160
+ GGML_METAL_ADD_KERNEL(mul_mat_q4_1_f32);
161
+ GGML_METAL_ADD_KERNEL(mul_mat_q2_k_f32);
162
+ GGML_METAL_ADD_KERNEL(mul_mat_q4_k_f32);
163
+ GGML_METAL_ADD_KERNEL(mul_mat_q6_k_f32);
164
+ GGML_METAL_ADD_KERNEL(rope);
165
+ GGML_METAL_ADD_KERNEL(cpy_f32_f16);
166
+ GGML_METAL_ADD_KERNEL(cpy_f32_f32);
167
+
168
+ #undef GGML_METAL_ADD_KERNEL
169
+ }
170
+
171
+ return ctx;
172
+ }
173
+
174
+ void ggml_metal_free(struct ggml_metal_context * ctx) {
175
+ fprintf(stderr, "%s: deallocating\n", __func__);
176
+
177
+ free(ctx);
178
+ }
179
+
180
+ // finds the Metal buffer that contains the tensor data on the GPU device
181
+ // the assumption is that there is 1-to-1 mapping between the host and device memory buffers, so we can find the
182
+ // Metal buffer based on the host memory pointer
183
+ //
184
+ static id<MTLBuffer> ggml_metal_get_buffer(struct ggml_metal_context * ctx, struct ggml_tensor * t, size_t * offs) {
185
+ //fprintf(stderr, "%s: data tensor '%16s', offs_data = %8ld, offs_eval = %8ld, offs_cach = %8ld\n", __func__, t->name, offs_data, offs_eval, offs_cach);
186
+
187
+ for (int i = 0; i < ctx->n_buffers; ++i) {
188
+ const int64_t ioffs = (int64_t) t->data - (int64_t) ctx->buffers[i].data;
189
+
190
+ if (ioffs >= 0 && ioffs < (int64_t) ctx->buffers[i].size) {
191
+ *offs = (size_t) ioffs;
192
+
193
+ //fprintf(stderr, "%s: '%s' tensor '%16s', offs = %8ld\n", __func__, ctx->buffers[i].name, t->name, *offs);
194
+
195
+ return ctx->buffers[i].metal;
196
+ }
197
+ }
198
+
199
+ fprintf(stderr, "%s: error: buffer is nil\n", __func__);
200
+
201
+ return nil;
202
+ }
203
+
204
+ bool ggml_metal_add_buffer(
205
+ struct ggml_metal_context * ctx,
206
+ const char * name,
207
+ void * data,
208
+ size_t size) {
209
+ if (ctx->n_buffers >= GGML_METAL_MAX_BUFFERS) {
210
+ fprintf(stderr, "%s: too many buffers\n", __func__);
211
+ return false;
212
+ }
213
+
214
+ if (data) {
215
+ // verify that the buffer does not overlap with any of the existing buffers
216
+ for (int i = 0; i < ctx->n_buffers; ++i) {
217
+ const int64_t ioffs = (int64_t) data - (int64_t) ctx->buffers[i].data;
218
+
219
+ if (ioffs >= 0 && ioffs < (int64_t) ctx->buffers[i].size) {
220
+ fprintf(stderr, "%s: error: buffer '%s' overlaps with '%s'\n", __func__, name, ctx->buffers[i].name);
221
+ return false;
222
+ }
223
+ }
224
+
225
+ size_t page_size = getpagesize();
226
+ size_t aligned_size = size;
227
+ if ((aligned_size % page_size) != 0) {
228
+ aligned_size += (page_size - (aligned_size % page_size));
229
+ }
230
+
231
+ ctx->buffers[ctx->n_buffers].name = name;
232
+ ctx->buffers[ctx->n_buffers].data = data;
233
+ ctx->buffers[ctx->n_buffers].size = size;
234
+
235
+ if (ctx->device.maxBufferLength < aligned_size) {
236
+ fprintf(stderr, "%s: buffer '%s' size %zu is larger than buffer maximum of %zu\n", __func__, name, aligned_size, ctx->device.maxBufferLength);
237
+ return false;
238
+ }
239
+ ctx->buffers[ctx->n_buffers].metal = [ctx->device newBufferWithBytesNoCopy:data length:aligned_size options:MTLResourceStorageModeShared deallocator:nil];
240
+
241
+ if (ctx->buffers[ctx->n_buffers].metal == nil) {
242
+ fprintf(stderr, "%s: failed to allocate '%-16s' buffer, size = %8.2f MB\n", __func__, name, aligned_size / 1024.0 / 1024.0);
243
+ return false;
244
+ } else {
245
+ fprintf(stderr, "%s: allocated '%-16s' buffer, size = %8.2f MB\n", __func__, name, aligned_size / 1024.0 / 1024.0);
246
+ }
247
+
248
+ ++ctx->n_buffers;
249
+ }
250
+
251
+ return true;
252
+ }
253
+
254
+ void ggml_metal_set_tensor(
255
+ struct ggml_metal_context * ctx,
256
+ struct ggml_tensor * t) {
257
+ metal_printf("%s: set input for tensor '%s'\n", __func__, t->name);
258
+
259
+ size_t offs;
260
+ id<MTLBuffer> id_dst = ggml_metal_get_buffer(ctx, t, &offs);
261
+
262
+ memcpy((void *) ((uint8_t *) id_dst.contents + offs), t->data, ggml_nbytes(t));
263
+ }
264
+
265
+ void ggml_metal_get_tensor(
266
+ struct ggml_metal_context * ctx,
267
+ struct ggml_tensor * t) {
268
+ metal_printf("%s: extract results for tensor '%s'\n", __func__, t->name);
269
+
270
+ size_t offs;
271
+ id<MTLBuffer> id_src = ggml_metal_get_buffer(ctx, t, &offs);
272
+
273
+ memcpy(t->data, (void *) ((uint8_t *) id_src.contents + offs), ggml_nbytes(t));
274
+ }
275
+
276
+ void ggml_metal_graph_compute(
277
+ struct ggml_metal_context * ctx,
278
+ struct ggml_cgraph * gf) {
279
+ metal_printf("%s: evaluating graph\n", __func__);
280
+
281
+ size_t offs_src0 = 0;
282
+ size_t offs_src1 = 0;
283
+ size_t offs_dst = 0;
284
+
285
+ id<MTLCommandBuffer> command_buffer = [ctx->queue commandBuffer];
286
+ id<MTLComputeCommandEncoder> encoder = nil;
287
+
288
+ for (int i = 0; i < gf->n_nodes; ++i) {
289
+ //metal_printf("%s: encoding node %3d, op = %8s\n", __func__, i, ggml_op_name(gf->nodes[i]->op));
290
+
291
+ struct ggml_tensor * src0 = gf->nodes[i]->src0;
292
+ struct ggml_tensor * src1 = gf->nodes[i]->src1;
293
+ struct ggml_tensor * dst = gf->nodes[i];
294
+
295
+ const int64_t ne00 = src0 ? src0->ne[0] : 0;
296
+ const int64_t ne01 = src0 ? src0->ne[1] : 0;
297
+ const int64_t ne02 = src0 ? src0->ne[2] : 0;
298
+ const int64_t ne03 = src0 ? src0->ne[3] : 0;
299
+
300
+ const uint64_t nb00 = src0 ? src0->nb[0] : 0;
301
+ const uint64_t nb01 = src0 ? src0->nb[1] : 0;
302
+ const uint64_t nb02 = src0 ? src0->nb[2] : 0;
303
+ const uint64_t nb03 = src0 ? src0->nb[3] : 0;
304
+
305
+ const int64_t ne10 = src1 ? src1->ne[0] : 0;
306
+ const int64_t ne11 = src1 ? src1->ne[1] : 0;
307
+ const int64_t ne12 = src1 ? src1->ne[2] : 0;
308
+ const int64_t ne13 = src1 ? src1->ne[3] : 0; UNUSED(ne13);
309
+
310
+ const uint64_t nb10 = src1 ? src1->nb[0] : 0;
311
+ const uint64_t nb11 = src1 ? src1->nb[1] : 0;
312
+ const uint64_t nb12 = src1 ? src1->nb[2] : 0;
313
+ const uint64_t nb13 = src1 ? src1->nb[3] : 0; UNUSED(nb13);
314
+
315
+ const int64_t ne0 = dst ? dst->ne[0] : 0;
316
+ const int64_t ne1 = dst ? dst->ne[1] : 0;
317
+ const int64_t ne2 = dst ? dst->ne[2] : 0;
318
+ const int64_t ne3 = dst ? dst->ne[3] : 0;
319
+
320
+ const uint64_t nb0 = dst ? dst->nb[0] : 0;
321
+ const uint64_t nb1 = dst ? dst->nb[1] : 0;
322
+ const uint64_t nb2 = dst ? dst->nb[2] : 0;
323
+ const uint64_t nb3 = dst ? dst->nb[3] : 0;
324
+
325
+ const enum ggml_type src0t = src0 ? src0->type : GGML_TYPE_COUNT;
326
+ const enum ggml_type src1t = src1 ? src1->type : GGML_TYPE_COUNT;
327
+ const enum ggml_type dstt = dst ? dst->type : GGML_TYPE_COUNT;
328
+
329
+ id<MTLBuffer> id_src0 = src0 ? ggml_metal_get_buffer(ctx, src0, &offs_src0) : nil;
330
+ id<MTLBuffer> id_src1 = src1 ? ggml_metal_get_buffer(ctx, src1, &offs_src1) : nil;
331
+ id<MTLBuffer> id_dst = dst ? ggml_metal_get_buffer(ctx, dst, &offs_dst) : nil;
332
+
333
+ //metal_printf("%s: op - %s\n", __func__, ggml_op_name(dst->op));
334
+ //if (src0) {
335
+ // metal_printf("%s: src0 - %4s [%5lld, %5lld, %5lld], %d, %s\n", __func__, ggml_type_name(src0t), ne00, ne01, ne02,
336
+ // ggml_is_contiguous(src0), src0->name);
337
+ //}
338
+ //if (src1) {
339
+ // metal_printf("%s: src1 - %4s [%5lld, %5lld, %5lld], %d, %s\n", __func__, ggml_type_name(src1t), ne10, ne11, ne12,
340
+ // ggml_is_contiguous(src1), src1->name);
341
+ //}
342
+ //if (dst) {
343
+ // metal_printf("%s: dst - %4s [%5lld, %5lld, %5lld], 1, %s\n", __func__, ggml_type_name(dstt), ne0, ne1, ne2,
344
+ // dst->name);
345
+ //}
346
+
347
+ switch (dst->op) {
348
+ case GGML_OP_RESHAPE:
349
+ case GGML_OP_VIEW:
350
+ case GGML_OP_TRANSPOSE:
351
+ case GGML_OP_PERMUTE:
352
+ {
353
+ // noop
354
+ } break;
355
+ case GGML_OP_ADD:
356
+ {
357
+ if (encoder == nil) {
358
+ encoder = [command_buffer computeCommandEncoder];
359
+ }
360
+
361
+ [encoder setComputePipelineState:ctx->pipeline_add];
362
+ [encoder setBuffer:id_src0 offset:offs_src0 atIndex:0];
363
+ [encoder setBuffer:id_src1 offset:offs_src1 atIndex:1];
364
+ [encoder setBuffer:id_dst offset:offs_dst atIndex:2];
365
+
366
+ const int64_t n = ggml_nelements(dst);
367
+
368
+ [encoder dispatchThreadgroups:MTLSizeMake(n, 1, 1) threadsPerThreadgroup:MTLSizeMake(1, 1, 1)];
369
+ } break;
370
+ case GGML_OP_MUL:
371
+ {
372
+ if (encoder == nil) {
373
+ encoder = [command_buffer computeCommandEncoder];
374
+ }
375
+
376
+ if (ggml_nelements(src1) == ne10) {
377
+ // src1 is a row
378
+ [encoder setComputePipelineState:ctx->pipeline_mul_row];
379
+ } else {
380
+ [encoder setComputePipelineState:ctx->pipeline_mul];
381
+ }
382
+ [encoder setBuffer:id_src0 offset:offs_src0 atIndex:0];
383
+ [encoder setBuffer:id_src1 offset:offs_src1 atIndex:1];
384
+ [encoder setBuffer:id_dst offset:offs_dst atIndex:2];
385
+ [encoder setBytes:&ne00 length:sizeof(ne00) atIndex:3];
386
+
387
+ const int64_t n = ggml_nelements(dst);
388
+
389
+ [encoder dispatchThreadgroups:MTLSizeMake(n, 1, 1) threadsPerThreadgroup:MTLSizeMake(1, 1, 1)];
390
+ } break;
391
+ case GGML_OP_SCALE:
392
+ {
393
+ if (encoder == nil) {
394
+ encoder = [command_buffer computeCommandEncoder];
395
+ }
396
+
397
+ const float scale = *(const float *) src1->data;
398
+
399
+ [encoder setComputePipelineState:ctx->pipeline_scale];
400
+ [encoder setBuffer:id_src0 offset:offs_src0 atIndex:0];
401
+ [encoder setBuffer:id_dst offset:offs_dst atIndex:1];
402
+ [encoder setBytes:&scale length:sizeof(scale) atIndex:2];
403
+
404
+ const int64_t n = ggml_nelements(dst);
405
+
406
+ [encoder dispatchThreadgroups:MTLSizeMake(n, 1, 1) threadsPerThreadgroup:MTLSizeMake(1, 1, 1)];
407
+ } break;
408
+ case GGML_OP_SILU:
409
+ {
410
+ if (encoder == nil) {
411
+ encoder = [command_buffer computeCommandEncoder];
412
+ }
413
+
414
+ [encoder setComputePipelineState:ctx->pipeline_silu];
415
+ [encoder setBuffer:id_src0 offset:offs_src0 atIndex:0];
416
+ [encoder setBuffer:id_dst offset:offs_dst atIndex:1];
417
+
418
+ const int64_t n = ggml_nelements(dst);
419
+
420
+ [encoder dispatchThreadgroups:MTLSizeMake(n, 1, 1) threadsPerThreadgroup:MTLSizeMake(1, 1, 1)];
421
+ } break;
422
+ case GGML_OP_RELU:
423
+ {
424
+ if (encoder == nil) {
425
+ encoder = [command_buffer computeCommandEncoder];
426
+ }
427
+
428
+ [encoder setComputePipelineState:ctx->pipeline_relu];
429
+ [encoder setBuffer:id_src0 offset:offs_src0 atIndex:0];
430
+ [encoder setBuffer:id_dst offset:offs_dst atIndex:1];
431
+
432
+ const int64_t n = ggml_nelements(dst);
433
+
434
+ [encoder dispatchThreadgroups:MTLSizeMake(n, 1, 1) threadsPerThreadgroup:MTLSizeMake(1, 1, 1)];
435
+ } break;
436
+ case GGML_OP_GELU:
437
+ {
438
+ if (encoder == nil) {
439
+ encoder = [command_buffer computeCommandEncoder];
440
+ }
441
+
442
+ [encoder setComputePipelineState:ctx->pipeline_gelu];
443
+ [encoder setBuffer:id_src0 offset:offs_src0 atIndex:0];
444
+ [encoder setBuffer:id_dst offset:offs_dst atIndex:1];
445
+
446
+ const int64_t n = ggml_nelements(dst);
447
+
448
+ [encoder dispatchThreadgroups:MTLSizeMake(n, 1, 1) threadsPerThreadgroup:MTLSizeMake(1, 1, 1)];
449
+ } break;
450
+ case GGML_OP_SOFT_MAX:
451
+ {
452
+ if (encoder == nil) {
453
+ encoder = [command_buffer computeCommandEncoder];
454
+ }
455
+
456
+ const int nth = 32;
457
+
458
+ [encoder setComputePipelineState:ctx->pipeline_soft_max];
459
+ [encoder setBuffer:id_src0 offset:offs_src0 atIndex:0];
460
+ [encoder setBuffer:id_dst offset:offs_dst atIndex:1];
461
+ [encoder setBytes:&ne00 length:sizeof(ne00) atIndex:2];
462
+ [encoder setBytes:&ne01 length:sizeof(ne01) atIndex:3];
463
+ [encoder setBytes:&ne02 length:sizeof(ne02) atIndex:4];
464
+ [encoder setThreadgroupMemoryLength:nth*sizeof(float) atIndex:0];
465
+
466
+ [encoder dispatchThreadgroups:MTLSizeMake(ne01, ne02, ne03) threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)];
467
+ } break;
468
+ case GGML_OP_DIAG_MASK_INF:
469
+ {
470
+ if (encoder == nil) {
471
+ encoder = [command_buffer computeCommandEncoder];
472
+ }
473
+
474
+ const int n_past = ((int32_t *)(src1->data))[0];
475
+
476
+ [encoder setComputePipelineState:ctx->pipeline_diag_mask_inf];
477
+ [encoder setBuffer:id_src0 offset:offs_src0 atIndex:0];
478
+ [encoder setBuffer:id_dst offset:offs_dst atIndex:1];
479
+ [encoder setBytes:&ne00 length:sizeof(ne00) atIndex:2];
480
+ [encoder setBytes:&ne01 length:sizeof(ne01) atIndex:3];
481
+ [encoder setBytes:&n_past length:sizeof(int) atIndex:4];
482
+
483
+ [encoder dispatchThreadgroups:MTLSizeMake(ne00, ne01, ne02) threadsPerThreadgroup:MTLSizeMake(1, 1, 1)];
484
+ } break;
485
+ case GGML_OP_MUL_MAT:
486
+ {
487
+ // TODO: needs to be updated after PR: https://github.com/ggerganov/ggml/pull/224
488
+
489
+ GGML_ASSERT(ne00 == ne10);
490
+ GGML_ASSERT(ne02 == ne12);
491
+
492
+ if (ggml_is_contiguous(src0) &&
493
+ ggml_is_contiguous(src1) &&
494
+ (src0t == GGML_TYPE_F32 || src0t == GGML_TYPE_F16) && ne11 > 1) {
495
+
496
+ if (encoder != nil) {
497
+ [encoder endEncoding];
498
+ encoder = nil;
499
+ }
500
+
501
+ MPSDataType src0dt = src0t == GGML_TYPE_F32 ? MPSDataTypeFloat32 : MPSDataTypeFloat16;
502
+ MPSDataType src1dt = src1t == GGML_TYPE_F32 ? MPSDataTypeFloat32 : MPSDataTypeFloat16;
503
+
504
+ // for F32 x F32 we use MPS
505
+ MPSMatrixDescriptor * desc0 = [MPSMatrixDescriptor
506
+ matrixDescriptorWithRows:ne01 columns:ne00 rowBytes:src0->nb[1] dataType:src0dt];
507
+
508
+ MPSMatrixDescriptor * desc1 = [MPSMatrixDescriptor
509
+ matrixDescriptorWithRows:ne11 columns:ne10 rowBytes:src1->nb[1] dataType:src1dt];
510
+
511
+ MPSMatrixDescriptor * desc = [MPSMatrixDescriptor
512
+ matrixDescriptorWithRows:ne1 columns:ne0 rowBytes:dst->nb[1] dataType:MPSDataTypeFloat32];
513
+
514
+ MPSMatrixMultiplication * mul = [[MPSMatrixMultiplication alloc]
515
+ initWithDevice:ctx->device transposeLeft:false transposeRight:true
516
+ resultRows:ne11 resultColumns:ne01 interiorColumns:ne00 alpha:1.0 beta:0.0];
517
+
518
+ // we need to do ne02 multiplications
519
+ // TODO: is there a way to do this in parallel - currently very slow ..
520
+ // TODO: might be possible to offload part of the computation to ANE using Accelerate's CBLAS
521
+ for (int64_t i02 = 0; i02 < ne02; ++i02) {
522
+ size_t offs_src0_cur = offs_src0 + i02*nb02;
523
+ size_t offs_src1_cur = offs_src1 + i02*nb12;
524
+ size_t offs_dst_cur = offs_dst + i02*nb2;
525
+
526
+ MPSMatrix * mat_src0 = [[MPSMatrix alloc] initWithBuffer:id_src0 offset:offs_src0_cur descriptor:desc0];
527
+ MPSMatrix * mat_src1 = [[MPSMatrix alloc] initWithBuffer:id_src1 offset:offs_src1_cur descriptor:desc1];
528
+ MPSMatrix * mat_dst = [[MPSMatrix alloc] initWithBuffer:id_dst offset:offs_dst_cur descriptor:desc ];
529
+
530
+ [mul encodeToCommandBuffer:command_buffer leftMatrix:mat_src1 rightMatrix:mat_src0 resultMatrix:mat_dst];
531
+ }
532
+ } else {
533
+ if (encoder == nil) {
534
+ encoder = [command_buffer computeCommandEncoder];
535
+ }
536
+
537
+ int nth0 = 32;
538
+ int nth1 = 1;
539
+
540
+ // use custom matrix x vector kernel
541
+ switch (src0t) {
542
+ case GGML_TYPE_F16:
543
+ {
544
+ GGML_ASSERT(ne02 == ne12);
545
+
546
+ nth0 = 64;
547
+ nth1 = 1;
548
+ [encoder setComputePipelineState:ctx->pipeline_mul_mat_f16_f32];
549
+ } break;
550
+ case GGML_TYPE_Q4_0:
551
+ {
552
+ GGML_ASSERT(ne02 == 1);
553
+ GGML_ASSERT(ne12 == 1);
554
+
555
+ nth0 = 8;
556
+ nth1 = 8;
557
+ [encoder setComputePipelineState:ctx->pipeline_mul_mat_q4_0_f32];
558
+ } break;
559
+ case GGML_TYPE_Q4_1:
560
+ {
561
+ GGML_ASSERT(ne02 == 1);
562
+ GGML_ASSERT(ne12 == 1);
563
+
564
+ nth0 = 8;
565
+ nth1 = 8;
566
+ [encoder setComputePipelineState:ctx->pipeline_mul_mat_q4_1_f32];
567
+ } break;
568
+ case GGML_TYPE_Q2_K:
569
+ {
570
+ GGML_ASSERT(ne02 == 1);
571
+ GGML_ASSERT(ne12 == 1);
572
+
573
+ nth0 = 4;
574
+ nth1 = 16;
575
+ [encoder setComputePipelineState:ctx->pipeline_mul_mat_q2_k_f32];
576
+ } break;
577
+ case GGML_TYPE_Q4_K:
578
+ {
579
+ GGML_ASSERT(ne02 == 1);
580
+ GGML_ASSERT(ne12 == 1);
581
+
582
+ nth0 = 4;
583
+ nth1 = 16;
584
+ [encoder setComputePipelineState:ctx->pipeline_mul_mat_q4_k_f32];
585
+ } break;
586
+ case GGML_TYPE_Q6_K:
587
+ {
588
+ GGML_ASSERT(ne02 == 1);
589
+ GGML_ASSERT(ne12 == 1);
590
+
591
+ nth0 = 4;
592
+ nth1 = 16;
593
+ [encoder setComputePipelineState:ctx->pipeline_mul_mat_q6_k_f32];
594
+ } break;
595
+ default:
596
+ {
597
+ fprintf(stderr, "Asserting on type %d\n",(int)src0t);
598
+ GGML_ASSERT(false && "not implemented");
599
+ }
600
+ };
601
+
602
+
603
+ [encoder setBuffer:id_src0 offset:offs_src0 atIndex:0];
604
+ [encoder setBuffer:id_src1 offset:offs_src1 atIndex:1];
605
+ [encoder setBuffer:id_dst offset:offs_dst atIndex:2];
606
+ [encoder setBytes:&ne00 length:sizeof(ne00) atIndex:3];
607
+ [encoder setBytes:&ne01 length:sizeof(ne01) atIndex:4];
608
+ [encoder setBytes:&nb00 length:sizeof(nb00) atIndex:5];
609
+ [encoder setBytes:&nb01 length:sizeof(nb01) atIndex:6];
610
+ [encoder setBytes:&nb02 length:sizeof(nb02) atIndex:7];
611
+ [encoder setBytes:&ne10 length:sizeof(ne10) atIndex:8];
612
+ [encoder setBytes:&ne11 length:sizeof(ne11) atIndex:9];
613
+ [encoder setBytes:&nb10 length:sizeof(nb10) atIndex:10];
614
+ [encoder setBytes:&nb11 length:sizeof(nb11) atIndex:11];
615
+ [encoder setBytes:&nb12 length:sizeof(nb12) atIndex:12];
616
+ [encoder setBytes:&ne0 length:sizeof(ne0) atIndex:13];
617
+ [encoder setBytes:&ne1 length:sizeof(ne1) atIndex:14];
618
+
619
+ if (src0t == GGML_TYPE_Q4_0 || src0t == GGML_TYPE_Q4_1) {
620
+ [encoder setThreadgroupMemoryLength:nth0*nth1*sizeof(float) atIndex:0];
621
+ [encoder dispatchThreadgroups:MTLSizeMake(ne01, ne11, 1) threadsPerThreadgroup:MTLSizeMake(nth0, nth1, 1)];
622
+ } else if (src0t == GGML_TYPE_Q2_K) {
623
+ [encoder setThreadgroupMemoryLength:nth0*nth1*sizeof(float) atIndex:0];
624
+ [encoder dispatchThreadgroups:MTLSizeMake(ne01, 1, 1) threadsPerThreadgroup:MTLSizeMake(nth0, nth1, 1)];
625
+ } else if (src0t == GGML_TYPE_Q4_K) {
626
+ [encoder setThreadgroupMemoryLength:nth0*nth1*sizeof(float) atIndex:0];
627
+ [encoder dispatchThreadgroups:MTLSizeMake(ne01, ne11, 1) threadsPerThreadgroup:MTLSizeMake(nth0, nth1, 1)];
628
+ } else if (src0t == GGML_TYPE_Q6_K) {
629
+ [encoder setThreadgroupMemoryLength:nth0*nth1*sizeof(float) atIndex:0];
630
+ [encoder dispatchThreadgroups:MTLSizeMake(ne01, ne11, 1) threadsPerThreadgroup:MTLSizeMake(nth0, nth1, 1)];
631
+ } else {
632
+ [encoder setThreadgroupMemoryLength:nth0*sizeof(float) atIndex:0];
633
+ [encoder dispatchThreadgroups:MTLSizeMake(ne01, ne11, ne12) threadsPerThreadgroup:MTLSizeMake(nth0, nth1, 1)];
634
+ }
635
+ }
636
+ } break;
637
+ case GGML_OP_GET_ROWS:
638
+ {
639
+ if (encoder == nil) {
640
+ encoder = [command_buffer computeCommandEncoder];
641
+ }
642
+
643
+ switch (src0->type) {
644
+ case GGML_TYPE_F16: [encoder setComputePipelineState:ctx->pipeline_get_rows_f16]; break;
645
+ case GGML_TYPE_Q4_0: [encoder setComputePipelineState:ctx->pipeline_get_rows_q4_0]; break;
646
+ case GGML_TYPE_Q4_1: [encoder setComputePipelineState:ctx->pipeline_get_rows_q4_1]; break;
647
+ case GGML_TYPE_Q2_K: [encoder setComputePipelineState:ctx->pipeline_get_rows_q2_k]; break;
648
+ case GGML_TYPE_Q4_K: [encoder setComputePipelineState:ctx->pipeline_get_rows_q4_k]; break;
649
+ case GGML_TYPE_Q6_K: [encoder setComputePipelineState:ctx->pipeline_get_rows_q6_k]; break;
650
+ default: GGML_ASSERT(false && "not implemented");
651
+ }
652
+
653
+ [encoder setBuffer:id_src0 offset:offs_src0 atIndex:0];
654
+ [encoder setBuffer:id_src1 offset:offs_src1 atIndex:1];
655
+ [encoder setBuffer:id_dst offset:offs_dst atIndex:2];
656
+ [encoder setBytes:&(src0->ne[0]) length:sizeof( int64_t) atIndex:3];
657
+ [encoder setBytes:&(src0->nb[1]) length:sizeof(uint64_t) atIndex:4];
658
+ [encoder setBytes:&(dst->nb[1]) length:sizeof(uint64_t) atIndex:5];
659
+
660
+ const int64_t n = ggml_nelements(src1);
661
+
662
+ [encoder dispatchThreadgroups:MTLSizeMake(n, 1, 1) threadsPerThreadgroup:MTLSizeMake(1, 1, 1)];
663
+ } break;
664
+ case GGML_OP_RMS_NORM:
665
+ {
666
+ if (encoder == nil) {
667
+ encoder = [command_buffer computeCommandEncoder];
668
+ }
669
+
670
+ const float eps = 1e-6f;
671
+
672
+ const int nth = 256;
673
+
674
+ [encoder setComputePipelineState:ctx->pipeline_rms_norm];
675
+ [encoder setBuffer:id_src0 offset:offs_src0 atIndex:0];
676
+ [encoder setBuffer:id_dst offset:offs_dst atIndex:1];
677
+ [encoder setBytes:&ne00 length:sizeof( int64_t) atIndex:2];
678
+ [encoder setBytes:&nb01 length:sizeof(uint64_t) atIndex:3];
679
+ [encoder setBytes:&eps length:sizeof( float) atIndex:4];
680
+ [encoder setThreadgroupMemoryLength:nth*sizeof(float) atIndex:0];
681
+
682
+ const int64_t nrows = ggml_nrows(src0);
683
+
684
+ [encoder dispatchThreadgroups:MTLSizeMake(nrows, 1, 1) threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)];
685
+ } break;
686
+ case GGML_OP_ROPE:
687
+ {
688
+ if (encoder == nil) {
689
+ encoder = [command_buffer computeCommandEncoder];
690
+ }
691
+
692
+ const int n_dims = ((int32_t *) src1->data)[1];
693
+ const int mode = ((int32_t *) src1->data)[2];
694
+
695
+ const int n_past = ((int32_t *)(src1->data))[0];
696
+
697
+ [encoder setComputePipelineState:ctx->pipeline_rope];
698
+ [encoder setBuffer:id_src0 offset:offs_src0 atIndex:0];
699
+ [encoder setBuffer:id_dst offset:offs_dst atIndex:1];
700
+ [encoder setBytes:&ne00 length:sizeof( int64_t) atIndex:2];
701
+ [encoder setBytes:&ne01 length:sizeof( int64_t) atIndex:3];
702
+ [encoder setBytes:&ne02 length:sizeof( int64_t) atIndex:4];
703
+ [encoder setBytes:&ne03 length:sizeof( int64_t) atIndex:5];
704
+ [encoder setBytes:&nb00 length:sizeof(uint64_t) atIndex:6];
705
+ [encoder setBytes:&nb01 length:sizeof(uint64_t) atIndex:7];
706
+ [encoder setBytes:&nb02 length:sizeof(uint64_t) atIndex:8];
707
+ [encoder setBytes:&nb03 length:sizeof(uint64_t) atIndex:9];
708
+ [encoder setBytes:&ne0 length:sizeof( int64_t) atIndex:10];
709
+ [encoder setBytes:&ne1 length:sizeof( int64_t) atIndex:11];
710
+ [encoder setBytes:&ne2 length:sizeof( int64_t) atIndex:12];
711
+ [encoder setBytes:&ne3 length:sizeof( int64_t) atIndex:13];
712
+ [encoder setBytes:&nb0 length:sizeof(uint64_t) atIndex:14];
713
+ [encoder setBytes:&nb1 length:sizeof(uint64_t) atIndex:15];
714
+ [encoder setBytes:&nb2 length:sizeof(uint64_t) atIndex:16];
715
+ [encoder setBytes:&nb3 length:sizeof(uint64_t) atIndex:17];
716
+ [encoder setBytes:&n_past length:sizeof( int) atIndex:18];
717
+ [encoder setBytes:&n_dims length:sizeof( int) atIndex:19];
718
+ [encoder setBytes:&mode length:sizeof( int) atIndex:20];
719
+
720
+ [encoder dispatchThreadgroups:MTLSizeMake(ne01, ne02, ne03) threadsPerThreadgroup:MTLSizeMake(1, 1, 1)];
721
+ } break;
722
+ case GGML_OP_CPY:
723
+ {
724
+ if (encoder == nil) {
725
+ encoder = [command_buffer computeCommandEncoder];
726
+ }
727
+
728
+ const int nth = 32;
729
+
730
+ switch (src0t) {
731
+ case GGML_TYPE_F32:
732
+ {
733
+ switch (dstt) {
734
+ case GGML_TYPE_F16: [encoder setComputePipelineState:ctx->pipeline_cpy_f32_f16]; break;
735
+ case GGML_TYPE_F32: [encoder setComputePipelineState:ctx->pipeline_cpy_f32_f32]; break;
736
+ default: GGML_ASSERT(false && "not implemented");
737
+ };
738
+ } break;
739
+ default: GGML_ASSERT(false && "not implemented");
740
+ }
741
+
742
+ [encoder setBuffer:id_src0 offset:offs_src0 atIndex:0];
743
+ [encoder setBuffer:id_dst offset:offs_dst atIndex:1];
744
+ [encoder setBytes:&ne00 length:sizeof( int64_t) atIndex:2];
745
+ [encoder setBytes:&ne01 length:sizeof( int64_t) atIndex:3];
746
+ [encoder setBytes:&ne02 length:sizeof( int64_t) atIndex:4];
747
+ [encoder setBytes:&ne03 length:sizeof( int64_t) atIndex:5];
748
+ [encoder setBytes:&nb00 length:sizeof(uint64_t) atIndex:6];
749
+ [encoder setBytes:&nb01 length:sizeof(uint64_t) atIndex:7];
750
+ [encoder setBytes:&nb02 length:sizeof(uint64_t) atIndex:8];
751
+ [encoder setBytes:&nb03 length:sizeof(uint64_t) atIndex:9];
752
+ [encoder setBytes:&ne0 length:sizeof( int64_t) atIndex:10];
753
+ [encoder setBytes:&ne1 length:sizeof( int64_t) atIndex:11];
754
+ [encoder setBytes:&ne2 length:sizeof( int64_t) atIndex:12];
755
+ [encoder setBytes:&ne3 length:sizeof( int64_t) atIndex:13];
756
+ [encoder setBytes:&nb0 length:sizeof(uint64_t) atIndex:14];
757
+ [encoder setBytes:&nb1 length:sizeof(uint64_t) atIndex:15];
758
+ [encoder setBytes:&nb2 length:sizeof(uint64_t) atIndex:16];
759
+ [encoder setBytes:&nb3 length:sizeof(uint64_t) atIndex:17];
760
+
761
+ [encoder dispatchThreadgroups:MTLSizeMake(ne01, ne02, ne03) threadsPerThreadgroup:MTLSizeMake(nth, 1, 1)];
762
+ } break;
763
+ default:
764
+ fprintf(stderr, "%s: node %3d, op = %8s not implemented\n", __func__, i, ggml_op_name(dst->op));
765
+ GGML_ASSERT(false);
766
+ }
767
+ }
768
+
769
+ if (encoder != nil) {
770
+ [encoder endEncoding];
771
+ encoder = nil;
772
+ }
773
+
774
+ [command_buffer commit];
775
+ [command_buffer waitUntilCompleted];
776
+
777
+ {
778
+ const double time_elapsed = [command_buffer GPUEndTime] - [command_buffer GPUStartTime];
779
+ UNUSED(time_elapsed);
780
+
781
+ metal_printf("%s: time elapsed = %f ms\n", __func__, time_elapsed * 1000.0);
782
+ }
783
+ }