diff --git a/ggml/src/gguf.cpp b/ggml/src/gguf.cpp index 9f9e4fe5d104..eb0262e421c9 100644 --- a/ggml/src/gguf.cpp +++ b/ggml/src/gguf.cpp @@ -682,7 +682,12 @@ static struct gguf_context * gguf_init_from_reader(const struct gguf_reader & gr } // check that the total number of elements is representable - if (ok && ((INT64_MAX/info.t.ne[1] <= info.t.ne[0]) || + // Zero-sized tensors are valid and must not be used as divisors. + bool has_zero_dim = false; + for (uint32_t j = 0; j < GGML_MAX_DIMS; ++j) { + has_zero_dim |= info.t.ne[j] == 0; + } + if (ok && !has_zero_dim && ((INT64_MAX/info.t.ne[1] <= info.t.ne[0]) || (INT64_MAX/info.t.ne[2] <= info.t.ne[0]*info.t.ne[1]) || (INT64_MAX/info.t.ne[3] <= info.t.ne[0]*info.t.ne[1]*info.t.ne[2]))) { diff --git a/tests/test-gguf.cpp b/tests/test-gguf.cpp index 2875dec806da..f4a47cbaf2f4 100644 --- a/tests/test-gguf.cpp +++ b/tests/test-gguf.cpp @@ -729,6 +729,59 @@ static bool handcrafted_check_tensor_data(const gguf_context * gguf_ctx, const u return ok; } +static FILE * get_zero_dim_tensor_file() { + FILE * file = tmpfile(); + if (!file) { + return nullptr; + } + + helper_write(file, GGUF_MAGIC, 4); + helper_write(file, uint32_t(GGUF_VERSION)); + helper_write(file, uint64_t(1)); // tensors + helper_write(file, uint64_t(0)); // key-value pairs + + const char name[] = "zero_dim"; + helper_write(file, uint64_t(sizeof(name) - 1)); + helper_write(file, name, sizeof(name) - 1); + helper_write(file, uint32_t(2)); + helper_write(file, int64_t(1)); + helper_write(file, int64_t(0)); + helper_write(file, int32_t(GGML_TYPE_I8)); + helper_write(file, uint64_t(0)); + + while (ftell(file) % GGUF_DEFAULT_ALIGNMENT != 0) { + helper_write(file, char(0)); + } + rewind(file); + return file; +} + +static std::pair test_zero_dim_tensor() { + FILE * file = get_zero_dim_tensor_file(); +#ifdef _WIN32 + if (!file) { + printf("failed to create tmpfile(), needs elevated privileges on Windows\n"); + return std::make_pair(0, 0); + } +#else + GGML_ASSERT(file); +#endif // _WIN32 + + struct gguf_init_params params = { + /*no_alloc =*/ true, + /*ctx =*/ nullptr, + }; + struct gguf_context * ctx = gguf_init_from_file_ptr(file, params); + const int tensor_id = ctx ? gguf_find_tensor(ctx, "zero_dim") : -1; + const int64_t * shape = tensor_id >= 0 ? gguf_get_tensor_ne(ctx, tensor_id) : nullptr; + const bool passed = ctx && tensor_id >= 0 && shape && shape[0] == 1 && shape[1] == 0; + + printf("%s: zero_dim_tensor_is_accepted: %s\n", __func__, passed ? "OK" : "FAIL"); + gguf_free(ctx); + fclose(file); + return std::make_pair(passed ? 1 : 0, 1); +} + static std::pair test_handcrafted_file(const unsigned int seed) { int npass = 0; int ntest = 0; @@ -1420,6 +1473,11 @@ int main(int argc, char ** argv) { npass += result.first; ntest += result.second; } + { + std::pair result = test_zero_dim_tensor(); + npass += result.first; + ntest += result.second; + } for (size_t i = 0; i < ggml_backend_dev_count(); ++i) { ggml_backend_dev_t dev = ggml_backend_dev_get(i);