ggml-metal.m 50 KB

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