cui-llama.rn 1.0.4 → 1.0.6

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (62) hide show
  1. package/README.md +35 -39
  2. package/android/src/main/CMakeLists.txt +11 -2
  3. package/android/src/main/java/com/rnllama/LlamaContext.java +24 -8
  4. package/android/src/main/java/com/rnllama/RNLlama.java +33 -1
  5. package/android/src/main/jni.cpp +62 -8
  6. package/android/src/newarch/java/com/rnllama/RNLlamaModule.java +5 -0
  7. package/android/src/oldarch/java/com/rnllama/RNLlamaModule.java +5 -0
  8. package/cpp/common.cpp +3237 -3231
  9. package/cpp/common.h +469 -468
  10. package/cpp/ggml-aarch64.c +2193 -2193
  11. package/cpp/ggml-aarch64.h +39 -39
  12. package/cpp/ggml-alloc.c +1036 -1042
  13. package/cpp/ggml-backend-impl.h +153 -153
  14. package/cpp/ggml-backend.c +2240 -2234
  15. package/cpp/ggml-backend.h +238 -238
  16. package/cpp/ggml-common.h +1833 -1829
  17. package/cpp/ggml-impl.h +755 -655
  18. package/cpp/ggml-metal.h +65 -65
  19. package/cpp/ggml-metal.m +3269 -3269
  20. package/cpp/ggml-quants.c +14872 -14860
  21. package/cpp/ggml-quants.h +132 -132
  22. package/cpp/ggml.c +22055 -22044
  23. package/cpp/ggml.h +2453 -2447
  24. package/cpp/llama-grammar.cpp +539 -0
  25. package/cpp/llama-grammar.h +39 -0
  26. package/cpp/llama-impl.h +26 -0
  27. package/cpp/llama-sampling.cpp +635 -0
  28. package/cpp/llama-sampling.h +56 -0
  29. package/cpp/llama-vocab.cpp +1721 -0
  30. package/cpp/llama-vocab.h +130 -0
  31. package/cpp/llama.cpp +19171 -21892
  32. package/cpp/llama.h +1240 -1217
  33. package/cpp/log.h +737 -737
  34. package/cpp/rn-llama.hpp +207 -29
  35. package/cpp/sampling.cpp +460 -460
  36. package/cpp/sgemm.cpp +1027 -1027
  37. package/cpp/sgemm.h +14 -14
  38. package/cpp/unicode.cpp +6 -0
  39. package/cpp/unicode.h +3 -0
  40. package/ios/RNLlama.mm +15 -6
  41. package/ios/RNLlamaContext.h +2 -8
  42. package/ios/RNLlamaContext.mm +41 -34
  43. package/lib/commonjs/NativeRNLlama.js.map +1 -1
  44. package/lib/commonjs/chat.js +37 -0
  45. package/lib/commonjs/chat.js.map +1 -0
  46. package/lib/commonjs/index.js +14 -1
  47. package/lib/commonjs/index.js.map +1 -1
  48. package/lib/module/NativeRNLlama.js.map +1 -1
  49. package/lib/module/chat.js +31 -0
  50. package/lib/module/chat.js.map +1 -0
  51. package/lib/module/index.js +14 -1
  52. package/lib/module/index.js.map +1 -1
  53. package/lib/typescript/NativeRNLlama.d.ts +5 -1
  54. package/lib/typescript/NativeRNLlama.d.ts.map +1 -1
  55. package/lib/typescript/chat.d.ts +10 -0
  56. package/lib/typescript/chat.d.ts.map +1 -0
  57. package/lib/typescript/index.d.ts +9 -2
  58. package/lib/typescript/index.d.ts.map +1 -1
  59. package/package.json +1 -1
  60. package/src/NativeRNLlama.ts +10 -1
  61. package/src/chat.ts +44 -0
  62. package/src/index.ts +31 -4
package/cpp/ggml-alloc.c CHANGED
@@ -1,1042 +1,1036 @@
1
- #include "ggml-alloc.h"
2
- #include "ggml-backend-impl.h"
3
- #include "ggml.h"
4
- #include "ggml-impl.h"
5
- #include <assert.h>
6
- #include <limits.h>
7
- #include <stdarg.h>
8
- #include <stdio.h>
9
- #include <stdlib.h>
10
- #include <string.h>
11
-
12
- #define MAX(a, b) ((a) > (b) ? (a) : (b))
13
- #define MAX_FREE_BLOCKS 256
14
-
15
- //#define LM_GGML_ALLOCATOR_DEBUG
16
-
17
- //#define AT_PRINTF(...) fprintf(stderr, __VA_ARGS__)
18
- #define AT_PRINTF(...)
19
-
20
-
21
- static bool lm_ggml_is_view(const struct lm_ggml_tensor * t) {
22
- return t->view_src != NULL;
23
- }
24
-
25
- static bool lm_ggml_are_same_layout(const struct lm_ggml_tensor * a, const struct lm_ggml_tensor * b) {
26
- if (a->type != b->type) {
27
- return false;
28
- }
29
- for (int i = 0; i < LM_GGML_MAX_DIMS; i++) {
30
- if (a->ne[i] != b->ne[i]) {
31
- return false;
32
- }
33
- if (a->nb[i] != b->nb[i]) {
34
- return false;
35
- }
36
- }
37
- return true;
38
- }
39
-
40
- static bool lm_ggml_op_can_inplace(enum lm_ggml_op op) {
41
- switch (op) {
42
- case LM_GGML_OP_SCALE:
43
- case LM_GGML_OP_DIAG_MASK_ZERO:
44
- case LM_GGML_OP_DIAG_MASK_INF:
45
- case LM_GGML_OP_ADD:
46
- case LM_GGML_OP_ADD1:
47
- case LM_GGML_OP_SUB:
48
- case LM_GGML_OP_MUL:
49
- case LM_GGML_OP_DIV:
50
- case LM_GGML_OP_SQR:
51
- case LM_GGML_OP_SQRT:
52
- case LM_GGML_OP_LOG:
53
- case LM_GGML_OP_UNARY:
54
- case LM_GGML_OP_ROPE:
55
- case LM_GGML_OP_RMS_NORM:
56
- case LM_GGML_OP_SOFT_MAX:
57
- return true;
58
-
59
- default:
60
- return false;
61
- }
62
- }
63
-
64
- static size_t aligned_offset(const void * buffer, size_t offset, size_t alignment) {
65
- assert(alignment && !(alignment & (alignment - 1))); // power of 2
66
- size_t align = (alignment - (((uintptr_t)buffer + offset) % alignment)) % alignment;
67
- return offset + align;
68
- }
69
-
70
- // tallocr
71
-
72
- struct lm_ggml_tallocr lm_ggml_tallocr_new(lm_ggml_backend_buffer_t buffer) {
73
- void * base = lm_ggml_backend_buffer_get_base(buffer);
74
- size_t align = lm_ggml_backend_buffer_get_alignment(buffer);
75
-
76
- assert(align && !(align & (align - 1))); // power of 2
77
-
78
- struct lm_ggml_tallocr talloc = (struct lm_ggml_tallocr) {
79
- /*.buffer = */ buffer,
80
- /*.base = */ base,
81
- /*.alignment = */ align,
82
- /*.offset = */ aligned_offset(base, 0, align),
83
- };
84
- return talloc;
85
- }
86
-
87
- void lm_ggml_tallocr_alloc(struct lm_ggml_tallocr * talloc, struct lm_ggml_tensor * tensor) {
88
- size_t size = lm_ggml_backend_buffer_get_alloc_size(talloc->buffer, tensor);
89
- size = LM_GGML_PAD(size, talloc->alignment);
90
-
91
- if (talloc->offset + size > lm_ggml_backend_buffer_get_size(talloc->buffer)) {
92
- fprintf(stderr, "%s: not enough space in the buffer to allocate %s (needed %zu, available %zu)\n",
93
- __func__, tensor->name, size, lm_ggml_backend_buffer_get_size(talloc->buffer) - talloc->offset);
94
- LM_GGML_ASSERT(!"not enough space in the buffer");
95
- return;
96
- }
97
-
98
- void * addr = (char *)lm_ggml_backend_buffer_get_base(talloc->buffer) + talloc->offset;
99
- talloc->offset += size;
100
-
101
- assert(((uintptr_t)addr % talloc->alignment) == 0);
102
-
103
- lm_ggml_backend_tensor_alloc(talloc->buffer, tensor, addr);
104
- }
105
-
106
- // dynamic tensor allocator
107
-
108
- struct free_block {
109
- size_t offset;
110
- size_t size;
111
- };
112
-
113
- struct lm_ggml_dyn_tallocr {
114
- size_t alignment;
115
- int n_free_blocks;
116
- struct free_block free_blocks[MAX_FREE_BLOCKS];
117
- size_t max_size;
118
-
119
- #ifdef LM_GGML_ALLOCATOR_DEBUG
120
- struct {
121
- const struct lm_ggml_tensor * tensor;
122
- size_t offset;
123
- } allocated_tensors[1024];
124
- #endif
125
- };
126
-
127
- #ifdef LM_GGML_ALLOCATOR_DEBUG
128
- static void add_allocated_tensor(struct lm_ggml_dyn_tallocr * alloc, size_t offset, const struct lm_ggml_tensor * tensor) {
129
- for (int i = 0; i < 1024; i++) {
130
- if (alloc->allocated_tensors[i].tensor == NULL) {
131
- alloc->allocated_tensors[i].tensor = tensor;
132
- alloc->allocated_tensors[i].offset = offset;
133
- return;
134
- }
135
- }
136
- LM_GGML_ASSERT(!"out of allocated_tensors");
137
- }
138
- static void remove_allocated_tensor(struct lm_ggml_dyn_tallocr * alloc, size_t offset, const struct lm_ggml_tensor * tensor) {
139
- for (int i = 0; i < 1024; i++) {
140
- if (alloc->allocated_tensors[i].offset == offset) {
141
- alloc->allocated_tensors[i].tensor = NULL;
142
- return;
143
- }
144
- }
145
- fprintf(stderr, "tried to free tensor %s not found\n", tensor->name);
146
- LM_GGML_ASSERT(!"tensor not found");
147
- }
148
- #endif
149
-
150
- static size_t lm_ggml_dyn_tallocr_alloc(struct lm_ggml_dyn_tallocr * alloc, size_t size, const struct lm_ggml_tensor * tensor) {
151
- size = aligned_offset(NULL, size, alloc->alignment);
152
-
153
- AT_PRINTF("%s: allocating %s (%zu bytes) - ", __func__, tensor->name, size);
154
-
155
- size_t max_avail = 0;
156
-
157
- // find the best fitting free block besides the last block
158
- int best_fit_block = -1;
159
- size_t best_fit_size = SIZE_MAX;
160
- for (int i = 0; i < alloc->n_free_blocks - 1; i++) {
161
- struct free_block * block = &alloc->free_blocks[i];
162
- max_avail = MAX(max_avail, block->size);
163
- if (block->size >= size && block->size <= best_fit_size) {
164
- best_fit_block = i;
165
- best_fit_size = block->size;
166
- }
167
- }
168
-
169
- if (best_fit_block == -1) {
170
- // the last block is our last resort
171
- struct free_block * block = &alloc->free_blocks[alloc->n_free_blocks - 1];
172
- max_avail = MAX(max_avail, block->size);
173
- if (block->size >= size) {
174
- best_fit_block = alloc->n_free_blocks - 1;
175
- } else {
176
- // this should never happen
177
- fprintf(stderr, "%s: not enough space in the buffer to allocate %zu bytes, largest block available %zu bytes\n",
178
- __func__, size, max_avail);
179
- LM_GGML_ASSERT(!"not enough space in the buffer");
180
- LM_GGML_UNREACHABLE();
181
- }
182
- }
183
-
184
- struct free_block * block = &alloc->free_blocks[best_fit_block];
185
- size_t offset = block->offset;
186
- block->offset = offset + size;
187
- block->size -= size;
188
- if (block->size == 0) {
189
- // remove block if empty
190
- alloc->n_free_blocks--;
191
- for (int j = best_fit_block; j < alloc->n_free_blocks; j++) {
192
- alloc->free_blocks[j] = alloc->free_blocks[j+1];
193
- }
194
- }
195
-
196
- AT_PRINTF("block %d, offset %zu\n", best_fit_block, offset);
197
-
198
- #ifdef LM_GGML_ALLOCATOR_DEBUG
199
- add_allocated_tensor(alloc, offset, tensor);
200
- size_t cur_max = offset + size;
201
- if (cur_max > alloc->max_size) {
202
- // sort allocated_tensors by offset
203
- for (int i = 0; i < 1024; i++) {
204
- for (int j = i + 1; j < 1024; j++) {
205
- if (alloc->allocated_tensors[i].offset > alloc->allocated_tensors[j].offset) {
206
- const struct lm_ggml_tensor * tmp_tensor = alloc->allocated_tensors[i].tensor;
207
- size_t tmp_offset = alloc->allocated_tensors[i].offset;
208
- alloc->allocated_tensors[i].tensor = alloc->allocated_tensors[j].tensor;
209
- alloc->allocated_tensors[i].offset = alloc->allocated_tensors[j].offset;
210
- alloc->allocated_tensors[j].tensor = tmp_tensor;
211
- alloc->allocated_tensors[j].offset = tmp_offset;
212
- }
213
- }
214
- }
215
- fprintf(stderr, "max_size = %.2f MB: tensors: ", cur_max / 1024.0 / 1024.0);
216
- for (int i = 0; i < 1024; i++) {
217
- if (alloc->allocated_tensors[i].tensor) {
218
- fprintf(stderr, "%s [%zx-%zx] (%.2f MB) ", alloc->allocated_tensors[i].tensor->name,
219
- alloc->allocated_tensors[i].offset,
220
- alloc->allocated_tensors[i].offset + lm_ggml_nbytes(alloc->allocated_tensors[i].tensor),
221
- lm_ggml_nbytes(alloc->allocated_tensors[i].tensor) / 1024.0 / 1024.0);
222
- }
223
- }
224
- fprintf(stderr, "\n");
225
- }
226
- #endif
227
-
228
- alloc->max_size = MAX(alloc->max_size, offset + size);
229
-
230
- return offset;
231
-
232
- LM_GGML_UNUSED(tensor);
233
- }
234
-
235
- // this is a very naive implementation, but for our case the number of free blocks should be very small
236
- static void lm_ggml_dyn_tallocr_free_tensor(struct lm_ggml_dyn_tallocr * alloc, size_t offset, size_t size, const struct lm_ggml_tensor * tensor) {
237
- size = aligned_offset(NULL, size, alloc->alignment);
238
-
239
- AT_PRINTF("%s: freeing %s at %zu (%zu bytes) - n_free_blocks = %d\n", __func__, tensor->name, offset, size, alloc->n_free_blocks);
240
-
241
- #ifdef LM_GGML_ALLOCATOR_DEBUG
242
- remove_allocated_tensor(alloc, offset, tensor);
243
- #endif
244
-
245
- // see if we can merge with an existing block
246
- for (int i = 0; i < alloc->n_free_blocks; i++) {
247
- struct free_block * block = &alloc->free_blocks[i];
248
- // check if ptr is at the end of the block
249
- if (block->offset + block->size == offset) {
250
- block->size += size;
251
- // check if we can merge with the next block
252
- if (i < alloc->n_free_blocks - 1 && block->offset + block->size == alloc->free_blocks[i+1].offset) {
253
- block->size += alloc->free_blocks[i+1].size;
254
- alloc->n_free_blocks--;
255
- for (int j = i+1; j < alloc->n_free_blocks; j++) {
256
- alloc->free_blocks[j] = alloc->free_blocks[j+1];
257
- }
258
- }
259
- return;
260
- }
261
- // check if ptr is at the beginning of the block
262
- if (offset + size == block->offset) {
263
- block->offset = offset;
264
- block->size += size;
265
- // check if we can merge with the previous block
266
- if (i > 0 && alloc->free_blocks[i-1].offset + alloc->free_blocks[i-1].size == block->offset) {
267
- alloc->free_blocks[i-1].size += block->size;
268
- alloc->n_free_blocks--;
269
- for (int j = i; j < alloc->n_free_blocks; j++) {
270
- alloc->free_blocks[j] = alloc->free_blocks[j+1];
271
- }
272
- }
273
- return;
274
- }
275
- }
276
- // otherwise, add a new block
277
- LM_GGML_ASSERT(alloc->n_free_blocks < MAX_FREE_BLOCKS && "out of free blocks");
278
- // insert the new block in the correct position to keep the array sorted by address (to make merging blocks faster)
279
- int insert_pos = 0;
280
- while (insert_pos < alloc->n_free_blocks && alloc->free_blocks[insert_pos].offset < offset) {
281
- insert_pos++;
282
- }
283
- // shift all blocks from insert_pos onward to make room for the new block
284
- for (int i = alloc->n_free_blocks; i > insert_pos; i--) {
285
- alloc->free_blocks[i] = alloc->free_blocks[i-1];
286
- }
287
- // insert the new block
288
- alloc->free_blocks[insert_pos].offset = offset;
289
- alloc->free_blocks[insert_pos].size = size;
290
- alloc->n_free_blocks++;
291
-
292
- LM_GGML_UNUSED(tensor);
293
- }
294
-
295
- static void lm_ggml_dyn_tallocr_reset(struct lm_ggml_dyn_tallocr * alloc) {
296
- alloc->n_free_blocks = 1;
297
- alloc->free_blocks[0].offset = 0;
298
- alloc->free_blocks[0].size = SIZE_MAX/2; // restrict maximum size of a measure allocator to half size_t max to avoid overflows
299
- alloc->max_size = 0;
300
- }
301
-
302
- static struct lm_ggml_dyn_tallocr * lm_ggml_dyn_tallocr_new(size_t alignment) {
303
- struct lm_ggml_dyn_tallocr * alloc = (struct lm_ggml_dyn_tallocr *)malloc(sizeof(struct lm_ggml_dyn_tallocr));
304
-
305
- *alloc = (struct lm_ggml_dyn_tallocr) {
306
- /*.alignment = */ alignment,
307
- /*.n_free_blocks = */ 0,
308
- /*.free_blocks = */ {{0}},
309
- /*.max_size = */ 0,
310
- #ifdef LM_GGML_ALLOCATOR_DEBUG
311
- /*.allocated_tensors = */ {{0}},
312
- #endif
313
- };
314
-
315
- lm_ggml_dyn_tallocr_reset(alloc);
316
-
317
- return alloc;
318
- }
319
-
320
- static void lm_ggml_dyn_tallocr_free(struct lm_ggml_dyn_tallocr * alloc) {
321
- free(alloc);
322
- }
323
-
324
- static size_t lm_ggml_dyn_tallocr_max_size(struct lm_ggml_dyn_tallocr * alloc) {
325
- return alloc->max_size;
326
- }
327
-
328
-
329
- /////////////////////////////////////
330
-
331
- // graph allocator
332
-
333
- struct hash_node {
334
- int n_children;
335
- int n_views;
336
- int buffer_id;
337
- size_t offset; // offset within the buffer
338
- bool allocated;
339
- };
340
-
341
- struct tensor_alloc {
342
- int buffer_id;
343
- size_t offset;
344
- size_t size_max; // 0 = pre-allocated, unused, or view
345
- };
346
-
347
- struct leaf_alloc {
348
- int buffer_id;
349
- struct tensor_alloc leaf;
350
- };
351
-
352
- struct node_alloc {
353
- struct tensor_alloc dst;
354
- struct tensor_alloc src[LM_GGML_MAX_SRC];
355
- };
356
-
357
- struct lm_ggml_gallocr {
358
- lm_ggml_backend_buffer_type_t * bufts; // [n_buffers]
359
- lm_ggml_backend_buffer_t * buffers; // [n_buffers]
360
- struct lm_ggml_dyn_tallocr ** buf_tallocs; // [n_buffers]
361
- int n_buffers;
362
-
363
- struct lm_ggml_hash_set hash_set;
364
- struct hash_node * hash_values; // [hash_set.size]
365
-
366
- struct node_alloc * node_allocs; // [n_nodes]
367
- int n_nodes;
368
-
369
- struct leaf_alloc * leaf_allocs; // [n_leafs]
370
- int n_leafs;
371
- };
372
-
373
- lm_ggml_gallocr_t lm_ggml_gallocr_new_n(lm_ggml_backend_buffer_type_t * bufts, int n_bufs) {
374
- lm_ggml_gallocr_t galloc = (lm_ggml_gallocr_t)calloc(1, sizeof(struct lm_ggml_gallocr));
375
- LM_GGML_ASSERT(galloc != NULL);
376
-
377
- galloc->bufts = calloc(n_bufs, sizeof(lm_ggml_backend_buffer_type_t));
378
- LM_GGML_ASSERT(galloc->bufts != NULL);
379
-
380
- galloc->buffers = calloc(n_bufs, sizeof(lm_ggml_backend_buffer_t));
381
- LM_GGML_ASSERT(galloc->buffers != NULL);
382
-
383
- galloc->buf_tallocs = calloc(n_bufs, sizeof(struct lm_ggml_dyn_tallocr *));
384
- LM_GGML_ASSERT(galloc->buf_tallocs != NULL);
385
-
386
- for (int i = 0; i < n_bufs; i++) {
387
- galloc->bufts[i] = bufts[i];
388
- galloc->buffers[i] = NULL;
389
-
390
- // check if the same buffer type is used multiple times and reuse the same allocator
391
- for (int j = 0; j < i; j++) {
392
- if (bufts[i] == bufts[j]) {
393
- galloc->buf_tallocs[i] = galloc->buf_tallocs[j];
394
- break;
395
- }
396
- }
397
-
398
- if (galloc->buf_tallocs[i] == NULL) {
399
- size_t alignment = lm_ggml_backend_buft_get_alignment(bufts[i]);
400
- galloc->buf_tallocs[i] = lm_ggml_dyn_tallocr_new(alignment);
401
- }
402
- }
403
- galloc->n_buffers = n_bufs;
404
-
405
- return galloc;
406
- }
407
-
408
- lm_ggml_gallocr_t lm_ggml_gallocr_new(lm_ggml_backend_buffer_type_t buft) {
409
- return lm_ggml_gallocr_new_n(&buft, 1);
410
- }
411
-
412
- void lm_ggml_gallocr_free(lm_ggml_gallocr_t galloc) {
413
- if (galloc == NULL) {
414
- return;
415
- }
416
-
417
- for (int i = 0; i < galloc->n_buffers; i++) {
418
- if (galloc->buffers != NULL) {
419
- // skip if already freed
420
- bool freed = false;
421
- for (int j = 0; j < i; j++) {
422
- if (galloc->buffers[j] == galloc->buffers[i]) {
423
- freed = true;
424
- break;
425
- }
426
- }
427
- if (!freed) {
428
- lm_ggml_backend_buffer_free(galloc->buffers[i]);
429
- }
430
- }
431
- if (galloc->buf_tallocs != NULL) {
432
- // skip if already freed
433
- bool freed = false;
434
- for (int j = 0; j < i; j++) {
435
- if (galloc->buf_tallocs[j] == galloc->buf_tallocs[i]) {
436
- freed = true;
437
- break;
438
- }
439
- }
440
- if (!freed) {
441
- lm_ggml_dyn_tallocr_free(galloc->buf_tallocs[i]);
442
- }
443
- }
444
- }
445
-
446
- free(galloc->hash_set.keys);
447
- free(galloc->hash_values);
448
- free(galloc->bufts);
449
- free(galloc->buffers);
450
- free(galloc->buf_tallocs);
451
- free(galloc->node_allocs);
452
- free(galloc->leaf_allocs);
453
- free(galloc);
454
- }
455
-
456
- typedef struct lm_ggml_gallocr * lm_ggml_gallocr_t;
457
-
458
- static struct hash_node * lm_ggml_gallocr_hash_get(lm_ggml_gallocr_t galloc, struct lm_ggml_tensor * t) {
459
- size_t i = lm_ggml_hash_find_or_insert(galloc->hash_set, t);
460
- return &galloc->hash_values[i];
461
- }
462
-
463
- static bool lm_ggml_gallocr_is_own(lm_ggml_gallocr_t galloc, struct lm_ggml_tensor * t) {
464
- return lm_ggml_gallocr_hash_get(galloc, t)->allocated;
465
- }
466
-
467
- static void lm_ggml_gallocr_set_node_offset(lm_ggml_gallocr_t galloc, struct lm_ggml_tensor * node, int buffer_id, size_t offset) {
468
- struct hash_node * hn = lm_ggml_gallocr_hash_get(galloc, node);
469
- hn->buffer_id = buffer_id;
470
- hn->offset = offset;
471
- hn->allocated = true;
472
- }
473
-
474
- static bool lm_ggml_gallocr_is_allocated(lm_ggml_gallocr_t galloc, struct lm_ggml_tensor * t) {
475
- return t->data != NULL || lm_ggml_gallocr_hash_get(galloc, t)->allocated;
476
- }
477
-
478
- static void lm_ggml_gallocr_allocate_node(lm_ggml_gallocr_t galloc, struct lm_ggml_tensor * node, int buffer_id) {
479
- struct hash_node * hn = lm_ggml_gallocr_hash_get(galloc, node);
480
-
481
- if (!lm_ggml_gallocr_is_allocated(galloc, node) && !lm_ggml_is_view(node)) {
482
- hn->allocated = true;
483
- assert(hn->offset == 0);
484
-
485
- // try to reuse a parent's buffer (inplace)
486
- if (lm_ggml_op_can_inplace(node->op)) {
487
- for (int i = 0; i < LM_GGML_MAX_SRC; i++) {
488
- struct lm_ggml_tensor * parent = node->src[i];
489
- if (parent == NULL) {
490
- continue;
491
- }
492
-
493
- // if the node's data is external, then we cannot re-use it
494
- if (!lm_ggml_gallocr_is_own(galloc, parent)) {
495
- AT_PRINTF("not reusing parent %s for %s as %p is external\n", parent->name, node->name, parent->data);
496
- continue;
497
- }
498
-
499
- // outputs cannot be reused
500
- if (parent->flags & LM_GGML_TENSOR_FLAG_OUTPUT || (parent->view_src != NULL && parent->view_src->flags & LM_GGML_TENSOR_FLAG_OUTPUT)) {
501
- AT_PRINTF("not reusing parent %s for %s as it is an output\n", parent->name, node->name);
502
- continue;
503
- }
504
-
505
- if (!lm_ggml_are_same_layout(node, parent)) {
506
- AT_PRINTF("not reusing parent %s for %s as layouts are different\n", parent->name, node->name);
507
- continue;
508
- }
509
-
510
- struct hash_node * p_hn = lm_ggml_gallocr_hash_get(galloc, parent);
511
- if (p_hn->n_children == 1 && p_hn->n_views == 0) {
512
- if (lm_ggml_is_view(parent)) {
513
- struct lm_ggml_tensor * view_src = parent->view_src;
514
- struct hash_node * view_src_hn = lm_ggml_gallocr_hash_get(galloc, view_src);
515
- if (view_src_hn->n_views == 1 && view_src_hn->n_children == 0 && view_src->data == parent->data) {
516
- AT_PRINTF("reusing view parent %s (%s) for %s\n", parent->name, view_src->name, node->name);
517
- assert(view_src_hn->offset == p_hn->offset);
518
- hn->buffer_id = p_hn->buffer_id;
519
- hn->offset = p_hn->offset;
520
- p_hn->allocated = false; // avoid freeing the parent
521
- view_src_hn->allocated = false;
522
- return;
523
- }
524
- } else {
525
- AT_PRINTF("reusing parent %s for %s\n", parent->name, node->name);
526
- hn->buffer_id = p_hn->buffer_id;
527
- hn->offset = p_hn->offset;
528
- p_hn->allocated = false; // avoid freeing the parent
529
- return;
530
- }
531
- }
532
- }
533
- }
534
- // allocate tensor from the buffer
535
- struct lm_ggml_dyn_tallocr * alloc = galloc->buf_tallocs[buffer_id];
536
- lm_ggml_backend_buffer_type_t buft = galloc->bufts[buffer_id];
537
- size_t size = lm_ggml_backend_buft_get_alloc_size(buft, node);
538
- size_t offset = lm_ggml_dyn_tallocr_alloc(alloc, size, node);
539
- hn->buffer_id = buffer_id;
540
- hn->offset = offset;
541
- return;
542
- }
543
- }
544
-
545
- static void lm_ggml_gallocr_free_node(lm_ggml_gallocr_t galloc, struct lm_ggml_tensor * node) {
546
- // graph outputs are never freed
547
- if (node->flags & LM_GGML_TENSOR_FLAG_OUTPUT) {
548
- AT_PRINTF("not freeing output %s\n", node->name);
549
- return;
550
- }
551
-
552
- struct hash_node * hn = lm_ggml_gallocr_hash_get(galloc, node);
553
- size_t offset = hn->offset;
554
- int buffer_id = hn->buffer_id;
555
- struct lm_ggml_dyn_tallocr * alloc = galloc->buf_tallocs[buffer_id];
556
- lm_ggml_backend_buffer_type_t buft = galloc->bufts[buffer_id];
557
- size_t size = lm_ggml_backend_buft_get_alloc_size(buft, node);
558
- lm_ggml_dyn_tallocr_free_tensor(alloc, offset, size, node);
559
- hn->allocated = false;
560
- }
561
-
562
- static int get_node_buffer_id(const int * node_buffer_ids, int i) {
563
- return node_buffer_ids ? node_buffer_ids[i] : 0;
564
- }
565
-
566
- static void lm_ggml_gallocr_alloc_graph_impl(lm_ggml_gallocr_t galloc, struct lm_ggml_cgraph * graph, const int * node_buffer_ids, const int * leaf_buffer_ids) {
567
- // clear hash tables
568
- memset(galloc->hash_set.keys, 0, galloc->hash_set.size * sizeof(struct lm_ggml_tensor *));
569
- memset(galloc->hash_values, 0, galloc->hash_set.size * sizeof(struct hash_node));
570
-
571
- // allocate leafs
572
- // these may be tensors that the application is not using in the graph, but may still want to allocate for other purposes
573
- for (int i = 0; i < graph->n_leafs; i++) {
574
- struct lm_ggml_tensor * leaf = graph->leafs[i];
575
- lm_ggml_gallocr_allocate_node(galloc, leaf, get_node_buffer_id(leaf_buffer_ids, i));
576
- }
577
-
578
- // count number of children and views
579
- // allocate other graph inputs and leafs first to avoid overwriting them
580
- for (int i = 0; i < graph->n_nodes; i++) {
581
- struct lm_ggml_tensor * node = graph->nodes[i];
582
-
583
- // TODO: better way to add external dependencies
584
- // LM_GGML_OP_NONE does not appear normally in the graph nodes, but is used by ggml-backend to add dependencies to
585
- // control when some tensors are allocated and freed. in this case, the dependencies are in `src`, but the node
586
- // itself is never used and should not be considered a dependency
587
- if (lm_ggml_is_view(node) && node->op != LM_GGML_OP_NONE) {
588
- struct lm_ggml_tensor * view_src = node->view_src;
589
- lm_ggml_gallocr_hash_get(galloc, view_src)->n_views += 1;
590
- }
591
-
592
- if (node->flags & LM_GGML_TENSOR_FLAG_INPUT) {
593
- lm_ggml_gallocr_allocate_node(galloc, graph->nodes[i], get_node_buffer_id(node_buffer_ids, i));
594
- }
595
-
596
- for (int j = 0; j < LM_GGML_MAX_SRC; j++) {
597
- struct lm_ggml_tensor * src = node->src[j];
598
- if (src == NULL) {
599
- continue;
600
- }
601
-
602
- lm_ggml_gallocr_hash_get(galloc, src)->n_children += 1;
603
-
604
- // allocate explicit inputs
605
- if (src->flags & LM_GGML_TENSOR_FLAG_INPUT) {
606
- lm_ggml_gallocr_allocate_node(galloc, src, get_node_buffer_id(node_buffer_ids, i));
607
- }
608
- }
609
- }
610
-
611
- // allocate tensors
612
- for (int i = 0; i < graph->n_nodes; i++) {
613
- struct lm_ggml_tensor * node = graph->nodes[i];
614
- int buffer_id = get_node_buffer_id(node_buffer_ids, i);
615
-
616
- // allocate parents (only leafs need to be allocated at this point)
617
- for (int j = 0; j < LM_GGML_MAX_SRC; j++) {
618
- struct lm_ggml_tensor * parent = node->src[j];
619
- if (parent == NULL) {
620
- continue;
621
- }
622
- lm_ggml_gallocr_allocate_node(galloc, parent, buffer_id);
623
- }
624
-
625
- // allocate node
626
- lm_ggml_gallocr_allocate_node(galloc, node, buffer_id);
627
-
628
- AT_PRINTF("exec: %s (%s) <= ", lm_ggml_op_desc(node), node->name);
629
- for (int j = 0; j < LM_GGML_MAX_SRC; j++) {
630
- struct lm_ggml_tensor * parent = node->src[j];
631
- if (parent == NULL) {
632
- continue;
633
- }
634
- AT_PRINTF("%s", parent->name);
635
- if (j < LM_GGML_MAX_SRC - 1 && node->src[j + 1] != NULL) {
636
- AT_PRINTF(", ");
637
- }
638
- }
639
- AT_PRINTF("\n");
640
-
641
- // update parents
642
- for (int j = 0; j < LM_GGML_MAX_SRC; j++) {
643
- struct lm_ggml_tensor * parent = node->src[j];
644
- if (parent == NULL) {
645
- continue;
646
- }
647
- struct hash_node * p_hn = lm_ggml_gallocr_hash_get(galloc, parent);
648
- p_hn->n_children -= 1;
649
-
650
- AT_PRINTF("parent %s: %d children, %d views, allocated: %d\n",
651
- parent->name, p_hn->n_children, p_hn->n_views, p_hn->allocated);
652
-
653
- if (p_hn->n_children == 0 && p_hn->n_views == 0) {
654
- if (lm_ggml_is_view(parent)) {
655
- struct lm_ggml_tensor * view_src = parent->view_src;
656
- struct hash_node * view_src_hn = lm_ggml_gallocr_hash_get(galloc, view_src);
657
- view_src_hn->n_views -= 1;
658
- AT_PRINTF("view_src %s: %d children, %d views\n",
659
- view_src->name, view_src_hn->n_children, view_src_hn->n_views);
660
- if (view_src_hn->n_views == 0 && view_src_hn->n_children == 0 && view_src_hn->allocated) {
661
- lm_ggml_gallocr_free_node(galloc, view_src);
662
- }
663
- }
664
- else if (p_hn->allocated) {
665
- lm_ggml_gallocr_free_node(galloc, parent);
666
- }
667
- }
668
- AT_PRINTF("\n");
669
- }
670
- }
671
- }
672
-
673
- bool lm_ggml_gallocr_reserve_n(lm_ggml_gallocr_t galloc, struct lm_ggml_cgraph * graph, const int * node_buffer_ids, const int * leaf_buffer_ids) {
674
- size_t hash_size = graph->visited_hash_table.size;
675
-
676
- // initialize hash table
677
- if (galloc->hash_set.size < hash_size) {
678
- free(galloc->hash_set.keys);
679
- free(galloc->hash_values);
680
- galloc->hash_set.size = hash_size;
681
- galloc->hash_set.keys = calloc(hash_size, sizeof(struct lm_ggml_tensor *));
682
- galloc->hash_values = calloc(hash_size, sizeof(struct hash_node));
683
- LM_GGML_ASSERT(galloc->hash_set.keys != NULL);
684
- LM_GGML_ASSERT(galloc->hash_values != NULL);
685
- } else {
686
- // reset hash table
687
- memset(galloc->hash_set.keys, 0, sizeof(struct lm_ggml_tensor *) * galloc->hash_set.size);
688
- memset(galloc->hash_values, 0, sizeof(struct hash_node) * galloc->hash_set.size);
689
- }
690
-
691
- // reset allocators
692
- for (int i = 0; i < galloc->n_buffers; i++) {
693
- lm_ggml_dyn_tallocr_reset(galloc->buf_tallocs[i]);
694
- }
695
-
696
- // allocate in hash table
697
- lm_ggml_gallocr_alloc_graph_impl(galloc, graph, node_buffer_ids, leaf_buffer_ids);
698
-
699
- // set the node_allocs from the hash table
700
- if (galloc->n_nodes < graph->n_nodes) {
701
- free(galloc->node_allocs);
702
- galloc->node_allocs = calloc(graph->n_nodes, sizeof(struct node_alloc));
703
- LM_GGML_ASSERT(galloc->node_allocs != NULL);
704
- }
705
- galloc->n_nodes = graph->n_nodes;
706
- for (int i = 0; i < graph->n_nodes; i++) {
707
- struct lm_ggml_tensor * node = graph->nodes[i];
708
- struct node_alloc * node_alloc = &galloc->node_allocs[i];
709
- if (node->view_src || node->data) {
710
- node_alloc->dst.buffer_id = -1;
711
- node_alloc->dst.offset = SIZE_MAX;
712
- node_alloc->dst.size_max = 0;
713
- } else {
714
- struct hash_node * hn = lm_ggml_gallocr_hash_get(galloc, node);
715
- node_alloc->dst.buffer_id = hn->buffer_id;
716
- node_alloc->dst.offset = hn->offset;
717
- node_alloc->dst.size_max = lm_ggml_backend_buft_get_alloc_size(galloc->bufts[hn->buffer_id], node);
718
- }
719
- for (int j = 0; j < LM_GGML_MAX_SRC; j++) {
720
- struct lm_ggml_tensor * src = node->src[j];
721
- if (!src || src->view_src || src->data) {
722
- node_alloc->src[j].buffer_id = -1;
723
- node_alloc->src[j].offset = SIZE_MAX;
724
- node_alloc->src[j].size_max = 0;
725
- } else {
726
- struct hash_node * hn = lm_ggml_gallocr_hash_get(galloc, src);
727
- node_alloc->src[j].buffer_id = hn->buffer_id;
728
- node_alloc->src[j].offset = hn->offset;
729
- node_alloc->src[j].size_max = lm_ggml_backend_buft_get_alloc_size(galloc->bufts[hn->buffer_id], src);
730
- }
731
- }
732
- }
733
- if (galloc->n_leafs < graph->n_leafs) {
734
- free(galloc->leaf_allocs);
735
- galloc->leaf_allocs = calloc(graph->n_leafs, sizeof(galloc->leaf_allocs[0]));
736
- LM_GGML_ASSERT(galloc->leaf_allocs != NULL);
737
- }
738
- galloc->n_leafs = graph->n_leafs;
739
- for (int i = 0; i < graph->n_leafs; i++) {
740
- struct lm_ggml_tensor * leaf = graph->leafs[i];
741
- struct hash_node * hn = lm_ggml_gallocr_hash_get(galloc, leaf);
742
- galloc->leaf_allocs[i].buffer_id = hn->buffer_id;
743
- if (leaf->view_src || leaf->data) {
744
- galloc->leaf_allocs[i].leaf.buffer_id = -1;
745
- galloc->leaf_allocs[i].leaf.offset = SIZE_MAX;
746
- galloc->leaf_allocs[i].leaf.size_max = 0;
747
- } else {
748
- galloc->leaf_allocs[i].leaf.buffer_id = hn->buffer_id;
749
- galloc->leaf_allocs[i].leaf.offset = hn->offset;
750
- galloc->leaf_allocs[i].leaf.size_max = lm_ggml_backend_buft_get_alloc_size(galloc->bufts[hn->buffer_id], leaf);
751
- }
752
- }
753
-
754
- // reallocate buffers if needed
755
- for (int i = 0; i < galloc->n_buffers; i++) {
756
- // if the buffer type is used multiple times, we reuse the same buffer
757
- for (int j = 0; j < i; j++) {
758
- if (galloc->buf_tallocs[j] == galloc->buf_tallocs[i]) {
759
- galloc->buffers[i] = galloc->buffers[j];
760
- break;
761
- }
762
- }
763
-
764
- size_t cur_size = galloc->buffers[i] ? lm_ggml_backend_buffer_get_size(galloc->buffers[i]) : 0;
765
- size_t new_size = lm_ggml_dyn_tallocr_max_size(galloc->buf_tallocs[i]);
766
-
767
- // even if there are no tensors allocated in this buffer, we still need to allocate it to initialize views
768
- if (new_size > cur_size || galloc->buffers[i] == NULL) {
769
- #ifndef NDEBUG
770
- fprintf(stderr, "%s: reallocating %s buffer from size %.02f MiB to %.02f MiB\n", __func__, lm_ggml_backend_buft_name(galloc->bufts[i]), cur_size / 1024.0 / 1024.0, new_size / 1024.0 / 1024.0);
771
- #endif
772
-
773
- lm_ggml_backend_buffer_free(galloc->buffers[i]);
774
- galloc->buffers[i] = lm_ggml_backend_buft_alloc_buffer(galloc->bufts[i], new_size);
775
- if (galloc->buffers[i] == NULL) {
776
- fprintf(stderr, "%s: failed to allocate %s buffer of size %zu\n", __func__, lm_ggml_backend_buft_name(galloc->bufts[i]), new_size);
777
- return false;
778
- }
779
- lm_ggml_backend_buffer_set_usage(galloc->buffers[i], LM_GGML_BACKEND_BUFFER_USAGE_COMPUTE);
780
- }
781
- }
782
-
783
- return true;
784
- }
785
-
786
- bool lm_ggml_gallocr_reserve(lm_ggml_gallocr_t galloc, struct lm_ggml_cgraph *graph) {
787
- return lm_ggml_gallocr_reserve_n(galloc, graph, NULL, NULL);
788
- }
789
-
790
- static void lm_ggml_gallocr_init_tensor(lm_ggml_gallocr_t galloc, struct lm_ggml_tensor * tensor, struct tensor_alloc * tensor_alloc) {
791
- int buffer_id = tensor_alloc->buffer_id;
792
- assert(tensor->data || tensor->view_src || lm_ggml_backend_buffer_get_alloc_size(galloc->buffers[buffer_id], tensor) <= tensor_alloc->size_max);
793
-
794
- if (tensor->view_src != NULL) {
795
- if (tensor->buffer == NULL) {
796
- assert(tensor_alloc->offset == SIZE_MAX);
797
- if (tensor->view_src->buffer == NULL) {
798
- // this tensor was allocated without ggml-backend
799
- return;
800
- }
801
- lm_ggml_backend_view_init(tensor);
802
- }
803
- } else {
804
- if (tensor->data == NULL) {
805
- assert(tensor_alloc->offset != SIZE_MAX);
806
- assert(lm_ggml_backend_buffer_get_alloc_size(galloc->buffers[buffer_id], tensor) <= tensor_alloc->size_max);
807
- void * base = lm_ggml_backend_buffer_get_base(galloc->buffers[buffer_id]);
808
- void * addr = (char *)base + tensor_alloc->offset;
809
- lm_ggml_backend_tensor_alloc(galloc->buffers[buffer_id], tensor, addr);
810
- } else {
811
- if (tensor->buffer == NULL) {
812
- // this tensor was allocated without ggml-backend
813
- return;
814
- }
815
- }
816
- }
817
- }
818
-
819
- static bool lm_ggml_gallocr_node_needs_realloc(lm_ggml_gallocr_t galloc, struct lm_ggml_tensor * node, struct tensor_alloc * talloc) {
820
- lm_ggml_backend_buffer_type_t buft = talloc->buffer_id != -1 ? galloc->bufts[talloc->buffer_id] : NULL;
821
- size_t node_size = (node->data || node->view_src) ? 0 : lm_ggml_backend_buft_get_alloc_size(buft, node);
822
- return talloc->size_max >= node_size;
823
- }
824
-
825
- static bool lm_ggml_gallocr_needs_realloc(lm_ggml_gallocr_t galloc, struct lm_ggml_cgraph * graph) {
826
- if (galloc->n_nodes != graph->n_nodes) {
827
- #ifndef NDEBUG
828
- fprintf(stderr, "%s: graph has different number of nodes\n", __func__);
829
- #endif
830
- return true;
831
- }
832
-
833
- if (galloc->n_leafs != graph->n_leafs) {
834
- #ifndef NDEBUG
835
- fprintf(stderr, "%s: graph has different number of leafs\n", __func__);
836
- #endif
837
- return true;
838
- }
839
-
840
- for (int i = 0; i < graph->n_nodes; i++) {
841
- struct lm_ggml_tensor * node = graph->nodes[i];
842
- struct node_alloc * node_alloc = &galloc->node_allocs[i];
843
-
844
- if (!lm_ggml_gallocr_node_needs_realloc(galloc, node, &node_alloc->dst)) {
845
- #ifndef NDEBUG
846
- fprintf(stderr, "%s: node %s is not valid\n", __func__, node->name);
847
- #endif
848
- return true;
849
- }
850
-
851
- for (int j = 0; j < LM_GGML_MAX_SRC; j++) {
852
- struct lm_ggml_tensor * src = node->src[j];
853
- if (src == NULL) {
854
- continue;
855
- }
856
- if (!lm_ggml_gallocr_node_needs_realloc(galloc, src, &node_alloc->src[j])) {
857
- #ifndef NDEBUG
858
- fprintf(stderr, "%s: src %d (%s) of node %s is not valid\n", __func__, j, src->name, node->name);
859
- #endif
860
- return true;
861
- }
862
- }
863
- }
864
-
865
- return false;
866
- }
867
-
868
- bool lm_ggml_gallocr_alloc_graph(lm_ggml_gallocr_t galloc, struct lm_ggml_cgraph * graph) {
869
- if (lm_ggml_gallocr_needs_realloc(galloc, graph)) {
870
- if (galloc->n_buffers == 1) {
871
- #ifndef NDEBUG
872
- fprintf(stderr, "%s: reallocating buffers automatically\n", __func__);
873
- #endif
874
- if (!lm_ggml_gallocr_reserve(galloc, graph)) {
875
- return false;
876
- }
877
- } else {
878
- #ifndef NDEBUG
879
- fprintf(stderr, "%s: cannot reallocate multi buffer graph automatically, call reserve\n", __func__);
880
- #endif
881
- return false;
882
- }
883
- }
884
-
885
- // reset buffers
886
- for (int i = 0; i < galloc->n_buffers; i++) {
887
- if (galloc->buffers[i] != NULL) {
888
- lm_ggml_backend_buffer_reset(galloc->buffers[i]);
889
- }
890
- }
891
-
892
- // allocate the graph tensors from the previous assignments
893
- // leafs
894
- for (int i = 0; i < graph->n_leafs; i++) {
895
- struct lm_ggml_tensor * leaf = graph->leafs[i];
896
- struct leaf_alloc * leaf_alloc = &galloc->leaf_allocs[i];
897
- lm_ggml_gallocr_init_tensor(galloc, leaf, &leaf_alloc->leaf);
898
- }
899
- // nodes
900
- for (int i = 0; i < graph->n_nodes; i++) {
901
- struct lm_ggml_tensor * node = graph->nodes[i];
902
- struct node_alloc * node_alloc = &galloc->node_allocs[i];
903
- for (int j = 0; j < LM_GGML_MAX_SRC; j++) {
904
- struct lm_ggml_tensor * src = node->src[j];
905
- if (src == NULL) {
906
- continue;
907
- }
908
- lm_ggml_gallocr_init_tensor(galloc, src, &node_alloc->src[j]);
909
- }
910
- lm_ggml_gallocr_init_tensor(galloc, node, &node_alloc->dst);
911
- }
912
-
913
- return true;
914
- }
915
-
916
- size_t lm_ggml_gallocr_get_buffer_size(lm_ggml_gallocr_t galloc, int buffer_id) {
917
- LM_GGML_ASSERT(buffer_id >= 0 && buffer_id < galloc->n_buffers);
918
-
919
- if (galloc->buffers[buffer_id] == NULL) {
920
- return 0;
921
- }
922
-
923
- for (int i = 0; i < buffer_id; i++) {
924
- if (galloc->buffers[i] == galloc->buffers[buffer_id]) {
925
- // this buffer is the same as a previous one due to the same buffer type being used multiple times
926
- // only return the buffer size the first time it appears to avoid double counting
927
- return 0;
928
- }
929
- }
930
-
931
- return lm_ggml_backend_buffer_get_size(galloc->buffers[buffer_id]);
932
- }
933
-
934
- // utils
935
-
936
- static bool alloc_tensor_range(struct lm_ggml_context * ctx,
937
- struct lm_ggml_tensor * first, struct lm_ggml_tensor * last,
938
- lm_ggml_backend_buffer_type_t buft, size_t size,
939
- lm_ggml_backend_buffer_t ** buffers, size_t * n_buffers) {
940
- lm_ggml_backend_buffer_t buffer = lm_ggml_backend_buft_alloc_buffer(buft, size);
941
- if (buffer == NULL) {
942
- #ifndef NDEBUG
943
- fprintf(stderr, "%s: failed to allocate %s buffer of size %zu\n", __func__, lm_ggml_backend_buft_name(buft), size);
944
- #endif
945
- for (size_t i = 0; i < *n_buffers; i++) {
946
- lm_ggml_backend_buffer_free((*buffers)[i]);
947
- }
948
- free(*buffers);
949
- return false;
950
- }
951
-
952
- struct lm_ggml_tallocr tallocr = lm_ggml_tallocr_new(buffer);
953
-
954
- for (struct lm_ggml_tensor * t = first; t != last; t = lm_ggml_get_next_tensor(ctx, t)) {
955
- if (t->data == NULL) {
956
- if (t->view_src == NULL) {
957
- lm_ggml_tallocr_alloc(&tallocr, t);
958
- } else if (t->buffer == NULL) {
959
- lm_ggml_backend_view_init(t);
960
- }
961
- } else {
962
- if (t->view_src != NULL && t->buffer == NULL) {
963
- // view of a pre-allocated tensor
964
- lm_ggml_backend_view_init(t);
965
- }
966
- }
967
- }
968
-
969
- *buffers = realloc(*buffers, sizeof(lm_ggml_backend_buffer_t) * (*n_buffers + 1));
970
- (*buffers)[(*n_buffers)++] = buffer;
971
-
972
- return true;
973
- }
974
-
975
- lm_ggml_backend_buffer_t lm_ggml_backend_alloc_ctx_tensors_from_buft(struct lm_ggml_context * ctx, lm_ggml_backend_buffer_type_t buft) {
976
- LM_GGML_ASSERT(lm_ggml_get_no_alloc(ctx) == true);
977
-
978
- size_t alignment = lm_ggml_backend_buft_get_alignment(buft);
979
- size_t max_size = lm_ggml_backend_buft_get_max_size(buft);
980
-
981
- lm_ggml_backend_buffer_t * buffers = NULL;
982
- size_t n_buffers = 0;
983
-
984
- size_t cur_buf_size = 0;
985
- struct lm_ggml_tensor * first = lm_ggml_get_first_tensor(ctx);
986
- for (struct lm_ggml_tensor * t = first; t != NULL; t = lm_ggml_get_next_tensor(ctx, t)) {
987
- size_t this_size = 0;
988
- if (t->data == NULL && t->view_src == NULL) {
989
- this_size = LM_GGML_PAD(lm_ggml_backend_buft_get_alloc_size(buft, t), alignment);
990
- }
991
-
992
- if (this_size > max_size) {
993
- fprintf(stderr, "%s: tensor %s is too large to fit in a %s buffer (tensor size: %zu, max buffer size: %zu)\n",
994
- __func__, t->name,
995
- lm_ggml_backend_buft_name(buft),
996
- this_size, max_size);
997
- for (size_t i = 0; i < n_buffers; i++) {
998
- lm_ggml_backend_buffer_free(buffers[i]);
999
- }
1000
- free(buffers);
1001
- return NULL;
1002
- }
1003
-
1004
- if ((cur_buf_size + this_size) > max_size) {
1005
- // allocate tensors in the current buffer
1006
- if (!alloc_tensor_range(ctx, first, t, buft, cur_buf_size, &buffers, &n_buffers)) {
1007
- return NULL;
1008
- }
1009
- first = t;
1010
- cur_buf_size = this_size;
1011
- } else {
1012
- cur_buf_size += this_size;
1013
- }
1014
- }
1015
-
1016
- // allocate remaining tensors
1017
- if (cur_buf_size > 0) {
1018
- if (!alloc_tensor_range(ctx, first, NULL, buft, cur_buf_size, &buffers, &n_buffers)) {
1019
- return NULL;
1020
- }
1021
- }
1022
-
1023
- if (n_buffers == 0) {
1024
- #ifndef NDEBUG
1025
- fprintf(stderr, "%s: all tensors in the context are already allocated\n", __func__);
1026
- #endif
1027
- return NULL;
1028
- }
1029
-
1030
- lm_ggml_backend_buffer_t buffer;
1031
- if (n_buffers == 1) {
1032
- buffer = buffers[0];
1033
- } else {
1034
- buffer = lm_ggml_backend_multi_buffer_alloc_buffer(buffers, n_buffers);
1035
- }
1036
- free(buffers);
1037
- return buffer;
1038
- }
1039
-
1040
- lm_ggml_backend_buffer_t lm_ggml_backend_alloc_ctx_tensors(struct lm_ggml_context * ctx, lm_ggml_backend_t backend) {
1041
- return lm_ggml_backend_alloc_ctx_tensors_from_buft(ctx, lm_ggml_backend_get_default_buffer_type(backend));
1042
- }
1
+ #include "ggml-alloc.h"
2
+ #include "ggml-backend-impl.h"
3
+ #include "ggml.h"
4
+ #include "ggml-impl.h"
5
+ #include <assert.h>
6
+ #include <limits.h>
7
+ #include <stdarg.h>
8
+ #include <stdio.h>
9
+ #include <stdlib.h>
10
+ #include <string.h>
11
+
12
+ #define MAX(a, b) ((a) > (b) ? (a) : (b))
13
+ #define MAX_FREE_BLOCKS 256
14
+
15
+ //#define LM_GGML_ALLOCATOR_DEBUG
16
+
17
+ //#define AT_PRINTF(...) fprintf(stderr, __VA_ARGS__)
18
+ #define AT_PRINTF(...)
19
+
20
+
21
+ static bool lm_ggml_is_view(const struct lm_ggml_tensor * t) {
22
+ return t->view_src != NULL;
23
+ }
24
+
25
+ static bool lm_ggml_are_same_layout(const struct lm_ggml_tensor * a, const struct lm_ggml_tensor * b) {
26
+ if (a->type != b->type) {
27
+ return false;
28
+ }
29
+ for (int i = 0; i < LM_GGML_MAX_DIMS; i++) {
30
+ if (a->ne[i] != b->ne[i]) {
31
+ return false;
32
+ }
33
+ if (a->nb[i] != b->nb[i]) {
34
+ return false;
35
+ }
36
+ }
37
+ return true;
38
+ }
39
+
40
+ static bool lm_ggml_op_can_inplace(enum lm_ggml_op op) {
41
+ switch (op) {
42
+ case LM_GGML_OP_SCALE:
43
+ case LM_GGML_OP_DIAG_MASK_ZERO:
44
+ case LM_GGML_OP_DIAG_MASK_INF:
45
+ case LM_GGML_OP_ADD:
46
+ case LM_GGML_OP_ADD1:
47
+ case LM_GGML_OP_SUB:
48
+ case LM_GGML_OP_MUL:
49
+ case LM_GGML_OP_DIV:
50
+ case LM_GGML_OP_SQR:
51
+ case LM_GGML_OP_SQRT:
52
+ case LM_GGML_OP_LOG:
53
+ case LM_GGML_OP_UNARY:
54
+ case LM_GGML_OP_ROPE:
55
+ case LM_GGML_OP_RMS_NORM:
56
+ case LM_GGML_OP_SOFT_MAX:
57
+ return true;
58
+
59
+ default:
60
+ return false;
61
+ }
62
+ }
63
+
64
+ static size_t aligned_offset(const void * buffer, size_t offset, size_t alignment) {
65
+ assert(alignment && !(alignment & (alignment - 1))); // power of 2
66
+ size_t align = (alignment - (((uintptr_t)buffer + offset) % alignment)) % alignment;
67
+ return offset + align;
68
+ }
69
+
70
+ // tallocr
71
+
72
+ struct lm_ggml_tallocr lm_ggml_tallocr_new(lm_ggml_backend_buffer_t buffer) {
73
+ void * base = lm_ggml_backend_buffer_get_base(buffer);
74
+ size_t align = lm_ggml_backend_buffer_get_alignment(buffer);
75
+
76
+ assert(align && !(align & (align - 1))); // power of 2
77
+
78
+ struct lm_ggml_tallocr talloc = (struct lm_ggml_tallocr) {
79
+ /*.buffer = */ buffer,
80
+ /*.base = */ base,
81
+ /*.alignment = */ align,
82
+ /*.offset = */ aligned_offset(base, 0, align),
83
+ };
84
+ return talloc;
85
+ }
86
+
87
+ void lm_ggml_tallocr_alloc(struct lm_ggml_tallocr * talloc, struct lm_ggml_tensor * tensor) {
88
+ size_t size = lm_ggml_backend_buffer_get_alloc_size(talloc->buffer, tensor);
89
+ size = LM_GGML_PAD(size, talloc->alignment);
90
+
91
+ if (talloc->offset + size > lm_ggml_backend_buffer_get_size(talloc->buffer)) {
92
+ fprintf(stderr, "%s: not enough space in the buffer to allocate %s (needed %zu, available %zu)\n",
93
+ __func__, tensor->name, size, lm_ggml_backend_buffer_get_size(talloc->buffer) - talloc->offset);
94
+ LM_GGML_ABORT("not enough space in the buffer");
95
+ }
96
+
97
+ void * addr = (char *)lm_ggml_backend_buffer_get_base(talloc->buffer) + talloc->offset;
98
+ talloc->offset += size;
99
+
100
+ assert(((uintptr_t)addr % talloc->alignment) == 0);
101
+
102
+ lm_ggml_backend_tensor_alloc(talloc->buffer, tensor, addr);
103
+ }
104
+
105
+ // dynamic tensor allocator
106
+
107
+ struct free_block {
108
+ size_t offset;
109
+ size_t size;
110
+ };
111
+
112
+ struct lm_ggml_dyn_tallocr {
113
+ size_t alignment;
114
+ int n_free_blocks;
115
+ struct free_block free_blocks[MAX_FREE_BLOCKS];
116
+ size_t max_size;
117
+
118
+ #ifdef LM_GGML_ALLOCATOR_DEBUG
119
+ struct {
120
+ const struct lm_ggml_tensor * tensor;
121
+ size_t offset;
122
+ } allocated_tensors[1024];
123
+ #endif
124
+ };
125
+
126
+ #ifdef LM_GGML_ALLOCATOR_DEBUG
127
+ static void add_allocated_tensor(struct lm_ggml_dyn_tallocr * alloc, size_t offset, const struct lm_ggml_tensor * tensor) {
128
+ for (int i = 0; i < 1024; i++) {
129
+ if (alloc->allocated_tensors[i].tensor == NULL) {
130
+ alloc->allocated_tensors[i].tensor = tensor;
131
+ alloc->allocated_tensors[i].offset = offset;
132
+ return;
133
+ }
134
+ }
135
+ LM_GGML_ABORT("out of allocated_tensors");
136
+ }
137
+ static void remove_allocated_tensor(struct lm_ggml_dyn_tallocr * alloc, size_t offset, const struct lm_ggml_tensor * tensor) {
138
+ for (int i = 0; i < 1024; i++) {
139
+ if (alloc->allocated_tensors[i].offset == offset) {
140
+ alloc->allocated_tensors[i].tensor = NULL;
141
+ return;
142
+ }
143
+ }
144
+ LM_GGML_ABORT("tried to free tensor %s not found\n", tensor->name);
145
+ }
146
+ #endif
147
+
148
+ static size_t lm_ggml_dyn_tallocr_alloc(struct lm_ggml_dyn_tallocr * alloc, size_t size, const struct lm_ggml_tensor * tensor) {
149
+ size = aligned_offset(NULL, size, alloc->alignment);
150
+
151
+ AT_PRINTF("%s: allocating %s (%zu bytes) - ", __func__, tensor->name, size);
152
+
153
+ size_t max_avail = 0;
154
+
155
+ // find the best fitting free block besides the last block
156
+ int best_fit_block = -1;
157
+ size_t best_fit_size = SIZE_MAX;
158
+ for (int i = 0; i < alloc->n_free_blocks - 1; i++) {
159
+ struct free_block * block = &alloc->free_blocks[i];
160
+ max_avail = MAX(max_avail, block->size);
161
+ if (block->size >= size && block->size <= best_fit_size) {
162
+ best_fit_block = i;
163
+ best_fit_size = block->size;
164
+ }
165
+ }
166
+
167
+ if (best_fit_block == -1) {
168
+ // the last block is our last resort
169
+ struct free_block * block = &alloc->free_blocks[alloc->n_free_blocks - 1];
170
+ max_avail = MAX(max_avail, block->size);
171
+ if (block->size >= size) {
172
+ best_fit_block = alloc->n_free_blocks - 1;
173
+ } else {
174
+ // this should never happen
175
+ fprintf(stderr, "%s: not enough space in the buffer to allocate %zu bytes, largest block available %zu bytes\n",
176
+ __func__, size, max_avail);
177
+ LM_GGML_ABORT("not enough space in the buffer");
178
+ }
179
+ }
180
+
181
+ struct free_block * block = &alloc->free_blocks[best_fit_block];
182
+ size_t offset = block->offset;
183
+ block->offset = offset + size;
184
+ block->size -= size;
185
+ if (block->size == 0) {
186
+ // remove block if empty
187
+ alloc->n_free_blocks--;
188
+ for (int j = best_fit_block; j < alloc->n_free_blocks; j++) {
189
+ alloc->free_blocks[j] = alloc->free_blocks[j+1];
190
+ }
191
+ }
192
+
193
+ AT_PRINTF("block %d, offset %zu\n", best_fit_block, offset);
194
+
195
+ #ifdef LM_GGML_ALLOCATOR_DEBUG
196
+ add_allocated_tensor(alloc, offset, tensor);
197
+ size_t cur_max = offset + size;
198
+ if (cur_max > alloc->max_size) {
199
+ // sort allocated_tensors by offset
200
+ for (int i = 0; i < 1024; i++) {
201
+ for (int j = i + 1; j < 1024; j++) {
202
+ if (alloc->allocated_tensors[i].offset > alloc->allocated_tensors[j].offset) {
203
+ const struct lm_ggml_tensor * tmp_tensor = alloc->allocated_tensors[i].tensor;
204
+ size_t tmp_offset = alloc->allocated_tensors[i].offset;
205
+ alloc->allocated_tensors[i].tensor = alloc->allocated_tensors[j].tensor;
206
+ alloc->allocated_tensors[i].offset = alloc->allocated_tensors[j].offset;
207
+ alloc->allocated_tensors[j].tensor = tmp_tensor;
208
+ alloc->allocated_tensors[j].offset = tmp_offset;
209
+ }
210
+ }
211
+ }
212
+ fprintf(stderr, "max_size = %.2f MB: tensors: ", cur_max / 1024.0 / 1024.0);
213
+ for (int i = 0; i < 1024; i++) {
214
+ if (alloc->allocated_tensors[i].tensor) {
215
+ fprintf(stderr, "%s [%zx-%zx] (%.2f MB) ", alloc->allocated_tensors[i].tensor->name,
216
+ alloc->allocated_tensors[i].offset,
217
+ alloc->allocated_tensors[i].offset + lm_ggml_nbytes(alloc->allocated_tensors[i].tensor),
218
+ lm_ggml_nbytes(alloc->allocated_tensors[i].tensor) / 1024.0 / 1024.0);
219
+ }
220
+ }
221
+ fprintf(stderr, "\n");
222
+ }
223
+ #endif
224
+
225
+ alloc->max_size = MAX(alloc->max_size, offset + size);
226
+
227
+ return offset;
228
+
229
+ LM_GGML_UNUSED(tensor);
230
+ }
231
+
232
+ // this is a very naive implementation, but for our case the number of free blocks should be very small
233
+ static void lm_ggml_dyn_tallocr_free_tensor(struct lm_ggml_dyn_tallocr * alloc, size_t offset, size_t size, const struct lm_ggml_tensor * tensor) {
234
+ size = aligned_offset(NULL, size, alloc->alignment);
235
+
236
+ AT_PRINTF("%s: freeing %s at %zu (%zu bytes) - n_free_blocks = %d\n", __func__, tensor->name, offset, size, alloc->n_free_blocks);
237
+
238
+ #ifdef LM_GGML_ALLOCATOR_DEBUG
239
+ remove_allocated_tensor(alloc, offset, tensor);
240
+ #endif
241
+
242
+ // see if we can merge with an existing block
243
+ for (int i = 0; i < alloc->n_free_blocks; i++) {
244
+ struct free_block * block = &alloc->free_blocks[i];
245
+ // check if ptr is at the end of the block
246
+ if (block->offset + block->size == offset) {
247
+ block->size += size;
248
+ // check if we can merge with the next block
249
+ if (i < alloc->n_free_blocks - 1 && block->offset + block->size == alloc->free_blocks[i+1].offset) {
250
+ block->size += alloc->free_blocks[i+1].size;
251
+ alloc->n_free_blocks--;
252
+ for (int j = i+1; j < alloc->n_free_blocks; j++) {
253
+ alloc->free_blocks[j] = alloc->free_blocks[j+1];
254
+ }
255
+ }
256
+ return;
257
+ }
258
+ // check if ptr is at the beginning of the block
259
+ if (offset + size == block->offset) {
260
+ block->offset = offset;
261
+ block->size += size;
262
+ // check if we can merge with the previous block
263
+ if (i > 0 && alloc->free_blocks[i-1].offset + alloc->free_blocks[i-1].size == block->offset) {
264
+ alloc->free_blocks[i-1].size += block->size;
265
+ alloc->n_free_blocks--;
266
+ for (int j = i; j < alloc->n_free_blocks; j++) {
267
+ alloc->free_blocks[j] = alloc->free_blocks[j+1];
268
+ }
269
+ }
270
+ return;
271
+ }
272
+ }
273
+ // otherwise, add a new block
274
+ LM_GGML_ASSERT(alloc->n_free_blocks < MAX_FREE_BLOCKS && "out of free blocks");
275
+ // insert the new block in the correct position to keep the array sorted by address (to make merging blocks faster)
276
+ int insert_pos = 0;
277
+ while (insert_pos < alloc->n_free_blocks && alloc->free_blocks[insert_pos].offset < offset) {
278
+ insert_pos++;
279
+ }
280
+ // shift all blocks from insert_pos onward to make room for the new block
281
+ for (int i = alloc->n_free_blocks; i > insert_pos; i--) {
282
+ alloc->free_blocks[i] = alloc->free_blocks[i-1];
283
+ }
284
+ // insert the new block
285
+ alloc->free_blocks[insert_pos].offset = offset;
286
+ alloc->free_blocks[insert_pos].size = size;
287
+ alloc->n_free_blocks++;
288
+
289
+ LM_GGML_UNUSED(tensor);
290
+ }
291
+
292
+ static void lm_ggml_dyn_tallocr_reset(struct lm_ggml_dyn_tallocr * alloc) {
293
+ alloc->n_free_blocks = 1;
294
+ alloc->free_blocks[0].offset = 0;
295
+ alloc->free_blocks[0].size = SIZE_MAX/2; // restrict maximum size of a measure allocator to half size_t max to avoid overflows
296
+ alloc->max_size = 0;
297
+ }
298
+
299
+ static struct lm_ggml_dyn_tallocr * lm_ggml_dyn_tallocr_new(size_t alignment) {
300
+ struct lm_ggml_dyn_tallocr * alloc = (struct lm_ggml_dyn_tallocr *)malloc(sizeof(struct lm_ggml_dyn_tallocr));
301
+
302
+ *alloc = (struct lm_ggml_dyn_tallocr) {
303
+ /*.alignment = */ alignment,
304
+ /*.n_free_blocks = */ 0,
305
+ /*.free_blocks = */ {{0}},
306
+ /*.max_size = */ 0,
307
+ #ifdef LM_GGML_ALLOCATOR_DEBUG
308
+ /*.allocated_tensors = */ {{0}},
309
+ #endif
310
+ };
311
+
312
+ lm_ggml_dyn_tallocr_reset(alloc);
313
+
314
+ return alloc;
315
+ }
316
+
317
+ static void lm_ggml_dyn_tallocr_free(struct lm_ggml_dyn_tallocr * alloc) {
318
+ free(alloc);
319
+ }
320
+
321
+ static size_t lm_ggml_dyn_tallocr_max_size(struct lm_ggml_dyn_tallocr * alloc) {
322
+ return alloc->max_size;
323
+ }
324
+
325
+
326
+ /////////////////////////////////////
327
+
328
+ // graph allocator
329
+
330
+ struct hash_node {
331
+ int n_children;
332
+ int n_views;
333
+ int buffer_id;
334
+ size_t offset; // offset within the buffer
335
+ bool allocated;
336
+ };
337
+
338
+ struct tensor_alloc {
339
+ int buffer_id;
340
+ size_t offset;
341
+ size_t size_max; // 0 = pre-allocated, unused, or view
342
+ };
343
+
344
+ struct leaf_alloc {
345
+ int buffer_id;
346
+ struct tensor_alloc leaf;
347
+ };
348
+
349
+ struct node_alloc {
350
+ struct tensor_alloc dst;
351
+ struct tensor_alloc src[LM_GGML_MAX_SRC];
352
+ };
353
+
354
+ struct lm_ggml_gallocr {
355
+ lm_ggml_backend_buffer_type_t * bufts; // [n_buffers]
356
+ lm_ggml_backend_buffer_t * buffers; // [n_buffers]
357
+ struct lm_ggml_dyn_tallocr ** buf_tallocs; // [n_buffers]
358
+ int n_buffers;
359
+
360
+ struct lm_ggml_hash_set hash_set;
361
+ struct hash_node * hash_values; // [hash_set.size]
362
+
363
+ struct node_alloc * node_allocs; // [n_nodes]
364
+ int n_nodes;
365
+
366
+ struct leaf_alloc * leaf_allocs; // [n_leafs]
367
+ int n_leafs;
368
+ };
369
+
370
+ lm_ggml_gallocr_t lm_ggml_gallocr_new_n(lm_ggml_backend_buffer_type_t * bufts, int n_bufs) {
371
+ lm_ggml_gallocr_t galloc = (lm_ggml_gallocr_t)calloc(1, sizeof(struct lm_ggml_gallocr));
372
+ LM_GGML_ASSERT(galloc != NULL);
373
+
374
+ galloc->bufts = calloc(n_bufs, sizeof(lm_ggml_backend_buffer_type_t));
375
+ LM_GGML_ASSERT(galloc->bufts != NULL);
376
+
377
+ galloc->buffers = calloc(n_bufs, sizeof(lm_ggml_backend_buffer_t));
378
+ LM_GGML_ASSERT(galloc->buffers != NULL);
379
+
380
+ galloc->buf_tallocs = calloc(n_bufs, sizeof(struct lm_ggml_dyn_tallocr *));
381
+ LM_GGML_ASSERT(galloc->buf_tallocs != NULL);
382
+
383
+ for (int i = 0; i < n_bufs; i++) {
384
+ galloc->bufts[i] = bufts[i];
385
+ galloc->buffers[i] = NULL;
386
+
387
+ // check if the same buffer type is used multiple times and reuse the same allocator
388
+ for (int j = 0; j < i; j++) {
389
+ if (bufts[i] == bufts[j]) {
390
+ galloc->buf_tallocs[i] = galloc->buf_tallocs[j];
391
+ break;
392
+ }
393
+ }
394
+
395
+ if (galloc->buf_tallocs[i] == NULL) {
396
+ size_t alignment = lm_ggml_backend_buft_get_alignment(bufts[i]);
397
+ galloc->buf_tallocs[i] = lm_ggml_dyn_tallocr_new(alignment);
398
+ }
399
+ }
400
+ galloc->n_buffers = n_bufs;
401
+
402
+ return galloc;
403
+ }
404
+
405
+ lm_ggml_gallocr_t lm_ggml_gallocr_new(lm_ggml_backend_buffer_type_t buft) {
406
+ return lm_ggml_gallocr_new_n(&buft, 1);
407
+ }
408
+
409
+ void lm_ggml_gallocr_free(lm_ggml_gallocr_t galloc) {
410
+ if (galloc == NULL) {
411
+ return;
412
+ }
413
+
414
+ for (int i = 0; i < galloc->n_buffers; i++) {
415
+ if (galloc->buffers != NULL) {
416
+ // skip if already freed
417
+ bool freed = false;
418
+ for (int j = 0; j < i; j++) {
419
+ if (galloc->buffers[j] == galloc->buffers[i]) {
420
+ freed = true;
421
+ break;
422
+ }
423
+ }
424
+ if (!freed) {
425
+ lm_ggml_backend_buffer_free(galloc->buffers[i]);
426
+ }
427
+ }
428
+ if (galloc->buf_tallocs != NULL) {
429
+ // skip if already freed
430
+ bool freed = false;
431
+ for (int j = 0; j < i; j++) {
432
+ if (galloc->buf_tallocs[j] == galloc->buf_tallocs[i]) {
433
+ freed = true;
434
+ break;
435
+ }
436
+ }
437
+ if (!freed) {
438
+ lm_ggml_dyn_tallocr_free(galloc->buf_tallocs[i]);
439
+ }
440
+ }
441
+ }
442
+
443
+ lm_ggml_hash_set_free(&galloc->hash_set);
444
+ free(galloc->hash_values);
445
+ free(galloc->bufts);
446
+ free(galloc->buffers);
447
+ free(galloc->buf_tallocs);
448
+ free(galloc->node_allocs);
449
+ free(galloc->leaf_allocs);
450
+ free(galloc);
451
+ }
452
+
453
+ typedef struct lm_ggml_gallocr * lm_ggml_gallocr_t;
454
+
455
+ static struct hash_node * lm_ggml_gallocr_hash_get(lm_ggml_gallocr_t galloc, struct lm_ggml_tensor * t) {
456
+ size_t i = lm_ggml_hash_find_or_insert(&galloc->hash_set, t);
457
+ return &galloc->hash_values[i];
458
+ }
459
+
460
+ static bool lm_ggml_gallocr_is_own(lm_ggml_gallocr_t galloc, struct lm_ggml_tensor * t) {
461
+ return lm_ggml_gallocr_hash_get(galloc, t)->allocated;
462
+ }
463
+
464
+ static void lm_ggml_gallocr_set_node_offset(lm_ggml_gallocr_t galloc, struct lm_ggml_tensor * node, int buffer_id, size_t offset) {
465
+ struct hash_node * hn = lm_ggml_gallocr_hash_get(galloc, node);
466
+ hn->buffer_id = buffer_id;
467
+ hn->offset = offset;
468
+ hn->allocated = true;
469
+ }
470
+
471
+ static bool lm_ggml_gallocr_is_allocated(lm_ggml_gallocr_t galloc, struct lm_ggml_tensor * t) {
472
+ return t->data != NULL || lm_ggml_gallocr_hash_get(galloc, t)->allocated;
473
+ }
474
+
475
+ static void lm_ggml_gallocr_allocate_node(lm_ggml_gallocr_t galloc, struct lm_ggml_tensor * node, int buffer_id) {
476
+ struct hash_node * hn = lm_ggml_gallocr_hash_get(galloc, node);
477
+
478
+ if (!lm_ggml_gallocr_is_allocated(galloc, node) && !lm_ggml_is_view(node)) {
479
+ hn->allocated = true;
480
+ assert(hn->offset == 0);
481
+
482
+ // try to reuse a parent's buffer (inplace)
483
+ if (lm_ggml_op_can_inplace(node->op)) {
484
+ for (int i = 0; i < LM_GGML_MAX_SRC; i++) {
485
+ struct lm_ggml_tensor * parent = node->src[i];
486
+ if (parent == NULL) {
487
+ continue;
488
+ }
489
+
490
+ // if the node's data is external, then we cannot re-use it
491
+ if (!lm_ggml_gallocr_is_own(galloc, parent)) {
492
+ AT_PRINTF("not reusing parent %s for %s as %p is external\n", parent->name, node->name, parent->data);
493
+ continue;
494
+ }
495
+
496
+ // outputs cannot be reused
497
+ if (parent->flags & LM_GGML_TENSOR_FLAG_OUTPUT || (parent->view_src != NULL && parent->view_src->flags & LM_GGML_TENSOR_FLAG_OUTPUT)) {
498
+ AT_PRINTF("not reusing parent %s for %s as it is an output\n", parent->name, node->name);
499
+ continue;
500
+ }
501
+
502
+ if (!lm_ggml_are_same_layout(node, parent)) {
503
+ AT_PRINTF("not reusing parent %s for %s as layouts are different\n", parent->name, node->name);
504
+ continue;
505
+ }
506
+
507
+ struct hash_node * p_hn = lm_ggml_gallocr_hash_get(galloc, parent);
508
+ if (p_hn->n_children == 1 && p_hn->n_views == 0) {
509
+ if (lm_ggml_is_view(parent)) {
510
+ struct lm_ggml_tensor * view_src = parent->view_src;
511
+ struct hash_node * view_src_hn = lm_ggml_gallocr_hash_get(galloc, view_src);
512
+ if (view_src_hn->n_views == 1 && view_src_hn->n_children == 0 && view_src->data == parent->data) {
513
+ AT_PRINTF("reusing view parent %s (%s) for %s\n", parent->name, view_src->name, node->name);
514
+ assert(view_src_hn->offset == p_hn->offset);
515
+ hn->buffer_id = p_hn->buffer_id;
516
+ hn->offset = p_hn->offset;
517
+ p_hn->allocated = false; // avoid freeing the parent
518
+ view_src_hn->allocated = false;
519
+ return;
520
+ }
521
+ } else {
522
+ AT_PRINTF("reusing parent %s for %s\n", parent->name, node->name);
523
+ hn->buffer_id = p_hn->buffer_id;
524
+ hn->offset = p_hn->offset;
525
+ p_hn->allocated = false; // avoid freeing the parent
526
+ return;
527
+ }
528
+ }
529
+ }
530
+ }
531
+ // allocate tensor from the buffer
532
+ struct lm_ggml_dyn_tallocr * alloc = galloc->buf_tallocs[buffer_id];
533
+ lm_ggml_backend_buffer_type_t buft = galloc->bufts[buffer_id];
534
+ size_t size = lm_ggml_backend_buft_get_alloc_size(buft, node);
535
+ size_t offset = lm_ggml_dyn_tallocr_alloc(alloc, size, node);
536
+ hn->buffer_id = buffer_id;
537
+ hn->offset = offset;
538
+ return;
539
+ }
540
+ }
541
+
542
+ static void lm_ggml_gallocr_free_node(lm_ggml_gallocr_t galloc, struct lm_ggml_tensor * node) {
543
+ // graph outputs are never freed
544
+ if (node->flags & LM_GGML_TENSOR_FLAG_OUTPUT) {
545
+ AT_PRINTF("not freeing output %s\n", node->name);
546
+ return;
547
+ }
548
+
549
+ struct hash_node * hn = lm_ggml_gallocr_hash_get(galloc, node);
550
+ size_t offset = hn->offset;
551
+ int buffer_id = hn->buffer_id;
552
+ struct lm_ggml_dyn_tallocr * alloc = galloc->buf_tallocs[buffer_id];
553
+ lm_ggml_backend_buffer_type_t buft = galloc->bufts[buffer_id];
554
+ size_t size = lm_ggml_backend_buft_get_alloc_size(buft, node);
555
+ lm_ggml_dyn_tallocr_free_tensor(alloc, offset, size, node);
556
+ hn->allocated = false;
557
+ }
558
+
559
+ static int get_node_buffer_id(const int * node_buffer_ids, int i) {
560
+ return node_buffer_ids ? node_buffer_ids[i] : 0;
561
+ }
562
+
563
+ static void lm_ggml_gallocr_alloc_graph_impl(lm_ggml_gallocr_t galloc, struct lm_ggml_cgraph * graph, const int * node_buffer_ids, const int * leaf_buffer_ids) {
564
+ // clear hash tables
565
+ lm_ggml_hash_set_reset(&galloc->hash_set);
566
+ memset(galloc->hash_values, 0, sizeof(struct hash_node) * galloc->hash_set.size);
567
+
568
+ // allocate leafs
569
+ // these may be tensors that the application is not using in the graph, but may still want to allocate for other purposes
570
+ for (int i = 0; i < graph->n_leafs; i++) {
571
+ struct lm_ggml_tensor * leaf = graph->leafs[i];
572
+ lm_ggml_gallocr_allocate_node(galloc, leaf, get_node_buffer_id(leaf_buffer_ids, i));
573
+ }
574
+
575
+ // count number of children and views
576
+ // allocate other graph inputs and leafs first to avoid overwriting them
577
+ for (int i = 0; i < graph->n_nodes; i++) {
578
+ struct lm_ggml_tensor * node = graph->nodes[i];
579
+
580
+ // TODO: better way to add external dependencies
581
+ // LM_GGML_OP_NONE does not appear normally in the graph nodes, but is used by ggml-backend to add dependencies to
582
+ // control when some tensors are allocated and freed. in this case, the dependencies are in `src`, but the node
583
+ // itself is never used and should not be considered a dependency
584
+ if (lm_ggml_is_view(node) && node->op != LM_GGML_OP_NONE) {
585
+ struct lm_ggml_tensor * view_src = node->view_src;
586
+ lm_ggml_gallocr_hash_get(galloc, view_src)->n_views += 1;
587
+ }
588
+
589
+ if (node->flags & LM_GGML_TENSOR_FLAG_INPUT) {
590
+ lm_ggml_gallocr_allocate_node(galloc, graph->nodes[i], get_node_buffer_id(node_buffer_ids, i));
591
+ }
592
+
593
+ for (int j = 0; j < LM_GGML_MAX_SRC; j++) {
594
+ struct lm_ggml_tensor * src = node->src[j];
595
+ if (src == NULL) {
596
+ continue;
597
+ }
598
+
599
+ lm_ggml_gallocr_hash_get(galloc, src)->n_children += 1;
600
+
601
+ // allocate explicit inputs
602
+ if (src->flags & LM_GGML_TENSOR_FLAG_INPUT) {
603
+ lm_ggml_gallocr_allocate_node(galloc, src, get_node_buffer_id(node_buffer_ids, i));
604
+ }
605
+ }
606
+ }
607
+
608
+ // allocate tensors
609
+ for (int i = 0; i < graph->n_nodes; i++) {
610
+ struct lm_ggml_tensor * node = graph->nodes[i];
611
+ int buffer_id = get_node_buffer_id(node_buffer_ids, i);
612
+
613
+ // allocate parents (only leafs need to be allocated at this point)
614
+ for (int j = 0; j < LM_GGML_MAX_SRC; j++) {
615
+ struct lm_ggml_tensor * parent = node->src[j];
616
+ if (parent == NULL) {
617
+ continue;
618
+ }
619
+ lm_ggml_gallocr_allocate_node(galloc, parent, buffer_id);
620
+ }
621
+
622
+ // allocate node
623
+ lm_ggml_gallocr_allocate_node(galloc, node, buffer_id);
624
+
625
+ AT_PRINTF("exec: %s (%s) <= ", lm_ggml_op_desc(node), node->name);
626
+ for (int j = 0; j < LM_GGML_MAX_SRC; j++) {
627
+ struct lm_ggml_tensor * parent = node->src[j];
628
+ if (parent == NULL) {
629
+ continue;
630
+ }
631
+ AT_PRINTF("%s", parent->name);
632
+ if (j < LM_GGML_MAX_SRC - 1 && node->src[j + 1] != NULL) {
633
+ AT_PRINTF(", ");
634
+ }
635
+ }
636
+ AT_PRINTF("\n");
637
+
638
+ // update parents
639
+ for (int j = 0; j < LM_GGML_MAX_SRC; j++) {
640
+ struct lm_ggml_tensor * parent = node->src[j];
641
+ if (parent == NULL) {
642
+ continue;
643
+ }
644
+ struct hash_node * p_hn = lm_ggml_gallocr_hash_get(galloc, parent);
645
+ p_hn->n_children -= 1;
646
+
647
+ AT_PRINTF("parent %s: %d children, %d views, allocated: %d\n",
648
+ parent->name, p_hn->n_children, p_hn->n_views, p_hn->allocated);
649
+
650
+ if (p_hn->n_children == 0 && p_hn->n_views == 0) {
651
+ if (lm_ggml_is_view(parent)) {
652
+ struct lm_ggml_tensor * view_src = parent->view_src;
653
+ struct hash_node * view_src_hn = lm_ggml_gallocr_hash_get(galloc, view_src);
654
+ view_src_hn->n_views -= 1;
655
+ AT_PRINTF("view_src %s: %d children, %d views\n",
656
+ view_src->name, view_src_hn->n_children, view_src_hn->n_views);
657
+ if (view_src_hn->n_views == 0 && view_src_hn->n_children == 0 && view_src_hn->allocated) {
658
+ lm_ggml_gallocr_free_node(galloc, view_src);
659
+ }
660
+ }
661
+ else if (p_hn->allocated) {
662
+ lm_ggml_gallocr_free_node(galloc, parent);
663
+ }
664
+ }
665
+ AT_PRINTF("\n");
666
+ }
667
+ }
668
+ }
669
+
670
+ bool lm_ggml_gallocr_reserve_n(lm_ggml_gallocr_t galloc, struct lm_ggml_cgraph * graph, const int * node_buffer_ids, const int * leaf_buffer_ids) {
671
+ size_t min_hash_size = graph->n_nodes + graph->n_leafs;
672
+ // add 25% margin to avoid hash collisions
673
+ min_hash_size += min_hash_size / 4;
674
+
675
+ // initialize hash table
676
+ if (galloc->hash_set.size < min_hash_size) {
677
+ lm_ggml_hash_set_free(&galloc->hash_set);
678
+ galloc->hash_set = lm_ggml_hash_set_new(min_hash_size);
679
+ LM_GGML_ASSERT(galloc->hash_set.keys != NULL);
680
+
681
+ free(galloc->hash_values);
682
+ galloc->hash_values = malloc(sizeof(struct hash_node) * galloc->hash_set.size);
683
+ LM_GGML_ASSERT(galloc->hash_values != NULL);
684
+ }
685
+
686
+ // reset allocators
687
+ for (int i = 0; i < galloc->n_buffers; i++) {
688
+ lm_ggml_dyn_tallocr_reset(galloc->buf_tallocs[i]);
689
+ }
690
+
691
+ // allocate in hash table
692
+ lm_ggml_gallocr_alloc_graph_impl(galloc, graph, node_buffer_ids, leaf_buffer_ids);
693
+
694
+ // set the node_allocs from the hash table
695
+ if (galloc->n_nodes < graph->n_nodes) {
696
+ free(galloc->node_allocs);
697
+ galloc->node_allocs = calloc(graph->n_nodes, sizeof(struct node_alloc));
698
+ LM_GGML_ASSERT(galloc->node_allocs != NULL);
699
+ }
700
+ galloc->n_nodes = graph->n_nodes;
701
+ for (int i = 0; i < graph->n_nodes; i++) {
702
+ struct lm_ggml_tensor * node = graph->nodes[i];
703
+ struct node_alloc * node_alloc = &galloc->node_allocs[i];
704
+ if (node->view_src || node->data) {
705
+ node_alloc->dst.buffer_id = -1;
706
+ node_alloc->dst.offset = SIZE_MAX;
707
+ node_alloc->dst.size_max = 0;
708
+ } else {
709
+ struct hash_node * hn = lm_ggml_gallocr_hash_get(galloc, node);
710
+ node_alloc->dst.buffer_id = hn->buffer_id;
711
+ node_alloc->dst.offset = hn->offset;
712
+ node_alloc->dst.size_max = lm_ggml_backend_buft_get_alloc_size(galloc->bufts[hn->buffer_id], node);
713
+ }
714
+ for (int j = 0; j < LM_GGML_MAX_SRC; j++) {
715
+ struct lm_ggml_tensor * src = node->src[j];
716
+ if (!src || src->view_src || src->data) {
717
+ node_alloc->src[j].buffer_id = -1;
718
+ node_alloc->src[j].offset = SIZE_MAX;
719
+ node_alloc->src[j].size_max = 0;
720
+ } else {
721
+ struct hash_node * hn = lm_ggml_gallocr_hash_get(galloc, src);
722
+ node_alloc->src[j].buffer_id = hn->buffer_id;
723
+ node_alloc->src[j].offset = hn->offset;
724
+ node_alloc->src[j].size_max = lm_ggml_backend_buft_get_alloc_size(galloc->bufts[hn->buffer_id], src);
725
+ }
726
+ }
727
+ }
728
+ if (galloc->n_leafs < graph->n_leafs) {
729
+ free(galloc->leaf_allocs);
730
+ galloc->leaf_allocs = calloc(graph->n_leafs, sizeof(galloc->leaf_allocs[0]));
731
+ LM_GGML_ASSERT(galloc->leaf_allocs != NULL);
732
+ }
733
+ galloc->n_leafs = graph->n_leafs;
734
+ for (int i = 0; i < graph->n_leafs; i++) {
735
+ struct lm_ggml_tensor * leaf = graph->leafs[i];
736
+ struct hash_node * hn = lm_ggml_gallocr_hash_get(galloc, leaf);
737
+ galloc->leaf_allocs[i].buffer_id = hn->buffer_id;
738
+ if (leaf->view_src || leaf->data) {
739
+ galloc->leaf_allocs[i].leaf.buffer_id = -1;
740
+ galloc->leaf_allocs[i].leaf.offset = SIZE_MAX;
741
+ galloc->leaf_allocs[i].leaf.size_max = 0;
742
+ } else {
743
+ galloc->leaf_allocs[i].leaf.buffer_id = hn->buffer_id;
744
+ galloc->leaf_allocs[i].leaf.offset = hn->offset;
745
+ galloc->leaf_allocs[i].leaf.size_max = lm_ggml_backend_buft_get_alloc_size(galloc->bufts[hn->buffer_id], leaf);
746
+ }
747
+ }
748
+
749
+ // reallocate buffers if needed
750
+ for (int i = 0; i < galloc->n_buffers; i++) {
751
+ // if the buffer type is used multiple times, we reuse the same buffer
752
+ for (int j = 0; j < i; j++) {
753
+ if (galloc->buf_tallocs[j] == galloc->buf_tallocs[i]) {
754
+ galloc->buffers[i] = galloc->buffers[j];
755
+ break;
756
+ }
757
+ }
758
+
759
+ size_t cur_size = galloc->buffers[i] ? lm_ggml_backend_buffer_get_size(galloc->buffers[i]) : 0;
760
+ size_t new_size = lm_ggml_dyn_tallocr_max_size(galloc->buf_tallocs[i]);
761
+
762
+ // even if there are no tensors allocated in this buffer, we still need to allocate it to initialize views
763
+ if (new_size > cur_size || galloc->buffers[i] == NULL) {
764
+ #ifndef NDEBUG
765
+ fprintf(stderr, "%s: reallocating %s buffer from size %.02f MiB to %.02f MiB\n", __func__, lm_ggml_backend_buft_name(galloc->bufts[i]), cur_size / 1024.0 / 1024.0, new_size / 1024.0 / 1024.0);
766
+ #endif
767
+
768
+ lm_ggml_backend_buffer_free(galloc->buffers[i]);
769
+ galloc->buffers[i] = lm_ggml_backend_buft_alloc_buffer(galloc->bufts[i], new_size);
770
+ if (galloc->buffers[i] == NULL) {
771
+ fprintf(stderr, "%s: failed to allocate %s buffer of size %zu\n", __func__, lm_ggml_backend_buft_name(galloc->bufts[i]), new_size);
772
+ return false;
773
+ }
774
+ lm_ggml_backend_buffer_set_usage(galloc->buffers[i], LM_GGML_BACKEND_BUFFER_USAGE_COMPUTE);
775
+ }
776
+ }
777
+
778
+ return true;
779
+ }
780
+
781
+ bool lm_ggml_gallocr_reserve(lm_ggml_gallocr_t galloc, struct lm_ggml_cgraph *graph) {
782
+ return lm_ggml_gallocr_reserve_n(galloc, graph, NULL, NULL);
783
+ }
784
+
785
+ static void lm_ggml_gallocr_init_tensor(lm_ggml_gallocr_t galloc, struct lm_ggml_tensor * tensor, struct tensor_alloc * tensor_alloc) {
786
+ int buffer_id = tensor_alloc->buffer_id;
787
+ assert(tensor->data || tensor->view_src || lm_ggml_backend_buffer_get_alloc_size(galloc->buffers[buffer_id], tensor) <= tensor_alloc->size_max);
788
+
789
+ if (tensor->view_src != NULL) {
790
+ if (tensor->buffer == NULL) {
791
+ assert(tensor_alloc->offset == SIZE_MAX);
792
+ if (tensor->view_src->buffer == NULL) {
793
+ // this tensor was allocated without ggml-backend
794
+ return;
795
+ }
796
+ lm_ggml_backend_view_init(tensor);
797
+ }
798
+ } else {
799
+ if (tensor->data == NULL) {
800
+ assert(tensor_alloc->offset != SIZE_MAX);
801
+ assert(lm_ggml_backend_buffer_get_alloc_size(galloc->buffers[buffer_id], tensor) <= tensor_alloc->size_max);
802
+ void * base = lm_ggml_backend_buffer_get_base(galloc->buffers[buffer_id]);
803
+ void * addr = (char *)base + tensor_alloc->offset;
804
+ lm_ggml_backend_tensor_alloc(galloc->buffers[buffer_id], tensor, addr);
805
+ } else {
806
+ if (tensor->buffer == NULL) {
807
+ // this tensor was allocated without ggml-backend
808
+ return;
809
+ }
810
+ }
811
+ }
812
+ }
813
+
814
+ static bool lm_ggml_gallocr_node_needs_realloc(lm_ggml_gallocr_t galloc, struct lm_ggml_tensor * node, struct tensor_alloc * talloc) {
815
+ size_t node_size = (node->data || node->view_src) ? 0 : lm_ggml_backend_buft_get_alloc_size(galloc->bufts[talloc->buffer_id], node);
816
+ return talloc->size_max >= node_size;
817
+ }
818
+
819
+ static bool lm_ggml_gallocr_needs_realloc(lm_ggml_gallocr_t galloc, struct lm_ggml_cgraph * graph) {
820
+ if (galloc->n_nodes != graph->n_nodes) {
821
+ #ifndef NDEBUG
822
+ fprintf(stderr, "%s: graph has different number of nodes\n", __func__);
823
+ #endif
824
+ return true;
825
+ }
826
+
827
+ if (galloc->n_leafs != graph->n_leafs) {
828
+ #ifndef NDEBUG
829
+ fprintf(stderr, "%s: graph has different number of leafs\n", __func__);
830
+ #endif
831
+ return true;
832
+ }
833
+
834
+ for (int i = 0; i < graph->n_nodes; i++) {
835
+ struct lm_ggml_tensor * node = graph->nodes[i];
836
+ struct node_alloc * node_alloc = &galloc->node_allocs[i];
837
+
838
+ if (!lm_ggml_gallocr_node_needs_realloc(galloc, node, &node_alloc->dst)) {
839
+ #ifndef NDEBUG
840
+ fprintf(stderr, "%s: node %s is not valid\n", __func__, node->name);
841
+ #endif
842
+ return true;
843
+ }
844
+
845
+ for (int j = 0; j < LM_GGML_MAX_SRC; j++) {
846
+ struct lm_ggml_tensor * src = node->src[j];
847
+ if (src == NULL) {
848
+ continue;
849
+ }
850
+ if (!lm_ggml_gallocr_node_needs_realloc(galloc, src, &node_alloc->src[j])) {
851
+ #ifndef NDEBUG
852
+ fprintf(stderr, "%s: src %d (%s) of node %s is not valid\n", __func__, j, src->name, node->name);
853
+ #endif
854
+ return true;
855
+ }
856
+ }
857
+ }
858
+
859
+ return false;
860
+ }
861
+
862
+ bool lm_ggml_gallocr_alloc_graph(lm_ggml_gallocr_t galloc, struct lm_ggml_cgraph * graph) {
863
+ if (lm_ggml_gallocr_needs_realloc(galloc, graph)) {
864
+ if (galloc->n_buffers == 1) {
865
+ #ifndef NDEBUG
866
+ fprintf(stderr, "%s: reallocating buffers automatically\n", __func__);
867
+ #endif
868
+ if (!lm_ggml_gallocr_reserve(galloc, graph)) {
869
+ return false;
870
+ }
871
+ } else {
872
+ #ifndef NDEBUG
873
+ fprintf(stderr, "%s: cannot reallocate multi buffer graph automatically, call reserve\n", __func__);
874
+ #endif
875
+ return false;
876
+ }
877
+ }
878
+
879
+ // reset buffers
880
+ for (int i = 0; i < galloc->n_buffers; i++) {
881
+ if (galloc->buffers[i] != NULL) {
882
+ lm_ggml_backend_buffer_reset(galloc->buffers[i]);
883
+ }
884
+ }
885
+
886
+ // allocate the graph tensors from the previous assignments
887
+ // leafs
888
+ for (int i = 0; i < graph->n_leafs; i++) {
889
+ struct lm_ggml_tensor * leaf = graph->leafs[i];
890
+ struct leaf_alloc * leaf_alloc = &galloc->leaf_allocs[i];
891
+ lm_ggml_gallocr_init_tensor(galloc, leaf, &leaf_alloc->leaf);
892
+ }
893
+ // nodes
894
+ for (int i = 0; i < graph->n_nodes; i++) {
895
+ struct lm_ggml_tensor * node = graph->nodes[i];
896
+ struct node_alloc * node_alloc = &galloc->node_allocs[i];
897
+ for (int j = 0; j < LM_GGML_MAX_SRC; j++) {
898
+ struct lm_ggml_tensor * src = node->src[j];
899
+ if (src == NULL) {
900
+ continue;
901
+ }
902
+ lm_ggml_gallocr_init_tensor(galloc, src, &node_alloc->src[j]);
903
+ }
904
+ lm_ggml_gallocr_init_tensor(galloc, node, &node_alloc->dst);
905
+ }
906
+
907
+ return true;
908
+ }
909
+
910
+ size_t lm_ggml_gallocr_get_buffer_size(lm_ggml_gallocr_t galloc, int buffer_id) {
911
+ LM_GGML_ASSERT(buffer_id >= 0 && buffer_id < galloc->n_buffers);
912
+
913
+ if (galloc->buffers[buffer_id] == NULL) {
914
+ return 0;
915
+ }
916
+
917
+ for (int i = 0; i < buffer_id; i++) {
918
+ if (galloc->buffers[i] == galloc->buffers[buffer_id]) {
919
+ // this buffer is the same as a previous one due to the same buffer type being used multiple times
920
+ // only return the buffer size the first time it appears to avoid double counting
921
+ return 0;
922
+ }
923
+ }
924
+
925
+ return lm_ggml_backend_buffer_get_size(galloc->buffers[buffer_id]);
926
+ }
927
+
928
+ // utils
929
+
930
+ static bool alloc_tensor_range(struct lm_ggml_context * ctx,
931
+ struct lm_ggml_tensor * first, struct lm_ggml_tensor * last,
932
+ lm_ggml_backend_buffer_type_t buft, size_t size,
933
+ lm_ggml_backend_buffer_t ** buffers, size_t * n_buffers) {
934
+ lm_ggml_backend_buffer_t buffer = lm_ggml_backend_buft_alloc_buffer(buft, size);
935
+ if (buffer == NULL) {
936
+ #ifndef NDEBUG
937
+ fprintf(stderr, "%s: failed to allocate %s buffer of size %zu\n", __func__, lm_ggml_backend_buft_name(buft), size);
938
+ #endif
939
+ for (size_t i = 0; i < *n_buffers; i++) {
940
+ lm_ggml_backend_buffer_free((*buffers)[i]);
941
+ }
942
+ free(*buffers);
943
+ return false;
944
+ }
945
+
946
+ struct lm_ggml_tallocr tallocr = lm_ggml_tallocr_new(buffer);
947
+
948
+ for (struct lm_ggml_tensor * t = first; t != last; t = lm_ggml_get_next_tensor(ctx, t)) {
949
+ if (t->data == NULL) {
950
+ if (t->view_src == NULL) {
951
+ lm_ggml_tallocr_alloc(&tallocr, t);
952
+ } else if (t->buffer == NULL) {
953
+ lm_ggml_backend_view_init(t);
954
+ }
955
+ } else {
956
+ if (t->view_src != NULL && t->buffer == NULL) {
957
+ // view of a pre-allocated tensor
958
+ lm_ggml_backend_view_init(t);
959
+ }
960
+ }
961
+ }
962
+
963
+ *buffers = realloc(*buffers, sizeof(lm_ggml_backend_buffer_t) * (*n_buffers + 1));
964
+ (*buffers)[(*n_buffers)++] = buffer;
965
+
966
+ return true;
967
+ }
968
+
969
+ lm_ggml_backend_buffer_t lm_ggml_backend_alloc_ctx_tensors_from_buft(struct lm_ggml_context * ctx, lm_ggml_backend_buffer_type_t buft) {
970
+ LM_GGML_ASSERT(lm_ggml_get_no_alloc(ctx) == true);
971
+
972
+ size_t alignment = lm_ggml_backend_buft_get_alignment(buft);
973
+ size_t max_size = lm_ggml_backend_buft_get_max_size(buft);
974
+
975
+ lm_ggml_backend_buffer_t * buffers = NULL;
976
+ size_t n_buffers = 0;
977
+
978
+ size_t cur_buf_size = 0;
979
+ struct lm_ggml_tensor * first = lm_ggml_get_first_tensor(ctx);
980
+ for (struct lm_ggml_tensor * t = first; t != NULL; t = lm_ggml_get_next_tensor(ctx, t)) {
981
+ size_t this_size = 0;
982
+ if (t->data == NULL && t->view_src == NULL) {
983
+ this_size = LM_GGML_PAD(lm_ggml_backend_buft_get_alloc_size(buft, t), alignment);
984
+ }
985
+
986
+ if (this_size > max_size) {
987
+ fprintf(stderr, "%s: tensor %s is too large to fit in a %s buffer (tensor size: %zu, max buffer size: %zu)\n",
988
+ __func__, t->name,
989
+ lm_ggml_backend_buft_name(buft),
990
+ this_size, max_size);
991
+ for (size_t i = 0; i < n_buffers; i++) {
992
+ lm_ggml_backend_buffer_free(buffers[i]);
993
+ }
994
+ free(buffers);
995
+ return NULL;
996
+ }
997
+
998
+ if ((cur_buf_size + this_size) > max_size) {
999
+ // allocate tensors in the current buffer
1000
+ if (!alloc_tensor_range(ctx, first, t, buft, cur_buf_size, &buffers, &n_buffers)) {
1001
+ return NULL;
1002
+ }
1003
+ first = t;
1004
+ cur_buf_size = this_size;
1005
+ } else {
1006
+ cur_buf_size += this_size;
1007
+ }
1008
+ }
1009
+
1010
+ // allocate remaining tensors
1011
+ if (cur_buf_size > 0) {
1012
+ if (!alloc_tensor_range(ctx, first, NULL, buft, cur_buf_size, &buffers, &n_buffers)) {
1013
+ return NULL;
1014
+ }
1015
+ }
1016
+
1017
+ if (n_buffers == 0) {
1018
+ #ifndef NDEBUG
1019
+ fprintf(stderr, "%s: all tensors in the context are already allocated\n", __func__);
1020
+ #endif
1021
+ return NULL;
1022
+ }
1023
+
1024
+ lm_ggml_backend_buffer_t buffer;
1025
+ if (n_buffers == 1) {
1026
+ buffer = buffers[0];
1027
+ } else {
1028
+ buffer = lm_ggml_backend_multi_buffer_alloc_buffer(buffers, n_buffers);
1029
+ }
1030
+ free(buffers);
1031
+ return buffer;
1032
+ }
1033
+
1034
+ lm_ggml_backend_buffer_t lm_ggml_backend_alloc_ctx_tensors(struct lm_ggml_context * ctx, lm_ggml_backend_t backend) {
1035
+ return lm_ggml_backend_alloc_ctx_tensors_from_buft(ctx, lm_ggml_backend_get_default_buffer_type(backend));
1036
+ }