cui-llama.rn 1.4.0 → 1.4.2

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 (108) hide show
  1. package/README.md +4 -23
  2. package/android/build.gradle +12 -3
  3. package/android/src/main/CMakeLists.txt +13 -7
  4. package/android/src/main/java/com/rnllama/LlamaContext.java +27 -20
  5. package/android/src/main/java/com/rnllama/RNLlama.java +5 -1
  6. package/android/src/main/jni.cpp +15 -12
  7. package/android/src/main/jniLibs/arm64-v8a/librnllama.so +0 -0
  8. package/android/src/main/jniLibs/arm64-v8a/librnllama_v8.so +0 -0
  9. package/android/src/main/jniLibs/arm64-v8a/librnllama_v8_2.so +0 -0
  10. package/android/src/main/jniLibs/arm64-v8a/librnllama_v8_2_dotprod.so +0 -0
  11. package/android/src/main/jniLibs/arm64-v8a/librnllama_v8_2_dotprod_i8mm.so +0 -0
  12. package/android/src/main/jniLibs/arm64-v8a/librnllama_v8_2_i8mm.so +0 -0
  13. package/android/src/main/jniLibs/x86_64/librnllama.so +0 -0
  14. package/android/src/main/jniLibs/x86_64/librnllama_x86_64.so +0 -0
  15. package/cpp/README.md +1 -1
  16. package/cpp/common.cpp +158 -267
  17. package/cpp/common.h +46 -12
  18. package/cpp/ggml-alloc.c +1042 -1037
  19. package/cpp/ggml-backend-impl.h +255 -256
  20. package/cpp/ggml-backend-reg.cpp +582 -582
  21. package/cpp/ggml-backend.cpp +2002 -2002
  22. package/cpp/ggml-backend.h +354 -352
  23. package/cpp/ggml-common.h +1853 -1853
  24. package/cpp/ggml-cpp.h +39 -39
  25. package/cpp/ggml-cpu-aarch64.cpp +4247 -4247
  26. package/cpp/ggml-cpu-aarch64.h +8 -8
  27. package/cpp/ggml-cpu-impl.h +386 -386
  28. package/cpp/ggml-cpu-quants.c +10920 -10839
  29. package/cpp/ggml-cpu-traits.cpp +36 -36
  30. package/cpp/ggml-cpu-traits.h +38 -38
  31. package/cpp/ggml-cpu.c +329 -60
  32. package/cpp/ggml-cpu.cpp +10 -2
  33. package/cpp/ggml-cpu.h +135 -135
  34. package/cpp/ggml-impl.h +567 -567
  35. package/cpp/ggml-metal-impl.h +17 -17
  36. package/cpp/ggml-metal.m +4884 -4884
  37. package/cpp/ggml-quants.c +5238 -5238
  38. package/cpp/ggml-threading.h +14 -14
  39. package/cpp/ggml.c +6514 -6448
  40. package/cpp/ggml.h +2194 -2163
  41. package/cpp/gguf.cpp +1329 -1325
  42. package/cpp/gguf.h +202 -202
  43. package/cpp/json-schema-to-grammar.cpp +1045 -1045
  44. package/cpp/json-schema-to-grammar.h +8 -8
  45. package/cpp/json.hpp +24766 -24766
  46. package/cpp/llama-adapter.cpp +347 -346
  47. package/cpp/llama-adapter.h +74 -73
  48. package/cpp/llama-arch.cpp +1487 -1434
  49. package/cpp/llama-arch.h +400 -395
  50. package/cpp/llama-batch.cpp +368 -368
  51. package/cpp/llama-batch.h +88 -88
  52. package/cpp/llama-chat.cpp +578 -567
  53. package/cpp/llama-chat.h +52 -51
  54. package/cpp/llama-context.cpp +1775 -1771
  55. package/cpp/llama-context.h +128 -128
  56. package/cpp/llama-cparams.cpp +1 -1
  57. package/cpp/llama-cparams.h +37 -37
  58. package/cpp/llama-cpp.h +30 -30
  59. package/cpp/llama-grammar.cpp +1139 -1139
  60. package/cpp/llama-grammar.h +143 -143
  61. package/cpp/llama-hparams.cpp +71 -71
  62. package/cpp/llama-hparams.h +139 -140
  63. package/cpp/llama-impl.cpp +167 -167
  64. package/cpp/llama-impl.h +61 -61
  65. package/cpp/llama-kv-cache.cpp +718 -718
  66. package/cpp/llama-kv-cache.h +218 -218
  67. package/cpp/llama-mmap.cpp +2 -1
  68. package/cpp/llama-mmap.h +67 -67
  69. package/cpp/llama-model-loader.cpp +1124 -1011
  70. package/cpp/llama-model-loader.h +167 -158
  71. package/cpp/llama-model.cpp +3997 -2202
  72. package/cpp/llama-model.h +370 -391
  73. package/cpp/llama-sampling.cpp +2408 -2406
  74. package/cpp/llama-sampling.h +32 -48
  75. package/cpp/llama-vocab.cpp +3247 -1982
  76. package/cpp/llama-vocab.h +125 -182
  77. package/cpp/llama.cpp +416 -2886
  78. package/cpp/llama.h +1323 -1285
  79. package/cpp/log.cpp +401 -401
  80. package/cpp/log.h +121 -121
  81. package/cpp/rn-llama.cpp +822 -0
  82. package/cpp/rn-llama.h +123 -0
  83. package/cpp/rn-llama.hpp +18 -12
  84. package/cpp/sampling.cpp +505 -500
  85. package/cpp/sgemm.cpp +2597 -2597
  86. package/cpp/speculative.cpp +277 -274
  87. package/cpp/speculative.h +28 -28
  88. package/cpp/unicode.cpp +2 -3
  89. package/ios/CMakeLists.txt +99 -0
  90. package/ios/RNLlama.h +5 -1
  91. package/ios/RNLlama.mm +2 -2
  92. package/ios/RNLlamaContext.h +8 -1
  93. package/ios/RNLlamaContext.mm +15 -11
  94. package/ios/rnllama.xcframework/Info.plist +74 -0
  95. package/jest/mock.js +3 -2
  96. package/lib/commonjs/NativeRNLlama.js.map +1 -1
  97. package/lib/commonjs/index.js +4 -2
  98. package/lib/commonjs/index.js.map +1 -1
  99. package/lib/module/NativeRNLlama.js.map +1 -1
  100. package/lib/module/index.js +4 -2
  101. package/lib/module/index.js.map +1 -1
  102. package/lib/typescript/NativeRNLlama.d.ts +5 -1
  103. package/lib/typescript/NativeRNLlama.d.ts.map +1 -1
  104. package/lib/typescript/index.d.ts.map +1 -1
  105. package/llama-rn.podspec +8 -2
  106. package/package.json +5 -2
  107. package/src/NativeRNLlama.ts +5 -1
  108. package/src/index.ts +9 -2
@@ -1,2002 +1,2002 @@
1
- // Note: porting this file to C++ is a work in progress
2
-
3
- #ifdef _WIN32
4
- #define WIN32_LEAN_AND_MEAN
5
- #ifndef NOMINMAX
6
- # define NOMINMAX
7
- #endif
8
- #include <windows.h>
9
- #endif
10
-
11
- #include "ggml-backend.h"
12
- #include "ggml-backend-impl.h"
13
- #include "ggml-alloc.h"
14
- #include "ggml-impl.h"
15
-
16
- #include <assert.h>
17
- #include <limits.h>
18
- #include <stdarg.h>
19
- #include <stdio.h>
20
- #include <stdlib.h>
21
- #include <string.h>
22
- #include <string>
23
- #include <vector>
24
-
25
- #ifdef __APPLE__
26
- #include <sys/types.h>
27
- #include <sys/sysctl.h>
28
- #endif
29
-
30
-
31
- // backend buffer type
32
-
33
- const char * lm_ggml_backend_buft_name(lm_ggml_backend_buffer_type_t buft) {
34
- return buft->iface.get_name(buft);
35
- }
36
-
37
- lm_ggml_backend_buffer_t lm_ggml_backend_buft_alloc_buffer(lm_ggml_backend_buffer_type_t buft, size_t size) {
38
- if (size == 0) {
39
- // return a dummy buffer for zero-sized allocations
40
- return lm_ggml_backend_buffer_init(buft, {}, NULL, 0);
41
- }
42
-
43
- return buft->iface.alloc_buffer(buft, size);
44
- }
45
-
46
- size_t lm_ggml_backend_buft_get_alignment(lm_ggml_backend_buffer_type_t buft) {
47
- return buft->iface.get_alignment(buft);
48
- }
49
-
50
- size_t lm_ggml_backend_buft_get_max_size(lm_ggml_backend_buffer_type_t buft) {
51
- // get_max_size is optional, defaults to SIZE_MAX
52
- if (buft->iface.get_max_size) {
53
- return buft->iface.get_max_size(buft);
54
- }
55
- return SIZE_MAX;
56
- }
57
-
58
- size_t lm_ggml_backend_buft_get_alloc_size(lm_ggml_backend_buffer_type_t buft, struct lm_ggml_tensor * tensor) {
59
- // get_alloc_size is optional, defaults to lm_ggml_nbytes
60
- if (buft->iface.get_alloc_size) {
61
- size_t size = buft->iface.get_alloc_size(buft, tensor);
62
- assert(size >= lm_ggml_nbytes(tensor));
63
- return size;
64
- }
65
- return lm_ggml_nbytes(tensor);
66
- }
67
-
68
- bool lm_ggml_backend_buft_is_host(lm_ggml_backend_buffer_type_t buft) {
69
- if (buft->iface.is_host) {
70
- return buft->iface.is_host(buft);
71
- }
72
- return false;
73
- }
74
-
75
- lm_ggml_backend_dev_t lm_ggml_backend_buft_get_device(lm_ggml_backend_buffer_type_t buft) {
76
- return buft->device;
77
- }
78
-
79
- // backend buffer
80
-
81
- lm_ggml_backend_buffer_t lm_ggml_backend_buffer_init(
82
- lm_ggml_backend_buffer_type_t buft,
83
- struct lm_ggml_backend_buffer_i iface,
84
- void * context,
85
- size_t size) {
86
- lm_ggml_backend_buffer_t buffer = new lm_ggml_backend_buffer {
87
- /* .interface = */ iface,
88
- /* .buft = */ buft,
89
- /* .context = */ context,
90
- /* .size = */ size,
91
- /* .usage = */ LM_GGML_BACKEND_BUFFER_USAGE_ANY
92
- };
93
-
94
- return buffer;
95
- }
96
-
97
- const char * lm_ggml_backend_buffer_name(lm_ggml_backend_buffer_t buffer) {
98
- return lm_ggml_backend_buft_name(lm_ggml_backend_buffer_get_type(buffer));
99
- }
100
-
101
- void lm_ggml_backend_buffer_free(lm_ggml_backend_buffer_t buffer) {
102
- if (buffer == NULL) {
103
- return;
104
- }
105
-
106
- if (buffer->iface.free_buffer != NULL) {
107
- buffer->iface.free_buffer(buffer);
108
- }
109
- delete buffer;
110
- }
111
-
112
- size_t lm_ggml_backend_buffer_get_size(lm_ggml_backend_buffer_t buffer) {
113
- return buffer->size;
114
- }
115
-
116
- void * lm_ggml_backend_buffer_get_base(lm_ggml_backend_buffer_t buffer) {
117
- // get_base is optional if the buffer is zero-sized
118
- if (buffer->size == 0) {
119
- return NULL;
120
- }
121
-
122
- void * base = buffer->iface.get_base(buffer);
123
-
124
- LM_GGML_ASSERT(base != NULL && "backend buffer base cannot be NULL");
125
-
126
- return base;
127
- }
128
-
129
- void lm_ggml_backend_buffer_init_tensor(lm_ggml_backend_buffer_t buffer, struct lm_ggml_tensor * tensor) {
130
- // init_tensor is optional
131
- if (buffer->iface.init_tensor) {
132
- buffer->iface.init_tensor(buffer, tensor);
133
- }
134
- }
135
-
136
- void lm_ggml_backend_buffer_clear(lm_ggml_backend_buffer_t buffer, uint8_t value) {
137
- // clear is optional if the buffer is zero-sized
138
- if (buffer->size == 0) {
139
- return;
140
- }
141
-
142
- buffer->iface.clear(buffer, value);
143
- }
144
-
145
- size_t lm_ggml_backend_buffer_get_alignment(lm_ggml_backend_buffer_t buffer) {
146
- return lm_ggml_backend_buft_get_alignment(lm_ggml_backend_buffer_get_type(buffer));
147
- }
148
-
149
- size_t lm_ggml_backend_buffer_get_max_size(lm_ggml_backend_buffer_t buffer) {
150
- return lm_ggml_backend_buft_get_max_size(lm_ggml_backend_buffer_get_type(buffer));
151
- }
152
-
153
- size_t lm_ggml_backend_buffer_get_alloc_size(lm_ggml_backend_buffer_t buffer, struct lm_ggml_tensor * tensor) {
154
- return lm_ggml_backend_buft_get_alloc_size(lm_ggml_backend_buffer_get_type(buffer), tensor);
155
- }
156
-
157
- bool lm_ggml_backend_buffer_is_host(lm_ggml_backend_buffer_t buffer) {
158
- return lm_ggml_backend_buft_is_host(lm_ggml_backend_buffer_get_type(buffer));
159
- }
160
-
161
- void lm_ggml_backend_buffer_set_usage(lm_ggml_backend_buffer_t buffer, enum lm_ggml_backend_buffer_usage usage) {
162
- buffer->usage = usage;
163
-
164
- // FIXME: add a generic callback to the buffer interface
165
- if (lm_ggml_backend_buffer_is_multi_buffer(buffer)) {
166
- lm_ggml_backend_multi_buffer_set_usage(buffer, usage);
167
- }
168
- }
169
-
170
- enum lm_ggml_backend_buffer_usage lm_ggml_backend_buffer_get_usage(lm_ggml_backend_buffer_t buffer) {
171
- return buffer->usage;
172
- }
173
-
174
- lm_ggml_backend_buffer_type_t lm_ggml_backend_buffer_get_type(lm_ggml_backend_buffer_t buffer) {
175
- return buffer->buft;
176
- }
177
-
178
- void lm_ggml_backend_buffer_reset(lm_ggml_backend_buffer_t buffer) {
179
- if (buffer->iface.reset) {
180
- buffer->iface.reset(buffer);
181
- }
182
- }
183
-
184
- bool lm_ggml_backend_buffer_copy_tensor(const struct lm_ggml_tensor * src, struct lm_ggml_tensor * dst) {
185
- lm_ggml_backend_buffer_t dst_buf = dst->view_src ? dst->view_src->buffer : dst->buffer;
186
- if (dst_buf->iface.cpy_tensor) {
187
- return dst_buf->iface.cpy_tensor(dst_buf, src, dst);
188
- }
189
- return false;
190
- }
191
-
192
- // backend
193
-
194
- lm_ggml_guid_t lm_ggml_backend_guid(lm_ggml_backend_t backend) {
195
- if (backend == NULL) {
196
- return NULL;
197
- }
198
- return backend->guid;
199
- }
200
-
201
- const char * lm_ggml_backend_name(lm_ggml_backend_t backend) {
202
- if (backend == NULL) {
203
- return "NULL";
204
- }
205
- return backend->iface.get_name(backend);
206
- }
207
-
208
- void lm_ggml_backend_free(lm_ggml_backend_t backend) {
209
- if (backend == NULL) {
210
- return;
211
- }
212
-
213
- backend->iface.free(backend);
214
- }
215
-
216
- lm_ggml_backend_buffer_type_t lm_ggml_backend_get_default_buffer_type(lm_ggml_backend_t backend) {
217
- return lm_ggml_backend_dev_buffer_type(backend->device);
218
- }
219
-
220
- lm_ggml_backend_buffer_t lm_ggml_backend_alloc_buffer(lm_ggml_backend_t backend, size_t size) {
221
- return lm_ggml_backend_buft_alloc_buffer(lm_ggml_backend_get_default_buffer_type(backend), size);
222
- }
223
-
224
- size_t lm_ggml_backend_get_alignment(lm_ggml_backend_t backend) {
225
- return lm_ggml_backend_buft_get_alignment(lm_ggml_backend_get_default_buffer_type(backend));
226
- }
227
-
228
- size_t lm_ggml_backend_get_max_size(lm_ggml_backend_t backend) {
229
- return lm_ggml_backend_buft_get_max_size(lm_ggml_backend_get_default_buffer_type(backend));
230
- }
231
-
232
- void lm_ggml_backend_tensor_set_async(lm_ggml_backend_t backend, struct lm_ggml_tensor * tensor, const void * data, size_t offset, size_t size) {
233
- LM_GGML_ASSERT(tensor->data != NULL && "tensor not allocated");
234
- LM_GGML_ASSERT(offset + size <= lm_ggml_nbytes(tensor) && "tensor write out of bounds");
235
-
236
- if (backend->iface.set_tensor_async == NULL) {
237
- lm_ggml_backend_tensor_set(tensor, data, offset, size);
238
- } else {
239
- backend->iface.set_tensor_async(backend, tensor, data, offset, size);
240
- }
241
- }
242
-
243
- void lm_ggml_backend_tensor_get_async(lm_ggml_backend_t backend, const struct lm_ggml_tensor * tensor, void * data, size_t offset, size_t size) {
244
- LM_GGML_ASSERT(tensor->data != NULL && "tensor not allocated");
245
- LM_GGML_ASSERT(offset + size <= lm_ggml_nbytes(tensor) && "tensor read out of bounds");
246
-
247
- if (backend->iface.get_tensor_async == NULL) {
248
- lm_ggml_backend_tensor_get(tensor, data, offset, size);
249
- } else {
250
- backend->iface.get_tensor_async(backend, tensor, data, offset, size);
251
- }
252
- }
253
-
254
- void lm_ggml_backend_tensor_set(struct lm_ggml_tensor * tensor, const void * data, size_t offset, size_t size) {
255
- LM_GGML_ASSERT(tensor);
256
- lm_ggml_backend_buffer_t buf = tensor->view_src ? tensor->view_src->buffer : tensor->buffer;
257
-
258
- if (size == 0) {
259
- return;
260
- }
261
-
262
- LM_GGML_ASSERT(buf != NULL && "tensor buffer not set");
263
- LM_GGML_ASSERT(tensor->data != NULL && "tensor not allocated");
264
- LM_GGML_ASSERT(offset + size <= lm_ggml_nbytes(tensor) && "tensor write out of bounds");
265
-
266
- buf->iface.set_tensor(buf, tensor, data, offset, size);
267
- }
268
-
269
- void lm_ggml_backend_tensor_get(const struct lm_ggml_tensor * tensor, void * data, size_t offset, size_t size) {
270
- LM_GGML_ASSERT(tensor);
271
- lm_ggml_backend_buffer_t buf = tensor->view_src ? tensor->view_src->buffer : tensor->buffer;
272
-
273
- if (size == 0) {
274
- return;
275
- }
276
-
277
- LM_GGML_ASSERT(buf != NULL && "tensor buffer not set");
278
- LM_GGML_ASSERT(tensor->data != NULL && "tensor not allocated");
279
- LM_GGML_ASSERT(offset + size <= lm_ggml_nbytes(tensor) && "tensor read out of bounds");
280
-
281
- buf->iface.get_tensor(buf, tensor, data, offset, size);
282
- }
283
-
284
- void lm_ggml_backend_tensor_memset(struct lm_ggml_tensor * tensor, uint8_t value, size_t offset, size_t size) {
285
- lm_ggml_backend_buffer_t buf = tensor->view_src ? tensor->view_src->buffer : tensor->buffer;
286
-
287
- if (size == 0) {
288
- return;
289
- }
290
-
291
- LM_GGML_ASSERT(buf != NULL && "tensor buffer not set");
292
- LM_GGML_ASSERT(tensor->data != NULL && "tensor not allocated");
293
- LM_GGML_ASSERT(offset + size <= lm_ggml_nbytes(tensor) && "tensor write out of bounds");
294
- LM_GGML_ASSERT(buf->iface.memset_tensor != NULL && "memset not implemented by backend buffer");
295
-
296
- buf->iface.memset_tensor(buf, tensor, value, offset, size);
297
- }
298
-
299
- void lm_ggml_backend_synchronize(lm_ggml_backend_t backend) {
300
- if (backend->iface.synchronize == NULL) {
301
- return;
302
- }
303
-
304
- backend->iface.synchronize(backend);
305
- }
306
-
307
- lm_ggml_backend_graph_plan_t lm_ggml_backend_graph_plan_create(lm_ggml_backend_t backend, struct lm_ggml_cgraph * cgraph) {
308
- LM_GGML_ASSERT(backend->iface.graph_plan_create != NULL);
309
-
310
- return backend->iface.graph_plan_create(backend, cgraph);
311
- }
312
-
313
- void lm_ggml_backend_graph_plan_free(lm_ggml_backend_t backend, lm_ggml_backend_graph_plan_t plan) {
314
- LM_GGML_ASSERT(backend->iface.graph_plan_free != NULL);
315
-
316
- backend->iface.graph_plan_free(backend, plan);
317
- }
318
-
319
- enum lm_ggml_status lm_ggml_backend_graph_plan_compute(lm_ggml_backend_t backend, lm_ggml_backend_graph_plan_t plan) {
320
- LM_GGML_ASSERT(backend->iface.graph_plan_compute != NULL);
321
-
322
- return backend->iface.graph_plan_compute(backend, plan);
323
- }
324
-
325
- enum lm_ggml_status lm_ggml_backend_graph_compute(lm_ggml_backend_t backend, struct lm_ggml_cgraph * cgraph) {
326
- enum lm_ggml_status err = lm_ggml_backend_graph_compute_async(backend, cgraph);
327
- lm_ggml_backend_synchronize(backend);
328
- return err;
329
- }
330
-
331
- enum lm_ggml_status lm_ggml_backend_graph_compute_async(lm_ggml_backend_t backend, struct lm_ggml_cgraph * cgraph) {
332
- return backend->iface.graph_compute(backend, cgraph);
333
- }
334
-
335
- bool lm_ggml_backend_supports_op(lm_ggml_backend_t backend, const struct lm_ggml_tensor * op) {
336
- return lm_ggml_backend_dev_supports_op(backend->device, op);
337
- }
338
-
339
- bool lm_ggml_backend_supports_buft(lm_ggml_backend_t backend, lm_ggml_backend_buffer_type_t buft) {
340
- return lm_ggml_backend_dev_supports_buft(backend->device, buft);
341
- }
342
-
343
- bool lm_ggml_backend_offload_op(lm_ggml_backend_t backend, const struct lm_ggml_tensor * op) {
344
- return lm_ggml_backend_dev_offload_op(backend->device, op);
345
- }
346
-
347
- lm_ggml_backend_dev_t lm_ggml_backend_get_device(lm_ggml_backend_t backend) {
348
- return backend->device;
349
- }
350
-
351
- // backend copy
352
-
353
- static bool lm_ggml_are_same_layout(const struct lm_ggml_tensor * a, const struct lm_ggml_tensor * b) {
354
- if (a->type != b->type) {
355
- return false;
356
- }
357
- for (int i = 0; i < LM_GGML_MAX_DIMS; i++) {
358
- if (a->ne[i] != b->ne[i]) {
359
- return false;
360
- }
361
- if (a->nb[i] != b->nb[i]) {
362
- return false;
363
- }
364
- }
365
- return true;
366
- }
367
-
368
- void lm_ggml_backend_tensor_copy(struct lm_ggml_tensor * src, struct lm_ggml_tensor * dst) {
369
- LM_GGML_ASSERT(lm_ggml_are_same_layout(src, dst) && "cannot copy tensors with different layouts");
370
-
371
- if (src == dst) {
372
- return;
373
- }
374
-
375
- if (lm_ggml_backend_buffer_is_host(src->buffer)) {
376
- lm_ggml_backend_tensor_set(dst, src->data, 0, lm_ggml_nbytes(src));
377
- } else if (lm_ggml_backend_buffer_is_host(dst->buffer)) {
378
- lm_ggml_backend_tensor_get(src, dst->data, 0, lm_ggml_nbytes(src));
379
- } else if (!lm_ggml_backend_buffer_copy_tensor(src, dst)) {
380
- #ifndef NDEBUG
381
- LM_GGML_LOG_DEBUG("%s: warning: slow copy from %s to %s\n", __func__, lm_ggml_backend_buffer_name(src->buffer), lm_ggml_backend_buffer_name(dst->buffer));
382
- #endif
383
- size_t nbytes = lm_ggml_nbytes(src);
384
- void * data = malloc(nbytes);
385
- lm_ggml_backend_tensor_get(src, data, 0, nbytes);
386
- lm_ggml_backend_tensor_set(dst, data, 0, nbytes);
387
- free(data);
388
- }
389
- }
390
-
391
- void lm_ggml_backend_tensor_copy_async(lm_ggml_backend_t backend_src, lm_ggml_backend_t backend_dst, struct lm_ggml_tensor * src, struct lm_ggml_tensor * dst) {
392
- LM_GGML_ASSERT(lm_ggml_are_same_layout(src, dst) && "cannot copy tensors with different layouts");
393
-
394
- if (src == dst) {
395
- return;
396
- }
397
-
398
- if (backend_dst->iface.cpy_tensor_async != NULL) {
399
- if (backend_dst->iface.cpy_tensor_async(backend_src, backend_dst, src, dst)) {
400
- return;
401
- }
402
- }
403
-
404
- // an async copy would normally happen after all the queued operations on both backends are completed
405
- // to simulate the same behavior, we need to synchronize both backends first, and do a blocking copy
406
- lm_ggml_backend_synchronize(backend_src);
407
- lm_ggml_backend_synchronize(backend_dst);
408
- lm_ggml_backend_tensor_copy(src, dst);
409
- }
410
-
411
- // events
412
-
413
- lm_ggml_backend_event_t lm_ggml_backend_event_new(lm_ggml_backend_dev_t device) {
414
- // null device is allowed for the transition period to the device interface
415
- if (device == NULL || device->iface.event_new == NULL) {
416
- return NULL;
417
- }
418
- return device->iface.event_new(device);
419
- }
420
-
421
- void lm_ggml_backend_event_free(lm_ggml_backend_event_t event) {
422
- if (event == NULL) {
423
- return;
424
- }
425
- event->device->iface.event_free(event->device, event);
426
- }
427
-
428
- void lm_ggml_backend_event_record(lm_ggml_backend_event_t event, lm_ggml_backend_t backend) {
429
- LM_GGML_ASSERT(backend->iface.event_record != NULL);
430
-
431
- backend->iface.event_record(backend, event);
432
- }
433
-
434
- void lm_ggml_backend_event_synchronize(lm_ggml_backend_event_t event) {
435
- LM_GGML_ASSERT(event->device->iface.event_synchronize);
436
-
437
- event->device->iface.event_synchronize(event->device, event);
438
- }
439
-
440
- void lm_ggml_backend_event_wait(lm_ggml_backend_t backend, lm_ggml_backend_event_t event) {
441
- LM_GGML_ASSERT(backend->iface.event_wait != NULL);
442
-
443
- backend->iface.event_wait(backend, event);
444
- }
445
-
446
- // Backend device
447
-
448
- const char * lm_ggml_backend_dev_name(lm_ggml_backend_dev_t device) {
449
- return device->iface.get_name(device);
450
- }
451
-
452
- const char * lm_ggml_backend_dev_description(lm_ggml_backend_dev_t device) {
453
- return device->iface.get_description(device);
454
- }
455
-
456
- void lm_ggml_backend_dev_memory(lm_ggml_backend_dev_t device, size_t * free, size_t * total) {
457
- device->iface.get_memory(device, free, total);
458
- }
459
-
460
- enum lm_ggml_backend_dev_type lm_ggml_backend_dev_type(lm_ggml_backend_dev_t device) {
461
- return device->iface.get_type(device);
462
- }
463
-
464
- void lm_ggml_backend_dev_get_props(lm_ggml_backend_dev_t device, struct lm_ggml_backend_dev_props * props) {
465
- memset(props, 0, sizeof(*props));
466
- device->iface.get_props(device, props);
467
- }
468
-
469
- lm_ggml_backend_reg_t lm_ggml_backend_dev_backend_reg(lm_ggml_backend_dev_t device) {
470
- return device->reg;
471
- }
472
-
473
- lm_ggml_backend_t lm_ggml_backend_dev_init(lm_ggml_backend_dev_t device, const char * params) {
474
- return device->iface.init_backend(device, params);
475
- }
476
-
477
- lm_ggml_backend_buffer_type_t lm_ggml_backend_dev_buffer_type(lm_ggml_backend_dev_t device) {
478
- return device->iface.get_buffer_type(device);
479
- }
480
-
481
- lm_ggml_backend_buffer_type_t lm_ggml_backend_dev_host_buffer_type(lm_ggml_backend_dev_t device) {
482
- if (device->iface.get_host_buffer_type == NULL) {
483
- return NULL;
484
- }
485
-
486
- return device->iface.get_host_buffer_type(device);
487
- }
488
-
489
- lm_ggml_backend_buffer_t lm_ggml_backend_dev_buffer_from_host_ptr(lm_ggml_backend_dev_t device, void * ptr, size_t size, size_t max_tensor_size) {
490
- return device->iface.buffer_from_host_ptr(device, ptr, size, max_tensor_size);
491
- }
492
-
493
- bool lm_ggml_backend_dev_supports_op(lm_ggml_backend_dev_t device, const struct lm_ggml_tensor * op) {
494
- return device->iface.supports_op(device, op);
495
- }
496
-
497
- bool lm_ggml_backend_dev_supports_buft(lm_ggml_backend_dev_t device, lm_ggml_backend_buffer_type_t buft) {
498
- return device->iface.supports_buft(device, buft);
499
- }
500
-
501
- bool lm_ggml_backend_dev_offload_op(lm_ggml_backend_dev_t device, const struct lm_ggml_tensor * op) {
502
- if (device->iface.offload_op != NULL) {
503
- return device->iface.offload_op(device, op);
504
- }
505
-
506
- return false;
507
- }
508
-
509
- // Backend (reg)
510
-
511
- const char * lm_ggml_backend_reg_name(lm_ggml_backend_reg_t reg) {
512
- return reg->iface.get_name(reg);
513
- }
514
-
515
- size_t lm_ggml_backend_reg_dev_count(lm_ggml_backend_reg_t reg) {
516
- return reg->iface.get_device_count(reg);
517
- }
518
-
519
- lm_ggml_backend_dev_t lm_ggml_backend_reg_dev_get(lm_ggml_backend_reg_t reg, size_t index) {
520
- return reg->iface.get_device(reg, index);
521
- }
522
-
523
- void * lm_ggml_backend_reg_get_proc_address(lm_ggml_backend_reg_t reg, const char * name) {
524
- if (!reg->iface.get_proc_address) {
525
- return NULL;
526
- }
527
- return reg->iface.get_proc_address(reg, name);
528
- }
529
-
530
- // multi-buffer buffer
531
-
532
- struct lm_ggml_backend_multi_buffer_context {
533
- lm_ggml_backend_buffer_t * buffers;
534
- size_t n_buffers;
535
- };
536
-
537
- static void lm_ggml_backend_multi_buffer_free_buffer(lm_ggml_backend_buffer_t buffer) {
538
- lm_ggml_backend_multi_buffer_context * ctx = (lm_ggml_backend_multi_buffer_context *) buffer->context;
539
- for (size_t i = 0; i < ctx->n_buffers; i++) {
540
- lm_ggml_backend_buffer_free(ctx->buffers[i]);
541
- }
542
-
543
- free(ctx->buffers);
544
- free(ctx);
545
- }
546
-
547
- static void lm_ggml_backend_multi_buffer_clear(lm_ggml_backend_buffer_t buffer, uint8_t value) {
548
- lm_ggml_backend_multi_buffer_context * ctx = (lm_ggml_backend_multi_buffer_context *) buffer->context;
549
- for (size_t i = 0; i < ctx->n_buffers; i++) {
550
- lm_ggml_backend_buffer_clear(ctx->buffers[i], value);
551
- }
552
- }
553
-
554
- static const struct lm_ggml_backend_buffer_i lm_ggml_backend_multi_buffer_i = {
555
- /* .free_buffer = */ lm_ggml_backend_multi_buffer_free_buffer,
556
- /* .get_base = */ NULL,
557
- /* .init_tensor = */ NULL,
558
- /* .memset_tensor = */ NULL,
559
- /* .set_tensor = */ NULL,
560
- /* .get_tensor = */ NULL,
561
- /* .cpy_tensor = */ NULL,
562
- /* .clear = */ lm_ggml_backend_multi_buffer_clear,
563
- /* .reset = */ NULL,
564
- };
565
-
566
- lm_ggml_backend_buffer_t lm_ggml_backend_multi_buffer_alloc_buffer(lm_ggml_backend_buffer_t * buffers, size_t n_buffers) {
567
- lm_ggml_backend_multi_buffer_context * ctx = (lm_ggml_backend_multi_buffer_context *) malloc(sizeof(struct lm_ggml_backend_multi_buffer_context));
568
- ctx->n_buffers = n_buffers;
569
- ctx->buffers = (lm_ggml_backend_buffer_t *) malloc(n_buffers * sizeof(lm_ggml_backend_buffer_t));
570
-
571
- LM_GGML_ASSERT(ctx->buffers != NULL);
572
-
573
- size_t total_size = 0;
574
- for (size_t i = 0; i < n_buffers; i++) {
575
- ctx->buffers[i] = buffers[i];
576
- total_size += lm_ggml_backend_buffer_get_size(buffers[i]);
577
- }
578
-
579
- return lm_ggml_backend_buffer_init(buffers[0]->buft, lm_ggml_backend_multi_buffer_i, ctx, total_size);
580
- }
581
-
582
- bool lm_ggml_backend_buffer_is_multi_buffer(lm_ggml_backend_buffer_t buffer) {
583
- return buffer->iface.free_buffer == lm_ggml_backend_multi_buffer_free_buffer;
584
- }
585
-
586
- void lm_ggml_backend_multi_buffer_set_usage(lm_ggml_backend_buffer_t buffer, enum lm_ggml_backend_buffer_usage usage) {
587
- LM_GGML_ASSERT(lm_ggml_backend_buffer_is_multi_buffer(buffer));
588
- lm_ggml_backend_multi_buffer_context * ctx = (lm_ggml_backend_multi_buffer_context *) buffer->context;
589
- for (size_t i = 0; i < ctx->n_buffers; i++) {
590
- lm_ggml_backend_buffer_set_usage(ctx->buffers[i], usage);
591
- }
592
- }
593
-
594
- // creates a copy of the tensor with the same memory layout
595
- static struct lm_ggml_tensor * lm_ggml_dup_tensor_layout(struct lm_ggml_context * ctx, const struct lm_ggml_tensor * tensor) {
596
- struct lm_ggml_tensor * dup = lm_ggml_dup_tensor(ctx, tensor);
597
- for (int i = 0; i < LM_GGML_MAX_DIMS; i++) {
598
- dup->nb[i] = tensor->nb[i];
599
- }
600
- return dup;
601
- }
602
-
603
- static bool lm_ggml_is_view_op(enum lm_ggml_op op) {
604
- return op == LM_GGML_OP_VIEW || op == LM_GGML_OP_RESHAPE || op == LM_GGML_OP_PERMUTE || op == LM_GGML_OP_TRANSPOSE;
605
- }
606
-
607
- // scheduler
608
-
609
- #ifndef LM_GGML_SCHED_MAX_BACKENDS
610
- #define LM_GGML_SCHED_MAX_BACKENDS 16
611
- #endif
612
-
613
- #ifndef LM_GGML_SCHED_MAX_SPLIT_INPUTS
614
- #define LM_GGML_SCHED_MAX_SPLIT_INPUTS LM_GGML_MAX_SRC
615
- #endif
616
-
617
- #ifndef LM_GGML_SCHED_MAX_COPIES
618
- #define LM_GGML_SCHED_MAX_COPIES 4
619
- #endif
620
-
621
- struct lm_ggml_backend_sched_split {
622
- int backend_id;
623
- int i_start;
624
- int i_end;
625
- struct lm_ggml_tensor * inputs[LM_GGML_SCHED_MAX_SPLIT_INPUTS];
626
- int n_inputs;
627
- // graph view of this split
628
- struct lm_ggml_cgraph graph;
629
- };
630
-
631
- struct lm_ggml_backend_sched {
632
- bool is_reset; // true if the scheduler has been reset since the last graph split
633
- bool is_alloc;
634
-
635
- int n_backends;
636
-
637
- lm_ggml_backend_t backends[LM_GGML_SCHED_MAX_BACKENDS];
638
- lm_ggml_backend_buffer_type_t bufts[LM_GGML_SCHED_MAX_BACKENDS];
639
- lm_ggml_gallocr_t galloc;
640
-
641
- // hash map of the nodes in the graph
642
- struct lm_ggml_hash_set hash_set;
643
- int * hv_tensor_backend_ids; // [hash_set.size]
644
- struct lm_ggml_tensor ** hv_tensor_copies; // [hash_set.size][n_backends][n_copies]
645
-
646
- int * node_backend_ids; // [graph_size]
647
- int * leaf_backend_ids; // [graph_size]
648
-
649
- int * prev_node_backend_ids; // [graph_size]
650
- int * prev_leaf_backend_ids; // [graph_size]
651
-
652
- // copy of the graph with modified inputs
653
- struct lm_ggml_cgraph graph;
654
-
655
- // graph splits
656
- struct lm_ggml_backend_sched_split * splits;
657
- int n_splits;
658
- int splits_capacity;
659
-
660
- // pipeline parallelism support
661
- int n_copies;
662
- int cur_copy;
663
- lm_ggml_backend_event_t events[LM_GGML_SCHED_MAX_BACKENDS][LM_GGML_SCHED_MAX_COPIES];
664
- struct lm_ggml_tensor * graph_inputs[LM_GGML_SCHED_MAX_SPLIT_INPUTS];
665
- int n_graph_inputs;
666
-
667
- struct lm_ggml_context * ctx;
668
-
669
- lm_ggml_backend_sched_eval_callback callback_eval;
670
- void * callback_eval_user_data;
671
-
672
- char * context_buffer;
673
- size_t context_buffer_size;
674
-
675
- int debug;
676
- };
677
-
678
- #define hash_id(tensor) lm_ggml_hash_find_or_insert(&sched->hash_set, tensor)
679
- #define tensor_backend_id(tensor) sched->hv_tensor_backend_ids[hash_id(tensor)]
680
- #define tensor_id_copy(id, backend_id, copy_id) sched->hv_tensor_copies[(id) * sched->n_backends * sched->n_copies + (backend_id) * sched->n_copies + (copy_id)]
681
- #define tensor_copy(tensor, backend_id, copy_id) tensor_id_copy(hash_id(tensor), backend_id, copy_id)
682
-
683
- // returns the priority of the backend, lower id is higher priority
684
- static int lm_ggml_backend_sched_backend_id(lm_ggml_backend_sched_t sched, lm_ggml_backend_t backend) {
685
- for (int i = 0; i < sched->n_backends; i++) {
686
- if (sched->backends[i] == backend) {
687
- return i;
688
- }
689
- }
690
- return -1;
691
- }
692
-
693
- static int lm_ggml_backend_sched_backend_from_buffer(lm_ggml_backend_sched_t sched, const struct lm_ggml_tensor * tensor, const struct lm_ggml_tensor * op) {
694
- lm_ggml_backend_buffer_t buffer = tensor->view_src ? tensor->view_src->buffer : tensor->buffer;
695
- if (buffer == NULL) {
696
- return -1;
697
- }
698
-
699
- // find highest prio backend that supports the buffer type and the op
700
- for (int i = 0; i < sched->n_backends; i++) {
701
- if (lm_ggml_backend_supports_buft(sched->backends[i], buffer->buft) &&
702
- lm_ggml_backend_supports_op(sched->backends[i], op)) {
703
- return i;
704
- }
705
- }
706
-
707
- #ifndef NDEBUG
708
- LM_GGML_LOG_DEBUG("%s: warning: no backend supports op %s with a weight with buffer type %s used in tensor %s, the weight will need to be copied\n",
709
- __func__, lm_ggml_op_desc(tensor), lm_ggml_backend_buffer_name(buffer), tensor->name);
710
- #endif
711
-
712
- return -1;
713
- }
714
-
715
- #if 0
716
- #define LM_GGML_SCHED_MAX_SPLITS_DEBUG 4096
717
- static char causes[LM_GGML_DEFAULT_GRAPH_SIZE*16 + LM_GGML_SCHED_MAX_SPLITS_DEBUG*LM_GGML_SCHED_MAX_SPLIT_INPUTS][128]; // debug only
718
- #define SET_CAUSE(node, ...) sprintf(causes[hash_id(node)], __VA_ARGS__)
719
- #define GET_CAUSE(node) causes[hash_id(node)]
720
- #else
721
- #define SET_CAUSE(node, ...)
722
- #define GET_CAUSE(node) ""
723
- #endif
724
-
725
- // returns the backend that should be used for the node based on the current locations
726
- static int lm_ggml_backend_sched_backend_id_from_cur(lm_ggml_backend_sched_t sched, struct lm_ggml_tensor * tensor) {
727
- // assign pre-allocated nodes to their backend
728
- int cur_backend_id = lm_ggml_backend_sched_backend_from_buffer(sched, tensor, tensor);
729
- if (cur_backend_id != -1) {
730
- SET_CAUSE(tensor, "1.dst");
731
- return cur_backend_id;
732
- }
733
-
734
- // view_src
735
- if (tensor->view_src != NULL) {
736
- cur_backend_id = lm_ggml_backend_sched_backend_from_buffer(sched, tensor->view_src, tensor);
737
- if (cur_backend_id != -1) {
738
- SET_CAUSE(tensor, "1.vsrc");
739
- return cur_backend_id;
740
- }
741
- }
742
-
743
- if (tensor->buffer || (tensor->view_src && tensor->view_src->buffer)) {
744
- // since the tensor is pre-allocated, it cannot be moved to another backend
745
- lm_ggml_backend_buffer_t buffer = tensor->view_src ? tensor->view_src->buffer : tensor->buffer;
746
- LM_GGML_ABORT("pre-allocated tensor (%s) in a buffer (%s) that cannot run the operation (%s)", tensor->name, lm_ggml_backend_buffer_name(buffer), lm_ggml_op_name(tensor->op));
747
- }
748
-
749
- // graph input
750
- if (tensor->flags & LM_GGML_TENSOR_FLAG_INPUT) {
751
- cur_backend_id = sched->n_backends - 1; // last backend (assumed CPU)
752
- SET_CAUSE(tensor, "1.inp");
753
- return cur_backend_id;
754
- }
755
-
756
- // operations with weights are preferably run on the same backend as the weights
757
- for (int i = 0; i < LM_GGML_MAX_SRC; i++) {
758
- const struct lm_ggml_tensor * src = tensor->src[i];
759
- if (src == NULL) {
760
- continue;
761
- }
762
- // skip ROPE since the rope freqs tensor is too small to choose a backend based on it
763
- // not an ideal solution
764
- if (tensor->op != LM_GGML_OP_ROPE && src->buffer != NULL && src->buffer->usage == LM_GGML_BACKEND_BUFFER_USAGE_WEIGHTS) {
765
- int src_backend_id = lm_ggml_backend_sched_backend_from_buffer(sched, src, tensor);
766
- // check if a backend with higher prio wants to offload the op
767
- if (src_backend_id == sched->n_backends - 1 && lm_ggml_backend_buffer_is_host(src->buffer)) {
768
- for (int b = 0; b < src_backend_id; b++) {
769
- if (lm_ggml_backend_supports_op(sched->backends[b], tensor) && lm_ggml_backend_offload_op(sched->backends[b], tensor)) {
770
- SET_CAUSE(tensor, "1.off");
771
- return b;
772
- }
773
- }
774
- }
775
- SET_CAUSE(tensor, "1.wgt%d", i);
776
- return src_backend_id;
777
- }
778
- }
779
-
780
- return -1;
781
- }
782
-
783
- static char * fmt_size(size_t size) {
784
- static char buffer[128];
785
- if (size >= 1024*1024) {
786
- snprintf(buffer, sizeof(buffer), "%zuM", size/1024/1024);
787
- } else {
788
- snprintf(buffer, sizeof(buffer), "%zuK", size/1024);
789
- }
790
- return buffer;
791
- }
792
-
793
- static void lm_ggml_backend_sched_print_assignments(lm_ggml_backend_sched_t sched, struct lm_ggml_cgraph * graph) {
794
- int cur_split = 0;
795
- for (int i = 0; i < graph->n_nodes; i++) {
796
- if (cur_split < sched->n_splits && i == sched->splits[cur_split].i_start) {
797
- lm_ggml_backend_t split_backend = sched->backends[sched->splits[cur_split].backend_id];
798
- LM_GGML_LOG_DEBUG("\n## SPLIT #%d: %s # %d inputs", cur_split, lm_ggml_backend_name(split_backend),
799
- sched->splits[cur_split].n_inputs);
800
- for (int j = 0; j < sched->splits[cur_split].n_inputs; j++) {
801
- if (j == 0) {
802
- LM_GGML_LOG_DEBUG(": ");
803
- }
804
- LM_GGML_LOG_DEBUG("[%s (%5.5s)] ", sched->splits[cur_split].inputs[j]->name,
805
- fmt_size(lm_ggml_nbytes(sched->splits[cur_split].inputs[j])));
806
- }
807
- LM_GGML_LOG_DEBUG("\n");
808
- cur_split++;
809
- }
810
- struct lm_ggml_tensor * node = graph->nodes[i];
811
- if (lm_ggml_is_view_op(node->op)) {
812
- continue;
813
- }
814
- if (sched->debug > 1) {
815
- lm_ggml_backend_t tensor_backend = lm_ggml_backend_sched_get_tensor_backend(sched, node);
816
- LM_GGML_LOG_DEBUG("node #%3d (%10.10s): %20.20s (%5.5s) [%5.5s %8.8s]:", i, lm_ggml_op_name(node->op), node->name,
817
- fmt_size(lm_ggml_nbytes(node)), tensor_backend ? lm_ggml_backend_name(tensor_backend) : "NULL", GET_CAUSE(node));
818
- for (int j = 0; j < LM_GGML_MAX_SRC; j++) {
819
- struct lm_ggml_tensor * src = node->src[j];
820
- if (src == NULL) {
821
- continue;
822
- }
823
- lm_ggml_backend_t src_backend = lm_ggml_backend_sched_get_tensor_backend(sched, src);
824
- LM_GGML_LOG_DEBUG(" %20.20s (%5.5s) [%5.5s %8.8s]", src->name,
825
- fmt_size(lm_ggml_nbytes(src)), src_backend ? lm_ggml_backend_name(src_backend) : "NULL", GET_CAUSE(src));
826
- }
827
- LM_GGML_LOG_DEBUG("\n");
828
- }
829
- }
830
- }
831
-
832
- static bool lm_ggml_backend_sched_buffer_supported(lm_ggml_backend_sched_t sched, struct lm_ggml_tensor * t, int backend_id) {
833
- lm_ggml_backend_buffer_t buf = t->view_src ? t->view_src->buffer : t->buffer;
834
- lm_ggml_backend_buffer_type_t buft = NULL;
835
-
836
- if (buf) {
837
- // the tensor is already allocated
838
- buft = buf->buft;
839
- } else {
840
- // see if the tensor already has a backend assigned, and use the buffer type of that backend
841
- int tensor_backend_id = tensor_backend_id(t);
842
- if (tensor_backend_id == -1 && t->view_src) {
843
- tensor_backend_id = tensor_backend_id(t->view_src);
844
- }
845
- if (tensor_backend_id != -1) {
846
- buft = sched->bufts[tensor_backend_id];
847
- }
848
- }
849
-
850
- return buft != NULL && lm_ggml_backend_supports_buft(sched->backends[backend_id], buft);
851
- }
852
-
853
- static void lm_ggml_backend_sched_set_if_supported(lm_ggml_backend_sched_t sched, struct lm_ggml_tensor * node, int cur_backend_id, int * node_backend_id) {
854
- if (lm_ggml_backend_supports_op(sched->backends[cur_backend_id], node)) {
855
- *node_backend_id = cur_backend_id;
856
- SET_CAUSE(node, "2.sup");
857
- }
858
- }
859
-
860
- // assigns backends to ops and splits the graph into subgraphs that can be computed on the same backend
861
- static void lm_ggml_backend_sched_split_graph(lm_ggml_backend_sched_t sched, struct lm_ggml_cgraph * graph) {
862
- // reset splits
863
- sched->n_splits = 0;
864
- sched->n_graph_inputs = 0;
865
- sched->is_reset = false;
866
-
867
- struct lm_ggml_init_params params = {
868
- /* .mem_size = */ sched->context_buffer_size,
869
- /* .mem_buffer = */ sched->context_buffer,
870
- /* .no_alloc = */ true
871
- };
872
-
873
- lm_ggml_free(sched->ctx);
874
-
875
- sched->ctx = lm_ggml_init(params);
876
- if (sched->ctx == NULL) {
877
- LM_GGML_ABORT("%s: failed to initialize context\n", __func__);
878
- }
879
-
880
- // pass 1: assign backends to ops with pre-allocated inputs
881
- for (int i = 0; i < graph->n_leafs; i++) {
882
- struct lm_ggml_tensor * leaf = graph->leafs[i];
883
- int * leaf_backend_id = &tensor_backend_id(leaf);
884
- // do not overwrite user assignments
885
- if (*leaf_backend_id == -1) {
886
- *leaf_backend_id = lm_ggml_backend_sched_backend_id_from_cur(sched, leaf);
887
- }
888
- }
889
-
890
- for (int i = 0; i < graph->n_nodes; i++) {
891
- struct lm_ggml_tensor * node = graph->nodes[i];
892
- int * node_backend_id = &tensor_backend_id(node);
893
- // do not overwrite user assignments
894
- if (*node_backend_id == -1) {
895
- *node_backend_id = lm_ggml_backend_sched_backend_id_from_cur(sched, node);
896
-
897
- #if 0
898
- // src
899
- if (node->op == LM_GGML_OP_NONE) {
900
- continue;
901
- }
902
-
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
- int * src_backend_id = &tensor_backend_id(src);
909
- if (*src_backend_id == -1) {
910
- *src_backend_id = lm_ggml_backend_sched_backend_id_from_cur(sched, src);
911
- }
912
- }
913
- #endif
914
- }
915
- }
916
-
917
- // pass 2: expand current backend assignments
918
- // assign the same backend to adjacent nodes
919
- // expand gpu backends (i.e. non last prio) up and down, ignoring cpu (the lowest priority backend)
920
- // thus, cpu will never be used unless weights are on cpu, or there are no gpu ops between cpu ops
921
- // ops unsupported by the backend being expanded will be left unassigned so that they can be assigned later when the locations of its inputs are known
922
- // expand gpu down
923
- {
924
- int cur_backend_id = -1;
925
- for (int i = 0; i < graph->n_nodes; i++) {
926
- struct lm_ggml_tensor * node = graph->nodes[i];
927
- if (lm_ggml_is_view_op(node->op)) {
928
- continue;
929
- }
930
- int * node_backend_id = &tensor_backend_id(node);
931
- if (*node_backend_id != -1) {
932
- if (*node_backend_id == sched->n_backends - 1) {
933
- // skip cpu (lowest prio backend)
934
- cur_backend_id = -1;
935
- } else {
936
- cur_backend_id = *node_backend_id;
937
- }
938
- } else if (cur_backend_id != -1) {
939
- lm_ggml_backend_sched_set_if_supported(sched, node, cur_backend_id, node_backend_id);
940
- }
941
- }
942
- }
943
- // expand gpu up
944
- {
945
- int cur_backend_id = -1;
946
- for (int i = graph->n_nodes - 1; i >= 0; i--) {
947
- struct lm_ggml_tensor * node = graph->nodes[i];
948
- if (lm_ggml_is_view_op(node->op)) {
949
- continue;
950
- }
951
- int * node_backend_id = &tensor_backend_id(node);
952
- if (*node_backend_id != -1) {
953
- if (*node_backend_id == sched->n_backends - 1) {
954
- // skip cpu (lowest prio backend)
955
- cur_backend_id = -1;
956
- } else {
957
- cur_backend_id = *node_backend_id;
958
- }
959
- } else if (cur_backend_id != -1) {
960
- lm_ggml_backend_sched_set_if_supported(sched, node, cur_backend_id, node_backend_id);
961
- }
962
- }
963
- }
964
- // expand rest down
965
- {
966
- int cur_backend_id = -1;
967
- for (int i = 0; i < graph->n_nodes; i++) {
968
- struct lm_ggml_tensor * node = graph->nodes[i];
969
- if (lm_ggml_is_view_op(node->op)) {
970
- continue;
971
- }
972
- int * node_backend_id = &tensor_backend_id(node);
973
- if (*node_backend_id != -1) {
974
- cur_backend_id = *node_backend_id;
975
- } else if (cur_backend_id != -1) {
976
- lm_ggml_backend_sched_set_if_supported(sched, node, cur_backend_id, node_backend_id);
977
- }
978
- }
979
- }
980
- // expand rest up
981
- {
982
- int cur_backend_id = -1;
983
- for (int i = graph->n_nodes - 1; i >= 0; i--) {
984
- struct lm_ggml_tensor * node = graph->nodes[i];
985
- if (lm_ggml_is_view_op(node->op)) {
986
- continue;
987
- }
988
- int * node_backend_id = &tensor_backend_id(node);
989
- if (*node_backend_id != -1) {
990
- cur_backend_id = *node_backend_id;
991
- } else if (cur_backend_id != -1) {
992
- lm_ggml_backend_sched_set_if_supported(sched, node, cur_backend_id, node_backend_id);
993
- }
994
- }
995
- }
996
-
997
- // pass 3: upgrade nodes to higher prio backends with compatible buffer types
998
- // if the tensor is already in the same buffer type (*) as another higher priority backend, we should move it there
999
- // however, we also need to verify that the sources are in compatible buffer types
1000
- // (*) the actual requirement is more relaxed, the buffer type of the backend should be supported by all the users of this tensor further down the graph
1001
- // however, this is slow to verify, so we have a more strict requirement that the buffer type is the same
1002
- // this is not uncommon since multiple backends can use host memory, with the same buffer type (eg. BLAS and CPU)
1003
- // additionally, set remaining unassigned nodes to the backend with the most supported inputs
1004
- // only nodes that could not be assigned during expansion due to the backend not supporting the op should be unassigned at this point
1005
- for (int i = 0; i < graph->n_nodes; i++) {
1006
- struct lm_ggml_tensor * node = graph->nodes[i];
1007
- if (lm_ggml_is_view_op(node->op)) {
1008
- continue;
1009
- }
1010
- int * node_backend_id = &tensor_backend_id(node);
1011
- if (*node_backend_id == -1) {
1012
- // unassigned node: find the backend with the most supported inputs
1013
- int n_supported_best = -1;
1014
- for (int b = 0; b < sched->n_backends; b++) {
1015
- if (lm_ggml_backend_supports_op(sched->backends[b], node)) {
1016
- int n_supported = 0;
1017
- for (int j = 0; j < LM_GGML_MAX_SRC; j++) {
1018
- struct lm_ggml_tensor * src = node->src[j];
1019
- if (src == NULL) {
1020
- continue;
1021
- }
1022
- if ((tensor_backend_id(src) != -1 || tensor_backend_id(src->view_src) != -1) && lm_ggml_backend_sched_buffer_supported(sched, src, b)) {
1023
- n_supported++;
1024
- }
1025
- }
1026
- if (n_supported > n_supported_best) {
1027
- n_supported_best = n_supported;
1028
- *node_backend_id = b;
1029
- SET_CAUSE(node, "3.best");
1030
- }
1031
- }
1032
- }
1033
- } else {
1034
- // assigned node: upgrade to higher prio backend if possible
1035
- for (int b = 0; b < *node_backend_id; b++) {
1036
- if (sched->bufts[b] == sched->bufts[*node_backend_id] && lm_ggml_backend_supports_op(sched->backends[b], node)) {
1037
- bool supported = true;
1038
- for (int j = 0; j < LM_GGML_MAX_SRC; j++) {
1039
- struct lm_ggml_tensor * src = node->src[j];
1040
- if (src == NULL) {
1041
- continue;
1042
- }
1043
- if (!lm_ggml_backend_sched_buffer_supported(sched, src, b)) {
1044
- supported = false;
1045
- break;
1046
- }
1047
- }
1048
- if (supported) {
1049
- *node_backend_id = b;
1050
- SET_CAUSE(node, "3.upg");
1051
- break;
1052
- }
1053
- }
1054
- }
1055
- }
1056
- }
1057
-
1058
- // pass 4: assign backends to remaining src from dst and view_src
1059
- for (int i = 0; i < graph->n_nodes; i++) {
1060
- struct lm_ggml_tensor * node = graph->nodes[i];
1061
- int * cur_backend_id = &tensor_backend_id(node);
1062
- if (node->view_src != NULL && *cur_backend_id == -1) {
1063
- *cur_backend_id = tensor_backend_id(node->view_src);
1064
- SET_CAUSE(node, "4.vsrc");
1065
- }
1066
- for (int j = 0; j < LM_GGML_MAX_SRC; j++) {
1067
- struct lm_ggml_tensor * src = node->src[j];
1068
- if (src == NULL) {
1069
- continue;
1070
- }
1071
- int * src_backend_id = &tensor_backend_id(src);
1072
- if (*src_backend_id == -1) {
1073
- if (src->view_src != NULL) {
1074
- // views are always on the same backend as the source
1075
- *src_backend_id = tensor_backend_id(src->view_src);
1076
- SET_CAUSE(src, "4.vsrc");
1077
- } else {
1078
- *src_backend_id = *cur_backend_id;
1079
- SET_CAUSE(src, "4.cur");
1080
- }
1081
- }
1082
- }
1083
- }
1084
-
1085
- // pass 5: split graph, find tensors that need to be copied
1086
- {
1087
- int i_split = 0;
1088
- struct lm_ggml_backend_sched_split * split = &sched->splits[0];
1089
- // find the backend of the first split, skipping view ops
1090
- int i = 0;
1091
- for (; i < graph->n_nodes; i++) {
1092
- struct lm_ggml_tensor * node = graph->nodes[i];
1093
- if (!lm_ggml_is_view_op(node->op)) {
1094
- split->backend_id = tensor_backend_id(node);
1095
- break;
1096
- }
1097
- }
1098
- split->i_start = 0;
1099
- split->n_inputs = 0;
1100
- int cur_backend_id = split->backend_id;
1101
- for (; i < graph->n_nodes; i++) {
1102
- struct lm_ggml_tensor * node = graph->nodes[i];
1103
-
1104
- if (lm_ggml_is_view_op(node->op)) {
1105
- continue;
1106
- }
1107
-
1108
- const int node_backend_id = tensor_backend_id(node);
1109
-
1110
- assert(node_backend_id != -1); // all nodes should be assigned by now
1111
-
1112
- // check if we should start a new split based on the sources of the current node
1113
- bool need_new_split = false;
1114
- if (node_backend_id == cur_backend_id && split->n_inputs > 0) {
1115
- for (int j = 0; j < LM_GGML_MAX_SRC; j++) {
1116
- struct lm_ggml_tensor * src = node->src[j];
1117
- if (src == NULL) {
1118
- continue;
1119
- }
1120
- // check if a weight is on a different and incompatible backend
1121
- // by starting a new split, the memory of the previously offloaded weights can be reused
1122
- if (src->buffer != NULL && src->buffer->usage == LM_GGML_BACKEND_BUFFER_USAGE_WEIGHTS) {
1123
- int src_backend_id = tensor_backend_id(src);
1124
- if (src_backend_id != cur_backend_id && !lm_ggml_backend_sched_buffer_supported(sched, src, cur_backend_id)) {
1125
- need_new_split = true;
1126
- break;
1127
- }
1128
- }
1129
- // check if the split has too many inputs
1130
- // FIXME: count the number of inputs instead of only checking when full
1131
- if (split->n_inputs == LM_GGML_SCHED_MAX_SPLIT_INPUTS) {
1132
- const size_t id = hash_id(src);
1133
- int src_backend_id = sched->hv_tensor_backend_ids[id];
1134
- bool supported = lm_ggml_backend_sched_buffer_supported(sched, src, cur_backend_id);
1135
- if (src_backend_id != cur_backend_id && tensor_id_copy(id, cur_backend_id, 0) == NULL && !supported) {
1136
- need_new_split = true;
1137
- break;
1138
- }
1139
- }
1140
- }
1141
- }
1142
-
1143
- if (node_backend_id != cur_backend_id || need_new_split) {
1144
- split->i_end = i;
1145
- i_split++;
1146
- if (i_split >= sched->splits_capacity) {
1147
- sched->splits_capacity *= 2;
1148
- sched->splits = (lm_ggml_backend_sched_split *)
1149
- realloc(sched->splits, sched->splits_capacity * sizeof(struct lm_ggml_backend_sched_split));
1150
- LM_GGML_ASSERT(sched->splits != NULL);
1151
- }
1152
- split = &sched->splits[i_split];
1153
- split->backend_id = node_backend_id;
1154
- split->i_start = i;
1155
- split->n_inputs = 0;
1156
- cur_backend_id = node_backend_id;
1157
- }
1158
-
1159
- // find inputs that are not on the same backend
1160
- for (int j = 0; j < LM_GGML_MAX_SRC; j++) {
1161
- struct lm_ggml_tensor * src = node->src[j];
1162
- if (src == NULL) {
1163
- continue;
1164
- }
1165
-
1166
- size_t src_id = hash_id(src);
1167
- const int src_backend_id = sched->hv_tensor_backend_ids[src_id];
1168
- assert(src_backend_id != -1); // all inputs should be assigned by now
1169
-
1170
- if (src->flags & LM_GGML_TENSOR_FLAG_INPUT && sched->n_copies > 1) {
1171
- if (tensor_id_copy(src_id, src_backend_id, 0) == NULL) {
1172
- lm_ggml_backend_t backend = sched->backends[src_backend_id];
1173
- for (int c = 0; c < sched->n_copies; c++) {
1174
- struct lm_ggml_tensor * tensor_copy;
1175
- if (c == sched->cur_copy) {
1176
- tensor_copy = src; // use the original tensor as the current copy
1177
- } else {
1178
- tensor_copy = lm_ggml_dup_tensor_layout(sched->ctx, src);
1179
- lm_ggml_format_name(tensor_copy, "%s#%s#%d", lm_ggml_backend_name(backend), src->name, c);
1180
- }
1181
- if (sched->n_copies > 1) {
1182
- lm_ggml_set_input(tensor_copy);
1183
- lm_ggml_set_output(tensor_copy); // prevent ggml-alloc from overwriting the tensor
1184
- }
1185
- tensor_id_copy(src_id, src_backend_id, c) = tensor_copy;
1186
- SET_CAUSE(tensor_copy, "4.cpy");
1187
- }
1188
- int n_graph_inputs = sched->n_graph_inputs++;
1189
- LM_GGML_ASSERT(n_graph_inputs < LM_GGML_SCHED_MAX_SPLIT_INPUTS);
1190
- sched->graph_inputs[n_graph_inputs] = src;
1191
- }
1192
- }
1193
-
1194
- if (src_backend_id != cur_backend_id && !lm_ggml_backend_sched_buffer_supported(sched, src, cur_backend_id)) {
1195
- // create a copy of the input in the split's backend
1196
- if (tensor_id_copy(src_id, cur_backend_id, 0) == NULL) {
1197
- lm_ggml_backend_t backend = sched->backends[cur_backend_id];
1198
- for (int c = 0; c < sched->n_copies; c++) {
1199
- struct lm_ggml_tensor * tensor_copy = lm_ggml_dup_tensor_layout(sched->ctx, src);
1200
- lm_ggml_format_name(tensor_copy, "%s#%s#%d", lm_ggml_backend_name(backend), src->name, c);
1201
- if (sched->n_copies > 1) {
1202
- lm_ggml_set_input(tensor_copy);
1203
- lm_ggml_set_output(tensor_copy); // prevent ggml-alloc from overwriting the tensor
1204
- }
1205
- tensor_id_copy(src_id, cur_backend_id, c) = tensor_copy;
1206
- SET_CAUSE(tensor_copy, "4.cpy");
1207
- }
1208
- int n_inputs = split->n_inputs++;
1209
- LM_GGML_ASSERT(n_inputs < LM_GGML_SCHED_MAX_SPLIT_INPUTS);
1210
- split->inputs[n_inputs] = src;
1211
- }
1212
- node->src[j] = tensor_id_copy(src_id, cur_backend_id, sched->cur_copy);
1213
- }
1214
- }
1215
- }
1216
- split->i_end = graph->n_nodes;
1217
- sched->n_splits = i_split + 1;
1218
- }
1219
-
1220
- if (sched->debug) {
1221
- lm_ggml_backend_sched_print_assignments(sched, graph);
1222
- }
1223
-
1224
- // swap node_backend_ids and leaf _backend_ids with prevs
1225
- {
1226
- int * tmp = sched->node_backend_ids;
1227
- sched->node_backend_ids = sched->prev_node_backend_ids;
1228
- sched->prev_node_backend_ids = tmp;
1229
-
1230
- tmp = sched->leaf_backend_ids;
1231
- sched->leaf_backend_ids = sched->prev_leaf_backend_ids;
1232
- sched->prev_leaf_backend_ids = tmp;
1233
- }
1234
-
1235
- int graph_size = std::max(graph->n_nodes, graph->n_leafs) + sched->n_splits*LM_GGML_SCHED_MAX_SPLIT_INPUTS*2*sched->n_copies;
1236
- if (sched->graph.size < graph_size) {
1237
- sched->graph.size = graph_size;
1238
- sched->graph.nodes = (lm_ggml_tensor **) realloc(sched->graph.nodes, graph_size * sizeof(struct lm_ggml_tensor *));
1239
- sched->graph.leafs = (lm_ggml_tensor **) realloc(sched->graph.leafs, graph_size * sizeof(struct lm_ggml_tensor *));
1240
- LM_GGML_ASSERT(sched->graph.nodes != NULL);
1241
- LM_GGML_ASSERT(sched->graph.leafs != NULL);
1242
- }
1243
- sched->graph.n_nodes = 0;
1244
- sched->graph.n_leafs = 0;
1245
-
1246
- struct lm_ggml_cgraph * graph_copy = &sched->graph;
1247
-
1248
- for (int i = 0; i < sched->n_splits; i++) {
1249
- struct lm_ggml_backend_sched_split * split = &sched->splits[i];
1250
- split->graph = lm_ggml_graph_view(graph, split->i_start, split->i_end);
1251
-
1252
- // add inputs to the graph copy so that they are allocated by ggml-alloc at the start of the split
1253
- for (int j = 0; j < split->n_inputs; j++) {
1254
- assert(graph_copy->size > (graph_copy->n_nodes + 1));
1255
-
1256
- struct lm_ggml_tensor * input = split->inputs[j];
1257
- const size_t input_id = hash_id(input);
1258
- struct lm_ggml_tensor * input_cpy = tensor_id_copy(input_id, split->backend_id, sched->cur_copy);
1259
-
1260
- // add a dependency to the input source so that it is not freed before the copy is done
1261
- struct lm_ggml_tensor * input_dep = lm_ggml_view_tensor(sched->ctx, input);
1262
- input_dep->src[0] = input;
1263
- sched->node_backend_ids[graph_copy->n_nodes] = sched->hv_tensor_backend_ids[input_id];
1264
- graph_copy->nodes[graph_copy->n_nodes++] = input_dep;
1265
-
1266
- // add a dependency to the input copy so that it is allocated at the start of the split
1267
- sched->node_backend_ids[graph_copy->n_nodes] = split->backend_id;
1268
- graph_copy->nodes[graph_copy->n_nodes++] = input_cpy;
1269
- }
1270
-
1271
- for (int j = split->i_start; j < split->i_end; j++) {
1272
- assert(graph_copy->size > graph_copy->n_nodes);
1273
- sched->node_backend_ids[graph_copy->n_nodes] = tensor_backend_id(graph->nodes[j]);
1274
- graph_copy->nodes[graph_copy->n_nodes++] = graph->nodes[j];
1275
- }
1276
- }
1277
-
1278
- if (sched->n_copies > 1) {
1279
- // add input copies as leafs so that they are allocated first
1280
- for (int i = 0; i < sched->n_graph_inputs; i++) {
1281
- struct lm_ggml_tensor * input = sched->graph_inputs[i];
1282
- size_t id = hash_id(input);
1283
- int backend_id = tensor_backend_id(input);
1284
- for (int c = 0; c < sched->n_copies; c++) {
1285
- struct lm_ggml_tensor * input_cpy = tensor_id_copy(id, backend_id, c);
1286
- sched->leaf_backend_ids[graph_copy->n_leafs] = backend_id;
1287
- assert(graph_copy->size > graph_copy->n_leafs);
1288
- graph_copy->leafs[graph_copy->n_leafs++] = input_cpy;
1289
- }
1290
- }
1291
-
1292
- for (int i = 0; i < sched->n_splits; i++) {
1293
- struct lm_ggml_backend_sched_split * split = &sched->splits[i];
1294
- int backend_id = split->backend_id;
1295
- for (int j = 0; j < split->n_inputs; j++) {
1296
- struct lm_ggml_tensor * input = split->inputs[j];
1297
- size_t id = hash_id(input);
1298
- for (int c = 0; c < sched->n_copies; c++) {
1299
- struct lm_ggml_tensor * input_cpy = tensor_id_copy(id, backend_id, c);
1300
- sched->leaf_backend_ids[graph_copy->n_leafs] = backend_id;
1301
- assert(graph_copy->size > graph_copy->n_leafs);
1302
- graph_copy->leafs[graph_copy->n_leafs++] = input_cpy;
1303
- }
1304
- }
1305
- }
1306
- }
1307
-
1308
- // add leafs from the original graph
1309
- for (int i = 0; i < graph->n_leafs; i++) {
1310
- struct lm_ggml_tensor * leaf = graph->leafs[i];
1311
- sched->leaf_backend_ids[graph_copy->n_leafs] = tensor_backend_id(leaf);
1312
- assert(graph_copy->size > graph_copy->n_leafs);
1313
- graph_copy->leafs[graph_copy->n_leafs++] = leaf;
1314
- }
1315
- }
1316
-
1317
- static bool lm_ggml_backend_sched_alloc_splits(lm_ggml_backend_sched_t sched) {
1318
- bool backend_ids_changed = false;
1319
- for (int i = 0; i < sched->graph.n_nodes; i++) {
1320
- if (sched->node_backend_ids[i] != sched->prev_node_backend_ids[i] &&
1321
- sched->bufts[sched->node_backend_ids[i]] != sched->bufts[sched->prev_node_backend_ids[i]]) {
1322
- backend_ids_changed = true;
1323
- break;
1324
- }
1325
- }
1326
- if (!backend_ids_changed) {
1327
- for (int i = 0; i < sched->graph.n_leafs; i++) {
1328
- if (sched->leaf_backend_ids[i] != sched->prev_leaf_backend_ids[i] &&
1329
- sched->bufts[sched->leaf_backend_ids[i]] != sched->bufts[sched->prev_leaf_backend_ids[i]]) {
1330
- backend_ids_changed = true;
1331
- break;
1332
- }
1333
- }
1334
- }
1335
-
1336
- // allocate graph
1337
- if (backend_ids_changed || !lm_ggml_gallocr_alloc_graph(sched->galloc, &sched->graph)) {
1338
- // the re-allocation may cause the split inputs to be moved to a different address
1339
- lm_ggml_backend_sched_synchronize(sched);
1340
- #ifndef NDEBUG
1341
- LM_GGML_LOG_DEBUG("%s: failed to allocate graph, reserving (backend_ids_changed = %d)\n", __func__, backend_ids_changed);
1342
- #endif
1343
- lm_ggml_gallocr_reserve_n(sched->galloc, &sched->graph, sched->node_backend_ids, sched->leaf_backend_ids);
1344
- if (!lm_ggml_gallocr_alloc_graph(sched->galloc, &sched->graph)) {
1345
- LM_GGML_LOG_ERROR("%s: failed to allocate graph\n", __func__);
1346
- return false;
1347
- }
1348
- }
1349
-
1350
- return true;
1351
- }
1352
-
1353
- static enum lm_ggml_status lm_ggml_backend_sched_compute_splits(lm_ggml_backend_sched_t sched) {
1354
- struct lm_ggml_backend_sched_split * splits = sched->splits;
1355
-
1356
- for (int i = 0; i < sched->n_splits; i++) {
1357
- struct lm_ggml_backend_sched_split * split = &splits[i];
1358
- int split_backend_id = split->backend_id;
1359
- lm_ggml_backend_t split_backend = sched->backends[split_backend_id];
1360
-
1361
- // copy the input tensors to the split backend
1362
- for (int j = 0; j < split->n_inputs; j++) {
1363
- lm_ggml_backend_t input_backend = lm_ggml_backend_sched_get_tensor_backend(sched, split->inputs[j]);
1364
- struct lm_ggml_tensor * input = split->inputs[j];
1365
- struct lm_ggml_tensor * input_cpy = tensor_copy(input, split_backend_id, sched->cur_copy);
1366
-
1367
- if (input->flags & LM_GGML_TENSOR_FLAG_INPUT) {
1368
- // inputs from the user must be copied immediately to prevent the user overwriting the data before the copy is done
1369
- if (sched->events[split_backend_id][sched->cur_copy] != NULL) {
1370
- lm_ggml_backend_event_synchronize(sched->events[split_backend_id][sched->cur_copy]);
1371
- } else {
1372
- lm_ggml_backend_synchronize(split_backend);
1373
- }
1374
- lm_ggml_backend_tensor_copy(input, input_cpy);
1375
- } else {
1376
- // wait for the split backend to finish using the input before overwriting it
1377
- if (sched->events[split_backend_id][sched->cur_copy] != NULL) {
1378
- lm_ggml_backend_event_wait(split_backend, sched->events[split_backend_id][sched->cur_copy]);
1379
- } else {
1380
- lm_ggml_backend_synchronize(split_backend);
1381
- }
1382
- // try async copy, but if not possible, we can still use a sync copy without synchronizing the dst backend, since we handle the synchronization here with multiple copies and events
1383
- // TODO: add public function to facilitate this, since applications do not have direct access to the backend interface
1384
- if (!split_backend->iface.cpy_tensor_async || !split_backend->iface.cpy_tensor_async(input_backend, split_backend, input, input_cpy)) {
1385
- lm_ggml_backend_synchronize(input_backend);
1386
- if (sched->events[split_backend_id][sched->cur_copy] != NULL) {
1387
- lm_ggml_backend_event_synchronize(sched->events[split_backend_id][sched->cur_copy]);
1388
- } else {
1389
- lm_ggml_backend_synchronize(split_backend);
1390
- }
1391
- lm_ggml_backend_tensor_copy(input, input_cpy);
1392
- }
1393
- }
1394
- }
1395
-
1396
- if (!sched->callback_eval) {
1397
- enum lm_ggml_status ec = lm_ggml_backend_graph_compute_async(split_backend, &split->graph);
1398
- if (ec != LM_GGML_STATUS_SUCCESS) {
1399
- return ec;
1400
- }
1401
- } else {
1402
- // similar to lm_ggml_backend_compare_graph_backend
1403
- for (int j0 = 0; j0 < split->graph.n_nodes; j0++) {
1404
- struct lm_ggml_tensor * t = split->graph.nodes[j0];
1405
-
1406
- // check if the user needs data from this node
1407
- bool need = sched->callback_eval(t, true, sched->callback_eval_user_data);
1408
-
1409
- int j1 = j0;
1410
-
1411
- // determine the range [j0, j1] of nodes that can be computed together
1412
- while (!need && j1 < split->graph.n_nodes - 1) {
1413
- t = split->graph.nodes[++j1];
1414
- need = sched->callback_eval(t, true, sched->callback_eval_user_data);
1415
- }
1416
-
1417
- struct lm_ggml_cgraph gv = lm_ggml_graph_view(&split->graph, j0, j1 + 1);
1418
-
1419
- enum lm_ggml_status ec = lm_ggml_backend_graph_compute_async(split_backend, &gv);
1420
- if (ec != LM_GGML_STATUS_SUCCESS) {
1421
- return ec;
1422
- }
1423
-
1424
- // TODO: pass backend to the callback, then the user can decide if they want to synchronize
1425
- lm_ggml_backend_synchronize(split_backend);
1426
-
1427
- if (need && !sched->callback_eval(t, false, sched->callback_eval_user_data)) {
1428
- break;
1429
- }
1430
-
1431
- j0 = j1;
1432
- }
1433
- }
1434
-
1435
- // record the event of this copy
1436
- if (split->n_inputs > 0) {
1437
- if (sched->events[split_backend_id][sched->cur_copy] != NULL) {
1438
- lm_ggml_backend_event_record(sched->events[split_backend_id][sched->cur_copy], split_backend);
1439
- }
1440
- }
1441
- }
1442
-
1443
- sched->cur_copy = (sched->cur_copy + 1) % sched->n_copies;
1444
-
1445
- return LM_GGML_STATUS_SUCCESS;
1446
- }
1447
-
1448
- lm_ggml_backend_sched_t lm_ggml_backend_sched_new(
1449
- lm_ggml_backend_t * backends,
1450
- lm_ggml_backend_buffer_type_t * bufts,
1451
- int n_backends,
1452
- size_t graph_size,
1453
- bool parallel) {
1454
- LM_GGML_ASSERT(n_backends > 0);
1455
- LM_GGML_ASSERT(n_backends <= LM_GGML_SCHED_MAX_BACKENDS);
1456
- LM_GGML_ASSERT(lm_ggml_backend_dev_type(lm_ggml_backend_get_device(backends[n_backends - 1])) == LM_GGML_BACKEND_DEVICE_TYPE_CPU);
1457
-
1458
- struct lm_ggml_backend_sched * sched = (lm_ggml_backend_sched *) calloc(1, sizeof(struct lm_ggml_backend_sched));
1459
-
1460
- const char * LM_GGML_SCHED_DEBUG = getenv("LM_GGML_SCHED_DEBUG");
1461
- sched->debug = LM_GGML_SCHED_DEBUG ? atoi(LM_GGML_SCHED_DEBUG) : 0;
1462
- sched->n_backends = n_backends;
1463
- sched->n_copies = parallel ? LM_GGML_SCHED_MAX_COPIES : 1;
1464
-
1465
- // initialize hash table
1466
- // FIXME: needs to be size*2 to account for leafs (do it in graph_split instead)
1467
- sched->hash_set = lm_ggml_hash_set_new(graph_size);
1468
- sched->hv_tensor_backend_ids = (int *) malloc(sched->hash_set.size * sizeof(sched->hv_tensor_backend_ids[0]));
1469
- sched->hv_tensor_copies = (lm_ggml_tensor **) malloc(sched->hash_set.size * sched->n_backends * sched->n_copies * sizeof(struct lm_ggml_tensor *));
1470
-
1471
- const size_t lm_ggml_sched_max_splits = graph_size; // at most there is one split for each node in the graph
1472
- const size_t nodes_size = graph_size + lm_ggml_sched_max_splits*LM_GGML_SCHED_MAX_SPLIT_INPUTS*2;
1473
- sched->node_backend_ids = (int *) calloc(nodes_size, sizeof(sched->node_backend_ids[0]));
1474
- sched->leaf_backend_ids = (int *) calloc(nodes_size, sizeof(sched->leaf_backend_ids[0]));
1475
- sched->prev_node_backend_ids = (int *) calloc(nodes_size, sizeof(sched->prev_node_backend_ids[0]));
1476
- sched->prev_leaf_backend_ids = (int *) calloc(nodes_size, sizeof(sched->prev_leaf_backend_ids[0]));
1477
-
1478
- sched->context_buffer_size = lm_ggml_sched_max_splits*LM_GGML_SCHED_MAX_SPLIT_INPUTS*2*sizeof(struct lm_ggml_tensor) + lm_ggml_graph_overhead_custom(graph_size, false);
1479
- sched->context_buffer = (char *) malloc(sched->context_buffer_size);
1480
-
1481
- const int initial_splits_capacity = 16;
1482
- sched->splits = (lm_ggml_backend_sched_split *) calloc(initial_splits_capacity, sizeof(sched->splits[0]));
1483
- sched->splits_capacity = initial_splits_capacity;
1484
-
1485
- for (int b = 0; b < n_backends; b++) {
1486
- sched->backends[b] = backends[b];
1487
- sched->bufts[b] = bufts ? bufts[b] : lm_ggml_backend_get_default_buffer_type(backends[b]);
1488
- LM_GGML_ASSERT(lm_ggml_backend_supports_buft(backends[b], sched->bufts[b]));
1489
-
1490
- if (sched->n_copies > 1) {
1491
- for (int c = 0; c < sched->n_copies; c++) {
1492
- sched->events[b][c] = lm_ggml_backend_event_new(backends[b]->device);
1493
- }
1494
- }
1495
- }
1496
-
1497
- sched->galloc = lm_ggml_gallocr_new_n(sched->bufts, n_backends);
1498
-
1499
- lm_ggml_backend_sched_reset(sched);
1500
-
1501
- return sched;
1502
- }
1503
-
1504
- void lm_ggml_backend_sched_free(lm_ggml_backend_sched_t sched) {
1505
- if (sched == NULL) {
1506
- return;
1507
- }
1508
- for (int b = 0; b < sched->n_backends; b++) {
1509
- for (int c = 0; c < sched->n_copies; c++) {
1510
- lm_ggml_backend_event_free(sched->events[b][c]);
1511
- }
1512
- }
1513
- lm_ggml_gallocr_free(sched->galloc);
1514
- lm_ggml_free(sched->ctx);
1515
- lm_ggml_hash_set_free(&sched->hash_set);
1516
- free(sched->splits);
1517
- free(sched->hv_tensor_backend_ids);
1518
- free(sched->hv_tensor_copies);
1519
- free(sched->node_backend_ids);
1520
- free(sched->leaf_backend_ids);
1521
- free(sched->prev_node_backend_ids);
1522
- free(sched->prev_leaf_backend_ids);
1523
- free(sched->context_buffer);
1524
- free(sched->graph.nodes);
1525
- free(sched->graph.leafs);
1526
- free(sched);
1527
- }
1528
-
1529
- void lm_ggml_backend_sched_reset(lm_ggml_backend_sched_t sched) {
1530
- // reset state for the next run
1531
- if (!sched->is_reset) {
1532
- lm_ggml_hash_set_reset(&sched->hash_set);
1533
- memset(sched->hv_tensor_backend_ids, -1, sched->hash_set.size * sizeof(sched->hv_tensor_backend_ids[0]));
1534
- memset(sched->hv_tensor_copies, 0, sched->hash_set.size * sched->n_backends * sched->n_copies * sizeof(struct lm_ggml_tensor *));
1535
- sched->is_reset = true;
1536
- }
1537
- sched->is_alloc = false;
1538
- }
1539
-
1540
- bool lm_ggml_backend_sched_reserve(lm_ggml_backend_sched_t sched, struct lm_ggml_cgraph * measure_graph) {
1541
- LM_GGML_ASSERT((int)sched->hash_set.size >= measure_graph->n_nodes + measure_graph->n_leafs);
1542
-
1543
- lm_ggml_backend_sched_split_graph(sched, measure_graph);
1544
-
1545
- lm_ggml_backend_sched_synchronize(sched);
1546
-
1547
- if (!lm_ggml_gallocr_reserve_n(sched->galloc, &sched->graph, sched->node_backend_ids, sched->leaf_backend_ids)) {
1548
- return false;
1549
- }
1550
-
1551
- lm_ggml_backend_sched_reset(sched);
1552
-
1553
- return true;
1554
- }
1555
-
1556
- bool lm_ggml_backend_sched_alloc_graph(lm_ggml_backend_sched_t sched, struct lm_ggml_cgraph * graph) {
1557
- LM_GGML_ASSERT((int)sched->hash_set.size >= graph->n_nodes + graph->n_leafs);
1558
-
1559
- lm_ggml_backend_sched_split_graph(sched, graph);
1560
-
1561
-
1562
- if (!lm_ggml_backend_sched_alloc_splits(sched)) {
1563
- return false;
1564
- }
1565
-
1566
- sched->is_alloc = true;
1567
-
1568
- return true;
1569
- }
1570
-
1571
- enum lm_ggml_status lm_ggml_backend_sched_graph_compute(lm_ggml_backend_sched_t sched, struct lm_ggml_cgraph * graph) {
1572
- enum lm_ggml_status err = lm_ggml_backend_sched_graph_compute_async(sched, graph);
1573
- lm_ggml_backend_sched_synchronize(sched);
1574
- return err;
1575
- }
1576
-
1577
- enum lm_ggml_status lm_ggml_backend_sched_graph_compute_async(lm_ggml_backend_sched_t sched, struct lm_ggml_cgraph * graph) {
1578
- if (!sched->is_reset && !sched->is_alloc) {
1579
- lm_ggml_backend_sched_reset(sched);
1580
- }
1581
-
1582
- if (!sched->is_alloc) {
1583
- if (!lm_ggml_backend_sched_alloc_graph(sched, graph)) {
1584
- return LM_GGML_STATUS_ALLOC_FAILED;
1585
- }
1586
- }
1587
-
1588
- return lm_ggml_backend_sched_compute_splits(sched);
1589
- }
1590
-
1591
- void lm_ggml_backend_sched_synchronize(lm_ggml_backend_sched_t sched) {
1592
- for (int i = 0; i < sched->n_backends; i++) {
1593
- lm_ggml_backend_synchronize(sched->backends[i]);
1594
- }
1595
- }
1596
-
1597
- void lm_ggml_backend_sched_set_eval_callback(lm_ggml_backend_sched_t sched, lm_ggml_backend_sched_eval_callback callback, void * user_data) {
1598
- sched->callback_eval = callback;
1599
- sched->callback_eval_user_data = user_data;
1600
- }
1601
-
1602
- int lm_ggml_backend_sched_get_n_splits(lm_ggml_backend_sched_t sched) {
1603
- return sched->n_splits;
1604
- }
1605
-
1606
- int lm_ggml_backend_sched_get_n_copies(lm_ggml_backend_sched_t sched) {
1607
- return sched->n_copies;
1608
- }
1609
-
1610
- int lm_ggml_backend_sched_get_n_backends(lm_ggml_backend_sched_t sched) {
1611
- return sched->n_backends;
1612
- }
1613
-
1614
- lm_ggml_backend_t lm_ggml_backend_sched_get_backend(lm_ggml_backend_sched_t sched, int i) {
1615
- LM_GGML_ASSERT(i >= 0 && i < sched->n_backends);
1616
- return sched->backends[i];
1617
- }
1618
-
1619
- size_t lm_ggml_backend_sched_get_buffer_size(lm_ggml_backend_sched_t sched, lm_ggml_backend_t backend) {
1620
- int backend_index = lm_ggml_backend_sched_backend_id(sched, backend);
1621
- LM_GGML_ASSERT(backend_index >= 0 && backend_index < sched->n_backends);
1622
-
1623
- return lm_ggml_gallocr_get_buffer_size(sched->galloc, backend_index);
1624
- }
1625
-
1626
- void lm_ggml_backend_sched_set_tensor_backend(lm_ggml_backend_sched_t sched, struct lm_ggml_tensor * node, lm_ggml_backend_t backend) {
1627
- int backend_index = lm_ggml_backend_sched_backend_id(sched, backend);
1628
- LM_GGML_ASSERT(backend_index >= 0 && backend_index < sched->n_backends);
1629
- tensor_backend_id(node) = backend_index;
1630
- SET_CAUSE(node, "usr");
1631
- sched->is_reset = false;
1632
- }
1633
-
1634
- lm_ggml_backend_t lm_ggml_backend_sched_get_tensor_backend(lm_ggml_backend_sched_t sched, struct lm_ggml_tensor * node) {
1635
- int backend_index = tensor_backend_id(node);
1636
- if (backend_index == -1) {
1637
- return NULL;
1638
- }
1639
- return sched->backends[backend_index];
1640
- }
1641
-
1642
- // utils
1643
-
1644
- void lm_ggml_backend_view_init(struct lm_ggml_tensor * tensor) {
1645
- LM_GGML_ASSERT(tensor->buffer == NULL);
1646
- LM_GGML_ASSERT(tensor->view_src != NULL);
1647
- LM_GGML_ASSERT(tensor->view_src->buffer != NULL);
1648
- LM_GGML_ASSERT(tensor->view_src->data != NULL);
1649
-
1650
- tensor->buffer = tensor->view_src->buffer;
1651
- tensor->data = (char *)tensor->view_src->data + tensor->view_offs;
1652
- lm_ggml_backend_buffer_init_tensor(tensor->buffer, tensor);
1653
- }
1654
-
1655
- void lm_ggml_backend_tensor_alloc(lm_ggml_backend_buffer_t buffer, struct lm_ggml_tensor * tensor, void * addr) {
1656
- LM_GGML_ASSERT(tensor->buffer == NULL);
1657
- LM_GGML_ASSERT(tensor->data == NULL);
1658
- LM_GGML_ASSERT(tensor->view_src == NULL);
1659
- LM_GGML_ASSERT(addr >= lm_ggml_backend_buffer_get_base(buffer));
1660
- LM_GGML_ASSERT((char *)addr + lm_ggml_backend_buffer_get_alloc_size(buffer, tensor) <=
1661
- (char *)lm_ggml_backend_buffer_get_base(buffer) + lm_ggml_backend_buffer_get_size(buffer));
1662
-
1663
- tensor->buffer = buffer;
1664
- tensor->data = addr;
1665
- lm_ggml_backend_buffer_init_tensor(buffer, tensor);
1666
- }
1667
-
1668
- static struct lm_ggml_tensor * graph_copy_dup_tensor(struct lm_ggml_hash_set hash_set, struct lm_ggml_tensor ** node_copies,
1669
- struct lm_ggml_context * ctx_allocated, struct lm_ggml_context * ctx_unallocated, struct lm_ggml_tensor * src) {
1670
-
1671
- LM_GGML_ASSERT(src != NULL);
1672
- LM_GGML_ASSERT(src->data && "graph must be allocated");
1673
-
1674
- size_t id = lm_ggml_hash_insert(&hash_set, src);
1675
- if (id == LM_GGML_HASHSET_ALREADY_EXISTS) {
1676
- return node_copies[lm_ggml_hash_find(&hash_set, src)];
1677
- }
1678
-
1679
- struct lm_ggml_tensor * dst = lm_ggml_dup_tensor_layout(src->data && !src->view_src ? ctx_allocated : ctx_unallocated, src);
1680
- if (src->view_src != NULL) {
1681
- dst->view_src = graph_copy_dup_tensor(hash_set, node_copies, ctx_allocated, ctx_unallocated, src->view_src);
1682
- dst->view_offs = src->view_offs;
1683
- }
1684
- dst->op = src->op;
1685
- memcpy(dst->op_params, src->op_params, sizeof(dst->op_params));
1686
- lm_ggml_set_name(dst, src->name);
1687
-
1688
- // copy src
1689
- for (int i = 0; i < LM_GGML_MAX_SRC; i++) {
1690
- struct lm_ggml_tensor * s = src->src[i];
1691
- if (s == NULL) {
1692
- continue;
1693
- }
1694
- dst->src[i] = graph_copy_dup_tensor(hash_set, node_copies, ctx_allocated, ctx_unallocated, s);
1695
- }
1696
-
1697
- node_copies[id] = dst;
1698
- return dst;
1699
- }
1700
-
1701
- static void graph_copy_init_tensor(struct lm_ggml_hash_set * hash_set, struct lm_ggml_tensor ** node_copies, bool * node_init, struct lm_ggml_tensor * src) {
1702
- size_t id = lm_ggml_hash_find(hash_set, src);
1703
- if (node_init[id]) {
1704
- return;
1705
- }
1706
- node_init[id] = true;
1707
-
1708
- struct lm_ggml_tensor * dst = node_copies[id];
1709
- if (dst->view_src != NULL) {
1710
- graph_copy_init_tensor(hash_set, node_copies, node_init, src->view_src);
1711
- lm_ggml_backend_view_init(dst);
1712
- }
1713
- else {
1714
- lm_ggml_backend_tensor_copy(src, dst);
1715
- }
1716
-
1717
- // init src
1718
- for (int i = 0; i < LM_GGML_MAX_SRC; i++) {
1719
- struct lm_ggml_tensor * s = src->src[i];
1720
- if (s == NULL) {
1721
- continue;
1722
- }
1723
- graph_copy_init_tensor(hash_set, node_copies, node_init, s);
1724
- }
1725
- }
1726
-
1727
- struct lm_ggml_backend_graph_copy lm_ggml_backend_graph_copy(lm_ggml_backend_t backend, struct lm_ggml_cgraph * graph) {
1728
- struct lm_ggml_hash_set hash_set = lm_ggml_hash_set_new(graph->visited_hash_set.size);
1729
- struct lm_ggml_tensor ** node_copies = (lm_ggml_tensor **) calloc(hash_set.size, sizeof(node_copies[0])); // NOLINT
1730
- bool * node_init = (bool *) calloc(hash_set.size, sizeof(node_init[0]));
1731
-
1732
- struct lm_ggml_init_params params = {
1733
- /* .mem_size = */ lm_ggml_tensor_overhead()*hash_set.size + lm_ggml_graph_overhead_custom(graph->size, false),
1734
- /* .mem_buffer = */ NULL,
1735
- /* .no_alloc = */ true
1736
- };
1737
-
1738
- struct lm_ggml_context * ctx_allocated = lm_ggml_init(params);
1739
- struct lm_ggml_context * ctx_unallocated = lm_ggml_init(params);
1740
-
1741
- if (ctx_allocated == NULL || ctx_unallocated == NULL) {
1742
- LM_GGML_LOG_ERROR("%s: failed to allocate context for graph copy\n", __func__);
1743
- lm_ggml_hash_set_free(&hash_set);
1744
- free(node_copies);
1745
- free(node_init);
1746
- lm_ggml_free(ctx_allocated);
1747
- lm_ggml_free(ctx_unallocated);
1748
- return {
1749
- /* .buffer = */ NULL,
1750
- /* .ctx_allocated = */ NULL,
1751
- /* .ctx_unallocated = */ NULL,
1752
- /* .graph = */ NULL,
1753
- };
1754
- }
1755
-
1756
- // dup nodes
1757
- for (int i = 0; i < graph->n_nodes; i++) {
1758
- struct lm_ggml_tensor * node = graph->nodes[i];
1759
- graph_copy_dup_tensor(hash_set, node_copies, ctx_allocated, ctx_unallocated, node);
1760
- }
1761
-
1762
- // allocate nodes
1763
- lm_ggml_backend_buffer_t buffer = lm_ggml_backend_alloc_ctx_tensors(ctx_allocated, backend);
1764
- if (buffer == NULL) {
1765
- LM_GGML_LOG_ERROR("%s: failed to allocate buffer for graph copy\n", __func__);
1766
- lm_ggml_hash_set_free(&hash_set);
1767
- free(node_copies);
1768
- free(node_init);
1769
- lm_ggml_free(ctx_allocated);
1770
- lm_ggml_free(ctx_unallocated);
1771
- return {
1772
- /* .buffer = */ NULL,
1773
- /* .ctx_allocated = */ NULL,
1774
- /* .ctx_unallocated = */ NULL,
1775
- /* .graph = */ NULL,
1776
- };
1777
- }
1778
-
1779
- //printf("copy buffer size: %zu MB\n", lm_ggml_backend_buffer_get_size(buffer) / 1024 / 1024);
1780
-
1781
- // copy data and init views
1782
- for (int i = 0; i < graph->n_nodes; i++) {
1783
- struct lm_ggml_tensor * node = graph->nodes[i];
1784
- graph_copy_init_tensor(&hash_set, node_copies, node_init, node);
1785
- }
1786
-
1787
- // build graph copy
1788
- struct lm_ggml_cgraph * graph_copy = lm_ggml_new_graph_custom(ctx_allocated, graph->size, false);
1789
- for (int i = 0; i < graph->n_nodes; i++) {
1790
- struct lm_ggml_tensor * node = graph->nodes[i];
1791
- struct lm_ggml_tensor * node_copy = node_copies[lm_ggml_hash_find(&hash_set, node)];
1792
- graph_copy->nodes[i] = node_copy;
1793
- }
1794
- graph_copy->n_nodes = graph->n_nodes;
1795
-
1796
- lm_ggml_hash_set_free(&hash_set);
1797
- free(node_copies);
1798
- free(node_init);
1799
-
1800
- return {
1801
- /* .buffer = */ buffer,
1802
- /* .ctx_allocated = */ ctx_allocated,
1803
- /* .ctx_unallocated = */ ctx_unallocated,
1804
- /* .graph = */ graph_copy,
1805
- };
1806
- }
1807
-
1808
- void lm_ggml_backend_graph_copy_free(struct lm_ggml_backend_graph_copy copy) {
1809
- lm_ggml_backend_buffer_free(copy.buffer);
1810
- lm_ggml_free(copy.ctx_allocated);
1811
- lm_ggml_free(copy.ctx_unallocated);
1812
- }
1813
-
1814
- bool lm_ggml_backend_compare_graph_backend(lm_ggml_backend_t backend1, lm_ggml_backend_t backend2, struct lm_ggml_cgraph * graph, lm_ggml_backend_eval_callback callback, void * user_data) {
1815
- struct lm_ggml_backend_graph_copy copy = lm_ggml_backend_graph_copy(backend2, graph);
1816
- if (copy.buffer == NULL) {
1817
- return false;
1818
- }
1819
-
1820
- struct lm_ggml_cgraph * g1 = graph;
1821
- struct lm_ggml_cgraph * g2 = copy.graph;
1822
-
1823
- assert(g1->n_nodes == g2->n_nodes);
1824
-
1825
- for (int i = 0; i < g1->n_nodes; i++) {
1826
- //printf("eval %d/%d\n", i, g1->n_nodes);
1827
- struct lm_ggml_tensor * t1 = g1->nodes[i];
1828
- struct lm_ggml_tensor * t2 = g2->nodes[i];
1829
-
1830
- assert(t1->op == t2->op && lm_ggml_are_same_layout(t1, t2));
1831
-
1832
- struct lm_ggml_cgraph g1v = lm_ggml_graph_view(g1, i, i + 1);
1833
- struct lm_ggml_cgraph g2v = lm_ggml_graph_view(g2, i, i + 1);
1834
-
1835
- lm_ggml_backend_graph_compute(backend1, &g1v);
1836
- lm_ggml_backend_graph_compute(backend2, &g2v);
1837
-
1838
- if (lm_ggml_is_view_op(t1->op)) {
1839
- continue;
1840
- }
1841
-
1842
- // compare results, calculate rms etc
1843
- if (!callback(i, t1, t2, user_data)) {
1844
- break;
1845
- }
1846
- }
1847
-
1848
- lm_ggml_backend_graph_copy_free(copy);
1849
-
1850
- return true;
1851
- }
1852
-
1853
- // CPU backend - buffer
1854
-
1855
- static void * lm_ggml_backend_cpu_buffer_get_base(lm_ggml_backend_buffer_t buffer) {
1856
- uintptr_t data = (uintptr_t)buffer->context;
1857
-
1858
- // align the buffer
1859
- if (data % TENSOR_ALIGNMENT != 0) {
1860
- data = LM_GGML_PAD(data, TENSOR_ALIGNMENT);
1861
- }
1862
-
1863
- return (void *)data;
1864
- }
1865
-
1866
- static void lm_ggml_backend_cpu_buffer_free_buffer(lm_ggml_backend_buffer_t buffer) {
1867
- lm_ggml_aligned_free(buffer->context, buffer->size);
1868
- }
1869
-
1870
- static void lm_ggml_backend_cpu_buffer_memset_tensor(lm_ggml_backend_buffer_t buffer, struct lm_ggml_tensor * tensor, uint8_t value, size_t offset, size_t size) {
1871
- memset((char *)tensor->data + offset, value, size);
1872
-
1873
- LM_GGML_UNUSED(buffer);
1874
- }
1875
-
1876
- static void lm_ggml_backend_cpu_buffer_set_tensor(lm_ggml_backend_buffer_t buffer, struct lm_ggml_tensor * tensor, const void * data, size_t offset, size_t size) {
1877
- memcpy((char *)tensor->data + offset, data, size);
1878
-
1879
- LM_GGML_UNUSED(buffer);
1880
- }
1881
-
1882
- static void lm_ggml_backend_cpu_buffer_get_tensor(lm_ggml_backend_buffer_t buffer, const struct lm_ggml_tensor * tensor, void * data, size_t offset, size_t size) {
1883
- memcpy(data, (const char *)tensor->data + offset, size);
1884
-
1885
- LM_GGML_UNUSED(buffer);
1886
- }
1887
-
1888
- static bool lm_ggml_backend_cpu_buffer_cpy_tensor(lm_ggml_backend_buffer_t buffer, const struct lm_ggml_tensor * src, struct lm_ggml_tensor * dst) {
1889
- if (lm_ggml_backend_buffer_is_host(src->buffer)) {
1890
- memcpy(dst->data, src->data, lm_ggml_nbytes(src));
1891
- return true;
1892
- }
1893
- return false;
1894
-
1895
- LM_GGML_UNUSED(buffer);
1896
- }
1897
-
1898
- static void lm_ggml_backend_cpu_buffer_clear(lm_ggml_backend_buffer_t buffer, uint8_t value) {
1899
- memset(buffer->context, value, buffer->size);
1900
- }
1901
-
1902
- static const struct lm_ggml_backend_buffer_i lm_ggml_backend_cpu_buffer_i = {
1903
- /* .free_buffer = */ lm_ggml_backend_cpu_buffer_free_buffer,
1904
- /* .get_base = */ lm_ggml_backend_cpu_buffer_get_base,
1905
- /* .init_tensor = */ NULL, // no initialization required
1906
- /* .memset_tensor = */ lm_ggml_backend_cpu_buffer_memset_tensor,
1907
- /* .set_tensor = */ lm_ggml_backend_cpu_buffer_set_tensor,
1908
- /* .get_tensor = */ lm_ggml_backend_cpu_buffer_get_tensor,
1909
- /* .cpy_tensor = */ lm_ggml_backend_cpu_buffer_cpy_tensor,
1910
- /* .clear = */ lm_ggml_backend_cpu_buffer_clear,
1911
- /* .reset = */ NULL,
1912
- };
1913
-
1914
- static const struct lm_ggml_backend_buffer_i lm_ggml_backend_cpu_buffer_from_ptr_i = {
1915
- /* .free_buffer = */ NULL, // ptr is not owned by the buffer, so it does not need to be freed
1916
- /* .get_base = */ lm_ggml_backend_cpu_buffer_get_base,
1917
- /* .init_tensor = */ NULL, // no initialization required
1918
- /* .memset_tensor = */ lm_ggml_backend_cpu_buffer_memset_tensor,
1919
- /* .set_tensor = */ lm_ggml_backend_cpu_buffer_set_tensor,
1920
- /* .get_tensor = */ lm_ggml_backend_cpu_buffer_get_tensor,
1921
- /* .cpy_tensor = */ lm_ggml_backend_cpu_buffer_cpy_tensor,
1922
- /* .clear = */ lm_ggml_backend_cpu_buffer_clear,
1923
- /* .reset = */ NULL,
1924
- };
1925
-
1926
- // CPU backend buffer type
1927
-
1928
- // this buffer type is defined here to make it available to all backends
1929
-
1930
- static const char * lm_ggml_backend_cpu_buffer_type_get_name(lm_ggml_backend_buffer_type_t buft) {
1931
- return "CPU";
1932
-
1933
- LM_GGML_UNUSED(buft);
1934
- }
1935
-
1936
- static lm_ggml_backend_buffer_t lm_ggml_backend_cpu_buffer_type_alloc_buffer(lm_ggml_backend_buffer_type_t buft, size_t size) {
1937
- void * data = lm_ggml_aligned_malloc(size);
1938
-
1939
- if (data == NULL) {
1940
- LM_GGML_LOG_ERROR("%s: failed to allocate buffer of size %zu\n", __func__, size);
1941
- return NULL;
1942
- }
1943
-
1944
- return lm_ggml_backend_buffer_init(buft, lm_ggml_backend_cpu_buffer_i, data, size);
1945
- }
1946
-
1947
- static size_t lm_ggml_backend_cpu_buffer_type_get_alignment(lm_ggml_backend_buffer_type_t buft) {
1948
- return TENSOR_ALIGNMENT;
1949
-
1950
- LM_GGML_UNUSED(buft);
1951
- }
1952
-
1953
- static bool lm_ggml_backend_cpu_buffer_type_is_host(lm_ggml_backend_buffer_type_t buft) {
1954
- return true;
1955
-
1956
- LM_GGML_UNUSED(buft);
1957
- }
1958
-
1959
- lm_ggml_backend_buffer_type_t lm_ggml_backend_cpu_buffer_type(void) {
1960
- static struct lm_ggml_backend_buffer_type lm_ggml_backend_cpu_buffer_type = {
1961
- /* .iface = */ {
1962
- /* .get_name = */ lm_ggml_backend_cpu_buffer_type_get_name,
1963
- /* .alloc_buffer = */ lm_ggml_backend_cpu_buffer_type_alloc_buffer,
1964
- /* .get_alignment = */ lm_ggml_backend_cpu_buffer_type_get_alignment,
1965
- /* .get_max_size = */ NULL, // defaults to SIZE_MAX
1966
- /* .get_alloc_size = */ NULL, // defaults to lm_ggml_nbytes
1967
- /* .is_host = */ lm_ggml_backend_cpu_buffer_type_is_host,
1968
- },
1969
- /* .device = */ NULL, // FIXME lm_ggml_backend_reg_dev_get(lm_ggml_backend_cpu_reg(), 0),
1970
- /* .context = */ NULL,
1971
- };
1972
-
1973
- return &lm_ggml_backend_cpu_buffer_type;
1974
- }
1975
-
1976
- static const char * lm_ggml_backend_cpu_buffer_from_ptr_type_get_name(lm_ggml_backend_buffer_type_t buft) {
1977
- return "CPU_Mapped";
1978
-
1979
- LM_GGML_UNUSED(buft);
1980
- }
1981
-
1982
- static lm_ggml_backend_buffer_type_t lm_ggml_backend_cpu_buffer_from_ptr_type(void) {
1983
- static struct lm_ggml_backend_buffer_type lm_ggml_backend_cpu_buffer_type = {
1984
- /* .iface = */ {
1985
- /* .get_name = */ lm_ggml_backend_cpu_buffer_from_ptr_type_get_name,
1986
- /* .alloc_buffer = */ lm_ggml_backend_cpu_buffer_type_alloc_buffer,
1987
- /* .get_alignment = */ lm_ggml_backend_cpu_buffer_type_get_alignment,
1988
- /* .get_max_size = */ NULL, // defaults to SIZE_MAX
1989
- /* .get_alloc_size = */ NULL, // defaults to lm_ggml_nbytes
1990
- /* .is_host = */ lm_ggml_backend_cpu_buffer_type_is_host,
1991
- },
1992
- /* .device = */ NULL, // FIXME lm_ggml_backend_reg_dev_get(lm_ggml_backend_cpu_reg(), 0),
1993
- /* .context = */ NULL,
1994
- };
1995
-
1996
- return &lm_ggml_backend_cpu_buffer_type;
1997
- }
1998
-
1999
- lm_ggml_backend_buffer_t lm_ggml_backend_cpu_buffer_from_ptr(void * ptr, size_t size) {
2000
- LM_GGML_ASSERT((uintptr_t)ptr % TENSOR_ALIGNMENT == 0 && "buffer pointer must be aligned");
2001
- return lm_ggml_backend_buffer_init(lm_ggml_backend_cpu_buffer_from_ptr_type(), lm_ggml_backend_cpu_buffer_from_ptr_i, ptr, size);
2002
- }
1
+ // Note: porting this file to C++ is a work in progress
2
+
3
+ #ifdef _WIN32
4
+ #define WIN32_LEAN_AND_MEAN
5
+ #ifndef NOMINMAX
6
+ # define NOMINMAX
7
+ #endif
8
+ #include <windows.h>
9
+ #endif
10
+
11
+ #include "ggml-backend.h"
12
+ #include "ggml-backend-impl.h"
13
+ #include "ggml-alloc.h"
14
+ #include "ggml-impl.h"
15
+
16
+ #include <assert.h>
17
+ #include <limits.h>
18
+ #include <stdarg.h>
19
+ #include <stdio.h>
20
+ #include <stdlib.h>
21
+ #include <string.h>
22
+ #include <string>
23
+ #include <vector>
24
+
25
+ #ifdef __APPLE__
26
+ #include <sys/types.h>
27
+ #include <sys/sysctl.h>
28
+ #endif
29
+
30
+
31
+ // backend buffer type
32
+
33
+ const char * lm_ggml_backend_buft_name(lm_ggml_backend_buffer_type_t buft) {
34
+ return buft->iface.get_name(buft);
35
+ }
36
+
37
+ lm_ggml_backend_buffer_t lm_ggml_backend_buft_alloc_buffer(lm_ggml_backend_buffer_type_t buft, size_t size) {
38
+ if (size == 0) {
39
+ // return a dummy buffer for zero-sized allocations
40
+ return lm_ggml_backend_buffer_init(buft, {}, NULL, 0);
41
+ }
42
+
43
+ return buft->iface.alloc_buffer(buft, size);
44
+ }
45
+
46
+ size_t lm_ggml_backend_buft_get_alignment(lm_ggml_backend_buffer_type_t buft) {
47
+ return buft->iface.get_alignment(buft);
48
+ }
49
+
50
+ size_t lm_ggml_backend_buft_get_max_size(lm_ggml_backend_buffer_type_t buft) {
51
+ // get_max_size is optional, defaults to SIZE_MAX
52
+ if (buft->iface.get_max_size) {
53
+ return buft->iface.get_max_size(buft);
54
+ }
55
+ return SIZE_MAX;
56
+ }
57
+
58
+ size_t lm_ggml_backend_buft_get_alloc_size(lm_ggml_backend_buffer_type_t buft, struct lm_ggml_tensor * tensor) {
59
+ // get_alloc_size is optional, defaults to lm_ggml_nbytes
60
+ if (buft->iface.get_alloc_size) {
61
+ size_t size = buft->iface.get_alloc_size(buft, tensor);
62
+ assert(size >= lm_ggml_nbytes(tensor));
63
+ return size;
64
+ }
65
+ return lm_ggml_nbytes(tensor);
66
+ }
67
+
68
+ bool lm_ggml_backend_buft_is_host(lm_ggml_backend_buffer_type_t buft) {
69
+ if (buft->iface.is_host) {
70
+ return buft->iface.is_host(buft);
71
+ }
72
+ return false;
73
+ }
74
+
75
+ lm_ggml_backend_dev_t lm_ggml_backend_buft_get_device(lm_ggml_backend_buffer_type_t buft) {
76
+ return buft->device;
77
+ }
78
+
79
+ // backend buffer
80
+
81
+ lm_ggml_backend_buffer_t lm_ggml_backend_buffer_init(
82
+ lm_ggml_backend_buffer_type_t buft,
83
+ struct lm_ggml_backend_buffer_i iface,
84
+ void * context,
85
+ size_t size) {
86
+ lm_ggml_backend_buffer_t buffer = new lm_ggml_backend_buffer {
87
+ /* .interface = */ iface,
88
+ /* .buft = */ buft,
89
+ /* .context = */ context,
90
+ /* .size = */ size,
91
+ /* .usage = */ LM_GGML_BACKEND_BUFFER_USAGE_ANY
92
+ };
93
+
94
+ return buffer;
95
+ }
96
+
97
+ const char * lm_ggml_backend_buffer_name(lm_ggml_backend_buffer_t buffer) {
98
+ return lm_ggml_backend_buft_name(lm_ggml_backend_buffer_get_type(buffer));
99
+ }
100
+
101
+ void lm_ggml_backend_buffer_free(lm_ggml_backend_buffer_t buffer) {
102
+ if (buffer == NULL) {
103
+ return;
104
+ }
105
+
106
+ if (buffer->iface.free_buffer != NULL) {
107
+ buffer->iface.free_buffer(buffer);
108
+ }
109
+ delete buffer;
110
+ }
111
+
112
+ size_t lm_ggml_backend_buffer_get_size(lm_ggml_backend_buffer_t buffer) {
113
+ return buffer->size;
114
+ }
115
+
116
+ void * lm_ggml_backend_buffer_get_base(lm_ggml_backend_buffer_t buffer) {
117
+ // get_base is optional if the buffer is zero-sized
118
+ if (buffer->size == 0) {
119
+ return NULL;
120
+ }
121
+
122
+ void * base = buffer->iface.get_base(buffer);
123
+
124
+ LM_GGML_ASSERT(base != NULL && "backend buffer base cannot be NULL");
125
+
126
+ return base;
127
+ }
128
+
129
+ void lm_ggml_backend_buffer_init_tensor(lm_ggml_backend_buffer_t buffer, struct lm_ggml_tensor * tensor) {
130
+ // init_tensor is optional
131
+ if (buffer->iface.init_tensor) {
132
+ buffer->iface.init_tensor(buffer, tensor);
133
+ }
134
+ }
135
+
136
+ void lm_ggml_backend_buffer_clear(lm_ggml_backend_buffer_t buffer, uint8_t value) {
137
+ // clear is optional if the buffer is zero-sized
138
+ if (buffer->size == 0) {
139
+ return;
140
+ }
141
+
142
+ buffer->iface.clear(buffer, value);
143
+ }
144
+
145
+ size_t lm_ggml_backend_buffer_get_alignment(lm_ggml_backend_buffer_t buffer) {
146
+ return lm_ggml_backend_buft_get_alignment(lm_ggml_backend_buffer_get_type(buffer));
147
+ }
148
+
149
+ size_t lm_ggml_backend_buffer_get_max_size(lm_ggml_backend_buffer_t buffer) {
150
+ return lm_ggml_backend_buft_get_max_size(lm_ggml_backend_buffer_get_type(buffer));
151
+ }
152
+
153
+ size_t lm_ggml_backend_buffer_get_alloc_size(lm_ggml_backend_buffer_t buffer, struct lm_ggml_tensor * tensor) {
154
+ return lm_ggml_backend_buft_get_alloc_size(lm_ggml_backend_buffer_get_type(buffer), tensor);
155
+ }
156
+
157
+ bool lm_ggml_backend_buffer_is_host(lm_ggml_backend_buffer_t buffer) {
158
+ return lm_ggml_backend_buft_is_host(lm_ggml_backend_buffer_get_type(buffer));
159
+ }
160
+
161
+ void lm_ggml_backend_buffer_set_usage(lm_ggml_backend_buffer_t buffer, enum lm_ggml_backend_buffer_usage usage) {
162
+ buffer->usage = usage;
163
+
164
+ // FIXME: add a generic callback to the buffer interface
165
+ if (lm_ggml_backend_buffer_is_multi_buffer(buffer)) {
166
+ lm_ggml_backend_multi_buffer_set_usage(buffer, usage);
167
+ }
168
+ }
169
+
170
+ enum lm_ggml_backend_buffer_usage lm_ggml_backend_buffer_get_usage(lm_ggml_backend_buffer_t buffer) {
171
+ return buffer->usage;
172
+ }
173
+
174
+ lm_ggml_backend_buffer_type_t lm_ggml_backend_buffer_get_type(lm_ggml_backend_buffer_t buffer) {
175
+ return buffer->buft;
176
+ }
177
+
178
+ void lm_ggml_backend_buffer_reset(lm_ggml_backend_buffer_t buffer) {
179
+ if (buffer->iface.reset) {
180
+ buffer->iface.reset(buffer);
181
+ }
182
+ }
183
+
184
+ bool lm_ggml_backend_buffer_copy_tensor(const struct lm_ggml_tensor * src, struct lm_ggml_tensor * dst) {
185
+ lm_ggml_backend_buffer_t dst_buf = dst->view_src ? dst->view_src->buffer : dst->buffer;
186
+ if (dst_buf->iface.cpy_tensor) {
187
+ return dst_buf->iface.cpy_tensor(dst_buf, src, dst);
188
+ }
189
+ return false;
190
+ }
191
+
192
+ // backend
193
+
194
+ lm_ggml_guid_t lm_ggml_backend_guid(lm_ggml_backend_t backend) {
195
+ if (backend == NULL) {
196
+ return NULL;
197
+ }
198
+ return backend->guid;
199
+ }
200
+
201
+ const char * lm_ggml_backend_name(lm_ggml_backend_t backend) {
202
+ if (backend == NULL) {
203
+ return "NULL";
204
+ }
205
+ return backend->iface.get_name(backend);
206
+ }
207
+
208
+ void lm_ggml_backend_free(lm_ggml_backend_t backend) {
209
+ if (backend == NULL) {
210
+ return;
211
+ }
212
+
213
+ backend->iface.free(backend);
214
+ }
215
+
216
+ lm_ggml_backend_buffer_type_t lm_ggml_backend_get_default_buffer_type(lm_ggml_backend_t backend) {
217
+ return lm_ggml_backend_dev_buffer_type(backend->device);
218
+ }
219
+
220
+ lm_ggml_backend_buffer_t lm_ggml_backend_alloc_buffer(lm_ggml_backend_t backend, size_t size) {
221
+ return lm_ggml_backend_buft_alloc_buffer(lm_ggml_backend_get_default_buffer_type(backend), size);
222
+ }
223
+
224
+ size_t lm_ggml_backend_get_alignment(lm_ggml_backend_t backend) {
225
+ return lm_ggml_backend_buft_get_alignment(lm_ggml_backend_get_default_buffer_type(backend));
226
+ }
227
+
228
+ size_t lm_ggml_backend_get_max_size(lm_ggml_backend_t backend) {
229
+ return lm_ggml_backend_buft_get_max_size(lm_ggml_backend_get_default_buffer_type(backend));
230
+ }
231
+
232
+ void lm_ggml_backend_tensor_set_async(lm_ggml_backend_t backend, struct lm_ggml_tensor * tensor, const void * data, size_t offset, size_t size) {
233
+ LM_GGML_ASSERT(tensor->data != NULL && "tensor not allocated");
234
+ LM_GGML_ASSERT(offset + size <= lm_ggml_nbytes(tensor) && "tensor write out of bounds");
235
+
236
+ if (backend->iface.set_tensor_async == NULL) {
237
+ lm_ggml_backend_tensor_set(tensor, data, offset, size);
238
+ } else {
239
+ backend->iface.set_tensor_async(backend, tensor, data, offset, size);
240
+ }
241
+ }
242
+
243
+ void lm_ggml_backend_tensor_get_async(lm_ggml_backend_t backend, const struct lm_ggml_tensor * tensor, void * data, size_t offset, size_t size) {
244
+ LM_GGML_ASSERT(tensor->data != NULL && "tensor not allocated");
245
+ LM_GGML_ASSERT(offset + size <= lm_ggml_nbytes(tensor) && "tensor read out of bounds");
246
+
247
+ if (backend->iface.get_tensor_async == NULL) {
248
+ lm_ggml_backend_tensor_get(tensor, data, offset, size);
249
+ } else {
250
+ backend->iface.get_tensor_async(backend, tensor, data, offset, size);
251
+ }
252
+ }
253
+
254
+ void lm_ggml_backend_tensor_set(struct lm_ggml_tensor * tensor, const void * data, size_t offset, size_t size) {
255
+ LM_GGML_ASSERT(tensor);
256
+ lm_ggml_backend_buffer_t buf = tensor->view_src ? tensor->view_src->buffer : tensor->buffer;
257
+
258
+ if (size == 0) {
259
+ return;
260
+ }
261
+
262
+ LM_GGML_ASSERT(buf != NULL && "tensor buffer not set");
263
+ LM_GGML_ASSERT(tensor->data != NULL && "tensor not allocated");
264
+ LM_GGML_ASSERT(offset + size <= lm_ggml_nbytes(tensor) && "tensor write out of bounds");
265
+
266
+ buf->iface.set_tensor(buf, tensor, data, offset, size);
267
+ }
268
+
269
+ void lm_ggml_backend_tensor_get(const struct lm_ggml_tensor * tensor, void * data, size_t offset, size_t size) {
270
+ LM_GGML_ASSERT(tensor);
271
+ lm_ggml_backend_buffer_t buf = tensor->view_src ? tensor->view_src->buffer : tensor->buffer;
272
+
273
+ if (size == 0) {
274
+ return;
275
+ }
276
+
277
+ LM_GGML_ASSERT(buf != NULL && "tensor buffer not set");
278
+ LM_GGML_ASSERT(tensor->data != NULL && "tensor not allocated");
279
+ LM_GGML_ASSERT(offset + size <= lm_ggml_nbytes(tensor) && "tensor read out of bounds");
280
+
281
+ buf->iface.get_tensor(buf, tensor, data, offset, size);
282
+ }
283
+
284
+ void lm_ggml_backend_tensor_memset(struct lm_ggml_tensor * tensor, uint8_t value, size_t offset, size_t size) {
285
+ lm_ggml_backend_buffer_t buf = tensor->view_src ? tensor->view_src->buffer : tensor->buffer;
286
+
287
+ if (size == 0) {
288
+ return;
289
+ }
290
+
291
+ LM_GGML_ASSERT(buf != NULL && "tensor buffer not set");
292
+ LM_GGML_ASSERT(tensor->data != NULL && "tensor not allocated");
293
+ LM_GGML_ASSERT(offset + size <= lm_ggml_nbytes(tensor) && "tensor write out of bounds");
294
+ LM_GGML_ASSERT(buf->iface.memset_tensor != NULL && "memset not implemented by backend buffer");
295
+
296
+ buf->iface.memset_tensor(buf, tensor, value, offset, size);
297
+ }
298
+
299
+ void lm_ggml_backend_synchronize(lm_ggml_backend_t backend) {
300
+ if (backend->iface.synchronize == NULL) {
301
+ return;
302
+ }
303
+
304
+ backend->iface.synchronize(backend);
305
+ }
306
+
307
+ lm_ggml_backend_graph_plan_t lm_ggml_backend_graph_plan_create(lm_ggml_backend_t backend, struct lm_ggml_cgraph * cgraph) {
308
+ LM_GGML_ASSERT(backend->iface.graph_plan_create != NULL);
309
+
310
+ return backend->iface.graph_plan_create(backend, cgraph);
311
+ }
312
+
313
+ void lm_ggml_backend_graph_plan_free(lm_ggml_backend_t backend, lm_ggml_backend_graph_plan_t plan) {
314
+ LM_GGML_ASSERT(backend->iface.graph_plan_free != NULL);
315
+
316
+ backend->iface.graph_plan_free(backend, plan);
317
+ }
318
+
319
+ enum lm_ggml_status lm_ggml_backend_graph_plan_compute(lm_ggml_backend_t backend, lm_ggml_backend_graph_plan_t plan) {
320
+ LM_GGML_ASSERT(backend->iface.graph_plan_compute != NULL);
321
+
322
+ return backend->iface.graph_plan_compute(backend, plan);
323
+ }
324
+
325
+ enum lm_ggml_status lm_ggml_backend_graph_compute(lm_ggml_backend_t backend, struct lm_ggml_cgraph * cgraph) {
326
+ enum lm_ggml_status err = lm_ggml_backend_graph_compute_async(backend, cgraph);
327
+ lm_ggml_backend_synchronize(backend);
328
+ return err;
329
+ }
330
+
331
+ enum lm_ggml_status lm_ggml_backend_graph_compute_async(lm_ggml_backend_t backend, struct lm_ggml_cgraph * cgraph) {
332
+ return backend->iface.graph_compute(backend, cgraph);
333
+ }
334
+
335
+ bool lm_ggml_backend_supports_op(lm_ggml_backend_t backend, const struct lm_ggml_tensor * op) {
336
+ return lm_ggml_backend_dev_supports_op(backend->device, op);
337
+ }
338
+
339
+ bool lm_ggml_backend_supports_buft(lm_ggml_backend_t backend, lm_ggml_backend_buffer_type_t buft) {
340
+ return lm_ggml_backend_dev_supports_buft(backend->device, buft);
341
+ }
342
+
343
+ bool lm_ggml_backend_offload_op(lm_ggml_backend_t backend, const struct lm_ggml_tensor * op) {
344
+ return lm_ggml_backend_dev_offload_op(backend->device, op);
345
+ }
346
+
347
+ lm_ggml_backend_dev_t lm_ggml_backend_get_device(lm_ggml_backend_t backend) {
348
+ return backend->device;
349
+ }
350
+
351
+ // backend copy
352
+
353
+ static bool lm_ggml_are_same_layout(const struct lm_ggml_tensor * a, const struct lm_ggml_tensor * b) {
354
+ if (a->type != b->type) {
355
+ return false;
356
+ }
357
+ for (int i = 0; i < LM_GGML_MAX_DIMS; i++) {
358
+ if (a->ne[i] != b->ne[i]) {
359
+ return false;
360
+ }
361
+ if (a->nb[i] != b->nb[i]) {
362
+ return false;
363
+ }
364
+ }
365
+ return true;
366
+ }
367
+
368
+ void lm_ggml_backend_tensor_copy(struct lm_ggml_tensor * src, struct lm_ggml_tensor * dst) {
369
+ LM_GGML_ASSERT(lm_ggml_are_same_layout(src, dst) && "cannot copy tensors with different layouts");
370
+
371
+ if (src == dst) {
372
+ return;
373
+ }
374
+
375
+ if (lm_ggml_backend_buffer_is_host(src->buffer)) {
376
+ lm_ggml_backend_tensor_set(dst, src->data, 0, lm_ggml_nbytes(src));
377
+ } else if (lm_ggml_backend_buffer_is_host(dst->buffer)) {
378
+ lm_ggml_backend_tensor_get(src, dst->data, 0, lm_ggml_nbytes(src));
379
+ } else if (!lm_ggml_backend_buffer_copy_tensor(src, dst)) {
380
+ #ifndef NDEBUG
381
+ LM_GGML_LOG_DEBUG("%s: warning: slow copy from %s to %s\n", __func__, lm_ggml_backend_buffer_name(src->buffer), lm_ggml_backend_buffer_name(dst->buffer));
382
+ #endif
383
+ size_t nbytes = lm_ggml_nbytes(src);
384
+ void * data = malloc(nbytes);
385
+ lm_ggml_backend_tensor_get(src, data, 0, nbytes);
386
+ lm_ggml_backend_tensor_set(dst, data, 0, nbytes);
387
+ free(data);
388
+ }
389
+ }
390
+
391
+ void lm_ggml_backend_tensor_copy_async(lm_ggml_backend_t backend_src, lm_ggml_backend_t backend_dst, struct lm_ggml_tensor * src, struct lm_ggml_tensor * dst) {
392
+ LM_GGML_ASSERT(lm_ggml_are_same_layout(src, dst) && "cannot copy tensors with different layouts");
393
+
394
+ if (src == dst) {
395
+ return;
396
+ }
397
+
398
+ if (backend_dst->iface.cpy_tensor_async != NULL) {
399
+ if (backend_dst->iface.cpy_tensor_async(backend_src, backend_dst, src, dst)) {
400
+ return;
401
+ }
402
+ }
403
+
404
+ // an async copy would normally happen after all the queued operations on both backends are completed
405
+ // to simulate the same behavior, we need to synchronize both backends first, and do a blocking copy
406
+ lm_ggml_backend_synchronize(backend_src);
407
+ lm_ggml_backend_synchronize(backend_dst);
408
+ lm_ggml_backend_tensor_copy(src, dst);
409
+ }
410
+
411
+ // events
412
+
413
+ lm_ggml_backend_event_t lm_ggml_backend_event_new(lm_ggml_backend_dev_t device) {
414
+ // null device is allowed for the transition period to the device interface
415
+ if (device == NULL || device->iface.event_new == NULL) {
416
+ return NULL;
417
+ }
418
+ return device->iface.event_new(device);
419
+ }
420
+
421
+ void lm_ggml_backend_event_free(lm_ggml_backend_event_t event) {
422
+ if (event == NULL) {
423
+ return;
424
+ }
425
+ event->device->iface.event_free(event->device, event);
426
+ }
427
+
428
+ void lm_ggml_backend_event_record(lm_ggml_backend_event_t event, lm_ggml_backend_t backend) {
429
+ LM_GGML_ASSERT(backend->iface.event_record != NULL);
430
+
431
+ backend->iface.event_record(backend, event);
432
+ }
433
+
434
+ void lm_ggml_backend_event_synchronize(lm_ggml_backend_event_t event) {
435
+ LM_GGML_ASSERT(event->device->iface.event_synchronize);
436
+
437
+ event->device->iface.event_synchronize(event->device, event);
438
+ }
439
+
440
+ void lm_ggml_backend_event_wait(lm_ggml_backend_t backend, lm_ggml_backend_event_t event) {
441
+ LM_GGML_ASSERT(backend->iface.event_wait != NULL);
442
+
443
+ backend->iface.event_wait(backend, event);
444
+ }
445
+
446
+ // Backend device
447
+
448
+ const char * lm_ggml_backend_dev_name(lm_ggml_backend_dev_t device) {
449
+ return device->iface.get_name(device);
450
+ }
451
+
452
+ const char * lm_ggml_backend_dev_description(lm_ggml_backend_dev_t device) {
453
+ return device->iface.get_description(device);
454
+ }
455
+
456
+ void lm_ggml_backend_dev_memory(lm_ggml_backend_dev_t device, size_t * free, size_t * total) {
457
+ device->iface.get_memory(device, free, total);
458
+ }
459
+
460
+ enum lm_ggml_backend_dev_type lm_ggml_backend_dev_type(lm_ggml_backend_dev_t device) {
461
+ return device->iface.get_type(device);
462
+ }
463
+
464
+ void lm_ggml_backend_dev_get_props(lm_ggml_backend_dev_t device, struct lm_ggml_backend_dev_props * props) {
465
+ memset(props, 0, sizeof(*props));
466
+ device->iface.get_props(device, props);
467
+ }
468
+
469
+ lm_ggml_backend_reg_t lm_ggml_backend_dev_backend_reg(lm_ggml_backend_dev_t device) {
470
+ return device->reg;
471
+ }
472
+
473
+ lm_ggml_backend_t lm_ggml_backend_dev_init(lm_ggml_backend_dev_t device, const char * params) {
474
+ return device->iface.init_backend(device, params);
475
+ }
476
+
477
+ lm_ggml_backend_buffer_type_t lm_ggml_backend_dev_buffer_type(lm_ggml_backend_dev_t device) {
478
+ return device->iface.get_buffer_type(device);
479
+ }
480
+
481
+ lm_ggml_backend_buffer_type_t lm_ggml_backend_dev_host_buffer_type(lm_ggml_backend_dev_t device) {
482
+ if (device->iface.get_host_buffer_type == NULL) {
483
+ return NULL;
484
+ }
485
+
486
+ return device->iface.get_host_buffer_type(device);
487
+ }
488
+
489
+ lm_ggml_backend_buffer_t lm_ggml_backend_dev_buffer_from_host_ptr(lm_ggml_backend_dev_t device, void * ptr, size_t size, size_t max_tensor_size) {
490
+ return device->iface.buffer_from_host_ptr(device, ptr, size, max_tensor_size);
491
+ }
492
+
493
+ bool lm_ggml_backend_dev_supports_op(lm_ggml_backend_dev_t device, const struct lm_ggml_tensor * op) {
494
+ return device->iface.supports_op(device, op);
495
+ }
496
+
497
+ bool lm_ggml_backend_dev_supports_buft(lm_ggml_backend_dev_t device, lm_ggml_backend_buffer_type_t buft) {
498
+ return device->iface.supports_buft(device, buft);
499
+ }
500
+
501
+ bool lm_ggml_backend_dev_offload_op(lm_ggml_backend_dev_t device, const struct lm_ggml_tensor * op) {
502
+ if (device->iface.offload_op != NULL) {
503
+ return device->iface.offload_op(device, op);
504
+ }
505
+
506
+ return false;
507
+ }
508
+
509
+ // Backend (reg)
510
+
511
+ const char * lm_ggml_backend_reg_name(lm_ggml_backend_reg_t reg) {
512
+ return reg->iface.get_name(reg);
513
+ }
514
+
515
+ size_t lm_ggml_backend_reg_dev_count(lm_ggml_backend_reg_t reg) {
516
+ return reg->iface.get_device_count(reg);
517
+ }
518
+
519
+ lm_ggml_backend_dev_t lm_ggml_backend_reg_dev_get(lm_ggml_backend_reg_t reg, size_t index) {
520
+ return reg->iface.get_device(reg, index);
521
+ }
522
+
523
+ void * lm_ggml_backend_reg_get_proc_address(lm_ggml_backend_reg_t reg, const char * name) {
524
+ if (!reg->iface.get_proc_address) {
525
+ return NULL;
526
+ }
527
+ return reg->iface.get_proc_address(reg, name);
528
+ }
529
+
530
+ // multi-buffer buffer
531
+
532
+ struct lm_ggml_backend_multi_buffer_context {
533
+ lm_ggml_backend_buffer_t * buffers;
534
+ size_t n_buffers;
535
+ };
536
+
537
+ static void lm_ggml_backend_multi_buffer_free_buffer(lm_ggml_backend_buffer_t buffer) {
538
+ lm_ggml_backend_multi_buffer_context * ctx = (lm_ggml_backend_multi_buffer_context *) buffer->context;
539
+ for (size_t i = 0; i < ctx->n_buffers; i++) {
540
+ lm_ggml_backend_buffer_free(ctx->buffers[i]);
541
+ }
542
+
543
+ free(ctx->buffers);
544
+ free(ctx);
545
+ }
546
+
547
+ static void lm_ggml_backend_multi_buffer_clear(lm_ggml_backend_buffer_t buffer, uint8_t value) {
548
+ lm_ggml_backend_multi_buffer_context * ctx = (lm_ggml_backend_multi_buffer_context *) buffer->context;
549
+ for (size_t i = 0; i < ctx->n_buffers; i++) {
550
+ lm_ggml_backend_buffer_clear(ctx->buffers[i], value);
551
+ }
552
+ }
553
+
554
+ static const struct lm_ggml_backend_buffer_i lm_ggml_backend_multi_buffer_i = {
555
+ /* .free_buffer = */ lm_ggml_backend_multi_buffer_free_buffer,
556
+ /* .get_base = */ NULL,
557
+ /* .init_tensor = */ NULL,
558
+ /* .memset_tensor = */ NULL,
559
+ /* .set_tensor = */ NULL,
560
+ /* .get_tensor = */ NULL,
561
+ /* .cpy_tensor = */ NULL,
562
+ /* .clear = */ lm_ggml_backend_multi_buffer_clear,
563
+ /* .reset = */ NULL,
564
+ };
565
+
566
+ lm_ggml_backend_buffer_t lm_ggml_backend_multi_buffer_alloc_buffer(lm_ggml_backend_buffer_t * buffers, size_t n_buffers) {
567
+ lm_ggml_backend_multi_buffer_context * ctx = (lm_ggml_backend_multi_buffer_context *) malloc(sizeof(struct lm_ggml_backend_multi_buffer_context));
568
+ ctx->n_buffers = n_buffers;
569
+ ctx->buffers = (lm_ggml_backend_buffer_t *) malloc(n_buffers * sizeof(lm_ggml_backend_buffer_t));
570
+
571
+ LM_GGML_ASSERT(ctx->buffers != NULL);
572
+
573
+ size_t total_size = 0;
574
+ for (size_t i = 0; i < n_buffers; i++) {
575
+ ctx->buffers[i] = buffers[i];
576
+ total_size += lm_ggml_backend_buffer_get_size(buffers[i]);
577
+ }
578
+
579
+ return lm_ggml_backend_buffer_init(buffers[0]->buft, lm_ggml_backend_multi_buffer_i, ctx, total_size);
580
+ }
581
+
582
+ bool lm_ggml_backend_buffer_is_multi_buffer(lm_ggml_backend_buffer_t buffer) {
583
+ return buffer->iface.free_buffer == lm_ggml_backend_multi_buffer_free_buffer;
584
+ }
585
+
586
+ void lm_ggml_backend_multi_buffer_set_usage(lm_ggml_backend_buffer_t buffer, enum lm_ggml_backend_buffer_usage usage) {
587
+ LM_GGML_ASSERT(lm_ggml_backend_buffer_is_multi_buffer(buffer));
588
+ lm_ggml_backend_multi_buffer_context * ctx = (lm_ggml_backend_multi_buffer_context *) buffer->context;
589
+ for (size_t i = 0; i < ctx->n_buffers; i++) {
590
+ lm_ggml_backend_buffer_set_usage(ctx->buffers[i], usage);
591
+ }
592
+ }
593
+
594
+ // creates a copy of the tensor with the same memory layout
595
+ static struct lm_ggml_tensor * lm_ggml_dup_tensor_layout(struct lm_ggml_context * ctx, const struct lm_ggml_tensor * tensor) {
596
+ struct lm_ggml_tensor * dup = lm_ggml_dup_tensor(ctx, tensor);
597
+ for (int i = 0; i < LM_GGML_MAX_DIMS; i++) {
598
+ dup->nb[i] = tensor->nb[i];
599
+ }
600
+ return dup;
601
+ }
602
+
603
+ static bool lm_ggml_is_view_op(enum lm_ggml_op op) {
604
+ return op == LM_GGML_OP_VIEW || op == LM_GGML_OP_RESHAPE || op == LM_GGML_OP_PERMUTE || op == LM_GGML_OP_TRANSPOSE;
605
+ }
606
+
607
+ // scheduler
608
+
609
+ #ifndef LM_GGML_SCHED_MAX_BACKENDS
610
+ #define LM_GGML_SCHED_MAX_BACKENDS 16
611
+ #endif
612
+
613
+ #ifndef LM_GGML_SCHED_MAX_SPLIT_INPUTS
614
+ #define LM_GGML_SCHED_MAX_SPLIT_INPUTS LM_GGML_MAX_SRC
615
+ #endif
616
+
617
+ #ifndef LM_GGML_SCHED_MAX_COPIES
618
+ #define LM_GGML_SCHED_MAX_COPIES 4
619
+ #endif
620
+
621
+ struct lm_ggml_backend_sched_split {
622
+ int backend_id;
623
+ int i_start;
624
+ int i_end;
625
+ struct lm_ggml_tensor * inputs[LM_GGML_SCHED_MAX_SPLIT_INPUTS];
626
+ int n_inputs;
627
+ // graph view of this split
628
+ struct lm_ggml_cgraph graph;
629
+ };
630
+
631
+ struct lm_ggml_backend_sched {
632
+ bool is_reset; // true if the scheduler has been reset since the last graph split
633
+ bool is_alloc;
634
+
635
+ int n_backends;
636
+
637
+ lm_ggml_backend_t backends[LM_GGML_SCHED_MAX_BACKENDS];
638
+ lm_ggml_backend_buffer_type_t bufts[LM_GGML_SCHED_MAX_BACKENDS];
639
+ lm_ggml_gallocr_t galloc;
640
+
641
+ // hash map of the nodes in the graph
642
+ struct lm_ggml_hash_set hash_set;
643
+ int * hv_tensor_backend_ids; // [hash_set.size]
644
+ struct lm_ggml_tensor ** hv_tensor_copies; // [hash_set.size][n_backends][n_copies]
645
+
646
+ int * node_backend_ids; // [graph_size]
647
+ int * leaf_backend_ids; // [graph_size]
648
+
649
+ int * prev_node_backend_ids; // [graph_size]
650
+ int * prev_leaf_backend_ids; // [graph_size]
651
+
652
+ // copy of the graph with modified inputs
653
+ struct lm_ggml_cgraph graph;
654
+
655
+ // graph splits
656
+ struct lm_ggml_backend_sched_split * splits;
657
+ int n_splits;
658
+ int splits_capacity;
659
+
660
+ // pipeline parallelism support
661
+ int n_copies;
662
+ int cur_copy;
663
+ lm_ggml_backend_event_t events[LM_GGML_SCHED_MAX_BACKENDS][LM_GGML_SCHED_MAX_COPIES];
664
+ struct lm_ggml_tensor * graph_inputs[LM_GGML_SCHED_MAX_SPLIT_INPUTS];
665
+ int n_graph_inputs;
666
+
667
+ struct lm_ggml_context * ctx;
668
+
669
+ lm_ggml_backend_sched_eval_callback callback_eval;
670
+ void * callback_eval_user_data;
671
+
672
+ char * context_buffer;
673
+ size_t context_buffer_size;
674
+
675
+ int debug;
676
+ };
677
+
678
+ #define hash_id(tensor) lm_ggml_hash_find_or_insert(&sched->hash_set, tensor)
679
+ #define tensor_backend_id(tensor) sched->hv_tensor_backend_ids[hash_id(tensor)]
680
+ #define tensor_id_copy(id, backend_id, copy_id) sched->hv_tensor_copies[(id) * sched->n_backends * sched->n_copies + (backend_id) * sched->n_copies + (copy_id)]
681
+ #define tensor_copy(tensor, backend_id, copy_id) tensor_id_copy(hash_id(tensor), backend_id, copy_id)
682
+
683
+ // returns the priority of the backend, lower id is higher priority
684
+ static int lm_ggml_backend_sched_backend_id(lm_ggml_backend_sched_t sched, lm_ggml_backend_t backend) {
685
+ for (int i = 0; i < sched->n_backends; i++) {
686
+ if (sched->backends[i] == backend) {
687
+ return i;
688
+ }
689
+ }
690
+ return -1;
691
+ }
692
+
693
+ static int lm_ggml_backend_sched_backend_from_buffer(lm_ggml_backend_sched_t sched, const struct lm_ggml_tensor * tensor, const struct lm_ggml_tensor * op) {
694
+ lm_ggml_backend_buffer_t buffer = tensor->view_src ? tensor->view_src->buffer : tensor->buffer;
695
+ if (buffer == NULL) {
696
+ return -1;
697
+ }
698
+
699
+ // find highest prio backend that supports the buffer type and the op
700
+ for (int i = 0; i < sched->n_backends; i++) {
701
+ if (lm_ggml_backend_supports_buft(sched->backends[i], buffer->buft) &&
702
+ lm_ggml_backend_supports_op(sched->backends[i], op)) {
703
+ return i;
704
+ }
705
+ }
706
+
707
+ #ifndef NDEBUG
708
+ LM_GGML_LOG_DEBUG("%s: warning: no backend supports op %s with a weight with buffer type %s used in tensor %s, the weight will need to be copied\n",
709
+ __func__, lm_ggml_op_desc(tensor), lm_ggml_backend_buffer_name(buffer), tensor->name);
710
+ #endif
711
+
712
+ return -1;
713
+ }
714
+
715
+ #if 0
716
+ #define LM_GGML_SCHED_MAX_SPLITS_DEBUG 4096
717
+ static char causes[LM_GGML_DEFAULT_GRAPH_SIZE*16 + LM_GGML_SCHED_MAX_SPLITS_DEBUG*LM_GGML_SCHED_MAX_SPLIT_INPUTS][128]; // debug only
718
+ #define SET_CAUSE(node, ...) sprintf(causes[hash_id(node)], __VA_ARGS__)
719
+ #define GET_CAUSE(node) causes[hash_id(node)]
720
+ #else
721
+ #define SET_CAUSE(node, ...)
722
+ #define GET_CAUSE(node) ""
723
+ #endif
724
+
725
+ // returns the backend that should be used for the node based on the current locations
726
+ static int lm_ggml_backend_sched_backend_id_from_cur(lm_ggml_backend_sched_t sched, struct lm_ggml_tensor * tensor) {
727
+ // assign pre-allocated nodes to their backend
728
+ int cur_backend_id = lm_ggml_backend_sched_backend_from_buffer(sched, tensor, tensor);
729
+ if (cur_backend_id != -1) {
730
+ SET_CAUSE(tensor, "1.dst");
731
+ return cur_backend_id;
732
+ }
733
+
734
+ // view_src
735
+ if (tensor->view_src != NULL) {
736
+ cur_backend_id = lm_ggml_backend_sched_backend_from_buffer(sched, tensor->view_src, tensor);
737
+ if (cur_backend_id != -1) {
738
+ SET_CAUSE(tensor, "1.vsrc");
739
+ return cur_backend_id;
740
+ }
741
+ }
742
+
743
+ if (tensor->buffer || (tensor->view_src && tensor->view_src->buffer)) {
744
+ // since the tensor is pre-allocated, it cannot be moved to another backend
745
+ lm_ggml_backend_buffer_t buffer = tensor->view_src ? tensor->view_src->buffer : tensor->buffer;
746
+ LM_GGML_ABORT("pre-allocated tensor (%s) in a buffer (%s) that cannot run the operation (%s)", tensor->name, lm_ggml_backend_buffer_name(buffer), lm_ggml_op_name(tensor->op));
747
+ }
748
+
749
+ // graph input
750
+ if (tensor->flags & LM_GGML_TENSOR_FLAG_INPUT) {
751
+ cur_backend_id = sched->n_backends - 1; // last backend (assumed CPU)
752
+ SET_CAUSE(tensor, "1.inp");
753
+ return cur_backend_id;
754
+ }
755
+
756
+ // operations with weights are preferably run on the same backend as the weights
757
+ for (int i = 0; i < LM_GGML_MAX_SRC; i++) {
758
+ const struct lm_ggml_tensor * src = tensor->src[i];
759
+ if (src == NULL) {
760
+ continue;
761
+ }
762
+ // skip ROPE since the rope freqs tensor is too small to choose a backend based on it
763
+ // not an ideal solution
764
+ if (tensor->op != LM_GGML_OP_ROPE && src->buffer != NULL && src->buffer->usage == LM_GGML_BACKEND_BUFFER_USAGE_WEIGHTS) {
765
+ int src_backend_id = lm_ggml_backend_sched_backend_from_buffer(sched, src, tensor);
766
+ // check if a backend with higher prio wants to offload the op
767
+ if (src_backend_id == sched->n_backends - 1 && lm_ggml_backend_buffer_is_host(src->buffer)) {
768
+ for (int b = 0; b < src_backend_id; b++) {
769
+ if (lm_ggml_backend_supports_op(sched->backends[b], tensor) && lm_ggml_backend_offload_op(sched->backends[b], tensor)) {
770
+ SET_CAUSE(tensor, "1.off");
771
+ return b;
772
+ }
773
+ }
774
+ }
775
+ SET_CAUSE(tensor, "1.wgt%d", i);
776
+ return src_backend_id;
777
+ }
778
+ }
779
+
780
+ return -1;
781
+ }
782
+
783
+ static char * fmt_size(size_t size) {
784
+ static char buffer[128];
785
+ if (size >= 1024*1024) {
786
+ snprintf(buffer, sizeof(buffer), "%zuM", size/1024/1024);
787
+ } else {
788
+ snprintf(buffer, sizeof(buffer), "%zuK", size/1024);
789
+ }
790
+ return buffer;
791
+ }
792
+
793
+ static void lm_ggml_backend_sched_print_assignments(lm_ggml_backend_sched_t sched, struct lm_ggml_cgraph * graph) {
794
+ int cur_split = 0;
795
+ for (int i = 0; i < graph->n_nodes; i++) {
796
+ if (cur_split < sched->n_splits && i == sched->splits[cur_split].i_start) {
797
+ lm_ggml_backend_t split_backend = sched->backends[sched->splits[cur_split].backend_id];
798
+ LM_GGML_LOG_DEBUG("\n## SPLIT #%d: %s # %d inputs", cur_split, lm_ggml_backend_name(split_backend),
799
+ sched->splits[cur_split].n_inputs);
800
+ for (int j = 0; j < sched->splits[cur_split].n_inputs; j++) {
801
+ if (j == 0) {
802
+ LM_GGML_LOG_DEBUG(": ");
803
+ }
804
+ LM_GGML_LOG_DEBUG("[%s (%5.5s)] ", sched->splits[cur_split].inputs[j]->name,
805
+ fmt_size(lm_ggml_nbytes(sched->splits[cur_split].inputs[j])));
806
+ }
807
+ LM_GGML_LOG_DEBUG("\n");
808
+ cur_split++;
809
+ }
810
+ struct lm_ggml_tensor * node = graph->nodes[i];
811
+ if (lm_ggml_is_view_op(node->op)) {
812
+ continue;
813
+ }
814
+ if (sched->debug > 1) {
815
+ lm_ggml_backend_t tensor_backend = lm_ggml_backend_sched_get_tensor_backend(sched, node);
816
+ LM_GGML_LOG_DEBUG("node #%3d (%10.10s): %20.20s (%5.5s) [%5.5s %8.8s]:", i, lm_ggml_op_name(node->op), node->name,
817
+ fmt_size(lm_ggml_nbytes(node)), tensor_backend ? lm_ggml_backend_name(tensor_backend) : "NULL", GET_CAUSE(node));
818
+ for (int j = 0; j < LM_GGML_MAX_SRC; j++) {
819
+ struct lm_ggml_tensor * src = node->src[j];
820
+ if (src == NULL) {
821
+ continue;
822
+ }
823
+ lm_ggml_backend_t src_backend = lm_ggml_backend_sched_get_tensor_backend(sched, src);
824
+ LM_GGML_LOG_DEBUG(" %20.20s (%5.5s) [%5.5s %8.8s]", src->name,
825
+ fmt_size(lm_ggml_nbytes(src)), src_backend ? lm_ggml_backend_name(src_backend) : "NULL", GET_CAUSE(src));
826
+ }
827
+ LM_GGML_LOG_DEBUG("\n");
828
+ }
829
+ }
830
+ }
831
+
832
+ static bool lm_ggml_backend_sched_buffer_supported(lm_ggml_backend_sched_t sched, struct lm_ggml_tensor * t, int backend_id) {
833
+ lm_ggml_backend_buffer_t buf = t->view_src ? t->view_src->buffer : t->buffer;
834
+ lm_ggml_backend_buffer_type_t buft = NULL;
835
+
836
+ if (buf) {
837
+ // the tensor is already allocated
838
+ buft = buf->buft;
839
+ } else {
840
+ // see if the tensor already has a backend assigned, and use the buffer type of that backend
841
+ int tensor_backend_id = tensor_backend_id(t);
842
+ if (tensor_backend_id == -1 && t->view_src) {
843
+ tensor_backend_id = tensor_backend_id(t->view_src);
844
+ }
845
+ if (tensor_backend_id != -1) {
846
+ buft = sched->bufts[tensor_backend_id];
847
+ }
848
+ }
849
+
850
+ return buft != NULL && lm_ggml_backend_supports_buft(sched->backends[backend_id], buft);
851
+ }
852
+
853
+ static void lm_ggml_backend_sched_set_if_supported(lm_ggml_backend_sched_t sched, struct lm_ggml_tensor * node, int cur_backend_id, int * node_backend_id) {
854
+ if (lm_ggml_backend_supports_op(sched->backends[cur_backend_id], node)) {
855
+ *node_backend_id = cur_backend_id;
856
+ SET_CAUSE(node, "2.sup");
857
+ }
858
+ }
859
+
860
+ // assigns backends to ops and splits the graph into subgraphs that can be computed on the same backend
861
+ static void lm_ggml_backend_sched_split_graph(lm_ggml_backend_sched_t sched, struct lm_ggml_cgraph * graph) {
862
+ // reset splits
863
+ sched->n_splits = 0;
864
+ sched->n_graph_inputs = 0;
865
+ sched->is_reset = false;
866
+
867
+ struct lm_ggml_init_params params = {
868
+ /* .mem_size = */ sched->context_buffer_size,
869
+ /* .mem_buffer = */ sched->context_buffer,
870
+ /* .no_alloc = */ true
871
+ };
872
+
873
+ lm_ggml_free(sched->ctx);
874
+
875
+ sched->ctx = lm_ggml_init(params);
876
+ if (sched->ctx == NULL) {
877
+ LM_GGML_ABORT("%s: failed to initialize context\n", __func__);
878
+ }
879
+
880
+ // pass 1: assign backends to ops with pre-allocated inputs
881
+ for (int i = 0; i < graph->n_leafs; i++) {
882
+ struct lm_ggml_tensor * leaf = graph->leafs[i];
883
+ int * leaf_backend_id = &tensor_backend_id(leaf);
884
+ // do not overwrite user assignments
885
+ if (*leaf_backend_id == -1) {
886
+ *leaf_backend_id = lm_ggml_backend_sched_backend_id_from_cur(sched, leaf);
887
+ }
888
+ }
889
+
890
+ for (int i = 0; i < graph->n_nodes; i++) {
891
+ struct lm_ggml_tensor * node = graph->nodes[i];
892
+ int * node_backend_id = &tensor_backend_id(node);
893
+ // do not overwrite user assignments
894
+ if (*node_backend_id == -1) {
895
+ *node_backend_id = lm_ggml_backend_sched_backend_id_from_cur(sched, node);
896
+
897
+ #if 0
898
+ // src
899
+ if (node->op == LM_GGML_OP_NONE) {
900
+ continue;
901
+ }
902
+
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
+ int * src_backend_id = &tensor_backend_id(src);
909
+ if (*src_backend_id == -1) {
910
+ *src_backend_id = lm_ggml_backend_sched_backend_id_from_cur(sched, src);
911
+ }
912
+ }
913
+ #endif
914
+ }
915
+ }
916
+
917
+ // pass 2: expand current backend assignments
918
+ // assign the same backend to adjacent nodes
919
+ // expand gpu backends (i.e. non last prio) up and down, ignoring cpu (the lowest priority backend)
920
+ // thus, cpu will never be used unless weights are on cpu, or there are no gpu ops between cpu ops
921
+ // ops unsupported by the backend being expanded will be left unassigned so that they can be assigned later when the locations of its inputs are known
922
+ // expand gpu down
923
+ {
924
+ int cur_backend_id = -1;
925
+ for (int i = 0; i < graph->n_nodes; i++) {
926
+ struct lm_ggml_tensor * node = graph->nodes[i];
927
+ if (lm_ggml_is_view_op(node->op)) {
928
+ continue;
929
+ }
930
+ int * node_backend_id = &tensor_backend_id(node);
931
+ if (*node_backend_id != -1) {
932
+ if (*node_backend_id == sched->n_backends - 1) {
933
+ // skip cpu (lowest prio backend)
934
+ cur_backend_id = -1;
935
+ } else {
936
+ cur_backend_id = *node_backend_id;
937
+ }
938
+ } else if (cur_backend_id != -1) {
939
+ lm_ggml_backend_sched_set_if_supported(sched, node, cur_backend_id, node_backend_id);
940
+ }
941
+ }
942
+ }
943
+ // expand gpu up
944
+ {
945
+ int cur_backend_id = -1;
946
+ for (int i = graph->n_nodes - 1; i >= 0; i--) {
947
+ struct lm_ggml_tensor * node = graph->nodes[i];
948
+ if (lm_ggml_is_view_op(node->op)) {
949
+ continue;
950
+ }
951
+ int * node_backend_id = &tensor_backend_id(node);
952
+ if (*node_backend_id != -1) {
953
+ if (*node_backend_id == sched->n_backends - 1) {
954
+ // skip cpu (lowest prio backend)
955
+ cur_backend_id = -1;
956
+ } else {
957
+ cur_backend_id = *node_backend_id;
958
+ }
959
+ } else if (cur_backend_id != -1) {
960
+ lm_ggml_backend_sched_set_if_supported(sched, node, cur_backend_id, node_backend_id);
961
+ }
962
+ }
963
+ }
964
+ // expand rest down
965
+ {
966
+ int cur_backend_id = -1;
967
+ for (int i = 0; i < graph->n_nodes; i++) {
968
+ struct lm_ggml_tensor * node = graph->nodes[i];
969
+ if (lm_ggml_is_view_op(node->op)) {
970
+ continue;
971
+ }
972
+ int * node_backend_id = &tensor_backend_id(node);
973
+ if (*node_backend_id != -1) {
974
+ cur_backend_id = *node_backend_id;
975
+ } else if (cur_backend_id != -1) {
976
+ lm_ggml_backend_sched_set_if_supported(sched, node, cur_backend_id, node_backend_id);
977
+ }
978
+ }
979
+ }
980
+ // expand rest up
981
+ {
982
+ int cur_backend_id = -1;
983
+ for (int i = graph->n_nodes - 1; i >= 0; i--) {
984
+ struct lm_ggml_tensor * node = graph->nodes[i];
985
+ if (lm_ggml_is_view_op(node->op)) {
986
+ continue;
987
+ }
988
+ int * node_backend_id = &tensor_backend_id(node);
989
+ if (*node_backend_id != -1) {
990
+ cur_backend_id = *node_backend_id;
991
+ } else if (cur_backend_id != -1) {
992
+ lm_ggml_backend_sched_set_if_supported(sched, node, cur_backend_id, node_backend_id);
993
+ }
994
+ }
995
+ }
996
+
997
+ // pass 3: upgrade nodes to higher prio backends with compatible buffer types
998
+ // if the tensor is already in the same buffer type (*) as another higher priority backend, we should move it there
999
+ // however, we also need to verify that the sources are in compatible buffer types
1000
+ // (*) the actual requirement is more relaxed, the buffer type of the backend should be supported by all the users of this tensor further down the graph
1001
+ // however, this is slow to verify, so we have a more strict requirement that the buffer type is the same
1002
+ // this is not uncommon since multiple backends can use host memory, with the same buffer type (eg. BLAS and CPU)
1003
+ // additionally, set remaining unassigned nodes to the backend with the most supported inputs
1004
+ // only nodes that could not be assigned during expansion due to the backend not supporting the op should be unassigned at this point
1005
+ for (int i = 0; i < graph->n_nodes; i++) {
1006
+ struct lm_ggml_tensor * node = graph->nodes[i];
1007
+ if (lm_ggml_is_view_op(node->op)) {
1008
+ continue;
1009
+ }
1010
+ int * node_backend_id = &tensor_backend_id(node);
1011
+ if (*node_backend_id == -1) {
1012
+ // unassigned node: find the backend with the most supported inputs
1013
+ int n_supported_best = -1;
1014
+ for (int b = 0; b < sched->n_backends; b++) {
1015
+ if (lm_ggml_backend_supports_op(sched->backends[b], node)) {
1016
+ int n_supported = 0;
1017
+ for (int j = 0; j < LM_GGML_MAX_SRC; j++) {
1018
+ struct lm_ggml_tensor * src = node->src[j];
1019
+ if (src == NULL) {
1020
+ continue;
1021
+ }
1022
+ if ((tensor_backend_id(src) != -1 || tensor_backend_id(src->view_src) != -1) && lm_ggml_backend_sched_buffer_supported(sched, src, b)) {
1023
+ n_supported++;
1024
+ }
1025
+ }
1026
+ if (n_supported > n_supported_best) {
1027
+ n_supported_best = n_supported;
1028
+ *node_backend_id = b;
1029
+ SET_CAUSE(node, "3.best");
1030
+ }
1031
+ }
1032
+ }
1033
+ } else {
1034
+ // assigned node: upgrade to higher prio backend if possible
1035
+ for (int b = 0; b < *node_backend_id; b++) {
1036
+ if (sched->bufts[b] == sched->bufts[*node_backend_id] && lm_ggml_backend_supports_op(sched->backends[b], node)) {
1037
+ bool supported = true;
1038
+ for (int j = 0; j < LM_GGML_MAX_SRC; j++) {
1039
+ struct lm_ggml_tensor * src = node->src[j];
1040
+ if (src == NULL) {
1041
+ continue;
1042
+ }
1043
+ if (!lm_ggml_backend_sched_buffer_supported(sched, src, b)) {
1044
+ supported = false;
1045
+ break;
1046
+ }
1047
+ }
1048
+ if (supported) {
1049
+ *node_backend_id = b;
1050
+ SET_CAUSE(node, "3.upg");
1051
+ break;
1052
+ }
1053
+ }
1054
+ }
1055
+ }
1056
+ }
1057
+
1058
+ // pass 4: assign backends to remaining src from dst and view_src
1059
+ for (int i = 0; i < graph->n_nodes; i++) {
1060
+ struct lm_ggml_tensor * node = graph->nodes[i];
1061
+ int * cur_backend_id = &tensor_backend_id(node);
1062
+ if (node->view_src != NULL && *cur_backend_id == -1) {
1063
+ *cur_backend_id = tensor_backend_id(node->view_src);
1064
+ SET_CAUSE(node, "4.vsrc");
1065
+ }
1066
+ for (int j = 0; j < LM_GGML_MAX_SRC; j++) {
1067
+ struct lm_ggml_tensor * src = node->src[j];
1068
+ if (src == NULL) {
1069
+ continue;
1070
+ }
1071
+ int * src_backend_id = &tensor_backend_id(src);
1072
+ if (*src_backend_id == -1) {
1073
+ if (src->view_src != NULL) {
1074
+ // views are always on the same backend as the source
1075
+ *src_backend_id = tensor_backend_id(src->view_src);
1076
+ SET_CAUSE(src, "4.vsrc");
1077
+ } else {
1078
+ *src_backend_id = *cur_backend_id;
1079
+ SET_CAUSE(src, "4.cur");
1080
+ }
1081
+ }
1082
+ }
1083
+ }
1084
+
1085
+ // pass 5: split graph, find tensors that need to be copied
1086
+ {
1087
+ int i_split = 0;
1088
+ struct lm_ggml_backend_sched_split * split = &sched->splits[0];
1089
+ // find the backend of the first split, skipping view ops
1090
+ int i = 0;
1091
+ for (; i < graph->n_nodes; i++) {
1092
+ struct lm_ggml_tensor * node = graph->nodes[i];
1093
+ if (!lm_ggml_is_view_op(node->op)) {
1094
+ split->backend_id = tensor_backend_id(node);
1095
+ break;
1096
+ }
1097
+ }
1098
+ split->i_start = 0;
1099
+ split->n_inputs = 0;
1100
+ int cur_backend_id = split->backend_id;
1101
+ for (; i < graph->n_nodes; i++) {
1102
+ struct lm_ggml_tensor * node = graph->nodes[i];
1103
+
1104
+ if (lm_ggml_is_view_op(node->op)) {
1105
+ continue;
1106
+ }
1107
+
1108
+ const int node_backend_id = tensor_backend_id(node);
1109
+
1110
+ assert(node_backend_id != -1); // all nodes should be assigned by now
1111
+
1112
+ // check if we should start a new split based on the sources of the current node
1113
+ bool need_new_split = false;
1114
+ if (node_backend_id == cur_backend_id && split->n_inputs > 0) {
1115
+ for (int j = 0; j < LM_GGML_MAX_SRC; j++) {
1116
+ struct lm_ggml_tensor * src = node->src[j];
1117
+ if (src == NULL) {
1118
+ continue;
1119
+ }
1120
+ // check if a weight is on a different and incompatible backend
1121
+ // by starting a new split, the memory of the previously offloaded weights can be reused
1122
+ if (src->buffer != NULL && src->buffer->usage == LM_GGML_BACKEND_BUFFER_USAGE_WEIGHTS) {
1123
+ int src_backend_id = tensor_backend_id(src);
1124
+ if (src_backend_id != cur_backend_id && !lm_ggml_backend_sched_buffer_supported(sched, src, cur_backend_id)) {
1125
+ need_new_split = true;
1126
+ break;
1127
+ }
1128
+ }
1129
+ // check if the split has too many inputs
1130
+ // FIXME: count the number of inputs instead of only checking when full
1131
+ if (split->n_inputs == LM_GGML_SCHED_MAX_SPLIT_INPUTS) {
1132
+ const size_t id = hash_id(src);
1133
+ int src_backend_id = sched->hv_tensor_backend_ids[id];
1134
+ bool supported = lm_ggml_backend_sched_buffer_supported(sched, src, cur_backend_id);
1135
+ if (src_backend_id != cur_backend_id && tensor_id_copy(id, cur_backend_id, 0) == NULL && !supported) {
1136
+ need_new_split = true;
1137
+ break;
1138
+ }
1139
+ }
1140
+ }
1141
+ }
1142
+
1143
+ if (node_backend_id != cur_backend_id || need_new_split) {
1144
+ split->i_end = i;
1145
+ i_split++;
1146
+ if (i_split >= sched->splits_capacity) {
1147
+ sched->splits_capacity *= 2;
1148
+ sched->splits = (lm_ggml_backend_sched_split *)
1149
+ realloc(sched->splits, sched->splits_capacity * sizeof(struct lm_ggml_backend_sched_split));
1150
+ LM_GGML_ASSERT(sched->splits != NULL);
1151
+ }
1152
+ split = &sched->splits[i_split];
1153
+ split->backend_id = node_backend_id;
1154
+ split->i_start = i;
1155
+ split->n_inputs = 0;
1156
+ cur_backend_id = node_backend_id;
1157
+ }
1158
+
1159
+ // find inputs that are not on the same backend
1160
+ for (int j = 0; j < LM_GGML_MAX_SRC; j++) {
1161
+ struct lm_ggml_tensor * src = node->src[j];
1162
+ if (src == NULL) {
1163
+ continue;
1164
+ }
1165
+
1166
+ size_t src_id = hash_id(src);
1167
+ const int src_backend_id = sched->hv_tensor_backend_ids[src_id];
1168
+ assert(src_backend_id != -1); // all inputs should be assigned by now
1169
+
1170
+ if (src->flags & LM_GGML_TENSOR_FLAG_INPUT && sched->n_copies > 1) {
1171
+ if (tensor_id_copy(src_id, src_backend_id, 0) == NULL) {
1172
+ lm_ggml_backend_t backend = sched->backends[src_backend_id];
1173
+ for (int c = 0; c < sched->n_copies; c++) {
1174
+ struct lm_ggml_tensor * tensor_copy;
1175
+ if (c == sched->cur_copy) {
1176
+ tensor_copy = src; // use the original tensor as the current copy
1177
+ } else {
1178
+ tensor_copy = lm_ggml_dup_tensor_layout(sched->ctx, src);
1179
+ lm_ggml_format_name(tensor_copy, "%s#%s#%d", lm_ggml_backend_name(backend), src->name, c);
1180
+ }
1181
+ if (sched->n_copies > 1) {
1182
+ lm_ggml_set_input(tensor_copy);
1183
+ lm_ggml_set_output(tensor_copy); // prevent ggml-alloc from overwriting the tensor
1184
+ }
1185
+ tensor_id_copy(src_id, src_backend_id, c) = tensor_copy;
1186
+ SET_CAUSE(tensor_copy, "4.cpy");
1187
+ }
1188
+ int n_graph_inputs = sched->n_graph_inputs++;
1189
+ LM_GGML_ASSERT(n_graph_inputs < LM_GGML_SCHED_MAX_SPLIT_INPUTS);
1190
+ sched->graph_inputs[n_graph_inputs] = src;
1191
+ }
1192
+ }
1193
+
1194
+ if (src_backend_id != cur_backend_id && !lm_ggml_backend_sched_buffer_supported(sched, src, cur_backend_id)) {
1195
+ // create a copy of the input in the split's backend
1196
+ if (tensor_id_copy(src_id, cur_backend_id, 0) == NULL) {
1197
+ lm_ggml_backend_t backend = sched->backends[cur_backend_id];
1198
+ for (int c = 0; c < sched->n_copies; c++) {
1199
+ struct lm_ggml_tensor * tensor_copy = lm_ggml_dup_tensor_layout(sched->ctx, src);
1200
+ lm_ggml_format_name(tensor_copy, "%s#%s#%d", lm_ggml_backend_name(backend), src->name, c);
1201
+ if (sched->n_copies > 1) {
1202
+ lm_ggml_set_input(tensor_copy);
1203
+ lm_ggml_set_output(tensor_copy); // prevent ggml-alloc from overwriting the tensor
1204
+ }
1205
+ tensor_id_copy(src_id, cur_backend_id, c) = tensor_copy;
1206
+ SET_CAUSE(tensor_copy, "4.cpy");
1207
+ }
1208
+ int n_inputs = split->n_inputs++;
1209
+ LM_GGML_ASSERT(n_inputs < LM_GGML_SCHED_MAX_SPLIT_INPUTS);
1210
+ split->inputs[n_inputs] = src;
1211
+ }
1212
+ node->src[j] = tensor_id_copy(src_id, cur_backend_id, sched->cur_copy);
1213
+ }
1214
+ }
1215
+ }
1216
+ split->i_end = graph->n_nodes;
1217
+ sched->n_splits = i_split + 1;
1218
+ }
1219
+
1220
+ if (sched->debug) {
1221
+ lm_ggml_backend_sched_print_assignments(sched, graph);
1222
+ }
1223
+
1224
+ // swap node_backend_ids and leaf _backend_ids with prevs
1225
+ {
1226
+ int * tmp = sched->node_backend_ids;
1227
+ sched->node_backend_ids = sched->prev_node_backend_ids;
1228
+ sched->prev_node_backend_ids = tmp;
1229
+
1230
+ tmp = sched->leaf_backend_ids;
1231
+ sched->leaf_backend_ids = sched->prev_leaf_backend_ids;
1232
+ sched->prev_leaf_backend_ids = tmp;
1233
+ }
1234
+
1235
+ int graph_size = std::max(graph->n_nodes, graph->n_leafs) + sched->n_splits*LM_GGML_SCHED_MAX_SPLIT_INPUTS*2*sched->n_copies;
1236
+ if (sched->graph.size < graph_size) {
1237
+ sched->graph.size = graph_size;
1238
+ sched->graph.nodes = (lm_ggml_tensor **) realloc(sched->graph.nodes, graph_size * sizeof(struct lm_ggml_tensor *));
1239
+ sched->graph.leafs = (lm_ggml_tensor **) realloc(sched->graph.leafs, graph_size * sizeof(struct lm_ggml_tensor *));
1240
+ LM_GGML_ASSERT(sched->graph.nodes != NULL);
1241
+ LM_GGML_ASSERT(sched->graph.leafs != NULL);
1242
+ }
1243
+ sched->graph.n_nodes = 0;
1244
+ sched->graph.n_leafs = 0;
1245
+
1246
+ struct lm_ggml_cgraph * graph_copy = &sched->graph;
1247
+
1248
+ for (int i = 0; i < sched->n_splits; i++) {
1249
+ struct lm_ggml_backend_sched_split * split = &sched->splits[i];
1250
+ split->graph = lm_ggml_graph_view(graph, split->i_start, split->i_end);
1251
+
1252
+ // add inputs to the graph copy so that they are allocated by ggml-alloc at the start of the split
1253
+ for (int j = 0; j < split->n_inputs; j++) {
1254
+ assert(graph_copy->size > (graph_copy->n_nodes + 1));
1255
+
1256
+ struct lm_ggml_tensor * input = split->inputs[j];
1257
+ const size_t input_id = hash_id(input);
1258
+ struct lm_ggml_tensor * input_cpy = tensor_id_copy(input_id, split->backend_id, sched->cur_copy);
1259
+
1260
+ // add a dependency to the input source so that it is not freed before the copy is done
1261
+ struct lm_ggml_tensor * input_dep = lm_ggml_view_tensor(sched->ctx, input);
1262
+ input_dep->src[0] = input;
1263
+ sched->node_backend_ids[graph_copy->n_nodes] = sched->hv_tensor_backend_ids[input_id];
1264
+ graph_copy->nodes[graph_copy->n_nodes++] = input_dep;
1265
+
1266
+ // add a dependency to the input copy so that it is allocated at the start of the split
1267
+ sched->node_backend_ids[graph_copy->n_nodes] = split->backend_id;
1268
+ graph_copy->nodes[graph_copy->n_nodes++] = input_cpy;
1269
+ }
1270
+
1271
+ for (int j = split->i_start; j < split->i_end; j++) {
1272
+ assert(graph_copy->size > graph_copy->n_nodes);
1273
+ sched->node_backend_ids[graph_copy->n_nodes] = tensor_backend_id(graph->nodes[j]);
1274
+ graph_copy->nodes[graph_copy->n_nodes++] = graph->nodes[j];
1275
+ }
1276
+ }
1277
+
1278
+ if (sched->n_copies > 1) {
1279
+ // add input copies as leafs so that they are allocated first
1280
+ for (int i = 0; i < sched->n_graph_inputs; i++) {
1281
+ struct lm_ggml_tensor * input = sched->graph_inputs[i];
1282
+ size_t id = hash_id(input);
1283
+ int backend_id = tensor_backend_id(input);
1284
+ for (int c = 0; c < sched->n_copies; c++) {
1285
+ struct lm_ggml_tensor * input_cpy = tensor_id_copy(id, backend_id, c);
1286
+ sched->leaf_backend_ids[graph_copy->n_leafs] = backend_id;
1287
+ assert(graph_copy->size > graph_copy->n_leafs);
1288
+ graph_copy->leafs[graph_copy->n_leafs++] = input_cpy;
1289
+ }
1290
+ }
1291
+
1292
+ for (int i = 0; i < sched->n_splits; i++) {
1293
+ struct lm_ggml_backend_sched_split * split = &sched->splits[i];
1294
+ int backend_id = split->backend_id;
1295
+ for (int j = 0; j < split->n_inputs; j++) {
1296
+ struct lm_ggml_tensor * input = split->inputs[j];
1297
+ size_t id = hash_id(input);
1298
+ for (int c = 0; c < sched->n_copies; c++) {
1299
+ struct lm_ggml_tensor * input_cpy = tensor_id_copy(id, backend_id, c);
1300
+ sched->leaf_backend_ids[graph_copy->n_leafs] = backend_id;
1301
+ assert(graph_copy->size > graph_copy->n_leafs);
1302
+ graph_copy->leafs[graph_copy->n_leafs++] = input_cpy;
1303
+ }
1304
+ }
1305
+ }
1306
+ }
1307
+
1308
+ // add leafs from the original graph
1309
+ for (int i = 0; i < graph->n_leafs; i++) {
1310
+ struct lm_ggml_tensor * leaf = graph->leafs[i];
1311
+ sched->leaf_backend_ids[graph_copy->n_leafs] = tensor_backend_id(leaf);
1312
+ assert(graph_copy->size > graph_copy->n_leafs);
1313
+ graph_copy->leafs[graph_copy->n_leafs++] = leaf;
1314
+ }
1315
+ }
1316
+
1317
+ static bool lm_ggml_backend_sched_alloc_splits(lm_ggml_backend_sched_t sched) {
1318
+ bool backend_ids_changed = false;
1319
+ for (int i = 0; i < sched->graph.n_nodes; i++) {
1320
+ if (sched->node_backend_ids[i] != sched->prev_node_backend_ids[i] &&
1321
+ sched->bufts[sched->node_backend_ids[i]] != sched->bufts[sched->prev_node_backend_ids[i]]) {
1322
+ backend_ids_changed = true;
1323
+ break;
1324
+ }
1325
+ }
1326
+ if (!backend_ids_changed) {
1327
+ for (int i = 0; i < sched->graph.n_leafs; i++) {
1328
+ if (sched->leaf_backend_ids[i] != sched->prev_leaf_backend_ids[i] &&
1329
+ sched->bufts[sched->leaf_backend_ids[i]] != sched->bufts[sched->prev_leaf_backend_ids[i]]) {
1330
+ backend_ids_changed = true;
1331
+ break;
1332
+ }
1333
+ }
1334
+ }
1335
+
1336
+ // allocate graph
1337
+ if (backend_ids_changed || !lm_ggml_gallocr_alloc_graph(sched->galloc, &sched->graph)) {
1338
+ // the re-allocation may cause the split inputs to be moved to a different address
1339
+ lm_ggml_backend_sched_synchronize(sched);
1340
+ #ifndef NDEBUG
1341
+ LM_GGML_LOG_DEBUG("%s: failed to allocate graph, reserving (backend_ids_changed = %d)\n", __func__, backend_ids_changed);
1342
+ #endif
1343
+ lm_ggml_gallocr_reserve_n(sched->galloc, &sched->graph, sched->node_backend_ids, sched->leaf_backend_ids);
1344
+ if (!lm_ggml_gallocr_alloc_graph(sched->galloc, &sched->graph)) {
1345
+ LM_GGML_LOG_ERROR("%s: failed to allocate graph\n", __func__);
1346
+ return false;
1347
+ }
1348
+ }
1349
+
1350
+ return true;
1351
+ }
1352
+
1353
+ static enum lm_ggml_status lm_ggml_backend_sched_compute_splits(lm_ggml_backend_sched_t sched) {
1354
+ struct lm_ggml_backend_sched_split * splits = sched->splits;
1355
+
1356
+ for (int i = 0; i < sched->n_splits; i++) {
1357
+ struct lm_ggml_backend_sched_split * split = &splits[i];
1358
+ int split_backend_id = split->backend_id;
1359
+ lm_ggml_backend_t split_backend = sched->backends[split_backend_id];
1360
+
1361
+ // copy the input tensors to the split backend
1362
+ for (int j = 0; j < split->n_inputs; j++) {
1363
+ lm_ggml_backend_t input_backend = lm_ggml_backend_sched_get_tensor_backend(sched, split->inputs[j]);
1364
+ struct lm_ggml_tensor * input = split->inputs[j];
1365
+ struct lm_ggml_tensor * input_cpy = tensor_copy(input, split_backend_id, sched->cur_copy);
1366
+
1367
+ if (input->flags & LM_GGML_TENSOR_FLAG_INPUT) {
1368
+ // inputs from the user must be copied immediately to prevent the user overwriting the data before the copy is done
1369
+ if (sched->events[split_backend_id][sched->cur_copy] != NULL) {
1370
+ lm_ggml_backend_event_synchronize(sched->events[split_backend_id][sched->cur_copy]);
1371
+ } else {
1372
+ lm_ggml_backend_synchronize(split_backend);
1373
+ }
1374
+ lm_ggml_backend_tensor_copy(input, input_cpy);
1375
+ } else {
1376
+ // wait for the split backend to finish using the input before overwriting it
1377
+ if (sched->events[split_backend_id][sched->cur_copy] != NULL) {
1378
+ lm_ggml_backend_event_wait(split_backend, sched->events[split_backend_id][sched->cur_copy]);
1379
+ } else {
1380
+ lm_ggml_backend_synchronize(split_backend);
1381
+ }
1382
+ // try async copy, but if not possible, we can still use a sync copy without synchronizing the dst backend, since we handle the synchronization here with multiple copies and events
1383
+ // TODO: add public function to facilitate this, since applications do not have direct access to the backend interface
1384
+ if (!split_backend->iface.cpy_tensor_async || !split_backend->iface.cpy_tensor_async(input_backend, split_backend, input, input_cpy)) {
1385
+ lm_ggml_backend_synchronize(input_backend);
1386
+ if (sched->events[split_backend_id][sched->cur_copy] != NULL) {
1387
+ lm_ggml_backend_event_synchronize(sched->events[split_backend_id][sched->cur_copy]);
1388
+ } else {
1389
+ lm_ggml_backend_synchronize(split_backend);
1390
+ }
1391
+ lm_ggml_backend_tensor_copy(input, input_cpy);
1392
+ }
1393
+ }
1394
+ }
1395
+
1396
+ if (!sched->callback_eval) {
1397
+ enum lm_ggml_status ec = lm_ggml_backend_graph_compute_async(split_backend, &split->graph);
1398
+ if (ec != LM_GGML_STATUS_SUCCESS) {
1399
+ return ec;
1400
+ }
1401
+ } else {
1402
+ // similar to lm_ggml_backend_compare_graph_backend
1403
+ for (int j0 = 0; j0 < split->graph.n_nodes; j0++) {
1404
+ struct lm_ggml_tensor * t = split->graph.nodes[j0];
1405
+
1406
+ // check if the user needs data from this node
1407
+ bool need = sched->callback_eval(t, true, sched->callback_eval_user_data);
1408
+
1409
+ int j1 = j0;
1410
+
1411
+ // determine the range [j0, j1] of nodes that can be computed together
1412
+ while (!need && j1 < split->graph.n_nodes - 1) {
1413
+ t = split->graph.nodes[++j1];
1414
+ need = sched->callback_eval(t, true, sched->callback_eval_user_data);
1415
+ }
1416
+
1417
+ struct lm_ggml_cgraph gv = lm_ggml_graph_view(&split->graph, j0, j1 + 1);
1418
+
1419
+ enum lm_ggml_status ec = lm_ggml_backend_graph_compute_async(split_backend, &gv);
1420
+ if (ec != LM_GGML_STATUS_SUCCESS) {
1421
+ return ec;
1422
+ }
1423
+
1424
+ // TODO: pass backend to the callback, then the user can decide if they want to synchronize
1425
+ lm_ggml_backend_synchronize(split_backend);
1426
+
1427
+ if (need && !sched->callback_eval(t, false, sched->callback_eval_user_data)) {
1428
+ break;
1429
+ }
1430
+
1431
+ j0 = j1;
1432
+ }
1433
+ }
1434
+
1435
+ // record the event of this copy
1436
+ if (split->n_inputs > 0) {
1437
+ if (sched->events[split_backend_id][sched->cur_copy] != NULL) {
1438
+ lm_ggml_backend_event_record(sched->events[split_backend_id][sched->cur_copy], split_backend);
1439
+ }
1440
+ }
1441
+ }
1442
+
1443
+ sched->cur_copy = (sched->cur_copy + 1) % sched->n_copies;
1444
+
1445
+ return LM_GGML_STATUS_SUCCESS;
1446
+ }
1447
+
1448
+ lm_ggml_backend_sched_t lm_ggml_backend_sched_new(
1449
+ lm_ggml_backend_t * backends,
1450
+ lm_ggml_backend_buffer_type_t * bufts,
1451
+ int n_backends,
1452
+ size_t graph_size,
1453
+ bool parallel) {
1454
+ LM_GGML_ASSERT(n_backends > 0);
1455
+ LM_GGML_ASSERT(n_backends <= LM_GGML_SCHED_MAX_BACKENDS);
1456
+ LM_GGML_ASSERT(lm_ggml_backend_dev_type(lm_ggml_backend_get_device(backends[n_backends - 1])) == LM_GGML_BACKEND_DEVICE_TYPE_CPU);
1457
+
1458
+ struct lm_ggml_backend_sched * sched = (lm_ggml_backend_sched *) calloc(1, sizeof(struct lm_ggml_backend_sched));
1459
+
1460
+ const char * LM_GGML_SCHED_DEBUG = getenv("LM_GGML_SCHED_DEBUG");
1461
+ sched->debug = LM_GGML_SCHED_DEBUG ? atoi(LM_GGML_SCHED_DEBUG) : 0;
1462
+ sched->n_backends = n_backends;
1463
+ sched->n_copies = parallel ? LM_GGML_SCHED_MAX_COPIES : 1;
1464
+
1465
+ // initialize hash table
1466
+ // FIXME: needs to be size*2 to account for leafs (do it in graph_split instead)
1467
+ sched->hash_set = lm_ggml_hash_set_new(graph_size);
1468
+ sched->hv_tensor_backend_ids = (int *) malloc(sched->hash_set.size * sizeof(sched->hv_tensor_backend_ids[0]));
1469
+ sched->hv_tensor_copies = (lm_ggml_tensor **) malloc(sched->hash_set.size * sched->n_backends * sched->n_copies * sizeof(struct lm_ggml_tensor *));
1470
+
1471
+ const size_t lm_ggml_sched_max_splits = graph_size; // at most there is one split for each node in the graph
1472
+ const size_t nodes_size = graph_size + lm_ggml_sched_max_splits*LM_GGML_SCHED_MAX_SPLIT_INPUTS*2;
1473
+ sched->node_backend_ids = (int *) calloc(nodes_size, sizeof(sched->node_backend_ids[0]));
1474
+ sched->leaf_backend_ids = (int *) calloc(nodes_size, sizeof(sched->leaf_backend_ids[0]));
1475
+ sched->prev_node_backend_ids = (int *) calloc(nodes_size, sizeof(sched->prev_node_backend_ids[0]));
1476
+ sched->prev_leaf_backend_ids = (int *) calloc(nodes_size, sizeof(sched->prev_leaf_backend_ids[0]));
1477
+
1478
+ sched->context_buffer_size = lm_ggml_sched_max_splits*LM_GGML_SCHED_MAX_SPLIT_INPUTS*2*sizeof(struct lm_ggml_tensor) + lm_ggml_graph_overhead_custom(graph_size, false);
1479
+ sched->context_buffer = (char *) malloc(sched->context_buffer_size);
1480
+
1481
+ const int initial_splits_capacity = 16;
1482
+ sched->splits = (lm_ggml_backend_sched_split *) calloc(initial_splits_capacity, sizeof(sched->splits[0]));
1483
+ sched->splits_capacity = initial_splits_capacity;
1484
+
1485
+ for (int b = 0; b < n_backends; b++) {
1486
+ sched->backends[b] = backends[b];
1487
+ sched->bufts[b] = bufts ? bufts[b] : lm_ggml_backend_get_default_buffer_type(backends[b]);
1488
+ LM_GGML_ASSERT(lm_ggml_backend_supports_buft(backends[b], sched->bufts[b]));
1489
+
1490
+ if (sched->n_copies > 1) {
1491
+ for (int c = 0; c < sched->n_copies; c++) {
1492
+ sched->events[b][c] = lm_ggml_backend_event_new(backends[b]->device);
1493
+ }
1494
+ }
1495
+ }
1496
+
1497
+ sched->galloc = lm_ggml_gallocr_new_n(sched->bufts, n_backends);
1498
+
1499
+ lm_ggml_backend_sched_reset(sched);
1500
+
1501
+ return sched;
1502
+ }
1503
+
1504
+ void lm_ggml_backend_sched_free(lm_ggml_backend_sched_t sched) {
1505
+ if (sched == NULL) {
1506
+ return;
1507
+ }
1508
+ for (int b = 0; b < sched->n_backends; b++) {
1509
+ for (int c = 0; c < sched->n_copies; c++) {
1510
+ lm_ggml_backend_event_free(sched->events[b][c]);
1511
+ }
1512
+ }
1513
+ lm_ggml_gallocr_free(sched->galloc);
1514
+ lm_ggml_free(sched->ctx);
1515
+ lm_ggml_hash_set_free(&sched->hash_set);
1516
+ free(sched->splits);
1517
+ free(sched->hv_tensor_backend_ids);
1518
+ free(sched->hv_tensor_copies);
1519
+ free(sched->node_backend_ids);
1520
+ free(sched->leaf_backend_ids);
1521
+ free(sched->prev_node_backend_ids);
1522
+ free(sched->prev_leaf_backend_ids);
1523
+ free(sched->context_buffer);
1524
+ free(sched->graph.nodes);
1525
+ free(sched->graph.leafs);
1526
+ free(sched);
1527
+ }
1528
+
1529
+ void lm_ggml_backend_sched_reset(lm_ggml_backend_sched_t sched) {
1530
+ // reset state for the next run
1531
+ if (!sched->is_reset) {
1532
+ lm_ggml_hash_set_reset(&sched->hash_set);
1533
+ memset(sched->hv_tensor_backend_ids, -1, sched->hash_set.size * sizeof(sched->hv_tensor_backend_ids[0]));
1534
+ memset(sched->hv_tensor_copies, 0, sched->hash_set.size * sched->n_backends * sched->n_copies * sizeof(struct lm_ggml_tensor *));
1535
+ sched->is_reset = true;
1536
+ }
1537
+ sched->is_alloc = false;
1538
+ }
1539
+
1540
+ bool lm_ggml_backend_sched_reserve(lm_ggml_backend_sched_t sched, struct lm_ggml_cgraph * measure_graph) {
1541
+ LM_GGML_ASSERT((int)sched->hash_set.size >= measure_graph->n_nodes + measure_graph->n_leafs);
1542
+
1543
+ lm_ggml_backend_sched_split_graph(sched, measure_graph);
1544
+
1545
+ lm_ggml_backend_sched_synchronize(sched);
1546
+
1547
+ if (!lm_ggml_gallocr_reserve_n(sched->galloc, &sched->graph, sched->node_backend_ids, sched->leaf_backend_ids)) {
1548
+ return false;
1549
+ }
1550
+
1551
+ lm_ggml_backend_sched_reset(sched);
1552
+
1553
+ return true;
1554
+ }
1555
+
1556
+ bool lm_ggml_backend_sched_alloc_graph(lm_ggml_backend_sched_t sched, struct lm_ggml_cgraph * graph) {
1557
+ LM_GGML_ASSERT((int)sched->hash_set.size >= graph->n_nodes + graph->n_leafs);
1558
+
1559
+ lm_ggml_backend_sched_split_graph(sched, graph);
1560
+
1561
+
1562
+ if (!lm_ggml_backend_sched_alloc_splits(sched)) {
1563
+ return false;
1564
+ }
1565
+
1566
+ sched->is_alloc = true;
1567
+
1568
+ return true;
1569
+ }
1570
+
1571
+ enum lm_ggml_status lm_ggml_backend_sched_graph_compute(lm_ggml_backend_sched_t sched, struct lm_ggml_cgraph * graph) {
1572
+ enum lm_ggml_status err = lm_ggml_backend_sched_graph_compute_async(sched, graph);
1573
+ lm_ggml_backend_sched_synchronize(sched);
1574
+ return err;
1575
+ }
1576
+
1577
+ enum lm_ggml_status lm_ggml_backend_sched_graph_compute_async(lm_ggml_backend_sched_t sched, struct lm_ggml_cgraph * graph) {
1578
+ if (!sched->is_reset && !sched->is_alloc) {
1579
+ lm_ggml_backend_sched_reset(sched);
1580
+ }
1581
+
1582
+ if (!sched->is_alloc) {
1583
+ if (!lm_ggml_backend_sched_alloc_graph(sched, graph)) {
1584
+ return LM_GGML_STATUS_ALLOC_FAILED;
1585
+ }
1586
+ }
1587
+
1588
+ return lm_ggml_backend_sched_compute_splits(sched);
1589
+ }
1590
+
1591
+ void lm_ggml_backend_sched_synchronize(lm_ggml_backend_sched_t sched) {
1592
+ for (int i = 0; i < sched->n_backends; i++) {
1593
+ lm_ggml_backend_synchronize(sched->backends[i]);
1594
+ }
1595
+ }
1596
+
1597
+ void lm_ggml_backend_sched_set_eval_callback(lm_ggml_backend_sched_t sched, lm_ggml_backend_sched_eval_callback callback, void * user_data) {
1598
+ sched->callback_eval = callback;
1599
+ sched->callback_eval_user_data = user_data;
1600
+ }
1601
+
1602
+ int lm_ggml_backend_sched_get_n_splits(lm_ggml_backend_sched_t sched) {
1603
+ return sched->n_splits;
1604
+ }
1605
+
1606
+ int lm_ggml_backend_sched_get_n_copies(lm_ggml_backend_sched_t sched) {
1607
+ return sched->n_copies;
1608
+ }
1609
+
1610
+ int lm_ggml_backend_sched_get_n_backends(lm_ggml_backend_sched_t sched) {
1611
+ return sched->n_backends;
1612
+ }
1613
+
1614
+ lm_ggml_backend_t lm_ggml_backend_sched_get_backend(lm_ggml_backend_sched_t sched, int i) {
1615
+ LM_GGML_ASSERT(i >= 0 && i < sched->n_backends);
1616
+ return sched->backends[i];
1617
+ }
1618
+
1619
+ size_t lm_ggml_backend_sched_get_buffer_size(lm_ggml_backend_sched_t sched, lm_ggml_backend_t backend) {
1620
+ int backend_index = lm_ggml_backend_sched_backend_id(sched, backend);
1621
+ LM_GGML_ASSERT(backend_index >= 0 && backend_index < sched->n_backends);
1622
+
1623
+ return lm_ggml_gallocr_get_buffer_size(sched->galloc, backend_index);
1624
+ }
1625
+
1626
+ void lm_ggml_backend_sched_set_tensor_backend(lm_ggml_backend_sched_t sched, struct lm_ggml_tensor * node, lm_ggml_backend_t backend) {
1627
+ int backend_index = lm_ggml_backend_sched_backend_id(sched, backend);
1628
+ LM_GGML_ASSERT(backend_index >= 0 && backend_index < sched->n_backends);
1629
+ tensor_backend_id(node) = backend_index;
1630
+ SET_CAUSE(node, "usr");
1631
+ sched->is_reset = false;
1632
+ }
1633
+
1634
+ lm_ggml_backend_t lm_ggml_backend_sched_get_tensor_backend(lm_ggml_backend_sched_t sched, struct lm_ggml_tensor * node) {
1635
+ int backend_index = tensor_backend_id(node);
1636
+ if (backend_index == -1) {
1637
+ return NULL;
1638
+ }
1639
+ return sched->backends[backend_index];
1640
+ }
1641
+
1642
+ // utils
1643
+
1644
+ void lm_ggml_backend_view_init(struct lm_ggml_tensor * tensor) {
1645
+ LM_GGML_ASSERT(tensor->buffer == NULL);
1646
+ LM_GGML_ASSERT(tensor->view_src != NULL);
1647
+ LM_GGML_ASSERT(tensor->view_src->buffer != NULL);
1648
+ LM_GGML_ASSERT(tensor->view_src->data != NULL);
1649
+
1650
+ tensor->buffer = tensor->view_src->buffer;
1651
+ tensor->data = (char *)tensor->view_src->data + tensor->view_offs;
1652
+ lm_ggml_backend_buffer_init_tensor(tensor->buffer, tensor);
1653
+ }
1654
+
1655
+ void lm_ggml_backend_tensor_alloc(lm_ggml_backend_buffer_t buffer, struct lm_ggml_tensor * tensor, void * addr) {
1656
+ LM_GGML_ASSERT(tensor->buffer == NULL);
1657
+ LM_GGML_ASSERT(tensor->data == NULL);
1658
+ LM_GGML_ASSERT(tensor->view_src == NULL);
1659
+ LM_GGML_ASSERT(addr >= lm_ggml_backend_buffer_get_base(buffer));
1660
+ LM_GGML_ASSERT((char *)addr + lm_ggml_backend_buffer_get_alloc_size(buffer, tensor) <=
1661
+ (char *)lm_ggml_backend_buffer_get_base(buffer) + lm_ggml_backend_buffer_get_size(buffer));
1662
+
1663
+ tensor->buffer = buffer;
1664
+ tensor->data = addr;
1665
+ lm_ggml_backend_buffer_init_tensor(buffer, tensor);
1666
+ }
1667
+
1668
+ static struct lm_ggml_tensor * graph_copy_dup_tensor(struct lm_ggml_hash_set hash_set, struct lm_ggml_tensor ** node_copies,
1669
+ struct lm_ggml_context * ctx_allocated, struct lm_ggml_context * ctx_unallocated, struct lm_ggml_tensor * src) {
1670
+
1671
+ LM_GGML_ASSERT(src != NULL);
1672
+ LM_GGML_ASSERT(src->data && "graph must be allocated");
1673
+
1674
+ size_t id = lm_ggml_hash_insert(&hash_set, src);
1675
+ if (id == LM_GGML_HASHSET_ALREADY_EXISTS) {
1676
+ return node_copies[lm_ggml_hash_find(&hash_set, src)];
1677
+ }
1678
+
1679
+ struct lm_ggml_tensor * dst = lm_ggml_dup_tensor_layout(src->data && !src->view_src ? ctx_allocated : ctx_unallocated, src);
1680
+ if (src->view_src != NULL) {
1681
+ dst->view_src = graph_copy_dup_tensor(hash_set, node_copies, ctx_allocated, ctx_unallocated, src->view_src);
1682
+ dst->view_offs = src->view_offs;
1683
+ }
1684
+ dst->op = src->op;
1685
+ memcpy(dst->op_params, src->op_params, sizeof(dst->op_params));
1686
+ lm_ggml_set_name(dst, src->name);
1687
+
1688
+ // copy src
1689
+ for (int i = 0; i < LM_GGML_MAX_SRC; i++) {
1690
+ struct lm_ggml_tensor * s = src->src[i];
1691
+ if (s == NULL) {
1692
+ continue;
1693
+ }
1694
+ dst->src[i] = graph_copy_dup_tensor(hash_set, node_copies, ctx_allocated, ctx_unallocated, s);
1695
+ }
1696
+
1697
+ node_copies[id] = dst;
1698
+ return dst;
1699
+ }
1700
+
1701
+ static void graph_copy_init_tensor(struct lm_ggml_hash_set * hash_set, struct lm_ggml_tensor ** node_copies, bool * node_init, struct lm_ggml_tensor * src) {
1702
+ size_t id = lm_ggml_hash_find(hash_set, src);
1703
+ if (node_init[id]) {
1704
+ return;
1705
+ }
1706
+ node_init[id] = true;
1707
+
1708
+ struct lm_ggml_tensor * dst = node_copies[id];
1709
+ if (dst->view_src != NULL) {
1710
+ graph_copy_init_tensor(hash_set, node_copies, node_init, src->view_src);
1711
+ lm_ggml_backend_view_init(dst);
1712
+ }
1713
+ else {
1714
+ lm_ggml_backend_tensor_copy(src, dst);
1715
+ }
1716
+
1717
+ // init src
1718
+ for (int i = 0; i < LM_GGML_MAX_SRC; i++) {
1719
+ struct lm_ggml_tensor * s = src->src[i];
1720
+ if (s == NULL) {
1721
+ continue;
1722
+ }
1723
+ graph_copy_init_tensor(hash_set, node_copies, node_init, s);
1724
+ }
1725
+ }
1726
+
1727
+ struct lm_ggml_backend_graph_copy lm_ggml_backend_graph_copy(lm_ggml_backend_t backend, struct lm_ggml_cgraph * graph) {
1728
+ struct lm_ggml_hash_set hash_set = lm_ggml_hash_set_new(graph->visited_hash_set.size);
1729
+ struct lm_ggml_tensor ** node_copies = (lm_ggml_tensor **) calloc(hash_set.size, sizeof(node_copies[0])); // NOLINT
1730
+ bool * node_init = (bool *) calloc(hash_set.size, sizeof(node_init[0]));
1731
+
1732
+ struct lm_ggml_init_params params = {
1733
+ /* .mem_size = */ lm_ggml_tensor_overhead()*hash_set.size + lm_ggml_graph_overhead_custom(graph->size, false),
1734
+ /* .mem_buffer = */ NULL,
1735
+ /* .no_alloc = */ true
1736
+ };
1737
+
1738
+ struct lm_ggml_context * ctx_allocated = lm_ggml_init(params);
1739
+ struct lm_ggml_context * ctx_unallocated = lm_ggml_init(params);
1740
+
1741
+ if (ctx_allocated == NULL || ctx_unallocated == NULL) {
1742
+ LM_GGML_LOG_ERROR("%s: failed to allocate context for graph copy\n", __func__);
1743
+ lm_ggml_hash_set_free(&hash_set);
1744
+ free(node_copies);
1745
+ free(node_init);
1746
+ lm_ggml_free(ctx_allocated);
1747
+ lm_ggml_free(ctx_unallocated);
1748
+ return {
1749
+ /* .buffer = */ NULL,
1750
+ /* .ctx_allocated = */ NULL,
1751
+ /* .ctx_unallocated = */ NULL,
1752
+ /* .graph = */ NULL,
1753
+ };
1754
+ }
1755
+
1756
+ // dup nodes
1757
+ for (int i = 0; i < graph->n_nodes; i++) {
1758
+ struct lm_ggml_tensor * node = graph->nodes[i];
1759
+ graph_copy_dup_tensor(hash_set, node_copies, ctx_allocated, ctx_unallocated, node);
1760
+ }
1761
+
1762
+ // allocate nodes
1763
+ lm_ggml_backend_buffer_t buffer = lm_ggml_backend_alloc_ctx_tensors(ctx_allocated, backend);
1764
+ if (buffer == NULL) {
1765
+ LM_GGML_LOG_ERROR("%s: failed to allocate buffer for graph copy\n", __func__);
1766
+ lm_ggml_hash_set_free(&hash_set);
1767
+ free(node_copies);
1768
+ free(node_init);
1769
+ lm_ggml_free(ctx_allocated);
1770
+ lm_ggml_free(ctx_unallocated);
1771
+ return {
1772
+ /* .buffer = */ NULL,
1773
+ /* .ctx_allocated = */ NULL,
1774
+ /* .ctx_unallocated = */ NULL,
1775
+ /* .graph = */ NULL,
1776
+ };
1777
+ }
1778
+
1779
+ //printf("copy buffer size: %zu MB\n", lm_ggml_backend_buffer_get_size(buffer) / 1024 / 1024);
1780
+
1781
+ // copy data and init views
1782
+ for (int i = 0; i < graph->n_nodes; i++) {
1783
+ struct lm_ggml_tensor * node = graph->nodes[i];
1784
+ graph_copy_init_tensor(&hash_set, node_copies, node_init, node);
1785
+ }
1786
+
1787
+ // build graph copy
1788
+ struct lm_ggml_cgraph * graph_copy = lm_ggml_new_graph_custom(ctx_allocated, graph->size, false);
1789
+ for (int i = 0; i < graph->n_nodes; i++) {
1790
+ struct lm_ggml_tensor * node = graph->nodes[i];
1791
+ struct lm_ggml_tensor * node_copy = node_copies[lm_ggml_hash_find(&hash_set, node)];
1792
+ graph_copy->nodes[i] = node_copy;
1793
+ }
1794
+ graph_copy->n_nodes = graph->n_nodes;
1795
+
1796
+ lm_ggml_hash_set_free(&hash_set);
1797
+ free(node_copies);
1798
+ free(node_init);
1799
+
1800
+ return {
1801
+ /* .buffer = */ buffer,
1802
+ /* .ctx_allocated = */ ctx_allocated,
1803
+ /* .ctx_unallocated = */ ctx_unallocated,
1804
+ /* .graph = */ graph_copy,
1805
+ };
1806
+ }
1807
+
1808
+ void lm_ggml_backend_graph_copy_free(struct lm_ggml_backend_graph_copy copy) {
1809
+ lm_ggml_backend_buffer_free(copy.buffer);
1810
+ lm_ggml_free(copy.ctx_allocated);
1811
+ lm_ggml_free(copy.ctx_unallocated);
1812
+ }
1813
+
1814
+ bool lm_ggml_backend_compare_graph_backend(lm_ggml_backend_t backend1, lm_ggml_backend_t backend2, struct lm_ggml_cgraph * graph, lm_ggml_backend_eval_callback callback, void * user_data) {
1815
+ struct lm_ggml_backend_graph_copy copy = lm_ggml_backend_graph_copy(backend2, graph);
1816
+ if (copy.buffer == NULL) {
1817
+ return false;
1818
+ }
1819
+
1820
+ struct lm_ggml_cgraph * g1 = graph;
1821
+ struct lm_ggml_cgraph * g2 = copy.graph;
1822
+
1823
+ assert(g1->n_nodes == g2->n_nodes);
1824
+
1825
+ for (int i = 0; i < g1->n_nodes; i++) {
1826
+ //printf("eval %d/%d\n", i, g1->n_nodes);
1827
+ struct lm_ggml_tensor * t1 = g1->nodes[i];
1828
+ struct lm_ggml_tensor * t2 = g2->nodes[i];
1829
+
1830
+ assert(t1->op == t2->op && lm_ggml_are_same_layout(t1, t2));
1831
+
1832
+ struct lm_ggml_cgraph g1v = lm_ggml_graph_view(g1, i, i + 1);
1833
+ struct lm_ggml_cgraph g2v = lm_ggml_graph_view(g2, i, i + 1);
1834
+
1835
+ lm_ggml_backend_graph_compute(backend1, &g1v);
1836
+ lm_ggml_backend_graph_compute(backend2, &g2v);
1837
+
1838
+ if (lm_ggml_is_view_op(t1->op)) {
1839
+ continue;
1840
+ }
1841
+
1842
+ // compare results, calculate rms etc
1843
+ if (!callback(i, t1, t2, user_data)) {
1844
+ break;
1845
+ }
1846
+ }
1847
+
1848
+ lm_ggml_backend_graph_copy_free(copy);
1849
+
1850
+ return true;
1851
+ }
1852
+
1853
+ // CPU backend - buffer
1854
+
1855
+ static void * lm_ggml_backend_cpu_buffer_get_base(lm_ggml_backend_buffer_t buffer) {
1856
+ uintptr_t data = (uintptr_t)buffer->context;
1857
+
1858
+ // align the buffer
1859
+ if (data % TENSOR_ALIGNMENT != 0) {
1860
+ data = LM_GGML_PAD(data, TENSOR_ALIGNMENT);
1861
+ }
1862
+
1863
+ return (void *)data;
1864
+ }
1865
+
1866
+ static void lm_ggml_backend_cpu_buffer_free_buffer(lm_ggml_backend_buffer_t buffer) {
1867
+ lm_ggml_aligned_free(buffer->context, buffer->size);
1868
+ }
1869
+
1870
+ static void lm_ggml_backend_cpu_buffer_memset_tensor(lm_ggml_backend_buffer_t buffer, struct lm_ggml_tensor * tensor, uint8_t value, size_t offset, size_t size) {
1871
+ memset((char *)tensor->data + offset, value, size);
1872
+
1873
+ LM_GGML_UNUSED(buffer);
1874
+ }
1875
+
1876
+ static void lm_ggml_backend_cpu_buffer_set_tensor(lm_ggml_backend_buffer_t buffer, struct lm_ggml_tensor * tensor, const void * data, size_t offset, size_t size) {
1877
+ memcpy((char *)tensor->data + offset, data, size);
1878
+
1879
+ LM_GGML_UNUSED(buffer);
1880
+ }
1881
+
1882
+ static void lm_ggml_backend_cpu_buffer_get_tensor(lm_ggml_backend_buffer_t buffer, const struct lm_ggml_tensor * tensor, void * data, size_t offset, size_t size) {
1883
+ memcpy(data, (const char *)tensor->data + offset, size);
1884
+
1885
+ LM_GGML_UNUSED(buffer);
1886
+ }
1887
+
1888
+ static bool lm_ggml_backend_cpu_buffer_cpy_tensor(lm_ggml_backend_buffer_t buffer, const struct lm_ggml_tensor * src, struct lm_ggml_tensor * dst) {
1889
+ if (lm_ggml_backend_buffer_is_host(src->buffer)) {
1890
+ memcpy(dst->data, src->data, lm_ggml_nbytes(src));
1891
+ return true;
1892
+ }
1893
+ return false;
1894
+
1895
+ LM_GGML_UNUSED(buffer);
1896
+ }
1897
+
1898
+ static void lm_ggml_backend_cpu_buffer_clear(lm_ggml_backend_buffer_t buffer, uint8_t value) {
1899
+ memset(buffer->context, value, buffer->size);
1900
+ }
1901
+
1902
+ static const struct lm_ggml_backend_buffer_i lm_ggml_backend_cpu_buffer_i = {
1903
+ /* .free_buffer = */ lm_ggml_backend_cpu_buffer_free_buffer,
1904
+ /* .get_base = */ lm_ggml_backend_cpu_buffer_get_base,
1905
+ /* .init_tensor = */ NULL, // no initialization required
1906
+ /* .memset_tensor = */ lm_ggml_backend_cpu_buffer_memset_tensor,
1907
+ /* .set_tensor = */ lm_ggml_backend_cpu_buffer_set_tensor,
1908
+ /* .get_tensor = */ lm_ggml_backend_cpu_buffer_get_tensor,
1909
+ /* .cpy_tensor = */ lm_ggml_backend_cpu_buffer_cpy_tensor,
1910
+ /* .clear = */ lm_ggml_backend_cpu_buffer_clear,
1911
+ /* .reset = */ NULL,
1912
+ };
1913
+
1914
+ static const struct lm_ggml_backend_buffer_i lm_ggml_backend_cpu_buffer_from_ptr_i = {
1915
+ /* .free_buffer = */ NULL, // ptr is not owned by the buffer, so it does not need to be freed
1916
+ /* .get_base = */ lm_ggml_backend_cpu_buffer_get_base,
1917
+ /* .init_tensor = */ NULL, // no initialization required
1918
+ /* .memset_tensor = */ lm_ggml_backend_cpu_buffer_memset_tensor,
1919
+ /* .set_tensor = */ lm_ggml_backend_cpu_buffer_set_tensor,
1920
+ /* .get_tensor = */ lm_ggml_backend_cpu_buffer_get_tensor,
1921
+ /* .cpy_tensor = */ lm_ggml_backend_cpu_buffer_cpy_tensor,
1922
+ /* .clear = */ lm_ggml_backend_cpu_buffer_clear,
1923
+ /* .reset = */ NULL,
1924
+ };
1925
+
1926
+ // CPU backend buffer type
1927
+
1928
+ // this buffer type is defined here to make it available to all backends
1929
+
1930
+ static const char * lm_ggml_backend_cpu_buffer_type_get_name(lm_ggml_backend_buffer_type_t buft) {
1931
+ return "CPU";
1932
+
1933
+ LM_GGML_UNUSED(buft);
1934
+ }
1935
+
1936
+ static lm_ggml_backend_buffer_t lm_ggml_backend_cpu_buffer_type_alloc_buffer(lm_ggml_backend_buffer_type_t buft, size_t size) {
1937
+ void * data = lm_ggml_aligned_malloc(size);
1938
+
1939
+ if (data == NULL) {
1940
+ LM_GGML_LOG_ERROR("%s: failed to allocate buffer of size %zu\n", __func__, size);
1941
+ return NULL;
1942
+ }
1943
+
1944
+ return lm_ggml_backend_buffer_init(buft, lm_ggml_backend_cpu_buffer_i, data, size);
1945
+ }
1946
+
1947
+ static size_t lm_ggml_backend_cpu_buffer_type_get_alignment(lm_ggml_backend_buffer_type_t buft) {
1948
+ return TENSOR_ALIGNMENT;
1949
+
1950
+ LM_GGML_UNUSED(buft);
1951
+ }
1952
+
1953
+ static bool lm_ggml_backend_cpu_buffer_type_is_host(lm_ggml_backend_buffer_type_t buft) {
1954
+ return true;
1955
+
1956
+ LM_GGML_UNUSED(buft);
1957
+ }
1958
+
1959
+ lm_ggml_backend_buffer_type_t lm_ggml_backend_cpu_buffer_type(void) {
1960
+ static struct lm_ggml_backend_buffer_type lm_ggml_backend_cpu_buffer_type = {
1961
+ /* .iface = */ {
1962
+ /* .get_name = */ lm_ggml_backend_cpu_buffer_type_get_name,
1963
+ /* .alloc_buffer = */ lm_ggml_backend_cpu_buffer_type_alloc_buffer,
1964
+ /* .get_alignment = */ lm_ggml_backend_cpu_buffer_type_get_alignment,
1965
+ /* .get_max_size = */ NULL, // defaults to SIZE_MAX
1966
+ /* .get_alloc_size = */ NULL, // defaults to lm_ggml_nbytes
1967
+ /* .is_host = */ lm_ggml_backend_cpu_buffer_type_is_host,
1968
+ },
1969
+ /* .device = */ NULL, // FIXME lm_ggml_backend_reg_dev_get(lm_ggml_backend_cpu_reg(), 0),
1970
+ /* .context = */ NULL,
1971
+ };
1972
+
1973
+ return &lm_ggml_backend_cpu_buffer_type;
1974
+ }
1975
+
1976
+ static const char * lm_ggml_backend_cpu_buffer_from_ptr_type_get_name(lm_ggml_backend_buffer_type_t buft) {
1977
+ return "CPU_Mapped";
1978
+
1979
+ LM_GGML_UNUSED(buft);
1980
+ }
1981
+
1982
+ static lm_ggml_backend_buffer_type_t lm_ggml_backend_cpu_buffer_from_ptr_type(void) {
1983
+ static struct lm_ggml_backend_buffer_type lm_ggml_backend_cpu_buffer_type = {
1984
+ /* .iface = */ {
1985
+ /* .get_name = */ lm_ggml_backend_cpu_buffer_from_ptr_type_get_name,
1986
+ /* .alloc_buffer = */ lm_ggml_backend_cpu_buffer_type_alloc_buffer,
1987
+ /* .get_alignment = */ lm_ggml_backend_cpu_buffer_type_get_alignment,
1988
+ /* .get_max_size = */ NULL, // defaults to SIZE_MAX
1989
+ /* .get_alloc_size = */ NULL, // defaults to lm_ggml_nbytes
1990
+ /* .is_host = */ lm_ggml_backend_cpu_buffer_type_is_host,
1991
+ },
1992
+ /* .device = */ NULL, // FIXME lm_ggml_backend_reg_dev_get(lm_ggml_backend_cpu_reg(), 0),
1993
+ /* .context = */ NULL,
1994
+ };
1995
+
1996
+ return &lm_ggml_backend_cpu_buffer_type;
1997
+ }
1998
+
1999
+ lm_ggml_backend_buffer_t lm_ggml_backend_cpu_buffer_from_ptr(void * ptr, size_t size) {
2000
+ LM_GGML_ASSERT((uintptr_t)ptr % TENSOR_ALIGNMENT == 0 && "buffer pointer must be aligned");
2001
+ return lm_ggml_backend_buffer_init(lm_ggml_backend_cpu_buffer_from_ptr_type(), lm_ggml_backend_cpu_buffer_from_ptr_i, ptr, size);
2002
+ }