char-fork/llama.cpp

git clone https://git.t4t.associates/char-fork/llama.cpp

HarrisonSecgguf : harden loader against malformed tensor dims and metadata types (#25596)d8a8beac2

char/aliceai-t5-moe
50.7 KiB1479 linesraw
1#include "ggml.h"
2#include "ggml-backend.h"
3#include "../ggml/src/ggml-impl.h"
4#include "gguf.h"
5
6#include <algorithm>
7#include <array>
8#include <cmath>
9#include <cstdint>
10#include <cstdio>
11#include <random>
12#include <string>
13#include <vector>
14
15constexpr int offset_has_kv      = 1000;
16constexpr int offset_has_tensors = 2000;
17constexpr int offset_has_data    = 3000;
18
19enum handcrafted_file_type {
20    HANDCRAFTED_HEADER_BAD_MAGIC           =  10,
21    HANDCRAFTED_HEADER_BAD_VERSION_0       =  15,
22    HANDCRAFTED_HEADER_BAD_VERSION_1       =  20,
23    HANDCRAFTED_HEADER_BAD_VERSION_FUTURE  =  30,
24    HANDCRAFTED_HEADER_BAD_N_TENSORS       =  40,
25    HANDCRAFTED_HEADER_BAD_N_KV            =  50,
26    HANDCRAFTED_HEADER_EMPTY               = 800,
27
28    HANDCRAFTED_KV_BAD_KEY_SIZE            =  10 + offset_has_kv,
29    HANDCRAFTED_KV_EMPTY_KEY               =  15 + offset_has_kv,
30    HANDCRAFTED_KV_BAD_TYPE                =  20 + offset_has_kv,
31    // HANDCRAFTED_KV_BAD_VALUE_SIZE          =  30 + offset_has_kv, // removed because it can result in allocations > 1 TB (default sanitizer limit)
32    HANDCRAFTED_KV_DUPLICATE_KEY           =  40 + offset_has_kv,
33    HANDCRAFTED_KV_BAD_ALIGN               =  50 + offset_has_kv,
34    HANDCRAFTED_KV_WRONG_TYPE_ALIGN        =  55 + offset_has_kv,
35    HANDCRAFTED_KV_SUCCESS                 = 800 + offset_has_kv,
36
37    HANDCRAFTED_TENSORS_BAD_NAME_SIZE      =  10 + offset_has_tensors,
38    HANDCRAFTED_TENSORS_BAD_N_DIMS         =  20 + offset_has_tensors,
39    HANDCRAFTED_TENSORS_BAD_SHAPE          =  30 + offset_has_tensors,
40    HANDCRAFTED_TENSORS_ZERO_DIM           =  35 + offset_has_tensors,
41    HANDCRAFTED_TENSORS_NE_TOO_BIG         =  40 + offset_has_tensors,
42    HANDCRAFTED_TENSORS_NBYTES_TOO_BIG     =  45 + offset_has_tensors,
43    HANDCRAFTED_TENSORS_BAD_TYPE           =  50 + offset_has_tensors,
44    HANDCRAFTED_TENSORS_BAD_OFFSET         =  60 + offset_has_tensors,
45    HANDCRAFTED_TENSORS_DUPLICATE_NAME     =  70 + offset_has_tensors,
46    HANDCRAFTED_TENSORS_BAD_ALIGN          =  75 + offset_has_tensors,
47    HANDCRAFTED_TENSORS_INCONSISTENT_ALIGN =  80 + offset_has_tensors,
48    HANDCRAFTED_TENSORS_SUCCESS            = 800 + offset_has_tensors,
49    HANDCRAFTED_TENSORS_CUSTOM_ALIGN       = 810 + offset_has_tensors,
50
51    HANDCRAFTED_DATA_NOT_ENOUGH_DATA       =  10 + offset_has_data,
52    HANDCRAFTED_DATA_BAD_ALIGN             =  15 + offset_has_data,
53    HANDCRAFTED_DATA_INCONSISTENT_ALIGN    =  20 + offset_has_data,
54    HANDCRAFTED_DATA_MEM_SIZE_OVERFLOW     =  30 + offset_has_data,
55    HANDCRAFTED_DATA_SUCCESS               = 800 + offset_has_data,
56    HANDCRAFTED_DATA_CUSTOM_ALIGN          = 810 + offset_has_data,
57};
58
59static std::string handcrafted_file_type_name(const enum handcrafted_file_type hft) {
60    switch (hft) {
61        case HANDCRAFTED_HEADER_BAD_MAGIC:           return "HEADER_BAD_MAGIC";
62        case HANDCRAFTED_HEADER_BAD_VERSION_0:       return "HEADER_BAD_VERSION_0";
63        case HANDCRAFTED_HEADER_BAD_VERSION_1:       return "HEADER_BAD_VERSION_1";
64        case HANDCRAFTED_HEADER_BAD_VERSION_FUTURE:  return "HEADER_BAD_VERSION_FUTURE";
65        case HANDCRAFTED_HEADER_BAD_N_KV:            return "HEADER_BAD_N_KV";
66        case HANDCRAFTED_HEADER_BAD_N_TENSORS:       return "HEADER_BAD_N_TENSORS";
67        case HANDCRAFTED_HEADER_EMPTY:               return "HEADER_EMPTY";
68
69        case HANDCRAFTED_KV_BAD_KEY_SIZE:            return "KV_BAD_KEY_SIZE";
70        case HANDCRAFTED_KV_EMPTY_KEY:               return "KV_EMPTY_KEY";
71        case HANDCRAFTED_KV_BAD_TYPE:                return "KV_BAD_TYPE";
72        case HANDCRAFTED_KV_DUPLICATE_KEY:           return "KV_DUPLICATE_KEY";
73        case HANDCRAFTED_KV_BAD_ALIGN:               return "KV_BAD_ALIGN";
74        case HANDCRAFTED_KV_WRONG_TYPE_ALIGN:        return "KV_WRONG_TYPE_ALIGN";
75        case HANDCRAFTED_KV_SUCCESS:                 return "KV_RANDOM_KV";
76
77        case HANDCRAFTED_TENSORS_BAD_NAME_SIZE:      return "TENSORS_BAD_NAME_SIZE";
78        case HANDCRAFTED_TENSORS_BAD_N_DIMS:         return "TENSORS_BAD_N_DIMS";
79        case HANDCRAFTED_TENSORS_BAD_SHAPE:          return "TENSORS_BAD_SHAPE";
80        case HANDCRAFTED_TENSORS_ZERO_DIM:           return "TENSORS_ZERO_DIM";
81        case HANDCRAFTED_TENSORS_NE_TOO_BIG:         return "TENSORS_NE_TOO_BIG";
82        case HANDCRAFTED_TENSORS_NBYTES_TOO_BIG:     return "TENSORS_NBYTES_TOO_BIG";
83        case HANDCRAFTED_TENSORS_BAD_TYPE:           return "TENSORS_BAD_TYPE";
84        case HANDCRAFTED_TENSORS_BAD_OFFSET:         return "TENSORS_BAD_OFFSET";
85        case HANDCRAFTED_TENSORS_DUPLICATE_NAME:     return "TENSORS_DUPLICATE_NAME";
86        case HANDCRAFTED_TENSORS_BAD_ALIGN:          return "TENSORS_BAD_ALIGN";
87        case HANDCRAFTED_TENSORS_INCONSISTENT_ALIGN: return "TENSORS_INCONSISTENT_ALIGN";
88        case HANDCRAFTED_TENSORS_SUCCESS:            return "TENSORS_SUCCESS";
89        case HANDCRAFTED_TENSORS_CUSTOM_ALIGN:       return "TENSORS_CUSTOM_ALIGN";
90
91        case HANDCRAFTED_DATA_NOT_ENOUGH_DATA:       return "DATA_NOT_ENOUGH_DATA";
92        case HANDCRAFTED_DATA_BAD_ALIGN:             return "DATA_BAD_ALIGN";
93        case HANDCRAFTED_DATA_INCONSISTENT_ALIGN:    return "DATA_INCONSISTENT_ALIGN";
94        case HANDCRAFTED_DATA_MEM_SIZE_OVERFLOW:     return "DATA_MEM_SIZE_OVERFLOW";
95        case HANDCRAFTED_DATA_SUCCESS:               return "DATA_SUCCESS";
96        case HANDCRAFTED_DATA_CUSTOM_ALIGN:          return "DATA_CUSTOM_ALIGN";
97    }
98    GGML_ABORT("fatal error");
99}
100
101static bool expect_context_not_null(const enum handcrafted_file_type hft) {
102    if (hft == HANDCRAFTED_TENSORS_ZERO_DIM) {
103        return true;
104    }
105    if (hft < offset_has_kv) {
106        return hft >= HANDCRAFTED_HEADER_EMPTY;
107    }
108    if (hft < offset_has_tensors) {
109        return hft >= HANDCRAFTED_KV_SUCCESS;
110    }
111    if (hft < offset_has_data) {
112        return hft >= HANDCRAFTED_TENSORS_SUCCESS;
113    }
114    return hft >= HANDCRAFTED_DATA_SUCCESS;
115}
116
117typedef std::pair<enum ggml_type, std::array<int64_t, GGML_MAX_DIMS>> tensor_config_t;
118
119static std::vector<tensor_config_t> get_tensor_configs(std::mt19937 & rng) {
120    std::vector<tensor_config_t> tensor_configs;
121    tensor_configs.reserve(100);
122
123    for (int i = 0; i < 100; ++i) {
124        const enum ggml_type type = ggml_type(rng() % GGML_TYPE_COUNT);
125        if (ggml_type_size(type) == 0) {
126            continue;
127        }
128
129        std::array<int64_t, GGML_MAX_DIMS> shape = {1, 1, 1, 1};
130        shape[0] = (1 + rng() % 10) * ggml_blck_size(type);
131        const int n_dims = 1 + rng() % GGML_MAX_DIMS;
132        for (int i = 1; i < n_dims; ++i) {
133            shape[i] = 1 + rng() % 10;
134        }
135
136        tensor_configs.push_back(std::make_pair(type, shape));
137    }
138
139    return tensor_configs;
140}
141
142static std::vector<std::pair<enum gguf_type, enum gguf_type>> get_kv_types(std::mt19937 rng) {
143    std::vector<std::pair<enum gguf_type, enum gguf_type>> kv_types;
144    kv_types.reserve(100);
145
146    for (int i = 0; i < 100; ++i) {
147        const gguf_type type = gguf_type(rng() % GGUF_TYPE_COUNT);
148
149        if (type == GGUF_TYPE_ARRAY) {
150            const gguf_type type_arr = gguf_type(rng() % GGUF_TYPE_COUNT);
151            if (type_arr == GGUF_TYPE_ARRAY) {
152                continue;
153            }
154            kv_types.push_back(std::make_pair(type, type_arr));
155            continue;
156        }
157
158        kv_types.push_back(std::make_pair(type, gguf_type(-1)));
159    }
160    std::shuffle(kv_types.begin(), kv_types.end(), rng);
161
162    return kv_types;
163}
164
165template <typename T>
166static void helper_write(FILE * file, const T & val) {
167    GGML_ASSERT(fwrite(&val, 1, sizeof(val), file) == sizeof(val));
168}
169
170static void helper_write(FILE * file, const void * data, const size_t nbytes) {
171    GGML_ASSERT(fwrite(data, 1, nbytes, file) == nbytes);
172}
173
174static std::vector<uint8_t> read_file_to_buffer(FILE * file) {
175    GGML_ASSERT(file != nullptr);
176    GGML_ASSERT(fseek(file, 0, SEEK_END) == 0);
177
178    const long size = ftell(file);
179    GGML_ASSERT(size >= 0);
180
181    rewind(file);
182
183    std::vector<uint8_t> data(static_cast<size_t>(size));
184    GGML_ASSERT(fread(data.data(), 1, data.size(), file) == data.size());
185
186    rewind(file);
187    return data;
188}
189
190struct callback_reader_data {
191    const uint8_t * data;
192    size_t size;
193};
194
195static size_t read_buffer_callback(void * userdata, void * output, uint64_t offset, size_t len) {
196    GGML_ASSERT(len > 0);
197
198    const callback_reader_data & reader = *static_cast<callback_reader_data *>(userdata);
199
200    if (offset > reader.size || len > reader.size - offset) {
201        return 0;
202    }
203
204    const size_t data_offset = static_cast<size_t>(offset);
205    const size_t nread = std::min(len, reader.size - data_offset);
206    memcpy(static_cast<uint8_t *>(output), reader.data + data_offset, nread);
207    return nread;
208}
209
210static FILE * get_handcrafted_file(const unsigned int seed, const enum handcrafted_file_type hft, const int extra_bytes = 0) {
211    FILE * file = tmpfile();
212
213    if (!file) {
214        return file;
215    }
216
217    std::mt19937 rng(seed);
218    uint32_t alignment = GGUF_DEFAULT_ALIGNMENT;
219
220    if (hft == HANDCRAFTED_HEADER_BAD_MAGIC) {
221        const char bad_magic[4] = {'F', 'U', 'G', 'G'};
222        helper_write(file, bad_magic, sizeof(bad_magic));
223    } else {
224        helper_write(file, GGUF_MAGIC, 4);
225    }
226
227    if (hft == HANDCRAFTED_HEADER_BAD_VERSION_0) {
228        const uint32_t version = 0;
229        helper_write(file, version);
230    } else if (hft == HANDCRAFTED_HEADER_BAD_VERSION_1) {
231        const uint32_t version = 1;
232        helper_write(file, version);
233    } else if (hft == HANDCRAFTED_HEADER_BAD_VERSION_FUTURE) {
234        const uint32_t version = GGUF_VERSION + 1;
235        helper_write(file, version);
236    } else {
237        const uint32_t version = GGUF_VERSION;
238        helper_write(file, version);
239    }
240
241    std::vector<tensor_config_t> tensor_configs;
242    if (hft >= offset_has_tensors) {
243        tensor_configs = get_tensor_configs(rng);
244    }
245
246    if (hft == HANDCRAFTED_DATA_MEM_SIZE_OVERFLOW) {
247        tensor_configs.resize(2);
248
249        tensor_configs[0] = { GGML_TYPE_I8, { 0x7FFFFFFFFFFFFFC0, 1, 1, 1 } };
250        tensor_configs[1] = { GGML_TYPE_I8, { 0x7FFFFFFFFFFFFFC0, 1, 1, 1 } };
251    }
252
253    if (hft == HANDCRAFTED_HEADER_BAD_N_TENSORS) {
254        const uint64_t n_tensors = -1;
255        helper_write(file, n_tensors);
256    } else {
257        const uint64_t n_tensors = tensor_configs.size();
258        helper_write(file, n_tensors);
259    }
260
261    std::vector<std::pair<enum gguf_type, enum gguf_type>> kv_types;
262    if (hft >= offset_has_kv) {
263        kv_types = get_kv_types(rng);
264    }
265    {
266        uint64_t n_kv = kv_types.size();
267        if (hft == HANDCRAFTED_KV_BAD_ALIGN        || hft == HANDCRAFTED_KV_WRONG_TYPE_ALIGN ||
268            hft == HANDCRAFTED_TENSORS_BAD_ALIGN   || hft == HANDCRAFTED_TENSORS_CUSTOM_ALIGN ||
269            hft == HANDCRAFTED_DATA_BAD_ALIGN      || hft == HANDCRAFTED_DATA_CUSTOM_ALIGN) {
270
271            n_kv += 1;
272        } else if (hft == HANDCRAFTED_HEADER_BAD_N_KV) {
273            n_kv = -1;
274        }
275        helper_write(file, n_kv);
276    }
277
278    if (hft < offset_has_kv) {
279        while (ftell(file) % alignment != 0) {
280            const char pad = 0;
281            helper_write(file, pad);
282        }
283
284        for (int i = 0; i < extra_bytes; ++i) {
285            const char tmp = 0;
286            helper_write(file, tmp);
287        }
288        rewind(file);
289        return file;
290    }
291
292    for (int i = 0; i < int(kv_types.size()); ++i) {
293        const enum gguf_type type     = gguf_type(hft == HANDCRAFTED_KV_BAD_TYPE ? GGUF_TYPE_COUNT : kv_types[i].first);
294        const enum gguf_type type_arr = gguf_type(hft == HANDCRAFTED_KV_BAD_TYPE ? GGUF_TYPE_COUNT : kv_types[i].second);
295
296        const std::string key = hft == HANDCRAFTED_KV_EMPTY_KEY
297            ? ""
298            : "my_key_" + std::to_string((hft == HANDCRAFTED_KV_DUPLICATE_KEY ? i/2 : i));
299
300        if (hft == HANDCRAFTED_KV_BAD_KEY_SIZE) {
301            const uint64_t n = -1;
302            helper_write(file, n);
303        } else {
304            const uint64_t n = key.length();
305            helper_write(file, n);
306        }
307        helper_write(file, key.data(), key.length());
308
309        {
310            const int32_t type32 = int32_t(type);
311            helper_write(file, type32);
312        }
313
314        uint32_t data[16];
315        for (int j = 0; j < 16; ++j) {
316            data[j] = rng();
317            if (type == GGUF_TYPE_STRING || type_arr == GGUF_TYPE_STRING) {
318                data[j] |= 0x01010101; // avoid random null-termination of string
319            }
320        }
321
322        if (type == GGUF_TYPE_STRING) {
323            const uint64_t n = rng() % sizeof(data);
324            helper_write(file, n);
325            helper_write(file, data, n);
326            continue;
327        }
328
329        if (type == GGUF_TYPE_ARRAY) {
330            {
331                const int32_t type32 = int32_t(type_arr);
332                helper_write(file, type32);
333            }
334            if (type_arr == GGUF_TYPE_STRING) {
335                const uint64_t nstr = rng() % (16 + 1);
336                helper_write(file, nstr);
337                for (uint64_t istr = 0; istr < nstr; ++istr) {
338                    const uint64_t n = rng() % (sizeof(uint32_t) + 1);
339                    helper_write(file, n);
340                    helper_write(file, &data[istr], n);
341                }
342                continue;
343            }
344            const size_t type_size = gguf_type_size(type_arr);
345            const uint64_t n = (rng() % sizeof(data)) / type_size;
346            helper_write(file, n);
347            helper_write(file, &data, n*type_size);
348            continue;
349        }
350
351        helper_write(file, data, hft == HANDCRAFTED_KV_BAD_TYPE ? 1 : gguf_type_size(type));
352    }
353
354    if (hft == HANDCRAFTED_KV_BAD_ALIGN        || hft == HANDCRAFTED_KV_WRONG_TYPE_ALIGN ||
355        hft == HANDCRAFTED_TENSORS_BAD_ALIGN   || hft == HANDCRAFTED_TENSORS_CUSTOM_ALIGN ||
356        hft == HANDCRAFTED_DATA_BAD_ALIGN      || hft == HANDCRAFTED_DATA_CUSTOM_ALIGN) {
357
358        const uint64_t n = strlen(GGUF_KEY_GENERAL_ALIGNMENT);
359        helper_write(file, n);
360        helper_write(file, GGUF_KEY_GENERAL_ALIGNMENT, n);
361
362        // HANDCRAFTED_KV_WRONG_TYPE_ALIGN declares general.alignment with a non-UINT32 type,
363        // which the loader must reject cleanly instead of aborting on an assertion
364        const int32_t type = hft == HANDCRAFTED_KV_WRONG_TYPE_ALIGN ? int32_t(GGUF_TYPE_INT32) : int32_t(GGUF_TYPE_UINT32);
365        helper_write(file, type);
366
367        alignment = expect_context_not_null(hft) ? 1 : 13;
368        helper_write(file, alignment);
369    }
370
371    if (hft < offset_has_tensors) {
372        while (ftell(file) % alignment != 0) {
373            const char pad = 0;
374            helper_write(file, pad);
375        }
376
377        for (int i = 0; i < extra_bytes; ++i) {
378            const char tmp = 0;
379            helper_write(file, tmp);
380        }
381        rewind(file);
382        return file;
383    }
384
385    if (hft == HANDCRAFTED_TENSORS_INCONSISTENT_ALIGN || hft == HANDCRAFTED_DATA_INCONSISTENT_ALIGN) {
386        alignment = 1;
387    }
388
389    uint64_t offset = 0;
390    for (int i = 0; i < int(tensor_configs.size()); ++i) {
391        const ggml_type                          type  = hft == HANDCRAFTED_TENSORS_NBYTES_TOO_BIG ? GGML_TYPE_I64 : tensor_configs[i].first;
392        const std::array<int64_t, GGML_MAX_DIMS> shape = tensor_configs[i].second;
393
394        std::string name = "my_tensor";
395        if (hft != HANDCRAFTED_TENSORS_DUPLICATE_NAME) {
396            name += "_" + std::to_string(i);
397        }
398        if (hft == HANDCRAFTED_TENSORS_BAD_NAME_SIZE) {
399            name += "_with_a_very_long_name_which_is_longer_than_what_is_allowed_for_ggml_tensors";
400            GGML_ASSERT(name.length() >= GGML_MAX_NAME);
401        }
402        {
403            const uint64_t n = name.length();
404            helper_write(file, n);
405        }
406        helper_write(file, name.data(), name.length());
407
408        uint32_t n_dims = (hft == HANDCRAFTED_TENSORS_NE_TOO_BIG || hft == HANDCRAFTED_TENSORS_NBYTES_TOO_BIG) ? 2 : 1;
409        for (int i = GGML_MAX_DIMS-1; i >= 1; --i) {
410            if (shape[i] != 1) {
411                n_dims = i + 1;
412                break;
413            }
414        }
415        if (hft == HANDCRAFTED_TENSORS_ZERO_DIM) {
416            n_dims = 2;
417        }
418        if (hft == HANDCRAFTED_TENSORS_BAD_N_DIMS) {
419            const uint32_t n_dims_bad = GGML_MAX_DIMS + 1;
420            helper_write(file, n_dims_bad);
421        } else {
422            helper_write(file, n_dims);
423        }
424
425        if (hft == HANDCRAFTED_TENSORS_BAD_SHAPE) {
426            const int64_t bad_dim = -1;
427            for (uint32_t j = 0; j < n_dims; ++j) {
428                helper_write(file, bad_dim);
429            }
430        } else if (hft == HANDCRAFTED_TENSORS_ZERO_DIM) {
431            const int64_t zero_shape[2] = { shape[0], 0 };
432            helper_write(file, zero_shape, 2*sizeof(int64_t));
433        } else if (hft == HANDCRAFTED_TENSORS_NE_TOO_BIG){
434            const int64_t big_dim = 4*int64_t(INT32_MAX);
435            for (uint32_t j = 0; j < n_dims; ++j) {
436                helper_write(file, big_dim);
437            }
438        } else if (hft == HANDCRAFTED_TENSORS_NBYTES_TOO_BIG){
439            const size_t  big_ne  = SIZE_MAX/ggml_type_size(type);
440            const int64_t big_dim = GGML_PAD(int64_t(1.01f*std::pow(big_ne, 1.0f/n_dims)) + 1, ggml_blck_size(type));
441            for (uint32_t j = 0; j < n_dims; ++j) {
442                helper_write(file, big_dim);
443            }
444        } else {
445            helper_write(file, shape.data(), n_dims*sizeof(int64_t));
446        }
447
448        {
449            const int32_t type32 = hft == HANDCRAFTED_TENSORS_BAD_TYPE ? GGML_TYPE_COUNT : int32_t(type);
450            helper_write(file, type32);
451        }
452
453        if (hft == HANDCRAFTED_TENSORS_BAD_OFFSET) {
454            const uint64_t bad_offset = -1;
455            helper_write(file, bad_offset);
456        } else {
457            helper_write(file, offset);
458        }
459
460        int64_t ne = shape[0];
461        for (uint32_t i = 1; i < n_dims; ++i) {
462            ne *= shape[i];
463        }
464        if (hft == HANDCRAFTED_TENSORS_ZERO_DIM) {
465            ne = 0;
466        }
467
468        offset += GGML_PAD(ggml_row_size(type, ne), (uint64_t) alignment);
469    }
470
471    while (ftell(file) % alignment != 0) {
472        const char pad = 0;
473        helper_write(file, pad);
474    }
475
476    if (hft >= offset_has_data) {
477        rng.seed(seed + 1);
478        uint64_t nbytes = offset;
479        if (hft == HANDCRAFTED_DATA_NOT_ENOUGH_DATA) {
480            nbytes -= 1;
481        }
482        if (hft == HANDCRAFTED_DATA_MEM_SIZE_OVERFLOW) {
483            nbytes = 32;
484        }
485        for (uint64_t i = 0; i < nbytes; ++i) {
486            const uint8_t random_byte = i % 256;
487            helper_write(file, random_byte);
488        }
489    }
490
491    for (int i = 0; i < extra_bytes; ++i) {
492        const char tmp = 0;
493        helper_write(file, tmp);
494    }
495    rewind(file);
496    return file;
497}
498
499static bool handcrafted_check_header(const gguf_context * gguf_ctx, const unsigned int seed, const bool has_kv, const bool has_tensors, const bool alignment_defined) {
500    if (!gguf_ctx) {
501        return false;
502    }
503
504    std::mt19937 rng(seed);
505
506    std::vector<tensor_config_t> tensor_configs;
507    if (has_tensors) {
508        tensor_configs = get_tensor_configs(rng);
509    }
510    std::vector<std::pair<enum gguf_type, enum gguf_type>> kv_types;
511    if (has_kv) {
512        kv_types = get_kv_types(rng);
513    }
514
515    bool ok = true;
516
517    if (gguf_get_version(gguf_ctx) != GGUF_VERSION) {
518        ok = false;
519    }
520    if (gguf_get_n_tensors(gguf_ctx) != int(tensor_configs.size())) {
521        ok = false;
522    }
523    if (gguf_get_n_kv(gguf_ctx) != int(alignment_defined ? kv_types.size() + 1 : kv_types.size())) {
524        ok = false;
525    }
526
527    return ok;
528}
529
530static bool handcrafted_check_kv(const gguf_context * gguf_ctx, const unsigned int seed, const bool has_tensors, const bool alignment_defined) {
531    if (!gguf_ctx) {
532        return false;
533    }
534
535    std::mt19937 rng(seed);
536
537    std::vector<tensor_config_t> tensor_configs;
538    if (has_tensors) {
539        tensor_configs = get_tensor_configs(rng);
540    }
541
542    std::vector<std::pair<enum gguf_type, enum gguf_type>> kv_types = get_kv_types(rng);
543
544    bool ok = true;
545
546    for (int i = 0; i < int(kv_types.size()); ++i) {
547        const enum gguf_type type     = gguf_type(kv_types[i].first);
548        const enum gguf_type type_arr = gguf_type(kv_types[i].second);
549
550        const std::string key = "my_key_" + std::to_string(i);
551
552        uint32_t data[16];
553        for (int j = 0; j < 16; ++j) {
554            data[j] = rng();
555            if (type == GGUF_TYPE_STRING || type_arr == GGUF_TYPE_STRING) {
556                data[j] |= 0x01010101; // avoid random null-termination of string
557            }
558        }
559
560        const char * data8 = reinterpret_cast<const char *>(data);
561        const int id = gguf_find_key(gguf_ctx, key.c_str());
562
563        if (type == GGUF_TYPE_STRING) {
564            const char * str = gguf_get_val_str(gguf_ctx, id);
565            const uint64_t n = strlen(str);
566            const uint64_t n_expected = rng() % sizeof(data);
567            if (n != n_expected) {
568                ok = false;
569                continue;
570            }
571            if (!std::equal(str, str + n, data8)) {
572                ok = false;
573            }
574            continue;
575        }
576
577        if (type == GGUF_TYPE_ARRAY) {
578            const size_t type_size = gguf_type_size(type_arr);
579            const uint64_t arr_n = gguf_get_arr_n(gguf_ctx, id);
580
581            if (type_arr == GGUF_TYPE_STRING) {
582                const uint64_t nstr_expected = rng() % (16 + 1);
583                if (arr_n != nstr_expected) {
584                    ok = false;
585                    continue;
586                }
587                for (uint64_t istr = 0; istr < nstr_expected; ++istr) {
588                    const char * str = gguf_get_arr_str(gguf_ctx, id, istr);
589                    const uint64_t n = strlen(str);
590                    const uint64_t n_expected = rng() % (sizeof(uint32_t) + 1);
591
592                    if (n != n_expected) {
593                        ok = false;
594                        continue;
595                    }
596                    const char * str_expected = reinterpret_cast<const char *>(&data[istr]);
597                    if (strncmp(str, str_expected, n) != 0) {
598                        ok = false;
599                        continue;
600                    }
601                }
602                continue;
603            }
604
605            const uint64_t arr_n_expected = (rng() % sizeof(data)) / type_size;
606            if (arr_n != arr_n_expected) {
607                ok = false;
608                continue;
609            }
610
611            const char * data_gguf = reinterpret_cast<const char *>(gguf_get_arr_data(gguf_ctx, id));
612
613            if (type_arr == GGUF_TYPE_BOOL) {
614                for (size_t arr_i = 0; arr_i < arr_n; ++arr_i) {
615                    if (bool(data8[arr_i]) != bool(data_gguf[arr_i])) {
616                        ok = false;
617                    }
618                }
619                continue;
620            }
621
622            if (!std::equal(data8, data8 + arr_n*type_size, data_gguf)) {
623                ok = false;
624            }
625            continue;
626        }
627
628        const char * data_gguf = reinterpret_cast<const char *>(gguf_get_val_data(gguf_ctx, id));
629
630        if (type == GGUF_TYPE_BOOL) {
631            if (bool(*data8) != bool(*data_gguf)) {
632                ok = false;
633            }
634            continue;
635        }
636
637        if (!std::equal(data8, data8 + gguf_type_size(type), data_gguf)) {
638            ok = false;
639        }
640    }
641
642    const uint32_t expected_alignment = alignment_defined ? 1 : GGUF_DEFAULT_ALIGNMENT;
643    if (gguf_get_alignment(gguf_ctx) != expected_alignment) {
644        ok = false;
645    }
646
647    return ok;
648}
649
650static bool handcrafted_check_tensors(const gguf_context * gguf_ctx, const unsigned int seed) {
651    if (!gguf_ctx) {
652        return false;
653    }
654
655    std::mt19937 rng(seed);
656
657    std::vector<tensor_config_t> tensor_configs = get_tensor_configs(rng);
658
659    // Call get_kv_types to get the same RNG state:
660    get_kv_types(rng);
661
662    bool ok = true;
663
664    const int id_alignment = gguf_find_key(gguf_ctx, GGUF_KEY_GENERAL_ALIGNMENT);
665    const uint32_t alignment = id_alignment >= 0 ? gguf_get_val_u32(gguf_ctx, id_alignment) : GGUF_DEFAULT_ALIGNMENT;
666
667    uint64_t expected_offset = 0;
668    for (int i = 0; i < int(tensor_configs.size()); ++i) {
669        const ggml_type                          type  = tensor_configs[i].first;
670        const std::array<int64_t, GGML_MAX_DIMS> shape = tensor_configs[i].second;
671
672        const std::string name = "my_tensor_" + std::to_string(i);
673        const int id = gguf_find_tensor(gguf_ctx, name.c_str());
674
675        if (id >= 0) {
676            if (std::string(gguf_get_tensor_name(gguf_ctx, id)) != name) {
677                ok = false;
678            }
679
680            if (gguf_get_tensor_type(gguf_ctx, id) != type) {
681                ok = false;
682            }
683
684            const int64_t * ne = gguf_get_tensor_ne(gguf_ctx, id);
685            for (int j = 0; j < GGML_MAX_DIMS; ++j) {
686                if (ne[j] != shape[j]) {
687                    ok = false;
688                }
689            }
690        } else {
691            ok = false;
692            continue;
693        }
694
695        const size_t offset = gguf_get_tensor_offset(gguf_ctx, id);
696
697        if (offset != expected_offset) {
698            ok = false;
699        }
700
701        int64_t ne = shape[0];
702        for (size_t j = 1; j < GGML_MAX_DIMS; ++j) {
703            ne *= shape[j];
704        }
705        expected_offset += GGML_PAD(ggml_row_size(type, ne), alignment);
706    }
707
708    return ok;
709}
710
711static bool handcrafted_check_tensor_data(const gguf_context * gguf_ctx, const unsigned int seed, FILE * file) {
712    if (!gguf_ctx) {
713        return false;
714    }
715
716    std::mt19937 rng(seed);
717
718    std::vector<tensor_config_t> tensor_configs = get_tensor_configs(rng);
719
720    bool ok = true;
721
722    for (int i = 0; i < int(tensor_configs.size()); ++i) {
723        const ggml_type                          type  = tensor_configs[i].first;
724        const std::array<int64_t, GGML_MAX_DIMS> shape = tensor_configs[i].second;
725
726        int64_t ne = shape[0];
727        for (size_t j = 1; j < GGML_MAX_DIMS; ++j) {
728            ne *= shape[j];
729        }
730        const size_t size = ggml_row_size(type, ne);
731
732        const std::string name = "my_tensor_" + std::to_string(i);
733        const size_t offset = gguf_get_tensor_offset(gguf_ctx, gguf_find_tensor(gguf_ctx, name.c_str()));
734
735        std::vector<uint8_t> data(size);
736        GGML_ASSERT(fseek(file, gguf_get_data_offset(gguf_ctx) + offset, SEEK_SET) == 0);
737        GGML_ASSERT(fread(data.data(), 1, data.size(), file) == data.size());
738
739        for (size_t j = 0; j < size; ++j) {
740            const uint8_t expected_byte = (j + offset) % 256;
741            if (data[j] != expected_byte) {
742                ok = false;
743            }
744        }
745    }
746
747    return ok;
748}
749
750static std::pair<int, int> test_handcrafted_file(const unsigned int seed) {
751    int npass = 0;
752    int ntest = 0;
753
754    const std::vector<handcrafted_file_type> hfts = {
755        HANDCRAFTED_HEADER_BAD_MAGIC,
756        HANDCRAFTED_HEADER_BAD_VERSION_0,
757        HANDCRAFTED_HEADER_BAD_VERSION_1,
758        HANDCRAFTED_HEADER_BAD_VERSION_FUTURE,
759        HANDCRAFTED_HEADER_BAD_N_KV,
760        HANDCRAFTED_HEADER_BAD_N_TENSORS,
761        HANDCRAFTED_HEADER_EMPTY,
762
763        HANDCRAFTED_KV_BAD_KEY_SIZE,
764        HANDCRAFTED_KV_EMPTY_KEY,
765        HANDCRAFTED_KV_BAD_TYPE,
766        HANDCRAFTED_KV_DUPLICATE_KEY,
767        HANDCRAFTED_KV_BAD_ALIGN,
768        HANDCRAFTED_KV_WRONG_TYPE_ALIGN,
769        HANDCRAFTED_KV_SUCCESS,
770
771        HANDCRAFTED_TENSORS_BAD_NAME_SIZE,
772        HANDCRAFTED_TENSORS_BAD_N_DIMS,
773        HANDCRAFTED_TENSORS_BAD_SHAPE,
774        HANDCRAFTED_TENSORS_ZERO_DIM,
775        HANDCRAFTED_TENSORS_NE_TOO_BIG,
776        HANDCRAFTED_TENSORS_NBYTES_TOO_BIG,
777        HANDCRAFTED_TENSORS_BAD_TYPE,
778        HANDCRAFTED_TENSORS_BAD_OFFSET,
779        HANDCRAFTED_TENSORS_DUPLICATE_NAME,
780        HANDCRAFTED_TENSORS_BAD_ALIGN,
781        HANDCRAFTED_TENSORS_INCONSISTENT_ALIGN,
782        HANDCRAFTED_TENSORS_SUCCESS,
783        HANDCRAFTED_TENSORS_CUSTOM_ALIGN,
784
785        HANDCRAFTED_DATA_NOT_ENOUGH_DATA,
786        HANDCRAFTED_DATA_BAD_ALIGN,
787        HANDCRAFTED_DATA_INCONSISTENT_ALIGN,
788        HANDCRAFTED_DATA_MEM_SIZE_OVERFLOW,
789        HANDCRAFTED_DATA_SUCCESS,
790        HANDCRAFTED_DATA_CUSTOM_ALIGN,
791    };
792
793    for (enum handcrafted_file_type hft : hfts) {
794        printf("%s: handcrafted_file_type=%s\n", __func__, handcrafted_file_type_name(hft).c_str());
795        FILE * file = get_handcrafted_file(seed, hft);
796
797#ifdef _WIN32
798        if (!file) {
799            printf("failed to create tmpfile(), needs elevated privileges on Windows");
800            printf("skipping tests");
801            continue;
802        }
803#else
804        GGML_ASSERT(file);
805#endif // _WIN32
806
807        struct ggml_context * ctx = nullptr;
808        struct gguf_init_params gguf_params = {
809            /*no_alloc =*/ false,
810            /*ctx      =*/ hft >= offset_has_data ? &ctx : nullptr,
811        };
812
813        struct gguf_context * gguf_ctx = gguf_init_from_file_ptr(file, gguf_params);
814
815        if (expect_context_not_null(hft)) {
816            printf("%s:   - context_not_null: ", __func__);
817        } else {
818            printf("%s:   - context_null: ", __func__);
819        }
820        if (bool(gguf_ctx) == expect_context_not_null(hft)) {
821            printf("\033[1;32mOK\033[0m\n");
822            npass++;
823        } else {
824            printf("\033[1;31mFAIL\033[0m\n");
825        }
826        ntest++;
827
828        if (hft >= offset_has_data && !expect_context_not_null(hft)) {
829            printf("%s:   - no_dangling_ggml_context_pointer: ", __func__);
830            if (ctx) {
831                printf("\033[1;31mFAIL\033[0m\n");
832            } else {
833                printf("\033[1;32mOK\033[0m\n");
834                npass++;
835            }
836            ntest++;
837        }
838
839        const bool alignment_defined = hft == HANDCRAFTED_TENSORS_CUSTOM_ALIGN || hft == HANDCRAFTED_DATA_CUSTOM_ALIGN;
840
841        if (expect_context_not_null(hft)) {
842            printf("%s:   - check_header: ", __func__);
843            if (handcrafted_check_header(gguf_ctx, seed, hft >= offset_has_kv, hft >= offset_has_tensors, alignment_defined)) {
844                printf("\033[1;32mOK\033[0m\n");
845                npass++;
846            } else {
847                printf("\033[1;31mFAIL\033[0m\n");
848            }
849            ntest++;
850        }
851
852        if (expect_context_not_null(hft) && hft >= offset_has_kv) {
853            printf("%s:   - check_kv: ", __func__);
854            if (handcrafted_check_kv(gguf_ctx, seed, hft >= offset_has_tensors, alignment_defined)) {
855                printf("\033[1;32mOK\033[0m\n");
856                npass++;
857            } else {
858                printf("\033[1;31mFAIL\033[0m\n");
859            }
860            ntest++;
861        }
862
863        // HANDCRAFTED_TENSORS_ZERO_DIM deliberately mangles the tensor shapes to 0 elements,
864        // so only assert that it loads without crashing; skip the exact-geometry comparison.
865        if (expect_context_not_null(hft) && hft >= offset_has_tensors && hft != HANDCRAFTED_TENSORS_ZERO_DIM) {
866            printf("%s:   - check_tensors: ", __func__);
867            if (handcrafted_check_tensors(gguf_ctx, seed)) {
868                printf("\033[1;32mOK\033[0m\n");
869                npass++;
870            } else {
871                printf("\033[1;31mFAIL\033[0m\n");
872            }
873            ntest++;
874        }
875
876        if (expect_context_not_null(hft) && hft >= offset_has_data) {
877            printf("%s:   - check_tensor_data: ", __func__);
878            if (handcrafted_check_tensor_data(gguf_ctx, seed, file)) {
879                printf("\033[1;32mOK\033[0m\n");
880                npass++;
881            } else {
882                printf("\033[1;31mFAIL\033[0m\n");
883            }
884            ntest++;
885        }
886
887        fclose(file);
888        if (gguf_ctx) {
889            ggml_free(ctx);
890            gguf_free(gguf_ctx);
891        }
892        printf("\n");
893    }
894
895
896    return std::make_pair(npass, ntest);
897}
898
899struct random_gguf_context_result {
900    struct gguf_context * gguf_ctx;
901    struct ggml_context * ctx;
902    ggml_backend_buffer_t buffer;
903};
904
905static struct random_gguf_context_result get_random_gguf_context(ggml_backend_t backend, const unsigned int seed) {
906    std::mt19937 rng(seed);
907
908    struct gguf_context * gguf_ctx = gguf_init_empty();
909
910    for (int i = 0; i < 256; ++i) {
911        const std::string key = "my_key_" + std::to_string(rng() % 1024);
912        const enum gguf_type type = gguf_type(rng() % GGUF_TYPE_COUNT);
913
914        switch (type) {
915            case GGUF_TYPE_UINT8:   gguf_set_val_u8  (gguf_ctx, key.c_str(), rng() % (1 <<  7));             break;
916            case GGUF_TYPE_INT8:    gguf_set_val_i8  (gguf_ctx, key.c_str(), rng() % (1 <<  7) - (1 <<  6)); break;
917            case GGUF_TYPE_UINT16:  gguf_set_val_u16 (gguf_ctx, key.c_str(), rng() % (1 << 15));             break;
918            case GGUF_TYPE_INT16:   gguf_set_val_i16 (gguf_ctx, key.c_str(), rng() % (1 << 15) - (1 << 14)); break;
919            case GGUF_TYPE_UINT32:  gguf_set_val_u32 (gguf_ctx, key.c_str(), rng());                         break;
920            case GGUF_TYPE_INT32:   gguf_set_val_i32 (gguf_ctx, key.c_str(), rng()             - (1 << 30)); break;
921            case GGUF_TYPE_FLOAT32: gguf_set_val_f32 (gguf_ctx, key.c_str(), rng() % 1024      - 512);       break;
922            case GGUF_TYPE_BOOL:    gguf_set_val_bool(gguf_ctx, key.c_str(), rng() % 2 == 0);                break;
923            case GGUF_TYPE_STRING:  gguf_set_val_str (gguf_ctx, key.c_str(), std::to_string(rng()).c_str()); break;
924            case GGUF_TYPE_UINT64:  gguf_set_val_u64 (gguf_ctx, key.c_str(), rng());                         break;
925            case GGUF_TYPE_INT64:   gguf_set_val_i64 (gguf_ctx, key.c_str(), rng()             - (1 << 30)); break;
926            case GGUF_TYPE_FLOAT64: gguf_set_val_f32 (gguf_ctx, key.c_str(), rng() % 1024      - 512);       break;
927            case GGUF_TYPE_ARRAY: {
928                const enum gguf_type type_arr = gguf_type(rng() % GGUF_TYPE_COUNT);
929                const uint64_t ne = rng() % 1024;
930
931                switch (type_arr) {
932                    case GGUF_TYPE_UINT8:
933                    case GGUF_TYPE_INT8:
934                    case GGUF_TYPE_UINT16:
935                    case GGUF_TYPE_INT16:
936                    case GGUF_TYPE_UINT32:
937                    case GGUF_TYPE_INT32:
938                    case GGUF_TYPE_FLOAT32:
939                    case GGUF_TYPE_BOOL:
940                    case GGUF_TYPE_UINT64:
941                    case GGUF_TYPE_INT64:
942                    case GGUF_TYPE_FLOAT64: {
943                        const size_t nbytes = ne*gguf_type_size(type_arr);
944                        std::vector<uint32_t> random_data((nbytes + sizeof(uint32_t) - 1) / sizeof(uint32_t));
945                        for (size_t j = 0; j < random_data.size(); ++j) {
946                            random_data[j] = rng();
947                            if (type_arr == GGUF_TYPE_BOOL) {
948                                random_data[j] &= 0x01010101; // the sanitizer complains if booleans are not 0 or 1
949                            }
950                        }
951                        gguf_set_arr_data(gguf_ctx, key.c_str(), type_arr, random_data.data(), ne);
952                    } break;
953                    case GGUF_TYPE_STRING: {
954                        std::vector<std::string>  data_cpp(ne);
955                        std::vector<const char *> data_c(ne);
956                        for (size_t j = 0; j < data_cpp.size(); ++j) {
957                            data_cpp[j] = std::to_string(rng());
958                            data_c[j]   = data_cpp[j].c_str();
959                        }
960                        gguf_set_arr_str(gguf_ctx, key.c_str(), data_c.data(), ne);
961                    } break;
962                    case GGUF_TYPE_ARRAY: {
963                        break; // not supported
964                    }
965                    case GGUF_TYPE_COUNT:
966                    default: {
967                        GGML_ABORT("fatal error");
968                    }
969                }
970            } break;
971            case GGUF_TYPE_COUNT:
972            default: {
973                GGML_ABORT("fatal error");
974            }
975        }
976    }
977
978    struct ggml_init_params ggml_params = {
979        /*.mem_size   =*/ 256*ggml_tensor_overhead(),
980        /*.mem_buffer =*/ nullptr,
981        /*.no_alloc   =*/ true,
982    };
983    struct ggml_context * ctx = ggml_init(ggml_params);
984
985    for (int i = 0; i < 256; ++i) {
986        const std::string name = "my_tensor_" + std::to_string(i);
987        const enum ggml_type type = ggml_type(rng() % GGML_TYPE_COUNT);
988        const size_t type_size = ggml_type_size(type);
989
990        if (type_size == 0) {
991            continue;
992        }
993
994        const int n_dims = 1 + rng() % GGML_MAX_DIMS;
995        int64_t ne[GGML_MAX_DIMS];
996        ne[0] = (1 + rng() % 10) * ggml_blck_size(type);
997        for (int j = 1; j < n_dims; ++j) {
998            ne[j] = 1 + rng() % 10;
999        }
1000
1001        struct ggml_tensor * tensor = ggml_new_tensor(ctx, type, n_dims, ne);
1002        ggml_set_name(tensor, name.c_str());
1003    }
1004
1005    ggml_backend_buffer_t buf = ggml_backend_alloc_ctx_tensors(ctx, backend);
1006    for (struct ggml_tensor * t = ggml_get_first_tensor(ctx); t != nullptr; t = ggml_get_next_tensor(ctx, t)) {
1007        const size_t nbytes = ggml_nbytes(t);
1008        std::vector<uint32_t> random_data((nbytes + sizeof(uint32_t) - 1) / sizeof(uint32_t));
1009        for (size_t j = 0; j < random_data.size(); ++j) {
1010            random_data[j] = rng();
1011        }
1012        ggml_backend_tensor_set(t, random_data.data(), 0, nbytes);
1013
1014        gguf_add_tensor(gguf_ctx, t);
1015    }
1016
1017    return {gguf_ctx, ctx, buf};
1018}
1019
1020static bool all_kv_in_other(const gguf_context * ctx, const gguf_context * other) {
1021    bool ok = true;
1022
1023    const int n_kv = gguf_get_n_kv(ctx);
1024    for (int id = 0; id < n_kv; ++id) {
1025        const char * name = gguf_get_key(ctx, id);
1026
1027        const int idx_other = gguf_find_key(other, name);
1028        if (idx_other < 0) {
1029            ok = false;
1030            continue;
1031        }
1032
1033        const gguf_type type = gguf_get_kv_type(ctx, id);
1034        if (type != gguf_get_kv_type(other, idx_other)) {
1035            ok = false;
1036            continue;
1037        }
1038
1039        if (type == GGUF_TYPE_ARRAY) {
1040            const size_t arr_n = gguf_get_arr_n(ctx, id);
1041            if (arr_n != gguf_get_arr_n(other, idx_other)) {
1042                ok = false;
1043                continue;
1044            }
1045
1046            const gguf_type type_arr = gguf_get_arr_type(ctx, id);
1047            if (type_arr != gguf_get_arr_type(other, idx_other)) {
1048                ok = false;
1049                continue;
1050            }
1051
1052            if (type_arr == GGUF_TYPE_BOOL) {
1053                const int8_t * data       = reinterpret_cast<const int8_t *>(gguf_get_arr_data(ctx,   id));
1054                const int8_t * data_other = reinterpret_cast<const int8_t *>(gguf_get_arr_data(other, idx_other));
1055                for (size_t arr_i = 0; arr_i < arr_n; ++arr_i) {
1056                    if (bool(data[arr_i]) != bool(data_other[arr_i])) {
1057                        ok = false;
1058                    }
1059                }
1060                continue;
1061            }
1062
1063            if (type_arr == GGUF_TYPE_STRING) {
1064                for (size_t arr_i = 0; arr_i < arr_n; ++arr_i) {
1065                    const std::string str       = gguf_get_arr_str(ctx,   id,       arr_i);
1066                    const std::string str_other = gguf_get_arr_str(other, idx_other, arr_i);
1067                    if (str != str_other) {
1068                        ok = false;
1069                    }
1070                }
1071                continue;
1072            }
1073
1074            const int8_t * data       = reinterpret_cast<const int8_t *>(gguf_get_arr_data(ctx,   id));
1075            const int8_t * data_other = reinterpret_cast<const int8_t *>(gguf_get_arr_data(other, idx_other));
1076            if (!std::equal(data, data + arr_n*gguf_type_size(type_arr), data_other)) {
1077                ok = false;
1078            }
1079            continue;
1080        }
1081
1082        if (type == GGUF_TYPE_STRING) {
1083            const std::string str       = gguf_get_val_str(ctx,   id);
1084            const std::string str_other = gguf_get_val_str(other, idx_other);
1085            if (str != str_other) {
1086                ok = false;
1087            }
1088            continue;
1089        }
1090
1091        const char * data       = reinterpret_cast<const char *>(gguf_get_val_data(ctx,   id));
1092        const char * data_other = reinterpret_cast<const char *>(gguf_get_val_data(other, idx_other));
1093        if (!std::equal(data, data + gguf_type_size(type), data_other)) {
1094            ok = false;
1095        }
1096    }
1097
1098    return ok;
1099}
1100
1101static bool all_tensors_in_other(const gguf_context * ctx, const gguf_context * other) {
1102    bool ok = true;
1103
1104    const int n_tensors = gguf_get_n_tensors(ctx);
1105    for (int id = 0; id < n_tensors; ++id) {
1106        const std::string name = gguf_get_tensor_name(ctx, id);
1107
1108        const int idx_other = gguf_find_tensor(other, name.c_str());
1109        if (id != idx_other) {
1110            ok = false;
1111            if (idx_other < 0) {
1112                continue;
1113            }
1114        }
1115
1116        const ggml_type type = gguf_get_tensor_type(ctx, id);
1117        if (type != gguf_get_tensor_type(other, id)) {
1118            ok = false;
1119        }
1120
1121        const size_t offset = gguf_get_tensor_offset(ctx, id);
1122        if (offset != gguf_get_tensor_offset(other, id)) {
1123            ok = false;
1124        }
1125    }
1126
1127    return ok;
1128}
1129
1130static bool same_tensor_data(const struct ggml_context * orig, const struct ggml_context * read) {
1131    bool ok = true;
1132
1133    struct ggml_tensor * t_orig = ggml_get_first_tensor(orig);
1134    struct ggml_tensor * t_read = ggml_get_first_tensor(read);
1135
1136    if (std::string(t_read->name) != "GGUF tensor data binary blob") {
1137        return false;
1138    }
1139    t_read = ggml_get_next_tensor(read, t_read);
1140
1141    while (t_orig) {
1142        if (!t_read) {
1143            ok = false;
1144            break;
1145        }
1146
1147        const size_t nbytes = ggml_nbytes(t_orig);
1148        if (ggml_nbytes(t_read) != nbytes) {
1149            ok = false;
1150            break;
1151        }
1152        std::vector<char> data_orig(nbytes);
1153        ggml_backend_tensor_get(t_orig, data_orig.data(), 0, nbytes);
1154        if (!std::equal(data_orig.data(), data_orig.data() + nbytes, reinterpret_cast<const char *>(t_read->data))) {
1155            ok = false;
1156        }
1157
1158        t_orig = ggml_get_next_tensor(orig, t_orig);
1159        t_read = ggml_get_next_tensor(read, t_read);
1160    }
1161    if (t_read) {
1162        ok = false;
1163    }
1164
1165    return ok;
1166}
1167
1168enum roundtrip_read_mode {
1169    ROUNDTRIP_READ_MODE_FILE,
1170    ROUNDTRIP_READ_MODE_BUFFER,
1171    ROUNDTRIP_READ_MODE_CALLBACK,
1172};
1173
1174static const char * roundtrip_read_mode_name(const roundtrip_read_mode mode) {
1175    switch (mode) {
1176        case ROUNDTRIP_READ_MODE_FILE:     return "file";
1177        case ROUNDTRIP_READ_MODE_BUFFER:   return "buffer";
1178        case ROUNDTRIP_READ_MODE_CALLBACK: return "callback";
1179    }
1180
1181    GGML_ABORT("fatal error");
1182}
1183
1184static std::pair<int, int> test_roundtrip(
1185        ggml_backend_dev_t dev, const unsigned int seed, const bool only_meta,
1186        const roundtrip_read_mode read_mode) {
1187    ggml_backend_t backend = ggml_backend_dev_init(dev, nullptr);
1188    printf("%s: device=%s, backend=%s, only_meta=%s, read_mode=%s\n",
1189        __func__, ggml_backend_dev_description(dev), ggml_backend_name(backend),
1190        only_meta ? "yes" : "no", roundtrip_read_mode_name(read_mode));
1191
1192    int npass = 0;
1193    int ntest = 0;
1194
1195    struct gguf_context * gguf_ctx_0;
1196    struct ggml_context * ctx_0;
1197    ggml_backend_buffer_t bbuf;
1198    {
1199        struct random_gguf_context_result result = get_random_gguf_context(backend, seed);
1200        gguf_ctx_0 = result.gguf_ctx;
1201        ctx_0      = result.ctx;
1202        bbuf       = result.buffer;
1203    }
1204
1205    FILE * file = tmpfile();
1206
1207#ifdef _WIN32
1208    if (!file) {
1209        printf("failed to create tmpfile(), needs elevated privileges on Windows");
1210        printf("skipping tests");
1211        return std::make_pair(0, 0);
1212    }
1213#else
1214    GGML_ASSERT(file);
1215#endif // _WIN32
1216
1217    gguf_write_to_file_ptr(gguf_ctx_0, file, only_meta);
1218    rewind(file);
1219
1220    struct ggml_context * ctx_1 = nullptr;
1221    struct gguf_init_params gguf_params = {
1222        /*no_alloc =*/ false,
1223        /*ctx      =*/ only_meta ? nullptr : &ctx_1,
1224    };
1225    struct gguf_context * gguf_ctx_1 = nullptr;
1226    const std::vector<uint8_t> data = read_mode == ROUNDTRIP_READ_MODE_FILE
1227        ? std::vector<uint8_t>()
1228        : read_file_to_buffer(file);
1229
1230    if (read_mode == ROUNDTRIP_READ_MODE_BUFFER) {
1231        gguf_ctx_1 = gguf_init_from_buffer(data.data(), data.size(), gguf_params);
1232    } else if (read_mode == ROUNDTRIP_READ_MODE_CALLBACK) {
1233        callback_reader_data reader = {
1234            /*.data = */ data.data(),
1235            /*.size = */ data.size(),
1236        };
1237        gguf_ctx_1 = gguf_init_from_callback(read_buffer_callback, &reader, 4096, 4ull << 30 /* 4GB */, gguf_params);
1238    } else {
1239        gguf_ctx_1 = gguf_init_from_file_ptr(file, gguf_params);
1240    }
1241
1242    printf("%s: same_version: ", __func__);
1243    if (gguf_get_version(gguf_ctx_0) == gguf_get_version(gguf_ctx_1)) {
1244        printf("\033[1;32mOK\033[0m\n");
1245        npass++;
1246    } else {
1247        printf("\033[1;31mFAIL\033[0m\n");
1248    }
1249    ntest++;
1250
1251    printf("%s: same_n_kv: ", __func__);
1252    if (gguf_get_n_kv(gguf_ctx_0) == gguf_get_n_kv(gguf_ctx_1)) {
1253        printf("\033[1;32mOK\033[0m\n");
1254        npass++;
1255    } else {
1256        printf("\033[1;31mFAIL\033[0m\n");
1257    }
1258    ntest++;
1259
1260    printf("%s: same_n_tensors: ", __func__);
1261    if (gguf_get_n_tensors(gguf_ctx_0) == gguf_get_n_tensors(gguf_ctx_1)) {
1262        printf("\033[1;32mOK\033[0m\n");
1263        npass++;
1264    } else {
1265        printf("\033[1;31mFAIL\033[0m\n");
1266    }
1267    ntest++;
1268
1269    printf("%s: all_orig_kv_in_read: ", __func__);
1270    if (all_kv_in_other(gguf_ctx_0, gguf_ctx_1)) {
1271        printf("\033[1;32mOK\033[0m\n");
1272        npass++;
1273    } else {
1274        printf("\033[1;31mFAIL\033[0m\n");
1275    }
1276    ntest++;
1277
1278    printf("%s: all_read_kv_in_orig: ", __func__);
1279    if (all_kv_in_other(gguf_ctx_1, gguf_ctx_0)) {
1280        printf("\033[1;32mOK\033[0m\n");
1281        npass++;
1282    } else {
1283        printf("\033[1;31mFAIL\033[0m\n");
1284    }
1285    ntest++;
1286
1287    printf("%s: all_orig_tensors_in_read: ", __func__);
1288    if (all_tensors_in_other(gguf_ctx_0, gguf_ctx_1)) {
1289        printf("\033[1;32mOK\033[0m\n");
1290        npass++;
1291    } else {
1292        printf("\033[1;31mFAIL\033[0m\n");
1293    }
1294    ntest++;
1295
1296    printf("%s: all_read_tensors_in_orig: ", __func__);
1297    if (all_tensors_in_other(gguf_ctx_1, gguf_ctx_0)) {
1298        printf("\033[1;32mOK\033[0m\n");
1299        npass++;
1300    } else {
1301        printf("\033[1;31mFAIL\033[0m\n");
1302    }
1303    ntest++;
1304
1305    if (!only_meta) {
1306        printf("%s: same_tensor_data: ", __func__);
1307        if (same_tensor_data(ctx_0, ctx_1)) {
1308            printf("\033[1;32mOK\033[0m\n");
1309            npass++;
1310        } else {
1311            printf("\033[1;31mFAIL\033[0m\n");
1312        }
1313        ntest++;
1314    }
1315
1316    ggml_backend_buffer_free(bbuf);
1317    ggml_free(ctx_0);
1318    ggml_free(ctx_1);
1319    gguf_free(gguf_ctx_0);
1320    gguf_free(gguf_ctx_1);
1321    ggml_backend_free(backend);
1322    fclose(file);
1323
1324    printf("\n");
1325    return std::make_pair(npass, ntest);
1326}
1327
1328static std::pair<int, int> test_gguf_set_kv(ggml_backend_dev_t dev, const unsigned int seed) {
1329    ggml_backend_t backend = ggml_backend_dev_init(dev, nullptr);
1330    printf("%s: device=%s, backend=%s\n", __func__, ggml_backend_dev_description(dev), ggml_backend_name(backend));
1331
1332    int npass = 0;
1333    int ntest = 0;
1334
1335    struct gguf_context * gguf_ctx_0;
1336    struct ggml_context * ctx_0;
1337    ggml_backend_buffer_t bbuf_0;
1338    {
1339        struct random_gguf_context_result result = get_random_gguf_context(backend, seed);
1340        gguf_ctx_0 = result.gguf_ctx;
1341        ctx_0      = result.ctx;
1342        bbuf_0     = result.buffer;
1343    }
1344
1345    struct gguf_context * gguf_ctx_1;
1346    struct ggml_context * ctx_1;
1347    ggml_backend_buffer_t bbuf_1;
1348    {
1349        struct random_gguf_context_result result = get_random_gguf_context(backend, seed + 1);
1350        gguf_ctx_1 = result.gguf_ctx;
1351        ctx_1      = result.ctx;
1352        bbuf_1     = result.buffer;
1353    }
1354
1355    struct gguf_context * gguf_ctx_2 = gguf_init_empty();
1356
1357    gguf_set_kv(gguf_ctx_1, gguf_ctx_0);
1358    gguf_set_kv(gguf_ctx_2, gguf_ctx_0);
1359
1360    printf("%s: same_n_kv: ", __func__);
1361    if (gguf_get_n_kv(gguf_ctx_0) == gguf_get_n_kv(gguf_ctx_2)) {
1362        printf("\033[1;32mOK\033[0m\n");
1363        npass++;
1364    } else {
1365        printf("\033[1;31mFAIL\033[0m\n");
1366    }
1367    ntest++;
1368
1369    printf("%s: all_kv_0_in_1: ", __func__);
1370    if (all_kv_in_other(gguf_ctx_0, gguf_ctx_1)) {
1371        printf("\033[1;32mOK\033[0m\n");
1372        npass++;
1373    } else {
1374        printf("\033[1;31mFAIL\033[0m\n");
1375    }
1376    ntest++;
1377
1378    printf("%s: all_kv_0_in_2: ", __func__);
1379    if (all_kv_in_other(gguf_ctx_0, gguf_ctx_2)) {
1380        printf("\033[1;32mOK\033[0m\n");
1381        npass++;
1382    } else {
1383        printf("\033[1;31mFAIL\033[0m\n");
1384    }
1385    ntest++;
1386
1387    gguf_set_kv(gguf_ctx_0, gguf_ctx_1);
1388
1389    printf("%s: same_n_kv_after_double_copy: ", __func__);
1390    if (gguf_get_n_kv(gguf_ctx_0) == gguf_get_n_kv(gguf_ctx_1)) {
1391        printf("\033[1;32mOK\033[0m\n");
1392        npass++;
1393    } else {
1394        printf("\033[1;31mFAIL\033[0m\n");
1395    }
1396    ntest++;
1397
1398    printf("%s: all_kv_1_in_0_after_double_copy: ", __func__);
1399    if (all_kv_in_other(gguf_ctx_1, gguf_ctx_0)) {
1400        printf("\033[1;32mOK\033[0m\n");
1401        npass++;
1402    } else {
1403        printf("\033[1;31mFAIL\033[0m\n");
1404    }
1405    ntest++;
1406
1407    ggml_backend_buffer_free(bbuf_0);
1408    ggml_backend_buffer_free(bbuf_1);
1409    ggml_free(ctx_0);
1410    ggml_free(ctx_1);
1411    gguf_free(gguf_ctx_0);
1412    gguf_free(gguf_ctx_1);
1413    gguf_free(gguf_ctx_2);
1414    ggml_backend_free(backend);
1415
1416    printf("\n");
1417    return std::make_pair(npass, ntest);
1418}
1419
1420static void print_usage() {
1421    printf("usage: test-gguf [seed]\n");
1422    printf("  if no seed is unspecified then a random seed is used\n");
1423}
1424
1425int main(int argc, char ** argv) {
1426    if (argc > 2) {
1427        print_usage();
1428        return 1;
1429    }
1430
1431    std::random_device rd;
1432    const unsigned int seed = argc < 2 ? rd() : std::stoi(argv[1]);
1433
1434    // Initialize ggml backends early so the prints aren't interleaved with the test results:
1435    ggml_backend_dev_count();
1436    fprintf(stderr, "\n");
1437
1438    int npass = 0;
1439    int ntest = 0;
1440    {
1441        std::pair<int, int> result = test_handcrafted_file(seed);
1442        npass += result.first;
1443        ntest += result.second;
1444    }
1445
1446    for (size_t i = 0; i < ggml_backend_dev_count(); ++i) {
1447        ggml_backend_dev_t dev = ggml_backend_dev_get(i);
1448
1449        for (bool only_meta : {true, false}) {
1450            std::pair<int, int> result = test_roundtrip(dev, seed, only_meta, ROUNDTRIP_READ_MODE_FILE);
1451            npass += result.first;
1452            ntest += result.second;
1453        }
1454        {
1455            std::pair<int, int> result = test_roundtrip(dev, seed, /*only_meta=*/false, ROUNDTRIP_READ_MODE_BUFFER);
1456            npass += result.first;
1457            ntest += result.second;
1458        }
1459        {
1460            std::pair<int, int> result = test_roundtrip(dev, seed, /*only_meta=*/false, ROUNDTRIP_READ_MODE_CALLBACK);
1461            npass += result.first;
1462            ntest += result.second;
1463        }
1464
1465        {
1466            std::pair<int, int> result = test_gguf_set_kv(dev, seed);
1467            npass += result.first;
1468            ntest += result.second;
1469        }
1470    }
1471
1472    printf("%d/%d tests passed\n", npass, ntest);
1473    if (npass != ntest) {
1474        printf("\033[1;31mFAIL\033[0m\n");
1475        return 1;
1476    }
1477    printf("\033[1;32mOK\033[0m\n");
1478    return 0;
1479}