From f8007e2a390cf694eb0b4c0d667e9ef46c4524cc Mon Sep 17 00:00:00 2001 From: Nidhin Date: Fri, 10 Jul 2026 16:43:49 +0530 Subject: [PATCH 01/13] ggml : add fabric vector index C API Add a standalone ggml vector index API and flat f32 TurboVec POC implementation for fabric embedding lookup. Includes .tvim v1 persistence, duplicate-safe atomic adds, scalar top-k search, and focused test coverage. --- ggml/CMakeLists.txt | 1 + ggml/include/ggml-vector-index.h | 130 +++++++ ggml/src/CMakeLists.txt | 2 + ggml/src/ggml-vector-index.cpp | 562 +++++++++++++++++++++++++++++++ tests/CMakeLists.txt | 9 + tests/test-vector-index.cpp | 175 ++++++++++ 6 files changed, 879 insertions(+) create mode 100644 ggml/include/ggml-vector-index.h create mode 100644 ggml/src/ggml-vector-index.cpp create mode 100644 tests/test-vector-index.cpp diff --git a/ggml/CMakeLists.txt b/ggml/CMakeLists.txt index 1f534ad7230d..c1390966712a 100644 --- a/ggml/CMakeLists.txt +++ b/ggml/CMakeLists.txt @@ -336,6 +336,7 @@ set(GGML_PUBLIC_HEADERS include/ggml-rpc.h include/ggml-virtgpu.h include/ggml-sycl.h + include/ggml-vector-index.h include/ggml-vulkan.h include/ggml-webgpu.h include/ggml-zendnn.h diff --git a/ggml/include/ggml-vector-index.h b/ggml/include/ggml-vector-index.h new file mode 100644 index 000000000000..217c58d74696 --- /dev/null +++ b/ggml/include/ggml-vector-index.h @@ -0,0 +1,130 @@ +#pragma once +// +// ggml-vector-index: TurboQuant-style ANN vector-index C API. +// +// POC NOTE +// -------- +// This is the public C API for fabric's vector index. The implementation +// under `ggml/src/ggml-vector-index.cpp` is intentionally naive (full f32 +// storage, scalar dot-product, min-heap top-k). The C API itself is final: +// downstream Bare addon bindings + JS wrappers depend on this signature and +// the on-disk file format documented at the bottom of this header. +// Quantization / SIMD / GPU kernels are future work and will be swapped in +// behind the same API without breaking callers. +// +// Threading: instances are NOT thread-safe. Callers must serialize access +// to a given handle. Multiple handles can be used concurrently. +// +// Endianness: persistence format is fixed little-endian. The POC asserts a +// LE host; we do not currently support big-endian platforms. + +#include "ggml.h" + +#include + +#ifdef __cplusplus +extern "C" { +#endif + +// Opaque handle to a vector index instance. +struct ggml_vec_index; +typedef struct ggml_vec_index ggml_vec_index_t; + +// Error codes returned from int-valued APIs. 0 = OK. Negative = failure. +// `_remove` is the exception: it returns 1 on removal and 0 on miss. +enum ggml_vec_index_error { + GGML_VEC_INDEX_OK = 0, + GGML_VEC_INDEX_E_INVALID_DIM = -1, + GGML_VEC_INDEX_E_INVALID_ARG = -2, + GGML_VEC_INDEX_E_DUPLICATE = -3, + GGML_VEC_INDEX_E_IO = -4, + GGML_VEC_INDEX_E_BAD_MAGIC = -5, + GGML_VEC_INDEX_E_BAD_VERSION = -6, + GGML_VEC_INDEX_E_OOM = -7, + GGML_VEC_INDEX_E_INTERNAL = -99, +}; + +// Lifecycle. +// +// `dim` must be > 0. `bit_width` is reserved for future quantization; the +// POC accepts any value in [1, 32] but ignores it (storage is always full +// f32). Returns NULL on bad args. +GGML_API ggml_vec_index_t * ggml_vec_index_create(int dim, int bit_width); + +GGML_API void ggml_vec_index_free(ggml_vec_index_t * idx); + +// Mutation. +// +// Adds `n` vectors of length `dim` each (row-major, contiguous in `vectors`), +// associating each with the corresponding `ids[i]` (caller-owned external id). +// Returns 0 on success. Returns GGML_VEC_INDEX_E_DUPLICATE if any id already +// exists in the index; in that case the index is unchanged (atomic add). +GGML_API int ggml_vec_index_add( + ggml_vec_index_t * idx, + const float * vectors, + int n, + const uint64_t * ids); + +// Removes the entry for `id` via swap-with-last (slot indices are NOT stable +// across removes; external ids ARE). Returns 1 if removed, 0 if not present, +// negative on error. +GGML_API int ggml_vec_index_remove(ggml_vec_index_t * idx, uint64_t id); + +// Returns 1 if the id is in the index, 0 otherwise. Read-only. +GGML_API int ggml_vec_index_contains(const ggml_vec_index_t * idx, uint64_t id); + +// No-op for the POC. Placeholder for future cache warming / codebook +// resolution after a bulk add. Reserved as a mutating op (warm-up may +// materialize derived state inside the index). +GGML_API void ggml_vec_index_prepare(ggml_vec_index_t * idx); + +// Top-k search. `queries` is `n_q * dim` row-major. `out_scores` and +// `out_ids` are caller-allocated buffers of size `n_q * k`. Each row is +// sorted descending by score (higher = closer / more similar). If the index +// holds fewer than k entries, the remaining slots in each row are filled +// with sentinel values: -FLT_MAX for scores, UINT64_MAX for ids. Read-only +// against the index (does not mutate state). +// +// Score semantics: scalar full-precision dot product. Callers that want +// cosine similarity must L2-normalize their vectors before insert AND +// before query; the index does NOT normalize internally. +GGML_API int ggml_vec_index_search( + const ggml_vec_index_t * idx, + const float * queries, + int n_q, + int k, + float * out_scores, + uint64_t * out_ids); + +// Persistence. Format is .tvim version 1; see bottom of this header. +GGML_API int ggml_vec_index_write( + ggml_vec_index_t * idx, + const char * path); + +// Returns NULL on failure (caller can inspect errno for I/O specifics). +GGML_API ggml_vec_index_t * ggml_vec_index_load(const char * path); + +// Stats. +GGML_API int ggml_vec_index_len(const ggml_vec_index_t * idx); +GGML_API int ggml_vec_index_dim(const ggml_vec_index_t * idx); +GGML_API int ggml_vec_index_bit_width(const ggml_vec_index_t * idx); + +// File format (.tvim version 1, all little-endian): +// +// offset size field +// ------ ----- ------------------------------------------------------- +// 0 4 magic = "TVPI" (bytes 0x54, 0x56, 0x50, 0x49) +// 4 1 version = 1 +// 5 1 bit_width +// 6 2 reserved (zero) +// 8 4 dim (uint32) +// 12 4 n_vectors (uint32) +// 16 N*D*4 vectors (float32, row-major) +// ... N*8 ids (uint64) +// +// Where N = n_vectors and D = dim. There is no checksum in v1; future +// versions may append one without breaking back-compat (version-gated). + +#ifdef __cplusplus +} +#endif diff --git a/ggml/src/CMakeLists.txt b/ggml/src/CMakeLists.txt index 4abc1e301dfe..1ebaa351c9f1 100644 --- a/ggml/src/CMakeLists.txt +++ b/ggml/src/CMakeLists.txt @@ -191,6 +191,7 @@ add_library(ggml-base ../include/ggml-backend.h ../include/ggml-cpp.h ../include/ggml-opt.h + ../include/ggml-vector-index.h ../include/gguf.h ../include/uint8-buff-stream.h ggml.c @@ -204,6 +205,7 @@ add_library(ggml-base ggml-tbq-quants.c ggml-quants.c ggml-quants.h + ggml-vector-index.cpp gguf.cpp uint8-buff-stream.cpp) diff --git a/ggml/src/ggml-vector-index.cpp b/ggml/src/ggml-vector-index.cpp new file mode 100644 index 000000000000..a38076f99ea4 --- /dev/null +++ b/ggml/src/ggml-vector-index.cpp @@ -0,0 +1,562 @@ +// ggml-vector-index.cpp - POC scalar implementation of the fabric vector +// index C API declared in `ggml/include/ggml-vector-index.h`. +// +// Storage: full f32 vectors as a contiguous std::vector. ID map uses +// std::unordered_map for lookup and a parallel vector for +// the slot->id reverse map. Remove uses swap-with-last. +// +// Search: naive scalar dot product across all slots + min-heap of size k. +// No SIMD, no GPU. Correctness over speed; the optimization phase will swap +// the storage layout and kernel without touching the C API. + +#include "ggml-vector-index.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace { + +constexpr uint8_t kTvimMagic[4] = { 'T', 'V', 'P', 'I' }; +constexpr uint8_t kTvimVersion = 1; +constexpr size_t kTvimHeaderSize = 16; + +static_assert(sizeof(float) == sizeof(uint32_t), "ggml-vector-index requires float32"); + +void put_u32_le(uint8_t * dst, uint32_t v) { + dst[0] = static_cast(v >> 0); + dst[1] = static_cast(v >> 8); + dst[2] = static_cast(v >> 16); + dst[3] = static_cast(v >> 24); +} + +void put_u64_le(uint8_t * dst, uint64_t v) { + for (int i = 0; i < 8; ++i) { + dst[i] = static_cast(v >> (8 * i)); + } +} + +uint32_t get_u32_le(const uint8_t * src) { + return (static_cast(src[0]) << 0) | + (static_cast(src[1]) << 8) | + (static_cast(src[2]) << 16) | + (static_cast(src[3]) << 24); +} + +uint64_t get_u64_le(const uint8_t * src) { + uint64_t v = 0; + for (int i = 0; i < 8; ++i) { + v |= static_cast(src[i]) << (8 * i); + } + return v; +} + +uint32_t float_to_u32(float v) { + uint32_t bits; + std::memcpy(&bits, &v, sizeof(bits)); + return bits; +} + +float u32_to_float(uint32_t bits) { + float v; + std::memcpy(&v, &bits, sizeof(v)); + return v; +} + +bool write_u32_le(std::ofstream & f, uint32_t v) { + uint8_t bytes[4]; + put_u32_le(bytes, v); + f.write(reinterpret_cast(bytes), sizeof(bytes)); + return static_cast(f); +} + +bool write_u64_le(std::ofstream & f, uint64_t v) { + uint8_t bytes[8]; + put_u64_le(bytes, v); + f.write(reinterpret_cast(bytes), sizeof(bytes)); + return static_cast(f); +} + +bool read_u32_le(std::ifstream & f, uint32_t & v) { + uint8_t bytes[4]; + f.read(reinterpret_cast(bytes), sizeof(bytes)); + if (!f) { + return false; + } + v = get_u32_le(bytes); + return true; +} + +bool read_u64_le(std::ifstream & f, uint64_t & v) { + uint8_t bytes[8]; + f.read(reinterpret_cast(bytes), sizeof(bytes)); + if (!f) { + return false; + } + v = get_u64_le(bytes); + return true; +} + +// Top-k via min-heap of (score, id). The heap holds at most `k` candidates; +// each new score is compared against the smallest in the heap. +struct ScoreId { + float score; + uint64_t id; +}; + +struct MinHeapCmp { + bool operator()(const ScoreId & a, const ScoreId & b) const { + // Min-heap by score (smallest score at the top). + return a.score > b.score; + } +}; + +} // namespace + +// Lifetime-managed instance state. Lives behind the opaque +// `ggml_vec_index_t` typedef. +struct ggml_vec_index { + int dim = 0; + int bit_width = 32; + + // Flat row-major storage: `data[slot * dim + i]` is component i of vec slot. + std::vector data; + + // slot -> external id (parallel to logical slot index). + std::vector slot_to_id; + + // external id -> slot. + std::unordered_map id_to_slot; +}; + +// --------------------------------------------------------------------------- +// Lifecycle +// --------------------------------------------------------------------------- + +ggml_vec_index_t * ggml_vec_index_create(int dim, int bit_width) { + try { + if (dim <= 0) { + return nullptr; + } + if (bit_width <= 0 || bit_width > 32) { + // POC ignores bit_width but still validates the range so callers + // surface bad config early. + return nullptr; + } + auto * idx = new (std::nothrow) ggml_vec_index(); + if (idx == nullptr) { + return nullptr; + } + idx->dim = dim; + idx->bit_width = bit_width; + return idx; + } catch (...) { + return nullptr; + } +} + +void ggml_vec_index_free(ggml_vec_index_t * idx) { + delete idx; +} + +// --------------------------------------------------------------------------- +// Mutation +// --------------------------------------------------------------------------- + +int ggml_vec_index_add( + ggml_vec_index_t * idx, + const float * vectors, + int n, + const uint64_t * ids) { + + size_t base_slot = 0; + size_t dim_sz = 0; + bool resized = false; + + try { + if (idx == nullptr || vectors == nullptr || ids == nullptr) { + return GGML_VEC_INDEX_E_INVALID_ARG; + } + if (n < 0) { + return GGML_VEC_INDEX_E_INVALID_ARG; + } + if (n == 0) { + return GGML_VEC_INDEX_OK; + } + + // Atomic add: detect duplicates first (against existing AND in-batch), + // bail before mutating any state. + for (int i = 0; i < n; ++i) { + if (idx->id_to_slot.find(ids[i]) != idx->id_to_slot.end()) { + return GGML_VEC_INDEX_E_DUPLICATE; + } + for (int j = i + 1; j < n; ++j) { + if (ids[i] == ids[j]) { + return GGML_VEC_INDEX_E_DUPLICATE; + } + } + } + + base_slot = idx->slot_to_id.size(); + dim_sz = static_cast(idx->dim); + const size_t n_sz = static_cast(n); + if (n_sz > std::numeric_limits::max() - base_slot) { + return GGML_VEC_INDEX_E_INVALID_ARG; + } + const size_t new_slots = base_slot + n_sz; + if (dim_sz != 0 && new_slots > std::numeric_limits::max() / dim_sz) { + return GGML_VEC_INDEX_E_INVALID_ARG; + } + + idx->data.resize(new_slots * dim_sz); + resized = true; + idx->slot_to_id.resize(new_slots); + idx->id_to_slot.reserve(new_slots); + + for (int i = 0; i < n; ++i) { + const size_t slot = base_slot + static_cast(i); + std::memcpy( + idx->data.data() + slot * dim_sz, + vectors + static_cast(i) * dim_sz, + dim_sz * sizeof(float)); + idx->slot_to_id[slot] = ids[i]; + idx->id_to_slot.emplace(ids[i], slot); + } + } catch (const std::bad_alloc &) { + if (idx != nullptr && resized) { + for (int i = 0; i < n; ++i) { + idx->id_to_slot.erase(ids[i]); + } + idx->data.resize(base_slot * dim_sz); + idx->slot_to_id.resize(base_slot); + } + return GGML_VEC_INDEX_E_OOM; + } catch (...) { + if (idx != nullptr && resized) { + for (int i = 0; i < n; ++i) { + idx->id_to_slot.erase(ids[i]); + } + idx->data.resize(base_slot * dim_sz); + idx->slot_to_id.resize(base_slot); + } + return GGML_VEC_INDEX_E_INTERNAL; + } + return GGML_VEC_INDEX_OK; +} + +int ggml_vec_index_remove(ggml_vec_index_t * idx, uint64_t id) { + try { + if (idx == nullptr) { + return GGML_VEC_INDEX_E_INVALID_ARG; + } + auto it = idx->id_to_slot.find(id); + if (it == idx->id_to_slot.end()) { + return 0; + } + const size_t slot = it->second; + const size_t last = idx->slot_to_id.size() - 1; + const size_t dim_sz = static_cast(idx->dim); + + if (slot != last) { + // Move last vector into the freed slot and update its id mapping. + std::memcpy( + idx->data.data() + slot * dim_sz, + idx->data.data() + last * dim_sz, + dim_sz * sizeof(float)); + const uint64_t moved_id = idx->slot_to_id[last]; + idx->slot_to_id[slot] = moved_id; + idx->id_to_slot[moved_id] = slot; + } + + idx->slot_to_id.pop_back(); + idx->data.resize(last * dim_sz); + idx->id_to_slot.erase(it); + return 1; + } catch (...) { + return GGML_VEC_INDEX_E_INTERNAL; + } +} + +int ggml_vec_index_contains(const ggml_vec_index_t * idx, uint64_t id) { + if (idx == nullptr) { + return 0; + } + return idx->id_to_slot.count(id) != 0 ? 1 : 0; +} + +void ggml_vec_index_prepare(ggml_vec_index_t * /*idx*/) { + // POC no-op. Future: warm caches, materialize codebooks, etc. +} + +// --------------------------------------------------------------------------- +// Search +// --------------------------------------------------------------------------- + +namespace { + +// Scalar dot product of two `dim`-length f32 vectors. +inline float dot(const float * a, const float * b, int dim) { + float acc = 0.0f; + for (int i = 0; i < dim; ++i) { + acc += a[i] * b[i]; + } + return acc; +} + +// Run a single query against all slots, write top-k into out_scores/out_ids. +// If the index holds fewer than k entries, pad with sentinels. +void search_one( + const ggml_vec_index_t & idx, + const float * query, + int k, + float * out_scores, + uint64_t * out_ids) { + + const int dim = idx.dim; + const size_t n_slots = idx.slot_to_id.size(); + + std::priority_queue, MinHeapCmp> heap; + + for (size_t slot = 0; slot < n_slots; ++slot) { + const float s = dot( + query, idx.data.data() + slot * static_cast(dim), dim); + if (heap.size() < static_cast(k)) { + heap.push({ s, idx.slot_to_id[slot] }); + } else if (s > heap.top().score) { + heap.pop(); + heap.push({ s, idx.slot_to_id[slot] }); + } + } + + // Drain the heap into a temporary descending list. + std::vector drained; + drained.reserve(heap.size()); + while (!heap.empty()) { + drained.push_back(heap.top()); + heap.pop(); + } + std::reverse(drained.begin(), drained.end()); // now descending by score + + for (int i = 0; i < k; ++i) { + if (static_cast(i) < drained.size()) { + out_scores[i] = drained[i].score; + out_ids[i] = drained[i].id; + } else { + out_scores[i] = -FLT_MAX; + out_ids[i] = UINT64_MAX; + } + } +} + +} // namespace + +int ggml_vec_index_search( + const ggml_vec_index_t * idx, + const float * queries, + int n_q, + int k, + float * out_scores, + uint64_t * out_ids) { + + if (idx == nullptr || queries == nullptr || + out_scores == nullptr || out_ids == nullptr) { + return GGML_VEC_INDEX_E_INVALID_ARG; + } + if (n_q < 0 || k <= 0) { + return GGML_VEC_INDEX_E_INVALID_ARG; + } + if (n_q == 0) { + return GGML_VEC_INDEX_OK; + } + + try { + const int dim = idx->dim; + const size_t n_q_sz = static_cast(n_q); + const size_t k_sz = static_cast(k); + const size_t dim_sz = static_cast(dim); + if ((dim_sz != 0 && n_q_sz > std::numeric_limits::max() / dim_sz) || + n_q_sz > std::numeric_limits::max() / k_sz) { + return GGML_VEC_INDEX_E_INVALID_ARG; + } + + for (int q = 0; q < n_q; ++q) { + search_one( + *idx, + queries + static_cast(q) * static_cast(dim), + k, + out_scores + static_cast(q) * static_cast(k), + out_ids + static_cast(q) * static_cast(k)); + } + } catch (const std::bad_alloc &) { + return GGML_VEC_INDEX_E_OOM; + } catch (...) { + return GGML_VEC_INDEX_E_INTERNAL; + } + return GGML_VEC_INDEX_OK; +} + +// --------------------------------------------------------------------------- +// Persistence +// --------------------------------------------------------------------------- + +int ggml_vec_index_write(ggml_vec_index_t * idx, const char * path) { + try { + if (idx == nullptr || path == nullptr) { + return GGML_VEC_INDEX_E_INVALID_ARG; + } + if (idx->slot_to_id.size() > std::numeric_limits::max()) { + return GGML_VEC_INDEX_E_INVALID_ARG; + } + const size_t n = idx->slot_to_id.size(); + const size_t dim_sz = static_cast(idx->dim); + if (dim_sz != 0 && n > std::numeric_limits::max() / dim_sz) { + return GGML_VEC_INDEX_E_INTERNAL; + } + if (idx->data.size() != n * dim_sz) { + return GGML_VEC_INDEX_E_INTERNAL; + } + + std::ofstream f(path, std::ios::binary | std::ios::trunc); + if (!f.is_open()) { + return GGML_VEC_INDEX_E_IO; + } + + // Header: 16 bytes. Layout matches the comment block in the header file. + uint8_t header[kTvimHeaderSize] = {}; + std::memcpy(header, kTvimMagic, 4); + header[4] = kTvimVersion; + header[5] = static_cast(idx->bit_width); + header[6] = 0; + header[7] = 0; + const uint32_t dim_le = static_cast(idx->dim); + const uint32_t n_le = static_cast(idx->slot_to_id.size()); + put_u32_le(header + 8, dim_le); + put_u32_le(header + 12, n_le); + + f.write(reinterpret_cast(header), sizeof(header)); + if (!f) { return GGML_VEC_INDEX_E_IO; } + + for (float v : idx->data) { + if (!write_u32_le(f, float_to_u32(v))) { + return GGML_VEC_INDEX_E_IO; + } + } + + for (uint64_t id : idx->slot_to_id) { + if (!write_u64_le(f, id)) { + return GGML_VEC_INDEX_E_IO; + } + } + + f.flush(); + if (!f) { return GGML_VEC_INDEX_E_IO; } + return GGML_VEC_INDEX_OK; + } catch (const std::bad_alloc &) { + return GGML_VEC_INDEX_E_OOM; + } catch (...) { + return GGML_VEC_INDEX_E_INTERNAL; + } +} + +ggml_vec_index_t * ggml_vec_index_load(const char * path) { + try { + if (path == nullptr) { + return nullptr; + } + std::ifstream f(path, std::ios::binary); + if (!f.is_open()) { + return nullptr; + } + + uint8_t header[kTvimHeaderSize] = {}; + f.read(reinterpret_cast(header), sizeof(header)); + if (!f || f.gcount() != static_cast(sizeof(header))) { + return nullptr; + } + if (std::memcmp(header, kTvimMagic, 4) != 0) { + return nullptr; + } + if (header[4] != kTvimVersion) { + return nullptr; + } + + const int bit_width = static_cast(header[5]); + const uint32_t dim_le = get_u32_le(header + 8); + const uint32_t n_le = get_u32_le(header + 12); + if (dim_le == 0 || dim_le > static_cast(std::numeric_limits::max())) { + return nullptr; + } + const int dim = static_cast(dim_le); + + std::unique_ptr idx( + ggml_vec_index_create(dim, bit_width), + ggml_vec_index_free); + if (idx == nullptr) { + return nullptr; + } + const size_t dim_sz = static_cast(dim); + const size_t n = static_cast(n_le); + if (n != 0 && dim_sz > std::numeric_limits::max() / n) { + return nullptr; + } + + idx->data.resize(n * dim_sz); + idx->slot_to_id.resize(n); + idx->id_to_slot.reserve(n); + + for (float & v : idx->data) { + uint32_t bits = 0; + if (!read_u32_le(f, bits)) { + return nullptr; + } + v = u32_to_float(bits); + } + + for (uint64_t & id : idx->slot_to_id) { + if (!read_u64_le(f, id)) { + return nullptr; + } + } + + for (size_t slot = 0; slot < n; ++slot) { + const uint64_t id = idx->slot_to_id[slot]; + const bool inserted = + idx->id_to_slot.emplace(id, slot).second; + if (!inserted) { + // Duplicate id in persisted file: corrupted. + return nullptr; + } + } + + return idx.release(); + } catch (...) { + return nullptr; + } +} + +// --------------------------------------------------------------------------- +// Stats +// --------------------------------------------------------------------------- + +int ggml_vec_index_len(const ggml_vec_index_t * idx) { + return idx ? static_cast(idx->slot_to_id.size()) : 0; +} + +int ggml_vec_index_dim(const ggml_vec_index_t * idx) { + return idx ? idx->dim : 0; +} + +int ggml_vec_index_bit_width(const ggml_vec_index_t * idx) { + return idx ? idx->bit_width : 0; +} diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index b61f1581cfb1..e75fe7be1fd2 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -144,6 +144,15 @@ if (NOT WIN32) ) endif() +# POC vector-index smoke test: no model required, no llama runtime touched. +add_executable(test-vector-index test-vector-index.cpp) +target_compile_features(test-vector-index PRIVATE cxx_std_17) +target_link_libraries(test-vector-index PRIVATE ggml) +if (LLAMA_TESTS_INSTALL) + install(TARGETS test-vector-index RUNTIME) +endif() +llama_test(test-vector-index) + if (LLAMA_LLGUIDANCE) llama_build_and_test(test-grammar-llguidance.cpp ARGS ${PROJECT_SOURCE_DIR}/models/ggml-vocab-llama-bpe.gguf) endif () diff --git a/tests/test-vector-index.cpp b/tests/test-vector-index.cpp new file mode 100644 index 000000000000..bc633b98f668 --- /dev/null +++ b/tests/test-vector-index.cpp @@ -0,0 +1,175 @@ +// test-vector-index.cpp - standalone C-API smoke test for the POC vector +// index. Exercises lifecycle, add, search, remove, contains, write, load, +// search-after-load. No model, no llama; only the new ggml-vector-index +// public C API. + +#include "ggml-vector-index.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace { + +constexpr int kDim = 4; + +#define CHECK(cond) \ + do { \ + if (!(cond)) { \ + std::fprintf(stderr, "FAIL %s:%d: %s\n", __FILE__, __LINE__, #cond);\ + std::exit(1); \ + } \ + } while (0) + +std::vector normalize(std::vector v) { + double sumsq = 0.0; + for (float x : v) sumsq += static_cast(x) * x; + const float n = static_cast(std::sqrt(sumsq)); + if (n > 0.0f) for (float & x : v) x /= n; + return v; +} + +} // namespace + +int main() { + auto * idx = ggml_vec_index_create(kDim, /*bit_width=*/32); + CHECK(idx != nullptr); + CHECK(ggml_vec_index_dim(idx) == kDim); + CHECK(ggml_vec_index_len(idx) == 0); + CHECK(ggml_vec_index_bit_width(idx) == 32); + + // Add 4 well-separated unit vectors. IDs are non-trivial uint64 to + // catch sign-extension or BigInt round-trip bugs at the JS boundary + // when this codepath is later exercised from Bare. + std::vector vecs; + std::vector ids = { + 42ULL, + (1ULL << 40) + 7ULL, + (1ULL << 62) + 11ULL, + UINT64_MAX - 13ULL, + }; + std::vector> seeds = { + normalize({1.0f, 0.0f, 0.0f, 0.0f}), + normalize({0.0f, 1.0f, 0.0f, 0.0f}), + normalize({0.0f, 0.0f, 1.0f, 0.0f}), + normalize({0.0f, 0.0f, 0.0f, 1.0f}), + }; + for (const auto & s : seeds) { + vecs.insert(vecs.end(), s.begin(), s.end()); + } + CHECK(ggml_vec_index_add( + idx, vecs.data(), static_cast(ids.size()), ids.data()) == 0); + CHECK(ggml_vec_index_len(idx) == 4); + CHECK(ggml_vec_index_contains(idx, ids[0]) == 1); + CHECK(ggml_vec_index_contains(idx, 999ULL) == 0); + + // Duplicate add must fail without mutating state. + { + const std::vector dup_ids = { ids[0] }; + std::vector dup_vec(seeds[0]); + CHECK(ggml_vec_index_add(idx, dup_vec.data(), 1, dup_ids.data()) + == GGML_VEC_INDEX_E_DUPLICATE); + CHECK(ggml_vec_index_len(idx) == 4); + } + + // In-batch duplicate ids must also fail atomically. + { + const uint64_t new_id = (1ULL << 50) + 123ULL; + const std::vector dup_ids = { new_id, new_id }; + std::vector dup_vecs; + dup_vecs.insert(dup_vecs.end(), seeds[0].begin(), seeds[0].end()); + dup_vecs.insert(dup_vecs.end(), seeds[1].begin(), seeds[1].end()); + CHECK(ggml_vec_index_add(idx, dup_vecs.data(), 2, dup_ids.data()) + == GGML_VEC_INDEX_E_DUPLICATE); + CHECK(ggml_vec_index_len(idx) == 4); + CHECK(ggml_vec_index_contains(idx, new_id) == 0); + } + + // Top-1 of querying with each unit vector should retrieve itself with + // score very close to 1.0 (full f32, no quantization noise). + { + std::array scores{}; + std::array out_ids{}; + for (size_t i = 0; i < seeds.size(); ++i) { + CHECK(ggml_vec_index_search( + idx, seeds[i].data(), 1, /*k=*/1, + scores.data(), out_ids.data()) == 0); + CHECK(out_ids[0] == ids[i]); + CHECK(std::fabs(scores[0] - 1.0f) < 1e-5f); + } + } + + // Top-k > len returns sentinel-padded tail. + { + std::array scores{}; + std::array out_ids{}; + CHECK(ggml_vec_index_search( + idx, seeds[0].data(), 1, /*k=*/8, + scores.data(), out_ids.data()) == 0); + CHECK(out_ids[0] == ids[0]); + // Tail entries (positions 4..7) use sentinel score/id values. + for (int i = 4; i < 8; ++i) { + CHECK(scores[i] == -FLT_MAX); + CHECK(out_ids[i] == UINT64_MAX); + } + } + + // Remove + search: the removed id must no longer surface. + { + CHECK(ggml_vec_index_remove(idx, ids[1]) == 1); + CHECK(ggml_vec_index_remove(idx, ids[1]) == 0); // already gone + CHECK(ggml_vec_index_len(idx) == 3); + CHECK(ggml_vec_index_contains(idx, ids[1]) == 0); + + std::array scores{}; + std::array out_ids{}; + CHECK(ggml_vec_index_search( + idx, seeds[1].data(), 1, /*k=*/3, + scores.data(), out_ids.data()) == 0); + for (int i = 0; i < 3; ++i) { + CHECK(out_ids[i] != ids[1]); + } + } + + // Persistence round-trip: write, free, load, re-query. + const auto tmp = std::filesystem::temp_directory_path() / + "ggml-vector-index-test.tvim"; + const std::string path = tmp.string(); + CHECK(ggml_vec_index_write(idx, path.c_str()) == 0); + + ggml_vec_index_free(idx); + + auto * loaded = ggml_vec_index_load(path.c_str()); + CHECK(loaded != nullptr); + CHECK(ggml_vec_index_dim(loaded) == kDim); + CHECK(ggml_vec_index_len(loaded) == 3); + CHECK(ggml_vec_index_bit_width(loaded) == 32); + CHECK(ggml_vec_index_contains(loaded, ids[0]) == 1); + CHECK(ggml_vec_index_contains(loaded, ids[1]) == 0); // stayed deleted + CHECK(ggml_vec_index_contains(loaded, ids[2]) == 1); + CHECK(ggml_vec_index_contains(loaded, ids[3]) == 1); + + // Top-1 self-match after reload. + { + std::array scores{}; + std::array out_ids{}; + CHECK(ggml_vec_index_search( + loaded, seeds[0].data(), 1, /*k=*/1, + scores.data(), out_ids.data()) == 0); + CHECK(out_ids[0] == ids[0]); + CHECK(std::fabs(scores[0] - 1.0f) < 1e-5f); + } + + ggml_vec_index_free(loaded); + std::filesystem::remove(path); + + std::printf("test-vector-index: OK\n"); + return 0; +} From a9badb70eb25fdc910159bd7fba9d1784d89e041 Mon Sep 17 00:00:00 2001 From: Nidhin Date: Fri, 10 Jul 2026 20:51:24 +0530 Subject: [PATCH 02/13] ggml : add production vector index persistence and q8 search Adds the fabric vector index API with f32/q8 storage, NEON-accelerated q8 search, crash-safe checksummed .tvim persistence, and regression coverage for corruption and migration paths. Assisted-by: GPT-5.5 --- ggml/include/ggml-vector-index.h | 78 +-- ggml/src/ggml-vector-index.cpp | 842 ++++++++++++++++++++++++++++--- tests/CMakeLists.txt | 9 +- tests/bench-vector-index.cpp | 217 ++++++++ tests/test-vector-index.cpp | 422 +++++++++++++++- 5 files changed, 1459 insertions(+), 109 deletions(-) create mode 100644 tests/bench-vector-index.cpp diff --git a/ggml/include/ggml-vector-index.h b/ggml/include/ggml-vector-index.h index 217c58d74696..75ea28f8241a 100644 --- a/ggml/include/ggml-vector-index.h +++ b/ggml/include/ggml-vector-index.h @@ -1,22 +1,16 @@ #pragma once // -// ggml-vector-index: TurboQuant-style ANN vector-index C API. +// ggml-vector-index: TurboVec-style vector-index C API. // -// POC NOTE -// -------- // This is the public C API for fabric's vector index. The implementation -// under `ggml/src/ggml-vector-index.cpp` is intentionally naive (full f32 -// storage, scalar dot-product, min-heap top-k). The C API itself is final: -// downstream Bare addon bindings + JS wrappers depend on this signature and -// the on-disk file format documented at the bottom of this header. -// Quantization / SIMD / GPU kernels are future work and will be swapped in -// behind the same API without breaking callers. +// under `ggml/src/ggml-vector-index.cpp` supports full f32 storage +// (`bit_width=32`) and production q8 storage (`bit_width=8`) with CPU search +// directly against quantized codes. ARM builds use NEON when available. // // Threading: instances are NOT thread-safe. Callers must serialize access // to a given handle. Multiple handles can be used concurrently. // -// Endianness: persistence format is fixed little-endian. The POC asserts a -// LE host; we do not currently support big-endian platforms. +// Endianness: persistence format is fixed little-endian. #include "ggml.h" @@ -46,9 +40,9 @@ enum ggml_vec_index_error { // Lifecycle. // -// `dim` must be > 0. `bit_width` is reserved for future quantization; the -// POC accepts any value in [1, 32] but ignores it (storage is always full -// f32). Returns NULL on bad args. +// `dim` must be > 0. `bit_width` must be 8 or 32. `bit_width=8` stores +// per-vector symmetric q8 codes with one f32 scale per vector. `bit_width=32` +// stores full f32 vectors. Returns NULL on bad args. GGML_API ggml_vec_index_t * ggml_vec_index_create(int dim, int bit_width); GGML_API void ggml_vec_index_free(ggml_vec_index_t * idx); @@ -59,6 +53,7 @@ GGML_API void ggml_vec_index_free(ggml_vec_index_t * idx); // associating each with the corresponding `ids[i]` (caller-owned external id). // Returns 0 on success. Returns GGML_VEC_INDEX_E_DUPLICATE if any id already // exists in the index; in that case the index is unchanged (atomic add). +// All vector components must be finite. GGML_API int ggml_vec_index_add( ggml_vec_index_t * idx, const float * vectors, @@ -73,9 +68,8 @@ GGML_API int ggml_vec_index_remove(ggml_vec_index_t * idx, uint64_t id); // Returns 1 if the id is in the index, 0 otherwise. Read-only. GGML_API int ggml_vec_index_contains(const ggml_vec_index_t * idx, uint64_t id); -// No-op for the POC. Placeholder for future cache warming / codebook -// resolution after a bulk add. Reserved as a mutating op (warm-up may -// materialize derived state inside the index). +// Placeholder for cache warming / codebook resolution after a bulk add. +// Currently a no-op. GGML_API void ggml_vec_index_prepare(ggml_vec_index_t * idx); // Top-k search. `queries` is `n_q * dim` row-major. `out_scores` and @@ -85,9 +79,12 @@ GGML_API void ggml_vec_index_prepare(ggml_vec_index_t * idx); // with sentinel values: -FLT_MAX for scores, UINT64_MAX for ids. Read-only // against the index (does not mutate state). // -// Score semantics: scalar full-precision dot product. Callers that want -// cosine similarity must L2-normalize their vectors before insert AND -// before query; the index does NOT normalize internally. +// Score semantics: dot product. For f32 storage this is a full-precision dot +// product. For q8 storage, the query remains f32 and each indexed component is +// scored as `q8_code * per_vector_scale` without expanding the stored matrix +// back to f32. Callers that want cosine similarity must L2-normalize their +// vectors before insert AND before query; the index does NOT normalize +// internally. All query components must be finite. GGML_API int ggml_vec_index_search( const ggml_vec_index_t * idx, const float * queries, @@ -96,12 +93,14 @@ GGML_API int ggml_vec_index_search( float * out_scores, uint64_t * out_ids); -// Persistence. Format is .tvim version 1; see bottom of this header. +// Persistence. Format is .tvim version 2; see bottom of this header. GGML_API int ggml_vec_index_write( ggml_vec_index_t * idx, const char * path); -// Returns NULL on failure (caller can inspect errno for I/O specifics). +// Loads v2 files and migrates v1 f32 snapshots. Legacy bit_width=8 snapshots +// are quantized to q8; all other legacy bit widths migrate to f32/32-bit. +// Returns NULL on failure. GGML_API ggml_vec_index_t * ggml_vec_index_load(const char * path); // Stats. @@ -109,21 +108,42 @@ GGML_API int ggml_vec_index_len(const ggml_vec_index_t * idx); GGML_API int ggml_vec_index_dim(const ggml_vec_index_t * idx); GGML_API int ggml_vec_index_bit_width(const ggml_vec_index_t * idx); -// File format (.tvim version 1, all little-endian): +// File format (.tvim version 2, all little-endian): // // offset size field // ------ ----- ------------------------------------------------------- // 0 4 magic = "TVPI" (bytes 0x54, 0x56, 0x50, 0x49) -// 4 1 version = 1 -// 5 1 bit_width -// 6 2 reserved (zero) +// 4 1 version = 2 +// 5 1 bit_width (8 or 32) +// 6 1 storage kind (1 = f32, 2 = q8) +// 7 1 flags (bit 0 = checksum trailer present) // 8 4 dim (uint32) // 12 4 n_vectors (uint32) -// 16 N*D*4 vectors (float32, row-major) +// 16 4 qparam_type (0 = none, 1 = per-vector f32 scale) +// 20 4 qparam_bytes_per_vector (0 or 4) +// 24 4 bytes_per_component (1 or 4) +// 28 4 reserved (zero) +// 32 ... qparams: +// - f32: empty +// - q8: N float32 scales +// ... ... vectors: +// - f32: N*D float32 values, row-major +// - q8: N*D int8 codes, row-major // ... N*8 ids (uint64) +// ... 4 header CRC32C, when flag bit 0 is set +// ... 4 qparams CRC32C, when flag bit 0 is set +// ... 4 vectors CRC32C, when flag bit 0 is set +// ... 4 ids CRC32C, when flag bit 0 is set // -// Where N = n_vectors and D = dim. There is no checksum in v1; future -// versions may append one without breaking back-compat (version-gated). +// Where N = n_vectors and D = dim. q8 uses symmetric per-vector quantization: +// scale = max(abs(v)) / 127, code = round(v / scale) clamped to [-127, 127]. +// Zero vectors use scale = 1 and all-zero codes. Each CRC32C covers exactly its +// corresponding serialized section; the header CRC covers bytes [0, 32), and +// the CRC32C of an empty section is zero. +// Legacy v2 files with flags=0 and no checksum trailer remain readable. +// Writers emit checksummed v2 files. Readers reject unknown versions and v2 +// flag bits; they also accept legacy v1 f32 snapshots. Legacy bit_width=8 +// snapshots migrate to q8, while all other legacy widths migrate to f32. #ifdef __cplusplus } diff --git a/ggml/src/ggml-vector-index.cpp b/ggml/src/ggml-vector-index.cpp index a38076f99ea4..f18370c09f7d 100644 --- a/ggml/src/ggml-vector-index.cpp +++ b/ggml/src/ggml-vector-index.cpp @@ -1,36 +1,72 @@ -// ggml-vector-index.cpp - POC scalar implementation of the fabric vector +// ggml-vector-index.cpp - CPU implementation of the fabric vector // index C API declared in `ggml/include/ggml-vector-index.h`. // -// Storage: full f32 vectors as a contiguous std::vector. ID map uses +// Storage: full f32 vectors or per-vector symmetric q8 codes. ID map uses // std::unordered_map for lookup and a parallel vector for // the slot->id reverse map. Remove uses swap-with-last. // -// Search: naive scalar dot product across all slots + min-heap of size k. -// No SIMD, no GPU. Correctness over speed; the optimization phase will swap -// the storage layout and kernel without touching the C API. +// Search: dot product across all slots + min-heap of size k. q8 search scores +// directly against stored codes and per-vector scales, with ARM NEON when +// available and a scalar fallback. #include "ggml-vector-index.h" #include +#include #include #include #include +#include #include #include +#include #include +#include #include #include #include #include #include +#include #include +#include #include +#if defined(__ARM_NEON) || defined(__ARM_NEON__) +#include +#define GGML_VEC_INDEX_USE_NEON 1 +#else +#define GGML_VEC_INDEX_USE_NEON 0 +#endif + +#ifndef _WIN32 +#include +#include +#include +#else +#ifndef NOMINMAX +#define NOMINMAX +#endif +#include +#include +#include +#include +#include +#endif + namespace { constexpr uint8_t kTvimMagic[4] = { 'T', 'V', 'P', 'I' }; -constexpr uint8_t kTvimVersion = 1; -constexpr size_t kTvimHeaderSize = 16; +constexpr uint8_t kTvimVersionV1 = 1; +constexpr uint8_t kTvimVersion = 2; +constexpr uint8_t kStorageF32 = 1; +constexpr uint8_t kStorageQ8 = 2; +constexpr uint8_t kFlagCRC32C = 1; +constexpr uint32_t kQParamNone = 0; +constexpr uint32_t kQParamScaleF32 = 1; +constexpr size_t kTvimV1HeaderSize = 16; +constexpr size_t kTvimHeaderSize = 32; +constexpr size_t kTvimChecksumSize = 16; static_assert(sizeof(float) == sizeof(uint32_t), "ggml-vector-index requires float32"); @@ -74,37 +110,82 @@ float u32_to_float(uint32_t bits) { return v; } -bool write_u32_le(std::ofstream & f, uint32_t v) { +bool write_u32_le(std::FILE * f, uint32_t v) { uint8_t bytes[4]; put_u32_le(bytes, v); - f.write(reinterpret_cast(bytes), sizeof(bytes)); - return static_cast(f); + return std::fwrite(bytes, 1, sizeof(bytes), f) == sizeof(bytes); } -bool write_u64_le(std::ofstream & f, uint64_t v) { +bool read_u32_le(std::ifstream & f, uint32_t & v) { + uint8_t bytes[4]; + f.read(reinterpret_cast(bytes), sizeof(bytes)); + if (!f) { + return false; + } + v = get_u32_le(bytes); + return true; +} + +bool read_u64_le(std::ifstream & f, uint64_t & v) { uint8_t bytes[8]; - put_u64_le(bytes, v); - f.write(reinterpret_cast(bytes), sizeof(bytes)); - return static_cast(f); + f.read(reinterpret_cast(bytes), sizeof(bytes)); + if (!f) { + return false; + } + v = get_u64_le(bytes); + return true; } -bool read_u32_le(std::ifstream & f, uint32_t & v) { +uint32_t crc32c_update(uint32_t crc, const void * data, size_t size) { + const auto * bytes = static_cast(data); + for (size_t i = 0; i < size; ++i) { + crc ^= bytes[i]; + for (int bit = 0; bit < 8; ++bit) { + crc = (crc >> 1) ^ (0x82f63b78u & (0u - (crc & 1u))); + } + } + return crc; +} + +bool read_u32_le_crc(std::ifstream & f, uint32_t & v, uint32_t & crc) { uint8_t bytes[4]; f.read(reinterpret_cast(bytes), sizeof(bytes)); if (!f) { return false; } v = get_u32_le(bytes); + crc = crc32c_update(crc, bytes, sizeof(bytes)); return true; } -bool read_u64_le(std::ifstream & f, uint64_t & v) { +bool read_u64_le_crc(std::ifstream & f, uint64_t & v, uint32_t & crc) { uint8_t bytes[8]; f.read(reinterpret_cast(bytes), sizeof(bytes)); if (!f) { return false; } v = get_u64_le(bytes); + crc = crc32c_update(crc, bytes, sizeof(bytes)); + return true; +} + +bool write_u32_le_crc(std::FILE * f, uint32_t v, uint32_t & crc) { + uint8_t bytes[4]; + put_u32_le(bytes, v); + if (std::fwrite(bytes, 1, sizeof(bytes), f) != sizeof(bytes)) { + return false; + } + crc = crc32c_update(crc, bytes, sizeof(bytes)); + return true; +} + +bool write_u64_le_crc(std::FILE * f, uint64_t v, uint32_t & crc) { + uint8_t bytes[8]; + put_u64_le(bytes, v); + if (std::fwrite(bytes, 1, sizeof(bytes), f) != sizeof(bytes)) { + return false; + } + crc = crc32c_update(crc, bytes, sizeof(bytes)); return true; } @@ -122,6 +203,235 @@ struct MinHeapCmp { } }; +bool is_supported_bit_width(int bit_width) { + return bit_width == 8 || bit_width == 32; +} + +bool all_finite(const float * values, size_t n) { + for (size_t i = 0; i < n; ++i) { + if (!std::isfinite(values[i])) { + return false; + } + } + return true; +} + +bool checked_mul_u64(uint64_t a, uint64_t b, uint64_t & out) { + if (a != 0 && b > std::numeric_limits::max() / a) { + return false; + } + out = a * b; + return true; +} + +bool checked_add_u64(uint64_t a, uint64_t b, uint64_t & out) { + if (b > std::numeric_limits::max() - a) { + return false; + } + out = a + b; + return true; +} + +bool expected_file_size( + uint64_t header_size, + uint64_t n, + uint64_t dim, + uint64_t qparam_bytes, + uint64_t component_bytes, + uint64_t & size) { + uint64_t qparams = 0; + uint64_t components = 0; + uint64_t ids = 0; + uint64_t total = header_size; + if (!checked_mul_u64(n, qparam_bytes, qparams) || + !checked_mul_u64(n, dim, components) || + !checked_mul_u64(components, component_bytes, components) || + !checked_mul_u64(n, sizeof(uint64_t), ids) || + !checked_add_u64(total, qparams, total) || + !checked_add_u64(total, components, total) || + !checked_add_u64(total, ids, total)) { + return false; + } + size = total; + return true; +} + +struct TempFile { + std::FILE * stream = nullptr; +#ifdef _WIN32 + std::wstring path; +#else + std::string path; +#endif + + TempFile() = default; + ~TempFile(); + TempFile(const TempFile &) = delete; + TempFile & operator=(const TempFile &) = delete; +}; + +#ifdef _WIN32 + +bool utf8_to_wide(const char * src, std::wstring & dst) { + const int size = MultiByteToWideChar( + CP_UTF8, MB_ERR_INVALID_CHARS, src, -1, nullptr, 0); + if (size <= 0) { + return false; + } + std::vector buffer(static_cast(size)); + if (MultiByteToWideChar( + CP_UTF8, MB_ERR_INVALID_CHARS, src, -1, buffer.data(), size) != size) { + return false; + } + dst.assign(buffer.data()); + return true; +} + +bool open_temp_file(const char * path, TempFile & temp) { + std::wstring destination; + if (!utf8_to_wide(path, destination)) { + return false; + } + static std::atomic sequence{ 0 }; + for (int attempt = 0; attempt < 128; ++attempt) { + wchar_t suffix[96]; + std::swprintf( + suffix, + sizeof(suffix) / sizeof(suffix[0]), + L".tmp.%lu.%lu.%llu", + static_cast(GetCurrentProcessId()), + static_cast(GetCurrentThreadId()), + static_cast(sequence.fetch_add(1))); + temp.path = destination + suffix; + + int fd = -1; + const errno_t error = _wsopen_s( + &fd, + temp.path.c_str(), + _O_CREAT | _O_EXCL | _O_RDWR | _O_BINARY | _O_NOINHERIT, + _SH_DENYRW, + _S_IREAD | _S_IWRITE); + if (error == 0) { + temp.stream = _wfdopen(fd, L"w+b"); + if (temp.stream != nullptr) { + return true; + } + _close(fd); + _wremove(temp.path.c_str()); + return false; + } + if (error != EEXIST) { + return false; + } + } + return false; +} + +#else + +bool open_temp_file(const char * path, TempFile & temp) { + std::string pattern = std::string(path) + ".tmp.XXXXXX"; + std::vector mutable_path(pattern.begin(), pattern.end()); + mutable_path.push_back('\0'); + const int fd = ::mkstemp(mutable_path.data()); + if (fd < 0) { + return false; + } + temp.path.assign(mutable_path.data()); + + struct stat destination_stat; + if (::stat(path, &destination_stat) == 0 && + ::fchmod(fd, destination_stat.st_mode & 0777) != 0) { + ::close(fd); + std::remove(temp.path.c_str()); + return false; + } + + temp.stream = ::fdopen(fd, "w+b"); + if (temp.stream == nullptr) { + ::close(fd); + std::remove(temp.path.c_str()); + return false; + } + return true; +} + +#endif + +void remove_temp_file(TempFile & temp) { + if (temp.path.empty()) { + return; + } +#ifdef _WIN32 + _wremove(temp.path.c_str()); +#else + std::remove(temp.path.c_str()); +#endif + temp.path.clear(); +} + +TempFile::~TempFile() { + if (stream != nullptr) { + std::fclose(stream); + } + remove_temp_file(*this); +} + +bool flush_and_sync(std::FILE * stream) { + if (std::fflush(stream) != 0) { + return false; + } +#ifdef _WIN32 + return _commit(_fileno(stream)) == 0; +#elif defined(__APPLE__) + const int fd = ::fileno(stream); + return ::fcntl(fd, F_FULLFSYNC) == 0 || ::fsync(fd) == 0; +#else + return ::fsync(::fileno(stream)) == 0; +#endif +} + +bool fsync_parent_dir(const char * path) { +#ifdef _WIN32 + (void) path; + return true; +#else + std::filesystem::path parent = std::filesystem::path(path).parent_path(); + if (parent.empty()) { + parent = "."; + } + const int fd = ::open(parent.c_str(), O_RDONLY); + if (fd < 0) { + return false; + } + const bool ok = ::fsync(fd) == 0; + ::close(fd); + return ok; +#endif +} + +bool rename_overwrite(const TempFile & temp, const char * dst) { +#ifdef _WIN32 + std::wstring destination; + if (!utf8_to_wide(dst, destination)) { + return false; + } + if (GetFileAttributesW(destination.c_str()) != INVALID_FILE_ATTRIBUTES) { + return ReplaceFileW( + destination.c_str(), + temp.path.c_str(), + nullptr, + REPLACEFILE_WRITE_THROUGH, + nullptr, + nullptr) != 0; + } + return MoveFileExW( + temp.path.c_str(), destination.c_str(), MOVEFILE_WRITE_THROUGH) != 0; +#else + return std::rename(temp.path.c_str(), dst) == 0; +#endif +} + } // namespace // Lifetime-managed instance state. Lives behind the opaque @@ -130,9 +440,13 @@ struct ggml_vec_index { int dim = 0; int bit_width = 32; - // Flat row-major storage: `data[slot * dim + i]` is component i of vec slot. + // Flat row-major f32 storage for bit_width=32. std::vector data; + // Flat row-major q8 storage for bit_width=8 plus one scale per vector. + std::vector q8_data; + std::vector q8_scale; + // slot -> external id (parallel to logical slot index). std::vector slot_to_id; @@ -140,6 +454,38 @@ struct ggml_vec_index { std::unordered_map id_to_slot; }; +static bool is_q8(const ggml_vec_index & idx) { + return idx.bit_width == 8; +} + +static uint8_t storage_kind(const ggml_vec_index & idx) { + return is_q8(idx) ? kStorageQ8 : kStorageF32; +} + +static void quantize_q8_row(const float * src, int8_t * dst, int dim, float & scale) { + float max_abs = 0.0f; + for (int i = 0; i < dim; ++i) { + max_abs = std::max(max_abs, std::fabs(src[i])); + } + + if (max_abs == 0.0f) { + scale = 1.0f; + std::memset(dst, 0, static_cast(dim)); + return; + } + + scale = max_abs / 127.0f; + if (scale == 0.0f) { + scale = max_abs; + } + for (int i = 0; i < dim; ++i) { + const float scaled = src[i] / scale; + int q = static_cast(std::nearbyint(scaled)); + q = std::max(-127, std::min(127, q)); + dst[i] = static_cast(q); + } +} + // --------------------------------------------------------------------------- // Lifecycle // --------------------------------------------------------------------------- @@ -149,9 +495,7 @@ ggml_vec_index_t * ggml_vec_index_create(int dim, int bit_width) { if (dim <= 0) { return nullptr; } - if (bit_width <= 0 || bit_width > 32) { - // POC ignores bit_width but still validates the range so callers - // surface bad config early. + if (!is_supported_bit_width(bit_width)) { return nullptr; } auto * idx = new (std::nothrow) ggml_vec_index(); @@ -184,6 +528,22 @@ int ggml_vec_index_add( size_t dim_sz = 0; bool resized = false; + auto rollback = [&]() noexcept { + if (idx == nullptr || !resized) { + return; + } + for (int i = 0; i < n; ++i) { + idx->id_to_slot.erase(ids[i]); + } + if (is_q8(*idx)) { + idx->q8_data.resize(base_slot * dim_sz); + idx->q8_scale.resize(base_slot); + } else { + idx->data.resize(base_slot * dim_sz); + } + idx->slot_to_id.resize(base_slot); + }; + try { if (idx == nullptr || vectors == nullptr || ids == nullptr) { return GGML_VEC_INDEX_E_INVALID_ARG; @@ -197,14 +557,14 @@ int ggml_vec_index_add( // Atomic add: detect duplicates first (against existing AND in-batch), // bail before mutating any state. + std::unordered_set batch_ids; + batch_ids.reserve(static_cast(n)); for (int i = 0; i < n; ++i) { if (idx->id_to_slot.find(ids[i]) != idx->id_to_slot.end()) { return GGML_VEC_INDEX_E_DUPLICATE; } - for (int j = i + 1; j < n; ++j) { - if (ids[i] == ids[j]) { - return GGML_VEC_INDEX_E_DUPLICATE; - } + if (!batch_ids.insert(ids[i]).second) { + return GGML_VEC_INDEX_E_DUPLICATE; } } @@ -218,38 +578,43 @@ int ggml_vec_index_add( if (dim_sz != 0 && new_slots > std::numeric_limits::max() / dim_sz) { return GGML_VEC_INDEX_E_INVALID_ARG; } + if (!all_finite(vectors, n_sz * dim_sz)) { + return GGML_VEC_INDEX_E_INVALID_ARG; + } - idx->data.resize(new_slots * dim_sz); resized = true; + if (is_q8(*idx)) { + idx->q8_data.resize(new_slots * dim_sz); + idx->q8_scale.resize(new_slots); + } else { + idx->data.resize(new_slots * dim_sz); + } idx->slot_to_id.resize(new_slots); idx->id_to_slot.reserve(new_slots); for (int i = 0; i < n; ++i) { const size_t slot = base_slot + static_cast(i); - std::memcpy( - idx->data.data() + slot * dim_sz, - vectors + static_cast(i) * dim_sz, - dim_sz * sizeof(float)); + const float * src = vectors + static_cast(i) * dim_sz; + if (is_q8(*idx)) { + quantize_q8_row( + src, + idx->q8_data.data() + slot * dim_sz, + idx->dim, + idx->q8_scale[slot]); + } else { + std::memcpy( + idx->data.data() + slot * dim_sz, + src, + dim_sz * sizeof(float)); + } idx->slot_to_id[slot] = ids[i]; idx->id_to_slot.emplace(ids[i], slot); } } catch (const std::bad_alloc &) { - if (idx != nullptr && resized) { - for (int i = 0; i < n; ++i) { - idx->id_to_slot.erase(ids[i]); - } - idx->data.resize(base_slot * dim_sz); - idx->slot_to_id.resize(base_slot); - } + rollback(); return GGML_VEC_INDEX_E_OOM; } catch (...) { - if (idx != nullptr && resized) { - for (int i = 0; i < n; ++i) { - idx->id_to_slot.erase(ids[i]); - } - idx->data.resize(base_slot * dim_sz); - idx->slot_to_id.resize(base_slot); - } + rollback(); return GGML_VEC_INDEX_E_INTERNAL; } return GGML_VEC_INDEX_OK; @@ -270,17 +635,30 @@ int ggml_vec_index_remove(ggml_vec_index_t * idx, uint64_t id) { if (slot != last) { // Move last vector into the freed slot and update its id mapping. - std::memcpy( - idx->data.data() + slot * dim_sz, - idx->data.data() + last * dim_sz, - dim_sz * sizeof(float)); + if (is_q8(*idx)) { + std::memcpy( + idx->q8_data.data() + slot * dim_sz, + idx->q8_data.data() + last * dim_sz, + dim_sz * sizeof(int8_t)); + idx->q8_scale[slot] = idx->q8_scale[last]; + } else { + std::memcpy( + idx->data.data() + slot * dim_sz, + idx->data.data() + last * dim_sz, + dim_sz * sizeof(float)); + } const uint64_t moved_id = idx->slot_to_id[last]; idx->slot_to_id[slot] = moved_id; idx->id_to_slot[moved_id] = slot; } idx->slot_to_id.pop_back(); - idx->data.resize(last * dim_sz); + if (is_q8(*idx)) { + idx->q8_data.resize(last * dim_sz); + idx->q8_scale.resize(last); + } else { + idx->data.resize(last * dim_sz); + } idx->id_to_slot.erase(it); return 1; } catch (...) { @@ -296,7 +674,7 @@ int ggml_vec_index_contains(const ggml_vec_index_t * idx, uint64_t id) { } void ggml_vec_index_prepare(ggml_vec_index_t * /*idx*/) { - // POC no-op. Future: warm caches, materialize codebooks, etc. + // Future: warm caches, materialize codebooks, etc. } // --------------------------------------------------------------------------- @@ -314,6 +692,59 @@ inline float dot(const float * a, const float * b, int dim) { return acc; } +inline float dot_q8_scalar(const float * query, const int8_t * codes, float scale, int dim) { + float acc = 0.0f; + for (int i = 0; i < dim; ++i) { + const float value = static_cast(codes[i]) * scale; + acc += query[i] * value; + } + return acc; +} + +#if GGML_VEC_INDEX_USE_NEON + +inline float horizontal_sum(float32x4_t v) { +#if defined(__aarch64__) + return vaddvq_f32(v); +#else + const float32x2_t sum2 = vadd_f32(vget_low_f32(v), vget_high_f32(v)); + return vget_lane_f32(vpadd_f32(sum2, sum2), 0); +#endif +} + +inline float dot_q8_neon(const float * query, const int8_t * codes, float scale, int dim) { + float32x4_t acc0 = vdupq_n_f32(0.0f); + float32x4_t acc1 = vdupq_n_f32(0.0f); + + int i = 0; + for (; i + 8 <= dim; i += 8) { + const int16x8_t q16 = vmovl_s8(vld1_s8(codes + i)); + const float32x4_t q0 = + vmulq_n_f32(vcvtq_f32_s32(vmovl_s16(vget_low_s16(q16))), scale); + const float32x4_t q1 = + vmulq_n_f32(vcvtq_f32_s32(vmovl_s16(vget_high_s16(q16))), scale); + acc0 = vmlaq_f32(acc0, vld1q_f32(query + i), q0); + acc1 = vmlaq_f32(acc1, vld1q_f32(query + i + 4), q1); + } + + float acc = horizontal_sum(acc0) + horizontal_sum(acc1); + for (; i < dim; ++i) { + const float value = static_cast(codes[i]) * scale; + acc += query[i] * value; + } + return acc; +} + +#endif + +inline float dot_q8(const float * query, const int8_t * codes, float scale, int dim) { +#if GGML_VEC_INDEX_USE_NEON + return dot_q8_neon(query, codes, scale, dim); +#else + return dot_q8_scalar(query, codes, scale, dim); +#endif +} + // Run a single query against all slots, write top-k into out_scores/out_ids. // If the index holds fewer than k entries, pad with sentinels. void search_one( @@ -329,8 +760,16 @@ void search_one( std::priority_queue, MinHeapCmp> heap; for (size_t slot = 0; slot < n_slots; ++slot) { - const float s = dot( - query, idx.data.data() + slot * static_cast(dim), dim); + const float s = is_q8(idx) ? + dot_q8( + query, + idx.q8_data.data() + slot * static_cast(dim), + idx.q8_scale[slot], + dim) : + dot( + query, + idx.data.data() + slot * static_cast(dim), + dim); if (heap.size() < static_cast(k)) { heap.push({ s, idx.slot_to_id[slot] }); } else if (s > heap.top().score) { @@ -389,6 +828,9 @@ int ggml_vec_index_search( n_q_sz > std::numeric_limits::max() / k_sz) { return GGML_VEC_INDEX_E_INVALID_ARG; } + if (!all_finite(queries, n_q_sz * dim_sz)) { + return GGML_VEC_INDEX_E_INVALID_ARG; + } for (int q = 0; q < n_q; ++q) { search_one( @@ -423,44 +865,106 @@ int ggml_vec_index_write(ggml_vec_index_t * idx, const char * path) { if (dim_sz != 0 && n > std::numeric_limits::max() / dim_sz) { return GGML_VEC_INDEX_E_INTERNAL; } - if (idx->data.size() != n * dim_sz) { + if (is_q8(*idx)) { + if (idx->q8_data.size() != n * dim_sz || idx->q8_scale.size() != n) { + return GGML_VEC_INDEX_E_INTERNAL; + } + } else if (idx->data.size() != n * dim_sz) { return GGML_VEC_INDEX_E_INTERNAL; } - std::ofstream f(path, std::ios::binary | std::ios::trunc); - if (!f.is_open()) { + TempFile temp; + if (!open_temp_file(path, temp)) { return GGML_VEC_INDEX_E_IO; } + auto fail_io = [&]() { + if (temp.stream != nullptr) { + std::fclose(temp.stream); + temp.stream = nullptr; + } + remove_temp_file(temp); + return GGML_VEC_INDEX_E_IO; + }; + std::FILE * f = temp.stream; - // Header: 16 bytes. Layout matches the comment block in the header file. + // Header: 32 bytes. Layout matches the comment block in the header file. uint8_t header[kTvimHeaderSize] = {}; std::memcpy(header, kTvimMagic, 4); header[4] = kTvimVersion; header[5] = static_cast(idx->bit_width); - header[6] = 0; - header[7] = 0; + header[6] = storage_kind(*idx); + header[7] = kFlagCRC32C; const uint32_t dim_le = static_cast(idx->dim); const uint32_t n_le = static_cast(idx->slot_to_id.size()); put_u32_le(header + 8, dim_le); put_u32_le(header + 12, n_le); + put_u32_le(header + 16, is_q8(*idx) ? kQParamScaleF32 : kQParamNone); + put_u32_le(header + 20, is_q8(*idx) ? 4u : 0u); + put_u32_le(header + 24, is_q8(*idx) ? 1u : 4u); + put_u32_le(header + 28, 0); - f.write(reinterpret_cast(header), sizeof(header)); - if (!f) { return GGML_VEC_INDEX_E_IO; } + if (std::fwrite(header, 1, sizeof(header), f) != sizeof(header)) { + return fail_io(); + } + + uint32_t header_crc = crc32c_update(0xffffffffu, header, sizeof(header)); + uint32_t qparams_crc = 0xffffffffu; + uint32_t vectors_crc = 0xffffffffu; + uint32_t ids_crc = 0xffffffffu; + if (is_q8(*idx)) { + for (float scale : idx->q8_scale) { + if (!write_u32_le_crc(f, float_to_u32(scale), qparams_crc)) { + return fail_io(); + } + } - for (float v : idx->data) { - if (!write_u32_le(f, float_to_u32(v))) { - return GGML_VEC_INDEX_E_IO; + if (!idx->q8_data.empty()) { + if (std::fwrite( + idx->q8_data.data(), + sizeof(int8_t), + idx->q8_data.size(), + f) != idx->q8_data.size()) { + return fail_io(); + } + vectors_crc = + crc32c_update(vectors_crc, idx->q8_data.data(), idx->q8_data.size()); + } + } else { + for (float v : idx->data) { + if (!write_u32_le_crc(f, float_to_u32(v), vectors_crc)) { + return fail_io(); + } } } for (uint64_t id : idx->slot_to_id) { - if (!write_u64_le(f, id)) { - return GGML_VEC_INDEX_E_IO; + if (!write_u64_le_crc(f, id, ids_crc)) { + return fail_io(); } } - f.flush(); - if (!f) { return GGML_VEC_INDEX_E_IO; } + if (!write_u32_le(f, header_crc ^ 0xffffffffu) || + !write_u32_le(f, qparams_crc ^ 0xffffffffu) || + !write_u32_le(f, vectors_crc ^ 0xffffffffu) || + !write_u32_le(f, ids_crc ^ 0xffffffffu)) { + return fail_io(); + } + if (!flush_and_sync(f)) { + return fail_io(); + } + const int close_result = std::fclose(temp.stream); + temp.stream = nullptr; + if (close_result != 0) { + remove_temp_file(temp); + return GGML_VEC_INDEX_E_IO; + } + if (!rename_overwrite(temp, path)) { + return fail_io(); + } + temp.path.clear(); + if (!fsync_parent_dir(path)) { + return GGML_VEC_INDEX_E_IO; + } return GGML_VEC_INDEX_OK; } catch (const std::bad_alloc &) { return GGML_VEC_INDEX_E_OOM; @@ -474,29 +978,111 @@ ggml_vec_index_t * ggml_vec_index_load(const char * path) { if (path == nullptr) { return nullptr; } - std::ifstream f(path, std::ios::binary); + std::ifstream f; +#ifdef _WIN32 + std::wstring wide_path; + if (!utf8_to_wide(path, wide_path)) { + return nullptr; + } + f.open(std::filesystem::path(wide_path), std::ios::binary); +#else + f.open(path, std::ios::binary); +#endif if (!f.is_open()) { return nullptr; } uint8_t header[kTvimHeaderSize] = {}; - f.read(reinterpret_cast(header), sizeof(header)); - if (!f || f.gcount() != static_cast(sizeof(header))) { + f.read(reinterpret_cast(header), kTvimV1HeaderSize); + if (!f || f.gcount() != static_cast(kTvimV1HeaderSize)) { return nullptr; } if (std::memcmp(header, kTvimMagic, 4) != 0) { return nullptr; } - if (header[4] != kTvimVersion) { + + const uint8_t version = header[4]; + if (version != kTvimVersionV1 && version != kTvimVersion) { + return nullptr; + } + if (version == kTvimVersion) { + f.read( + reinterpret_cast(header + kTvimV1HeaderSize), + kTvimHeaderSize - kTvimV1HeaderSize); + if (!f || + f.gcount() != static_cast( + kTvimHeaderSize - kTvimV1HeaderSize)) { + return nullptr; + } + } + + const uint8_t flags = version == kTvimVersion ? header[7] : 0; + if ((version == kTvimVersionV1 && (header[6] != 0 || header[7] != 0)) || + (version == kTvimVersion && + ((flags & ~kFlagCRC32C) != 0 || get_u32_le(header + 28) != 0))) { return nullptr; } - const int bit_width = static_cast(header[5]); + const int serialized_bit_width = static_cast(header[5]); + if ((version == kTvimVersionV1 && + (serialized_bit_width <= 0 || serialized_bit_width > 32)) || + (version == kTvimVersion && !is_supported_bit_width(serialized_bit_width))) { + return nullptr; + } + const int bit_width = + version == kTvimVersionV1 && serialized_bit_width != 8 ? + 32 : serialized_bit_width; + const uint8_t kind = header[6]; const uint32_t dim_le = get_u32_le(header + 8); const uint32_t n_le = get_u32_le(header + 12); + const uint32_t qparam_type = + version == kTvimVersion ? get_u32_le(header + 16) : kQParamNone; + const uint32_t qparam_bytes = + version == kTvimVersion ? get_u32_le(header + 20) : 0; + const uint32_t comp_bytes = + version == kTvimVersion ? get_u32_le(header + 24) : 4; if (dim_le == 0 || dim_le > static_cast(std::numeric_limits::max())) { return nullptr; } + if (version == kTvimVersion && + ((bit_width == 8 && (kind != kStorageQ8 || qparam_type != kQParamScaleF32 || + qparam_bytes != 4 || comp_bytes != 1)) || + (bit_width == 32 && (kind != kStorageF32 || qparam_type != kQParamNone || + qparam_bytes != 0 || comp_bytes != 4)))) { + return nullptr; + } + + uint64_t expected_size = 0; + if (!expected_file_size( + version == kTvimVersion ? kTvimHeaderSize : kTvimV1HeaderSize, + n_le, + dim_le, + qparam_bytes, + comp_bytes, + expected_size)) { + return nullptr; + } + if ((flags & kFlagCRC32C) != 0 && + !checked_add_u64(expected_size, kTvimChecksumSize, expected_size)) { + return nullptr; + } + if (expected_size > + static_cast(std::numeric_limits::max())) { + return nullptr; + } + f.seekg(0, std::ios::end); + const std::streamoff actual_size = f.tellg(); + if (!f || actual_size != static_cast(expected_size)) { + return nullptr; + } + f.seekg( + static_cast( + version == kTvimVersion ? kTvimHeaderSize : kTvimV1HeaderSize), + std::ios::beg); + if (!f) { + return nullptr; + } + const int dim = static_cast(dim_le); std::unique_ptr idx( @@ -511,20 +1097,120 @@ ggml_vec_index_t * ggml_vec_index_load(const char * path) { return nullptr; } - idx->data.resize(n * dim_sz); + if (is_q8(*idx)) { + idx->q8_data.resize(n * dim_sz); + idx->q8_scale.resize(n); + } else { + idx->data.resize(n * dim_sz); + } idx->slot_to_id.resize(n); idx->id_to_slot.reserve(n); - for (float & v : idx->data) { - uint32_t bits = 0; - if (!read_u32_le(f, bits)) { - return nullptr; + const bool checksummed = (flags & kFlagCRC32C) != 0; + uint32_t header_crc = + crc32c_update(0xffffffffu, header, kTvimHeaderSize); + uint32_t qparams_crc = 0xffffffffu; + uint32_t vectors_crc = 0xffffffffu; + uint32_t ids_crc = 0xffffffffu; + + if (version == kTvimVersionV1 && is_q8(*idx)) { + std::vector row(dim_sz); + for (size_t slot = 0; slot < n; ++slot) { + for (float & v : row) { + uint32_t bits = 0; + if (!read_u32_le(f, bits)) { + return nullptr; + } + v = u32_to_float(bits); + if (!std::isfinite(v)) { + return nullptr; + } + } + quantize_q8_row( + row.data(), + idx->q8_data.data() + slot * dim_sz, + dim, + idx->q8_scale[slot]); + } + } else if (is_q8(*idx)) { + for (float & scale : idx->q8_scale) { + uint32_t bits = 0; + const bool read_ok = checksummed ? + read_u32_le_crc(f, bits, qparams_crc) : + read_u32_le(f, bits); + if (!read_ok) { + return nullptr; + } + scale = u32_to_float(bits); + if (!std::isfinite(scale) || scale <= 0.0f) { + return nullptr; + } + } + + if (!idx->q8_data.empty()) { + if (idx->q8_data.size() > + static_cast(std::numeric_limits::max())) { + return nullptr; + } + f.read( + reinterpret_cast(idx->q8_data.data()), + static_cast(idx->q8_data.size() * sizeof(int8_t))); + if (!f) { + return nullptr; + } + if (checksummed) { + vectors_crc = crc32c_update( + vectors_crc, idx->q8_data.data(), idx->q8_data.size()); + } + } + for (size_t slot = 0; slot < n; ++slot) { + const float scale = idx->q8_scale[slot]; + const int8_t * row = idx->q8_data.data() + slot * dim_sz; + for (size_t i = 0; i < dim_sz; ++i) { + if (row[i] == std::numeric_limits::min() || + !std::isfinite(static_cast(row[i]) * scale)) { + return nullptr; + } + } + } + } else { + for (float & v : idx->data) { + uint32_t bits = 0; + const bool read_ok = checksummed ? + read_u32_le_crc(f, bits, vectors_crc) : + read_u32_le(f, bits); + if (!read_ok) { + return nullptr; + } + v = u32_to_float(bits); + if (!std::isfinite(v)) { + return nullptr; + } } - v = u32_to_float(bits); } for (uint64_t & id : idx->slot_to_id) { - if (!read_u64_le(f, id)) { + const bool read_ok = checksummed ? + read_u64_le_crc(f, id, ids_crc) : + read_u64_le(f, id); + if (!read_ok) { + return nullptr; + } + } + + if (checksummed) { + uint32_t expected_header_crc = 0; + uint32_t expected_qparams_crc = 0; + uint32_t expected_vectors_crc = 0; + uint32_t expected_ids_crc = 0; + if (!read_u32_le(f, expected_header_crc) || + !read_u32_le(f, expected_qparams_crc) || + !read_u32_le(f, expected_vectors_crc) || + !read_u32_le(f, expected_ids_crc) || + (header_crc ^ 0xffffffffu) != expected_header_crc || + (qparams_crc ^ 0xffffffffu) != expected_qparams_crc || + (vectors_crc ^ 0xffffffffu) != expected_vectors_crc || + (ids_crc ^ 0xffffffffu) != expected_ids_crc) { return nullptr; } } diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index e75fe7be1fd2..36fa743c8716 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -144,7 +144,7 @@ if (NOT WIN32) ) endif() -# POC vector-index smoke test: no model required, no llama runtime touched. +# Vector-index regression test: no model required, no llama runtime touched. add_executable(test-vector-index test-vector-index.cpp) target_compile_features(test-vector-index PRIVATE cxx_std_17) target_link_libraries(test-vector-index PRIVATE ggml) @@ -153,6 +153,13 @@ if (LLAMA_TESTS_INSTALL) endif() llama_test(test-vector-index) +add_executable(bench-vector-index bench-vector-index.cpp) +target_compile_features(bench-vector-index PRIVATE cxx_std_17) +target_link_libraries(bench-vector-index PRIVATE ggml) +if (LLAMA_TESTS_INSTALL) + install(TARGETS bench-vector-index RUNTIME) +endif() + if (LLAMA_LLGUIDANCE) llama_build_and_test(test-grammar-llguidance.cpp ARGS ${PROJECT_SOURCE_DIR}/models/ggml-vocab-llama-bpe.gguf) endif () diff --git a/tests/bench-vector-index.cpp b/tests/bench-vector-index.cpp new file mode 100644 index 000000000000..b588466f97ed --- /dev/null +++ b/tests/bench-vector-index.cpp @@ -0,0 +1,217 @@ +// bench-vector-index.cpp - f32 vs q8 vector-index quality and latency smoke bench. + +#include "ggml-vector-index.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace { + +#define CHECK(cond) \ + do { \ + if (!(cond)) { \ + std::fprintf(stderr, "FAIL %s:%d: %s\n", __FILE__, __LINE__, #cond);\ + std::exit(1); \ + } \ + } while (0) + +struct BenchConfig { + int n_vec = 2048; + int dim = 256; + int n_query = 64; + int k = 10; + int warmups = 2; + int repeats = 7; +}; + +struct TimedSearch { + double ms = 0.0; + std::vector scores; + std::vector ids; +}; + +const char * q8_kernel_name() { +#if defined(__ARM_NEON) || defined(__ARM_NEON__) + return "arm-neon"; +#else + return "scalar"; +#endif +} + +std::vector make_normalized_vectors(int n, int dim, uint32_t seed) { + std::mt19937 rng(seed); + std::normal_distribution dist(0.0f, 1.0f); + + std::vector vectors(static_cast(n) * static_cast(dim)); + for (int row = 0; row < n; ++row) { + float norm2 = 0.0f; + float * v = vectors.data() + static_cast(row) * static_cast(dim); + for (int i = 0; i < dim; ++i) { + v[i] = dist(rng); + norm2 += v[i] * v[i]; + } + const float inv_norm = norm2 > 0.0f ? 1.0f / std::sqrt(norm2) : 1.0f; + for (int i = 0; i < dim; ++i) { + v[i] *= inv_norm; + } + } + return vectors; +} + +TimedSearch run_search( + const ggml_vec_index_t * idx, + const std::vector & queries, + int n_query, + int k, + int warmups, + int repeats) { + TimedSearch result; + result.scores.resize(static_cast(n_query) * static_cast(k)); + result.ids.resize(static_cast(n_query) * static_cast(k)); + + for (int i = 0; i < warmups; ++i) { + CHECK(ggml_vec_index_search( + idx, + queries.data(), + n_query, + k, + result.scores.data(), + result.ids.data()) == GGML_VEC_INDEX_OK); + } + + std::vector times; + times.reserve(static_cast(repeats)); + for (int i = 0; i < repeats; ++i) { + const auto t0 = std::chrono::steady_clock::now(); + CHECK(ggml_vec_index_search( + idx, + queries.data(), + n_query, + k, + result.scores.data(), + result.ids.data()) == GGML_VEC_INDEX_OK); + const auto t1 = std::chrono::steady_clock::now(); + times.push_back(std::chrono::duration(t1 - t0).count()); + } + std::sort(times.begin(), times.end()); + result.ms = times[times.size() / 2]; + return result; +} + +float dot_exact( + const std::vector & vectors, + const float * query, + uint64_t id, + int dim) { + const size_t row = static_cast(id - 1); + const float * v = vectors.data() + row * static_cast(dim); + float acc = 0.0f; + for (int i = 0; i < dim; ++i) { + acc += query[i] * v[i]; + } + return acc; +} + +std::filesystem::path write_index_file(ggml_vec_index_t * idx, const char * name) { + const auto path = std::filesystem::temp_directory_path() / name; + CHECK(ggml_vec_index_write(idx, path.string().c_str()) == GGML_VEC_INDEX_OK); + return path; +} + +} // namespace + +int main() { + const BenchConfig cfg; + + std::vector vectors = make_normalized_vectors(cfg.n_vec, cfg.dim, 0xdeadbeef); + std::vector queries = make_normalized_vectors(cfg.n_query, cfg.dim, 0xc001d00d); + std::vector ids(static_cast(cfg.n_vec)); + for (int i = 0; i < cfg.n_vec; ++i) { + ids[static_cast(i)] = static_cast(i) + 1; + } + + ggml_vec_index_t * f32 = ggml_vec_index_create(cfg.dim, 32); + ggml_vec_index_t * q8 = ggml_vec_index_create(cfg.dim, 8); + CHECK(f32 != nullptr); + CHECK(q8 != nullptr); + CHECK(ggml_vec_index_add(f32, vectors.data(), cfg.n_vec, ids.data()) == GGML_VEC_INDEX_OK); + CHECK(ggml_vec_index_add(q8, vectors.data(), cfg.n_vec, ids.data()) == GGML_VEC_INDEX_OK); + + const TimedSearch f32_res = run_search( + f32, queries, cfg.n_query, cfg.k, cfg.warmups, cfg.repeats); + const TimedSearch q8_res = run_search( + q8, queries, cfg.n_query, cfg.k, cfg.warmups, cfg.repeats); + + double mean_abs_drift = 0.0; + double max_abs_drift = 0.0; + int drift_count = 0; + int overlap = 0; + + for (int q = 0; q < cfg.n_query; ++q) { + std::unordered_set f32_topk; + for (int j = 0; j < cfg.k; ++j) { + f32_topk.insert(f32_res.ids[static_cast(q) * cfg.k + j]); + } + + const float * query = queries.data() + static_cast(q) * static_cast(cfg.dim); + for (int j = 0; j < cfg.k; ++j) { + const size_t pos = static_cast(q) * static_cast(cfg.k) + static_cast(j); + const uint64_t id = q8_res.ids[pos]; + if (f32_topk.count(id) != 0) { + ++overlap; + } + const float exact = dot_exact(vectors, query, id, cfg.dim); + const double drift = std::fabs(static_cast(exact) - q8_res.scores[pos]); + mean_abs_drift += drift; + max_abs_drift = std::max(max_abs_drift, drift); + ++drift_count; + } + } + mean_abs_drift /= static_cast(drift_count); + const double recall_at_k = static_cast(overlap) / + static_cast(cfg.n_query * cfg.k); + + const auto f32_path = write_index_file(f32, "ggml-vector-index-bench-f32.tvim"); + const auto q8_path = write_index_file(q8, "ggml-vector-index-bench-q8.tvim"); + const uintmax_t f32_file_size = std::filesystem::file_size(f32_path); + const uintmax_t q8_file_size = std::filesystem::file_size(q8_path); + std::filesystem::remove(f32_path); + std::filesystem::remove(q8_path); + + const size_t f32_memory_bytes = + static_cast(cfg.n_vec) * static_cast(cfg.dim) * sizeof(float) + + static_cast(cfg.n_vec) * sizeof(uint64_t); + const size_t q8_memory_bytes = + static_cast(cfg.n_vec) * static_cast(cfg.dim) * sizeof(int8_t) + + static_cast(cfg.n_vec) * sizeof(float) + + static_cast(cfg.n_vec) * sizeof(uint64_t); + + std::printf("bench-vector-index\n"); + std::printf(" q8 kernel=%s\n", q8_kernel_name()); + std::printf(" n_vec=%d dim=%d n_query=%d k=%d warmups=%d repeats=%d\n", + cfg.n_vec, cfg.dim, cfg.n_query, cfg.k, cfg.warmups, cfg.repeats); + std::printf(" estimated memory: f32=%zu bytes q8=%zu bytes ratio=%.3f\n", + f32_memory_bytes, q8_memory_bytes, + static_cast(q8_memory_bytes) / static_cast(f32_memory_bytes)); + std::printf(" file size: f32=%llu bytes q8=%llu bytes ratio=%.3f\n", + static_cast(f32_file_size), + static_cast(q8_file_size), + static_cast(q8_file_size) / static_cast(f32_file_size)); + std::printf(" median latency: f32=%.3f ms q8=%.3f ms ratio=%.3f\n", + f32_res.ms, q8_res.ms, q8_res.ms / f32_res.ms); + std::printf(" quality: recall@%d=%.4f mean_abs_score_drift=%.6f max_abs_score_drift=%.6f\n", + cfg.k, recall_at_k, mean_abs_drift, max_abs_drift); + + ggml_vec_index_free(f32); + ggml_vec_index_free(q8); + return 0; +} diff --git a/tests/test-vector-index.cpp b/tests/test-vector-index.cpp index bc633b98f668..799dabbbf784 100644 --- a/tests/test-vector-index.cpp +++ b/tests/test-vector-index.cpp @@ -1,10 +1,11 @@ -// test-vector-index.cpp - standalone C-API smoke test for the POC vector +// test-vector-index.cpp - standalone C-API smoke test for the vector // index. Exercises lifecycle, add, search, remove, contains, write, load, // search-after-load. No model, no llama; only the new ggml-vector-index // public C API. #include "ggml-vector-index.h" +#include #include #include #include @@ -12,10 +13,17 @@ #include #include #include +#include #include +#include +#include #include #include +#ifndef _WIN32 +#include +#endif + namespace { constexpr int kDim = 4; @@ -36,6 +44,110 @@ std::vector normalize(std::vector v) { return v; } +float q8_dot_reference(const std::vector & vector, const std::vector & query) { + CHECK(vector.size() == query.size()); + + float max_abs = 0.0f; + for (float value : vector) { + max_abs = std::max(max_abs, std::fabs(value)); + } + const float scale = max_abs == 0.0f ? 1.0f : max_abs / 127.0f; + + float acc = 0.0f; + for (size_t i = 0; i < vector.size(); ++i) { + int code = max_abs == 0.0f ? + 0 : static_cast(std::nearbyint(vector[i] / scale)); + code = std::max(-127, std::min(127, code)); + acc += query[i] * (static_cast(code) * scale); + } + return acc; +} + +uint8_t read_file_byte(const std::string & path, std::streamoff offset) { + std::ifstream f(path, std::ios::binary); + CHECK(f.is_open()); + f.seekg(offset); + char c = 0; + f.read(&c, 1); + CHECK(f.good()); + return static_cast(c); +} + +std::vector read_file_bytes(const std::string & path) { + std::ifstream f(path, std::ios::binary); + CHECK(f.is_open()); + const auto size = std::filesystem::file_size(path); + std::vector bytes(static_cast(size)); + if (!bytes.empty()) { + f.read(reinterpret_cast(bytes.data()), static_cast(bytes.size())); + CHECK(f.gcount() == static_cast(bytes.size())); + } + return bytes; +} + +void write_file_bytes(const std::string & path, const std::vector & bytes) { + std::ofstream f(path, std::ios::binary | std::ios::trunc); + CHECK(f.is_open()); + if (!bytes.empty()) { + f.write(reinterpret_cast(bytes.data()), static_cast(bytes.size())); + } + CHECK(static_cast(f)); +} + +void append_u32_le(std::vector & bytes, uint32_t value) { + for (int i = 0; i < 4; ++i) { + bytes.push_back(static_cast(value >> (8 * i))); + } +} + +void append_u64_le(std::vector & bytes, uint64_t value) { + for (int i = 0; i < 8; ++i) { + bytes.push_back(static_cast(value >> (8 * i))); + } +} + +void append_f32_le(std::vector & bytes, float value) { + uint32_t bits = 0; + std::memcpy(&bits, &value, sizeof(bits)); + append_u32_le(bytes, bits); +} + +void write_v1_index( + const std::string & path, + int dim, + int bit_width, + const std::vector & vectors, + const std::vector & ids) { + CHECK(vectors.size() == ids.size() * static_cast(dim)); + + std::vector bytes = { 'T', 'V', 'P', 'I', 1, + static_cast(bit_width), 0, 0 }; + append_u32_le(bytes, static_cast(dim)); + append_u32_le(bytes, static_cast(ids.size())); + for (float value : vectors) { + append_f32_le(bytes, value); + } + for (uint64_t id : ids) { + append_u64_le(bytes, id); + } + write_file_bytes(path, bytes); +} + +template +void expect_corrupt_load_fails( + const std::string & source_path, + const std::string & corrupt_path, + Fn mutate) { + std::vector bytes = read_file_bytes(source_path); + mutate(bytes); + write_file_bytes(corrupt_path, bytes); + + auto * bad = ggml_vec_index_load(corrupt_path.c_str()); + CHECK(bad == nullptr); + ggml_vec_index_free(bad); + std::filesystem::remove(corrupt_path); +} + } // namespace int main() { @@ -45,6 +157,17 @@ int main() { CHECK(ggml_vec_index_len(idx) == 0); CHECK(ggml_vec_index_bit_width(idx) == 32); + // Non-finite vectors are rejected without mutation. + { + const std::array bad_vector = { + 1.0f, 0.0f, std::numeric_limits::infinity(), 0.0f, + }; + const uint64_t bad_id = 777ULL; + CHECK(ggml_vec_index_add(idx, bad_vector.data(), 1, &bad_id) + == GGML_VEC_INDEX_E_INVALID_ARG); + CHECK(ggml_vec_index_len(idx) == 0); + } + // Add 4 well-separated unit vectors. IDs are non-trivial uint64 to // catch sign-extension or BigInt round-trip bugs at the JS boundary // when this codepath is later exercised from Bare. @@ -70,6 +193,18 @@ int main() { CHECK(ggml_vec_index_contains(idx, ids[0]) == 1); CHECK(ggml_vec_index_contains(idx, 999ULL) == 0); + // Non-finite queries are rejected before search. + { + const std::array bad_query = { + 1.0f, std::numeric_limits::quiet_NaN(), 0.0f, 0.0f, + }; + std::array scores{}; + std::array out_ids{}; + CHECK(ggml_vec_index_search( + idx, bad_query.data(), 1, /*k=*/1, + scores.data(), out_ids.data()) == GGML_VEC_INDEX_E_INVALID_ARG); + } + // Duplicate add must fail without mutating state. { const std::vector dup_ids = { ids[0] }; @@ -143,6 +278,13 @@ int main() { "ggml-vector-index-test.tvim"; const std::string path = tmp.string(); CHECK(ggml_vec_index_write(idx, path.c_str()) == 0); +#ifndef _WIN32 + CHECK(::chmod(path.c_str(), 0600) == 0); + CHECK(ggml_vec_index_write(idx, path.c_str()) == 0); + struct stat persisted_stat; + CHECK(::stat(path.c_str(), &persisted_stat) == 0); + CHECK((persisted_stat.st_mode & 0777) == 0600); +#endif ggml_vec_index_free(idx); @@ -170,6 +312,284 @@ int main() { ggml_vec_index_free(loaded); std::filesystem::remove(path); + // v1 f32 snapshots migrate into current f32 or q8 storage. + { + const std::vector v1_ids = { + (1ULL << 37) + 1ULL, + (1ULL << 37) + 2ULL, + }; + std::vector v1_vectors; + v1_vectors.insert(v1_vectors.end(), seeds[0].begin(), seeds[0].end()); + v1_vectors.insert(v1_vectors.end(), seeds[1].begin(), seeds[1].end()); + + for (int bit_width : { 32, 8, 4 }) { + const auto v1_tmp = std::filesystem::temp_directory_path() / + ("ggml-vector-index-v1-" + std::to_string(bit_width) + ".tvim"); + const std::string v1_path = v1_tmp.string(); + write_v1_index(v1_path, kDim, bit_width, v1_vectors, v1_ids); + + auto * v1 = ggml_vec_index_load(v1_path.c_str()); + CHECK(v1 != nullptr); + CHECK(ggml_vec_index_dim(v1) == kDim); + CHECK(ggml_vec_index_len(v1) == 2); + CHECK(ggml_vec_index_bit_width(v1) == (bit_width == 8 ? 8 : 32)); + + std::array scores{}; + std::array out_ids{}; + CHECK(ggml_vec_index_search( + v1, seeds[1].data(), 1, /*k=*/1, + scores.data(), out_ids.data()) == GGML_VEC_INDEX_OK); + CHECK(out_ids[0] == v1_ids[1]); + CHECK(std::fabs(scores[0] - 1.0f) < 1e-5f); + + ggml_vec_index_free(v1); + std::filesystem::remove(v1_path); + } + } + + // q8 score parity for a dimension that exercises the SIMD tail. + { + constexpr int tail_dim = 13; + const std::vector tail_vector = { + -1.0f, 0.75f, -0.5f, 0.25f, 0.125f, -0.875f, 0.625f, + -0.375f, 0.9f, -0.7f, 0.3f, -0.2f, 0.05f, + }; + const std::vector tail_query = { + 0.2f, -0.4f, 0.6f, -0.8f, 1.0f, 0.3f, -0.5f, + 0.7f, -0.9f, 0.11f, -0.22f, 0.33f, -0.44f, + }; + const uint64_t tail_id = (1ULL << 55) + 321ULL; + + auto * tail_idx = ggml_vec_index_create(tail_dim, /*bit_width=*/8); + CHECK(tail_idx != nullptr); + CHECK(ggml_vec_index_add( + tail_idx, tail_vector.data(), 1, &tail_id) == GGML_VEC_INDEX_OK); + + std::array scores{}; + std::array out_ids{}; + CHECK(ggml_vec_index_search( + tail_idx, tail_query.data(), 1, /*k=*/1, + scores.data(), out_ids.data()) == GGML_VEC_INDEX_OK); + CHECK(out_ids[0] == tail_id); + + const float expected = q8_dot_reference(tail_vector, tail_query); + const float tolerance = 1e-5f * std::max(1.0f, std::fabs(expected)); + CHECK(std::fabs(scores[0] - expected) <= tolerance); + + ggml_vec_index_free(tail_idx); + + std::vector zero_vector(tail_dim, 0.0f); + auto * zero_idx = ggml_vec_index_create(tail_dim, /*bit_width=*/8); + CHECK(zero_idx != nullptr); + CHECK(ggml_vec_index_add( + zero_idx, zero_vector.data(), 1, &tail_id) == GGML_VEC_INDEX_OK); + CHECK(ggml_vec_index_search( + zero_idx, tail_query.data(), 1, /*k=*/1, + scores.data(), out_ids.data()) == GGML_VEC_INDEX_OK); + CHECK(scores[0] == 0.0f); + ggml_vec_index_free(zero_idx); + } + + // Applying the q8 scale after accumulation can overflow even when the + // dequantized dot product is finite. + { + constexpr int overflow_dim = 8; + const std::vector small_vector(overflow_dim, 1e-30f); + const std::vector large_query(overflow_dim, 1e38f); + const uint64_t overflow_id = 123456789ULL; + + auto * overflow_idx = ggml_vec_index_create(overflow_dim, /*bit_width=*/8); + CHECK(overflow_idx != nullptr); + CHECK(ggml_vec_index_add( + overflow_idx, small_vector.data(), 1, &overflow_id) == GGML_VEC_INDEX_OK); + + std::array scores{}; + std::array out_ids{}; + CHECK(ggml_vec_index_search( + overflow_idx, large_query.data(), 1, /*k=*/1, + scores.data(), out_ids.data()) == GGML_VEC_INDEX_OK); + const float expected = static_cast( + overflow_dim * static_cast(small_vector[0]) * large_query[0]); + CHECK(out_ids[0] == overflow_id); + CHECK(std::isfinite(scores[0])); + CHECK(std::fabs(scores[0] - expected) <= std::fabs(expected) * 1e-5f); + + ggml_vec_index_free(overflow_idx); + } + + // Header metadata is protected even when all payload sections are empty. + { + auto * empty_idx = ggml_vec_index_create(kDim, /*bit_width=*/8); + CHECK(empty_idx != nullptr); + const std::string empty_path = + (std::filesystem::temp_directory_path() / + "ggml-vector-index-empty-test.tvim").string(); + const std::string corrupt_path = + (std::filesystem::temp_directory_path() / + "ggml-vector-index-empty-corrupt-test.tvim").string(); + CHECK(ggml_vec_index_write(empty_idx, empty_path.c_str()) == GGML_VEC_INDEX_OK); + expect_corrupt_load_fails( + empty_path, corrupt_path, [](std::vector & bytes) { + bytes[8] += 1; // valid dimension change must fail the header CRC + }); + ggml_vec_index_free(empty_idx); + std::filesystem::remove(empty_path); + } + + // q8 production path: stores quantized codes, searches directly against + // q8 storage, and persists as .tvim v2 with q8 metadata. + { + auto * q8 = ggml_vec_index_create(kDim, /*bit_width=*/8); + CHECK(q8 != nullptr); + CHECK(ggml_vec_index_dim(q8) == kDim); + CHECK(ggml_vec_index_bit_width(q8) == 8); + + const std::vector q8_ids = { + (1ULL << 33) + 99ULL, + (1ULL << 48) + 77ULL, + }; + std::vector q8_vecs; + q8_vecs.insert(q8_vecs.end(), seeds[0].begin(), seeds[0].end()); + q8_vecs.insert(q8_vecs.end(), seeds[2].begin(), seeds[2].end()); + CHECK(ggml_vec_index_add(q8, q8_vecs.data(), 2, q8_ids.data()) == 0); + + std::array scores{}; + std::array out_ids{}; + CHECK(ggml_vec_index_search( + q8, seeds[2].data(), 1, /*k=*/4, + scores.data(), out_ids.data()) == 0); + CHECK(out_ids[0] == q8_ids[1]); + CHECK(std::fabs(scores[0] - 1.0f) < 1e-5f); + CHECK(scores[2] == -FLT_MAX); + CHECK(out_ids[2] == UINT64_MAX); + CHECK(scores[3] == -FLT_MAX); + CHECK(out_ids[3] == UINT64_MAX); + + const auto q8_tmp = std::filesystem::temp_directory_path() / + "ggml-vector-index-q8-test.tvim"; + const std::string q8_path = q8_tmp.string(); + CHECK(ggml_vec_index_write(q8, q8_path.c_str()) == 0); + CHECK(read_file_byte(q8_path, 4) == 2); // .tvim v2 + CHECK(read_file_byte(q8_path, 5) == 8); // q8 bit width + CHECK(read_file_byte(q8_path, 6) == 2); // q8 storage kind + CHECK(read_file_byte(q8_path, 7) == 1); // checksum trailer present + + const auto corrupt_tmp = std::filesystem::temp_directory_path() / + "ggml-vector-index-corrupt-test.tvim"; + const std::string corrupt_path = corrupt_tmp.string(); + + // Legacy v2 files without a checksum remain readable. + const auto legacy_v2_tmp = std::filesystem::temp_directory_path() / + "ggml-vector-index-legacy-v2-test.tvim"; + const std::string legacy_v2_path = legacy_v2_tmp.string(); + std::vector legacy_v2 = read_file_bytes(q8_path); + legacy_v2[7] = 0; + legacy_v2.resize(legacy_v2.size() - 4 * sizeof(uint32_t)); + write_file_bytes(legacy_v2_path, legacy_v2); + auto * legacy_loaded = ggml_vec_index_load(legacy_v2_path.c_str()); + CHECK(legacy_loaded != nullptr); + CHECK(ggml_vec_index_len(legacy_loaded) == 2); + ggml_vec_index_free(legacy_loaded); + expect_corrupt_load_fails( + legacy_v2_path, corrupt_path, [](std::vector & bytes) { + bytes[32] = 0; + bytes[33] = 0; + bytes[34] = 0; + bytes[35] = 0; // q8 scale must be positive and finite + }); + expect_corrupt_load_fails( + legacy_v2_path, corrupt_path, [](std::vector & bytes) { + bytes[40] = 0x80; // q8 codes are restricted to [-127, 127] + }); + expect_corrupt_load_fails( + legacy_v2_path, corrupt_path, [](std::vector & bytes) { + const size_t id_offset = + 32 + 2 * sizeof(float) + 2 * kDim * sizeof(int8_t); + for (size_t i = 0; i < sizeof(uint64_t); ++i) { + bytes[id_offset + sizeof(uint64_t) + i] = bytes[id_offset + i]; + } + }); + std::filesystem::remove(legacy_v2_path); + + expect_corrupt_load_fails(q8_path, corrupt_path, [](std::vector & bytes) { + bytes[0] = 'X'; + }); + expect_corrupt_load_fails(q8_path, corrupt_path, [](std::vector & bytes) { + bytes[4] = 99; // unsupported version + }); + expect_corrupt_load_fails(q8_path, corrupt_path, [](std::vector & bytes) { + bytes[6] = 1; // storage kind does not match bit_width=8 + }); + expect_corrupt_load_fails(q8_path, corrupt_path, [](std::vector & bytes) { + bytes[7] |= 0x80; // unknown flags are rejected + }); + expect_corrupt_load_fails(q8_path, corrupt_path, [](std::vector & bytes) { + bytes[16] = 0; // qparam_type must be scale-f32 for q8 + }); + expect_corrupt_load_fails(q8_path, corrupt_path, [](std::vector & bytes) { + bytes[28] = 1; // reserved u32 must be zero + }); + expect_corrupt_load_fails(q8_path, corrupt_path, [](std::vector & bytes) { + bytes[32] ^= 1; // q8 scale payload bit flip + }); + expect_corrupt_load_fails(q8_path, corrupt_path, [](std::vector & bytes) { + bytes[40] ^= 1; // q8 code payload bit flip + }); + expect_corrupt_load_fails(q8_path, corrupt_path, [](std::vector & bytes) { + bytes[48] ^= 1; // id payload bit flip + }); + expect_corrupt_load_fails(q8_path, corrupt_path, [](std::vector & bytes) { + bytes[bytes.size() - 16] ^= 1; // header checksum bit flip + }); + expect_corrupt_load_fails(q8_path, corrupt_path, [](std::vector & bytes) { + bytes[bytes.size() - 12] ^= 1; // qparams checksum bit flip + }); + expect_corrupt_load_fails(q8_path, corrupt_path, [](std::vector & bytes) { + bytes[bytes.size() - 8] ^= 1; // vectors checksum bit flip + }); + expect_corrupt_load_fails(q8_path, corrupt_path, [](std::vector & bytes) { + bytes[bytes.size() - 4] ^= 1; // ids checksum bit flip + }); + expect_corrupt_load_fails(q8_path, corrupt_path, [](std::vector & bytes) { + bytes.resize(35); // truncated q8 scales + }); + expect_corrupt_load_fails(q8_path, corrupt_path, [](std::vector & bytes) { + bytes.resize(43); // truncated q8 codes + }); + expect_corrupt_load_fails(q8_path, corrupt_path, [](std::vector & bytes) { + bytes.resize(bytes.size() - 1); // truncated ids + }); + expect_corrupt_load_fails(q8_path, corrupt_path, [](std::vector & bytes) { + bytes.push_back(0); // trailing data is not part of the declared file + }); + expect_corrupt_load_fails(q8_path, corrupt_path, [](std::vector & bytes) { + bytes[12] = 0xff; + bytes[13] = 0xff; + bytes[14] = 0xff; + bytes[15] = 0xff; // impossible vector count for this file size + }); + ggml_vec_index_free(q8); + + auto * q8_loaded = ggml_vec_index_load(q8_path.c_str()); + CHECK(q8_loaded != nullptr); + CHECK(ggml_vec_index_dim(q8_loaded) == kDim); + CHECK(ggml_vec_index_len(q8_loaded) == 2); + CHECK(ggml_vec_index_bit_width(q8_loaded) == 8); + CHECK(ggml_vec_index_contains(q8_loaded, q8_ids[0]) == 1); + CHECK(ggml_vec_index_contains(q8_loaded, q8_ids[1]) == 1); + + scores = {}; + out_ids = {}; + CHECK(ggml_vec_index_search( + q8_loaded, seeds[0].data(), 1, /*k=*/1, + scores.data(), out_ids.data()) == 0); + CHECK(out_ids[0] == q8_ids[0]); + CHECK(std::fabs(scores[0] - 1.0f) < 1e-5f); + + ggml_vec_index_free(q8_loaded); + std::filesystem::remove(q8_path); + } + std::printf("test-vector-index: OK\n"); return 0; } From aaffdf7175d59049111e808278cf7a95c8401bf9 Mon Sep 17 00:00:00 2001 From: Nidhin Date: Mon, 13 Jul 2026 15:58:11 +0530 Subject: [PATCH 03/13] ggml : harden vector index fault handling Adds isolated fault-injection coverage for vector index OOM and interrupted-write paths, including partial batch rollback, failed snapshot preservation, and temp-file cleanup. Assisted-by: GPT-5.5 --- ggml/src/ggml-vector-index.cpp | 79 ++++++++++++-- tests/CMakeLists.txt | 13 +++ tests/test-vector-index-faults.cpp | 163 +++++++++++++++++++++++++++++ 3 files changed, 246 insertions(+), 9 deletions(-) create mode 100644 tests/test-vector-index-faults.cpp diff --git a/ggml/src/ggml-vector-index.cpp b/ggml/src/ggml-vector-index.cpp index f18370c09f7d..059050df7770 100644 --- a/ggml/src/ggml-vector-index.cpp +++ b/ggml/src/ggml-vector-index.cpp @@ -54,6 +54,13 @@ #include #endif +#ifdef GGML_VEC_INDEX_TEST_HOOKS +extern "C" { +void ggml_vec_index_test_set_oom_countdown(int64_t countdown); +void ggml_vec_index_test_set_write_fail_after(int64_t bytes); +} +#endif + namespace { constexpr uint8_t kTvimMagic[4] = { 'T', 'V', 'P', 'I' }; @@ -70,6 +77,45 @@ constexpr size_t kTvimChecksumSize = 16; static_assert(sizeof(float) == sizeof(uint32_t), "ggml-vector-index requires float32"); +#ifdef GGML_VEC_INDEX_TEST_HOOKS +std::atomic g_test_oom_countdown{ -1 }; +std::atomic g_test_write_fail_after{ -1 }; + +void test_maybe_throw_bad_alloc() { + const int64_t remaining = g_test_oom_countdown.load(); + if (remaining < 0) { + return; + } + if (g_test_oom_countdown.fetch_sub(1) == 0) { + throw std::bad_alloc(); + } +} + +bool test_consume_write_bytes(size_t n) { + const int64_t current = g_test_write_fail_after.load(); + if (current < 0) { + return true; + } + if (static_cast(current) < n) { + g_test_write_fail_after.store(0); + return false; + } + g_test_write_fail_after.fetch_sub(static_cast(n)); + return true; +} +#else +#define test_maybe_throw_bad_alloc() ((void) 0) +#endif + +inline bool write_bytes(std::FILE * f, const void * data, size_t size) { +#ifdef GGML_VEC_INDEX_TEST_HOOKS + if (!test_consume_write_bytes(size)) { + return false; + } +#endif + return std::fwrite(data, 1, size, f) == size; +} + void put_u32_le(uint8_t * dst, uint32_t v) { dst[0] = static_cast(v >> 0); dst[1] = static_cast(v >> 8); @@ -113,7 +159,7 @@ float u32_to_float(uint32_t bits) { bool write_u32_le(std::FILE * f, uint32_t v) { uint8_t bytes[4]; put_u32_le(bytes, v); - return std::fwrite(bytes, 1, sizeof(bytes), f) == sizeof(bytes); + return write_bytes(f, bytes, sizeof(bytes)); } bool read_u32_le(std::ifstream & f, uint32_t & v) { @@ -172,7 +218,7 @@ bool read_u64_le_crc(std::ifstream & f, uint64_t & v, uint32_t & crc) { bool write_u32_le_crc(std::FILE * f, uint32_t v, uint32_t & crc) { uint8_t bytes[4]; put_u32_le(bytes, v); - if (std::fwrite(bytes, 1, sizeof(bytes), f) != sizeof(bytes)) { + if (!write_bytes(f, bytes, sizeof(bytes))) { return false; } crc = crc32c_update(crc, bytes, sizeof(bytes)); @@ -182,7 +228,7 @@ bool write_u32_le_crc(std::FILE * f, uint32_t v, uint32_t & crc) { bool write_u64_le_crc(std::FILE * f, uint64_t v, uint32_t & crc) { uint8_t bytes[8]; put_u64_le(bytes, v); - if (std::fwrite(bytes, 1, sizeof(bytes), f) != sizeof(bytes)) { + if (!write_bytes(f, bytes, sizeof(bytes))) { return false; } crc = crc32c_update(crc, bytes, sizeof(bytes)); @@ -434,6 +480,18 @@ bool rename_overwrite(const TempFile & temp, const char * dst) { } // namespace +#ifdef GGML_VEC_INDEX_TEST_HOOKS +extern "C" { +void ggml_vec_index_test_set_oom_countdown(int64_t countdown) { + g_test_oom_countdown.store(countdown); +} + +void ggml_vec_index_test_set_write_fail_after(int64_t bytes) { + g_test_write_fail_after.store(bytes); +} +} +#endif + // Lifetime-managed instance state. Lives behind the opaque // `ggml_vec_index_t` typedef. struct ggml_vec_index { @@ -557,6 +615,7 @@ int ggml_vec_index_add( // Atomic add: detect duplicates first (against existing AND in-batch), // bail before mutating any state. + test_maybe_throw_bad_alloc(); std::unordered_set batch_ids; batch_ids.reserve(static_cast(n)); for (int i = 0; i < n; ++i) { @@ -590,6 +649,7 @@ int ggml_vec_index_add( idx->data.resize(new_slots * dim_sz); } idx->slot_to_id.resize(new_slots); + test_maybe_throw_bad_alloc(); idx->id_to_slot.reserve(new_slots); for (int i = 0; i < n; ++i) { @@ -608,6 +668,7 @@ int ggml_vec_index_add( dim_sz * sizeof(float)); } idx->slot_to_id[slot] = ids[i]; + test_maybe_throw_bad_alloc(); idx->id_to_slot.emplace(ids[i], slot); } } catch (const std::bad_alloc &) { @@ -757,6 +818,7 @@ void search_one( const int dim = idx.dim; const size_t n_slots = idx.slot_to_id.size(); + test_maybe_throw_bad_alloc(); std::priority_queue, MinHeapCmp> heap; for (size_t slot = 0; slot < n_slots; ++slot) { @@ -873,10 +935,12 @@ int ggml_vec_index_write(ggml_vec_index_t * idx, const char * path) { return GGML_VEC_INDEX_E_INTERNAL; } + test_maybe_throw_bad_alloc(); TempFile temp; if (!open_temp_file(path, temp)) { return GGML_VEC_INDEX_E_IO; } + test_maybe_throw_bad_alloc(); auto fail_io = [&]() { if (temp.stream != nullptr) { std::fclose(temp.stream); @@ -903,7 +967,7 @@ int ggml_vec_index_write(ggml_vec_index_t * idx, const char * path) { put_u32_le(header + 24, is_q8(*idx) ? 1u : 4u); put_u32_le(header + 28, 0); - if (std::fwrite(header, 1, sizeof(header), f) != sizeof(header)) { + if (!write_bytes(f, header, sizeof(header))) { return fail_io(); } @@ -919,11 +983,7 @@ int ggml_vec_index_write(ggml_vec_index_t * idx, const char * path) { } if (!idx->q8_data.empty()) { - if (std::fwrite( - idx->q8_data.data(), - sizeof(int8_t), - idx->q8_data.size(), - f) != idx->q8_data.size()) { + if (!write_bytes(f, idx->q8_data.data(), idx->q8_data.size())) { return fail_io(); } vectors_crc = @@ -1097,6 +1157,7 @@ ggml_vec_index_t * ggml_vec_index_load(const char * path) { return nullptr; } + test_maybe_throw_bad_alloc(); if (is_q8(*idx)) { idx->q8_data.resize(n * dim_sz); idx->q8_scale.resize(n); diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 36fa743c8716..0c74f0a5ac2a 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -153,6 +153,19 @@ if (LLAMA_TESTS_INSTALL) endif() llama_test(test-vector-index) +add_executable( + test-vector-index-faults + test-vector-index-faults.cpp + ${PROJECT_SOURCE_DIR}/ggml/src/ggml-vector-index.cpp +) +target_compile_features(test-vector-index-faults PRIVATE cxx_std_17) +target_compile_definitions(test-vector-index-faults PRIVATE GGML_VEC_INDEX_TEST_HOOKS) +target_include_directories(test-vector-index-faults PRIVATE ${PROJECT_SOURCE_DIR}/ggml/include) +if (NOT WIN32) + target_link_libraries(test-vector-index-faults PRIVATE m) +endif() +llama_test(test-vector-index-faults) + add_executable(bench-vector-index bench-vector-index.cpp) target_compile_features(bench-vector-index PRIVATE cxx_std_17) target_link_libraries(bench-vector-index PRIVATE ggml) diff --git a/tests/test-vector-index-faults.cpp b/tests/test-vector-index-faults.cpp new file mode 100644 index 000000000000..99944e30d1f0 --- /dev/null +++ b/tests/test-vector-index-faults.cpp @@ -0,0 +1,163 @@ +#include "ggml-vector-index.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +extern "C" void ggml_vec_index_test_set_oom_countdown(int64_t countdown); +extern "C" void ggml_vec_index_test_set_write_fail_after(int64_t bytes); + +namespace { + +#define CHECK(cond) \ + do { \ + if (!(cond)) { \ + std::fprintf(stderr, "FAIL %s:%d: %s\n", __FILE__, __LINE__, #cond);\ + std::exit(1); \ + } \ + } while (0) + +void reset_fault_hooks() { + ggml_vec_index_test_set_oom_countdown(-1); + ggml_vec_index_test_set_write_fail_after(-1); +} + +std::vector read_file_bytes(const std::string & path) { + std::ifstream f(path, std::ios::binary); + CHECK(f.is_open()); + const auto size = std::filesystem::file_size(path); + std::vector bytes(static_cast(size)); + if (!bytes.empty()) { + f.read(reinterpret_cast(bytes.data()), static_cast(bytes.size())); + CHECK(f.gcount() == static_cast(bytes.size())); + } + return bytes; +} + +std::string temp_file_prefix(const std::string & path) { + return std::filesystem::path(path).filename().string() + ".tmp."; +} + +void remove_temp_siblings(const std::string & path) { + const std::filesystem::path p(path); + const std::filesystem::path parent = + p.parent_path().empty() ? std::filesystem::path(".") : p.parent_path(); + const std::string prefix = temp_file_prefix(path); + for (const auto & entry : std::filesystem::directory_iterator(parent)) { + const std::string name = entry.path().filename().string(); + if (name.compare(0, prefix.size(), prefix) == 0) { + std::filesystem::remove(entry.path()); + } + } +} + +void expect_no_temp_siblings(const std::string & path) { + const std::filesystem::path p(path); + const std::filesystem::path parent = + p.parent_path().empty() ? std::filesystem::path(".") : p.parent_path(); + const std::string prefix = temp_file_prefix(path); + for (const auto & entry : std::filesystem::directory_iterator(parent)) { + const std::string name = entry.path().filename().string(); + CHECK(name.compare(0, prefix.size(), prefix) != 0); + } +} + +} // namespace + +int main() { + constexpr int dim = 4; + const std::array base_vectors = { + 1.0f, 0.0f, 0.0f, 0.0f, + 0.0f, 1.0f, 0.0f, 0.0f, + }; + const std::array base_ids = { 101, 102 }; + + auto * idx = ggml_vec_index_create(dim, /*bit_width=*/32); + CHECK(idx != nullptr); + CHECK(ggml_vec_index_add( + idx, base_vectors.data(), 2, base_ids.data()) == GGML_VEC_INDEX_OK); + + { + const std::array new_vectors = { + 0.0f, 0.0f, 1.0f, 0.0f, + 0.0f, 0.0f, 0.0f, 1.0f, + }; + const std::array new_ids = { 201, 202 }; + + ggml_vec_index_test_set_oom_countdown(0); + CHECK(ggml_vec_index_add( + idx, new_vectors.data(), 2, new_ids.data()) == GGML_VEC_INDEX_E_OOM); + reset_fault_hooks(); + CHECK(ggml_vec_index_len(idx) == 2); + + // Fail before the second map insertion, after the first was committed. + ggml_vec_index_test_set_oom_countdown(3); + CHECK(ggml_vec_index_add( + idx, new_vectors.data(), 2, new_ids.data()) == GGML_VEC_INDEX_E_OOM); + reset_fault_hooks(); + CHECK(ggml_vec_index_len(idx) == 2); + CHECK(ggml_vec_index_contains(idx, new_ids[0]) == 0); + CHECK(ggml_vec_index_contains(idx, new_ids[1]) == 0); + } + + { + std::array scores{}; + std::array ids{}; + ggml_vec_index_test_set_oom_countdown(0); + CHECK(ggml_vec_index_search( + idx, base_vectors.data(), 1, /*k=*/1, + scores.data(), ids.data()) == GGML_VEC_INDEX_E_OOM); + reset_fault_hooks(); + } + + const std::string path = + (std::filesystem::temp_directory_path() / + "ggml-vector-index-fault-test.tvim").string(); + std::filesystem::remove(path); + remove_temp_siblings(path); + CHECK(ggml_vec_index_write(idx, path.c_str()) == GGML_VEC_INDEX_OK); + const std::vector old_snapshot = read_file_bytes(path); + + const std::array extra_vector = { 0.5f, 0.5f, 0.5f, 0.5f }; + const uint64_t extra_id = 301; + CHECK(ggml_vec_index_add( + idx, extra_vector.data(), 1, &extra_id) == GGML_VEC_INDEX_OK); + + ggml_vec_index_test_set_oom_countdown(0); + CHECK(ggml_vec_index_write(idx, path.c_str()) == GGML_VEC_INDEX_E_OOM); + reset_fault_hooks(); + CHECK(read_file_bytes(path) == old_snapshot); + expect_no_temp_siblings(path); + + // The first checkpoint passes; the second fails after temp creation. + ggml_vec_index_test_set_oom_countdown(1); + CHECK(ggml_vec_index_write(idx, path.c_str()) == GGML_VEC_INDEX_E_OOM); + reset_fault_hooks(); + CHECK(read_file_bytes(path) == old_snapshot); + expect_no_temp_siblings(path); + + ggml_vec_index_test_set_write_fail_after(40); + CHECK(ggml_vec_index_write(idx, path.c_str()) == GGML_VEC_INDEX_E_IO); + reset_fault_hooks(); + CHECK(read_file_bytes(path) == old_snapshot); + expect_no_temp_siblings(path); + + ggml_vec_index_test_set_oom_countdown(0); + CHECK(ggml_vec_index_load(path.c_str()) == nullptr); + reset_fault_hooks(); + + auto * loaded = ggml_vec_index_load(path.c_str()); + CHECK(loaded != nullptr); + CHECK(ggml_vec_index_len(loaded) == 2); + ggml_vec_index_free(loaded); + ggml_vec_index_free(idx); + std::filesystem::remove(path); + + std::printf("test-vector-index-faults: OK\n"); + return 0; +} From 3a928cf584e0d1b3e98323c422f0d15218bf7b6b Mon Sep 17 00:00:00 2001 From: Nidhin Date: Mon, 13 Jul 2026 17:40:53 +0530 Subject: [PATCH 04/13] ggml : add AVX2 q8 vector index search Adds runtime-dispatched AVX2 acceleration for q8 vector index search on x86 while preserving scalar fallback and the existing ARM NEON path. Assisted-by: GPT-5.5 --- ggml/include/ggml-vector-index.h | 3 +- ggml/src/ggml-vector-index.cpp | 72 +++++++++++++++++++++++++++++++- tests/bench-vector-index.cpp | 5 +++ 3 files changed, 77 insertions(+), 3 deletions(-) diff --git a/ggml/include/ggml-vector-index.h b/ggml/include/ggml-vector-index.h index 75ea28f8241a..433136863a6a 100644 --- a/ggml/include/ggml-vector-index.h +++ b/ggml/include/ggml-vector-index.h @@ -5,7 +5,8 @@ // This is the public C API for fabric's vector index. The implementation // under `ggml/src/ggml-vector-index.cpp` supports full f32 storage // (`bit_width=32`) and production q8 storage (`bit_width=8`) with CPU search -// directly against quantized codes. ARM builds use NEON when available. +// directly against quantized codes. ARM builds use NEON when available; x86 +// builds use AVX2 at runtime when supported. // // Threading: instances are NOT thread-safe. Callers must serialize access // to a given handle. Multiple handles can be used concurrently. diff --git a/ggml/src/ggml-vector-index.cpp b/ggml/src/ggml-vector-index.cpp index 059050df7770..baef4c8ccb6d 100644 --- a/ggml/src/ggml-vector-index.cpp +++ b/ggml/src/ggml-vector-index.cpp @@ -6,8 +6,8 @@ // the slot->id reverse map. Remove uses swap-with-last. // // Search: dot product across all slots + min-heap of size k. q8 search scores -// directly against stored codes and per-vector scales, with ARM NEON when -// available and a scalar fallback. +// directly against stored codes and per-vector scales, with ARM NEON or x86 +// AVX2 when available and a scalar fallback. #include "ggml-vector-index.h" @@ -39,6 +39,19 @@ #define GGML_VEC_INDEX_USE_NEON 0 #endif +#if !GGML_VEC_INDEX_USE_NEON && (defined(__AVX2__) || ((defined(__GNUC__) || defined(__clang__)) && (defined(__x86_64__) || defined(__i386__)))) +#include +#define GGML_VEC_INDEX_USE_AVX2 1 +#if defined(__AVX2__) +#define GGML_VEC_INDEX_AVX2_ATTR +#else +#define GGML_VEC_INDEX_AVX2_ATTR __attribute__((target("avx2"))) +#endif +#else +#define GGML_VEC_INDEX_USE_AVX2 0 +#define GGML_VEC_INDEX_AVX2_ATTR +#endif + #ifndef _WIN32 #include #include @@ -798,9 +811,64 @@ inline float dot_q8_neon(const float * query, const int8_t * codes, float scale, #endif +#if GGML_VEC_INDEX_USE_AVX2 + +GGML_VEC_INDEX_AVX2_ATTR inline float horizontal_sum_avx2(__m256 v) { + const __m128 lo = _mm256_castps256_ps128(v); + const __m128 hi = _mm256_extractf128_ps(v, 1); + __m128 sum = _mm_add_ps(lo, hi); + sum = _mm_hadd_ps(sum, sum); + sum = _mm_hadd_ps(sum, sum); + return _mm_cvtss_f32(sum); +} + +GGML_VEC_INDEX_AVX2_ATTR inline float dot_q8_avx2( + const float * query, + const int8_t * codes, + float scale, + int dim) { + const __m256 scale_v = _mm256_set1_ps(scale); + __m256 acc_v = _mm256_setzero_ps(); + + int i = 0; + for (; i + 8 <= dim; i += 8) { + const __m128i q8 = _mm_loadl_epi64(reinterpret_cast(codes + i)); + const __m256 q = _mm256_mul_ps( + _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(q8)), + scale_v); + acc_v = _mm256_add_ps(acc_v, _mm256_mul_ps(_mm256_loadu_ps(query + i), q)); + } + + float acc = horizontal_sum_avx2(acc_v); + for (; i < dim; ++i) { + const float value = static_cast(codes[i]) * scale; + acc += query[i] * value; + } + return acc; +} + +bool cpu_has_avx2() { +#if defined(__AVX2__) + return true; +#elif defined(__GNUC__) || defined(__clang__) + __builtin_cpu_init(); + return __builtin_cpu_supports("avx2"); +#else + return false; +#endif +} + +#endif + inline float dot_q8(const float * query, const int8_t * codes, float scale, int dim) { #if GGML_VEC_INDEX_USE_NEON return dot_q8_neon(query, codes, scale, dim); +#elif GGML_VEC_INDEX_USE_AVX2 + static const bool has_avx2 = cpu_has_avx2(); + if (has_avx2) { + return dot_q8_avx2(query, codes, scale, dim); + } + return dot_q8_scalar(query, codes, scale, dim); #else return dot_q8_scalar(query, codes, scale, dim); #endif diff --git a/tests/bench-vector-index.cpp b/tests/bench-vector-index.cpp index b588466f97ed..d36a1f7e0b1b 100644 --- a/tests/bench-vector-index.cpp +++ b/tests/bench-vector-index.cpp @@ -42,6 +42,11 @@ struct TimedSearch { const char * q8_kernel_name() { #if defined(__ARM_NEON) || defined(__ARM_NEON__) return "arm-neon"; +#elif defined(__AVX2__) + return "avx2"; +#elif (defined(__GNUC__) || defined(__clang__)) && (defined(__x86_64__) || defined(__i386__)) + __builtin_cpu_init(); + return __builtin_cpu_supports("avx2") ? "avx2" : "scalar"; #else return "scalar"; #endif From 44a6ec425e3924e47c7e0a2386af2eaff1040bcf Mon Sep 17 00:00:00 2001 From: Nidhin Date: Tue, 14 Jul 2026 18:12:51 +0530 Subject: [PATCH 05/13] ggml : add production vector index features Adds quantized q8/q4 storage, crash-safe persistence with checksums, filtered and prepared-filtered search, delta logs with compaction, IVF-flat ANN search, SIMD q8/q4 kernels, and thread-safe concurrent reads. Assisted-by: GPT-5.5 --- ggml/include/ggml-vector-index.h | 153 ++- ggml/src/ggml-vector-index.cpp | 1716 ++++++++++++++++++++++++++-- tests/bench-vector-index.cpp | 498 +++++++- tests/test-vector-index-faults.cpp | 53 + tests/test-vector-index.cpp | 582 +++++++++- 5 files changed, 2836 insertions(+), 166 deletions(-) diff --git a/ggml/include/ggml-vector-index.h b/ggml/include/ggml-vector-index.h index 433136863a6a..52da1f930eaf 100644 --- a/ggml/include/ggml-vector-index.h +++ b/ggml/include/ggml-vector-index.h @@ -4,12 +4,16 @@ // // This is the public C API for fabric's vector index. The implementation // under `ggml/src/ggml-vector-index.cpp` supports full f32 storage -// (`bit_width=32`) and production q8 storage (`bit_width=8`) with CPU search -// directly against quantized codes. ARM builds use NEON when available; x86 -// builds use AVX2 at runtime when supported. +// (`bit_width=32`), production q8 storage (`bit_width=8`), and packed q4 +// storage (`bit_width=4`) with CPU search directly against quantized codes. +// q8 uses NEON/AVX2 when available; q4 uses NEON when available. // -// Threading: instances are NOT thread-safe. Callers must serialize access -// to a given handle. Multiple handles can be used concurrently. +// Threading: read-only APIs on the same handle can run concurrently. Mutations, +// persistence writes, compaction, and IVF builds are serialized with reads and +// with each other. The caller must still keep the handle alive for the duration +// of every API call. Prepared filter handles must also remain alive for the +// full duration of any `ggml_vec_index_search_prepared_filtered` call using +// them; do not free a filter concurrently with a search that uses it. // // Endianness: persistence format is fixed little-endian. @@ -25,6 +29,11 @@ extern "C" { struct ggml_vec_index; typedef struct ggml_vec_index ggml_vec_index_t; +// Prepared filtered-search handle. Valid only for the index generation it was +// created from; any successful add/remove invalidates existing filters. +struct ggml_vec_index_filter; +typedef struct ggml_vec_index_filter ggml_vec_index_filter_t; + // Error codes returned from int-valued APIs. 0 = OK. Negative = failure. // `_remove` is the exception: it returns 1 on removal and 0 on miss. enum ggml_vec_index_error { @@ -41,9 +50,9 @@ enum ggml_vec_index_error { // Lifecycle. // -// `dim` must be > 0. `bit_width` must be 8 or 32. `bit_width=8` stores -// per-vector symmetric q8 codes with one f32 scale per vector. `bit_width=32` -// stores full f32 vectors. Returns NULL on bad args. +// `dim` must be > 0. `bit_width` must be 4, 8, or 32. `bit_width=4` and +// `bit_width=8` store per-vector symmetric quantized codes with one f32 scale +// per vector. `bit_width=32` stores full f32 vectors. Returns NULL on bad args. GGML_API ggml_vec_index_t * ggml_vec_index_create(int dim, int bit_width); GGML_API void ggml_vec_index_free(ggml_vec_index_t * idx); @@ -66,13 +75,38 @@ GGML_API int ggml_vec_index_add( // negative on error. GGML_API int ggml_vec_index_remove(ggml_vec_index_t * idx, uint64_t id); +// Logged mutations for incremental persistence. These update `idx` and append +// a durable delta record to `delta_path`. Replay the log on top of a full .tvim +// snapshot with `ggml_vec_index_load_with_delta`. +GGML_API int ggml_vec_index_add_logged( + ggml_vec_index_t * idx, + const float * vectors, + int n, + const uint64_t * ids, + const char * delta_path); + +GGML_API int ggml_vec_index_remove_logged( + ggml_vec_index_t * idx, + uint64_t id, + const char * delta_path); + // Returns 1 if the id is in the index, 0 otherwise. Read-only. GGML_API int ggml_vec_index_contains(const ggml_vec_index_t * idx, uint64_t id); // Placeholder for cache warming / codebook resolution after a bulk add. -// Currently a no-op. +// Currently a no-op. Use `ggml_vec_index_build_ivf` to build ANN state. GGML_API void ggml_vec_index_prepare(ggml_vec_index_t * idx); +// Builds an in-memory IVF-flat approximate nearest-neighbor structure. This is +// not persisted in .tvim files; call again after loading if ANN search is +// needed. Successful add/remove calls invalidate the IVF structure. +// `n_lists` is capped to the current index length. `n_iter` controls centroid +// refinement; 0 uses deterministic initial centroids only. +GGML_API int ggml_vec_index_build_ivf( + ggml_vec_index_t * idx, + int n_lists, + int n_iter); + // Top-k search. `queries` is `n_q * dim` row-major. `out_scores` and // `out_ids` are caller-allocated buffers of size `n_q * k`. Each row is // sorted descending by score (higher = closer / more similar). If the index @@ -81,8 +115,8 @@ GGML_API void ggml_vec_index_prepare(ggml_vec_index_t * idx); // against the index (does not mutate state). // // Score semantics: dot product. For f32 storage this is a full-precision dot -// product. For q8 storage, the query remains f32 and each indexed component is -// scored as `q8_code * per_vector_scale` without expanding the stored matrix +// product. For q4/q8 storage, the query remains f32 and each indexed component +// is scored as `q_code * per_vector_scale` without expanding the stored matrix // back to f32. Callers that want cosine similarity must L2-normalize their // vectors before insert AND before query; the index does NOT normalize // internally. All query components must be finite. @@ -94,6 +128,52 @@ GGML_API int ggml_vec_index_search( float * out_scores, uint64_t * out_ids); +// Filtered top-k search. Only entries whose ids appear in `allowed_ids` are +// considered. Missing ids are ignored; duplicate filter ids are treated once. +// `allowed_ids` may be NULL only when `n_allowed == 0`, which produces only +// sentinel results. The same filter is applied to every query row. +GGML_API int ggml_vec_index_search_filtered( + const ggml_vec_index_t * idx, + const float * queries, + int n_q, + int k, + const uint64_t * allowed_ids, + int n_allowed, + float * out_scores, + uint64_t * out_ids); + +// Prepared filtered search. Creating a filter maps, sorts, and deduplicates +// `allowed_ids` once, so callers can reuse it for repeated searches over the +// same allowlist. Stale filters return GGML_VEC_INDEX_E_INVALID_ARG. +GGML_API ggml_vec_index_filter_t * ggml_vec_index_filter_create( + const ggml_vec_index_t * idx, + const uint64_t * allowed_ids, + int n_allowed); + +GGML_API void ggml_vec_index_filter_free(ggml_vec_index_filter_t * filter); + +GGML_API int ggml_vec_index_search_prepared_filtered( + const ggml_vec_index_t * idx, + const ggml_vec_index_filter_t * filter, + const float * queries, + int n_q, + int k, + float * out_scores, + uint64_t * out_ids); + +// IVF-flat ANN top-k search. `ggml_vec_index_build_ivf` must have been called +// after the most recent mutation. `nprobe` controls how many centroid lists are +// searched; higher values improve recall and lower the latency win. If nprobe +// is greater than the number of built lists, all lists are searched. +GGML_API int ggml_vec_index_search_ivf( + const ggml_vec_index_t * idx, + const float * queries, + int n_q, + int k, + int nprobe, + float * out_scores, + uint64_t * out_ids); + // Persistence. Format is .tvim version 2; see bottom of this header. GGML_API int ggml_vec_index_write( ggml_vec_index_t * idx, @@ -104,6 +184,19 @@ GGML_API int ggml_vec_index_write( // Returns NULL on failure. GGML_API ggml_vec_index_t * ggml_vec_index_load(const char * path); +// Loads a full .tvim snapshot and replays an append-only delta log. Missing +// delta logs are treated as empty. +GGML_API ggml_vec_index_t * ggml_vec_index_load_with_delta( + const char * snapshot_path, + const char * delta_path); + +// Compacts incremental persistence by writing a full snapshot and replacing +// the delta log with an empty matching .tvid header. +GGML_API int ggml_vec_index_compact_delta( + ggml_vec_index_t * idx, + const char * snapshot_path, + const char * delta_path); + // Stats. GGML_API int ggml_vec_index_len(const ggml_vec_index_t * idx); GGML_API int ggml_vec_index_dim(const ggml_vec_index_t * idx); @@ -115,21 +208,22 @@ GGML_API int ggml_vec_index_bit_width(const ggml_vec_index_t * idx); // ------ ----- ------------------------------------------------------- // 0 4 magic = "TVPI" (bytes 0x54, 0x56, 0x50, 0x49) // 4 1 version = 2 -// 5 1 bit_width (8 or 32) -// 6 1 storage kind (1 = f32, 2 = q8) +// 5 1 bit_width (4, 8, or 32) +// 6 1 storage kind (1 = f32, 2 = q8, 3 = q4) // 7 1 flags (bit 0 = checksum trailer present) // 8 4 dim (uint32) // 12 4 n_vectors (uint32) // 16 4 qparam_type (0 = none, 1 = per-vector f32 scale) // 20 4 qparam_bytes_per_vector (0 or 4) -// 24 4 bytes_per_component (1 or 4) +// 24 4 bytes_per_component (0 for packed q4, 1 for q8, 4 for f32) // 28 4 reserved (zero) // 32 ... qparams: // - f32: empty -// - q8: N float32 scales +// - q4/q8: N float32 scales // ... ... vectors: // - f32: N*D float32 values, row-major // - q8: N*D int8 codes, row-major +// - q4: N*ceil(D/2) packed unsigned nibbles, row-major // ... N*8 ids (uint64) // ... 4 header CRC32C, when flag bit 0 is set // ... 4 qparams CRC32C, when flag bit 0 is set @@ -138,13 +232,40 @@ GGML_API int ggml_vec_index_bit_width(const ggml_vec_index_t * idx); // // Where N = n_vectors and D = dim. q8 uses symmetric per-vector quantization: // scale = max(abs(v)) / 127, code = round(v / scale) clamped to [-127, 127]. -// Zero vectors use scale = 1 and all-zero codes. Each CRC32C covers exactly its +// q4 uses scale = max(abs(v)) / 7, code = round(v / scale) clamped to [-7, 7], +// stored as unsigned nibble `code + 8` (0 is invalid). Zero vectors use +// scale = 1 and all-zero dequantized codes. Each CRC32C covers exactly its // corresponding serialized section; the header CRC covers bytes [0, 32), and // the CRC32C of an empty section is zero. // Legacy v2 files with flags=0 and no checksum trailer remain readable. // Writers emit checksummed v2 files. Readers reject unknown versions and v2 // flag bits; they also accept legacy v1 f32 snapshots. Legacy bit_width=8 // snapshots migrate to q8, while all other legacy widths migrate to f32. +// +// Delta log (.tvid version 1, all little-endian): +// +// file header: +// 0 4 magic = "TVDL" +// 4 1 version = 1 +// 5 1 bit_width (4, 8, or 32) +// 6 2 reserved (zero) +// 8 4 dim (uint32) +// 12 4 base snapshot state CRC32C +// +// record header: +// 0 1 op (1 = add, 2 = remove) +// 1 3 reserved (zero) +// 4 4 n (add count; remove uses 1) +// 8 8 payload bytes +// 16 4 CRC32C over record header bytes [0, 16), state CRC, and payload +// 20 4 state CRC32C after applying this record +// +// add payload: N uint64 ids, then N*D float32 vectors +// remove payload: one uint64 id +// +// The base snapshot CRC binds the log to the snapshot state it extends. Record +// state CRCs let loading validate the final replay state and recognize a +// compacted snapshot when a process crashed before replacing the old delta log. #ifdef __cplusplus } diff --git a/ggml/src/ggml-vector-index.cpp b/ggml/src/ggml-vector-index.cpp index baef4c8ccb6d..3291c2d42aa3 100644 --- a/ggml/src/ggml-vector-index.cpp +++ b/ggml/src/ggml-vector-index.cpp @@ -25,8 +25,10 @@ #include #include #include +#include #include #include +#include #include #include #include @@ -81,12 +83,19 @@ constexpr uint8_t kTvimVersionV1 = 1; constexpr uint8_t kTvimVersion = 2; constexpr uint8_t kStorageF32 = 1; constexpr uint8_t kStorageQ8 = 2; +constexpr uint8_t kStorageQ4 = 3; constexpr uint8_t kFlagCRC32C = 1; constexpr uint32_t kQParamNone = 0; constexpr uint32_t kQParamScaleF32 = 1; constexpr size_t kTvimV1HeaderSize = 16; constexpr size_t kTvimHeaderSize = 32; constexpr size_t kTvimChecksumSize = 16; +constexpr uint8_t kTvidMagic[4] = { 'T', 'V', 'D', 'L' }; +constexpr uint8_t kTvidVersion = 1; +constexpr uint8_t kTvidOpAdd = 1; +constexpr uint8_t kTvidOpRemove = 2; +constexpr size_t kTvidHeaderSize = 16; +constexpr size_t kTvidRecordHeaderSize = 24; static_assert(sizeof(float) == sizeof(uint32_t), "ggml-vector-index requires float32"); @@ -263,7 +272,7 @@ struct MinHeapCmp { }; bool is_supported_bit_width(int bit_width) { - return bit_width == 8 || bit_width == 32; + return bit_width == 4 || bit_width == 8 || bit_width == 32; } bool all_finite(const float * values, size_t n) { @@ -302,9 +311,20 @@ bool expected_file_size( uint64_t components = 0; uint64_t ids = 0; uint64_t total = header_size; + if (component_bytes == 0) { + uint64_t row_bytes = 0; + if (!checked_add_u64(dim, 1, row_bytes)) { + return false; + } + row_bytes /= 2; + if (!checked_mul_u64(n, row_bytes, components)) { + return false; + } + } else if (!checked_mul_u64(n, dim, components) || + !checked_mul_u64(components, component_bytes, components)) { + return false; + } if (!checked_mul_u64(n, qparam_bytes, qparams) || - !checked_mul_u64(n, dim, components) || - !checked_mul_u64(components, component_bytes, components) || !checked_mul_u64(n, sizeof(uint64_t), ids) || !checked_add_u64(total, qparams, total) || !checked_add_u64(total, components, total) || @@ -508,8 +528,11 @@ void ggml_vec_index_test_set_write_fail_after(int64_t bytes) { // Lifetime-managed instance state. Lives behind the opaque // `ggml_vec_index_t` typedef. struct ggml_vec_index { + mutable std::shared_mutex mutex; + int dim = 0; int bit_width = 32; + uint64_t generation = 0; // Flat row-major f32 storage for bit_width=32. std::vector data; @@ -518,19 +541,63 @@ struct ggml_vec_index { std::vector q8_data; std::vector q8_scale; + // Packed row-major q4 storage for bit_width=4 plus one scale per vector. + std::vector q4_data; + std::vector q4_scale; + // slot -> external id (parallel to logical slot index). std::vector slot_to_id; // external id -> slot. std::unordered_map id_to_slot; + + // In-memory IVF-flat ANN structure. Rebuilt explicitly after mutations. + uint64_t ivf_generation = std::numeric_limits::max(); + int ivf_n_lists = 0; + std::vector ivf_centroids; + std::vector> ivf_lists; +}; + +struct ggml_vec_index_filter { + int dim = 0; + int bit_width = 32; + uint64_t generation = 0; + std::vector slots; }; +static void invalidate_ivf(ggml_vec_index & idx) { + idx.ivf_generation = std::numeric_limits::max(); + idx.ivf_n_lists = 0; + idx.ivf_centroids.clear(); + idx.ivf_lists.clear(); +} + static bool is_q8(const ggml_vec_index & idx) { return idx.bit_width == 8; } +static bool is_q4(const ggml_vec_index & idx) { + return idx.bit_width == 4; +} + +static bool is_quantized(const ggml_vec_index & idx) { + return is_q4(idx) || is_q8(idx); +} + static uint8_t storage_kind(const ggml_vec_index & idx) { - return is_q8(idx) ? kStorageQ8 : kStorageF32; + return is_q4(idx) ? kStorageQ4 : (is_q8(idx) ? kStorageQ8 : kStorageF32); +} + +static size_t q4_row_bytes(size_t dim) { + return (dim + 1) / 2; +} + +static uint8_t q4_encode(int q) { + return static_cast(q + 8); +} + +static int q4_decode(uint8_t nibble) { + return static_cast(nibble) - 8; } static void quantize_q8_row(const float * src, int8_t * dst, int dim, float & scale) { @@ -557,6 +624,515 @@ static void quantize_q8_row(const float * src, int8_t * dst, int dim, float & sc } } +static void quantize_q4_row(const float * src, uint8_t * dst, int dim, float & scale) { + float max_abs = 0.0f; + for (int i = 0; i < dim; ++i) { + max_abs = std::max(max_abs, std::fabs(src[i])); + } + + std::memset(dst, 0x88, q4_row_bytes(static_cast(dim))); + if (max_abs == 0.0f) { + scale = 1.0f; + return; + } + + scale = max_abs / 7.0f; + if (scale == 0.0f) { + scale = max_abs; + } + for (int i = 0; i < dim; ++i) { + const float scaled = src[i] / scale; + int q = static_cast(std::nearbyint(scaled)); + q = std::max(-7, std::min(7, q)); + const uint8_t code = q4_encode(q); + uint8_t & byte = dst[static_cast(i) / 2]; + if ((i & 1) == 0) { + byte = static_cast((byte & 0xf0u) | code); + } else { + byte = static_cast((byte & 0x0fu) | (code << 4)); + } + } +} + +namespace { + +uint32_t crc32c_update_u32(uint32_t crc, uint32_t v) { + uint8_t bytes[4]; + put_u32_le(bytes, v); + return crc32c_update(crc, bytes, sizeof(bytes)); +} + +uint32_t crc32c_update_u64(uint32_t crc, uint64_t v) { + uint8_t bytes[8]; + put_u64_le(bytes, v); + return crc32c_update(crc, bytes, sizeof(bytes)); +} + +uint32_t index_state_crc32c(const ggml_vec_index & idx) { + uint32_t crc = 0xffffffffu; + crc = crc32c_update_u32(crc, static_cast(idx.dim)); + crc = crc32c_update_u32(crc, static_cast(idx.bit_width)); + crc = crc32c_update_u32(crc, static_cast(storage_kind(idx))); + crc = crc32c_update_u64(crc, static_cast(idx.slot_to_id.size())); + if (is_q4(idx)) { + for (float scale : idx.q4_scale) { + crc = crc32c_update_u32(crc, float_to_u32(scale)); + } + if (!idx.q4_data.empty()) { + crc = crc32c_update(crc, idx.q4_data.data(), idx.q4_data.size()); + } + } else if (is_q8(idx)) { + for (float scale : idx.q8_scale) { + crc = crc32c_update_u32(crc, float_to_u32(scale)); + } + if (!idx.q8_data.empty()) { + crc = crc32c_update(crc, idx.q8_data.data(), idx.q8_data.size()); + } + } else { + for (float v : idx.data) { + crc = crc32c_update_u32(crc, float_to_u32(v)); + } + } + for (uint64_t id : idx.slot_to_id) { + crc = crc32c_update_u64(crc, id); + } + return crc ^ 0xffffffffu; +} + +uint32_t index_state_crc32c_after_remove(const ggml_vec_index & idx, uint64_t id) { + const auto it = idx.id_to_slot.find(id); + if (it == idx.id_to_slot.end()) { + return index_state_crc32c(idx); + } + const size_t removed = it->second; + const size_t last = idx.slot_to_id.size() - 1; + const size_t dim_sz = static_cast(idx.dim); + + uint32_t crc = 0xffffffffu; + crc = crc32c_update_u32(crc, static_cast(idx.dim)); + crc = crc32c_update_u32(crc, static_cast(idx.bit_width)); + crc = crc32c_update_u32(crc, static_cast(storage_kind(idx))); + crc = crc32c_update_u64(crc, static_cast(last)); + + auto logical_slot = [&](size_t out_slot) { + return out_slot == removed ? last : out_slot; + }; + + if (is_q4(idx)) { + const size_t row_bytes = q4_row_bytes(dim_sz); + for (size_t slot = 0; slot < last; ++slot) { + crc = crc32c_update_u32(crc, float_to_u32(idx.q4_scale[logical_slot(slot)])); + } + for (size_t slot = 0; slot < last; ++slot) { + const size_t src_slot = logical_slot(slot); + crc = crc32c_update( + crc, + idx.q4_data.data() + src_slot * row_bytes, + row_bytes * sizeof(uint8_t)); + } + } else if (is_q8(idx)) { + for (size_t slot = 0; slot < last; ++slot) { + crc = crc32c_update_u32(crc, float_to_u32(idx.q8_scale[logical_slot(slot)])); + } + for (size_t slot = 0; slot < last; ++slot) { + const size_t src_slot = logical_slot(slot); + crc = crc32c_update( + crc, + idx.q8_data.data() + src_slot * dim_sz, + dim_sz * sizeof(int8_t)); + } + } else { + for (size_t slot = 0; slot < last; ++slot) { + const size_t src_slot = logical_slot(slot); + for (size_t i = 0; i < dim_sz; ++i) { + crc = crc32c_update_u32(crc, float_to_u32(idx.data[src_slot * dim_sz + i])); + } + } + } + for (size_t slot = 0; slot < last; ++slot) { + crc = crc32c_update_u64(crc, idx.slot_to_id[logical_slot(slot)]); + } + return crc ^ 0xffffffffu; +} + +bool filesystem_path_from_utf8(const char * path, std::filesystem::path & out) { +#ifdef _WIN32 + std::wstring wide; + if (!utf8_to_wide(path, wide)) { + return false; + } + out = std::filesystem::path(wide); +#else + out = std::filesystem::path(path); +#endif + return true; +} + +bool open_append_file(const char * path, std::FILE ** out) { +#ifdef _WIN32 + std::wstring wide; + if (!utf8_to_wide(path, wide)) { + return false; + } + *out = _wfopen(wide.c_str(), L"a+b"); +#else + *out = std::fopen(path, "a+b"); +#endif + return *out != nullptr; +} + +void fill_delta_header(const ggml_vec_index & idx, uint32_t base_crc, uint8_t * header) { + std::memset(header, 0, kTvidHeaderSize); + std::memcpy(header, kTvidMagic, 4); + header[4] = kTvidVersion; + header[5] = static_cast(idx.bit_width); + put_u32_le(header + 8, static_cast(idx.dim)); + put_u32_le(header + 12, base_crc); +} + +bool validate_delta_header( + const char * path, + const ggml_vec_index & idx, + uint64_t & size, + uint32_t & base_crc) { + std::filesystem::path fs_path; + if (!filesystem_path_from_utf8(path, fs_path)) { + return false; + } + if (!std::filesystem::exists(fs_path)) { + size = 0; + base_crc = index_state_crc32c(idx); + return true; + } + size = static_cast(std::filesystem::file_size(fs_path)); + if (size == 0) { + base_crc = index_state_crc32c(idx); + return true; + } + if (size < kTvidHeaderSize) { + return false; + } + + std::ifstream f(fs_path, std::ios::binary); + if (!f.is_open()) { + return false; + } + uint8_t header[kTvidHeaderSize] = {}; + f.read(reinterpret_cast(header), sizeof(header)); + if (!f) { + return false; + } + if (std::memcmp(header, kTvidMagic, 4) != 0 || + header[4] != kTvidVersion || + header[5] != static_cast(idx.bit_width) || + header[6] != 0 || + header[7] != 0 || + get_u32_le(header + 8) != static_cast(idx.dim)) { + return false; + } + base_crc = get_u32_le(header + 12); + return true; +} + +bool truncate_file_to(const char * path, uint64_t size) { + try { + std::filesystem::path fs_path; + if (!filesystem_path_from_utf8(path, fs_path)) { + return false; + } + std::filesystem::resize_file(fs_path, size); + return true; + } catch (...) { + return false; + } +} + +bool inspect_delta_log_tail( + const char * path, + const ggml_vec_index & idx, + uint32_t & last_state_crc, + uint64_t & complete_size) { + std::filesystem::path fs_path; + if (!filesystem_path_from_utf8(path, fs_path)) { + return false; + } + const uint64_t file_size = static_cast(std::filesystem::file_size(fs_path)); + if (file_size < kTvidHeaderSize) { + return false; + } + + std::ifstream f(fs_path, std::ios::binary); + if (!f.is_open()) { + return false; + } + + uint8_t header[kTvidHeaderSize] = {}; + f.read(reinterpret_cast(header), sizeof(header)); + if (!f || + std::memcmp(header, kTvidMagic, 4) != 0 || + header[4] != kTvidVersion || + header[5] != static_cast(idx.bit_width) || + header[6] != 0 || + header[7] != 0 || + get_u32_le(header + 8) != static_cast(idx.dim)) { + return false; + } + + last_state_crc = get_u32_le(header + 12); + complete_size = kTvidHeaderSize; + uint64_t offset = kTvidHeaderSize; + while (offset < file_size) { + uint8_t record[kTvidRecordHeaderSize] = {}; + f.read(reinterpret_cast(record), sizeof(record)); + if (f.gcount() != static_cast(sizeof(record))) { + return true; + } + offset += kTvidRecordHeaderSize; + + const uint8_t op = record[0]; + const uint64_t payload_bytes = get_u64_le(record + 8); + const uint32_t expected_crc = get_u32_le(record + 16); + const uint32_t state_crc = get_u32_le(record + 20); + if (record[1] != 0 || record[2] != 0 || record[3] != 0 || + (op != kTvidOpAdd && op != kTvidOpRemove)) { + return false; + } + if (payload_bytes > file_size - offset) { + return true; + } + + if (payload_bytes > static_cast(std::numeric_limits::max()) || + payload_bytes > static_cast(std::numeric_limits::max())) { + return false; + } + std::vector payload(static_cast(payload_bytes)); + if (!payload.empty()) { + f.read( + reinterpret_cast(payload.data()), + static_cast(payload.size())); + } + if (!f) { + return true; + } + offset += payload_bytes; + + uint32_t crc = crc32c_update(0xffffffffu, record, 16); + crc = crc32c_update(crc, record + 20, 4); + if (!payload.empty()) { + crc = crc32c_update(crc, payload.data(), payload.size()); + } + if ((crc ^ 0xffffffffu) != expected_crc) { + return false; + } + last_state_crc = state_crc; + complete_size = offset; + } + return true; +} + +int append_delta_record( + const ggml_vec_index & idx, + const char * delta_path, + uint8_t op, + uint32_t n, + uint32_t base_crc_for_new_log, + uint32_t state_crc, + const std::vector & payload) { + if (delta_path == nullptr) { + return GGML_VEC_INDEX_E_INVALID_ARG; + } + uint64_t old_size = 0; + uint32_t existing_base_crc = 0; + if (!validate_delta_header(delta_path, idx, old_size, existing_base_crc)) { + return GGML_VEC_INDEX_E_IO; + } + if (old_size != 0) { + uint32_t tail_crc = 0; + uint64_t complete_size = 0; + if (!inspect_delta_log_tail(delta_path, idx, tail_crc, complete_size)) { + return GGML_VEC_INDEX_E_IO; + } + if (tail_crc != base_crc_for_new_log) { + return GGML_VEC_INDEX_E_IO; + } + if (complete_size != old_size) { + if (!truncate_file_to(delta_path, complete_size)) { + return GGML_VEC_INDEX_E_INTERNAL; + } + old_size = complete_size; + } + } + (void) existing_base_crc; + + std::FILE * f = nullptr; + if (!open_append_file(delta_path, &f)) { + return GGML_VEC_INDEX_E_IO; + } + auto close_file = [&]() { + if (f != nullptr) { + std::fclose(f); + f = nullptr; + } + }; + auto fail_io = [&]() { + close_file(); + return truncate_file_to(delta_path, old_size) ? + GGML_VEC_INDEX_E_IO : + GGML_VEC_INDEX_E_INTERNAL; + }; + + if (old_size == 0) { + uint8_t header[kTvidHeaderSize] = {}; + fill_delta_header(idx, base_crc_for_new_log, header); + if (!write_bytes(f, header, sizeof(header))) { + return fail_io(); + } + } + + uint8_t record[kTvidRecordHeaderSize] = {}; + record[0] = op; + put_u32_le(record + 4, n); + put_u64_le(record + 8, static_cast(payload.size())); + put_u32_le(record + 20, state_crc); + uint32_t crc = crc32c_update(0xffffffffu, record, 16); + crc = crc32c_update(crc, record + 20, 4); + if (!payload.empty()) { + crc = crc32c_update(crc, payload.data(), payload.size()); + } + put_u32_le(record + 16, crc ^ 0xffffffffu); + + if (!write_bytes(f, record, sizeof(record)) || + (!payload.empty() && !write_bytes(f, payload.data(), payload.size())) || + !flush_and_sync(f)) { + return fail_io(); + } + const int close_result = std::fclose(f); + f = nullptr; + if (close_result != 0 || !fsync_parent_dir(delta_path)) { + return truncate_file_to(delta_path, old_size) ? + GGML_VEC_INDEX_E_IO : + GGML_VEC_INDEX_E_INTERNAL; + } + return GGML_VEC_INDEX_OK; +} + +int write_empty_delta_log(const ggml_vec_index & idx, const char * delta_path) { + if (delta_path == nullptr) { + return GGML_VEC_INDEX_E_INVALID_ARG; + } + try { + TempFile temp; + if (!open_temp_file(delta_path, temp)) { + return GGML_VEC_INDEX_E_IO; + } + auto fail_io = [&]() { + if (temp.stream != nullptr) { + std::fclose(temp.stream); + temp.stream = nullptr; + } + remove_temp_file(temp); + return GGML_VEC_INDEX_E_IO; + }; + + uint8_t header[kTvidHeaderSize] = {}; + fill_delta_header(idx, index_state_crc32c(idx), header); + if (!write_bytes(temp.stream, header, sizeof(header)) || + !flush_and_sync(temp.stream)) { + return fail_io(); + } + const int close_result = std::fclose(temp.stream); + temp.stream = nullptr; + if (close_result != 0) { + remove_temp_file(temp); + return GGML_VEC_INDEX_E_IO; + } + if (!rename_overwrite(temp, delta_path)) { + return fail_io(); + } + temp.path.clear(); + if (!fsync_parent_dir(delta_path)) { + return GGML_VEC_INDEX_E_IO; + } + return GGML_VEC_INDEX_OK; + } catch (const std::bad_alloc &) { + return GGML_VEC_INDEX_E_OOM; + } catch (...) { + return GGML_VEC_INDEX_E_INTERNAL; + } +} + +bool validate_logged_add_args( + const ggml_vec_index_t * idx, + const float * vectors, + int n, + const uint64_t * ids) { + if (idx == nullptr || vectors == nullptr || ids == nullptr || n < 0) { + return false; + } + if (n == 0) { + return true; + } + const size_t n_sz = static_cast(n); + const size_t dim_sz = static_cast(idx->dim); + if (dim_sz != 0 && n_sz > std::numeric_limits::max() / dim_sz) { + return false; + } + return all_finite(vectors, n_sz * dim_sz); +} + +int check_logged_add_duplicates( + const ggml_vec_index_t * idx, + int n, + const uint64_t * ids) { + std::unordered_set batch_ids; + batch_ids.reserve(static_cast(n)); + for (int i = 0; i < n; ++i) { + if (idx->id_to_slot.find(ids[i]) != idx->id_to_slot.end()) { + return GGML_VEC_INDEX_E_DUPLICATE; + } + if (!batch_ids.insert(ids[i]).second) { + return GGML_VEC_INDEX_E_DUPLICATE; + } + } + return GGML_VEC_INDEX_OK; +} + +bool build_add_delta_payload( + const ggml_vec_index_t * idx, + const float * vectors, + int n, + const uint64_t * ids, + std::vector & payload) { + const size_t n_sz = static_cast(n); + const size_t dim_sz = static_cast(idx->dim); + const size_t id_bytes = n_sz * sizeof(uint64_t); + const size_t vector_count = n_sz * dim_sz; + if (vector_count > (std::numeric_limits::max() - id_bytes) / sizeof(uint32_t)) { + return false; + } + payload.clear(); + payload.reserve(id_bytes + vector_count * sizeof(uint32_t)); + for (int i = 0; i < n; ++i) { + uint8_t bytes[8]; + put_u64_le(bytes, ids[i]); + payload.insert(payload.end(), bytes, bytes + sizeof(bytes)); + } + for (size_t i = 0; i < vector_count; ++i) { + uint8_t bytes[4]; + put_u32_le(bytes, float_to_u32(vectors[i])); + payload.insert(payload.end(), bytes, bytes + sizeof(bytes)); + } + return true; +} + +std::vector build_remove_delta_payload(uint64_t id) { + std::vector payload(sizeof(uint64_t)); + put_u64_le(payload.data(), id); + return payload; +} + +} // namespace + // --------------------------------------------------------------------------- // Lifecycle // --------------------------------------------------------------------------- @@ -589,7 +1165,7 @@ void ggml_vec_index_free(ggml_vec_index_t * idx) { // Mutation // --------------------------------------------------------------------------- -int ggml_vec_index_add( +static int ggml_vec_index_add_unlocked( ggml_vec_index_t * idx, const float * vectors, int n, @@ -606,7 +1182,10 @@ int ggml_vec_index_add( for (int i = 0; i < n; ++i) { idx->id_to_slot.erase(ids[i]); } - if (is_q8(*idx)) { + if (is_q4(*idx)) { + idx->q4_data.resize(base_slot * q4_row_bytes(dim_sz)); + idx->q4_scale.resize(base_slot); + } else if (is_q8(*idx)) { idx->q8_data.resize(base_slot * dim_sz); idx->q8_scale.resize(base_slot); } else { @@ -655,7 +1234,10 @@ int ggml_vec_index_add( } resized = true; - if (is_q8(*idx)) { + if (is_q4(*idx)) { + idx->q4_data.resize(new_slots * q4_row_bytes(dim_sz)); + idx->q4_scale.resize(new_slots); + } else if (is_q8(*idx)) { idx->q8_data.resize(new_slots * dim_sz); idx->q8_scale.resize(new_slots); } else { @@ -668,7 +1250,13 @@ int ggml_vec_index_add( for (int i = 0; i < n; ++i) { const size_t slot = base_slot + static_cast(i); const float * src = vectors + static_cast(i) * dim_sz; - if (is_q8(*idx)) { + if (is_q4(*idx)) { + quantize_q4_row( + src, + idx->q4_data.data() + slot * q4_row_bytes(dim_sz), + idx->dim, + idx->q4_scale[slot]); + } else if (is_q8(*idx)) { quantize_q8_row( src, idx->q8_data.data() + slot * dim_sz, @@ -684,6 +1272,8 @@ int ggml_vec_index_add( test_maybe_throw_bad_alloc(); idx->id_to_slot.emplace(ids[i], slot); } + ++idx->generation; + invalidate_ivf(*idx); } catch (const std::bad_alloc &) { rollback(); return GGML_VEC_INDEX_E_OOM; @@ -694,7 +1284,23 @@ int ggml_vec_index_add( return GGML_VEC_INDEX_OK; } -int ggml_vec_index_remove(ggml_vec_index_t * idx, uint64_t id) { +int ggml_vec_index_add( + ggml_vec_index_t * idx, + const float * vectors, + int n, + const uint64_t * ids) { + if (idx == nullptr) { + return GGML_VEC_INDEX_E_INVALID_ARG; + } + try { + std::unique_lock lock(idx->mutex); + return ggml_vec_index_add_unlocked(idx, vectors, n, ids); + } catch (...) { + return GGML_VEC_INDEX_E_INTERNAL; + } +} + +static int ggml_vec_index_remove_unlocked(ggml_vec_index_t * idx, uint64_t id) { try { if (idx == nullptr) { return GGML_VEC_INDEX_E_INVALID_ARG; @@ -709,7 +1315,14 @@ int ggml_vec_index_remove(ggml_vec_index_t * idx, uint64_t id) { if (slot != last) { // Move last vector into the freed slot and update its id mapping. - if (is_q8(*idx)) { + if (is_q4(*idx)) { + const size_t row_bytes = q4_row_bytes(dim_sz); + std::memcpy( + idx->q4_data.data() + slot * row_bytes, + idx->q4_data.data() + last * row_bytes, + row_bytes * sizeof(uint8_t)); + idx->q4_scale[slot] = idx->q4_scale[last]; + } else if (is_q8(*idx)) { std::memcpy( idx->q8_data.data() + slot * dim_sz, idx->q8_data.data() + last * dim_sz, @@ -727,24 +1340,151 @@ int ggml_vec_index_remove(ggml_vec_index_t * idx, uint64_t id) { } idx->slot_to_id.pop_back(); - if (is_q8(*idx)) { + if (is_q4(*idx)) { + idx->q4_data.resize(last * q4_row_bytes(dim_sz)); + idx->q4_scale.resize(last); + } else if (is_q8(*idx)) { idx->q8_data.resize(last * dim_sz); idx->q8_scale.resize(last); } else { idx->data.resize(last * dim_sz); } idx->id_to_slot.erase(it); + ++idx->generation; + invalidate_ivf(*idx); return 1; } catch (...) { return GGML_VEC_INDEX_E_INTERNAL; } } +int ggml_vec_index_remove(ggml_vec_index_t * idx, uint64_t id) { + if (idx == nullptr) { + return GGML_VEC_INDEX_E_INVALID_ARG; + } + try { + std::unique_lock lock(idx->mutex); + return ggml_vec_index_remove_unlocked(idx, id); + } catch (...) { + return GGML_VEC_INDEX_E_INTERNAL; + } +} + +int ggml_vec_index_add_logged( + ggml_vec_index_t * idx, + const float * vectors, + int n, + const uint64_t * ids, + const char * delta_path) { + bool added = false; + std::unique_lock lock; + try { + if (idx == nullptr || delta_path == nullptr) { + return GGML_VEC_INDEX_E_INVALID_ARG; + } + lock = std::unique_lock(idx->mutex); + if (!validate_logged_add_args(idx, vectors, n, ids)) { + return GGML_VEC_INDEX_E_INVALID_ARG; + } + if (n == 0) { + return GGML_VEC_INDEX_OK; + } + + const int duplicate_status = check_logged_add_duplicates(idx, n, ids); + if (duplicate_status != GGML_VEC_INDEX_OK) { + return duplicate_status; + } + + std::vector payload; + if (!build_add_delta_payload(idx, vectors, n, ids, payload)) { + return GGML_VEC_INDEX_E_INVALID_ARG; + } + + const uint32_t base_crc = index_state_crc32c(*idx); + const int add_status = ggml_vec_index_add_unlocked(idx, vectors, n, ids); + if (add_status != GGML_VEC_INDEX_OK) { + return add_status; + } + added = true; + + const int append_status = append_delta_record( + *idx, + delta_path, + kTvidOpAdd, + static_cast(n), + base_crc, + index_state_crc32c(*idx), + payload); + if (append_status != GGML_VEC_INDEX_OK) { + if (append_status == GGML_VEC_INDEX_E_IO) { + for (int i = 0; i < n; ++i) { + ggml_vec_index_remove_unlocked(idx, ids[i]); + } + } + return append_status; + } + return GGML_VEC_INDEX_OK; + } catch (const std::bad_alloc &) { + if (added) { + for (int i = 0; i < n; ++i) { + ggml_vec_index_remove_unlocked(idx, ids[i]); + } + } + return GGML_VEC_INDEX_E_OOM; + } catch (...) { + if (added) { + for (int i = 0; i < n; ++i) { + ggml_vec_index_remove_unlocked(idx, ids[i]); + } + } + return GGML_VEC_INDEX_E_INTERNAL; + } +} + +int ggml_vec_index_remove_logged( + ggml_vec_index_t * idx, + uint64_t id, + const char * delta_path) { + try { + if (idx == nullptr || delta_path == nullptr) { + return GGML_VEC_INDEX_E_INVALID_ARG; + } + std::unique_lock lock(idx->mutex); + if (idx->id_to_slot.count(id) == 0) { + return 0; + } + const std::vector payload = build_remove_delta_payload(id); + const uint32_t base_crc = index_state_crc32c(*idx); + const uint32_t post_remove_crc = index_state_crc32c_after_remove(*idx, id); + const int append_status = append_delta_record( + *idx, + delta_path, + kTvidOpRemove, + 1, + base_crc, + post_remove_crc, + payload); + if (append_status != GGML_VEC_INDEX_OK) { + return append_status; + } + return ggml_vec_index_remove_unlocked(idx, id); + } catch (const std::bad_alloc &) { + return GGML_VEC_INDEX_E_OOM; + } catch (...) { + return GGML_VEC_INDEX_E_INTERNAL; + } +} + int ggml_vec_index_contains(const ggml_vec_index_t * idx, uint64_t id) { if (idx == nullptr) { return 0; } - return idx->id_to_slot.count(id) != 0 ? 1 : 0; + try { + std::shared_lock lock(idx->mutex); + return idx->id_to_slot.count(id) != 0 ? 1 : 0; + } catch (...) { + return 0; + } } void ggml_vec_index_prepare(ggml_vec_index_t * /*idx*/) { @@ -775,6 +1515,19 @@ inline float dot_q8_scalar(const float * query, const int8_t * codes, float scal return acc; } +inline float dot_q4_scalar(const float * query, const uint8_t * codes, float scale, int dim) { + float acc = 0.0f; + for (int i = 0; i < dim; ++i) { + const uint8_t byte = codes[static_cast(i) / 2]; + const uint8_t nibble = (i & 1) == 0 ? + static_cast(byte & 0x0f) : + static_cast(byte >> 4); + const float value = static_cast(q4_decode(nibble)) * scale; + acc += query[i] * value; + } + return acc; +} + #if GGML_VEC_INDEX_USE_NEON inline float horizontal_sum(float32x4_t v) { @@ -809,6 +1562,54 @@ inline float dot_q8_neon(const float * query, const int8_t * codes, float scale, return acc; } +inline void dot_q4_neon_accum8( + const float * query, + uint8x8_t codes, + float scale, + float32x4_t & acc0, + float32x4_t & acc1) { + const int16x8_t q16 = vsubq_s16( + vreinterpretq_s16_u16(vmovl_u8(codes)), + vdupq_n_s16(8)); + const float32x4_t q0 = + vmulq_n_f32(vcvtq_f32_s32(vmovl_s16(vget_low_s16(q16))), scale); + const float32x4_t q1 = + vmulq_n_f32(vcvtq_f32_s32(vmovl_s16(vget_high_s16(q16))), scale); + acc0 = vmlaq_f32(acc0, vld1q_f32(query), q0); + acc1 = vmlaq_f32(acc1, vld1q_f32(query + 4), q1); +} + +inline float dot_q4_neon(const float * query, const uint8_t * codes, float scale, int dim) { + float32x4_t acc0 = vdupq_n_f32(0.0f); + float32x4_t acc1 = vdupq_n_f32(0.0f); + float32x4_t acc2 = vdupq_n_f32(0.0f); + float32x4_t acc3 = vdupq_n_f32(0.0f); + + int i = 0; + for (; i + 16 <= dim; i += 16) { + const uint8x8_t packed = vld1_u8(codes + static_cast(i) / 2); + const uint8x8_t low = vand_u8(packed, vdup_n_u8(0x0f)); + const uint8x8_t high = vshr_n_u8(packed, 4); + const uint8x8x2_t zipped = vzip_u8(low, high); + + dot_q4_neon_accum8(query + i, zipped.val[0], scale, acc0, acc1); + dot_q4_neon_accum8(query + i + 8, zipped.val[1], scale, acc2, acc3); + } + + float acc = + horizontal_sum(acc0) + horizontal_sum(acc1) + + horizontal_sum(acc2) + horizontal_sum(acc3); + for (; i < dim; ++i) { + const uint8_t byte = codes[static_cast(i) / 2]; + const uint8_t nibble = (i & 1) == 0 ? + static_cast(byte & 0x0f) : + static_cast(byte >> 4); + const float value = static_cast(q4_decode(nibble)) * scale; + acc += query[i] * value; + } + return acc; +} + #endif #if GGML_VEC_INDEX_USE_AVX2 @@ -858,83 +1659,401 @@ bool cpu_has_avx2() { #endif } -#endif +#endif + +inline float dot_q8(const float * query, const int8_t * codes, float scale, int dim) { +#if GGML_VEC_INDEX_USE_NEON + return dot_q8_neon(query, codes, scale, dim); +#elif GGML_VEC_INDEX_USE_AVX2 + static const bool has_avx2 = cpu_has_avx2(); + if (has_avx2) { + return dot_q8_avx2(query, codes, scale, dim); + } + return dot_q8_scalar(query, codes, scale, dim); +#else + return dot_q8_scalar(query, codes, scale, dim); +#endif +} + +inline float dot_q4(const float * query, const uint8_t * codes, float scale, int dim) { +#if GGML_VEC_INDEX_USE_NEON + return dot_q4_neon(query, codes, scale, dim); +#else + return dot_q4_scalar(query, codes, scale, dim); +#endif +} + +inline float score_slot(const ggml_vec_index_t & idx, const float * query, size_t slot) { + const int dim = idx.dim; + return is_q4(idx) ? + dot_q4( + query, + idx.q4_data.data() + slot * q4_row_bytes(static_cast(dim)), + idx.q4_scale[slot], + dim) : + is_q8(idx) ? + dot_q8( + query, + idx.q8_data.data() + slot * static_cast(dim), + idx.q8_scale[slot], + dim) : + dot( + query, + idx.data.data() + slot * static_cast(dim), + dim); +} + +void decode_slot_to_f32(const ggml_vec_index_t & idx, size_t slot, float * dst) { + const int dim = idx.dim; + if (is_q4(idx)) { + const uint8_t * codes = + idx.q4_data.data() + slot * q4_row_bytes(static_cast(dim)); + const float scale = idx.q4_scale[slot]; + for (int i = 0; i < dim; ++i) { + const uint8_t byte = codes[static_cast(i) / 2]; + const uint8_t nibble = (i & 1) == 0 ? + static_cast(byte & 0x0f) : + static_cast(byte >> 4); + dst[i] = static_cast(q4_decode(nibble)) * scale; + } + } else if (is_q8(idx)) { + const int8_t * codes = idx.q8_data.data() + slot * static_cast(dim); + const float scale = idx.q8_scale[slot]; + for (int i = 0; i < dim; ++i) { + dst[i] = static_cast(codes[i]) * scale; + } + } else { + std::memcpy( + dst, + idx.data.data() + slot * static_cast(dim), + static_cast(dim) * sizeof(float)); + } +} + +size_t best_centroid(const float * query, const std::vector & centroids, int n_lists, int dim) { + size_t best = 0; + float best_score = -FLT_MAX; + for (int list = 0; list < n_lists; ++list) { + const float s = dot(query, centroids.data() + static_cast(list) * dim, dim); + if (s > best_score) { + best_score = s; + best = static_cast(list); + } + } + return best; +} + +// Run a single query against all slots, write top-k into out_scores/out_ids. +// If the index holds fewer than k entries, pad with sentinels. +void search_one( + const ggml_vec_index_t & idx, + const float * query, + int k, + float * out_scores, + uint64_t * out_ids, + const std::vector * allowed_slots = nullptr) { + + const size_t n_slots = idx.slot_to_id.size(); + + test_maybe_throw_bad_alloc(); + std::priority_queue, MinHeapCmp> heap; + + auto visit_slot = [&](size_t slot) { + const float s = score_slot(idx, query, slot); + if (heap.size() < static_cast(k)) { + heap.push({ s, idx.slot_to_id[slot] }); + } else if (s > heap.top().score) { + heap.pop(); + heap.push({ s, idx.slot_to_id[slot] }); + } + }; + + if (allowed_slots != nullptr) { + for (size_t slot : *allowed_slots) { + if (slot < n_slots) { + visit_slot(slot); + } + } + } else { + for (size_t slot = 0; slot < n_slots; ++slot) { + visit_slot(slot); + } + } + + // Drain the heap into a temporary descending list. + std::vector drained; + drained.reserve(heap.size()); + while (!heap.empty()) { + drained.push_back(heap.top()); + heap.pop(); + } + std::reverse(drained.begin(), drained.end()); // now descending by score + + for (int i = 0; i < k; ++i) { + if (static_cast(i) < drained.size()) { + out_scores[i] = drained[i].score; + out_ids[i] = drained[i].id; + } else { + out_scores[i] = -FLT_MAX; + out_ids[i] = UINT64_MAX; + } + } +} + +std::vector allowed_slots_for_ids( + const ggml_vec_index_t & idx, + const uint64_t * allowed_ids, + int n_allowed) { + std::vector slots; + slots.reserve(static_cast(n_allowed)); + for (int i = 0; i < n_allowed; ++i) { + const auto it = idx.id_to_slot.find(allowed_ids[i]); + if (it != idx.id_to_slot.end()) { + slots.push_back(it->second); + } + } + std::sort(slots.begin(), slots.end()); + slots.erase(std::unique(slots.begin(), slots.end()), slots.end()); + return slots; +} + +} // namespace + +static int ggml_vec_index_build_ivf_unlocked(ggml_vec_index_t * idx, int n_lists, int n_iter) { + try { + if (idx == nullptr || n_lists <= 0 || n_iter < 0) { + return GGML_VEC_INDEX_E_INVALID_ARG; + } + + const size_t n_slots = idx->slot_to_id.size(); + const int dim = idx->dim; + if (n_slots == 0) { + invalidate_ivf(*idx); + idx->ivf_generation = idx->generation; + return GGML_VEC_INDEX_OK; + } + + const int actual_lists = static_cast( + std::min(static_cast(n_lists), n_slots)); + const size_t dim_sz = static_cast(dim); + test_maybe_throw_bad_alloc(); + + std::vector centroids(static_cast(actual_lists) * dim_sz); + std::vector next_centroids(centroids.size()); + std::vector counts(static_cast(actual_lists)); + std::vector row(dim_sz); + std::vector> lists(static_cast(actual_lists)); + + for (int list = 0; list < actual_lists; ++list) { + const size_t slot = static_cast(list) * n_slots / + static_cast(actual_lists); + float * centroid = centroids.data() + static_cast(list) * dim_sz; + decode_slot_to_f32(*idx, slot, centroid); + } + + for (int iter = 0; iter < n_iter; ++iter) { + std::fill(next_centroids.begin(), next_centroids.end(), 0.0f); + std::fill(counts.begin(), counts.end(), 0); + + for (size_t slot = 0; slot < n_slots; ++slot) { + decode_slot_to_f32(*idx, slot, row.data()); + const size_t list = best_centroid(row.data(), centroids, actual_lists, dim); + float * dst = next_centroids.data() + list * dim_sz; + for (int i = 0; i < dim; ++i) { + dst[i] += row[static_cast(i)]; + } + ++counts[list]; + } + + for (int list = 0; list < actual_lists; ++list) { + float * centroid = centroids.data() + static_cast(list) * dim_sz; + if (counts[static_cast(list)] == 0) { + continue; + } + const float inv_count = 1.0f / + static_cast(counts[static_cast(list)]); + const float * src = + next_centroids.data() + static_cast(list) * dim_sz; + for (int i = 0; i < dim; ++i) { + centroid[i] = src[static_cast(i)] * inv_count; + } + } + } + + for (size_t slot = 0; slot < n_slots; ++slot) { + decode_slot_to_f32(*idx, slot, row.data()); + const size_t list = best_centroid(row.data(), centroids, actual_lists, dim); + lists[list].push_back(slot); + } + + idx->ivf_centroids = std::move(centroids); + idx->ivf_lists = std::move(lists); + idx->ivf_n_lists = actual_lists; + idx->ivf_generation = idx->generation; + return GGML_VEC_INDEX_OK; + } catch (const std::bad_alloc &) { + return GGML_VEC_INDEX_E_OOM; + } catch (...) { + return GGML_VEC_INDEX_E_INTERNAL; + } +} + +int ggml_vec_index_build_ivf(ggml_vec_index_t * idx, int n_lists, int n_iter) { + if (idx == nullptr) { + return GGML_VEC_INDEX_E_INVALID_ARG; + } + try { + std::unique_lock lock(idx->mutex); + return ggml_vec_index_build_ivf_unlocked(idx, n_lists, n_iter); + } catch (...) { + return GGML_VEC_INDEX_E_INTERNAL; + } +} + +static int ggml_vec_index_search_impl( + const ggml_vec_index_t * idx, + const float * queries, + int n_q, + int k, + bool filtered, + const uint64_t * allowed_ids, + int n_allowed, + const ggml_vec_index_filter_t * prepared_filter, + float * out_scores, + uint64_t * out_ids) { + + if (idx == nullptr || queries == nullptr || + out_scores == nullptr || out_ids == nullptr) { + return GGML_VEC_INDEX_E_INVALID_ARG; + } + if (n_q < 0 || k <= 0 || + (filtered && prepared_filter == nullptr && + (n_allowed < 0 || (n_allowed > 0 && allowed_ids == nullptr)))) { + return GGML_VEC_INDEX_E_INVALID_ARG; + } + if (n_q == 0) { + return GGML_VEC_INDEX_OK; + } + + try { + std::shared_lock lock(idx->mutex); + const int dim = idx->dim; + const size_t n_q_sz = static_cast(n_q); + const size_t k_sz = static_cast(k); + const size_t dim_sz = static_cast(dim); + if ((dim_sz != 0 && n_q_sz > std::numeric_limits::max() / dim_sz) || + n_q_sz > std::numeric_limits::max() / k_sz) { + return GGML_VEC_INDEX_E_INVALID_ARG; + } + if (!all_finite(queries, n_q_sz * dim_sz)) { + return GGML_VEC_INDEX_E_INVALID_ARG; + } + + std::vector allowed_slots; + const std::vector * allowed_ptr = nullptr; + if (prepared_filter != nullptr) { + if (prepared_filter->dim != idx->dim || + prepared_filter->bit_width != idx->bit_width || + prepared_filter->generation != idx->generation) { + return GGML_VEC_INDEX_E_INVALID_ARG; + } + allowed_ptr = &prepared_filter->slots; + } else if (filtered) { + allowed_slots = allowed_slots_for_ids(*idx, allowed_ids, n_allowed); + allowed_ptr = &allowed_slots; + } -inline float dot_q8(const float * query, const int8_t * codes, float scale, int dim) { -#if GGML_VEC_INDEX_USE_NEON - return dot_q8_neon(query, codes, scale, dim); -#elif GGML_VEC_INDEX_USE_AVX2 - static const bool has_avx2 = cpu_has_avx2(); - if (has_avx2) { - return dot_q8_avx2(query, codes, scale, dim); + for (int q = 0; q < n_q; ++q) { + search_one( + *idx, + queries + static_cast(q) * static_cast(dim), + k, + out_scores + static_cast(q) * static_cast(k), + out_ids + static_cast(q) * static_cast(k), + allowed_ptr); + } + } catch (const std::bad_alloc &) { + return GGML_VEC_INDEX_E_OOM; + } catch (...) { + return GGML_VEC_INDEX_E_INTERNAL; } - return dot_q8_scalar(query, codes, scale, dim); -#else - return dot_q8_scalar(query, codes, scale, dim); -#endif + return GGML_VEC_INDEX_OK; } -// Run a single query against all slots, write top-k into out_scores/out_ids. -// If the index holds fewer than k entries, pad with sentinels. -void search_one( - const ggml_vec_index_t & idx, - const float * query, +int ggml_vec_index_search( + const ggml_vec_index_t * idx, + const float * queries, + int n_q, int k, float * out_scores, uint64_t * out_ids) { + return ggml_vec_index_search_impl( + idx, queries, n_q, k, false, nullptr, 0, nullptr, out_scores, out_ids); +} - const int dim = idx.dim; - const size_t n_slots = idx.slot_to_id.size(); - - test_maybe_throw_bad_alloc(); - std::priority_queue, MinHeapCmp> heap; +int ggml_vec_index_search_filtered( + const ggml_vec_index_t * idx, + const float * queries, + int n_q, + int k, + const uint64_t * allowed_ids, + int n_allowed, + float * out_scores, + uint64_t * out_ids) { + return ggml_vec_index_search_impl( + idx, queries, n_q, k, true, allowed_ids, n_allowed, nullptr, out_scores, out_ids); +} - for (size_t slot = 0; slot < n_slots; ++slot) { - const float s = is_q8(idx) ? - dot_q8( - query, - idx.q8_data.data() + slot * static_cast(dim), - idx.q8_scale[slot], - dim) : - dot( - query, - idx.data.data() + slot * static_cast(dim), - dim); - if (heap.size() < static_cast(k)) { - heap.push({ s, idx.slot_to_id[slot] }); - } else if (s > heap.top().score) { - heap.pop(); - heap.push({ s, idx.slot_to_id[slot] }); +ggml_vec_index_filter_t * ggml_vec_index_filter_create( + const ggml_vec_index_t * idx, + const uint64_t * allowed_ids, + int n_allowed) { + try { + if (idx == nullptr || n_allowed < 0 || + (n_allowed > 0 && allowed_ids == nullptr)) { + return nullptr; + } + std::shared_lock lock(idx->mutex); + auto * filter = new (std::nothrow) ggml_vec_index_filter(); + if (filter == nullptr) { + return nullptr; } + std::unique_ptr owned(filter); + owned->dim = idx->dim; + owned->bit_width = idx->bit_width; + owned->generation = idx->generation; + owned->slots = allowed_slots_for_ids(*idx, allowed_ids, n_allowed); + return owned.release(); + } catch (...) { + return nullptr; } +} - // Drain the heap into a temporary descending list. - std::vector drained; - drained.reserve(heap.size()); - while (!heap.empty()) { - drained.push_back(heap.top()); - heap.pop(); - } - std::reverse(drained.begin(), drained.end()); // now descending by score +void ggml_vec_index_filter_free(ggml_vec_index_filter_t * filter) { + delete filter; +} - for (int i = 0; i < k; ++i) { - if (static_cast(i) < drained.size()) { - out_scores[i] = drained[i].score; - out_ids[i] = drained[i].id; - } else { - out_scores[i] = -FLT_MAX; - out_ids[i] = UINT64_MAX; - } +int ggml_vec_index_search_prepared_filtered( + const ggml_vec_index_t * idx, + const ggml_vec_index_filter_t * filter, + const float * queries, + int n_q, + int k, + float * out_scores, + uint64_t * out_ids) { + if (filter == nullptr) { + return GGML_VEC_INDEX_E_INVALID_ARG; } + return ggml_vec_index_search_impl( + idx, queries, n_q, k, true, nullptr, 0, filter, out_scores, out_ids); } -} // namespace - -int ggml_vec_index_search( +int ggml_vec_index_search_ivf( const ggml_vec_index_t * idx, const float * queries, int n_q, int k, + int nprobe, float * out_scores, uint64_t * out_ids) { @@ -942,7 +2061,7 @@ int ggml_vec_index_search( out_scores == nullptr || out_ids == nullptr) { return GGML_VEC_INDEX_E_INVALID_ARG; } - if (n_q < 0 || k <= 0) { + if (n_q < 0 || k <= 0 || nprobe <= 0) { return GGML_VEC_INDEX_E_INVALID_ARG; } if (n_q == 0) { @@ -950,9 +2069,10 @@ int ggml_vec_index_search( } try { + std::shared_lock lock(idx->mutex); const int dim = idx->dim; const size_t n_q_sz = static_cast(n_q); - const size_t k_sz = static_cast(k); + const size_t k_sz = static_cast(k); const size_t dim_sz = static_cast(dim); if ((dim_sz != 0 && n_q_sz > std::numeric_limits::max() / dim_sz) || n_q_sz > std::numeric_limits::max() / k_sz) { @@ -961,28 +2081,68 @@ int ggml_vec_index_search( if (!all_finite(queries, n_q_sz * dim_sz)) { return GGML_VEC_INDEX_E_INVALID_ARG; } + if (idx->ivf_generation != idx->generation || + idx->ivf_n_lists < 0 || + static_cast(idx->ivf_n_lists) != idx->ivf_lists.size() || + idx->ivf_centroids.size() != static_cast(idx->ivf_n_lists) * dim_sz) { + return GGML_VEC_INDEX_E_INVALID_ARG; + } + const int probe_count = std::min(nprobe, idx->ivf_n_lists); for (int q = 0; q < n_q; ++q) { - search_one( - *idx, - queries + static_cast(q) * static_cast(dim), - k, - out_scores + static_cast(q) * static_cast(k), - out_ids + static_cast(q) * static_cast(k)); + const float * query = queries + static_cast(q) * dim_sz; + float * scores = out_scores + static_cast(q) * k_sz; + uint64_t * ids = out_ids + static_cast(q) * k_sz; + + if (idx->ivf_n_lists == 0) { + const std::vector empty_slots; + search_one(*idx, query, k, scores, ids, &empty_slots); + continue; + } + + std::vector centroid_scores; + centroid_scores.reserve(static_cast(idx->ivf_n_lists)); + for (int list = 0; list < idx->ivf_n_lists; ++list) { + const float score = dot( + query, + idx->ivf_centroids.data() + static_cast(list) * dim_sz, + dim); + centroid_scores.push_back({ score, static_cast(list) }); + } + std::sort( + centroid_scores.begin(), + centroid_scores.end(), + [](const ScoreId & a, const ScoreId & b) { + return a.score > b.score; + }); + + size_t candidate_count = 0; + for (int probe = 0; probe < probe_count; ++probe) { + candidate_count += + idx->ivf_lists[static_cast(centroid_scores[probe].id)].size(); + } + std::vector candidate_slots; + candidate_slots.reserve(candidate_count); + for (int probe = 0; probe < probe_count; ++probe) { + const auto & list = + idx->ivf_lists[static_cast(centroid_scores[probe].id)]; + candidate_slots.insert(candidate_slots.end(), list.begin(), list.end()); + } + search_one(*idx, query, k, scores, ids, &candidate_slots); } + return GGML_VEC_INDEX_OK; } catch (const std::bad_alloc &) { return GGML_VEC_INDEX_E_OOM; } catch (...) { return GGML_VEC_INDEX_E_INTERNAL; } - return GGML_VEC_INDEX_OK; } // --------------------------------------------------------------------------- // Persistence // --------------------------------------------------------------------------- -int ggml_vec_index_write(ggml_vec_index_t * idx, const char * path) { +static int ggml_vec_index_write_unlocked(ggml_vec_index_t * idx, const char * path) { try { if (idx == nullptr || path == nullptr) { return GGML_VEC_INDEX_E_INVALID_ARG; @@ -995,7 +2155,11 @@ int ggml_vec_index_write(ggml_vec_index_t * idx, const char * path) { if (dim_sz != 0 && n > std::numeric_limits::max() / dim_sz) { return GGML_VEC_INDEX_E_INTERNAL; } - if (is_q8(*idx)) { + if (is_q4(*idx)) { + if (idx->q4_data.size() != n * q4_row_bytes(dim_sz) || idx->q4_scale.size() != n) { + return GGML_VEC_INDEX_E_INTERNAL; + } + } else if (is_q8(*idx)) { if (idx->q8_data.size() != n * dim_sz || idx->q8_scale.size() != n) { return GGML_VEC_INDEX_E_INTERNAL; } @@ -1030,9 +2194,9 @@ int ggml_vec_index_write(ggml_vec_index_t * idx, const char * path) { const uint32_t n_le = static_cast(idx->slot_to_id.size()); put_u32_le(header + 8, dim_le); put_u32_le(header + 12, n_le); - put_u32_le(header + 16, is_q8(*idx) ? kQParamScaleF32 : kQParamNone); - put_u32_le(header + 20, is_q8(*idx) ? 4u : 0u); - put_u32_le(header + 24, is_q8(*idx) ? 1u : 4u); + put_u32_le(header + 16, is_quantized(*idx) ? kQParamScaleF32 : kQParamNone); + put_u32_le(header + 20, is_quantized(*idx) ? 4u : 0u); + put_u32_le(header + 24, is_q4(*idx) ? 0u : (is_q8(*idx) ? 1u : 4u)); put_u32_le(header + 28, 0); if (!write_bytes(f, header, sizeof(header))) { @@ -1043,14 +2207,24 @@ int ggml_vec_index_write(ggml_vec_index_t * idx, const char * path) { uint32_t qparams_crc = 0xffffffffu; uint32_t vectors_crc = 0xffffffffu; uint32_t ids_crc = 0xffffffffu; - if (is_q8(*idx)) { - for (float scale : idx->q8_scale) { + if (is_quantized(*idx)) { + const std::vector & scales = is_q4(*idx) ? idx->q4_scale : idx->q8_scale; + for (float scale : scales) { if (!write_u32_le_crc(f, float_to_u32(scale), qparams_crc)) { return fail_io(); } } - if (!idx->q8_data.empty()) { + if (is_q4(*idx)) { + if (!idx->q4_data.empty() && + !write_bytes(f, idx->q4_data.data(), idx->q4_data.size())) { + return fail_io(); + } + if (!idx->q4_data.empty()) { + vectors_crc = + crc32c_update(vectors_crc, idx->q4_data.data(), idx->q4_data.size()); + } + } else if (!idx->q8_data.empty()) { if (!write_bytes(f, idx->q8_data.data(), idx->q8_data.size())) { return fail_io(); } @@ -1101,6 +2275,18 @@ int ggml_vec_index_write(ggml_vec_index_t * idx, const char * path) { } } +int ggml_vec_index_write(ggml_vec_index_t * idx, const char * path) { + if (idx == nullptr) { + return GGML_VEC_INDEX_E_INVALID_ARG; + } + try { + std::unique_lock lock(idx->mutex); + return ggml_vec_index_write_unlocked(idx, path); + } catch (...) { + return GGML_VEC_INDEX_E_INTERNAL; + } +} + ggml_vec_index_t * ggml_vec_index_load(const char * path) { try { if (path == nullptr) { @@ -1173,7 +2359,9 @@ ggml_vec_index_t * ggml_vec_index_load(const char * path) { return nullptr; } if (version == kTvimVersion && - ((bit_width == 8 && (kind != kStorageQ8 || qparam_type != kQParamScaleF32 || + ((bit_width == 4 && (kind != kStorageQ4 || qparam_type != kQParamScaleF32 || + qparam_bytes != 4 || comp_bytes != 0)) || + (bit_width == 8 && (kind != kStorageQ8 || qparam_type != kQParamScaleF32 || qparam_bytes != 4 || comp_bytes != 1)) || (bit_width == 32 && (kind != kStorageF32 || qparam_type != kQParamNone || qparam_bytes != 0 || comp_bytes != 4)))) { @@ -1226,7 +2414,10 @@ ggml_vec_index_t * ggml_vec_index_load(const char * path) { } test_maybe_throw_bad_alloc(); - if (is_q8(*idx)) { + if (is_q4(*idx)) { + idx->q4_data.resize(n * q4_row_bytes(dim_sz)); + idx->q4_scale.resize(n); + } else if (is_q8(*idx)) { idx->q8_data.resize(n * dim_sz); idx->q8_scale.resize(n); } else { @@ -1242,7 +2433,7 @@ ggml_vec_index_t * ggml_vec_index_load(const char * path) { uint32_t vectors_crc = 0xffffffffu; uint32_t ids_crc = 0xffffffffu; - if (version == kTvimVersionV1 && is_q8(*idx)) { + if (version == kTvimVersionV1 && is_quantized(*idx)) { std::vector row(dim_sz); for (size_t slot = 0; slot < n; ++slot) { for (float & v : row) { @@ -1255,14 +2446,23 @@ ggml_vec_index_t * ggml_vec_index_load(const char * path) { return nullptr; } } - quantize_q8_row( - row.data(), - idx->q8_data.data() + slot * dim_sz, - dim, - idx->q8_scale[slot]); + if (is_q4(*idx)) { + quantize_q4_row( + row.data(), + idx->q4_data.data() + slot * q4_row_bytes(dim_sz), + dim, + idx->q4_scale[slot]); + } else { + quantize_q8_row( + row.data(), + idx->q8_data.data() + slot * dim_sz, + dim, + idx->q8_scale[slot]); + } } - } else if (is_q8(*idx)) { - for (float & scale : idx->q8_scale) { + } else if (is_quantized(*idx)) { + std::vector & scales = is_q4(*idx) ? idx->q4_scale : idx->q8_scale; + for (float & scale : scales) { uint32_t bits = 0; const bool read_ok = checksummed ? read_u32_le_crc(f, bits, qparams_crc) : @@ -1276,30 +2476,67 @@ ggml_vec_index_t * ggml_vec_index_load(const char * path) { } } - if (!idx->q8_data.empty()) { - if (idx->q8_data.size() > - static_cast(std::numeric_limits::max())) { - return nullptr; - } - f.read( - reinterpret_cast(idx->q8_data.data()), - static_cast(idx->q8_data.size() * sizeof(int8_t))); - if (!f) { - return nullptr; + std::vector * q4_data = is_q4(*idx) ? &idx->q4_data : nullptr; + if (is_q4(*idx)) { + if (!q4_data->empty()) { + if (q4_data->size() > + static_cast(std::numeric_limits::max())) { + return nullptr; + } + f.read( + reinterpret_cast(q4_data->data()), + static_cast(q4_data->size() * sizeof(uint8_t))); + if (!f) { + return nullptr; + } + if (checksummed) { + vectors_crc = crc32c_update(vectors_crc, q4_data->data(), q4_data->size()); + } } - if (checksummed) { - vectors_crc = crc32c_update( - vectors_crc, idx->q8_data.data(), idx->q8_data.size()); + const size_t row_bytes = q4_row_bytes(dim_sz); + for (size_t slot = 0; slot < n; ++slot) { + const float scale = idx->q4_scale[slot]; + const uint8_t * row = idx->q4_data.data() + slot * row_bytes; + for (size_t i = 0; i < dim_sz; ++i) { + const uint8_t byte = row[i / 2]; + const uint8_t nibble = (i & 1) == 0 ? + static_cast(byte & 0x0f) : + static_cast(byte >> 4); + if (nibble == 0 || + !std::isfinite(static_cast(q4_decode(nibble)) * scale)) { + return nullptr; + } + } + if ((dim_sz & 1) != 0 && (row[row_bytes - 1] >> 4) != 8) { + return nullptr; + } } - } - for (size_t slot = 0; slot < n; ++slot) { - const float scale = idx->q8_scale[slot]; - const int8_t * row = idx->q8_data.data() + slot * dim_sz; - for (size_t i = 0; i < dim_sz; ++i) { - if (row[i] == std::numeric_limits::min() || - !std::isfinite(static_cast(row[i]) * scale)) { + } else { + if (!idx->q8_data.empty()) { + if (idx->q8_data.size() > + static_cast(std::numeric_limits::max())) { return nullptr; } + f.read( + reinterpret_cast(idx->q8_data.data()), + static_cast(idx->q8_data.size() * sizeof(int8_t))); + if (!f) { + return nullptr; + } + if (checksummed) { + vectors_crc = crc32c_update( + vectors_crc, idx->q8_data.data(), idx->q8_data.size()); + } + } + for (size_t slot = 0; slot < n; ++slot) { + const float scale = idx->q8_scale[slot]; + const int8_t * row = idx->q8_data.data() + slot * dim_sz; + for (size_t i = 0; i < dim_sz; ++i) { + if (row[i] == std::numeric_limits::min() || + !std::isfinite(static_cast(row[i]) * scale)) { + return nullptr; + } + } } } } else { @@ -1360,18 +2597,259 @@ ggml_vec_index_t * ggml_vec_index_load(const char * path) { } } +namespace { + +bool expected_add_delta_payload_size(uint64_t n, uint64_t dim, uint64_t & size) { + uint64_t id_bytes = 0; + uint64_t values = 0; + uint64_t vector_bytes = 0; + if (!checked_mul_u64(n, sizeof(uint64_t), id_bytes) || + !checked_mul_u64(n, dim, values) || + !checked_mul_u64(values, sizeof(uint32_t), vector_bytes) || + !checked_add_u64(id_bytes, vector_bytes, size)) { + return false; + } + return true; +} + +bool read_delta_payload(std::ifstream & f, uint64_t payload_bytes, std::vector & payload) { + if (payload_bytes > static_cast(std::numeric_limits::max()) || + payload_bytes > static_cast(std::numeric_limits::max())) { + return false; + } + payload.resize(static_cast(payload_bytes)); + if (!payload.empty()) { + f.read( + reinterpret_cast(payload.data()), + static_cast(payload.size())); + if (!f) { + return false; + } + } + return true; +} + +bool replay_add_delta( + ggml_vec_index_t * idx, + uint32_t n, + const std::vector & payload) { + if (n == 0 || n > static_cast(std::numeric_limits::max())) { + return false; + } + uint64_t expected_payload = 0; + if (!expected_add_delta_payload_size(n, static_cast(idx->dim), expected_payload) || + expected_payload != payload.size()) { + return false; + } + + const size_t n_sz = static_cast(n); + const size_t dim_sz = static_cast(idx->dim); + std::vector ids(n_sz); + std::vector vectors(n_sz * dim_sz); + + const uint8_t * ptr = payload.data(); + for (size_t i = 0; i < n_sz; ++i) { + ids[i] = get_u64_le(ptr); + ptr += sizeof(uint64_t); + } + for (float & v : vectors) { + v = u32_to_float(get_u32_le(ptr)); + if (!std::isfinite(v)) { + return false; + } + ptr += sizeof(uint32_t); + } + + const int status = ggml_vec_index_add( + idx, + vectors.data(), + static_cast(n), + ids.data()); + return status == GGML_VEC_INDEX_OK; +} + +bool replay_remove_delta(ggml_vec_index_t * idx, uint32_t n, const std::vector & payload) { + if (n != 1 || payload.size() != sizeof(uint64_t)) { + return false; + } + const uint64_t id = get_u64_le(payload.data()); + return ggml_vec_index_remove(idx, id) >= 0; +} + +bool replay_delta_log(ggml_vec_index_t * idx, const char * delta_path) { + std::filesystem::path fs_path; + if (!filesystem_path_from_utf8(delta_path, fs_path)) { + return false; + } + if (!std::filesystem::exists(fs_path) || std::filesystem::file_size(fs_path) == 0) { + return true; + } + const uint64_t file_size = static_cast(std::filesystem::file_size(fs_path)); + if (file_size < kTvidHeaderSize) { + return false; + } + + std::ifstream f(fs_path, std::ios::binary); + if (!f.is_open()) { + return false; + } + uint8_t header[kTvidHeaderSize] = {}; + f.read(reinterpret_cast(header), sizeof(header)); + if (!f || + std::memcmp(header, kTvidMagic, 4) != 0 || + header[4] != kTvidVersion || + header[5] != static_cast(idx->bit_width) || + header[6] != 0 || + header[7] != 0 || + get_u32_le(header + 8) != static_cast(idx->dim)) { + return false; + } + + const uint32_t base_crc = get_u32_le(header + 12); + const uint32_t snapshot_crc = index_state_crc32c(*idx); + const bool apply_records = snapshot_crc == base_crc; + uint32_t last_state_crc = base_crc; + + uint64_t offset = kTvidHeaderSize; + while (offset < file_size) { + uint8_t record[kTvidRecordHeaderSize] = {}; + f.read(reinterpret_cast(record), sizeof(record)); + if (f.gcount() == 0 && f.eof()) { + break; + } + if (f.gcount() != static_cast(sizeof(record))) { + break; // torn trailing record header + } + offset += kTvidRecordHeaderSize; + + const uint8_t op = record[0]; + const uint32_t n = get_u32_le(record + 4); + const uint64_t payload_bytes = get_u64_le(record + 8); + const uint32_t expected_crc = get_u32_le(record + 16); + const uint32_t state_crc = get_u32_le(record + 20); + if (record[1] != 0 || record[2] != 0 || record[3] != 0 || + (op != kTvidOpAdd && op != kTvidOpRemove)) { + return false; + } + if (payload_bytes > file_size - offset) { + break; // torn trailing record payload + } + + std::vector payload; + if (!read_delta_payload(f, payload_bytes, payload)) { + break; + } + offset += payload_bytes; + + uint32_t crc = crc32c_update(0xffffffffu, record, 16); + crc = crc32c_update(crc, record + 20, 4); + if (!payload.empty()) { + crc = crc32c_update(crc, payload.data(), payload.size()); + } + if ((crc ^ 0xffffffffu) != expected_crc) { + return false; + } + + if (apply_records) { + if (op == kTvidOpAdd) { + if (!replay_add_delta(idx, n, payload)) { + return false; + } + } else { + if (!replay_remove_delta(idx, n, payload)) { + return false; + } + } + } + last_state_crc = state_crc; + } + + return apply_records ? + index_state_crc32c(*idx) == last_state_crc : + snapshot_crc == last_state_crc; +} + +} // namespace + +ggml_vec_index_t * ggml_vec_index_load_with_delta( + const char * snapshot_path, + const char * delta_path) { + try { + if (snapshot_path == nullptr || delta_path == nullptr) { + return nullptr; + } + std::unique_ptr idx( + ggml_vec_index_load(snapshot_path), + ggml_vec_index_free); + if (idx == nullptr) { + return nullptr; + } + if (!replay_delta_log(idx.get(), delta_path)) { + return nullptr; + } + return idx.release(); + } catch (...) { + return nullptr; + } +} + +int ggml_vec_index_compact_delta( + ggml_vec_index_t * idx, + const char * snapshot_path, + const char * delta_path) { + try { + if (idx == nullptr || snapshot_path == nullptr || delta_path == nullptr) { + return GGML_VEC_INDEX_E_INVALID_ARG; + } + std::unique_lock lock(idx->mutex); + const int write_status = ggml_vec_index_write_unlocked(idx, snapshot_path); + if (write_status != GGML_VEC_INDEX_OK) { + return write_status; + } + return write_empty_delta_log(*idx, delta_path); + } catch (const std::bad_alloc &) { + return GGML_VEC_INDEX_E_OOM; + } catch (...) { + return GGML_VEC_INDEX_E_INTERNAL; + } +} + // --------------------------------------------------------------------------- // Stats // --------------------------------------------------------------------------- int ggml_vec_index_len(const ggml_vec_index_t * idx) { - return idx ? static_cast(idx->slot_to_id.size()) : 0; + if (idx == nullptr) { + return 0; + } + try { + std::shared_lock lock(idx->mutex); + return static_cast(idx->slot_to_id.size()); + } catch (...) { + return 0; + } } int ggml_vec_index_dim(const ggml_vec_index_t * idx) { - return idx ? idx->dim : 0; + if (idx == nullptr) { + return 0; + } + try { + std::shared_lock lock(idx->mutex); + return idx->dim; + } catch (...) { + return 0; + } } int ggml_vec_index_bit_width(const ggml_vec_index_t * idx) { - return idx ? idx->bit_width : 0; + if (idx == nullptr) { + return 0; + } + try { + std::shared_lock lock(idx->mutex); + return idx->bit_width; + } catch (...) { + return 0; + } } diff --git a/tests/bench-vector-index.cpp b/tests/bench-vector-index.cpp index d36a1f7e0b1b..f3b873090c45 100644 --- a/tests/bench-vector-index.cpp +++ b/tests/bench-vector-index.cpp @@ -9,6 +9,7 @@ #include #include #include +#include #include #include #include @@ -31,6 +32,10 @@ struct BenchConfig { int k = 10; int warmups = 2; int repeats = 7; + int delta_ops = 256; + int ivf_lists = 64; + int ivf_iters = 4; + int ivf_nprobe = 4; }; struct TimedSearch { @@ -39,6 +44,24 @@ struct TimedSearch { std::vector ids; }; +template +double median_time_ms(int warmups, int repeats, Fn fn) { + for (int i = 0; i < warmups; ++i) { + fn(); + } + + std::vector times; + times.reserve(static_cast(repeats)); + for (int i = 0; i < repeats; ++i) { + const auto t0 = std::chrono::steady_clock::now(); + fn(); + const auto t1 = std::chrono::steady_clock::now(); + times.push_back(std::chrono::duration(t1 - t0).count()); + } + std::sort(times.begin(), times.end()); + return times[times.size() / 2]; +} + const char * q8_kernel_name() { #if defined(__ARM_NEON) || defined(__ARM_NEON__) return "arm-neon"; @@ -52,6 +75,14 @@ const char * q8_kernel_name() { #endif } +const char * q4_kernel_name() { +#if defined(__ARM_NEON) || defined(__ARM_NEON__) + return "arm-neon"; +#else + return "scalar"; +#endif +} + std::vector make_normalized_vectors(int n, int dim, uint32_t seed) { std::mt19937 rng(seed); std::normal_distribution dist(0.0f, 1.0f); @@ -83,7 +114,7 @@ TimedSearch run_search( result.scores.resize(static_cast(n_query) * static_cast(k)); result.ids.resize(static_cast(n_query) * static_cast(k)); - for (int i = 0; i < warmups; ++i) { + result.ms = median_time_ms(warmups, repeats, [&]() { CHECK(ggml_vec_index_search( idx, queries.data(), @@ -91,24 +122,83 @@ TimedSearch run_search( k, result.scores.data(), result.ids.data()) == GGML_VEC_INDEX_OK); - } + }); + return result; +} - std::vector times; - times.reserve(static_cast(repeats)); - for (int i = 0; i < repeats; ++i) { - const auto t0 = std::chrono::steady_clock::now(); - CHECK(ggml_vec_index_search( +TimedSearch run_ivf_search( + const ggml_vec_index_t * idx, + const std::vector & queries, + int n_query, + int k, + int nprobe, + int warmups, + int repeats) { + TimedSearch result; + result.scores.resize(static_cast(n_query) * static_cast(k)); + result.ids.resize(static_cast(n_query) * static_cast(k)); + + result.ms = median_time_ms(warmups, repeats, [&]() { + CHECK(ggml_vec_index_search_ivf( idx, queries.data(), n_query, k, + nprobe, result.scores.data(), result.ids.data()) == GGML_VEC_INDEX_OK); - const auto t1 = std::chrono::steady_clock::now(); - times.push_back(std::chrono::duration(t1 - t0).count()); - } - std::sort(times.begin(), times.end()); - result.ms = times[times.size() / 2]; + }); + return result; +} + +TimedSearch run_filtered_search( + const ggml_vec_index_t * idx, + const std::vector & queries, + int n_query, + int k, + const std::vector & allowed_ids, + int warmups, + int repeats) { + TimedSearch result; + result.scores.resize(static_cast(n_query) * static_cast(k)); + result.ids.resize(static_cast(n_query) * static_cast(k)); + + result.ms = median_time_ms(warmups, repeats, [&]() { + CHECK(ggml_vec_index_search_filtered( + idx, + queries.data(), + n_query, + k, + allowed_ids.data(), + static_cast(allowed_ids.size()), + result.scores.data(), + result.ids.data()) == GGML_VEC_INDEX_OK); + }); + return result; +} + +TimedSearch run_prepared_filtered_search( + const ggml_vec_index_t * idx, + const ggml_vec_index_filter_t * filter, + const std::vector & queries, + int n_query, + int k, + int warmups, + int repeats) { + TimedSearch result; + result.scores.resize(static_cast(n_query) * static_cast(k)); + result.ids.resize(static_cast(n_query) * static_cast(k)); + + result.ms = median_time_ms(warmups, repeats, [&]() { + CHECK(ggml_vec_index_search_prepared_filtered( + idx, + filter, + queries.data(), + n_query, + k, + result.scores.data(), + result.ids.data()) == GGML_VEC_INDEX_OK); + }); return result; } @@ -132,6 +222,144 @@ std::filesystem::path write_index_file(ggml_vec_index_t * idx, const char * name return path; } +std::vector read_file_bytes(const std::filesystem::path & path) { + std::ifstream f(path, std::ios::binary); + CHECK(f.is_open()); + const auto size = std::filesystem::file_size(path); + std::vector bytes(static_cast(size)); + if (!bytes.empty()) { + f.read(reinterpret_cast(bytes.data()), static_cast(bytes.size())); + CHECK(f.gcount() == static_cast(bytes.size())); + } + return bytes; +} + +void write_file_bytes(const std::filesystem::path & path, const std::vector & bytes) { + std::ofstream f(path, std::ios::binary | std::ios::trunc); + CHECK(f.is_open()); + if (!bytes.empty()) { + f.write(reinterpret_cast(bytes.data()), static_cast(bytes.size())); + } + CHECK(static_cast(f)); +} + +std::vector make_allowlist(int n_vec, int n_allowed) { + std::vector allowed; + allowed.reserve(static_cast(n_allowed)); + for (int i = 0; i < n_allowed; ++i) { + const int row = (i * 37) % n_vec; + allowed.push_back(static_cast(row) + 1); + } + std::sort(allowed.begin(), allowed.end()); + allowed.erase(std::unique(allowed.begin(), allowed.end()), allowed.end()); + return allowed; +} + +double recall_against(const TimedSearch & exact, const TimedSearch & candidate, int n_query, int k) { + int overlap = 0; + for (int q = 0; q < n_query; ++q) { + std::unordered_set exact_topk; + for (int j = 0; j < k; ++j) { + exact_topk.insert(exact.ids[static_cast(q) * k + j]); + } + for (int j = 0; j < k; ++j) { + const uint64_t id = candidate.ids[static_cast(q) * k + j]; + if (exact_topk.count(id) != 0) { + ++overlap; + } + } + } + return static_cast(overlap) / static_cast(n_query * k); +} + +struct DeltaBenchResult { + double snapshot_load_ms = 0.0; + double replay_load_ms = 0.0; + double compact_ms = 0.0; + double post_compact_load_ms = 0.0; + uintmax_t snapshot_bytes_before = 0; + uintmax_t delta_bytes_before = 0; + uintmax_t snapshot_bytes_after = 0; + uintmax_t delta_bytes_after = 0; +}; + +DeltaBenchResult run_delta_bench( + int bit_width, + const std::vector & vectors, + const std::vector & ids, + const BenchConfig & cfg, + const char * snapshot_name, + const char * delta_name) { + CHECK(cfg.delta_ops > 0 && cfg.delta_ops * 2 < cfg.n_vec); + const int base_n = cfg.n_vec - cfg.delta_ops; + const std::filesystem::path snapshot_path = + std::filesystem::temp_directory_path() / snapshot_name; + const std::filesystem::path delta_path = + std::filesystem::temp_directory_path() / delta_name; + std::filesystem::remove(snapshot_path); + std::filesystem::remove(delta_path); + + ggml_vec_index_t * idx = ggml_vec_index_create(cfg.dim, bit_width); + CHECK(idx != nullptr); + CHECK(ggml_vec_index_add(idx, vectors.data(), base_n, ids.data()) == GGML_VEC_INDEX_OK); + CHECK(ggml_vec_index_write(idx, snapshot_path.string().c_str()) == GGML_VEC_INDEX_OK); + + for (int i = 0; i < cfg.delta_ops; ++i) { + const int row = base_n + i; + CHECK(ggml_vec_index_add_logged( + idx, + vectors.data() + static_cast(row) * static_cast(cfg.dim), + 1, + ids.data() + row, + delta_path.string().c_str()) == GGML_VEC_INDEX_OK); + } + for (int i = 0; i < cfg.delta_ops / 2; ++i) { + CHECK(ggml_vec_index_remove_logged( + idx, + ids[static_cast(i)], + delta_path.string().c_str()) == 1); + } + + DeltaBenchResult result; + result.snapshot_bytes_before = std::filesystem::file_size(snapshot_path); + result.delta_bytes_before = std::filesystem::file_size(delta_path); + const std::vector dirty_delta = read_file_bytes(delta_path); + + result.snapshot_load_ms = median_time_ms(cfg.warmups, cfg.repeats, [&]() { + ggml_vec_index_t * loaded = ggml_vec_index_load(snapshot_path.string().c_str()); + CHECK(loaded != nullptr); + ggml_vec_index_free(loaded); + }); + result.replay_load_ms = median_time_ms(cfg.warmups, cfg.repeats, [&]() { + ggml_vec_index_t * loaded = ggml_vec_index_load_with_delta( + snapshot_path.string().c_str(), + delta_path.string().c_str()); + CHECK(loaded != nullptr); + ggml_vec_index_free(loaded); + }); + result.compact_ms = median_time_ms(0, cfg.repeats, [&]() { + write_file_bytes(delta_path, dirty_delta); + CHECK(ggml_vec_index_compact_delta( + idx, + snapshot_path.string().c_str(), + delta_path.string().c_str()) == GGML_VEC_INDEX_OK); + }); + result.snapshot_bytes_after = std::filesystem::file_size(snapshot_path); + result.delta_bytes_after = std::filesystem::file_size(delta_path); + result.post_compact_load_ms = median_time_ms(cfg.warmups, cfg.repeats, [&]() { + ggml_vec_index_t * loaded = ggml_vec_index_load_with_delta( + snapshot_path.string().c_str(), + delta_path.string().c_str()); + CHECK(loaded != nullptr); + ggml_vec_index_free(loaded); + }); + + ggml_vec_index_free(idx); + std::filesystem::remove(snapshot_path); + std::filesystem::remove(delta_path); + return result; +} + } // namespace int main() { @@ -146,20 +374,100 @@ int main() { ggml_vec_index_t * f32 = ggml_vec_index_create(cfg.dim, 32); ggml_vec_index_t * q8 = ggml_vec_index_create(cfg.dim, 8); + ggml_vec_index_t * q4 = ggml_vec_index_create(cfg.dim, 4); CHECK(f32 != nullptr); CHECK(q8 != nullptr); + CHECK(q4 != nullptr); CHECK(ggml_vec_index_add(f32, vectors.data(), cfg.n_vec, ids.data()) == GGML_VEC_INDEX_OK); CHECK(ggml_vec_index_add(q8, vectors.data(), cfg.n_vec, ids.data()) == GGML_VEC_INDEX_OK); + CHECK(ggml_vec_index_add(q4, vectors.data(), cfg.n_vec, ids.data()) == GGML_VEC_INDEX_OK); + + const double f32_ivf_build_ms = median_time_ms(cfg.warmups, cfg.repeats, [&]() { + CHECK(ggml_vec_index_build_ivf(f32, cfg.ivf_lists, cfg.ivf_iters) == GGML_VEC_INDEX_OK); + }); + const double q8_ivf_build_ms = median_time_ms(cfg.warmups, cfg.repeats, [&]() { + CHECK(ggml_vec_index_build_ivf(q8, cfg.ivf_lists, cfg.ivf_iters) == GGML_VEC_INDEX_OK); + }); + const double q4_ivf_build_ms = median_time_ms(cfg.warmups, cfg.repeats, [&]() { + CHECK(ggml_vec_index_build_ivf(q4, cfg.ivf_lists, cfg.ivf_iters) == GGML_VEC_INDEX_OK); + }); const TimedSearch f32_res = run_search( f32, queries, cfg.n_query, cfg.k, cfg.warmups, cfg.repeats); const TimedSearch q8_res = run_search( q8, queries, cfg.n_query, cfg.k, cfg.warmups, cfg.repeats); + const TimedSearch q4_res = run_search( + q4, queries, cfg.n_query, cfg.k, cfg.warmups, cfg.repeats); + const TimedSearch f32_ivf_res = run_ivf_search( + f32, queries, cfg.n_query, cfg.k, cfg.ivf_nprobe, cfg.warmups, cfg.repeats); + const TimedSearch q8_ivf_res = run_ivf_search( + q8, queries, cfg.n_query, cfg.k, cfg.ivf_nprobe, cfg.warmups, cfg.repeats); + const TimedSearch q4_ivf_res = run_ivf_search( + q4, queries, cfg.n_query, cfg.k, cfg.ivf_nprobe, cfg.warmups, cfg.repeats); + + const std::vector filter_sizes = { + 32, + 128, + 512, + cfg.n_vec / 2, + }; + std::vector> allowlists; + std::vector f32_filtered; + std::vector q8_filtered; + std::vector f32_prepared_filtered; + std::vector q8_prepared_filtered; + for (int requested : filter_sizes) { + const int n_allowed = std::min(requested, cfg.n_vec); + allowlists.push_back(make_allowlist(cfg.n_vec, n_allowed)); + f32_filtered.push_back(run_filtered_search( + f32, + queries, + cfg.n_query, + cfg.k, + allowlists.back(), + cfg.warmups, + cfg.repeats)); + q8_filtered.push_back(run_filtered_search( + q8, + queries, + cfg.n_query, + cfg.k, + allowlists.back(), + cfg.warmups, + cfg.repeats)); + ggml_vec_index_filter_t * f32_filter = ggml_vec_index_filter_create( + f32, allowlists.back().data(), static_cast(allowlists.back().size())); + ggml_vec_index_filter_t * q8_filter = ggml_vec_index_filter_create( + q8, allowlists.back().data(), static_cast(allowlists.back().size())); + CHECK(f32_filter != nullptr); + CHECK(q8_filter != nullptr); + f32_prepared_filtered.push_back(run_prepared_filtered_search( + f32, + f32_filter, + queries, + cfg.n_query, + cfg.k, + cfg.warmups, + cfg.repeats)); + q8_prepared_filtered.push_back(run_prepared_filtered_search( + q8, + q8_filter, + queries, + cfg.n_query, + cfg.k, + cfg.warmups, + cfg.repeats)); + ggml_vec_index_filter_free(f32_filter); + ggml_vec_index_filter_free(q8_filter); + } - double mean_abs_drift = 0.0; - double max_abs_drift = 0.0; + double q8_mean_abs_drift = 0.0; + double q8_max_abs_drift = 0.0; + double q4_mean_abs_drift = 0.0; + double q4_max_abs_drift = 0.0; + int q8_overlap = 0; + int q4_overlap = 0; int drift_count = 0; - int overlap = 0; for (int q = 0; q < cfg.n_query; ++q) { std::unordered_set f32_topk; @@ -172,25 +480,68 @@ int main() { const size_t pos = static_cast(q) * static_cast(cfg.k) + static_cast(j); const uint64_t id = q8_res.ids[pos]; if (f32_topk.count(id) != 0) { - ++overlap; + ++q8_overlap; } const float exact = dot_exact(vectors, query, id, cfg.dim); const double drift = std::fabs(static_cast(exact) - q8_res.scores[pos]); - mean_abs_drift += drift; - max_abs_drift = std::max(max_abs_drift, drift); + q8_mean_abs_drift += drift; + q8_max_abs_drift = std::max(q8_max_abs_drift, drift); + + const uint64_t q4_id = q4_res.ids[pos]; + if (f32_topk.count(q4_id) != 0) { + ++q4_overlap; + } + const float q4_exact = dot_exact(vectors, query, q4_id, cfg.dim); + const double q4_drift = std::fabs(static_cast(q4_exact) - q4_res.scores[pos]); + q4_mean_abs_drift += q4_drift; + q4_max_abs_drift = std::max(q4_max_abs_drift, q4_drift); ++drift_count; } } - mean_abs_drift /= static_cast(drift_count); - const double recall_at_k = static_cast(overlap) / - static_cast(cfg.n_query * cfg.k); + q8_mean_abs_drift /= static_cast(drift_count); + q4_mean_abs_drift /= static_cast(drift_count); + const double q8_recall_at_k = static_cast(q8_overlap) / + static_cast(cfg.n_query * cfg.k); + const double q4_recall_at_k = static_cast(q4_overlap) / + static_cast(cfg.n_query * cfg.k); + const double f32_ivf_recall_at_k = + recall_against(f32_res, f32_ivf_res, cfg.n_query, cfg.k); + const double q8_ivf_recall_at_k = + recall_against(f32_res, q8_ivf_res, cfg.n_query, cfg.k); + const double q4_ivf_recall_at_k = + recall_against(f32_res, q4_ivf_res, cfg.n_query, cfg.k); const auto f32_path = write_index_file(f32, "ggml-vector-index-bench-f32.tvim"); const auto q8_path = write_index_file(q8, "ggml-vector-index-bench-q8.tvim"); + const auto q4_path = write_index_file(q4, "ggml-vector-index-bench-q4.tvim"); const uintmax_t f32_file_size = std::filesystem::file_size(f32_path); const uintmax_t q8_file_size = std::filesystem::file_size(q8_path); + const uintmax_t q4_file_size = std::filesystem::file_size(q4_path); std::filesystem::remove(f32_path); std::filesystem::remove(q8_path); + std::filesystem::remove(q4_path); + + const DeltaBenchResult f32_delta = run_delta_bench( + 32, + vectors, + ids, + cfg, + "ggml-vector-index-bench-delta-f32.tvim", + "ggml-vector-index-bench-delta-f32.tvid"); + const DeltaBenchResult q8_delta = run_delta_bench( + 8, + vectors, + ids, + cfg, + "ggml-vector-index-bench-delta-q8.tvim", + "ggml-vector-index-bench-delta-q8.tvid"); + const DeltaBenchResult q4_delta = run_delta_bench( + 4, + vectors, + ids, + cfg, + "ggml-vector-index-bench-delta-q4.tvim", + "ggml-vector-index-bench-delta-q4.tvid"); const size_t f32_memory_bytes = static_cast(cfg.n_vec) * static_cast(cfg.dim) * sizeof(float) + @@ -199,24 +550,111 @@ int main() { static_cast(cfg.n_vec) * static_cast(cfg.dim) * sizeof(int8_t) + static_cast(cfg.n_vec) * sizeof(float) + static_cast(cfg.n_vec) * sizeof(uint64_t); + const size_t q4_memory_bytes = + static_cast(cfg.n_vec) * ((static_cast(cfg.dim) + 1) / 2) + + static_cast(cfg.n_vec) * sizeof(float) + + static_cast(cfg.n_vec) * sizeof(uint64_t); std::printf("bench-vector-index\n"); std::printf(" q8 kernel=%s\n", q8_kernel_name()); + std::printf(" q4 kernel=%s\n", q4_kernel_name()); std::printf(" n_vec=%d dim=%d n_query=%d k=%d warmups=%d repeats=%d\n", cfg.n_vec, cfg.dim, cfg.n_query, cfg.k, cfg.warmups, cfg.repeats); - std::printf(" estimated memory: f32=%zu bytes q8=%zu bytes ratio=%.3f\n", - f32_memory_bytes, q8_memory_bytes, - static_cast(q8_memory_bytes) / static_cast(f32_memory_bytes)); - std::printf(" file size: f32=%llu bytes q8=%llu bytes ratio=%.3f\n", + std::printf(" estimated memory: f32=%zu bytes q8=%zu bytes q4=%zu bytes q8/f32=%.3f q4/f32=%.3f\n", + f32_memory_bytes, q8_memory_bytes, q4_memory_bytes, + static_cast(q8_memory_bytes) / static_cast(f32_memory_bytes), + static_cast(q4_memory_bytes) / static_cast(f32_memory_bytes)); + std::printf(" file size: f32=%llu bytes q8=%llu bytes q4=%llu bytes q8/f32=%.3f q4/f32=%.3f\n", static_cast(f32_file_size), static_cast(q8_file_size), - static_cast(q8_file_size) / static_cast(f32_file_size)); - std::printf(" median latency: f32=%.3f ms q8=%.3f ms ratio=%.3f\n", - f32_res.ms, q8_res.ms, q8_res.ms / f32_res.ms); - std::printf(" quality: recall@%d=%.4f mean_abs_score_drift=%.6f max_abs_score_drift=%.6f\n", - cfg.k, recall_at_k, mean_abs_drift, max_abs_drift); + static_cast(q4_file_size), + static_cast(q8_file_size) / static_cast(f32_file_size), + static_cast(q4_file_size) / static_cast(f32_file_size)); + std::printf(" median latency: f32=%.3f ms q8=%.3f ms q4=%.3f ms q8/f32=%.3f q4/f32=%.3f\n", + f32_res.ms, q8_res.ms, q4_res.ms, q8_res.ms / f32_res.ms, q4_res.ms / f32_res.ms); + std::printf( + " ivf build: lists=%d iters=%d f32=%.3f ms q8=%.3f ms q4=%.3f ms\n", + cfg.ivf_lists, + cfg.ivf_iters, + f32_ivf_build_ms, + q8_ivf_build_ms, + q4_ivf_build_ms); + std::printf( + " ivf latency: nprobe=%d f32=%.3f ms q8=%.3f ms q4=%.3f ms f32/full=%.3f q8/full=%.3f q4/full=%.3f\n", + cfg.ivf_nprobe, + f32_ivf_res.ms, + q8_ivf_res.ms, + q4_ivf_res.ms, + f32_ivf_res.ms / f32_res.ms, + q8_ivf_res.ms / q8_res.ms, + q4_ivf_res.ms / q4_res.ms); + std::printf( + " ivf recall: f32@%d=%.4f q8@%d=%.4f q4@%d=%.4f against exact f32\n", + cfg.k, + f32_ivf_recall_at_k, + cfg.k, + q8_ivf_recall_at_k, + cfg.k, + q4_ivf_recall_at_k); + for (size_t i = 0; i < allowlists.size(); ++i) { + std::printf( + " filtered latency: allowed=%zu f32=%.3f ms q8=%.3f ms f32/prepared=%.3f ms q8/prepared=%.3f ms\n", + allowlists[i].size(), + f32_filtered[i].ms, + q8_filtered[i].ms, + f32_prepared_filtered[i].ms, + q8_prepared_filtered[i].ms); + std::printf( + " filtered ratio: allowed=%zu f32/full=%.3f q8/full=%.3f f32/prep_speedup=%.3f q8/prep_speedup=%.3f\n", + allowlists[i].size(), + f32_filtered[i].ms / f32_res.ms, + q8_filtered[i].ms / q8_res.ms, + f32_filtered[i].ms / f32_prepared_filtered[i].ms, + q8_filtered[i].ms / q8_prepared_filtered[i].ms); + } + std::printf( + " delta load f32: snapshot=%.3f ms replay=%.3f ms compact=%.3f ms post_compact=%.3f ms\n", + f32_delta.snapshot_load_ms, + f32_delta.replay_load_ms, + f32_delta.compact_ms, + f32_delta.post_compact_load_ms); + std::printf( + " delta bytes f32: snapshot_before=%llu delta_before=%llu snapshot_after=%llu delta_after=%llu\n", + static_cast(f32_delta.snapshot_bytes_before), + static_cast(f32_delta.delta_bytes_before), + static_cast(f32_delta.snapshot_bytes_after), + static_cast(f32_delta.delta_bytes_after)); + std::printf( + " delta load q8: snapshot=%.3f ms replay=%.3f ms compact=%.3f ms post_compact=%.3f ms\n", + q8_delta.snapshot_load_ms, + q8_delta.replay_load_ms, + q8_delta.compact_ms, + q8_delta.post_compact_load_ms); + std::printf( + " delta bytes q8: snapshot_before=%llu delta_before=%llu snapshot_after=%llu delta_after=%llu\n", + static_cast(q8_delta.snapshot_bytes_before), + static_cast(q8_delta.delta_bytes_before), + static_cast(q8_delta.snapshot_bytes_after), + static_cast(q8_delta.delta_bytes_after)); + std::printf( + " delta load q4: snapshot=%.3f ms replay=%.3f ms compact=%.3f ms post_compact=%.3f ms\n", + q4_delta.snapshot_load_ms, + q4_delta.replay_load_ms, + q4_delta.compact_ms, + q4_delta.post_compact_load_ms); + std::printf( + " delta bytes q4: snapshot_before=%llu delta_before=%llu snapshot_after=%llu delta_after=%llu\n", + static_cast(q4_delta.snapshot_bytes_before), + static_cast(q4_delta.delta_bytes_before), + static_cast(q4_delta.snapshot_bytes_after), + static_cast(q4_delta.delta_bytes_after)); + std::printf(" quality q8: recall@%d=%.4f mean_abs_score_drift=%.6f max_abs_score_drift=%.6f\n", + cfg.k, q8_recall_at_k, q8_mean_abs_drift, q8_max_abs_drift); + std::printf(" quality q4: recall@%d=%.4f mean_abs_score_drift=%.6f max_abs_score_drift=%.6f\n", + cfg.k, q4_recall_at_k, q4_mean_abs_drift, q4_max_abs_drift); ggml_vec_index_free(f32); ggml_vec_index_free(q8); + ggml_vec_index_free(q4); return 0; } diff --git a/tests/test-vector-index-faults.cpp b/tests/test-vector-index-faults.cpp index 99944e30d1f0..7bd5594fee4a 100644 --- a/tests/test-vector-index-faults.cpp +++ b/tests/test-vector-index-faults.cpp @@ -115,6 +115,26 @@ int main() { reset_fault_hooks(); } + { + std::array scores{}; + std::array ids{}; + CHECK(ggml_vec_index_build_ivf(idx, /*n_lists=*/2, /*n_iter=*/1) + == GGML_VEC_INDEX_OK); + CHECK(ggml_vec_index_search_ivf( + idx, base_vectors.data(), 1, /*k=*/1, /*nprobe=*/2, + scores.data(), ids.data()) == GGML_VEC_INDEX_OK); + + ggml_vec_index_test_set_oom_countdown(0); + CHECK(ggml_vec_index_build_ivf(idx, /*n_lists=*/2, /*n_iter=*/1) + == GGML_VEC_INDEX_E_OOM); + reset_fault_hooks(); + + CHECK(ggml_vec_index_search_ivf( + idx, base_vectors.data(), 1, /*k=*/1, /*nprobe=*/2, + scores.data(), ids.data()) == GGML_VEC_INDEX_OK); + CHECK(ids[0] == base_ids[0]); + } + const std::string path = (std::filesystem::temp_directory_path() / "ggml-vector-index-fault-test.tvim").string(); @@ -155,8 +175,41 @@ int main() { CHECK(loaded != nullptr); CHECK(ggml_vec_index_len(loaded) == 2); ggml_vec_index_free(loaded); + + const std::string delta_path = + (std::filesystem::temp_directory_path() / + "ggml-vector-index-fault-test.tvid").string(); + std::filesystem::remove(delta_path); + + const std::array logged_vector = { 0.0f, 0.0f, 1.0f, 0.0f }; + const uint64_t logged_id = 401; + ggml_vec_index_test_set_write_fail_after(8); + CHECK(ggml_vec_index_add_logged( + idx, logged_vector.data(), 1, &logged_id, delta_path.c_str()) == + GGML_VEC_INDEX_E_IO); + reset_fault_hooks(); + CHECK(ggml_vec_index_contains(idx, logged_id) == 0); + CHECK(ggml_vec_index_len(idx) == 3); + if (std::filesystem::exists(delta_path)) { + CHECK(std::filesystem::file_size(delta_path) == 0); + } + + CHECK(ggml_vec_index_add_logged( + idx, logged_vector.data(), 1, &logged_id, delta_path.c_str()) == + GGML_VEC_INDEX_OK); + CHECK(ggml_vec_index_contains(idx, logged_id) == 1); + const std::vector old_delta = read_file_bytes(delta_path); + + ggml_vec_index_test_set_write_fail_after(8); + CHECK(ggml_vec_index_remove_logged(idx, logged_id, delta_path.c_str()) == + GGML_VEC_INDEX_E_IO); + reset_fault_hooks(); + CHECK(ggml_vec_index_contains(idx, logged_id) == 1); + CHECK(read_file_bytes(delta_path) == old_delta); + ggml_vec_index_free(idx); std::filesystem::remove(path); + std::filesystem::remove(delta_path); std::printf("test-vector-index-faults: OK\n"); return 0; diff --git a/tests/test-vector-index.cpp b/tests/test-vector-index.cpp index 799dabbbf784..329b6db10380 100644 --- a/tests/test-vector-index.cpp +++ b/tests/test-vector-index.cpp @@ -8,6 +8,7 @@ #include #include #include +#include #include #include #include @@ -18,6 +19,7 @@ #include #include #include +#include #include #ifndef _WIN32 @@ -63,6 +65,25 @@ float q8_dot_reference(const std::vector & vector, const std::vector & vector, const std::vector & query) { + CHECK(vector.size() == query.size()); + + float max_abs = 0.0f; + for (float value : vector) { + max_abs = std::max(max_abs, std::fabs(value)); + } + const float scale = max_abs == 0.0f ? 1.0f : max_abs / 7.0f; + + float acc = 0.0f; + for (size_t i = 0; i < vector.size(); ++i) { + int code = max_abs == 0.0f ? + 0 : static_cast(std::nearbyint(vector[i] / scale)); + code = std::max(-7, std::min(7, code)); + acc += query[i] * (static_cast(code) * scale); + } + return acc; +} + uint8_t read_file_byte(const std::string & path, std::streamoff offset) { std::ifstream f(path, std::ios::binary); CHECK(f.is_open()); @@ -94,6 +115,15 @@ void write_file_bytes(const std::string & path, const std::vector & byt CHECK(static_cast(f)); } +void append_file_bytes(const std::string & path, const std::vector & bytes) { + std::ofstream f(path, std::ios::binary | std::ios::app); + CHECK(f.is_open()); + if (!bytes.empty()) { + f.write(reinterpret_cast(bytes.data()), static_cast(bytes.size())); + } + CHECK(static_cast(f)); +} + void append_u32_le(std::vector & bytes, uint32_t value) { for (int i = 0; i < 4; ++i) { bytes.push_back(static_cast(value >> (8 * i))); @@ -227,6 +257,140 @@ int main() { CHECK(ggml_vec_index_contains(idx, new_id) == 0); } + // IVF-flat ANN search is explicit and stale builds are rejected. + { + auto * ann = ggml_vec_index_create(kDim, /*bit_width=*/32); + CHECK(ann != nullptr); + CHECK(ggml_vec_index_add( + ann, vecs.data(), static_cast(ids.size()), ids.data()) == GGML_VEC_INDEX_OK); + + const std::vector query = normalize({0.9f, 0.3f, 0.1f, -0.2f}); + std::array exact_scores{}; + std::array ann_scores{}; + std::array exact_ids{}; + std::array ann_ids{}; + + CHECK(ggml_vec_index_search_ivf( + ann, query.data(), 1, /*k=*/1, /*nprobe=*/1, + ann_scores.data(), ann_ids.data()) == GGML_VEC_INDEX_E_INVALID_ARG); + CHECK(ggml_vec_index_build_ivf(ann, /*n_lists=*/0, /*n_iter=*/1) + == GGML_VEC_INDEX_E_INVALID_ARG); + CHECK(ggml_vec_index_build_ivf(ann, /*n_lists=*/16, /*n_iter=*/3) + == GGML_VEC_INDEX_OK); + + CHECK(ggml_vec_index_search( + ann, query.data(), 1, /*k=*/4, + exact_scores.data(), exact_ids.data()) == GGML_VEC_INDEX_OK); + CHECK(ggml_vec_index_search_ivf( + ann, query.data(), 1, /*k=*/4, /*nprobe=*/16, + ann_scores.data(), ann_ids.data()) == GGML_VEC_INDEX_OK); + for (int i = 0; i < 4; ++i) { + CHECK(ann_ids[i] == exact_ids[i]); + CHECK(std::fabs(ann_scores[i] - exact_scores[i]) < 1e-5f); + } + + CHECK(ggml_vec_index_search_ivf( + ann, seeds[0].data(), 1, /*k=*/1, /*nprobe=*/1, + ann_scores.data(), ann_ids.data()) == GGML_VEC_INDEX_OK); + CHECK(ann_ids[0] == ids[0]); + + const uint64_t ann_new_id = 9999991ULL; + CHECK(ggml_vec_index_add(ann, seeds[3].data(), 1, &ann_new_id) + == GGML_VEC_INDEX_OK); + CHECK(ggml_vec_index_search_ivf( + ann, query.data(), 1, /*k=*/1, /*nprobe=*/1, + ann_scores.data(), ann_ids.data()) == GGML_VEC_INDEX_E_INVALID_ARG); + + CHECK(ggml_vec_index_build_ivf(ann, /*n_lists=*/16, /*n_iter=*/3) + == GGML_VEC_INDEX_OK); + CHECK(ggml_vec_index_search_ivf( + ann, query.data(), 1, /*k=*/1, /*nprobe=*/16, + ann_scores.data(), ann_ids.data()) == GGML_VEC_INDEX_OK); + ggml_vec_index_free(ann); + } + + // Read-only APIs on one handle can run concurrently. + { + constexpr int n_rows = 16; + std::vector rows; + std::vector row_ids; + rows.reserve(static_cast(n_rows) * kDim); + row_ids.reserve(n_rows); + for (int row = 0; row < n_rows; ++row) { + const std::vector v = normalize({ + static_cast((row % 5) - 2), + static_cast(((row + 1) % 7) - 3), + static_cast(((row * 3) % 11) - 5), + 1.0f, + }); + rows.insert(rows.end(), v.begin(), v.end()); + row_ids.push_back(static_cast(7000 + row)); + } + + auto * concurrent = ggml_vec_index_create(kDim, /*bit_width=*/8); + CHECK(concurrent != nullptr); + CHECK(ggml_vec_index_add(concurrent, rows.data(), n_rows, row_ids.data()) + == GGML_VEC_INDEX_OK); + CHECK(ggml_vec_index_build_ivf(concurrent, /*n_lists=*/4, /*n_iter=*/2) + == GGML_VEC_INDEX_OK); + + const std::array allowed = { + row_ids[0], row_ids[2], row_ids[4], row_ids[6], row_ids[8], + }; + ggml_vec_index_filter_t * filter = ggml_vec_index_filter_create( + concurrent, allowed.data(), static_cast(allowed.size())); + CHECK(filter != nullptr); + + std::atomic ready{ 0 }; + std::atomic start{ false }; + std::vector threads; + for (int t = 0; t < 8; ++t) { + threads.emplace_back([&, t]() { + std::array query = { + rows[static_cast((t % n_rows) * kDim + 0)], + rows[static_cast((t % n_rows) * kDim + 1)], + rows[static_cast((t % n_rows) * kDim + 2)], + rows[static_cast((t % n_rows) * kDim + 3)], + }; + std::array scores{}; + std::array out_ids{}; + ready.fetch_add(1); + while (!start.load()) { + std::this_thread::yield(); + } + for (int iter = 0; iter < 200; ++iter) { + CHECK(ggml_vec_index_search( + concurrent, query.data(), 1, /*k=*/3, + scores.data(), out_ids.data()) == GGML_VEC_INDEX_OK); + CHECK(ggml_vec_index_search_filtered( + concurrent, query.data(), 1, /*k=*/3, + allowed.data(), static_cast(allowed.size()), + scores.data(), out_ids.data()) == GGML_VEC_INDEX_OK); + CHECK(ggml_vec_index_search_prepared_filtered( + concurrent, filter, query.data(), 1, /*k=*/3, + scores.data(), out_ids.data()) == GGML_VEC_INDEX_OK); + CHECK(ggml_vec_index_search_ivf( + concurrent, query.data(), 1, /*k=*/3, /*nprobe=*/4, + scores.data(), out_ids.data()) == GGML_VEC_INDEX_OK); + CHECK(ggml_vec_index_contains(concurrent, row_ids[static_cast(t % n_rows)]) == 1); + CHECK(ggml_vec_index_len(concurrent) == n_rows); + CHECK(ggml_vec_index_dim(concurrent) == kDim); + CHECK(ggml_vec_index_bit_width(concurrent) == 8); + } + }); + } + while (ready.load() != 8) { + std::this_thread::yield(); + } + start.store(true); + for (std::thread & thread : threads) { + thread.join(); + } + + ggml_vec_index_filter_free(filter); + ggml_vec_index_free(concurrent); + } + // Top-1 of querying with each unit vector should retrieve itself with // score very close to 1.0 (full f32, no quantization noise). { @@ -256,6 +420,83 @@ int main() { } } + // Filtered search only considers ids present in the allowlist. Missing + // and duplicate filter ids do not produce duplicate result rows. + { + const uint64_t missing_id = (1ULL << 60) + 99ULL; + const std::array allowed = { + ids[2], missing_id, ids[0], ids[0], + }; + std::array scores{}; + std::array out_ids{}; + CHECK(ggml_vec_index_search_filtered( + idx, seeds[0].data(), 1, /*k=*/3, + allowed.data(), static_cast(allowed.size()), + scores.data(), out_ids.data()) == GGML_VEC_INDEX_OK); + CHECK(out_ids[0] == ids[0]); + CHECK(out_ids[1] == ids[2]); + CHECK(out_ids[2] == UINT64_MAX); + CHECK(scores[2] == -FLT_MAX); + + CHECK(ggml_vec_index_search_filtered( + idx, seeds[0].data(), 1, /*k=*/2, + nullptr, /*n_allowed=*/0, + scores.data(), out_ids.data()) == GGML_VEC_INDEX_OK); + CHECK(out_ids[0] == UINT64_MAX); + CHECK(out_ids[1] == UINT64_MAX); + CHECK(scores[0] == -FLT_MAX); + CHECK(scores[1] == -FLT_MAX); + + CHECK(ggml_vec_index_search_filtered( + idx, seeds[0].data(), 1, /*k=*/1, + nullptr, /*n_allowed=*/1, + scores.data(), out_ids.data()) == GGML_VEC_INDEX_E_INVALID_ARG); + + auto * filter = ggml_vec_index_filter_create( + idx, allowed.data(), static_cast(allowed.size())); + CHECK(filter != nullptr); + scores = {}; + out_ids = {}; + CHECK(ggml_vec_index_search_prepared_filtered( + idx, filter, seeds[0].data(), 1, /*k=*/3, + scores.data(), out_ids.data()) == GGML_VEC_INDEX_OK); + CHECK(out_ids[0] == ids[0]); + CHECK(out_ids[1] == ids[2]); + CHECK(out_ids[2] == UINT64_MAX); + CHECK(scores[2] == -FLT_MAX); + ggml_vec_index_filter_free(filter); + + auto * empty_filter = ggml_vec_index_filter_create( + idx, nullptr, /*n_allowed=*/0); + CHECK(empty_filter != nullptr); + CHECK(ggml_vec_index_search_prepared_filtered( + idx, empty_filter, seeds[0].data(), 1, /*k=*/2, + scores.data(), out_ids.data()) == GGML_VEC_INDEX_OK); + CHECK(out_ids[0] == UINT64_MAX); + CHECK(out_ids[1] == UINT64_MAX); + CHECK(scores[0] == -FLT_MAX); + CHECK(scores[1] == -FLT_MAX); + ggml_vec_index_filter_free(empty_filter); + + CHECK(ggml_vec_index_filter_create( + idx, nullptr, /*n_allowed=*/1) == nullptr); + CHECK(ggml_vec_index_search_prepared_filtered( + idx, nullptr, seeds[0].data(), 1, /*k=*/1, + scores.data(), out_ids.data()) == GGML_VEC_INDEX_E_INVALID_ARG); + + auto * stale_filter = ggml_vec_index_filter_create( + idx, allowed.data(), static_cast(allowed.size())); + CHECK(stale_filter != nullptr); + const uint64_t stale_new_id = (1ULL << 60) + 100ULL; + CHECK(ggml_vec_index_add( + idx, seeds[3].data(), 1, &stale_new_id) == GGML_VEC_INDEX_OK); + CHECK(ggml_vec_index_search_prepared_filtered( + idx, stale_filter, seeds[0].data(), 1, /*k=*/1, + scores.data(), out_ids.data()) == GGML_VEC_INDEX_E_INVALID_ARG); + CHECK(ggml_vec_index_remove(idx, stale_new_id) == 1); + ggml_vec_index_filter_free(stale_filter); + } + // Remove + search: the removed id must no longer surface. { CHECK(ggml_vec_index_remove(idx, ids[1]) == 1); @@ -312,7 +553,152 @@ int main() { ggml_vec_index_free(loaded); std::filesystem::remove(path); - // v1 f32 snapshots migrate into current f32 or q8 storage. + // Incremental persistence: replay add/remove deltas on top of a snapshot. + { + const std::string snapshot_path = + (std::filesystem::temp_directory_path() / + "ggml-vector-index-delta-base.tvim").string(); + const std::string delta_path = + (std::filesystem::temp_directory_path() / + "ggml-vector-index-delta-log.tvid").string(); + const std::string missing_delta_path = + (std::filesystem::temp_directory_path() / + "ggml-vector-index-missing-delta-log.tvid").string(); + const std::string mismatched_snapshot_path = + (std::filesystem::temp_directory_path() / + "ggml-vector-index-delta-mismatch.tvim").string(); + const std::string corrupt_delta_path = + (std::filesystem::temp_directory_path() / + "ggml-vector-index-delta-corrupt.tvid").string(); + std::filesystem::remove(snapshot_path); + std::filesystem::remove(delta_path); + std::filesystem::remove(missing_delta_path); + std::filesystem::remove(mismatched_snapshot_path); + std::filesystem::remove(corrupt_delta_path); + + auto * base = ggml_vec_index_create(kDim, /*bit_width=*/32); + CHECK(base != nullptr); + std::vector base_vecs; + base_vecs.insert(base_vecs.end(), seeds[0].begin(), seeds[0].end()); + base_vecs.insert(base_vecs.end(), seeds[1].begin(), seeds[1].end()); + CHECK(ggml_vec_index_add(base, base_vecs.data(), 2, ids.data()) == GGML_VEC_INDEX_OK); + CHECK(ggml_vec_index_write(base, snapshot_path.c_str()) == GGML_VEC_INDEX_OK); + + auto * base_only = ggml_vec_index_load_with_delta( + snapshot_path.c_str(), missing_delta_path.c_str()); + CHECK(base_only != nullptr); + CHECK(ggml_vec_index_len(base_only) == 2); + ggml_vec_index_free(base_only); + + const uint64_t delta_id = (1ULL << 41) + 7ULL; + CHECK(ggml_vec_index_add_logged( + base, seeds[2].data(), 1, &delta_id, delta_path.c_str()) == GGML_VEC_INDEX_OK); + CHECK(ggml_vec_index_remove_logged( + base, ids[0], delta_path.c_str()) == 1); + CHECK(ggml_vec_index_add_logged( + base, seeds[2].data(), 1, &delta_id, delta_path.c_str()) == + GGML_VEC_INDEX_E_DUPLICATE); + + auto * replayed = ggml_vec_index_load_with_delta( + snapshot_path.c_str(), delta_path.c_str()); + CHECK(replayed != nullptr); + CHECK(ggml_vec_index_len(replayed) == 2); + CHECK(ggml_vec_index_contains(replayed, ids[0]) == 0); + CHECK(ggml_vec_index_contains(replayed, ids[1]) == 1); + CHECK(ggml_vec_index_contains(replayed, delta_id) == 1); + + std::array scores{}; + std::array out_ids{}; + CHECK(ggml_vec_index_search( + replayed, seeds[2].data(), 1, /*k=*/2, + scores.data(), out_ids.data()) == GGML_VEC_INDEX_OK); + CHECK(out_ids[0] == delta_id); + + std::vector corrupt_delta = read_file_bytes(delta_path); + corrupt_delta[16 + 20] ^= 1; // state CRC is covered by record CRC + write_file_bytes(corrupt_delta_path, corrupt_delta); + auto * corrupt_delta_loaded = ggml_vec_index_load_with_delta( + snapshot_path.c_str(), corrupt_delta_path.c_str()); + CHECK(corrupt_delta_loaded == nullptr); + ggml_vec_index_free(corrupt_delta_loaded); + + auto * mismatch = ggml_vec_index_create(kDim, /*bit_width=*/32); + CHECK(mismatch != nullptr); + std::vector mismatch_vecs; + mismatch_vecs.insert(mismatch_vecs.end(), seeds[0].begin(), seeds[0].end()); + mismatch_vecs.insert(mismatch_vecs.end(), seeds[3].begin(), seeds[3].end()); + CHECK(ggml_vec_index_add( + mismatch, mismatch_vecs.data(), 2, ids.data()) == GGML_VEC_INDEX_OK); + CHECK(ggml_vec_index_write( + mismatch, mismatched_snapshot_path.c_str()) == GGML_VEC_INDEX_OK); + auto * mismatch_loaded = ggml_vec_index_load_with_delta( + mismatched_snapshot_path.c_str(), delta_path.c_str()); + CHECK(mismatch_loaded == nullptr); + ggml_vec_index_free(mismatch_loaded); + const uint64_t mismatch_new_id = (1ULL << 41) + 9ULL; + CHECK(ggml_vec_index_add_logged( + mismatch, seeds[2].data(), 1, &mismatch_new_id, delta_path.c_str()) == + GGML_VEC_INDEX_E_IO); + CHECK(ggml_vec_index_contains(mismatch, mismatch_new_id) == 0); + ggml_vec_index_free(mismatch); + + const std::vector pre_compact_delta = read_file_bytes(delta_path); + CHECK(ggml_vec_index_compact_delta( + base, snapshot_path.c_str(), delta_path.c_str()) == GGML_VEC_INDEX_OK); + CHECK(std::filesystem::file_size(delta_path) == 16); + + auto * compacted = ggml_vec_index_load_with_delta( + snapshot_path.c_str(), delta_path.c_str()); + CHECK(compacted != nullptr); + CHECK(ggml_vec_index_len(compacted) == 2); + CHECK(ggml_vec_index_contains(compacted, ids[0]) == 0); + CHECK(ggml_vec_index_contains(compacted, ids[1]) == 1); + CHECK(ggml_vec_index_contains(compacted, delta_id) == 1); + ggml_vec_index_free(compacted); + + // Crash window: the compacted snapshot is durable but the old log + // survived. Replay must remain idempotent. + write_file_bytes(delta_path, pre_compact_delta); + auto * compacted_with_old_log = ggml_vec_index_load_with_delta( + snapshot_path.c_str(), delta_path.c_str()); + CHECK(compacted_with_old_log != nullptr); + CHECK(ggml_vec_index_len(compacted_with_old_log) == 2); + CHECK(ggml_vec_index_contains(compacted_with_old_log, ids[0]) == 0); + CHECK(ggml_vec_index_contains(compacted_with_old_log, ids[1]) == 1); + CHECK(ggml_vec_index_contains(compacted_with_old_log, delta_id) == 1); + ggml_vec_index_free(compacted_with_old_log); + + CHECK(ggml_vec_index_compact_delta( + base, snapshot_path.c_str(), delta_path.c_str()) == GGML_VEC_INDEX_OK); + const uint64_t post_compact_id = (1ULL << 41) + 8ULL; + CHECK(ggml_vec_index_add_logged( + base, seeds[3].data(), 1, &post_compact_id, delta_path.c_str()) == + GGML_VEC_INDEX_OK); + auto * replayed_after_compact = ggml_vec_index_load_with_delta( + snapshot_path.c_str(), delta_path.c_str()); + CHECK(replayed_after_compact != nullptr); + CHECK(ggml_vec_index_len(replayed_after_compact) == 3); + CHECK(ggml_vec_index_contains(replayed_after_compact, post_compact_id) == 1); + ggml_vec_index_free(replayed_after_compact); + + append_file_bytes(delta_path, { 0x01, 0x00, 0x00 }); + auto * replayed_with_torn_tail = ggml_vec_index_load_with_delta( + snapshot_path.c_str(), delta_path.c_str()); + CHECK(replayed_with_torn_tail != nullptr); + CHECK(ggml_vec_index_contains(replayed_with_torn_tail, delta_id) == 1); + CHECK(ggml_vec_index_contains(replayed_with_torn_tail, post_compact_id) == 1); + + ggml_vec_index_free(replayed_with_torn_tail); + ggml_vec_index_free(replayed); + ggml_vec_index_free(base); + std::filesystem::remove(snapshot_path); + std::filesystem::remove(delta_path); + std::filesystem::remove(mismatched_snapshot_path); + std::filesystem::remove(corrupt_delta_path); + } + + // v1 f32 snapshots migrate to q8 only for legacy bit_width=8; other + // legacy widths, including bit_width=4, migrate to f32. { const std::vector v1_ids = { (1ULL << 37) + 1ULL, @@ -417,6 +803,160 @@ int main() { ggml_vec_index_free(overflow_idx); } + // q4 production path: packed nibbles with one f32 scale per vector. + { + constexpr int tail_dim = 13; + const std::vector tail_vector = { + -1.0f, 0.75f, -0.5f, 0.25f, 0.125f, -0.875f, 0.625f, + -0.375f, 0.9f, -0.7f, 0.3f, -0.2f, 0.05f, + }; + const std::vector tail_query = { + 0.2f, -0.4f, 0.6f, -0.8f, 1.0f, 0.3f, -0.5f, + 0.7f, -0.9f, 0.11f, -0.22f, 0.33f, -0.44f, + }; + const uint64_t q4_id = (1ULL << 55) + 654ULL; + + auto * q4 = ggml_vec_index_create(tail_dim, /*bit_width=*/4); + CHECK(q4 != nullptr); + CHECK(ggml_vec_index_bit_width(q4) == 4); + CHECK(ggml_vec_index_add(q4, tail_vector.data(), 1, &q4_id) == GGML_VEC_INDEX_OK); + + std::array scores{}; + std::array out_ids{}; + CHECK(ggml_vec_index_search( + q4, tail_query.data(), 1, /*k=*/2, + scores.data(), out_ids.data()) == GGML_VEC_INDEX_OK); + CHECK(out_ids[0] == q4_id); + CHECK(out_ids[1] == UINT64_MAX); + CHECK(scores[1] == -FLT_MAX); + + const float expected = q4_dot_reference(tail_vector, tail_query); + const float tolerance = 1e-5f * std::max(1.0f, std::fabs(expected)); + CHECK(std::fabs(scores[0] - expected) <= tolerance); + + const std::string q4_path = + (std::filesystem::temp_directory_path() / + "ggml-vector-index-q4-test.tvim").string(); + const std::string corrupt_path = + (std::filesystem::temp_directory_path() / + "ggml-vector-index-q4-corrupt-test.tvim").string(); + CHECK(ggml_vec_index_write(q4, q4_path.c_str()) == GGML_VEC_INDEX_OK); + CHECK(read_file_byte(q4_path, 4) == 2); // .tvim v2 + CHECK(read_file_byte(q4_path, 5) == 4); // q4 bit width + CHECK(read_file_byte(q4_path, 6) == 3); // q4 storage kind + CHECK(read_file_byte(q4_path, 24) == 0); // packed components + + auto * q4_loaded = ggml_vec_index_load(q4_path.c_str()); + CHECK(q4_loaded != nullptr); + CHECK(ggml_vec_index_bit_width(q4_loaded) == 4); + CHECK(ggml_vec_index_len(q4_loaded) == 1); + scores = {}; + out_ids = {}; + CHECK(ggml_vec_index_search( + q4_loaded, tail_query.data(), 1, /*k=*/1, + scores.data(), out_ids.data()) == GGML_VEC_INDEX_OK); + CHECK(out_ids[0] == q4_id); + CHECK(std::fabs(scores[0] - expected) <= tolerance); + + constexpr size_t q4_vector_offset = 32 + sizeof(float); + constexpr size_t q4_row_bytes = (tail_dim + 1) / 2; + expect_corrupt_load_fails(q4_path, corrupt_path, [](std::vector & bytes) { + bytes[q4_vector_offset] = + static_cast(bytes[q4_vector_offset] & 0xf0u); // low nibble 0 is invalid + }); + expect_corrupt_load_fails(q4_path, corrupt_path, [](std::vector & bytes) { + const size_t tail_byte = q4_vector_offset + q4_row_bytes - 1; + bytes[tail_byte] = + static_cast((bytes[tail_byte] & 0x0fu) | 0x90u); // odd tail high nibble must be zero-code + }); + + ggml_vec_index_free(q4_loaded); + ggml_vec_index_free(q4); + std::filesystem::remove(q4_path); + } + + // q4 parity for a dimension that exercises the optimized loop and tail. + { + constexpr int q4_dim = 33; + std::vector q4_vector(q4_dim); + std::vector q4_query(q4_dim); + for (int i = 0; i < q4_dim; ++i) { + q4_vector[static_cast(i)] = + static_cast((i % 11) - 5) / 5.0f; + q4_query[static_cast(i)] = + static_cast(((i * 7) % 13) - 6) / 7.0f; + } + const uint64_t q4_id = (1ULL << 56) + 123ULL; + auto * q4 = ggml_vec_index_create(q4_dim, /*bit_width=*/4); + CHECK(q4 != nullptr); + CHECK(ggml_vec_index_add(q4, q4_vector.data(), 1, &q4_id) == GGML_VEC_INDEX_OK); + + std::array scores{}; + std::array out_ids{}; + CHECK(ggml_vec_index_search( + q4, q4_query.data(), 1, /*k=*/1, + scores.data(), out_ids.data()) == GGML_VEC_INDEX_OK); + CHECK(out_ids[0] == q4_id); + + const float expected = q4_dot_reference(q4_vector, q4_query); + const float tolerance = 1e-5f * std::max(1.0f, std::fabs(expected)); + CHECK(std::fabs(scores[0] - expected) <= tolerance); + ggml_vec_index_free(q4); + } + + // Delta replay keeps q8 storage q8 and reuses the normal q8 quantization path. + { + const std::string snapshot_path = + (std::filesystem::temp_directory_path() / + "ggml-vector-index-q8-delta-base.tvim").string(); + const std::string delta_path = + (std::filesystem::temp_directory_path() / + "ggml-vector-index-q8-delta-log.tvid").string(); + std::filesystem::remove(snapshot_path); + std::filesystem::remove(delta_path); + + const uint64_t base_id = (1ULL << 42) + 1ULL; + const uint64_t delta_id = (1ULL << 42) + 2ULL; + auto * q8_delta = ggml_vec_index_create(kDim, /*bit_width=*/8); + CHECK(q8_delta != nullptr); + CHECK(ggml_vec_index_add( + q8_delta, seeds[0].data(), 1, &base_id) == GGML_VEC_INDEX_OK); + CHECK(ggml_vec_index_write(q8_delta, snapshot_path.c_str()) == GGML_VEC_INDEX_OK); + CHECK(ggml_vec_index_add_logged( + q8_delta, seeds[3].data(), 1, &delta_id, delta_path.c_str()) == + GGML_VEC_INDEX_OK); + + auto * replayed_q8 = ggml_vec_index_load_with_delta( + snapshot_path.c_str(), delta_path.c_str()); + CHECK(replayed_q8 != nullptr); + CHECK(ggml_vec_index_bit_width(replayed_q8) == 8); + CHECK(ggml_vec_index_len(replayed_q8) == 2); + + std::array scores{}; + std::array out_ids{}; + CHECK(ggml_vec_index_search( + replayed_q8, seeds[3].data(), 1, /*k=*/1, + scores.data(), out_ids.data()) == GGML_VEC_INDEX_OK); + CHECK(out_ids[0] == delta_id); + CHECK(std::fabs(scores[0] - 1.0f) < 1e-5f); + + CHECK(ggml_vec_index_compact_delta( + q8_delta, snapshot_path.c_str(), delta_path.c_str()) == GGML_VEC_INDEX_OK); + CHECK(std::filesystem::file_size(delta_path) == 16); + auto * compacted_q8 = ggml_vec_index_load_with_delta( + snapshot_path.c_str(), delta_path.c_str()); + CHECK(compacted_q8 != nullptr); + CHECK(ggml_vec_index_bit_width(compacted_q8) == 8); + CHECK(ggml_vec_index_len(compacted_q8) == 2); + CHECK(ggml_vec_index_contains(compacted_q8, delta_id) == 1); + + ggml_vec_index_free(compacted_q8); + ggml_vec_index_free(replayed_q8); + ggml_vec_index_free(q8_delta); + std::filesystem::remove(snapshot_path); + std::filesystem::remove(delta_path); + } + // Header metadata is protected even when all payload sections are empty. { auto * empty_idx = ggml_vec_index_create(kDim, /*bit_width=*/8); @@ -465,6 +1005,46 @@ int main() { CHECK(scores[3] == -FLT_MAX); CHECK(out_ids[3] == UINT64_MAX); + const uint64_t q8_missing_id = (1ULL << 59) + 17ULL; + const std::array q8_allowed = { + q8_missing_id, q8_ids[0], q8_ids[0], + }; + scores = {}; + out_ids = {}; + CHECK(ggml_vec_index_search_filtered( + q8, seeds[2].data(), 1, /*k=*/3, + q8_allowed.data(), static_cast(q8_allowed.size()), + scores.data(), out_ids.data()) == GGML_VEC_INDEX_OK); + CHECK(out_ids[0] == q8_ids[0]); + CHECK(out_ids[1] == UINT64_MAX); + CHECK(out_ids[2] == UINT64_MAX); + CHECK(scores[1] == -FLT_MAX); + CHECK(scores[2] == -FLT_MAX); + + CHECK(ggml_vec_index_search_filtered( + q8, seeds[2].data(), 1, /*k=*/2, + nullptr, /*n_allowed=*/0, + scores.data(), out_ids.data()) == GGML_VEC_INDEX_OK); + CHECK(out_ids[0] == UINT64_MAX); + CHECK(out_ids[1] == UINT64_MAX); + CHECK(scores[0] == -FLT_MAX); + CHECK(scores[1] == -FLT_MAX); + + auto * q8_filter = ggml_vec_index_filter_create( + q8, q8_allowed.data(), static_cast(q8_allowed.size())); + CHECK(q8_filter != nullptr); + scores = {}; + out_ids = {}; + CHECK(ggml_vec_index_search_prepared_filtered( + q8, q8_filter, seeds[2].data(), 1, /*k=*/3, + scores.data(), out_ids.data()) == GGML_VEC_INDEX_OK); + CHECK(out_ids[0] == q8_ids[0]); + CHECK(out_ids[1] == UINT64_MAX); + CHECK(out_ids[2] == UINT64_MAX); + CHECK(scores[1] == -FLT_MAX); + CHECK(scores[2] == -FLT_MAX); + ggml_vec_index_filter_free(q8_filter); + const auto q8_tmp = std::filesystem::temp_directory_path() / "ggml-vector-index-q8-test.tvim"; const std::string q8_path = q8_tmp.string(); From 8b24b6c72c11ee1aa044068e4f89d150d58e107d Mon Sep 17 00:00:00 2001 From: Nidhin Date: Tue, 14 Jul 2026 22:36:54 +0530 Subject: [PATCH 06/13] llama : add production vector index features Adds q8/q4 storage and search, filtered search, IVF-flat ANN, incremental persistence, mmap loading, tombstone deletes, and in-memory compaction. Covers crash-safe snapshots, checksums, delta replay/compaction, SIMD kernels, fault tests, and vector index benchmarks. Assisted-by: GPT-5.5 --- ggml/include/ggml-vector-index.h | 19 +- ggml/src/ggml-vector-index.cpp | 760 ++++++++++++++++++++++++++----- tests/bench-vector-index.cpp | 98 ++++ tests/test-vector-index.cpp | 258 +++++++++++ 4 files changed, 1015 insertions(+), 120 deletions(-) diff --git a/ggml/include/ggml-vector-index.h b/ggml/include/ggml-vector-index.h index 52da1f930eaf..2ab010869851 100644 --- a/ggml/include/ggml-vector-index.h +++ b/ggml/include/ggml-vector-index.h @@ -70,11 +70,16 @@ GGML_API int ggml_vec_index_add( int n, const uint64_t * ids); -// Removes the entry for `id` via swap-with-last (slot indices are NOT stable -// across removes; external ids ARE). Returns 1 if removed, 0 if not present, -// negative on error. +// Removes the entry for `id` by marking its internal slot deleted. Physical +// storage is compacted only when writing a snapshot. Returns 1 if removed, +// 0 if not present, negative on error. GGML_API int ggml_vec_index_remove(ggml_vec_index_t * idx, uint64_t id); +// Physically removes deleted slots from in-memory storage. This does not write +// to disk. If any slots are removed, prepared filters and IVF state are +// invalidated. Returns 0 on success, negative on error. +GGML_API int ggml_vec_index_compact(ggml_vec_index_t * idx); + // Logged mutations for incremental persistence. These update `idx` and append // a durable delta record to `delta_path`. Replay the log on top of a full .tvim // snapshot with `ggml_vec_index_load_with_delta`. @@ -184,6 +189,14 @@ GGML_API int ggml_vec_index_write( // Returns NULL on failure. GGML_API ggml_vec_index_t * ggml_vec_index_load(const char * path); +// Loads a v2 .tvim snapshot with its vector section memory-mapped read-only. +// IDs and quantization scales are copied into memory for lookup and scoring. +// Mutating APIs return GGML_VEC_INDEX_E_INVALID_ARG on mmap-backed handles. +// `ggml_vec_index_write` can snapshot mmap-backed handles, but callers should +// write to a different path than the mapped source file. +// Returns NULL on failure or unsupported file format. +GGML_API ggml_vec_index_t * ggml_vec_index_load_mmap(const char * path); + // Loads a full .tvim snapshot and replays an append-only delta log. Missing // delta logs are treated as empty. GGML_API ggml_vec_index_t * ggml_vec_index_load_with_delta( diff --git a/ggml/src/ggml-vector-index.cpp b/ggml/src/ggml-vector-index.cpp index 3291c2d42aa3..cee1bd6a56fa 100644 --- a/ggml/src/ggml-vector-index.cpp +++ b/ggml/src/ggml-vector-index.cpp @@ -3,7 +3,7 @@ // // Storage: full f32 vectors or per-vector symmetric q8 codes. ID map uses // std::unordered_map for lookup and a parallel vector for -// the slot->id reverse map. Remove uses swap-with-last. +// the slot->id reverse map. Remove marks slots deleted; snapshots compact live rows. // // Search: dot product across all slots + min-heap of size k. q8 search scores // directly against stored codes and per-vector scales, with ARM NEON or x86 @@ -56,6 +56,7 @@ #ifndef _WIN32 #include +#include #include #include #else @@ -99,6 +100,53 @@ constexpr size_t kTvidRecordHeaderSize = 24; static_assert(sizeof(float) == sizeof(uint32_t), "ggml-vector-index requires float32"); +struct MappedFile { + void * data = nullptr; + size_t size = 0; +#ifdef _WIN32 + HANDLE file = INVALID_HANDLE_VALUE; + HANDLE mapping = nullptr; +#else + int fd = -1; +#endif + + MappedFile() = default; + ~MappedFile(); + MappedFile(const MappedFile &) = delete; + MappedFile & operator=(const MappedFile &) = delete; +}; + +void close_mapped_file(MappedFile & mapped) { +#ifdef _WIN32 + if (mapped.data != nullptr) { + UnmapViewOfFile(mapped.data); + mapped.data = nullptr; + } + if (mapped.mapping != nullptr) { + CloseHandle(mapped.mapping); + mapped.mapping = nullptr; + } + if (mapped.file != INVALID_HANDLE_VALUE) { + CloseHandle(mapped.file); + mapped.file = INVALID_HANDLE_VALUE; + } +#else + if (mapped.data != nullptr) { + munmap(mapped.data, mapped.size); + mapped.data = nullptr; + } + if (mapped.fd >= 0) { + close(mapped.fd); + mapped.fd = -1; + } +#endif + mapped.size = 0; +} + +MappedFile::~MappedFile() { + close_mapped_file(*this); +} + #ifdef GGML_VEC_INDEX_TEST_HOOKS std::atomic g_test_oom_countdown{ -1 }; std::atomic g_test_write_fail_after{ -1 }; @@ -166,6 +214,11 @@ uint64_t get_u64_le(const uint8_t * src) { return v; } +bool host_is_little_endian() { + const uint16_t value = 1; + return *reinterpret_cast(&value) == 1; +} + uint32_t float_to_u32(float v) { uint32_t bits; std::memcpy(&bits, &v, sizeof(bits)); @@ -406,6 +459,42 @@ bool open_temp_file(const char * path, TempFile & temp) { return false; } +bool map_file_readonly(const char * path, MappedFile & mapped) { + std::wstring wide; + if (!utf8_to_wide(path, wide)) { + return false; + } + mapped.file = CreateFileW( + wide.c_str(), + GENERIC_READ, + FILE_SHARE_READ, + nullptr, + OPEN_EXISTING, + FILE_ATTRIBUTE_NORMAL, + nullptr); + if (mapped.file == INVALID_HANDLE_VALUE) { + return false; + } + LARGE_INTEGER size; + if (!GetFileSizeEx(mapped.file, &size) || size.QuadPart <= 0 || + static_cast(size.QuadPart) > std::numeric_limits::max()) { + close_mapped_file(mapped); + return false; + } + mapped.mapping = CreateFileMappingW(mapped.file, nullptr, PAGE_READONLY, 0, 0, nullptr); + if (mapped.mapping == nullptr) { + close_mapped_file(mapped); + return false; + } + mapped.data = MapViewOfFile(mapped.mapping, FILE_MAP_READ, 0, 0, 0); + if (mapped.data == nullptr) { + close_mapped_file(mapped); + return false; + } + mapped.size = static_cast(size.QuadPart); + return true; +} + #else bool open_temp_file(const char * path, TempFile & temp) { @@ -435,6 +524,29 @@ bool open_temp_file(const char * path, TempFile & temp) { return true; } +bool map_file_readonly(const char * path, MappedFile & mapped) { + mapped.fd = ::open(path, O_RDONLY); + if (mapped.fd < 0) { + return false; + } + struct stat st; + if (::fstat(mapped.fd, &st) != 0 || st.st_size <= 0 || + static_cast(st.st_size) > std::numeric_limits::max()) { + close_mapped_file(mapped); + return false; + } + mapped.size = static_cast(st.st_size); + mapped.data = mmap(nullptr, mapped.size, PROT_READ, MAP_PRIVATE, mapped.fd, 0); + if (mapped.data == MAP_FAILED) { + mapped.data = nullptr; + close_mapped_file(mapped); + return false; + } + ::close(mapped.fd); + mapped.fd = -1; + return true; +} + #endif void remove_temp_file(TempFile & temp) { @@ -533,6 +645,13 @@ struct ggml_vec_index { int dim = 0; int bit_width = 32; uint64_t generation = 0; + bool read_only_mmap = false; + + std::unique_ptr mapped_file; + size_t mapped_vector_bytes = 0; + const float * mapped_data = nullptr; + const int8_t * mapped_q8_data = nullptr; + const uint8_t * mapped_q4_data = nullptr; // Flat row-major f32 storage for bit_width=32. std::vector data; @@ -547,6 +666,8 @@ struct ggml_vec_index { // slot -> external id (parallel to logical slot index). std::vector slot_to_id; + std::vector slot_active; + size_t n_active = 0; // external id -> slot. std::unordered_map id_to_slot; @@ -592,6 +713,56 @@ static size_t q4_row_bytes(size_t dim) { return (dim + 1) / 2; } +static size_t vector_bytes(const ggml_vec_index & idx) { + const size_t n = idx.slot_to_id.size(); + const size_t dim_sz = static_cast(idx.dim); + if (is_q4(idx)) { + return n * q4_row_bytes(dim_sz); + } + if (is_q8(idx)) { + return n * dim_sz * sizeof(int8_t); + } + return n * dim_sz * sizeof(float); +} + +static bool slot_is_active(const ggml_vec_index & idx, size_t slot) { + return slot < idx.slot_active.size() && idx.slot_active[slot] != 0; +} + +static size_t active_count(const ggml_vec_index & idx) { + return idx.n_active; +} + +static const float * f32_data_ptr(const ggml_vec_index & idx) { + return idx.mapped_data != nullptr ? idx.mapped_data : idx.data.data(); +} + +static const int8_t * q8_data_ptr(const ggml_vec_index & idx) { + return idx.mapped_q8_data != nullptr ? idx.mapped_q8_data : idx.q8_data.data(); +} + +static const uint8_t * q4_data_ptr(const ggml_vec_index & idx) { + return idx.mapped_q4_data != nullptr ? idx.mapped_q4_data : idx.q4_data.data(); +} + +static bool has_vector_storage(const ggml_vec_index & idx) { + const size_t bytes = vector_bytes(idx); + if (idx.read_only_mmap) { + return idx.mapped_vector_bytes == bytes && + (bytes == 0 || + idx.mapped_data != nullptr || + idx.mapped_q8_data != nullptr || + idx.mapped_q4_data != nullptr); + } + if (is_q4(idx)) { + return idx.q4_data.size() == bytes; + } + if (is_q8(idx)) { + return idx.q8_data.size() == bytes; + } + return idx.data.size() == bytes / sizeof(float); +} + static uint8_t q4_encode(int q) { return static_cast(q + 8); } @@ -673,28 +844,49 @@ uint32_t index_state_crc32c(const ggml_vec_index & idx) { crc = crc32c_update_u32(crc, static_cast(idx.dim)); crc = crc32c_update_u32(crc, static_cast(idx.bit_width)); crc = crc32c_update_u32(crc, static_cast(storage_kind(idx))); - crc = crc32c_update_u64(crc, static_cast(idx.slot_to_id.size())); + crc = crc32c_update_u64(crc, static_cast(active_count(idx))); if (is_q4(idx)) { - for (float scale : idx.q4_scale) { - crc = crc32c_update_u32(crc, float_to_u32(scale)); + const size_t row_bytes = q4_row_bytes(static_cast(idx.dim)); + const uint8_t * data = q4_data_ptr(idx); + for (size_t slot = 0; slot < idx.slot_to_id.size(); ++slot) { + if (slot_is_active(idx, slot)) { + crc = crc32c_update_u32(crc, float_to_u32(idx.q4_scale[slot])); + } } - if (!idx.q4_data.empty()) { - crc = crc32c_update(crc, idx.q4_data.data(), idx.q4_data.size()); + for (size_t slot = 0; slot < idx.slot_to_id.size(); ++slot) { + if (slot_is_active(idx, slot)) { + crc = crc32c_update(crc, data + slot * row_bytes, row_bytes); + } } } else if (is_q8(idx)) { - for (float scale : idx.q8_scale) { - crc = crc32c_update_u32(crc, float_to_u32(scale)); + const size_t dim_sz = static_cast(idx.dim); + const int8_t * data = q8_data_ptr(idx); + for (size_t slot = 0; slot < idx.slot_to_id.size(); ++slot) { + if (slot_is_active(idx, slot)) { + crc = crc32c_update_u32(crc, float_to_u32(idx.q8_scale[slot])); + } } - if (!idx.q8_data.empty()) { - crc = crc32c_update(crc, idx.q8_data.data(), idx.q8_data.size()); + for (size_t slot = 0; slot < idx.slot_to_id.size(); ++slot) { + if (slot_is_active(idx, slot)) { + crc = crc32c_update(crc, data + slot * dim_sz, dim_sz * sizeof(int8_t)); + } } } else { - for (float v : idx.data) { - crc = crc32c_update_u32(crc, float_to_u32(v)); + const float * data = f32_data_ptr(idx); + const size_t dim_sz = static_cast(idx.dim); + for (size_t slot = 0; slot < idx.slot_to_id.size(); ++slot) { + if (!slot_is_active(idx, slot)) { + continue; + } + for (size_t i = 0; i < dim_sz; ++i) { + crc = crc32c_update_u32(crc, float_to_u32(data[slot * dim_sz + i])); + } } } - for (uint64_t id : idx.slot_to_id) { - crc = crc32c_update_u64(crc, id); + for (size_t slot = 0; slot < idx.slot_to_id.size(); ++slot) { + if (slot_is_active(idx, slot)) { + crc = crc32c_update_u64(crc, idx.slot_to_id[slot]); + } } return crc ^ 0xffffffffu; } @@ -705,52 +897,54 @@ uint32_t index_state_crc32c_after_remove(const ggml_vec_index & idx, uint64_t id return index_state_crc32c(idx); } const size_t removed = it->second; - const size_t last = idx.slot_to_id.size() - 1; const size_t dim_sz = static_cast(idx.dim); uint32_t crc = 0xffffffffu; crc = crc32c_update_u32(crc, static_cast(idx.dim)); crc = crc32c_update_u32(crc, static_cast(idx.bit_width)); crc = crc32c_update_u32(crc, static_cast(storage_kind(idx))); - crc = crc32c_update_u64(crc, static_cast(last)); - - auto logical_slot = [&](size_t out_slot) { - return out_slot == removed ? last : out_slot; - }; + crc = crc32c_update_u64(crc, static_cast(active_count(idx) - 1)); if (is_q4(idx)) { const size_t row_bytes = q4_row_bytes(dim_sz); - for (size_t slot = 0; slot < last; ++slot) { - crc = crc32c_update_u32(crc, float_to_u32(idx.q4_scale[logical_slot(slot)])); + const uint8_t * data = q4_data_ptr(idx); + for (size_t slot = 0; slot < idx.slot_to_id.size(); ++slot) { + if (slot_is_active(idx, slot) && slot != removed) { + crc = crc32c_update_u32(crc, float_to_u32(idx.q4_scale[slot])); + } } - for (size_t slot = 0; slot < last; ++slot) { - const size_t src_slot = logical_slot(slot); - crc = crc32c_update( - crc, - idx.q4_data.data() + src_slot * row_bytes, - row_bytes * sizeof(uint8_t)); + for (size_t slot = 0; slot < idx.slot_to_id.size(); ++slot) { + if (slot_is_active(idx, slot) && slot != removed) { + crc = crc32c_update(crc, data + slot * row_bytes, row_bytes); + } } } else if (is_q8(idx)) { - for (size_t slot = 0; slot < last; ++slot) { - crc = crc32c_update_u32(crc, float_to_u32(idx.q8_scale[logical_slot(slot)])); + const int8_t * data = q8_data_ptr(idx); + for (size_t slot = 0; slot < idx.slot_to_id.size(); ++slot) { + if (slot_is_active(idx, slot) && slot != removed) { + crc = crc32c_update_u32(crc, float_to_u32(idx.q8_scale[slot])); + } } - for (size_t slot = 0; slot < last; ++slot) { - const size_t src_slot = logical_slot(slot); - crc = crc32c_update( - crc, - idx.q8_data.data() + src_slot * dim_sz, - dim_sz * sizeof(int8_t)); + for (size_t slot = 0; slot < idx.slot_to_id.size(); ++slot) { + if (slot_is_active(idx, slot) && slot != removed) { + crc = crc32c_update(crc, data + slot * dim_sz, dim_sz * sizeof(int8_t)); + } } } else { - for (size_t slot = 0; slot < last; ++slot) { - const size_t src_slot = logical_slot(slot); + const float * data = f32_data_ptr(idx); + for (size_t slot = 0; slot < idx.slot_to_id.size(); ++slot) { + if (!slot_is_active(idx, slot) || slot == removed) { + continue; + } for (size_t i = 0; i < dim_sz; ++i) { - crc = crc32c_update_u32(crc, float_to_u32(idx.data[src_slot * dim_sz + i])); + crc = crc32c_update_u32(crc, float_to_u32(data[slot * dim_sz + i])); } } } - for (size_t slot = 0; slot < last; ++slot) { - crc = crc32c_update_u64(crc, idx.slot_to_id[logical_slot(slot)]); + for (size_t slot = 0; slot < idx.slot_to_id.size(); ++slot) { + if (slot_is_active(idx, slot) && slot != removed) { + crc = crc32c_update_u64(crc, idx.slot_to_id[slot]); + } } return crc ^ 0xffffffffu; } @@ -1192,12 +1386,17 @@ static int ggml_vec_index_add_unlocked( idx->data.resize(base_slot * dim_sz); } idx->slot_to_id.resize(base_slot); + idx->slot_active.resize(base_slot); + idx->n_active = idx->id_to_slot.size(); }; try { if (idx == nullptr || vectors == nullptr || ids == nullptr) { return GGML_VEC_INDEX_E_INVALID_ARG; } + if (idx->read_only_mmap) { + return GGML_VEC_INDEX_E_INVALID_ARG; + } if (n < 0) { return GGML_VEC_INDEX_E_INVALID_ARG; } @@ -1244,6 +1443,7 @@ static int ggml_vec_index_add_unlocked( idx->data.resize(new_slots * dim_sz); } idx->slot_to_id.resize(new_slots); + idx->slot_active.resize(new_slots, 0); test_maybe_throw_bad_alloc(); idx->id_to_slot.reserve(new_slots); @@ -1269,9 +1469,11 @@ static int ggml_vec_index_add_unlocked( dim_sz * sizeof(float)); } idx->slot_to_id[slot] = ids[i]; + idx->slot_active[slot] = 1; test_maybe_throw_bad_alloc(); idx->id_to_slot.emplace(ids[i], slot); } + idx->n_active += n_sz; ++idx->generation; invalidate_ivf(*idx); } catch (const std::bad_alloc &) { @@ -1305,66 +1507,149 @@ static int ggml_vec_index_remove_unlocked(ggml_vec_index_t * idx, uint64_t id) { if (idx == nullptr) { return GGML_VEC_INDEX_E_INVALID_ARG; } + if (idx->read_only_mmap) { + return GGML_VEC_INDEX_E_INVALID_ARG; + } auto it = idx->id_to_slot.find(id); if (it == idx->id_to_slot.end()) { return 0; } - const size_t slot = it->second; - const size_t last = idx->slot_to_id.size() - 1; - const size_t dim_sz = static_cast(idx->dim); + const size_t slot = it->second; + if (!slot_is_active(*idx, slot)) { + idx->id_to_slot.erase(it); + return 0; + } + idx->slot_active[slot] = 0; + --idx->n_active; + idx->id_to_slot.erase(it); + ++idx->generation; + invalidate_ivf(*idx); + return 1; + } catch (...) { + return GGML_VEC_INDEX_E_INTERNAL; + } +} - if (slot != last) { - // Move last vector into the freed slot and update its id mapping. - if (is_q4(*idx)) { - const size_t row_bytes = q4_row_bytes(dim_sz); +int ggml_vec_index_remove(ggml_vec_index_t * idx, uint64_t id) { + if (idx == nullptr) { + return GGML_VEC_INDEX_E_INVALID_ARG; + } + try { + std::unique_lock lock(idx->mutex); + return ggml_vec_index_remove_unlocked(idx, id); + } catch (...) { + return GGML_VEC_INDEX_E_INTERNAL; + } +} + +static int ggml_vec_index_compact_unlocked(ggml_vec_index_t * idx) { + try { + if (idx == nullptr) { + return GGML_VEC_INDEX_E_INVALID_ARG; + } + if (idx->read_only_mmap) { + return GGML_VEC_INDEX_E_INVALID_ARG; + } + + const size_t n_slots = idx->slot_to_id.size(); + const size_t n_live = active_count(*idx); + if (n_live == n_slots) { + return GGML_VEC_INDEX_OK; + } + + const size_t dim_sz = static_cast(idx->dim); + if (dim_sz != 0 && n_live > std::numeric_limits::max() / dim_sz) { + return GGML_VEC_INDEX_E_INTERNAL; + } + + test_maybe_throw_bad_alloc(); + std::vector new_slot_to_id; + std::vector new_slot_active; + std::unordered_map new_id_to_slot; + new_slot_to_id.reserve(n_live); + new_slot_active.assign(n_live, 1); + new_id_to_slot.reserve(n_live); + + if (is_q4(*idx)) { + const size_t row_bytes = q4_row_bytes(dim_sz); + std::vector new_q4_data; + std::vector new_q4_scale; + new_q4_data.resize(n_live * row_bytes); + new_q4_scale.reserve(n_live); + for (size_t slot = 0; slot < n_slots; ++slot) { + if (!slot_is_active(*idx, slot)) { + continue; + } + const size_t out_slot = new_slot_to_id.size(); std::memcpy( + new_q4_data.data() + out_slot * row_bytes, idx->q4_data.data() + slot * row_bytes, - idx->q4_data.data() + last * row_bytes, row_bytes * sizeof(uint8_t)); - idx->q4_scale[slot] = idx->q4_scale[last]; - } else if (is_q8(*idx)) { + new_q4_scale.push_back(idx->q4_scale[slot]); + new_slot_to_id.push_back(idx->slot_to_id[slot]); + new_id_to_slot.emplace(idx->slot_to_id[slot], out_slot); + } + idx->q4_data.swap(new_q4_data); + idx->q4_scale.swap(new_q4_scale); + } else if (is_q8(*idx)) { + std::vector new_q8_data; + std::vector new_q8_scale; + new_q8_data.resize(n_live * dim_sz); + new_q8_scale.reserve(n_live); + for (size_t slot = 0; slot < n_slots; ++slot) { + if (!slot_is_active(*idx, slot)) { + continue; + } + const size_t out_slot = new_slot_to_id.size(); std::memcpy( + new_q8_data.data() + out_slot * dim_sz, idx->q8_data.data() + slot * dim_sz, - idx->q8_data.data() + last * dim_sz, dim_sz * sizeof(int8_t)); - idx->q8_scale[slot] = idx->q8_scale[last]; - } else { + new_q8_scale.push_back(idx->q8_scale[slot]); + new_slot_to_id.push_back(idx->slot_to_id[slot]); + new_id_to_slot.emplace(idx->slot_to_id[slot], out_slot); + } + idx->q8_data.swap(new_q8_data); + idx->q8_scale.swap(new_q8_scale); + } else { + std::vector new_data; + new_data.resize(n_live * dim_sz); + for (size_t slot = 0; slot < n_slots; ++slot) { + if (!slot_is_active(*idx, slot)) { + continue; + } + const size_t out_slot = new_slot_to_id.size(); std::memcpy( + new_data.data() + out_slot * dim_sz, idx->data.data() + slot * dim_sz, - idx->data.data() + last * dim_sz, dim_sz * sizeof(float)); + new_slot_to_id.push_back(idx->slot_to_id[slot]); + new_id_to_slot.emplace(idx->slot_to_id[slot], out_slot); } - const uint64_t moved_id = idx->slot_to_id[last]; - idx->slot_to_id[slot] = moved_id; - idx->id_to_slot[moved_id] = slot; + idx->data.swap(new_data); } - idx->slot_to_id.pop_back(); - if (is_q4(*idx)) { - idx->q4_data.resize(last * q4_row_bytes(dim_sz)); - idx->q4_scale.resize(last); - } else if (is_q8(*idx)) { - idx->q8_data.resize(last * dim_sz); - idx->q8_scale.resize(last); - } else { - idx->data.resize(last * dim_sz); - } - idx->id_to_slot.erase(it); + idx->slot_to_id.swap(new_slot_to_id); + idx->slot_active.swap(new_slot_active); + idx->id_to_slot.swap(new_id_to_slot); + idx->n_active = idx->slot_to_id.size(); ++idx->generation; invalidate_ivf(*idx); - return 1; + return GGML_VEC_INDEX_OK; + } catch (const std::bad_alloc &) { + return GGML_VEC_INDEX_E_OOM; } catch (...) { return GGML_VEC_INDEX_E_INTERNAL; } } -int ggml_vec_index_remove(ggml_vec_index_t * idx, uint64_t id) { +int ggml_vec_index_compact(ggml_vec_index_t * idx) { if (idx == nullptr) { return GGML_VEC_INDEX_E_INVALID_ARG; } try { std::unique_lock lock(idx->mutex); - return ggml_vec_index_remove_unlocked(idx, id); + return ggml_vec_index_compact_unlocked(idx); } catch (...) { return GGML_VEC_INDEX_E_INTERNAL; } @@ -1383,6 +1668,9 @@ int ggml_vec_index_add_logged( return GGML_VEC_INDEX_E_INVALID_ARG; } lock = std::unique_lock(idx->mutex); + if (idx->read_only_mmap) { + return GGML_VEC_INDEX_E_INVALID_ARG; + } if (!validate_logged_add_args(idx, vectors, n, ids)) { return GGML_VEC_INDEX_E_INVALID_ARG; } @@ -1450,6 +1738,9 @@ int ggml_vec_index_remove_logged( return GGML_VEC_INDEX_E_INVALID_ARG; } std::unique_lock lock(idx->mutex); + if (idx->read_only_mmap) { + return GGML_VEC_INDEX_E_INVALID_ARG; + } if (idx->id_to_slot.count(id) == 0) { return 0; } @@ -1688,18 +1979,18 @@ inline float score_slot(const ggml_vec_index_t & idx, const float * query, size_ return is_q4(idx) ? dot_q4( query, - idx.q4_data.data() + slot * q4_row_bytes(static_cast(dim)), + q4_data_ptr(idx) + slot * q4_row_bytes(static_cast(dim)), idx.q4_scale[slot], dim) : is_q8(idx) ? dot_q8( query, - idx.q8_data.data() + slot * static_cast(dim), + q8_data_ptr(idx) + slot * static_cast(dim), idx.q8_scale[slot], dim) : dot( query, - idx.data.data() + slot * static_cast(dim), + f32_data_ptr(idx) + slot * static_cast(dim), dim); } @@ -1707,7 +1998,7 @@ void decode_slot_to_f32(const ggml_vec_index_t & idx, size_t slot, float * dst) const int dim = idx.dim; if (is_q4(idx)) { const uint8_t * codes = - idx.q4_data.data() + slot * q4_row_bytes(static_cast(dim)); + q4_data_ptr(idx) + slot * q4_row_bytes(static_cast(dim)); const float scale = idx.q4_scale[slot]; for (int i = 0; i < dim; ++i) { const uint8_t byte = codes[static_cast(i) / 2]; @@ -1717,7 +2008,7 @@ void decode_slot_to_f32(const ggml_vec_index_t & idx, size_t slot, float * dst) dst[i] = static_cast(q4_decode(nibble)) * scale; } } else if (is_q8(idx)) { - const int8_t * codes = idx.q8_data.data() + slot * static_cast(dim); + const int8_t * codes = q8_data_ptr(idx) + slot * static_cast(dim); const float scale = idx.q8_scale[slot]; for (int i = 0; i < dim; ++i) { dst[i] = static_cast(codes[i]) * scale; @@ -1725,7 +2016,7 @@ void decode_slot_to_f32(const ggml_vec_index_t & idx, size_t slot, float * dst) } else { std::memcpy( dst, - idx.data.data() + slot * static_cast(dim), + f32_data_ptr(idx) + slot * static_cast(dim), static_cast(dim) * sizeof(float)); } } @@ -1759,6 +2050,9 @@ void search_one( std::priority_queue, MinHeapCmp> heap; auto visit_slot = [&](size_t slot) { + if (!slot_is_active(idx, slot)) { + return; + } const float s = score_slot(idx, query, slot); if (heap.size() < static_cast(k)) { heap.push({ s, idx.slot_to_id[slot] }); @@ -1808,7 +2102,7 @@ std::vector allowed_slots_for_ids( slots.reserve(static_cast(n_allowed)); for (int i = 0; i < n_allowed; ++i) { const auto it = idx.id_to_slot.find(allowed_ids[i]); - if (it != idx.id_to_slot.end()) { + if (it != idx.id_to_slot.end() && slot_is_active(idx, it->second)) { slots.push_back(it->second); } } @@ -1826,15 +2120,16 @@ static int ggml_vec_index_build_ivf_unlocked(ggml_vec_index_t * idx, int n_lists } const size_t n_slots = idx->slot_to_id.size(); + const size_t n_live = active_count(*idx); const int dim = idx->dim; - if (n_slots == 0) { + if (n_live == 0) { invalidate_ivf(*idx); idx->ivf_generation = idx->generation; return GGML_VEC_INDEX_OK; } const int actual_lists = static_cast( - std::min(static_cast(n_lists), n_slots)); + std::min(static_cast(n_lists), n_live)); const size_t dim_sz = static_cast(dim); test_maybe_throw_bad_alloc(); @@ -1843,10 +2138,17 @@ static int ggml_vec_index_build_ivf_unlocked(ggml_vec_index_t * idx, int n_lists std::vector counts(static_cast(actual_lists)); std::vector row(dim_sz); std::vector> lists(static_cast(actual_lists)); + std::vector active_slots; + active_slots.reserve(n_live); + for (size_t slot = 0; slot < n_slots; ++slot) { + if (slot_is_active(*idx, slot)) { + active_slots.push_back(slot); + } + } for (int list = 0; list < actual_lists; ++list) { - const size_t slot = static_cast(list) * n_slots / - static_cast(actual_lists); + const size_t slot = active_slots[static_cast(list) * active_slots.size() / + static_cast(actual_lists)]; float * centroid = centroids.data() + static_cast(list) * dim_sz; decode_slot_to_f32(*idx, slot, centroid); } @@ -1855,7 +2157,7 @@ static int ggml_vec_index_build_ivf_unlocked(ggml_vec_index_t * idx, int n_lists std::fill(next_centroids.begin(), next_centroids.end(), 0.0f); std::fill(counts.begin(), counts.end(), 0); - for (size_t slot = 0; slot < n_slots; ++slot) { + for (size_t slot : active_slots) { decode_slot_to_f32(*idx, slot, row.data()); const size_t list = best_centroid(row.data(), centroids, actual_lists, dim); float * dst = next_centroids.data() + list * dim_sz; @@ -1880,7 +2182,7 @@ static int ggml_vec_index_build_ivf_unlocked(ggml_vec_index_t * idx, int n_lists } } - for (size_t slot = 0; slot < n_slots; ++slot) { + for (size_t slot : active_slots) { decode_slot_to_f32(*idx, slot, row.data()); const size_t list = best_centroid(row.data(), centroids, actual_lists, dim); lists[list].push_back(slot); @@ -2147,23 +2449,18 @@ static int ggml_vec_index_write_unlocked(ggml_vec_index_t * idx, const char * pa if (idx == nullptr || path == nullptr) { return GGML_VEC_INDEX_E_INVALID_ARG; } - if (idx->slot_to_id.size() > std::numeric_limits::max()) { + if (active_count(*idx) > std::numeric_limits::max()) { return GGML_VEC_INDEX_E_INVALID_ARG; } - const size_t n = idx->slot_to_id.size(); + const size_t n = active_count(*idx); + const size_t n_slots = idx->slot_to_id.size(); const size_t dim_sz = static_cast(idx->dim); - if (dim_sz != 0 && n > std::numeric_limits::max() / dim_sz) { + if (dim_sz != 0 && n_slots > std::numeric_limits::max() / dim_sz) { return GGML_VEC_INDEX_E_INTERNAL; } - if (is_q4(*idx)) { - if (idx->q4_data.size() != n * q4_row_bytes(dim_sz) || idx->q4_scale.size() != n) { - return GGML_VEC_INDEX_E_INTERNAL; - } - } else if (is_q8(*idx)) { - if (idx->q8_data.size() != n * dim_sz || idx->q8_scale.size() != n) { - return GGML_VEC_INDEX_E_INTERNAL; - } - } else if (idx->data.size() != n * dim_sz) { + if (!has_vector_storage(*idx) || + (is_q4(*idx) && idx->q4_scale.size() != n_slots) || + (is_q8(*idx) && idx->q8_scale.size() != n_slots)) { return GGML_VEC_INDEX_E_INTERNAL; } @@ -2191,7 +2488,7 @@ static int ggml_vec_index_write_unlocked(ggml_vec_index_t * idx, const char * pa header[6] = storage_kind(*idx); header[7] = kFlagCRC32C; const uint32_t dim_le = static_cast(idx->dim); - const uint32_t n_le = static_cast(idx->slot_to_id.size()); + const uint32_t n_le = static_cast(n); put_u32_le(header + 8, dim_le); put_u32_le(header + 12, n_le); put_u32_le(header + 16, is_quantized(*idx) ? kQParamScaleF32 : kQParamNone); @@ -2209,38 +2506,63 @@ static int ggml_vec_index_write_unlocked(ggml_vec_index_t * idx, const char * pa uint32_t ids_crc = 0xffffffffu; if (is_quantized(*idx)) { const std::vector & scales = is_q4(*idx) ? idx->q4_scale : idx->q8_scale; - for (float scale : scales) { + for (size_t slot = 0; slot < n_slots; ++slot) { + if (!slot_is_active(*idx, slot)) { + continue; + } + const float scale = scales[slot]; if (!write_u32_le_crc(f, float_to_u32(scale), qparams_crc)) { return fail_io(); } } if (is_q4(*idx)) { - if (!idx->q4_data.empty() && - !write_bytes(f, idx->q4_data.data(), idx->q4_data.size())) { - return fail_io(); - } - if (!idx->q4_data.empty()) { - vectors_crc = - crc32c_update(vectors_crc, idx->q4_data.data(), idx->q4_data.size()); + const size_t row_bytes = q4_row_bytes(dim_sz); + const uint8_t * data = q4_data_ptr(*idx); + for (size_t slot = 0; slot < n_slots; ++slot) { + if (!slot_is_active(*idx, slot)) { + continue; + } + const uint8_t * row = data + slot * row_bytes; + if (!write_bytes(f, row, row_bytes)) { + return fail_io(); + } + vectors_crc = crc32c_update(vectors_crc, row, row_bytes); } - } else if (!idx->q8_data.empty()) { - if (!write_bytes(f, idx->q8_data.data(), idx->q8_data.size())) { - return fail_io(); + } else { + const size_t row_bytes = dim_sz * sizeof(int8_t); + const int8_t * data = q8_data_ptr(*idx); + for (size_t slot = 0; slot < n_slots; ++slot) { + if (!slot_is_active(*idx, slot)) { + continue; + } + const int8_t * row = data + slot * dim_sz; + if (!write_bytes(f, row, row_bytes)) { + return fail_io(); + } + vectors_crc = crc32c_update(vectors_crc, row, row_bytes); } - vectors_crc = - crc32c_update(vectors_crc, idx->q8_data.data(), idx->q8_data.size()); } } else { - for (float v : idx->data) { - if (!write_u32_le_crc(f, float_to_u32(v), vectors_crc)) { - return fail_io(); + const float * data = f32_data_ptr(*idx); + for (size_t slot = 0; slot < n_slots; ++slot) { + if (!slot_is_active(*idx, slot)) { + continue; + } + for (size_t i = 0; i < dim_sz; ++i) { + const float v = data[slot * dim_sz + i]; + if (!write_u32_le_crc(f, float_to_u32(v), vectors_crc)) { + return fail_io(); + } } } } - for (uint64_t id : idx->slot_to_id) { - if (!write_u64_le_crc(f, id, ids_crc)) { + for (size_t slot = 0; slot < n_slots; ++slot) { + if (!slot_is_active(*idx, slot)) { + continue; + } + if (!write_u64_le_crc(f, idx->slot_to_id[slot], ids_crc)) { return fail_io(); } } @@ -2424,6 +2746,8 @@ ggml_vec_index_t * ggml_vec_index_load(const char * path) { idx->data.resize(n * dim_sz); } idx->slot_to_id.resize(n); + idx->slot_active.assign(n, 1); + idx->n_active = n; idx->id_to_slot.reserve(n); const bool checksummed = (flags & kFlagCRC32C) != 0; @@ -2597,6 +2921,208 @@ ggml_vec_index_t * ggml_vec_index_load(const char * path) { } } +ggml_vec_index_t * ggml_vec_index_load_mmap(const char * path) { + try { + if (path == nullptr || !host_is_little_endian()) { + return nullptr; + } + + auto mapped = std::make_unique(); + if (!map_file_readonly(path, *mapped) || mapped->size < kTvimHeaderSize) { + return nullptr; + } + + const auto * bytes = static_cast(mapped->data); + if (std::memcmp(bytes, kTvimMagic, 4) != 0 || bytes[4] != kTvimVersion) { + return nullptr; + } + + const uint8_t flags = bytes[7]; + if ((flags & ~kFlagCRC32C) != 0 || get_u32_le(bytes + 28) != 0) { + return nullptr; + } + + const int bit_width = static_cast(bytes[5]); + if (!is_supported_bit_width(bit_width)) { + return nullptr; + } + const uint8_t kind = bytes[6]; + const uint32_t dim_le = get_u32_le(bytes + 8); + const uint32_t n_le = get_u32_le(bytes + 12); + const uint32_t qparam_type = get_u32_le(bytes + 16); + const uint32_t qparam_bytes = get_u32_le(bytes + 20); + const uint32_t comp_bytes = get_u32_le(bytes + 24); + if (dim_le == 0 || dim_le > static_cast(std::numeric_limits::max())) { + return nullptr; + } + if ((bit_width == 4 && (kind != kStorageQ4 || qparam_type != kQParamScaleF32 || + qparam_bytes != 4 || comp_bytes != 0)) || + (bit_width == 8 && (kind != kStorageQ8 || qparam_type != kQParamScaleF32 || + qparam_bytes != 4 || comp_bytes != 1)) || + (bit_width == 32 && (kind != kStorageF32 || qparam_type != kQParamNone || + qparam_bytes != 0 || comp_bytes != 4))) { + return nullptr; + } + + uint64_t expected_size = 0; + if (!expected_file_size( + kTvimHeaderSize, + n_le, + dim_le, + qparam_bytes, + comp_bytes, + expected_size)) { + return nullptr; + } + const bool checksummed = (flags & kFlagCRC32C) != 0; + if (checksummed && !checked_add_u64(expected_size, kTvimChecksumSize, expected_size)) { + return nullptr; + } + if (expected_size != static_cast(mapped->size)) { + return nullptr; + } + + uint64_t qparams_bytes_u64 = 0; + uint64_t vectors_bytes_u64 = 0; + uint64_t ids_bytes_u64 = 0; + uint64_t component_count_u64 = 0; + if (!checked_mul_u64(n_le, qparam_bytes, qparams_bytes_u64) || + !checked_mul_u64(n_le, sizeof(uint64_t), ids_bytes_u64)) { + return nullptr; + } + if (comp_bytes == 0) { + uint64_t row_bytes = 0; + if (!checked_add_u64(dim_le, 1, row_bytes)) { + return nullptr; + } + row_bytes /= 2; + if (!checked_mul_u64(n_le, row_bytes, vectors_bytes_u64)) { + return nullptr; + } + } else if (!checked_mul_u64(n_le, dim_le, component_count_u64) || + !checked_mul_u64(component_count_u64, comp_bytes, vectors_bytes_u64)) { + return nullptr; + } + const uint64_t qparams_offset = kTvimHeaderSize; + const uint64_t vectors_offset = qparams_offset + qparams_bytes_u64; + const uint64_t ids_offset = vectors_offset + vectors_bytes_u64; + const uint64_t checksums_offset = ids_offset + ids_bytes_u64; + + if (checksummed) { + const uint32_t header_crc = + crc32c_update(0xffffffffu, bytes, kTvimHeaderSize) ^ 0xffffffffu; + const uint32_t qparams_crc = + crc32c_update( + 0xffffffffu, + bytes + static_cast(qparams_offset), + static_cast(qparams_bytes_u64)) ^ 0xffffffffu; + const uint32_t vectors_crc = + crc32c_update( + 0xffffffffu, + bytes + static_cast(vectors_offset), + static_cast(vectors_bytes_u64)) ^ 0xffffffffu; + const uint32_t ids_crc = + crc32c_update( + 0xffffffffu, + bytes + static_cast(ids_offset), + static_cast(ids_bytes_u64)) ^ 0xffffffffu; + if (header_crc != get_u32_le(bytes + static_cast(checksums_offset)) || + qparams_crc != get_u32_le(bytes + static_cast(checksums_offset + 4)) || + vectors_crc != get_u32_le(bytes + static_cast(checksums_offset + 8)) || + ids_crc != get_u32_le(bytes + static_cast(checksums_offset + 12))) { + return nullptr; + } + } + + std::unique_ptr idx( + ggml_vec_index_create(static_cast(dim_le), bit_width), + ggml_vec_index_free); + if (idx == nullptr) { + return nullptr; + } + + const size_t n = static_cast(n_le); + const size_t dim_sz = static_cast(dim_le); + idx->slot_to_id.resize(n); + idx->slot_active.assign(n, 1); + idx->n_active = n; + idx->id_to_slot.reserve(n); + if (is_quantized(*idx)) { + std::vector & scales = is_q4(*idx) ? idx->q4_scale : idx->q8_scale; + scales.resize(n); + for (size_t slot = 0; slot < n; ++slot) { + const uint32_t bits = get_u32_le( + bytes + static_cast(qparams_offset) + slot * sizeof(uint32_t)); + scales[slot] = u32_to_float(bits); + if (!std::isfinite(scales[slot]) || scales[slot] <= 0.0f) { + return nullptr; + } + } + } + + if (bit_width == 4) { + const size_t row_bytes = q4_row_bytes(dim_sz); + const auto * q4 = bytes + static_cast(vectors_offset); + for (size_t slot = 0; slot < n; ++slot) { + const float scale = idx->q4_scale[slot]; + const uint8_t * row = q4 + slot * row_bytes; + for (size_t i = 0; i < dim_sz; ++i) { + const uint8_t byte = row[i / 2]; + const uint8_t nibble = (i & 1) == 0 ? + static_cast(byte & 0x0f) : + static_cast(byte >> 4); + if (nibble == 0 || + !std::isfinite(static_cast(q4_decode(nibble)) * scale)) { + return nullptr; + } + } + if ((dim_sz & 1) != 0 && (row[row_bytes - 1] >> 4) != 8) { + return nullptr; + } + } + idx->mapped_q4_data = q4; + } else if (bit_width == 8) { + const auto * q8 = reinterpret_cast( + bytes + static_cast(vectors_offset)); + for (size_t slot = 0; slot < n; ++slot) { + const float scale = idx->q8_scale[slot]; + const int8_t * row = q8 + slot * dim_sz; + for (size_t i = 0; i < dim_sz; ++i) { + if (row[i] == std::numeric_limits::min() || + !std::isfinite(static_cast(row[i]) * scale)) { + return nullptr; + } + } + } + idx->mapped_q8_data = q8; + } else { + const auto * f32 = reinterpret_cast( + bytes + static_cast(vectors_offset)); + const size_t count = n * dim_sz; + if (!all_finite(f32, count)) { + return nullptr; + } + idx->mapped_data = f32; + } + + const uint8_t * ids = bytes + static_cast(ids_offset); + for (size_t slot = 0; slot < n; ++slot) { + const uint64_t id = get_u64_le(ids + slot * sizeof(uint64_t)); + idx->slot_to_id[slot] = id; + if (!idx->id_to_slot.emplace(id, slot).second) { + return nullptr; + } + } + + idx->read_only_mmap = true; + idx->mapped_vector_bytes = static_cast(vectors_bytes_u64); + idx->mapped_file = std::move(mapped); + return idx.release(); + } catch (...) { + return nullptr; + } +} + namespace { bool expected_add_delta_payload_size(uint64_t n, uint64_t dim, uint64_t & size) { @@ -2824,7 +3350,7 @@ int ggml_vec_index_len(const ggml_vec_index_t * idx) { } try { std::shared_lock lock(idx->mutex); - return static_cast(idx->slot_to_id.size()); + return static_cast(active_count(*idx)); } catch (...) { return 0; } diff --git a/tests/bench-vector-index.cpp b/tests/bench-vector-index.cpp index f3b873090c45..d035a6b1e19e 100644 --- a/tests/bench-vector-index.cpp +++ b/tests/bench-vector-index.cpp @@ -36,6 +36,7 @@ struct BenchConfig { int ivf_lists = 64; int ivf_iters = 4; int ivf_nprobe = 4; + int delete_stride = 2; }; struct TimedSearch { @@ -283,6 +284,54 @@ struct DeltaBenchResult { uintmax_t delta_bytes_after = 0; }; +struct DeleteBenchResult { + int deleted = 0; + int live = 0; + double full_search_ms = 0.0; + double tombstone_search_ms = 0.0; + double compact_ms = 0.0; + double compacted_search_ms = 0.0; +}; + +DeleteBenchResult run_delete_bench( + int bit_width, + const std::vector & vectors, + const std::vector & ids, + const std::vector & queries, + const BenchConfig & cfg) { + CHECK(cfg.delete_stride > 0); + + ggml_vec_index_t * idx = ggml_vec_index_create(cfg.dim, bit_width); + CHECK(idx != nullptr); + CHECK(ggml_vec_index_add(idx, vectors.data(), cfg.n_vec, ids.data()) == GGML_VEC_INDEX_OK); + + DeleteBenchResult result; + result.full_search_ms = run_search( + idx, queries, cfg.n_query, cfg.k, cfg.warmups, cfg.repeats).ms; + + for (int row = 0; row < cfg.n_vec; row += cfg.delete_stride) { + CHECK(ggml_vec_index_remove(idx, ids[static_cast(row)]) == 1); + ++result.deleted; + } + result.live = ggml_vec_index_len(idx); + CHECK(result.live == cfg.n_vec - result.deleted); + + result.tombstone_search_ms = run_search( + idx, queries, cfg.n_query, cfg.k, cfg.warmups, cfg.repeats).ms; + + const auto t0 = std::chrono::steady_clock::now(); + CHECK(ggml_vec_index_compact(idx) == GGML_VEC_INDEX_OK); + const auto t1 = std::chrono::steady_clock::now(); + result.compact_ms = std::chrono::duration(t1 - t0).count(); + CHECK(ggml_vec_index_len(idx) == result.live); + + result.compacted_search_ms = run_search( + idx, queries, cfg.n_query, cfg.k, cfg.warmups, cfg.repeats).ms; + + ggml_vec_index_free(idx); + return result; +} + DeltaBenchResult run_delta_bench( int bit_width, const std::vector & vectors, @@ -517,6 +566,21 @@ int main() { const uintmax_t f32_file_size = std::filesystem::file_size(f32_path); const uintmax_t q8_file_size = std::filesystem::file_size(q8_path); const uintmax_t q4_file_size = std::filesystem::file_size(q4_path); + const double f32_mmap_load_ms = median_time_ms(cfg.warmups, cfg.repeats, [&]() { + ggml_vec_index_t * loaded = ggml_vec_index_load_mmap(f32_path.string().c_str()); + CHECK(loaded != nullptr); + ggml_vec_index_free(loaded); + }); + const double q8_mmap_load_ms = median_time_ms(cfg.warmups, cfg.repeats, [&]() { + ggml_vec_index_t * loaded = ggml_vec_index_load_mmap(q8_path.string().c_str()); + CHECK(loaded != nullptr); + ggml_vec_index_free(loaded); + }); + const double q4_mmap_load_ms = median_time_ms(cfg.warmups, cfg.repeats, [&]() { + ggml_vec_index_t * loaded = ggml_vec_index_load_mmap(q4_path.string().c_str()); + CHECK(loaded != nullptr); + ggml_vec_index_free(loaded); + }); std::filesystem::remove(f32_path); std::filesystem::remove(q8_path); std::filesystem::remove(q4_path); @@ -542,6 +606,9 @@ int main() { cfg, "ggml-vector-index-bench-delta-q4.tvim", "ggml-vector-index-bench-delta-q4.tvid"); + const DeleteBenchResult f32_delete = run_delete_bench(32, vectors, ids, queries, cfg); + const DeleteBenchResult q8_delete = run_delete_bench(8, vectors, ids, queries, cfg); + const DeleteBenchResult q4_delete = run_delete_bench(4, vectors, ids, queries, cfg); const size_t f32_memory_bytes = static_cast(cfg.n_vec) * static_cast(cfg.dim) * sizeof(float) + @@ -570,6 +637,8 @@ int main() { static_cast(q4_file_size), static_cast(q8_file_size) / static_cast(f32_file_size), static_cast(q4_file_size) / static_cast(f32_file_size)); + std::printf(" mmap load: f32=%.3f ms q8=%.3f ms q4=%.3f ms\n", + f32_mmap_load_ms, q8_mmap_load_ms, q4_mmap_load_ms); std::printf(" median latency: f32=%.3f ms q8=%.3f ms q4=%.3f ms q8/f32=%.3f q4/f32=%.3f\n", f32_res.ms, q8_res.ms, q4_res.ms, q8_res.ms / f32_res.ms, q4_res.ms / f32_res.ms); std::printf( @@ -648,6 +717,35 @@ int main() { static_cast(q4_delta.delta_bytes_before), static_cast(q4_delta.snapshot_bytes_after), static_cast(q4_delta.delta_bytes_after)); + std::printf( + " delete-heavy: deleted=%d live=%d stride=%d\n", + f32_delete.deleted, + f32_delete.live, + cfg.delete_stride); + std::printf( + " delete f32: full_before=%.3f ms tombstoned=%.3f ms compact=%.3f ms compacted=%.3f ms tomb/full_before=%.3f compacted/tomb=%.3f\n", + f32_delete.full_search_ms, + f32_delete.tombstone_search_ms, + f32_delete.compact_ms, + f32_delete.compacted_search_ms, + f32_delete.tombstone_search_ms / f32_delete.full_search_ms, + f32_delete.compacted_search_ms / f32_delete.tombstone_search_ms); + std::printf( + " delete q8: full_before=%.3f ms tombstoned=%.3f ms compact=%.3f ms compacted=%.3f ms tomb/full_before=%.3f compacted/tomb=%.3f\n", + q8_delete.full_search_ms, + q8_delete.tombstone_search_ms, + q8_delete.compact_ms, + q8_delete.compacted_search_ms, + q8_delete.tombstone_search_ms / q8_delete.full_search_ms, + q8_delete.compacted_search_ms / q8_delete.tombstone_search_ms); + std::printf( + " delete q4: full_before=%.3f ms tombstoned=%.3f ms compact=%.3f ms compacted=%.3f ms tomb/full_before=%.3f compacted/tomb=%.3f\n", + q4_delete.full_search_ms, + q4_delete.tombstone_search_ms, + q4_delete.compact_ms, + q4_delete.compacted_search_ms, + q4_delete.tombstone_search_ms / q4_delete.full_search_ms, + q4_delete.compacted_search_ms / q4_delete.tombstone_search_ms); std::printf(" quality q8: recall@%d=%.4f mean_abs_score_drift=%.6f max_abs_score_drift=%.6f\n", cfg.k, q8_recall_at_k, q8_mean_abs_drift, q8_max_abs_drift); std::printf(" quality q4: recall@%d=%.4f mean_abs_score_drift=%.6f max_abs_score_drift=%.6f\n", diff --git a/tests/test-vector-index.cpp b/tests/test-vector-index.cpp index 329b6db10380..7e21751c15a9 100644 --- a/tests/test-vector-index.cpp +++ b/tests/test-vector-index.cpp @@ -553,6 +553,134 @@ int main() { ggml_vec_index_free(loaded); std::filesystem::remove(path); + // Tombstone removal: later adds append, searches skip deleted slots, IVF + // rebuilds exclude them, and snapshots write only live rows. + { + for (int bit_width : { 32, 8, 4 }) { + const std::string tombstone_path = + (std::filesystem::temp_directory_path() / + ("ggml-vector-index-tombstone-" + std::to_string(bit_width) + ".tvim")).string(); + std::filesystem::remove(tombstone_path); + + auto * tombstone_idx = ggml_vec_index_create(kDim, bit_width); + CHECK(tombstone_idx != nullptr); + CHECK(ggml_vec_index_add( + tombstone_idx, vecs.data(), static_cast(ids.size()), ids.data()) == + GGML_VEC_INDEX_OK); + CHECK(ggml_vec_index_remove(tombstone_idx, ids[1]) == 1); + CHECK(ggml_vec_index_len(tombstone_idx) == 3); + + const uint64_t appended_id = + (1ULL << 40) + 123ULL + static_cast(bit_width); + CHECK(ggml_vec_index_add( + tombstone_idx, seeds[1].data(), 1, &appended_id) == GGML_VEC_INDEX_OK); + CHECK(ggml_vec_index_len(tombstone_idx) == 4); + CHECK(ggml_vec_index_contains(tombstone_idx, ids[1]) == 0); + CHECK(ggml_vec_index_contains(tombstone_idx, appended_id) == 1); + + std::array scores{}; + std::array out_ids{}; + CHECK(ggml_vec_index_search( + tombstone_idx, seeds[1].data(), 1, /*k=*/4, + scores.data(), out_ids.data()) == GGML_VEC_INDEX_OK); + for (uint64_t result_id : out_ids) { + CHECK(result_id != ids[1]); + } + CHECK(out_ids[0] == appended_id); + + const std::array allowed = { ids[1], appended_id }; + scores = {}; + out_ids = {}; + CHECK(ggml_vec_index_search_filtered( + tombstone_idx, seeds[1].data(), 1, /*k=*/2, + allowed.data(), static_cast(allowed.size()), + scores.data(), out_ids.data()) == GGML_VEC_INDEX_OK); + CHECK(out_ids[0] == appended_id); + CHECK(out_ids[1] == UINT64_MAX); + + auto * filter = ggml_vec_index_filter_create( + tombstone_idx, allowed.data(), static_cast(allowed.size())); + CHECK(filter != nullptr); + scores = {}; + out_ids = {}; + CHECK(ggml_vec_index_search_prepared_filtered( + tombstone_idx, filter, seeds[1].data(), 1, /*k=*/2, + scores.data(), out_ids.data()) == GGML_VEC_INDEX_OK); + CHECK(out_ids[0] == appended_id); + CHECK(out_ids[1] == UINT64_MAX); + ggml_vec_index_filter_free(filter); + + CHECK(ggml_vec_index_build_ivf(tombstone_idx, /*n_lists=*/8, /*n_iter=*/2) == + GGML_VEC_INDEX_OK); + scores = {}; + out_ids = {}; + CHECK(ggml_vec_index_search_ivf( + tombstone_idx, seeds[1].data(), 1, /*k=*/4, /*nprobe=*/8, + scores.data(), out_ids.data()) == GGML_VEC_INDEX_OK); + CHECK(out_ids[0] == appended_id); + for (uint64_t result_id : out_ids) { + CHECK(result_id != ids[1]); + } + + auto * stale_filter = ggml_vec_index_filter_create( + tombstone_idx, allowed.data(), static_cast(allowed.size())); + CHECK(stale_filter != nullptr); + CHECK(ggml_vec_index_compact(tombstone_idx) == GGML_VEC_INDEX_OK); + CHECK(ggml_vec_index_len(tombstone_idx) == 4); + CHECK(ggml_vec_index_contains(tombstone_idx, ids[1]) == 0); + CHECK(ggml_vec_index_contains(tombstone_idx, appended_id) == 1); + CHECK(ggml_vec_index_search_prepared_filtered( + tombstone_idx, stale_filter, seeds[1].data(), 1, /*k=*/2, + scores.data(), out_ids.data()) == GGML_VEC_INDEX_E_INVALID_ARG); + ggml_vec_index_filter_free(stale_filter); + CHECK(ggml_vec_index_search_ivf( + tombstone_idx, seeds[1].data(), 1, /*k=*/1, /*nprobe=*/1, + scores.data(), out_ids.data()) == GGML_VEC_INDEX_E_INVALID_ARG); + + scores = {}; + out_ids = {}; + CHECK(ggml_vec_index_search_filtered( + tombstone_idx, seeds[1].data(), 1, /*k=*/2, + allowed.data(), static_cast(allowed.size()), + scores.data(), out_ids.data()) == GGML_VEC_INDEX_OK); + CHECK(out_ids[0] == appended_id); + CHECK(out_ids[1] == UINT64_MAX); + CHECK(ggml_vec_index_build_ivf(tombstone_idx, /*n_lists=*/8, /*n_iter=*/2) == + GGML_VEC_INDEX_OK); + scores = {}; + out_ids = {}; + CHECK(ggml_vec_index_search_ivf( + tombstone_idx, seeds[1].data(), 1, /*k=*/4, /*nprobe=*/8, + scores.data(), out_ids.data()) == GGML_VEC_INDEX_OK); + CHECK(out_ids[0] == appended_id); + for (uint64_t result_id : out_ids) { + CHECK(result_id != ids[1]); + } + + CHECK(ggml_vec_index_write(tombstone_idx, tombstone_path.c_str()) == + GGML_VEC_INDEX_OK); + ggml_vec_index_free(tombstone_idx); + + auto * tombstone_loaded = ggml_vec_index_load(tombstone_path.c_str()); + CHECK(tombstone_loaded != nullptr); + CHECK(ggml_vec_index_bit_width(tombstone_loaded) == bit_width); + CHECK(ggml_vec_index_len(tombstone_loaded) == 4); + CHECK(ggml_vec_index_contains(tombstone_loaded, ids[1]) == 0); + CHECK(ggml_vec_index_contains(tombstone_loaded, appended_id) == 1); + scores = {}; + out_ids = {}; + CHECK(ggml_vec_index_search( + tombstone_loaded, seeds[1].data(), 1, /*k=*/4, + scores.data(), out_ids.data()) == GGML_VEC_INDEX_OK); + CHECK(out_ids[0] == appended_id); + for (uint64_t result_id : out_ids) { + CHECK(result_id != ids[1]); + } + ggml_vec_index_free(tombstone_loaded); + std::filesystem::remove(tombstone_path); + } + } + // Incremental persistence: replay add/remove deltas on top of a snapshot. { const std::string snapshot_path = @@ -697,6 +825,60 @@ int main() { std::filesystem::remove(corrupt_delta_path); } + // Delta replay supports tombstone delete followed by re-adding the same ID. + { + for (int bit_width : { 32, 8, 4 }) { + const std::string snapshot_path = + (std::filesystem::temp_directory_path() / + ("ggml-vector-index-delta-tombstone-" + + std::to_string(bit_width) + ".tvim")).string(); + const std::string delta_path = + (std::filesystem::temp_directory_path() / + ("ggml-vector-index-delta-tombstone-" + + std::to_string(bit_width) + ".tvid")).string(); + std::filesystem::remove(snapshot_path); + std::filesystem::remove(delta_path); + + auto * delta_tombstone = ggml_vec_index_create(kDim, bit_width); + CHECK(delta_tombstone != nullptr); + std::vector base_vecs; + base_vecs.insert(base_vecs.end(), seeds[0].begin(), seeds[0].end()); + base_vecs.insert(base_vecs.end(), seeds[1].begin(), seeds[1].end()); + CHECK(ggml_vec_index_add( + delta_tombstone, base_vecs.data(), 2, ids.data()) == GGML_VEC_INDEX_OK); + CHECK(ggml_vec_index_write( + delta_tombstone, snapshot_path.c_str()) == GGML_VEC_INDEX_OK); + + CHECK(ggml_vec_index_remove_logged( + delta_tombstone, ids[0], delta_path.c_str()) == 1); + CHECK(ggml_vec_index_add_logged( + delta_tombstone, seeds[2].data(), 1, &ids[0], delta_path.c_str()) == + GGML_VEC_INDEX_OK); + CHECK(ggml_vec_index_len(delta_tombstone) == 2); + CHECK(ggml_vec_index_contains(delta_tombstone, ids[0]) == 1); + + auto * replayed_tombstone = ggml_vec_index_load_with_delta( + snapshot_path.c_str(), delta_path.c_str()); + CHECK(replayed_tombstone != nullptr); + CHECK(ggml_vec_index_bit_width(replayed_tombstone) == bit_width); + CHECK(ggml_vec_index_len(replayed_tombstone) == 2); + CHECK(ggml_vec_index_contains(replayed_tombstone, ids[0]) == 1); + CHECK(ggml_vec_index_contains(replayed_tombstone, ids[1]) == 1); + + std::array scores{}; + std::array out_ids{}; + CHECK(ggml_vec_index_search( + replayed_tombstone, seeds[2].data(), 1, /*k=*/2, + scores.data(), out_ids.data()) == GGML_VEC_INDEX_OK); + CHECK(out_ids[0] == ids[0]); + + ggml_vec_index_free(replayed_tombstone); + ggml_vec_index_free(delta_tombstone); + std::filesystem::remove(snapshot_path); + std::filesystem::remove(delta_path); + } + } + // v1 f32 snapshots migrate to q8 only for legacy bit_width=8; other // legacy widths, including bit_width=4, migrate to f32. { @@ -716,6 +898,7 @@ int main() { auto * v1 = ggml_vec_index_load(v1_path.c_str()); CHECK(v1 != nullptr); + CHECK(ggml_vec_index_load_mmap(v1_path.c_str()) == nullptr); CHECK(ggml_vec_index_dim(v1) == kDim); CHECK(ggml_vec_index_len(v1) == 2); CHECK(ggml_vec_index_bit_width(v1) == (bit_width == 8 ? 8 : 32)); @@ -733,6 +916,81 @@ int main() { } } + // mmap loading maps the vector section read-only and keeps search parity. + { + const std::vector mmap_ids = { + (1ULL << 38) + 1ULL, + (1ULL << 38) + 2ULL, + (1ULL << 38) + 3ULL, + (1ULL << 38) + 4ULL, + }; + std::array normal_scores{}; + std::array mmap_scores{}; + std::array normal_ids{}; + std::array mmap_out_ids{}; + const std::vector query = normalize({0.5f, -0.25f, 0.75f, 0.125f}); + + for (int bit_width : { 32, 8, 4 }) { + const auto mmap_tmp = std::filesystem::temp_directory_path() / + ("ggml-vector-index-mmap-" + std::to_string(bit_width) + ".tvim"); + const auto mmap_copy_tmp = std::filesystem::temp_directory_path() / + ("ggml-vector-index-mmap-copy-" + std::to_string(bit_width) + ".tvim"); + const std::string mmap_path = mmap_tmp.string(); + const std::string mmap_copy_path = mmap_copy_tmp.string(); + std::filesystem::remove(mmap_path); + std::filesystem::remove(mmap_copy_path); + + auto * source = ggml_vec_index_create(kDim, bit_width); + CHECK(source != nullptr); + CHECK(ggml_vec_index_add( + source, vecs.data(), static_cast(mmap_ids.size()), mmap_ids.data()) == + GGML_VEC_INDEX_OK); + CHECK(ggml_vec_index_write(source, mmap_path.c_str()) == GGML_VEC_INDEX_OK); + + auto * normal = ggml_vec_index_load(mmap_path.c_str()); + auto * mapped = ggml_vec_index_load_mmap(mmap_path.c_str()); + CHECK(normal != nullptr); + CHECK(mapped != nullptr); + CHECK(ggml_vec_index_bit_width(mapped) == bit_width); + CHECK(ggml_vec_index_len(mapped) == static_cast(mmap_ids.size())); + + CHECK(ggml_vec_index_search( + normal, query.data(), 1, /*k=*/4, + normal_scores.data(), normal_ids.data()) == GGML_VEC_INDEX_OK); + CHECK(ggml_vec_index_search( + mapped, query.data(), 1, /*k=*/4, + mmap_scores.data(), mmap_out_ids.data()) == GGML_VEC_INDEX_OK); + for (int i = 0; i < 4; ++i) { + CHECK(mmap_out_ids[i] == normal_ids[i]); + CHECK(std::fabs(mmap_scores[i] - normal_scores[i]) <= 1e-6f); + } + + CHECK(ggml_vec_index_build_ivf(mapped, /*n_lists=*/2, /*n_iter=*/2) + == GGML_VEC_INDEX_OK); + CHECK(ggml_vec_index_search_ivf( + mapped, query.data(), 1, /*k=*/2, /*nprobe=*/2, + mmap_scores.data(), mmap_out_ids.data()) == GGML_VEC_INDEX_OK); + + const uint64_t new_id = (1ULL << 38) + 99ULL; + CHECK(ggml_vec_index_add(mapped, seeds[0].data(), 1, &new_id) + == GGML_VEC_INDEX_E_INVALID_ARG); + CHECK(ggml_vec_index_remove(mapped, mmap_ids[0]) + == GGML_VEC_INDEX_E_INVALID_ARG); + CHECK(ggml_vec_index_compact(mapped) == GGML_VEC_INDEX_E_INVALID_ARG); + CHECK(ggml_vec_index_write(mapped, mmap_copy_path.c_str()) == GGML_VEC_INDEX_OK); + auto * copied = ggml_vec_index_load(mmap_copy_path.c_str()); + CHECK(copied != nullptr); + CHECK(ggml_vec_index_len(copied) == static_cast(mmap_ids.size())); + + ggml_vec_index_free(copied); + ggml_vec_index_free(mapped); + ggml_vec_index_free(normal); + ggml_vec_index_free(source); + std::filesystem::remove(mmap_path); + std::filesystem::remove(mmap_copy_path); + } + } + // q8 score parity for a dimension that exercises the SIMD tail. { constexpr int tail_dim = 13; From 5b49384cb2d6f32a1f19efc919e95a37e478b37d Mon Sep 17 00:00:00 2001 From: Nidhin Date: Wed, 15 Jul 2026 01:33:43 +0530 Subject: [PATCH 07/13] tests : expand vector index benchmark coverage Adds mmap load timing, delete-heavy compaction metrics, multi-distribution q4/q8 quality checks, and benchmark-only q4 calibration experiments. Assisted-by: GPT-5.5 --- tests/bench-vector-index.cpp | 410 +++++++++++++++++++++++++++++++---- 1 file changed, 362 insertions(+), 48 deletions(-) diff --git a/tests/bench-vector-index.cpp b/tests/bench-vector-index.cpp index d035a6b1e19e..b4d9ae98fffb 100644 --- a/tests/bench-vector-index.cpp +++ b/tests/bench-vector-index.cpp @@ -10,6 +10,7 @@ #include #include #include +#include #include #include #include @@ -45,6 +46,36 @@ struct TimedSearch { std::vector ids; }; +struct QualityMetrics { + double recall_at_k = 0.0; + double mean_abs_score_drift = 0.0; + double max_abs_score_drift = 0.0; +}; + +struct QualityBenchResult { + const char * name = ""; + QualityMetrics q8; + QualityMetrics q4; + std::vector q4_calibrated; +}; + +struct Q4CalibrationMode { + const char * name = ""; + float percentile = 1.0f; + float rms_factor = 0.0f; +}; + +struct Q4SimulatedIndex { + int dim = 0; + std::vector codes; + std::vector scales; +}; + +struct ScoreId { + float score = 0.0f; + uint64_t id = 0; +}; + template double median_time_ms(int warmups, int repeats, Fn fn) { for (int i = 0; i < warmups; ++i) { @@ -104,6 +135,100 @@ std::vector make_normalized_vectors(int n, int dim, uint32_t seed) { return vectors; } +std::vector make_gaussian_vectors(int n, int dim, uint32_t seed, bool normalize_rows) { + std::mt19937 rng(seed); + std::normal_distribution dist(0.0f, 1.0f); + + std::vector vectors(static_cast(n) * static_cast(dim)); + for (int row = 0; row < n; ++row) { + float norm2 = 0.0f; + float * v = vectors.data() + static_cast(row) * static_cast(dim); + for (int i = 0; i < dim; ++i) { + v[i] = dist(rng); + norm2 += v[i] * v[i]; + } + if (normalize_rows) { + const float inv_norm = norm2 > 0.0f ? 1.0f / std::sqrt(norm2) : 1.0f; + for (int i = 0; i < dim; ++i) { + v[i] *= inv_norm; + } + } + } + return vectors; +} + +std::vector make_sparse_vectors(int n, int dim, uint32_t seed, int nnz) { + CHECK(nnz > 0 && nnz <= dim); + std::mt19937 rng(seed); + std::normal_distribution dist(0.0f, 1.0f); + std::uniform_int_distribution coord_dist(0, dim - 1); + + std::vector vectors(static_cast(n) * static_cast(dim), 0.0f); + for (int row = 0; row < n; ++row) { + float * v = vectors.data() + static_cast(row) * static_cast(dim); + for (int j = 0; j < nnz; ++j) { + v[coord_dist(rng)] += dist(rng); + } + float norm2 = 0.0f; + for (int i = 0; i < dim; ++i) { + norm2 += v[i] * v[i]; + } + const float inv_norm = norm2 > 0.0f ? 1.0f / std::sqrt(norm2) : 1.0f; + for (int i = 0; i < dim; ++i) { + v[i] *= inv_norm; + } + } + return vectors; +} + +std::vector make_cluster_centers(int dim, uint32_t seed, int n_clusters) { + CHECK(n_clusters > 0); + std::mt19937 rng(seed); + std::normal_distribution center_dist(0.0f, 1.0f); + std::vector centers(static_cast(n_clusters) * static_cast(dim)); + for (int cluster = 0; cluster < n_clusters; ++cluster) { + float norm2 = 0.0f; + float * center = centers.data() + static_cast(cluster) * static_cast(dim); + for (int i = 0; i < dim; ++i) { + center[i] = center_dist(rng); + norm2 += center[i] * center[i]; + } + const float inv_norm = norm2 > 0.0f ? 1.0f / std::sqrt(norm2) : 1.0f; + for (int i = 0; i < dim; ++i) { + center[i] *= inv_norm; + } + } + return centers; +} + +std::vector make_clustered_vectors( + int n, + int dim, + uint32_t seed, + const std::vector & centers, + int n_clusters) { + CHECK(n_clusters > 0); + CHECK(centers.size() == static_cast(n_clusters) * static_cast(dim)); + std::mt19937 rng(seed); + std::normal_distribution noise_dist(0.0f, 0.05f); + std::vector vectors(static_cast(n) * static_cast(dim)); + for (int row = 0; row < n; ++row) { + const float * center = + centers.data() + static_cast(row % n_clusters) * static_cast(dim); + float * v = vectors.data() + static_cast(row) * static_cast(dim); + float norm2 = 0.0f; + for (int i = 0; i < dim; ++i) { + v[i] = center[i] + noise_dist(rng); + norm2 += v[i] * v[i]; + } + const float inv_norm = norm2 > 0.0f ? 1.0f / std::sqrt(norm2) : 1.0f; + for (int i = 0; i < dim; ++i) { + v[i] *= inv_norm; + } + } + return vectors; +} + TimedSearch run_search( const ggml_vec_index_t * idx, const std::vector & queries, @@ -217,6 +342,110 @@ float dot_exact( return acc; } +float q4_reference_abs_for_mode(const float * row, int dim, const Q4CalibrationMode & mode) { + float max_abs = 0.0f; + float sumsq = 0.0f; + std::vector abs_values; + abs_values.reserve(static_cast(dim)); + for (int i = 0; i < dim; ++i) { + const float a = std::fabs(row[i]); + max_abs = std::max(max_abs, a); + sumsq += row[i] * row[i]; + abs_values.push_back(a); + } + if (max_abs == 0.0f) { + return 0.0f; + } + + if (mode.rms_factor > 0.0f) { + const float rms = std::sqrt(sumsq / static_cast(dim)); + return std::min(max_abs, mode.rms_factor * rms); + } + + if (mode.percentile < 1.0f) { + std::sort(abs_values.begin(), abs_values.end()); + const size_t rank = static_cast( + std::floor(mode.percentile * static_cast(dim - 1))); + return std::max(abs_values[rank], std::numeric_limits::min()); + } + + return max_abs; +} + +Q4SimulatedIndex build_simulated_q4( + const std::vector & vectors, + int n, + int dim, + const Q4CalibrationMode & mode) { + Q4SimulatedIndex sim; + sim.dim = dim; + sim.codes.resize(static_cast(n) * static_cast(dim)); + sim.scales.resize(static_cast(n)); + for (int row = 0; row < n; ++row) { + const float * src = vectors.data() + static_cast(row) * static_cast(dim); + const float ref_abs = q4_reference_abs_for_mode(src, dim, mode); + float scale = ref_abs == 0.0f ? 1.0f : ref_abs / 7.0f; + if (scale == 0.0f) { + scale = ref_abs; + } + sim.scales[static_cast(row)] = scale; + int8_t * dst = sim.codes.data() + static_cast(row) * static_cast(dim); + for (int i = 0; i < dim; ++i) { + const float scaled = src[i] / scale; + int q = static_cast(std::nearbyint(scaled)); + q = std::max(-7, std::min(7, q)); + dst[i] = static_cast(q); + } + } + return sim; +} + +float dot_simulated_q4(const Q4SimulatedIndex & sim, const float * query, int row) { + const int8_t * codes = + sim.codes.data() + static_cast(row) * static_cast(sim.dim); + const float scale = sim.scales[static_cast(row)]; + float acc = 0.0f; + for (int i = 0; i < sim.dim; ++i) { + acc += query[i] * (static_cast(codes[i]) * scale); + } + return acc; +} + +TimedSearch run_simulated_q4_search( + const Q4SimulatedIndex & sim, + const std::vector & ids, + const std::vector & queries, + int n_query, + int k) { + TimedSearch result; + result.scores.resize(static_cast(n_query) * static_cast(k)); + result.ids.resize(static_cast(n_query) * static_cast(k)); + std::vector candidates(ids.size()); + for (int q = 0; q < n_query; ++q) { + const float * query = queries.data() + static_cast(q) * static_cast(sim.dim); + for (size_t row = 0; row < ids.size(); ++row) { + candidates[row] = { + dot_simulated_q4(sim, query, static_cast(row)), + ids[row], + }; + } + std::partial_sort( + candidates.begin(), + candidates.begin() + k, + candidates.end(), + [](const ScoreId & a, const ScoreId & b) { + return a.score > b.score; + }); + for (int i = 0; i < k; ++i) { + const size_t out = static_cast(q) * static_cast(k) + + static_cast(i); + result.scores[out] = candidates[static_cast(i)].score; + result.ids[out] = candidates[static_cast(i)].id; + } + } + return result; +} + std::filesystem::path write_index_file(ggml_vec_index_t * idx, const char * name) { const auto path = std::filesystem::temp_directory_path() / name; CHECK(ggml_vec_index_write(idx, path.string().c_str()) == GGML_VEC_INDEX_OK); @@ -273,6 +502,91 @@ double recall_against(const TimedSearch & exact, const TimedSearch & candidate, return static_cast(overlap) / static_cast(n_query * k); } +QualityMetrics quality_against( + const TimedSearch & exact, + const TimedSearch & candidate, + const std::vector & vectors, + const std::vector & queries, + int n_query, + int k, + int dim) { + QualityMetrics metrics; + int overlap = 0; + int drift_count = 0; + for (int q = 0; q < n_query; ++q) { + std::unordered_set exact_topk; + for (int j = 0; j < k; ++j) { + exact_topk.insert(exact.ids[static_cast(q) * static_cast(k) + j]); + } + + const float * query = queries.data() + static_cast(q) * static_cast(dim); + for (int j = 0; j < k; ++j) { + const size_t pos = static_cast(q) * static_cast(k) + + static_cast(j); + const uint64_t id = candidate.ids[pos]; + if (exact_topk.count(id) != 0) { + ++overlap; + } + const float exact_score = dot_exact(vectors, query, id, dim); + const double drift = std::fabs(static_cast(exact_score) - candidate.scores[pos]); + metrics.mean_abs_score_drift += drift; + metrics.max_abs_score_drift = std::max(metrics.max_abs_score_drift, drift); + ++drift_count; + } + } + metrics.recall_at_k = static_cast(overlap) / + static_cast(n_query * k); + metrics.mean_abs_score_drift /= static_cast(drift_count); + return metrics; +} + +QualityBenchResult run_quality_bench( + const char * name, + const std::vector & vectors, + const std::vector & queries, + const std::vector & q4_calibration_modes, + const BenchConfig & cfg) { + std::vector ids(static_cast(cfg.n_vec)); + for (int i = 0; i < cfg.n_vec; ++i) { + ids[static_cast(i)] = static_cast(i) + 1; + } + + ggml_vec_index_t * f32 = ggml_vec_index_create(cfg.dim, 32); + ggml_vec_index_t * q8 = ggml_vec_index_create(cfg.dim, 8); + ggml_vec_index_t * q4 = ggml_vec_index_create(cfg.dim, 4); + CHECK(f32 != nullptr); + CHECK(q8 != nullptr); + CHECK(q4 != nullptr); + CHECK(ggml_vec_index_add(f32, vectors.data(), cfg.n_vec, ids.data()) == GGML_VEC_INDEX_OK); + CHECK(ggml_vec_index_add(q8, vectors.data(), cfg.n_vec, ids.data()) == GGML_VEC_INDEX_OK); + CHECK(ggml_vec_index_add(q4, vectors.data(), cfg.n_vec, ids.data()) == GGML_VEC_INDEX_OK); + + const TimedSearch f32_res = run_search( + f32, queries, cfg.n_query, cfg.k, cfg.warmups, cfg.repeats); + const TimedSearch q8_res = run_search( + q8, queries, cfg.n_query, cfg.k, cfg.warmups, cfg.repeats); + const TimedSearch q4_res = run_search( + q4, queries, cfg.n_query, cfg.k, cfg.warmups, cfg.repeats); + + QualityBenchResult result; + result.name = name; + result.q8 = quality_against(f32_res, q8_res, vectors, queries, cfg.n_query, cfg.k, cfg.dim); + result.q4 = quality_against(f32_res, q4_res, vectors, queries, cfg.n_query, cfg.k, cfg.dim); + result.q4_calibrated.reserve(q4_calibration_modes.size()); + for (const Q4CalibrationMode & mode : q4_calibration_modes) { + const Q4SimulatedIndex sim = build_simulated_q4(vectors, cfg.n_vec, cfg.dim, mode); + const TimedSearch sim_res = run_simulated_q4_search( + sim, ids, queries, cfg.n_query, cfg.k); + result.q4_calibrated.push_back( + quality_against(f32_res, sim_res, vectors, queries, cfg.n_query, cfg.k, cfg.dim)); + } + + ggml_vec_index_free(f32); + ggml_vec_index_free(q8); + ggml_vec_index_free(q4); + return result; +} + struct DeltaBenchResult { double snapshot_load_ms = 0.0; double replay_load_ms = 0.0; @@ -510,49 +824,6 @@ int main() { ggml_vec_index_filter_free(q8_filter); } - double q8_mean_abs_drift = 0.0; - double q8_max_abs_drift = 0.0; - double q4_mean_abs_drift = 0.0; - double q4_max_abs_drift = 0.0; - int q8_overlap = 0; - int q4_overlap = 0; - int drift_count = 0; - - for (int q = 0; q < cfg.n_query; ++q) { - std::unordered_set f32_topk; - for (int j = 0; j < cfg.k; ++j) { - f32_topk.insert(f32_res.ids[static_cast(q) * cfg.k + j]); - } - - const float * query = queries.data() + static_cast(q) * static_cast(cfg.dim); - for (int j = 0; j < cfg.k; ++j) { - const size_t pos = static_cast(q) * static_cast(cfg.k) + static_cast(j); - const uint64_t id = q8_res.ids[pos]; - if (f32_topk.count(id) != 0) { - ++q8_overlap; - } - const float exact = dot_exact(vectors, query, id, cfg.dim); - const double drift = std::fabs(static_cast(exact) - q8_res.scores[pos]); - q8_mean_abs_drift += drift; - q8_max_abs_drift = std::max(q8_max_abs_drift, drift); - - const uint64_t q4_id = q4_res.ids[pos]; - if (f32_topk.count(q4_id) != 0) { - ++q4_overlap; - } - const float q4_exact = dot_exact(vectors, query, q4_id, cfg.dim); - const double q4_drift = std::fabs(static_cast(q4_exact) - q4_res.scores[pos]); - q4_mean_abs_drift += q4_drift; - q4_max_abs_drift = std::max(q4_max_abs_drift, q4_drift); - ++drift_count; - } - } - q8_mean_abs_drift /= static_cast(drift_count); - q4_mean_abs_drift /= static_cast(drift_count); - const double q8_recall_at_k = static_cast(q8_overlap) / - static_cast(cfg.n_query * cfg.k); - const double q4_recall_at_k = static_cast(q4_overlap) / - static_cast(cfg.n_query * cfg.k); const double f32_ivf_recall_at_k = recall_against(f32_res, f32_ivf_res, cfg.n_query, cfg.k); const double q8_ivf_recall_at_k = @@ -609,6 +880,33 @@ int main() { const DeleteBenchResult f32_delete = run_delete_bench(32, vectors, ids, queries, cfg); const DeleteBenchResult q8_delete = run_delete_bench(8, vectors, ids, queries, cfg); const DeleteBenchResult q4_delete = run_delete_bench(4, vectors, ids, queries, cfg); + const std::vector q4_calibration_modes = { + { "p99_abs", 0.99f, 0.0f }, + { "p95_abs", 0.95f, 0.0f }, + { "rms_3", 1.0f, 3.0f }, + }; + std::vector quality_suite; + quality_suite.push_back(run_quality_bench( + "normalized_gaussian", vectors, queries, q4_calibration_modes, cfg)); + quality_suite.push_back(run_quality_bench( + "raw_gaussian", + make_gaussian_vectors(cfg.n_vec, cfg.dim, 0x12345678, false), + make_gaussian_vectors(cfg.n_query, cfg.dim, 0x87654321, false), + q4_calibration_modes, + cfg)); + quality_suite.push_back(run_quality_bench( + "sparse_16", + make_sparse_vectors(cfg.n_vec, cfg.dim, 0x51a2b3c4, 16), + make_sparse_vectors(cfg.n_query, cfg.dim, 0x15a2b3c4, 16), + q4_calibration_modes, + cfg)); + const std::vector cluster_centers = make_cluster_centers(cfg.dim, 0x0ddc0ffe, 64); + quality_suite.push_back(run_quality_bench( + "clustered_64", + make_clustered_vectors(cfg.n_vec, cfg.dim, 0x0ddc0ffd, cluster_centers, 64), + make_clustered_vectors(cfg.n_query, cfg.dim, 0x0ddc0fff, cluster_centers, 64), + q4_calibration_modes, + cfg)); const size_t f32_memory_bytes = static_cast(cfg.n_vec) * static_cast(cfg.dim) * sizeof(float) + @@ -746,11 +1044,27 @@ int main() { q4_delete.compacted_search_ms, q4_delete.tombstone_search_ms / q4_delete.full_search_ms, q4_delete.compacted_search_ms / q4_delete.tombstone_search_ms); - std::printf(" quality q8: recall@%d=%.4f mean_abs_score_drift=%.6f max_abs_score_drift=%.6f\n", - cfg.k, q8_recall_at_k, q8_mean_abs_drift, q8_max_abs_drift); - std::printf(" quality q4: recall@%d=%.4f mean_abs_score_drift=%.6f max_abs_score_drift=%.6f\n", - cfg.k, q4_recall_at_k, q4_mean_abs_drift, q4_max_abs_drift); - + for (const QualityBenchResult & quality : quality_suite) { + std::printf( + " quality %-19s q8_recall=%.4f q8_mean_drift=%.6f q8_max_drift=%.6f q4_recall=%.4f q4_mean_drift=%.6f q4_max_drift=%.6f\n", + quality.name, + quality.q8.recall_at_k, + quality.q8.mean_abs_score_drift, + quality.q8.max_abs_score_drift, + quality.q4.recall_at_k, + quality.q4.mean_abs_score_drift, + quality.q4.max_abs_score_drift); + for (size_t i = 0; i < q4_calibration_modes.size(); ++i) { + const QualityMetrics & metrics = quality.q4_calibrated[i]; + std::printf( + " q4cal %-19s mode=%-7s recall=%.4f mean_drift=%.6f max_drift=%.6f\n", + quality.name, + q4_calibration_modes[i].name, + metrics.recall_at_k, + metrics.mean_abs_score_drift, + metrics.max_abs_score_drift); + } + } ggml_vec_index_free(f32); ggml_vec_index_free(q8); ggml_vec_index_free(q4); From be7af402964e106041d6b2d3d7333e3e703c8c5a Mon Sep 17 00:00:00 2001 From: Nidhin Date: Thu, 16 Jul 2026 15:16:51 +0530 Subject: [PATCH 08/13] ggml: preserve vector index state on failed logged add --- ggml/src/ggml-vector-index.cpp | 114 +++++++++++++++++++++-------- tests/test-vector-index-faults.cpp | 40 ++++++++++ 2 files changed, 123 insertions(+), 31 deletions(-) diff --git a/ggml/src/ggml-vector-index.cpp b/ggml/src/ggml-vector-index.cpp index cee1bd6a56fa..87eaa5141831 100644 --- a/ggml/src/ggml-vector-index.cpp +++ b/ggml/src/ggml-vector-index.cpp @@ -74,6 +74,7 @@ extern "C" { void ggml_vec_index_test_set_oom_countdown(int64_t countdown); void ggml_vec_index_test_set_write_fail_after(int64_t bytes); +void ggml_vec_index_test_set_truncate_fail(int fail); } #endif @@ -150,6 +151,7 @@ MappedFile::~MappedFile() { #ifdef GGML_VEC_INDEX_TEST_HOOKS std::atomic g_test_oom_countdown{ -1 }; std::atomic g_test_write_fail_after{ -1 }; +std::atomic g_test_truncate_fail{ false }; void test_maybe_throw_bad_alloc() { const int64_t remaining = g_test_oom_countdown.load(); @@ -634,6 +636,10 @@ void ggml_vec_index_test_set_oom_countdown(int64_t countdown) { void ggml_vec_index_test_set_write_fail_after(int64_t bytes) { g_test_write_fail_after.store(bytes); } + +void ggml_vec_index_test_set_truncate_fail(int fail) { + g_test_truncate_fail.store(fail != 0); +} } #endif @@ -1030,6 +1036,11 @@ bool validate_delta_header( bool truncate_file_to(const char * path, uint64_t size) { try { +#ifdef GGML_VEC_INDEX_TEST_HOOKS + if (g_test_truncate_fail.load()) { + return false; + } +#endif std::filesystem::path fs_path; if (!filesystem_path_from_utf8(path, fs_path)) { return false; @@ -1124,6 +1135,25 @@ bool inspect_delta_log_tail( return true; } +bool delta_log_ends_at_state( + const char * path, + const ggml_vec_index & idx, + uint32_t state_crc) { + try { + std::filesystem::path fs_path; + if (!filesystem_path_from_utf8(path, fs_path)) { + return false; + } + uint32_t tail_crc = 0; + uint64_t complete_size = 0; + return inspect_delta_log_tail(path, idx, tail_crc, complete_size) && + complete_size == static_cast(std::filesystem::file_size(fs_path)) && + tail_crc == state_crc; + } catch (...) { + return false; + } +} + int append_delta_record( const ggml_vec_index & idx, const char * delta_path, @@ -1359,11 +1389,38 @@ void ggml_vec_index_free(ggml_vec_index_t * idx) { // Mutation // --------------------------------------------------------------------------- +static void rollback_appended_slots_unlocked( + ggml_vec_index_t * idx, + size_t base_slot, + const uint64_t * ids, + int n) noexcept { + if (idx == nullptr) { + return; + } + for (int i = 0; i < n; ++i) { + idx->id_to_slot.erase(ids[i]); + } + const size_t dim_sz = static_cast(idx->dim); + if (is_q4(*idx)) { + idx->q4_data.resize(base_slot * q4_row_bytes(dim_sz)); + idx->q4_scale.resize(base_slot); + } else if (is_q8(*idx)) { + idx->q8_data.resize(base_slot * dim_sz); + idx->q8_scale.resize(base_slot); + } else { + idx->data.resize(base_slot * dim_sz); + } + idx->slot_to_id.resize(base_slot); + idx->slot_active.resize(base_slot); + idx->n_active = idx->id_to_slot.size(); +} + static int ggml_vec_index_add_unlocked( ggml_vec_index_t * idx, const float * vectors, int n, - const uint64_t * ids) { + const uint64_t * ids, + bool finalize) { size_t base_slot = 0; size_t dim_sz = 0; @@ -1373,21 +1430,7 @@ static int ggml_vec_index_add_unlocked( if (idx == nullptr || !resized) { return; } - for (int i = 0; i < n; ++i) { - idx->id_to_slot.erase(ids[i]); - } - if (is_q4(*idx)) { - idx->q4_data.resize(base_slot * q4_row_bytes(dim_sz)); - idx->q4_scale.resize(base_slot); - } else if (is_q8(*idx)) { - idx->q8_data.resize(base_slot * dim_sz); - idx->q8_scale.resize(base_slot); - } else { - idx->data.resize(base_slot * dim_sz); - } - idx->slot_to_id.resize(base_slot); - idx->slot_active.resize(base_slot); - idx->n_active = idx->id_to_slot.size(); + rollback_appended_slots_unlocked(idx, base_slot, ids, n); }; try { @@ -1474,8 +1517,10 @@ static int ggml_vec_index_add_unlocked( idx->id_to_slot.emplace(ids[i], slot); } idx->n_active += n_sz; - ++idx->generation; - invalidate_ivf(*idx); + if (finalize) { + ++idx->generation; + invalidate_ivf(*idx); + } } catch (const std::bad_alloc &) { rollback(); return GGML_VEC_INDEX_E_OOM; @@ -1496,7 +1541,7 @@ int ggml_vec_index_add( } try { std::unique_lock lock(idx->mutex); - return ggml_vec_index_add_unlocked(idx, vectors, n, ids); + return ggml_vec_index_add_unlocked(idx, vectors, n, ids, true); } catch (...) { return GGML_VEC_INDEX_E_INTERNAL; } @@ -1662,6 +1707,7 @@ int ggml_vec_index_add_logged( const uint64_t * ids, const char * delta_path) { bool added = false; + size_t base_slot = 0; std::unique_lock lock; try { if (idx == nullptr || delta_path == nullptr) { @@ -1689,41 +1735,47 @@ int ggml_vec_index_add_logged( } const uint32_t base_crc = index_state_crc32c(*idx); - const int add_status = ggml_vec_index_add_unlocked(idx, vectors, n, ids); + base_slot = idx->slot_to_id.size(); + const int add_status = ggml_vec_index_add_unlocked(idx, vectors, n, ids, false); if (add_status != GGML_VEC_INDEX_OK) { return add_status; } added = true; + const uint32_t added_state_crc = index_state_crc32c(*idx); const int append_status = append_delta_record( *idx, delta_path, kTvidOpAdd, static_cast(n), base_crc, - index_state_crc32c(*idx), + added_state_crc, payload); if (append_status != GGML_VEC_INDEX_OK) { - if (append_status == GGML_VEC_INDEX_E_IO) { - for (int i = 0; i < n; ++i) { - ggml_vec_index_remove_unlocked(idx, ids[i]); - } + const bool record_is_complete = + append_status == GGML_VEC_INDEX_E_INTERNAL && + delta_log_ends_at_state(delta_path, *idx, added_state_crc); + if (record_is_complete) { + ++idx->generation; + invalidate_ivf(*idx); + } else { + rollback_appended_slots_unlocked(idx, base_slot, ids, n); } + added = false; return append_status; } + ++idx->generation; + invalidate_ivf(*idx); + added = false; return GGML_VEC_INDEX_OK; } catch (const std::bad_alloc &) { if (added) { - for (int i = 0; i < n; ++i) { - ggml_vec_index_remove_unlocked(idx, ids[i]); - } + rollback_appended_slots_unlocked(idx, base_slot, ids, n); } return GGML_VEC_INDEX_E_OOM; } catch (...) { if (added) { - for (int i = 0; i < n; ++i) { - ggml_vec_index_remove_unlocked(idx, ids[i]); - } + rollback_appended_slots_unlocked(idx, base_slot, ids, n); } return GGML_VEC_INDEX_E_INTERNAL; } diff --git a/tests/test-vector-index-faults.cpp b/tests/test-vector-index-faults.cpp index 7bd5594fee4a..806e596292d6 100644 --- a/tests/test-vector-index-faults.cpp +++ b/tests/test-vector-index-faults.cpp @@ -11,6 +11,7 @@ extern "C" void ggml_vec_index_test_set_oom_countdown(int64_t countdown); extern "C" void ggml_vec_index_test_set_write_fail_after(int64_t bytes); +extern "C" void ggml_vec_index_test_set_truncate_fail(int fail); namespace { @@ -25,6 +26,7 @@ namespace { void reset_fault_hooks() { ggml_vec_index_test_set_oom_countdown(-1); ggml_vec_index_test_set_write_fail_after(-1); + ggml_vec_index_test_set_truncate_fail(0); } std::vector read_file_bytes(const std::string & path) { @@ -194,12 +196,50 @@ int main() { CHECK(std::filesystem::file_size(delta_path) == 0); } + const uint64_t allowed_id = base_ids[0]; + std::array logged_scores{}; + std::array logged_out_ids{}; + auto * stale_filter = ggml_vec_index_filter_create(idx, &allowed_id, 1); + CHECK(stale_filter != nullptr); + CHECK(ggml_vec_index_build_ivf(idx, /*n_lists=*/2, /*n_iter=*/1) + == GGML_VEC_INDEX_OK); CHECK(ggml_vec_index_add_logged( idx, logged_vector.data(), 1, &logged_id, delta_path.c_str()) == GGML_VEC_INDEX_OK); CHECK(ggml_vec_index_contains(idx, logged_id) == 1); + CHECK(ggml_vec_index_search_prepared_filtered( + idx, stale_filter, base_vectors.data(), 1, /*k=*/1, + logged_scores.data(), logged_out_ids.data()) == GGML_VEC_INDEX_E_INVALID_ARG); + CHECK(ggml_vec_index_search_ivf( + idx, base_vectors.data(), 1, /*k=*/1, /*nprobe=*/2, + logged_scores.data(), logged_out_ids.data()) == GGML_VEC_INDEX_E_INVALID_ARG); + ggml_vec_index_filter_free(stale_filter); const std::vector old_delta = read_file_bytes(delta_path); + auto * rollback_filter = ggml_vec_index_filter_create(idx, &allowed_id, 1); + CHECK(rollback_filter != nullptr); + CHECK(ggml_vec_index_build_ivf(idx, /*n_lists=*/2, /*n_iter=*/1) + == GGML_VEC_INDEX_OK); + const uint64_t failed_internal_id = 402; + ggml_vec_index_test_set_write_fail_after(8); + ggml_vec_index_test_set_truncate_fail(1); + CHECK(ggml_vec_index_add_logged( + idx, logged_vector.data(), 1, &failed_internal_id, delta_path.c_str()) == + GGML_VEC_INDEX_E_INTERNAL); + reset_fault_hooks(); + CHECK(ggml_vec_index_contains(idx, failed_internal_id) == 0); + CHECK(ggml_vec_index_len(idx) == 4); + CHECK(read_file_bytes(delta_path) == old_delta); + CHECK(ggml_vec_index_search_prepared_filtered( + idx, rollback_filter, base_vectors.data(), 1, /*k=*/1, + logged_scores.data(), logged_out_ids.data()) == GGML_VEC_INDEX_OK); + CHECK(logged_out_ids[0] == allowed_id); + CHECK(ggml_vec_index_search_ivf( + idx, base_vectors.data(), 1, /*k=*/1, /*nprobe=*/2, + logged_scores.data(), logged_out_ids.data()) == GGML_VEC_INDEX_OK); + CHECK(logged_out_ids[0] == allowed_id); + ggml_vec_index_filter_free(rollback_filter); + ggml_vec_index_test_set_write_fail_after(8); CHECK(ggml_vec_index_remove_logged(idx, logged_id, delta_path.c_str()) == GGML_VEC_INDEX_E_IO); From 38b450b9f9d9c0e5afd3d928dbf3d64dd6cce09d Mon Sep 17 00:00:00 2001 From: Nidhin Date: Thu, 16 Jul 2026 16:33:53 +0530 Subject: [PATCH 09/13] ggml-vector-index: harden persistence and split library target Move the vector index implementation out of ggml-base into its own exported target, and tighten API invariants around prepared filters, reserved sentinel IDs, and delta-log compaction recovery. --- ggml/CMakeLists.txt | 4 +- ggml/include/ggml-vector-index.h | 10 ++--- ggml/src/CMakeLists.txt | 19 ++++++-- ggml/src/ggml-vector-index.cpp | 61 +++++++++++++++++++++++--- tests/CMakeLists.txt | 4 +- tests/test-vector-index.cpp | 74 ++++++++++++++++++++++++++++++-- 6 files changed, 148 insertions(+), 24 deletions(-) diff --git a/ggml/CMakeLists.txt b/ggml/CMakeLists.txt index c1390966712a..c1e09b8282d4 100644 --- a/ggml/CMakeLists.txt +++ b/ggml/CMakeLists.txt @@ -336,7 +336,6 @@ set(GGML_PUBLIC_HEADERS include/ggml-rpc.h include/ggml-virtgpu.h include/ggml-sycl.h - include/ggml-vector-index.h include/ggml-vulkan.h include/ggml-webgpu.h include/ggml-zendnn.h @@ -344,11 +343,12 @@ set(GGML_PUBLIC_HEADERS include/gguf.h) set_target_properties(ggml PROPERTIES PUBLIC_HEADER "${GGML_PUBLIC_HEADERS}") +set_target_properties(ggml-vector-index PROPERTIES PUBLIC_HEADER include/ggml-vector-index.h) #if (GGML_METAL) # set_target_properties(ggml PROPERTIES RESOURCE "${CMAKE_CURRENT_SOURCE_DIR}/src/ggml-metal.metal") #endif() install( - TARGETS ggml ggml-base + TARGETS ggml ggml-base ggml-vector-index EXPORT ggml-targets) install( diff --git a/ggml/include/ggml-vector-index.h b/ggml/include/ggml-vector-index.h index 2ab010869851..d5bcdea0e671 100644 --- a/ggml/include/ggml-vector-index.h +++ b/ggml/include/ggml-vector-index.h @@ -1,10 +1,9 @@ #pragma once // -// ggml-vector-index: TurboVec-style vector-index C API. +// ggml-vector-index: vector-index C API. // -// This is the public C API for fabric's vector index. The implementation -// under `ggml/src/ggml-vector-index.cpp` supports full f32 storage -// (`bit_width=32`), production q8 storage (`bit_width=8`), and packed q4 +// This public C API supports full f32 storage (`bit_width=32`), q8 storage +// (`bit_width=8`), and packed q4 // storage (`bit_width=4`) with CPU search directly against quantized codes. // q8 uses NEON/AVX2 when available; q4 uses NEON when available. // @@ -63,7 +62,8 @@ GGML_API void ggml_vec_index_free(ggml_vec_index_t * idx); // associating each with the corresponding `ids[i]` (caller-owned external id). // Returns 0 on success. Returns GGML_VEC_INDEX_E_DUPLICATE if any id already // exists in the index; in that case the index is unchanged (atomic add). -// All vector components must be finite. +// All vector components must be finite. UINT64_MAX is reserved for search +// result padding and is not a valid id. GGML_API int ggml_vec_index_add( ggml_vec_index_t * idx, const float * vectors, diff --git a/ggml/src/CMakeLists.txt b/ggml/src/CMakeLists.txt index 1ebaa351c9f1..52739415f5c3 100644 --- a/ggml/src/CMakeLists.txt +++ b/ggml/src/CMakeLists.txt @@ -191,7 +191,6 @@ add_library(ggml-base ../include/ggml-backend.h ../include/ggml-cpp.h ../include/ggml-opt.h - ../include/ggml-vector-index.h ../include/gguf.h ../include/uint8-buff-stream.h ggml.c @@ -205,7 +204,6 @@ add_library(ggml-base ggml-tbq-quants.c ggml-quants.c ggml-quants.h - ggml-vector-index.cpp gguf.cpp uint8-buff-stream.cpp) @@ -263,6 +261,19 @@ if (CMAKE_SYSTEM_NAME MATCHES "Linux") target_link_libraries(ggml PRIVATE dl) endif() +add_library(ggml-vector-index + ../include/ggml-vector-index.h + ggml-vector-index.cpp) +add_library(ggml::ggml-vector-index ALIAS ggml-vector-index) +add_library(ggml::vector-index ALIAS ggml-vector-index) + +set_target_properties(ggml-vector-index PROPERTIES + VERSION ${GGML_VERSION} + SOVERSION ${GGML_VERSION_MAJOR} +) + +target_link_libraries(ggml-vector-index PUBLIC ggml-base) + function(ggml_add_backend_library backend) if (GGML_BACKEND_DL) add_library(${backend} MODULE ${ARGN}) @@ -494,7 +505,7 @@ ggml_add_backend(Hexagon) ggml_add_backend(ZenDNN) ggml_add_backend(OPENVINO) -foreach (target ggml-base ggml) +foreach (target ggml-base ggml ggml-vector-index) target_include_directories(${target} PUBLIC $ $) target_compile_features (${target} PRIVATE c_std_11 cxx_std_17) # don't bump endforeach() @@ -516,7 +527,7 @@ if(CMAKE_SYSTEM_NAME MATCHES "visionOS") endif() if (BUILD_SHARED_LIBS OR GGML_BACKEND_DL) - foreach (target ggml-base ggml) + foreach (target ggml-base ggml ggml-vector-index) set_target_properties(${target} PROPERTIES POSITION_INDEPENDENT_CODE ON) target_compile_definitions(${target} PRIVATE GGML_BUILD) target_compile_definitions(${target} PUBLIC GGML_SHARED) diff --git a/ggml/src/ggml-vector-index.cpp b/ggml/src/ggml-vector-index.cpp index 87eaa5141831..930f78b2d6db 100644 --- a/ggml/src/ggml-vector-index.cpp +++ b/ggml/src/ggml-vector-index.cpp @@ -1,4 +1,4 @@ -// ggml-vector-index.cpp - CPU implementation of the fabric vector +// ggml-vector-index.cpp - CPU implementation of the vector // index C API declared in `ggml/include/ggml-vector-index.h`. // // Storage: full f32 vectors or per-vector symmetric q8 codes. ID map uses @@ -330,6 +330,10 @@ bool is_supported_bit_width(int bit_width) { return bit_width == 4 || bit_width == 8 || bit_width == 32; } +bool is_valid_id(uint64_t id) { + return id != UINT64_MAX; +} + bool all_finite(const float * values, size_t n) { for (size_t i = 0; i < n; ++i) { if (!std::isfinite(values[i])) { @@ -652,6 +656,8 @@ struct ggml_vec_index { int bit_width = 32; uint64_t generation = 0; bool read_only_mmap = false; + bool delta_log_rebase_pending = false; + uint32_t delta_log_rebase_crc = 0; std::unique_ptr mapped_file; size_t mapped_vector_bytes = 0; @@ -686,6 +692,7 @@ struct ggml_vec_index { }; struct ggml_vec_index_filter { + const ggml_vec_index_t * owner = nullptr; int dim = 0; int bit_width = 32; uint64_t generation = 0; @@ -1155,7 +1162,7 @@ bool delta_log_ends_at_state( } int append_delta_record( - const ggml_vec_index & idx, + ggml_vec_index & idx, const char * delta_path, uint8_t op, uint32_t n, @@ -1185,8 +1192,19 @@ int append_delta_record( } old_size = complete_size; } + if (idx.delta_log_rebase_pending && + idx.delta_log_rebase_crc == base_crc_for_new_log && + existing_base_crc != base_crc_for_new_log) { + // The snapshot already includes this log's records (e.g. crash after + // compacting the snapshot but before replacing the old delta log). + if (!truncate_file_to(delta_path, 0)) { + return GGML_VEC_INDEX_E_INTERNAL; + } + old_size = 0; + idx.delta_log_rebase_pending = false; + idx.delta_log_rebase_crc = 0; + } } - (void) existing_base_crc; std::FILE * f = nullptr; if (!open_append_file(delta_path, &f)) { @@ -1311,6 +1329,9 @@ int check_logged_add_duplicates( std::unordered_set batch_ids; batch_ids.reserve(static_cast(n)); for (int i = 0; i < n; ++i) { + if (!is_valid_id(ids[i])) { + return GGML_VEC_INDEX_E_INVALID_ARG; + } if (idx->id_to_slot.find(ids[i]) != idx->id_to_slot.end()) { return GGML_VEC_INDEX_E_DUPLICATE; } @@ -1453,6 +1474,9 @@ static int ggml_vec_index_add_unlocked( std::unordered_set batch_ids; batch_ids.reserve(static_cast(n)); for (int i = 0; i < n; ++i) { + if (!is_valid_id(ids[i])) { + return GGML_VEC_INDEX_E_INVALID_ARG; + } if (idx->id_to_slot.find(ids[i]) != idx->id_to_slot.end()) { return GGML_VEC_INDEX_E_DUPLICATE; } @@ -1552,6 +1576,9 @@ static int ggml_vec_index_remove_unlocked(ggml_vec_index_t * idx, uint64_t id) { if (idx == nullptr) { return GGML_VEC_INDEX_E_INVALID_ARG; } + if (!is_valid_id(id)) { + return GGML_VEC_INDEX_E_INVALID_ARG; + } if (idx->read_only_mmap) { return GGML_VEC_INDEX_E_INVALID_ARG; } @@ -1790,6 +1817,9 @@ int ggml_vec_index_remove_logged( return GGML_VEC_INDEX_E_INVALID_ARG; } std::unique_lock lock(idx->mutex); + if (!is_valid_id(id)) { + return GGML_VEC_INDEX_E_INVALID_ARG; + } if (idx->read_only_mmap) { return GGML_VEC_INDEX_E_INVALID_ARG; } @@ -2306,7 +2336,8 @@ static int ggml_vec_index_search_impl( std::vector allowed_slots; const std::vector * allowed_ptr = nullptr; if (prepared_filter != nullptr) { - if (prepared_filter->dim != idx->dim || + if (prepared_filter->owner != idx || + prepared_filter->dim != idx->dim || prepared_filter->bit_width != idx->bit_width || prepared_filter->generation != idx->generation) { return GGML_VEC_INDEX_E_INVALID_ARG; @@ -2373,6 +2404,7 @@ ggml_vec_index_filter_t * ggml_vec_index_filter_create( return nullptr; } std::unique_ptr owned(filter); + owned->owner = idx; owned->dim = idx->dim; owned->bit_width = idx->bit_width; owned->generation = idx->generation; @@ -2938,6 +2970,9 @@ ggml_vec_index_t * ggml_vec_index_load(const char * path) { if (!read_ok) { return nullptr; } + if (!is_valid_id(id)) { + return nullptr; + } } if (checksummed) { @@ -3160,6 +3195,9 @@ ggml_vec_index_t * ggml_vec_index_load_mmap(const char * path) { const uint8_t * ids = bytes + static_cast(ids_offset); for (size_t slot = 0; slot < n; ++slot) { const uint64_t id = get_u64_le(ids + slot * sizeof(uint64_t)); + if (!is_valid_id(id)) { + return nullptr; + } idx->slot_to_id[slot] = id; if (!idx->id_to_slot.emplace(id, slot).second) { return nullptr; @@ -3251,6 +3289,9 @@ bool replay_remove_delta(ggml_vec_index_t * idx, uint32_t n, const std::vector= 0; } @@ -3342,9 +3383,15 @@ bool replay_delta_log(ggml_vec_index_t * idx, const char * delta_path) { last_state_crc = state_crc; } - return apply_records ? - index_state_crc32c(*idx) == last_state_crc : - snapshot_crc == last_state_crc; + if (apply_records) { + return index_state_crc32c(*idx) == last_state_crc; + } + if (snapshot_crc != last_state_crc) { + return false; + } + idx->delta_log_rebase_pending = base_crc != snapshot_crc; + idx->delta_log_rebase_crc = snapshot_crc; + return true; } } // namespace diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 0c74f0a5ac2a..e9d16c3dedf6 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -147,7 +147,7 @@ endif() # Vector-index regression test: no model required, no llama runtime touched. add_executable(test-vector-index test-vector-index.cpp) target_compile_features(test-vector-index PRIVATE cxx_std_17) -target_link_libraries(test-vector-index PRIVATE ggml) +target_link_libraries(test-vector-index PRIVATE ggml-vector-index) if (LLAMA_TESTS_INSTALL) install(TARGETS test-vector-index RUNTIME) endif() @@ -168,7 +168,7 @@ llama_test(test-vector-index-faults) add_executable(bench-vector-index bench-vector-index.cpp) target_compile_features(bench-vector-index PRIVATE cxx_std_17) -target_link_libraries(bench-vector-index PRIVATE ggml) +target_link_libraries(bench-vector-index PRIVATE ggml-vector-index) if (LLAMA_TESTS_INSTALL) install(TARGETS bench-vector-index RUNTIME) endif() diff --git a/tests/test-vector-index.cpp b/tests/test-vector-index.cpp index 7e21751c15a9..cddee225c680 100644 --- a/tests/test-vector-index.cpp +++ b/tests/test-vector-index.cpp @@ -198,9 +198,21 @@ int main() { CHECK(ggml_vec_index_len(idx) == 0); } + // UINT64_MAX is reserved as the empty-result sentinel. + { + const std::array vector = { + 1.0f, 0.0f, 0.0f, 0.0f, + }; + const uint64_t reserved_id = UINT64_MAX; + CHECK(ggml_vec_index_add(idx, vector.data(), 1, &reserved_id) + == GGML_VEC_INDEX_E_INVALID_ARG); + CHECK(ggml_vec_index_remove(idx, reserved_id) + == GGML_VEC_INDEX_E_INVALID_ARG); + CHECK(ggml_vec_index_len(idx) == 0); + } + // Add 4 well-separated unit vectors. IDs are non-trivial uint64 to - // catch sign-extension or BigInt round-trip bugs at the JS boundary - // when this codepath is later exercised from Bare. + // catch sign-extension bugs when this codepath is called from bindings. std::vector vecs; std::vector ids = { 42ULL, @@ -464,6 +476,17 @@ int main() { CHECK(out_ids[1] == ids[2]); CHECK(out_ids[2] == UINT64_MAX); CHECK(scores[2] == -FLT_MAX); + + auto * other_idx = ggml_vec_index_create(kDim, /*bit_width=*/32); + CHECK(other_idx != nullptr); + CHECK(ggml_vec_index_add( + other_idx, vecs.data(), static_cast(ids.size()), ids.data()) == + GGML_VEC_INDEX_OK); + CHECK(ggml_vec_index_search_prepared_filtered( + other_idx, filter, seeds[0].data(), 1, /*k=*/1, + scores.data(), out_ids.data()) == GGML_VEC_INDEX_E_INVALID_ARG); + ggml_vec_index_free(other_idx); + ggml_vec_index_filter_free(filter); auto * empty_filter = ggml_vec_index_filter_create( @@ -527,6 +550,17 @@ int main() { CHECK((persisted_stat.st_mode & 0777) == 0600); #endif + const std::string reserved_id_path = + (std::filesystem::temp_directory_path() / + "ggml-vector-index-reserved-id-corrupt.tvim").string(); + expect_corrupt_load_fails(path, reserved_id_path, [](std::vector & bytes) { + constexpr size_t ids_offset = 32 + 3 * kDim * sizeof(float); + CHECK(bytes.size() >= ids_offset + sizeof(uint64_t)); + for (size_t i = 0; i < sizeof(uint64_t); ++i) { + bytes[ids_offset + i] = 0xff; + } + }); + ggml_vec_index_free(idx); auto * loaded = ggml_vec_index_load(path.c_str()); @@ -718,6 +752,16 @@ int main() { CHECK(ggml_vec_index_len(base_only) == 2); ggml_vec_index_free(base_only); + const uint64_t reserved_delta_id = UINT64_MAX; + CHECK(ggml_vec_index_add_logged( + base, seeds[2].data(), 1, &reserved_delta_id, delta_path.c_str()) == + GGML_VEC_INDEX_E_INVALID_ARG); + CHECK(ggml_vec_index_remove_logged( + base, reserved_delta_id, delta_path.c_str()) == + GGML_VEC_INDEX_E_INVALID_ARG); + CHECK(ggml_vec_index_len(base) == 2); + CHECK(!std::filesystem::exists(delta_path)); + const uint64_t delta_id = (1ULL << 41) + 7ULL; CHECK(ggml_vec_index_add_logged( base, seeds[2].data(), 1, &delta_id, delta_path.c_str()) == GGML_VEC_INDEX_OK); @@ -794,6 +838,20 @@ int main() { CHECK(ggml_vec_index_contains(compacted_with_old_log, ids[0]) == 0); CHECK(ggml_vec_index_contains(compacted_with_old_log, ids[1]) == 1); CHECK(ggml_vec_index_contains(compacted_with_old_log, delta_id) == 1); + + const uint64_t post_crash_compact_id = (1ULL << 41) + 10ULL; + CHECK(ggml_vec_index_add_logged( + compacted_with_old_log, seeds[3].data(), 1, &post_crash_compact_id, delta_path.c_str()) == + GGML_VEC_INDEX_OK); + auto * replayed_after_old_log_append = ggml_vec_index_load_with_delta( + snapshot_path.c_str(), delta_path.c_str()); + CHECK(replayed_after_old_log_append != nullptr); + CHECK(ggml_vec_index_len(replayed_after_old_log_append) == 3); + CHECK(ggml_vec_index_contains(replayed_after_old_log_append, ids[0]) == 0); + CHECK(ggml_vec_index_contains(replayed_after_old_log_append, ids[1]) == 1); + CHECK(ggml_vec_index_contains(replayed_after_old_log_append, delta_id) == 1); + CHECK(ggml_vec_index_contains(replayed_after_old_log_append, post_crash_compact_id) == 1); + ggml_vec_index_free(replayed_after_old_log_append); ggml_vec_index_free(compacted_with_old_log); CHECK(ggml_vec_index_compact_delta( @@ -1061,7 +1119,7 @@ int main() { ggml_vec_index_free(overflow_idx); } - // q4 production path: packed nibbles with one f32 scale per vector. + // q4 path: packed nibbles with one f32 scale per vector. { constexpr int tail_dim = 13; const std::vector tail_vector = { @@ -1234,7 +1292,7 @@ int main() { std::filesystem::remove(empty_path); } - // q8 production path: stores quantized codes, searches directly against + // q8 path: stores quantized codes, searches directly against // q8 storage, and persists as .tvim v2 with q8 metadata. { auto * q8 = ggml_vec_index_create(kDim, /*bit_width=*/8); @@ -1347,6 +1405,14 @@ int main() { bytes[id_offset + sizeof(uint64_t) + i] = bytes[id_offset + i]; } }); + expect_corrupt_load_fails( + legacy_v2_path, corrupt_path, [](std::vector & bytes) { + const size_t id_offset = + 32 + 2 * sizeof(float) + 2 * kDim * sizeof(int8_t); + for (size_t i = 0; i < sizeof(uint64_t); ++i) { + bytes[id_offset + i] = 0xff; + } + }); std::filesystem::remove(legacy_v2_path); expect_corrupt_load_fails(q8_path, corrupt_path, [](std::vector & bytes) { From b85119c47f6ccde431674f851806b0e6545f7f2f Mon Sep 17 00:00:00 2001 From: Nidhin Date: Thu, 16 Jul 2026 17:55:33 +0530 Subject: [PATCH 10/13] ggml-vector-index: harden logged persistence edge cases Fix delta-log consistency issues around committed mutation failures, shared log serialization, deterministic quantized replay, and unsafe compaction paths. Also tighten prepared filter validation and CMake package/linkage behavior. --- ggml/cmake/ggml-config.cmake.in | 7 + ggml/include/ggml-vector-index.h | 8 +- ggml/src/CMakeLists.txt | 2 +- ggml/src/ggml-vector-index.cpp | 234 ++++++++++++++++++++++++++++- tests/CMakeLists.txt | 5 +- tests/bench-vector-index.cpp | 18 ++- tests/test-vector-index-faults.cpp | 157 ++++++++++++++++++- tests/test-vector-index.cpp | 122 ++++++++++++++- 8 files changed, 540 insertions(+), 13 deletions(-) diff --git a/ggml/cmake/ggml-config.cmake.in b/ggml/cmake/ggml-config.cmake.in index 75f824225bd8..2d9c0edebdb3 100644 --- a/ggml/cmake/ggml-config.cmake.in +++ b/ggml/cmake/ggml-config.cmake.in @@ -120,6 +120,13 @@ set_and_check(GGML_LIB_DIR "@PACKAGE_GGML_LIB_INSTALL_DIR@") # Include the exported targets file include("${CMAKE_CURRENT_LIST_DIR}/ggml-targets.cmake") +if(TARGET ggml::ggml-vector-index AND NOT TARGET ggml::vector-index) + add_library(ggml::vector-index INTERFACE IMPORTED) + set_target_properties(ggml::vector-index + PROPERTIES + INTERFACE_LINK_LIBRARIES ggml::ggml-vector-index) +endif() + if(NOT TARGET ggml::ggml) find_package(Threads REQUIRED) diff --git a/ggml/include/ggml-vector-index.h b/ggml/include/ggml-vector-index.h index d5bcdea0e671..82c74c391606 100644 --- a/ggml/include/ggml-vector-index.h +++ b/ggml/include/ggml-vector-index.h @@ -28,8 +28,9 @@ extern "C" { struct ggml_vec_index; typedef struct ggml_vec_index ggml_vec_index_t; -// Prepared filtered-search handle. Valid only for the index generation it was -// created from; any successful add/remove invalidates existing filters. +// Prepared filtered-search handle. Valid only while the source index remains +// alive and at the generation it was created from; any successful add/remove +// invalidates existing filters. struct ggml_vec_index_filter; typedef struct ggml_vec_index_filter ggml_vec_index_filter_t; @@ -149,7 +150,8 @@ GGML_API int ggml_vec_index_search_filtered( // Prepared filtered search. Creating a filter maps, sorts, and deduplicates // `allowed_ids` once, so callers can reuse it for repeated searches over the -// same allowlist. Stale filters return GGML_VEC_INDEX_E_INVALID_ARG. +// same allowlist. The source index must outlive every filter created from it. +// Stale filters return GGML_VEC_INDEX_E_INVALID_ARG. GGML_API ggml_vec_index_filter_t * ggml_vec_index_filter_create( const ggml_vec_index_t * idx, const uint64_t * allowed_ids, diff --git a/ggml/src/CMakeLists.txt b/ggml/src/CMakeLists.txt index 52739415f5c3..687abfbe1d52 100644 --- a/ggml/src/CMakeLists.txt +++ b/ggml/src/CMakeLists.txt @@ -272,7 +272,7 @@ set_target_properties(ggml-vector-index PROPERTIES SOVERSION ${GGML_VERSION_MAJOR} ) -target_link_libraries(ggml-vector-index PUBLIC ggml-base) +target_link_libraries(ggml-vector-index PUBLIC ggml-base PRIVATE Threads::Threads) function(ggml_add_backend_library backend) if (GGML_BACKEND_DL) diff --git a/ggml/src/ggml-vector-index.cpp b/ggml/src/ggml-vector-index.cpp index 930f78b2d6db..dcc1334d8104 100644 --- a/ggml/src/ggml-vector-index.cpp +++ b/ggml/src/ggml-vector-index.cpp @@ -16,6 +16,7 @@ #include #include #include +#include #include #include #include @@ -30,6 +31,8 @@ #include #include #include +#include +#include #include #include #include @@ -75,6 +78,8 @@ extern "C" { void ggml_vec_index_test_set_oom_countdown(int64_t countdown); void ggml_vec_index_test_set_write_fail_after(int64_t bytes); void ggml_vec_index_test_set_truncate_fail(int fail); +void ggml_vec_index_test_set_parent_fsync_fail(int fail); +void ggml_vec_index_test_set_delta_append_wait_target(int target); } #endif @@ -152,6 +157,9 @@ MappedFile::~MappedFile() { std::atomic g_test_oom_countdown{ -1 }; std::atomic g_test_write_fail_after{ -1 }; std::atomic g_test_truncate_fail{ false }; +std::atomic g_test_parent_fsync_fail{ false }; +std::atomic g_test_delta_append_wait_target{ 0 }; +std::atomic g_test_delta_append_waiters{ 0 }; void test_maybe_throw_bad_alloc() { const int64_t remaining = g_test_oom_countdown.load(); @@ -175,8 +183,23 @@ bool test_consume_write_bytes(size_t n) { g_test_write_fail_after.fetch_sub(static_cast(n)); return true; } + +void test_wait_after_delta_validate() { + const int target = g_test_delta_append_wait_target.load(); + if (target <= 0) { + return; + } + + g_test_delta_append_waiters.fetch_add(1); + const auto deadline = std::chrono::steady_clock::now() + std::chrono::milliseconds(250); + while (g_test_delta_append_waiters.load() < target && + std::chrono::steady_clock::now() < deadline) { + std::this_thread::sleep_for(std::chrono::milliseconds(1)); + } +} #else #define test_maybe_throw_bad_alloc() ((void) 0) +#define test_wait_after_delta_validate() ((void) 0) #endif inline bool write_bytes(std::FILE * f, const void * data, size_t size) { @@ -589,6 +612,11 @@ bool flush_and_sync(std::FILE * stream) { } bool fsync_parent_dir(const char * path) { +#ifdef GGML_VEC_INDEX_TEST_HOOKS + if (g_test_parent_fsync_fail.load()) { + return false; + } +#endif #ifdef _WIN32 (void) path; return true; @@ -644,9 +672,20 @@ void ggml_vec_index_test_set_write_fail_after(int64_t bytes) { void ggml_vec_index_test_set_truncate_fail(int fail) { g_test_truncate_fail.store(fail != 0); } + +void ggml_vec_index_test_set_parent_fsync_fail(int fail) { + g_test_parent_fsync_fail.store(fail != 0); +} + +void ggml_vec_index_test_set_delta_append_wait_target(int target) { + g_test_delta_append_waiters.store(0); + g_test_delta_append_wait_target.store(target); +} } #endif +static std::atomic g_next_filter_cookie{ 1 }; + // Lifetime-managed instance state. Lives behind the opaque // `ggml_vec_index_t` typedef. struct ggml_vec_index { @@ -655,6 +694,7 @@ struct ggml_vec_index { int dim = 0; int bit_width = 32; uint64_t generation = 0; + uint64_t filter_cookie = 0; bool read_only_mmap = false; bool delta_log_rebase_pending = false; uint32_t delta_log_rebase_crc = 0; @@ -693,6 +733,7 @@ struct ggml_vec_index { struct ggml_vec_index_filter { const ggml_vec_index_t * owner = nullptr; + uint64_t owner_cookie = 0; int dim = 0; int bit_width = 32; uint64_t generation = 0; @@ -784,6 +825,22 @@ static int q4_decode(uint8_t nibble) { return static_cast(nibble) - 8; } +static int round_nearest_even(float value) { + const float lower_f = std::floor(value); + const float upper_f = lower_f + 1.0f; + const float lower_dist = value - lower_f; + const float upper_dist = upper_f - value; + if (lower_dist < upper_dist) { + return static_cast(lower_f); + } + if (upper_dist < lower_dist) { + return static_cast(upper_f); + } + + const int lower = static_cast(lower_f); + return (lower % 2) == 0 ? lower : static_cast(upper_f); +} + static void quantize_q8_row(const float * src, int8_t * dst, int dim, float & scale) { float max_abs = 0.0f; for (int i = 0; i < dim; ++i) { @@ -802,7 +859,7 @@ static void quantize_q8_row(const float * src, int8_t * dst, int dim, float & sc } for (int i = 0; i < dim; ++i) { const float scaled = src[i] / scale; - int q = static_cast(std::nearbyint(scaled)); + int q = round_nearest_even(scaled); q = std::max(-127, std::min(127, q)); dst[i] = static_cast(q); } @@ -826,7 +883,7 @@ static void quantize_q4_row(const float * src, uint8_t * dst, int dim, float & s } for (int i = 0; i < dim; ++i) { const float scaled = src[i] / scale; - int q = static_cast(std::nearbyint(scaled)); + int q = round_nearest_even(scaled); q = std::max(-7, std::min(7, q)); const uint8_t code = q4_encode(q); uint8_t & byte = dst[static_cast(i) / 2]; @@ -975,6 +1032,137 @@ bool filesystem_path_from_utf8(const char * path, std::filesystem::path & out) { return true; } +bool filesystem_paths_equal(const char * lhs, const char * rhs) { + if (std::strcmp(lhs, rhs) == 0) { + return true; + } + + std::filesystem::path lhs_path; + std::filesystem::path rhs_path; + if (!filesystem_path_from_utf8(lhs, lhs_path) || + !filesystem_path_from_utf8(rhs, rhs_path)) { + return false; + } + + std::error_code lhs_ec; + std::error_code rhs_ec; + std::filesystem::path lhs_resolved = std::filesystem::weakly_canonical(lhs_path, lhs_ec); + std::filesystem::path rhs_resolved = std::filesystem::weakly_canonical(rhs_path, rhs_ec); + if (lhs_ec || rhs_ec) { + lhs_ec.clear(); + rhs_ec.clear(); + lhs_resolved = std::filesystem::absolute(lhs_path, lhs_ec); + rhs_resolved = std::filesystem::absolute(rhs_path, rhs_ec); + if (lhs_ec || rhs_ec) { + return false; + } + } + return lhs_resolved.lexically_normal() == rhs_resolved.lexically_normal(); +} + +std::mutex & delta_log_process_mutex() { + static std::mutex mutex; + return mutex; +} + +bool delta_lock_path(const char * path, std::filesystem::path & out) { + std::filesystem::path fs_path; + if (!filesystem_path_from_utf8(path, fs_path)) { + return false; + } + + std::error_code ec; + out = std::filesystem::weakly_canonical(fs_path, ec); + if (ec) { + ec.clear(); + out = std::filesystem::absolute(fs_path, ec); + if (ec) { + return false; + } + } + out = out.lexically_normal(); + out += ".lock"; + return true; +} + +class DeltaLogLock { +public: + explicit DeltaLogLock(const char * path) : process_lock(delta_log_process_mutex()) { + std::filesystem::path lock_path; + if (!delta_lock_path(path, lock_path)) { + return; + } +#ifdef _WIN32 + file = CreateFileW( + lock_path.wstring().c_str(), + GENERIC_READ | GENERIC_WRITE, + FILE_SHARE_READ | FILE_SHARE_WRITE | FILE_SHARE_DELETE, + nullptr, + OPEN_ALWAYS, + FILE_ATTRIBUTE_NORMAL, + nullptr); + if (file == INVALID_HANDLE_VALUE) { + return; + } + OVERLAPPED overlapped = {}; + if (LockFileEx(file, LOCKFILE_EXCLUSIVE_LOCK, 0, MAXDWORD, MAXDWORD, &overlapped) == 0) { + CloseHandle(file); + file = INVALID_HANDLE_VALUE; + return; + } +#else + fd = ::open(lock_path.c_str(), O_CREAT | O_RDWR, 0666); + if (fd < 0) { + return; + } + struct flock lock = {}; + lock.l_type = F_WRLCK; + lock.l_whence = SEEK_SET; + while (::fcntl(fd, F_SETLKW, &lock) != 0) { + if (errno == EINTR) { + continue; + } + ::close(fd); + fd = -1; + return; + } +#endif + locked = true; + } + + ~DeltaLogLock() { + if (!locked) { + return; + } +#ifdef _WIN32 + OVERLAPPED overlapped = {}; + UnlockFileEx(file, 0, MAXDWORD, MAXDWORD, &overlapped); + CloseHandle(file); + file = INVALID_HANDLE_VALUE; +#else + struct flock lock = {}; + lock.l_type = F_UNLCK; + lock.l_whence = SEEK_SET; + (void) ::fcntl(fd, F_SETLK, &lock); + ::close(fd); + fd = -1; +#endif + } + + bool ok() const { + return locked; + } + +private: + std::unique_lock process_lock; + bool locked = false; +#ifdef _WIN32 + HANDLE file = INVALID_HANDLE_VALUE; +#else + int fd = -1; +#endif +}; + bool open_append_file(const char * path, std::FILE ** out) { #ifdef _WIN32 std::wstring wide; @@ -1172,6 +1360,11 @@ int append_delta_record( if (delta_path == nullptr) { return GGML_VEC_INDEX_E_INVALID_ARG; } + DeltaLogLock delta_lock(delta_path); + if (!delta_lock.ok()) { + return GGML_VEC_INDEX_E_IO; + } + uint64_t old_size = 0; uint32_t existing_base_crc = 0; if (!validate_delta_header(delta_path, idx, old_size, existing_base_crc)) { @@ -1205,6 +1398,7 @@ int append_delta_record( idx.delta_log_rebase_crc = 0; } } + test_wait_after_delta_validate(); std::FILE * f = nullptr; if (!open_append_file(delta_path, &f)) { @@ -1258,7 +1452,7 @@ int append_delta_record( return GGML_VEC_INDEX_OK; } -int write_empty_delta_log(const ggml_vec_index & idx, const char * delta_path) { +int write_empty_delta_log_unlocked(const ggml_vec_index & idx, const char * delta_path) { if (delta_path == nullptr) { return GGML_VEC_INDEX_E_INVALID_ARG; } @@ -1303,6 +1497,17 @@ int write_empty_delta_log(const ggml_vec_index & idx, const char * delta_path) { } } +int write_empty_delta_log(const ggml_vec_index & idx, const char * delta_path) { + if (delta_path == nullptr) { + return GGML_VEC_INDEX_E_INVALID_ARG; + } + DeltaLogLock delta_lock(delta_path); + if (!delta_lock.ok()) { + return GGML_VEC_INDEX_E_IO; + } + return write_empty_delta_log_unlocked(idx, delta_path); +} + bool validate_logged_add_args( const ggml_vec_index_t * idx, const float * vectors, @@ -1396,6 +1601,10 @@ ggml_vec_index_t * ggml_vec_index_create(int dim, int bit_width) { } idx->dim = dim; idx->bit_width = bit_width; + idx->filter_cookie = g_next_filter_cookie.fetch_add(1, std::memory_order_relaxed); + if (idx->filter_cookie == 0) { + idx->filter_cookie = g_next_filter_cookie.fetch_add(1, std::memory_order_relaxed); + } return idx; } catch (...) { return nullptr; @@ -1785,6 +1994,8 @@ int ggml_vec_index_add_logged( if (record_is_complete) { ++idx->generation; invalidate_ivf(*idx); + added = false; + return GGML_VEC_INDEX_OK; } else { rollback_appended_slots_unlocked(idx, base_slot, ids, n); } @@ -1838,6 +2049,12 @@ int ggml_vec_index_remove_logged( post_remove_crc, payload); if (append_status != GGML_VEC_INDEX_OK) { + const bool record_is_complete = + append_status == GGML_VEC_INDEX_E_INTERNAL && + delta_log_ends_at_state(delta_path, *idx, post_remove_crc); + if (record_is_complete) { + (void) ggml_vec_index_remove_unlocked(idx, id); + } return append_status; } return ggml_vec_index_remove_unlocked(idx, id); @@ -2337,6 +2554,7 @@ static int ggml_vec_index_search_impl( const std::vector * allowed_ptr = nullptr; if (prepared_filter != nullptr) { if (prepared_filter->owner != idx || + prepared_filter->owner_cookie != idx->filter_cookie || prepared_filter->dim != idx->dim || prepared_filter->bit_width != idx->bit_width || prepared_filter->generation != idx->generation) { @@ -2405,6 +2623,7 @@ ggml_vec_index_filter_t * ggml_vec_index_filter_create( } std::unique_ptr owned(filter); owned->owner = idx; + owned->owner_cookie = idx->filter_cookie; owned->dim = idx->dim; owned->bit_width = idx->bit_width; owned->generation = idx->generation; @@ -3426,12 +3645,19 @@ int ggml_vec_index_compact_delta( if (idx == nullptr || snapshot_path == nullptr || delta_path == nullptr) { return GGML_VEC_INDEX_E_INVALID_ARG; } + if (filesystem_paths_equal(snapshot_path, delta_path)) { + return GGML_VEC_INDEX_E_INVALID_ARG; + } std::unique_lock lock(idx->mutex); + DeltaLogLock delta_lock(delta_path); + if (!delta_lock.ok()) { + return GGML_VEC_INDEX_E_IO; + } const int write_status = ggml_vec_index_write_unlocked(idx, snapshot_path); if (write_status != GGML_VEC_INDEX_OK) { return write_status; } - return write_empty_delta_log(*idx, delta_path); + return write_empty_delta_log_unlocked(*idx, delta_path); } catch (const std::bad_alloc &) { return GGML_VEC_INDEX_E_OOM; } catch (...) { diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index e9d16c3dedf6..7fc14ac08083 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -145,9 +145,11 @@ if (NOT WIN32) endif() # Vector-index regression test: no model required, no llama runtime touched. +find_package(Threads REQUIRED) + add_executable(test-vector-index test-vector-index.cpp) target_compile_features(test-vector-index PRIVATE cxx_std_17) -target_link_libraries(test-vector-index PRIVATE ggml-vector-index) +target_link_libraries(test-vector-index PRIVATE ggml-vector-index Threads::Threads) if (LLAMA_TESTS_INSTALL) install(TARGETS test-vector-index RUNTIME) endif() @@ -161,6 +163,7 @@ add_executable( target_compile_features(test-vector-index-faults PRIVATE cxx_std_17) target_compile_definitions(test-vector-index-faults PRIVATE GGML_VEC_INDEX_TEST_HOOKS) target_include_directories(test-vector-index-faults PRIVATE ${PROJECT_SOURCE_DIR}/ggml/include) +target_link_libraries(test-vector-index-faults PRIVATE Threads::Threads) if (NOT WIN32) target_link_libraries(test-vector-index-faults PRIVATE m) endif() diff --git a/tests/bench-vector-index.cpp b/tests/bench-vector-index.cpp index b4d9ae98fffb..97491aae714c 100644 --- a/tests/bench-vector-index.cpp +++ b/tests/bench-vector-index.cpp @@ -76,6 +76,22 @@ struct ScoreId { uint64_t id = 0; }; +int round_nearest_even(float value) { + const float lower_f = std::floor(value); + const float upper_f = lower_f + 1.0f; + const float lower_dist = value - lower_f; + const float upper_dist = upper_f - value; + if (lower_dist < upper_dist) { + return static_cast(lower_f); + } + if (upper_dist < lower_dist) { + return static_cast(upper_f); + } + + const int lower = static_cast(lower_f); + return (lower % 2) == 0 ? lower : static_cast(upper_f); +} + template double median_time_ms(int warmups, int repeats, Fn fn) { for (int i = 0; i < warmups; ++i) { @@ -392,7 +408,7 @@ Q4SimulatedIndex build_simulated_q4( int8_t * dst = sim.codes.data() + static_cast(row) * static_cast(dim); for (int i = 0; i < dim; ++i) { const float scaled = src[i] / scale; - int q = static_cast(std::nearbyint(scaled)); + int q = round_nearest_even(scaled); q = std::max(-7, std::min(7, q)); dst[i] = static_cast(q); } diff --git a/tests/test-vector-index-faults.cpp b/tests/test-vector-index-faults.cpp index 806e596292d6..0a771724a9cf 100644 --- a/tests/test-vector-index-faults.cpp +++ b/tests/test-vector-index-faults.cpp @@ -1,17 +1,21 @@ #include "ggml-vector-index.h" #include +#include #include #include #include #include #include #include +#include #include extern "C" void ggml_vec_index_test_set_oom_countdown(int64_t countdown); extern "C" void ggml_vec_index_test_set_write_fail_after(int64_t bytes); extern "C" void ggml_vec_index_test_set_truncate_fail(int fail); +extern "C" void ggml_vec_index_test_set_parent_fsync_fail(int fail); +extern "C" void ggml_vec_index_test_set_delta_append_wait_target(int target); namespace { @@ -27,6 +31,8 @@ void reset_fault_hooks() { ggml_vec_index_test_set_oom_countdown(-1); ggml_vec_index_test_set_write_fail_after(-1); ggml_vec_index_test_set_truncate_fail(0); + ggml_vec_index_test_set_parent_fsync_fail(0); + ggml_vec_index_test_set_delta_append_wait_target(0); } std::vector read_file_bytes(const std::string & path) { @@ -240,12 +246,161 @@ int main() { CHECK(logged_out_ids[0] == allowed_id); ggml_vec_index_filter_free(rollback_filter); + auto * committed_add_filter = ggml_vec_index_filter_create(idx, &allowed_id, 1); + CHECK(committed_add_filter != nullptr); + CHECK(ggml_vec_index_build_ivf(idx, /*n_lists=*/2, /*n_iter=*/1) + == GGML_VEC_INDEX_OK); + const uint64_t committed_add_id = 403; + ggml_vec_index_test_set_parent_fsync_fail(1); + ggml_vec_index_test_set_truncate_fail(1); + CHECK(ggml_vec_index_add_logged( + idx, extra_vector.data(), 1, &committed_add_id, delta_path.c_str()) == + GGML_VEC_INDEX_OK); + reset_fault_hooks(); + CHECK(ggml_vec_index_contains(idx, committed_add_id) == 1); + CHECK(ggml_vec_index_len(idx) == 5); + CHECK(ggml_vec_index_search_prepared_filtered( + idx, committed_add_filter, base_vectors.data(), 1, /*k=*/1, + logged_scores.data(), logged_out_ids.data()) == GGML_VEC_INDEX_E_INVALID_ARG); + CHECK(ggml_vec_index_search_ivf( + idx, base_vectors.data(), 1, /*k=*/1, /*nprobe=*/2, + logged_scores.data(), logged_out_ids.data()) == GGML_VEC_INDEX_E_INVALID_ARG); + ggml_vec_index_filter_free(committed_add_filter); + const std::vector delta_after_committed_add = read_file_bytes(delta_path); + ggml_vec_index_test_set_write_fail_after(8); CHECK(ggml_vec_index_remove_logged(idx, logged_id, delta_path.c_str()) == GGML_VEC_INDEX_E_IO); reset_fault_hooks(); CHECK(ggml_vec_index_contains(idx, logged_id) == 1); - CHECK(read_file_bytes(delta_path) == old_delta); + CHECK(read_file_bytes(delta_path) == delta_after_committed_add); + + auto * remove_filter = ggml_vec_index_filter_create(idx, &logged_id, 1); + CHECK(remove_filter != nullptr); + CHECK(ggml_vec_index_build_ivf(idx, /*n_lists=*/2, /*n_iter=*/1) + == GGML_VEC_INDEX_OK); + ggml_vec_index_test_set_parent_fsync_fail(1); + ggml_vec_index_test_set_truncate_fail(1); + CHECK(ggml_vec_index_remove_logged(idx, logged_id, delta_path.c_str()) == + GGML_VEC_INDEX_E_INTERNAL); + reset_fault_hooks(); + CHECK(ggml_vec_index_contains(idx, logged_id) == 0); + CHECK(ggml_vec_index_len(idx) == 4); + CHECK(read_file_bytes(delta_path) != delta_after_committed_add); + CHECK(ggml_vec_index_search_prepared_filtered( + idx, remove_filter, logged_vector.data(), 1, /*k=*/1, + logged_scores.data(), logged_out_ids.data()) == GGML_VEC_INDEX_E_INVALID_ARG); + CHECK(ggml_vec_index_search_ivf( + idx, base_vectors.data(), 1, /*k=*/1, /*nprobe=*/2, + logged_scores.data(), logged_out_ids.data()) == GGML_VEC_INDEX_E_INVALID_ARG); + ggml_vec_index_filter_free(remove_filter); + + const std::string committed_add_snapshot_path = + (std::filesystem::temp_directory_path() / + "ggml-vector-index-committed-add-base.tvim").string(); + const std::string committed_add_delta_path = + (std::filesystem::temp_directory_path() / + "ggml-vector-index-committed-add-log.tvid").string(); + std::filesystem::remove(committed_add_snapshot_path); + std::filesystem::remove(committed_add_delta_path); + std::filesystem::remove(committed_add_delta_path + ".lock"); + + auto * committed_base = ggml_vec_index_create(dim, /*bit_width=*/32); + CHECK(committed_base != nullptr); + CHECK(ggml_vec_index_add( + committed_base, base_vectors.data(), 2, base_ids.data()) == GGML_VEC_INDEX_OK); + CHECK(ggml_vec_index_write(committed_base, committed_add_snapshot_path.c_str()) == + GGML_VEC_INDEX_OK); + const uint64_t committed_replay_id = 601; + ggml_vec_index_test_set_parent_fsync_fail(1); + ggml_vec_index_test_set_truncate_fail(1); + CHECK(ggml_vec_index_add_logged( + committed_base, extra_vector.data(), 1, + &committed_replay_id, committed_add_delta_path.c_str()) == + GGML_VEC_INDEX_OK); + reset_fault_hooks(); + CHECK(ggml_vec_index_contains(committed_base, committed_replay_id) == 1); + auto * committed_replayed = ggml_vec_index_load_with_delta( + committed_add_snapshot_path.c_str(), committed_add_delta_path.c_str()); + CHECK(committed_replayed != nullptr); + CHECK(ggml_vec_index_len(committed_replayed) == 3); + CHECK(ggml_vec_index_contains(committed_replayed, committed_replay_id) == 1); + ggml_vec_index_free(committed_replayed); + ggml_vec_index_free(committed_base); + std::filesystem::remove(committed_add_snapshot_path); + std::filesystem::remove(committed_add_delta_path); + std::filesystem::remove(committed_add_delta_path + ".lock"); + + const std::string shared_snapshot_path = + (std::filesystem::temp_directory_path() / + "ggml-vector-index-shared-delta-base.tvim").string(); + const std::string shared_delta_path = + (std::filesystem::temp_directory_path() / + "ggml-vector-index-shared-delta-log.tvid").string(); + std::filesystem::remove(shared_snapshot_path); + std::filesystem::remove(shared_delta_path); + std::filesystem::remove(shared_delta_path + ".lock"); + + auto * shared_base = ggml_vec_index_create(dim, /*bit_width=*/32); + CHECK(shared_base != nullptr); + CHECK(ggml_vec_index_add( + shared_base, base_vectors.data(), 2, base_ids.data()) == GGML_VEC_INDEX_OK); + CHECK(ggml_vec_index_write(shared_base, shared_snapshot_path.c_str()) == GGML_VEC_INDEX_OK); + ggml_vec_index_free(shared_base); + + auto * shared_a = ggml_vec_index_load(shared_snapshot_path.c_str()); + auto * shared_b = ggml_vec_index_load(shared_snapshot_path.c_str()); + CHECK(shared_a != nullptr); + CHECK(shared_b != nullptr); + const uint64_t shared_id_a = 501; + const uint64_t shared_id_b = 502; + std::atomic start_shared_appends{ false }; + int status_a = GGML_VEC_INDEX_E_INTERNAL; + int status_b = GGML_VEC_INDEX_E_INTERNAL; + ggml_vec_index_test_set_delta_append_wait_target(2); + std::thread thread_a([&]() { + while (!start_shared_appends.load()) { + std::this_thread::yield(); + } + status_a = ggml_vec_index_add_logged( + shared_a, logged_vector.data(), 1, &shared_id_a, shared_delta_path.c_str()); + }); + std::thread thread_b([&]() { + while (!start_shared_appends.load()) { + std::this_thread::yield(); + } + status_b = ggml_vec_index_add_logged( + shared_b, extra_vector.data(), 1, &shared_id_b, shared_delta_path.c_str()); + }); + start_shared_appends.store(true); + thread_a.join(); + thread_b.join(); + reset_fault_hooks(); + + CHECK((status_a == GGML_VEC_INDEX_OK && status_b == GGML_VEC_INDEX_E_IO) || + (status_b == GGML_VEC_INDEX_OK && status_a == GGML_VEC_INDEX_E_IO)); + CHECK(ggml_vec_index_contains(shared_a, shared_id_a) == + (status_a == GGML_VEC_INDEX_OK ? 1 : 0)); + CHECK(ggml_vec_index_contains(shared_b, shared_id_b) == + (status_b == GGML_VEC_INDEX_OK ? 1 : 0)); + + auto * shared_replayed = ggml_vec_index_load_with_delta( + shared_snapshot_path.c_str(), shared_delta_path.c_str()); + CHECK(shared_replayed != nullptr); + CHECK(ggml_vec_index_len(shared_replayed) == 3); + CHECK(ggml_vec_index_contains(shared_replayed, base_ids[0]) == 1); + CHECK(ggml_vec_index_contains(shared_replayed, base_ids[1]) == 1); + CHECK(ggml_vec_index_contains(shared_replayed, shared_id_a) == + (status_a == GGML_VEC_INDEX_OK ? 1 : 0)); + CHECK(ggml_vec_index_contains(shared_replayed, shared_id_b) == + (status_b == GGML_VEC_INDEX_OK ? 1 : 0)); + + ggml_vec_index_free(shared_replayed); + ggml_vec_index_free(shared_a); + ggml_vec_index_free(shared_b); + std::filesystem::remove(shared_snapshot_path); + std::filesystem::remove(shared_delta_path); + std::filesystem::remove(shared_delta_path + ".lock"); ggml_vec_index_free(idx); std::filesystem::remove(path); diff --git a/tests/test-vector-index.cpp b/tests/test-vector-index.cpp index cddee225c680..6a7e5fb32250 100644 --- a/tests/test-vector-index.cpp +++ b/tests/test-vector-index.cpp @@ -10,6 +10,7 @@ #include #include #include +#include #include #include #include @@ -46,6 +47,22 @@ std::vector normalize(std::vector v) { return v; } +int round_nearest_even(float value) { + const float lower_f = std::floor(value); + const float upper_f = lower_f + 1.0f; + const float lower_dist = value - lower_f; + const float upper_dist = upper_f - value; + if (lower_dist < upper_dist) { + return static_cast(lower_f); + } + if (upper_dist < lower_dist) { + return static_cast(upper_f); + } + + const int lower = static_cast(lower_f); + return (lower % 2) == 0 ? lower : static_cast(upper_f); +} + float q8_dot_reference(const std::vector & vector, const std::vector & query) { CHECK(vector.size() == query.size()); @@ -58,7 +75,7 @@ float q8_dot_reference(const std::vector & vector, const std::vector(std::nearbyint(vector[i] / scale)); + 0 : round_nearest_even(vector[i] / scale); code = std::max(-127, std::min(127, code)); acc += query[i] * (static_cast(code) * scale); } @@ -77,7 +94,7 @@ float q4_dot_reference(const std::vector & vector, const std::vector(std::nearbyint(vector[i] / scale)); + 0 : round_nearest_even(vector[i] / scale); code = std::max(-7, std::min(7, code)); acc += query[i] * (static_cast(code) * scale); } @@ -815,6 +832,18 @@ int main() { ggml_vec_index_free(mismatch); const std::vector pre_compact_delta = read_file_bytes(delta_path); + const std::vector pre_same_path_snapshot = read_file_bytes(snapshot_path); + CHECK(ggml_vec_index_compact_delta( + base, snapshot_path.c_str(), snapshot_path.c_str()) == + GGML_VEC_INDEX_E_INVALID_ARG); + CHECK(read_file_bytes(snapshot_path) == pre_same_path_snapshot); + auto * same_path_snapshot = ggml_vec_index_load(snapshot_path.c_str()); + CHECK(same_path_snapshot != nullptr); + CHECK(ggml_vec_index_len(same_path_snapshot) == 2); + CHECK(ggml_vec_index_contains(same_path_snapshot, ids[0]) == 1); + CHECK(ggml_vec_index_contains(same_path_snapshot, delta_id) == 0); + ggml_vec_index_free(same_path_snapshot); + CHECK(ggml_vec_index_compact_delta( base, snapshot_path.c_str(), delta_path.c_str()) == GGML_VEC_INDEX_OK); CHECK(std::filesystem::file_size(delta_path) == 16); @@ -1220,6 +1249,95 @@ int main() { ggml_vec_index_free(q4); } + // Quantization must not depend on the caller's active rounding mode. + { + const int saved_rounding_mode = std::fegetround(); + CHECK(saved_rounding_mode != -1); + + auto score_after_add_with_rounding = [&](int bit_width, int rounding_mode) { + const float max_code = bit_width == 8 ? 127.0f : 7.0f; + const std::array rounding_vector = { + max_code, 2.5f, 0.0f, 0.0f, + }; + const std::array rounding_query = { + 0.0f, 1.0f, 0.0f, 0.0f, + }; + const uint64_t rounding_id = + (1ULL << 57) + static_cast(bit_width) + + static_cast(rounding_mode); + CHECK(std::fesetround(rounding_mode) == 0); + auto * rounding_idx = ggml_vec_index_create(kDim, bit_width); + CHECK(rounding_idx != nullptr); + CHECK(ggml_vec_index_add( + rounding_idx, rounding_vector.data(), 1, &rounding_id) == + GGML_VEC_INDEX_OK); + CHECK(std::fesetround(saved_rounding_mode) == 0); + + std::array scores{}; + std::array out_ids{}; + CHECK(ggml_vec_index_search( + rounding_idx, rounding_query.data(), 1, /*k=*/1, + scores.data(), out_ids.data()) == GGML_VEC_INDEX_OK); + CHECK(out_ids[0] == rounding_id); + ggml_vec_index_free(rounding_idx); + return scores[0]; + }; + + for (int bit_width : { 8, 4 }) { + const float downward_score = score_after_add_with_rounding(bit_width, FE_DOWNWARD); + const float upward_score = score_after_add_with_rounding(bit_width, FE_UPWARD); + CHECK(downward_score == upward_score); + CHECK(downward_score == 2.0f); + + const std::string snapshot_path = + (std::filesystem::temp_directory_path() / + ("ggml-vector-index-rounding-" + std::to_string(bit_width) + ".tvim")).string(); + const std::string delta_path = + (std::filesystem::temp_directory_path() / + ("ggml-vector-index-rounding-" + std::to_string(bit_width) + ".tvid")).string(); + std::filesystem::remove(snapshot_path); + std::filesystem::remove(delta_path); + std::filesystem::remove(delta_path + ".lock"); + + const float max_code = bit_width == 8 ? 127.0f : 7.0f; + const std::array rounding_vector = { + max_code, 2.5f, 0.0f, 0.0f, + }; + const std::array rounding_query = { + 0.0f, 1.0f, 0.0f, 0.0f, + }; + const uint64_t rounding_id = (1ULL << 58) + static_cast(bit_width); + auto * logged_rounding = ggml_vec_index_create(kDim, bit_width); + CHECK(logged_rounding != nullptr); + CHECK(ggml_vec_index_write(logged_rounding, snapshot_path.c_str()) == + GGML_VEC_INDEX_OK); + CHECK(std::fesetround(FE_DOWNWARD) == 0); + CHECK(ggml_vec_index_add_logged( + logged_rounding, rounding_vector.data(), 1, &rounding_id, + delta_path.c_str()) == GGML_VEC_INDEX_OK); + CHECK(std::fesetround(FE_UPWARD) == 0); + auto * replayed_rounding = ggml_vec_index_load_with_delta( + snapshot_path.c_str(), delta_path.c_str()); + CHECK(std::fesetround(saved_rounding_mode) == 0); + CHECK(replayed_rounding != nullptr); + + std::array scores{}; + std::array out_ids{}; + CHECK(ggml_vec_index_search( + replayed_rounding, rounding_query.data(), 1, /*k=*/1, + scores.data(), out_ids.data()) == GGML_VEC_INDEX_OK); + CHECK(out_ids[0] == rounding_id); + CHECK(scores[0] == 2.0f); + + ggml_vec_index_free(replayed_rounding); + ggml_vec_index_free(logged_rounding); + std::filesystem::remove(snapshot_path); + std::filesystem::remove(delta_path); + std::filesystem::remove(delta_path + ".lock"); + } + CHECK(std::fesetround(saved_rounding_mode) == 0); + } + // Delta replay keeps q8 storage q8 and reuses the normal q8 quantization path. { const std::string snapshot_path = From 6db17cc9e928ba1731a32c1355ee75794cd6e0b4 Mon Sep 17 00:00:00 2001 From: Nidhin Date: Thu, 16 Jul 2026 18:27:49 +0530 Subject: [PATCH 11/13] ggml-vector-index: fix logged mutation durability edge cases Return success when logged mutations or snapshot replacements have already committed despite late durability checks, and cache validated delta tails to avoid rescanning the full log on repeated appends. --- ggml/src/ggml-vector-index.cpp | 98 +++++++++++++++++++---- tests/test-vector-index-faults.cpp | 120 ++++++++++++++++++++++++++++- 2 files changed, 201 insertions(+), 17 deletions(-) diff --git a/ggml/src/ggml-vector-index.cpp b/ggml/src/ggml-vector-index.cpp index dcc1334d8104..efbddfa6e637 100644 --- a/ggml/src/ggml-vector-index.cpp +++ b/ggml/src/ggml-vector-index.cpp @@ -80,6 +80,8 @@ void ggml_vec_index_test_set_write_fail_after(int64_t bytes); void ggml_vec_index_test_set_truncate_fail(int fail); void ggml_vec_index_test_set_parent_fsync_fail(int fail); void ggml_vec_index_test_set_delta_append_wait_target(int target); +void ggml_vec_index_test_reset_delta_tail_scan_count(void); +int64_t ggml_vec_index_test_get_delta_tail_scan_count(void); } #endif @@ -160,6 +162,7 @@ std::atomic g_test_truncate_fail{ false }; std::atomic g_test_parent_fsync_fail{ false }; std::atomic g_test_delta_append_wait_target{ 0 }; std::atomic g_test_delta_append_waiters{ 0 }; +std::atomic g_test_delta_tail_scan_count{ 0 }; void test_maybe_throw_bad_alloc() { const int64_t remaining = g_test_oom_countdown.load(); @@ -681,6 +684,14 @@ void ggml_vec_index_test_set_delta_append_wait_target(int target) { g_test_delta_append_waiters.store(0); g_test_delta_append_wait_target.store(target); } + +void ggml_vec_index_test_reset_delta_tail_scan_count(void) { + g_test_delta_tail_scan_count.store(0); +} + +int64_t ggml_vec_index_test_get_delta_tail_scan_count(void) { + return g_test_delta_tail_scan_count.load(); +} } #endif @@ -698,6 +709,10 @@ struct ggml_vec_index { bool read_only_mmap = false; bool delta_log_rebase_pending = false; uint32_t delta_log_rebase_crc = 0; + bool delta_tail_cache_valid = false; + std::string delta_tail_cache_path; + uint64_t delta_tail_cache_size = 0; + uint32_t delta_tail_cache_crc = 0; std::unique_ptr mapped_file; size_t mapped_vector_bytes = 0; @@ -1060,6 +1075,49 @@ bool filesystem_paths_equal(const char * lhs, const char * rhs) { return lhs_resolved.lexically_normal() == rhs_resolved.lexically_normal(); } +void invalidate_delta_tail_cache(ggml_vec_index & idx) noexcept { + idx.delta_tail_cache_valid = false; + idx.delta_tail_cache_size = 0; + idx.delta_tail_cache_crc = 0; +} + +void set_delta_tail_cache( + ggml_vec_index & idx, + const char * path, + uint64_t size, + uint32_t tail_crc) noexcept { + try { + idx.delta_tail_cache_path = path; + idx.delta_tail_cache_size = size; + idx.delta_tail_cache_crc = tail_crc; + idx.delta_tail_cache_valid = true; + } catch (...) { + invalidate_delta_tail_cache(idx); + } +} + +bool get_delta_tail_cache( + ggml_vec_index & idx, + const char * path, + uint64_t file_size, + uint32_t & tail_crc, + uint64_t & complete_size) { + if (!idx.delta_tail_cache_valid || + idx.delta_tail_cache_size != file_size || + idx.delta_tail_cache_path.empty()) { + return false; + } + const bool same_path = + idx.delta_tail_cache_path == path || + filesystem_paths_equal(idx.delta_tail_cache_path.c_str(), path); + if (!same_path) { + return false; + } + tail_crc = idx.delta_tail_cache_crc; + complete_size = idx.delta_tail_cache_size; + return true; +} + std::mutex & delta_log_process_mutex() { static std::mutex mutex; return mutex; @@ -1252,6 +1310,9 @@ bool inspect_delta_log_tail( const ggml_vec_index & idx, uint32_t & last_state_crc, uint64_t & complete_size) { +#ifdef GGML_VEC_INDEX_TEST_HOOKS + g_test_delta_tail_scan_count.fetch_add(1); +#endif std::filesystem::path fs_path; if (!filesystem_path_from_utf8(path, fs_path)) { return false; @@ -1373,14 +1434,17 @@ int append_delta_record( if (old_size != 0) { uint32_t tail_crc = 0; uint64_t complete_size = 0; - if (!inspect_delta_log_tail(delta_path, idx, tail_crc, complete_size)) { + if (!get_delta_tail_cache(idx, delta_path, old_size, tail_crc, complete_size) && + !inspect_delta_log_tail(delta_path, idx, tail_crc, complete_size)) { return GGML_VEC_INDEX_E_IO; } if (tail_crc != base_crc_for_new_log) { + invalidate_delta_tail_cache(idx); return GGML_VEC_INDEX_E_IO; } if (complete_size != old_size) { if (!truncate_file_to(delta_path, complete_size)) { + invalidate_delta_tail_cache(idx); return GGML_VEC_INDEX_E_INTERNAL; } old_size = complete_size; @@ -1391,11 +1455,13 @@ int append_delta_record( // The snapshot already includes this log's records (e.g. crash after // compacting the snapshot but before replacing the old delta log). if (!truncate_file_to(delta_path, 0)) { + invalidate_delta_tail_cache(idx); return GGML_VEC_INDEX_E_INTERNAL; } old_size = 0; idx.delta_log_rebase_pending = false; idx.delta_log_rebase_crc = 0; + invalidate_delta_tail_cache(idx); } } test_wait_after_delta_validate(); @@ -1412,9 +1478,9 @@ int append_delta_record( }; auto fail_io = [&]() { close_file(); - return truncate_file_to(delta_path, old_size) ? - GGML_VEC_INDEX_E_IO : - GGML_VEC_INDEX_E_INTERNAL; + const bool truncated = truncate_file_to(delta_path, old_size); + invalidate_delta_tail_cache(idx); + return truncated ? GGML_VEC_INDEX_E_IO : GGML_VEC_INDEX_E_INTERNAL; }; if (old_size == 0) { @@ -1445,10 +1511,16 @@ int append_delta_record( const int close_result = std::fclose(f); f = nullptr; if (close_result != 0 || !fsync_parent_dir(delta_path)) { - return truncate_file_to(delta_path, old_size) ? - GGML_VEC_INDEX_E_IO : - GGML_VEC_INDEX_E_INTERNAL; - } + const bool truncated = truncate_file_to(delta_path, old_size); + invalidate_delta_tail_cache(idx); + return truncated ? GGML_VEC_INDEX_E_IO : GGML_VEC_INDEX_E_INTERNAL; + } + const uint64_t written_size = + old_size + + (old_size == 0 ? kTvidHeaderSize : 0) + + kTvidRecordHeaderSize + + static_cast(payload.size()); + set_delta_tail_cache(idx, delta_path, written_size, state_crc); return GGML_VEC_INDEX_OK; } @@ -1486,9 +1558,7 @@ int write_empty_delta_log_unlocked(const ggml_vec_index & idx, const char * delt return fail_io(); } temp.path.clear(); - if (!fsync_parent_dir(delta_path)) { - return GGML_VEC_INDEX_E_IO; - } + (void) fsync_parent_dir(delta_path); return GGML_VEC_INDEX_OK; } catch (const std::bad_alloc &) { return GGML_VEC_INDEX_E_OOM; @@ -2053,7 +2123,7 @@ int ggml_vec_index_remove_logged( append_status == GGML_VEC_INDEX_E_INTERNAL && delta_log_ends_at_state(delta_path, *idx, post_remove_crc); if (record_is_complete) { - (void) ggml_vec_index_remove_unlocked(idx, id); + return ggml_vec_index_remove_unlocked(idx, id); } return append_status; } @@ -2889,9 +2959,7 @@ static int ggml_vec_index_write_unlocked(ggml_vec_index_t * idx, const char * pa return fail_io(); } temp.path.clear(); - if (!fsync_parent_dir(path)) { - return GGML_VEC_INDEX_E_IO; - } + (void) fsync_parent_dir(path); return GGML_VEC_INDEX_OK; } catch (const std::bad_alloc &) { return GGML_VEC_INDEX_E_OOM; diff --git a/tests/test-vector-index-faults.cpp b/tests/test-vector-index-faults.cpp index 0a771724a9cf..f7afda2f068a 100644 --- a/tests/test-vector-index-faults.cpp +++ b/tests/test-vector-index-faults.cpp @@ -16,6 +16,8 @@ extern "C" void ggml_vec_index_test_set_write_fail_after(int64_t bytes); extern "C" void ggml_vec_index_test_set_truncate_fail(int fail); extern "C" void ggml_vec_index_test_set_parent_fsync_fail(int fail); extern "C" void ggml_vec_index_test_set_delta_append_wait_target(int target); +extern "C" void ggml_vec_index_test_reset_delta_tail_scan_count(void); +extern "C" int64_t ggml_vec_index_test_get_delta_tail_scan_count(void); namespace { @@ -33,6 +35,7 @@ void reset_fault_hooks() { ggml_vec_index_test_set_truncate_fail(0); ggml_vec_index_test_set_parent_fsync_fail(0); ggml_vec_index_test_set_delta_append_wait_target(0); + ggml_vec_index_test_reset_delta_tail_scan_count(); } std::vector read_file_bytes(const std::string & path) { @@ -184,6 +187,40 @@ int main() { CHECK(ggml_vec_index_len(loaded) == 2); ggml_vec_index_free(loaded); + { + const std::string parent_fsync_path = + (std::filesystem::temp_directory_path() / + "ggml-vector-index-parent-fsync-test.tvim").string(); + std::filesystem::remove(parent_fsync_path); + + auto * parent_fsync_idx = ggml_vec_index_create(dim, /*bit_width=*/32); + CHECK(parent_fsync_idx != nullptr); + CHECK(ggml_vec_index_add( + parent_fsync_idx, base_vectors.data(), 2, base_ids.data()) == GGML_VEC_INDEX_OK); + CHECK(ggml_vec_index_write(parent_fsync_idx, parent_fsync_path.c_str()) == + GGML_VEC_INDEX_OK); + const std::vector before_parent_fsync = + read_file_bytes(parent_fsync_path); + + const uint64_t parent_fsync_id = 302; + CHECK(ggml_vec_index_add( + parent_fsync_idx, extra_vector.data(), 1, &parent_fsync_id) == + GGML_VEC_INDEX_OK); + ggml_vec_index_test_set_parent_fsync_fail(1); + CHECK(ggml_vec_index_write(parent_fsync_idx, parent_fsync_path.c_str()) == + GGML_VEC_INDEX_OK); + reset_fault_hooks(); + CHECK(read_file_bytes(parent_fsync_path) != before_parent_fsync); + + auto * parent_fsync_loaded = ggml_vec_index_load(parent_fsync_path.c_str()); + CHECK(parent_fsync_loaded != nullptr); + CHECK(ggml_vec_index_len(parent_fsync_loaded) == 3); + CHECK(ggml_vec_index_contains(parent_fsync_loaded, parent_fsync_id) == 1); + ggml_vec_index_free(parent_fsync_loaded); + ggml_vec_index_free(parent_fsync_idx); + std::filesystem::remove(parent_fsync_path); + } + const std::string delta_path = (std::filesystem::temp_directory_path() / "ggml-vector-index-fault-test.tvid").string(); @@ -281,8 +318,7 @@ int main() { == GGML_VEC_INDEX_OK); ggml_vec_index_test_set_parent_fsync_fail(1); ggml_vec_index_test_set_truncate_fail(1); - CHECK(ggml_vec_index_remove_logged(idx, logged_id, delta_path.c_str()) == - GGML_VEC_INDEX_E_INTERNAL); + CHECK(ggml_vec_index_remove_logged(idx, logged_id, delta_path.c_str()) == 1); reset_fault_hooks(); CHECK(ggml_vec_index_contains(idx, logged_id) == 0); CHECK(ggml_vec_index_len(idx) == 4); @@ -331,6 +367,86 @@ int main() { std::filesystem::remove(committed_add_delta_path); std::filesystem::remove(committed_add_delta_path + ".lock"); + const std::string compact_parent_snapshot_path = + (std::filesystem::temp_directory_path() / + "ggml-vector-index-compact-parent-fsync-base.tvim").string(); + const std::string compact_parent_delta_path = + (std::filesystem::temp_directory_path() / + "ggml-vector-index-compact-parent-fsync-log.tvid").string(); + std::filesystem::remove(compact_parent_snapshot_path); + std::filesystem::remove(compact_parent_delta_path); + std::filesystem::remove(compact_parent_delta_path + ".lock"); + + auto * compact_parent = ggml_vec_index_create(dim, /*bit_width=*/32); + CHECK(compact_parent != nullptr); + CHECK(ggml_vec_index_add( + compact_parent, base_vectors.data(), 2, base_ids.data()) == GGML_VEC_INDEX_OK); + CHECK(ggml_vec_index_write(compact_parent, compact_parent_snapshot_path.c_str()) == + GGML_VEC_INDEX_OK); + const uint64_t compact_parent_id = 602; + CHECK(ggml_vec_index_add_logged( + compact_parent, extra_vector.data(), 1, + &compact_parent_id, compact_parent_delta_path.c_str()) == GGML_VEC_INDEX_OK); + ggml_vec_index_test_set_parent_fsync_fail(1); + CHECK(ggml_vec_index_compact_delta( + compact_parent, + compact_parent_snapshot_path.c_str(), + compact_parent_delta_path.c_str()) == GGML_VEC_INDEX_OK); + reset_fault_hooks(); + CHECK(std::filesystem::file_size(compact_parent_delta_path) == 16); + auto * compact_parent_replayed = ggml_vec_index_load_with_delta( + compact_parent_snapshot_path.c_str(), compact_parent_delta_path.c_str()); + CHECK(compact_parent_replayed != nullptr); + CHECK(ggml_vec_index_len(compact_parent_replayed) == 3); + CHECK(ggml_vec_index_contains(compact_parent_replayed, compact_parent_id) == 1); + ggml_vec_index_free(compact_parent_replayed); + ggml_vec_index_free(compact_parent); + std::filesystem::remove(compact_parent_snapshot_path); + std::filesystem::remove(compact_parent_delta_path); + std::filesystem::remove(compact_parent_delta_path + ".lock"); + + const std::string cached_tail_snapshot_path = + (std::filesystem::temp_directory_path() / + "ggml-vector-index-cached-tail-base.tvim").string(); + const std::string cached_tail_delta_path = + (std::filesystem::temp_directory_path() / + "ggml-vector-index-cached-tail-log.tvid").string(); + std::filesystem::remove(cached_tail_snapshot_path); + std::filesystem::remove(cached_tail_delta_path); + std::filesystem::remove(cached_tail_delta_path + ".lock"); + + auto * cached_tail = ggml_vec_index_create(dim, /*bit_width=*/32); + CHECK(cached_tail != nullptr); + CHECK(ggml_vec_index_add( + cached_tail, base_vectors.data(), 2, base_ids.data()) == GGML_VEC_INDEX_OK); + CHECK(ggml_vec_index_write(cached_tail, cached_tail_snapshot_path.c_str()) == + GGML_VEC_INDEX_OK); + const uint64_t cached_tail_id_a = 701; + const uint64_t cached_tail_id_b = 702; + ggml_vec_index_test_reset_delta_tail_scan_count(); + CHECK(ggml_vec_index_add_logged( + cached_tail, logged_vector.data(), 1, + &cached_tail_id_a, cached_tail_delta_path.c_str()) == GGML_VEC_INDEX_OK); + CHECK(ggml_vec_index_add_logged( + cached_tail, extra_vector.data(), 1, + &cached_tail_id_b, cached_tail_delta_path.c_str()) == GGML_VEC_INDEX_OK); + CHECK(ggml_vec_index_remove_logged( + cached_tail, base_ids[0], cached_tail_delta_path.c_str()) == 1); + CHECK(ggml_vec_index_test_get_delta_tail_scan_count() == 0); + + auto * cached_tail_replayed = ggml_vec_index_load_with_delta( + cached_tail_snapshot_path.c_str(), cached_tail_delta_path.c_str()); + CHECK(cached_tail_replayed != nullptr); + CHECK(ggml_vec_index_len(cached_tail_replayed) == 3); + CHECK(ggml_vec_index_contains(cached_tail_replayed, base_ids[0]) == 0); + CHECK(ggml_vec_index_contains(cached_tail_replayed, cached_tail_id_a) == 1); + CHECK(ggml_vec_index_contains(cached_tail_replayed, cached_tail_id_b) == 1); + ggml_vec_index_free(cached_tail_replayed); + ggml_vec_index_free(cached_tail); + std::filesystem::remove(cached_tail_snapshot_path); + std::filesystem::remove(cached_tail_delta_path); + std::filesystem::remove(cached_tail_delta_path + ".lock"); + const std::string shared_snapshot_path = (std::filesystem::temp_directory_path() / "ggml-vector-index-shared-delta-base.tvim").string(); From c2d3db2f62857eb474aa9e750088ac1c2d54f232 Mon Sep 17 00:00:00 2001 From: Nidhin Date: Thu, 16 Jul 2026 21:38:33 +0530 Subject: [PATCH 12/13] Fix vector index persistence and coverage gaps Tighten mmap snapshot safety, delta durability/error handling, IVF probing, length limits, and package exports while expanding vector-index tests and benchmark coverage for edge cases, quantized persistence, and filtered q4 paths. --- ggml/cmake/ggml-config.cmake.in | 16 ++- ggml/include/ggml-vector-index.h | 30 ++++- ggml/src/CMakeLists.txt | 2 +- ggml/src/ggml-vector-index.cpp | 97 +++++++++++--- tests/CMakeLists.txt | 21 +-- tests/bench-vector-index.cpp | 34 ++++- tests/test-vector-index-faults.cpp | 97 +++++++++++++- tests/test-vector-index.cpp | 203 +++++++++++++++++++++++++++++ 8 files changed, 458 insertions(+), 42 deletions(-) diff --git a/ggml/cmake/ggml-config.cmake.in b/ggml/cmake/ggml-config.cmake.in index 2d9c0edebdb3..cae4feab2973 100644 --- a/ggml/cmake/ggml-config.cmake.in +++ b/ggml/cmake/ggml-config.cmake.in @@ -208,11 +208,25 @@ if(NOT TARGET ggml::ggml) PROPERTIES INTERFACE_LINK_LIBRARIES "${GGML_INTERFACE_LINK_LIBRARIES}") +endif() + +if(NOT TARGET ggml::all) + set(_ggml_all_targets "") + if (NOT GGML_BACKEND_DL) + foreach(_ggml_backend ${GGML_AVAILABLE_BACKENDS}) + if(TARGET ggml::${_ggml_backend}) + list(APPEND _ggml_all_targets ggml::${_ggml_backend}) + endif() + endforeach() + endif() + if(TARGET ggml::vector-index) + list(APPEND _ggml_all_targets ggml::vector-index) + endif() add_library(ggml::all INTERFACE IMPORTED) set_target_properties(ggml::all PROPERTIES INTERFACE_LINK_LIBRARIES "${_ggml_all_targets}") - + unset(_ggml_all_targets) endif() check_required_components(ggml) diff --git a/ggml/include/ggml-vector-index.h b/ggml/include/ggml-vector-index.h index 82c74c391606..66109617997e 100644 --- a/ggml/include/ggml-vector-index.h +++ b/ggml/include/ggml-vector-index.h @@ -16,10 +16,24 @@ // // Endianness: persistence format is fixed little-endian. -#include "ggml.h" - #include +#ifndef GGML_API +# ifdef GGML_SHARED +# if defined(_WIN32) && !defined(__MINGW32__) +# ifdef GGML_BUILD +# define GGML_API __declspec(dllexport) extern +# else +# define GGML_API __declspec(dllimport) extern +# endif +# else +# define GGML_API __attribute__ ((visibility ("default"))) extern +# endif +# else +# define GGML_API extern +# endif +#endif + #ifdef __cplusplus extern "C" { #endif @@ -38,7 +52,6 @@ typedef struct ggml_vec_index_filter ggml_vec_index_filter_t; // `_remove` is the exception: it returns 1 on removal and 0 on miss. enum ggml_vec_index_error { GGML_VEC_INDEX_OK = 0, - GGML_VEC_INDEX_E_INVALID_DIM = -1, GGML_VEC_INDEX_E_INVALID_ARG = -2, GGML_VEC_INDEX_E_DUPLICATE = -3, GGML_VEC_INDEX_E_IO = -4, @@ -64,7 +77,7 @@ GGML_API void ggml_vec_index_free(ggml_vec_index_t * idx); // Returns 0 on success. Returns GGML_VEC_INDEX_E_DUPLICATE if any id already // exists in the index; in that case the index is unchanged (atomic add). // All vector components must be finite. UINT64_MAX is reserved for search -// result padding and is not a valid id. +// result padding and is not a valid id. Live index length is capped at INT_MAX. GGML_API int ggml_vec_index_add( ggml_vec_index_t * idx, const float * vectors, @@ -91,12 +104,15 @@ GGML_API int ggml_vec_index_add_logged( const uint64_t * ids, const char * delta_path); +// Same return convention as `ggml_vec_index_remove`: 1 if removed, 0 if not +// present, negative on error. GGML_API int ggml_vec_index_remove_logged( ggml_vec_index_t * idx, uint64_t id, const char * delta_path); -// Returns 1 if the id is in the index, 0 otherwise. Read-only. +// Returns 1 if the id is in the index, 0 otherwise. NULL handles return 0. +// Read-only. GGML_API int ggml_vec_index_contains(const ggml_vec_index_t * idx, uint64_t id); // Placeholder for cache warming / codebook resolution after a bulk add. @@ -194,7 +210,7 @@ GGML_API ggml_vec_index_t * ggml_vec_index_load(const char * path); // Loads a v2 .tvim snapshot with its vector section memory-mapped read-only. // IDs and quantization scales are copied into memory for lookup and scoring. // Mutating APIs return GGML_VEC_INDEX_E_INVALID_ARG on mmap-backed handles. -// `ggml_vec_index_write` can snapshot mmap-backed handles, but callers should +// `ggml_vec_index_write` can snapshot mmap-backed handles, but callers must // write to a different path than the mapped source file. // Returns NULL on failure or unsupported file format. GGML_API ggml_vec_index_t * ggml_vec_index_load_mmap(const char * path); @@ -212,7 +228,7 @@ GGML_API int ggml_vec_index_compact_delta( const char * snapshot_path, const char * delta_path); -// Stats. +// Stats. NULL handles return 0. GGML_API int ggml_vec_index_len(const ggml_vec_index_t * idx); GGML_API int ggml_vec_index_dim(const ggml_vec_index_t * idx); GGML_API int ggml_vec_index_bit_width(const ggml_vec_index_t * idx); diff --git a/ggml/src/CMakeLists.txt b/ggml/src/CMakeLists.txt index 687abfbe1d52..7876be7ec7e4 100644 --- a/ggml/src/CMakeLists.txt +++ b/ggml/src/CMakeLists.txt @@ -272,7 +272,7 @@ set_target_properties(ggml-vector-index PROPERTIES SOVERSION ${GGML_VERSION_MAJOR} ) -target_link_libraries(ggml-vector-index PUBLIC ggml-base PRIVATE Threads::Threads) +target_link_libraries(ggml-vector-index PRIVATE Threads::Threads) function(ggml_add_backend_library backend) if (GGML_BACKEND_DL) diff --git a/ggml/src/ggml-vector-index.cpp b/ggml/src/ggml-vector-index.cpp index efbddfa6e637..1c216eeeb798 100644 --- a/ggml/src/ggml-vector-index.cpp +++ b/ggml/src/ggml-vector-index.cpp @@ -105,6 +105,7 @@ constexpr uint8_t kTvidOpAdd = 1; constexpr uint8_t kTvidOpRemove = 2; constexpr size_t kTvidHeaderSize = 16; constexpr size_t kTvidRecordHeaderSize = 24; +constexpr size_t kMaxIndexLen = static_cast(std::numeric_limits::max()); static_assert(sizeof(float) == sizeof(uint32_t), "ggml-vector-index requires float32"); @@ -715,6 +716,7 @@ struct ggml_vec_index { uint32_t delta_tail_cache_crc = 0; std::unique_ptr mapped_file; + std::string mapped_source_path; size_t mapped_vector_bytes = 0; const float * mapped_data = nullptr; const int8_t * mapped_q8_data = nullptr; @@ -1118,11 +1120,6 @@ bool get_delta_tail_cache( return true; } -std::mutex & delta_log_process_mutex() { - static std::mutex mutex; - return mutex; -} - bool delta_lock_path(const char * path, std::filesystem::path & out) { std::filesystem::path fs_path; if (!filesystem_path_from_utf8(path, fs_path)) { @@ -1143,13 +1140,45 @@ bool delta_lock_path(const char * path, std::filesystem::path & out) { return true; } +std::shared_ptr delta_log_process_mutex_for(const std::filesystem::path & lock_path) { + static std::mutex registry_mutex; + static std::unordered_map> registry; + + // POSIX advisory locks are process-owned, so same-process threads need an + // in-memory companion lock. Key it by path to avoid serializing every log. + const std::string key = lock_path.u8string(); + std::lock_guard guard(registry_mutex); + const auto found = registry.find(key); + if (found != registry.end()) { + std::shared_ptr mutex = found->second.lock(); + if (mutex != nullptr) { + return mutex; + } + registry.erase(found); + } + + for (auto it = registry.begin(); it != registry.end();) { + if (it->second.expired()) { + it = registry.erase(it); + } else { + ++it; + } + } + + std::shared_ptr mutex = std::make_shared(); + registry.emplace(key, mutex); + return mutex; +} + class DeltaLogLock { public: - explicit DeltaLogLock(const char * path) : process_lock(delta_log_process_mutex()) { + explicit DeltaLogLock(const char * path) { std::filesystem::path lock_path; if (!delta_lock_path(path, lock_path)) { return; } + process_mutex = delta_log_process_mutex_for(lock_path); + process_lock = std::unique_lock(*process_mutex); #ifdef _WIN32 file = CreateFileW( lock_path.wstring().c_str(), @@ -1212,6 +1241,7 @@ class DeltaLogLock { } private: + std::shared_ptr process_mutex; std::unique_lock process_lock; bool locked = false; #ifdef _WIN32 @@ -1558,7 +1588,9 @@ int write_empty_delta_log_unlocked(const ggml_vec_index & idx, const char * delt return fail_io(); } temp.path.clear(); - (void) fsync_parent_dir(delta_path); + if (!fsync_parent_dir(delta_path)) { + return GGML_VEC_INDEX_E_IO; + } return GGML_VEC_INDEX_OK; } catch (const std::bad_alloc &) { return GGML_VEC_INDEX_E_OOM; @@ -1767,6 +1799,9 @@ static int ggml_vec_index_add_unlocked( base_slot = idx->slot_to_id.size(); dim_sz = static_cast(idx->dim); const size_t n_sz = static_cast(n); + if (n_sz > kMaxIndexLen || active_count(*idx) > kMaxIndexLen - n_sz) { + return GGML_VEC_INDEX_E_INVALID_ARG; + } if (n_sz > std::numeric_limits::max() - base_slot) { return GGML_VEC_INDEX_E_INVALID_ARG; } @@ -2791,16 +2826,25 @@ int ggml_vec_index_search_ivf( return a.score > b.score; }); + std::vector selected_lists; + selected_lists.reserve(static_cast(probe_count)); size_t candidate_count = 0; - for (int probe = 0; probe < probe_count; ++probe) { - candidate_count += - idx->ivf_lists[static_cast(centroid_scores[probe].id)].size(); + for (const ScoreId & centroid : centroid_scores) { + const size_t list_id = static_cast(centroid.id); + const auto & list = idx->ivf_lists[list_id]; + if (list.empty()) { + continue; + } + selected_lists.push_back(list_id); + candidate_count += list.size(); + if (selected_lists.size() == static_cast(probe_count)) { + break; + } } std::vector candidate_slots; candidate_slots.reserve(candidate_count); - for (int probe = 0; probe < probe_count; ++probe) { - const auto & list = - idx->ivf_lists[static_cast(centroid_scores[probe].id)]; + for (size_t list_id : selected_lists) { + const auto & list = idx->ivf_lists[list_id]; candidate_slots.insert(candidate_slots.end(), list.begin(), list.end()); } search_one(*idx, query, k, scores, ids, &candidate_slots); @@ -2822,6 +2866,11 @@ static int ggml_vec_index_write_unlocked(ggml_vec_index_t * idx, const char * pa if (idx == nullptr || path == nullptr) { return GGML_VEC_INDEX_E_INVALID_ARG; } + if (idx->read_only_mmap && + !idx->mapped_source_path.empty() && + filesystem_paths_equal(idx->mapped_source_path.c_str(), path)) { + return GGML_VEC_INDEX_E_INVALID_ARG; + } if (active_count(*idx) > std::numeric_limits::max()) { return GGML_VEC_INDEX_E_INVALID_ARG; } @@ -2959,7 +3008,9 @@ static int ggml_vec_index_write_unlocked(ggml_vec_index_t * idx, const char * pa return fail_io(); } temp.path.clear(); - (void) fsync_parent_dir(path); + if (!fsync_parent_dir(path)) { + return GGML_VEC_INDEX_E_IO; + } return GGML_VEC_INDEX_OK; } catch (const std::bad_alloc &) { return GGML_VEC_INDEX_E_OOM; @@ -3102,6 +3153,9 @@ ggml_vec_index_t * ggml_vec_index_load(const char * path) { } const size_t dim_sz = static_cast(dim); const size_t n = static_cast(n_le); + if (n > kMaxIndexLen) { + return nullptr; + } if (n != 0 && dim_sz > std::numeric_limits::max() / n) { return nullptr; } @@ -3416,6 +3470,9 @@ ggml_vec_index_t * ggml_vec_index_load_mmap(const char * path) { } const size_t n = static_cast(n_le); + if (n > kMaxIndexLen) { + return nullptr; + } const size_t dim_sz = static_cast(dim_le); idx->slot_to_id.resize(n); idx->slot_active.assign(n, 1); @@ -3492,6 +3549,7 @@ ggml_vec_index_t * ggml_vec_index_load_mmap(const char * path) { } idx->read_only_mmap = true; + idx->mapped_source_path = path; idx->mapped_vector_bytes = static_cast(vectors_bytes_u64); idx->mapped_file = std::move(mapped); return idx.release(); @@ -3725,7 +3783,13 @@ int ggml_vec_index_compact_delta( if (write_status != GGML_VEC_INDEX_OK) { return write_status; } - return write_empty_delta_log_unlocked(*idx, delta_path); + const int delta_status = write_empty_delta_log_unlocked(*idx, delta_path); + if (delta_status != GGML_VEC_INDEX_OK) { + return delta_status; + } + idx->delta_log_rebase_pending = false; + idx->delta_log_rebase_crc = 0; + return GGML_VEC_INDEX_OK; } catch (const std::bad_alloc &) { return GGML_VEC_INDEX_E_OOM; } catch (...) { @@ -3743,7 +3807,8 @@ int ggml_vec_index_len(const ggml_vec_index_t * idx) { } try { std::shared_lock lock(idx->mutex); - return static_cast(active_count(*idx)); + const size_t n = active_count(*idx); + return n > kMaxIndexLen ? std::numeric_limits::max() : static_cast(n); } catch (...) { return 0; } diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 7fc14ac08083..35566699463e 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -155,18 +155,18 @@ if (LLAMA_TESTS_INSTALL) endif() llama_test(test-vector-index) -add_executable( - test-vector-index-faults - test-vector-index-faults.cpp - ${PROJECT_SOURCE_DIR}/ggml/src/ggml-vector-index.cpp -) -target_compile_features(test-vector-index-faults PRIVATE cxx_std_17) -target_compile_definitions(test-vector-index-faults PRIVATE GGML_VEC_INDEX_TEST_HOOKS) -target_include_directories(test-vector-index-faults PRIVATE ${PROJECT_SOURCE_DIR}/ggml/include) -target_link_libraries(test-vector-index-faults PRIVATE Threads::Threads) +add_library(test-vector-index-fault-hooks STATIC ${PROJECT_SOURCE_DIR}/ggml/src/ggml-vector-index.cpp) +target_compile_features(test-vector-index-fault-hooks PRIVATE cxx_std_17) +target_compile_definitions(test-vector-index-fault-hooks PRIVATE GGML_VEC_INDEX_TEST_HOOKS) +target_include_directories(test-vector-index-fault-hooks PUBLIC ${PROJECT_SOURCE_DIR}/ggml/include) +target_link_libraries(test-vector-index-fault-hooks PRIVATE Threads::Threads) if (NOT WIN32) - target_link_libraries(test-vector-index-faults PRIVATE m) + target_link_libraries(test-vector-index-fault-hooks PRIVATE m) endif() + +add_executable(test-vector-index-faults test-vector-index-faults.cpp) +target_compile_features(test-vector-index-faults PRIVATE cxx_std_17) +target_link_libraries(test-vector-index-faults PRIVATE test-vector-index-fault-hooks) llama_test(test-vector-index-faults) add_executable(bench-vector-index bench-vector-index.cpp) @@ -175,6 +175,7 @@ target_link_libraries(bench-vector-index PRIVATE ggml-vector-index) if (LLAMA_TESTS_INSTALL) install(TARGETS bench-vector-index RUNTIME) endif() +llama_test(bench-vector-index LABEL bench) if (LLAMA_LLGUIDANCE) llama_build_and_test(test-grammar-llguidance.cpp ARGS ${PROJECT_SOURCE_DIR}/models/ggml-vocab-llama-bpe.gguf) diff --git a/tests/bench-vector-index.cpp b/tests/bench-vector-index.cpp index 97491aae714c..113cbbf8c096 100644 --- a/tests/bench-vector-index.cpp +++ b/tests/bench-vector-index.cpp @@ -793,8 +793,10 @@ int main() { std::vector> allowlists; std::vector f32_filtered; std::vector q8_filtered; + std::vector q4_filtered; std::vector f32_prepared_filtered; std::vector q8_prepared_filtered; + std::vector q4_prepared_filtered; for (int requested : filter_sizes) { const int n_allowed = std::min(requested, cfg.n_vec); allowlists.push_back(make_allowlist(cfg.n_vec, n_allowed)); @@ -814,12 +816,23 @@ int main() { allowlists.back(), cfg.warmups, cfg.repeats)); + q4_filtered.push_back(run_filtered_search( + q4, + queries, + cfg.n_query, + cfg.k, + allowlists.back(), + cfg.warmups, + cfg.repeats)); ggml_vec_index_filter_t * f32_filter = ggml_vec_index_filter_create( f32, allowlists.back().data(), static_cast(allowlists.back().size())); ggml_vec_index_filter_t * q8_filter = ggml_vec_index_filter_create( q8, allowlists.back().data(), static_cast(allowlists.back().size())); + ggml_vec_index_filter_t * q4_filter = ggml_vec_index_filter_create( + q4, allowlists.back().data(), static_cast(allowlists.back().size())); CHECK(f32_filter != nullptr); CHECK(q8_filter != nullptr); + CHECK(q4_filter != nullptr); f32_prepared_filtered.push_back(run_prepared_filtered_search( f32, f32_filter, @@ -836,8 +849,17 @@ int main() { cfg.k, cfg.warmups, cfg.repeats)); + q4_prepared_filtered.push_back(run_prepared_filtered_search( + q4, + q4_filter, + queries, + cfg.n_query, + cfg.k, + cfg.warmups, + cfg.repeats)); ggml_vec_index_filter_free(f32_filter); ggml_vec_index_filter_free(q8_filter); + ggml_vec_index_filter_free(q4_filter); } const double f32_ivf_recall_at_k = @@ -981,19 +1003,23 @@ int main() { q4_ivf_recall_at_k); for (size_t i = 0; i < allowlists.size(); ++i) { std::printf( - " filtered latency: allowed=%zu f32=%.3f ms q8=%.3f ms f32/prepared=%.3f ms q8/prepared=%.3f ms\n", + " filtered latency: allowed=%zu f32=%.3f ms q8=%.3f ms q4=%.3f ms f32/prepared=%.3f ms q8/prepared=%.3f ms q4/prepared=%.3f ms\n", allowlists[i].size(), f32_filtered[i].ms, q8_filtered[i].ms, + q4_filtered[i].ms, f32_prepared_filtered[i].ms, - q8_prepared_filtered[i].ms); + q8_prepared_filtered[i].ms, + q4_prepared_filtered[i].ms); std::printf( - " filtered ratio: allowed=%zu f32/full=%.3f q8/full=%.3f f32/prep_speedup=%.3f q8/prep_speedup=%.3f\n", + " filtered ratio: allowed=%zu f32/full=%.3f q8/full=%.3f q4/full=%.3f f32/prep_speedup=%.3f q8/prep_speedup=%.3f q4/prep_speedup=%.3f\n", allowlists[i].size(), f32_filtered[i].ms / f32_res.ms, q8_filtered[i].ms / q8_res.ms, + q4_filtered[i].ms / q4_res.ms, f32_filtered[i].ms / f32_prepared_filtered[i].ms, - q8_filtered[i].ms / q8_prepared_filtered[i].ms); + q8_filtered[i].ms / q8_prepared_filtered[i].ms, + q4_filtered[i].ms / q4_prepared_filtered[i].ms); } std::printf( " delta load f32: snapshot=%.3f ms replay=%.3f ms compact=%.3f ms post_compact=%.3f ms\n", diff --git a/tests/test-vector-index-faults.cpp b/tests/test-vector-index-faults.cpp index f7afda2f068a..ec4230fed3a3 100644 --- a/tests/test-vector-index-faults.cpp +++ b/tests/test-vector-index-faults.cpp @@ -78,6 +78,94 @@ void expect_no_temp_siblings(const std::string & path) { } } +void test_quantized_logged_faults(int bit_width) { + constexpr int dim = 4; + const std::array base_vectors = { + 1.0f, 0.0f, 0.0f, 0.0f, + 0.0f, 1.0f, 0.0f, 0.0f, + }; + const std::array base_ids = { + static_cast(800 + bit_width), + static_cast(900 + bit_width), + }; + const std::array logged_vector = { 0.25f, -0.5f, 0.75f, -1.0f }; + const std::array extra_vector = { -0.125f, 0.375f, -0.625f, 0.875f }; + + const std::string suffix = std::to_string(bit_width); + const std::string snapshot_path = + (std::filesystem::temp_directory_path() / + ("ggml-vector-index-quant-fault-base-" + suffix + ".tvim")).string(); + const std::string delta_path = + (std::filesystem::temp_directory_path() / + ("ggml-vector-index-quant-fault-log-" + suffix + ".tvid")).string(); + std::filesystem::remove(snapshot_path); + std::filesystem::remove(delta_path); + std::filesystem::remove(delta_path + ".lock"); + + auto * idx = ggml_vec_index_create(dim, bit_width); + CHECK(idx != nullptr); + CHECK(ggml_vec_index_add( + idx, base_vectors.data(), static_cast(base_ids.size()), base_ids.data()) == + GGML_VEC_INDEX_OK); + CHECK(ggml_vec_index_write(idx, snapshot_path.c_str()) == GGML_VEC_INDEX_OK); + + const uint64_t failed_initial_id = static_cast(1000 + bit_width); + ggml_vec_index_test_set_write_fail_after(8); + CHECK(ggml_vec_index_add_logged( + idx, logged_vector.data(), 1, &failed_initial_id, delta_path.c_str()) == + GGML_VEC_INDEX_E_IO); + reset_fault_hooks(); + CHECK(ggml_vec_index_contains(idx, failed_initial_id) == 0); + CHECK(ggml_vec_index_len(idx) == 2); + if (std::filesystem::exists(delta_path)) { + CHECK(std::filesystem::file_size(delta_path) == 0); + } + + const uint64_t logged_id = static_cast(1100 + bit_width); + CHECK(ggml_vec_index_add_logged( + idx, logged_vector.data(), 1, &logged_id, delta_path.c_str()) == + GGML_VEC_INDEX_OK); + CHECK(ggml_vec_index_contains(idx, logged_id) == 1); + CHECK(ggml_vec_index_len(idx) == 3); + const std::vector old_delta = read_file_bytes(delta_path); + + const uint64_t failed_rollback_id = static_cast(1200 + bit_width); + ggml_vec_index_test_set_write_fail_after(8); + ggml_vec_index_test_set_truncate_fail(1); + CHECK(ggml_vec_index_add_logged( + idx, extra_vector.data(), 1, &failed_rollback_id, delta_path.c_str()) == + GGML_VEC_INDEX_E_INTERNAL); + reset_fault_hooks(); + CHECK(ggml_vec_index_contains(idx, failed_rollback_id) == 0); + CHECK(ggml_vec_index_len(idx) == 3); + CHECK(read_file_bytes(delta_path) == old_delta); + + ggml_vec_index_test_set_write_fail_after(8); + CHECK(ggml_vec_index_remove_logged(idx, logged_id, delta_path.c_str()) == + GGML_VEC_INDEX_E_IO); + reset_fault_hooks(); + CHECK(ggml_vec_index_contains(idx, logged_id) == 1); + CHECK(read_file_bytes(delta_path) == old_delta); + + CHECK(ggml_vec_index_remove_logged(idx, base_ids[0], delta_path.c_str()) == 1); + CHECK(ggml_vec_index_contains(idx, base_ids[0]) == 0); + CHECK(ggml_vec_index_len(idx) == 2); + + auto * replayed = ggml_vec_index_load_with_delta(snapshot_path.c_str(), delta_path.c_str()); + CHECK(replayed != nullptr); + CHECK(ggml_vec_index_bit_width(replayed) == bit_width); + CHECK(ggml_vec_index_len(replayed) == 2); + CHECK(ggml_vec_index_contains(replayed, base_ids[0]) == 0); + CHECK(ggml_vec_index_contains(replayed, base_ids[1]) == 1); + CHECK(ggml_vec_index_contains(replayed, logged_id) == 1); + + ggml_vec_index_free(replayed); + ggml_vec_index_free(idx); + std::filesystem::remove(snapshot_path); + std::filesystem::remove(delta_path); + std::filesystem::remove(delta_path + ".lock"); +} + } // namespace int main() { @@ -208,7 +296,7 @@ int main() { GGML_VEC_INDEX_OK); ggml_vec_index_test_set_parent_fsync_fail(1); CHECK(ggml_vec_index_write(parent_fsync_idx, parent_fsync_path.c_str()) == - GGML_VEC_INDEX_OK); + GGML_VEC_INDEX_E_IO); reset_fault_hooks(); CHECK(read_file_bytes(parent_fsync_path) != before_parent_fsync); @@ -391,9 +479,9 @@ int main() { CHECK(ggml_vec_index_compact_delta( compact_parent, compact_parent_snapshot_path.c_str(), - compact_parent_delta_path.c_str()) == GGML_VEC_INDEX_OK); + compact_parent_delta_path.c_str()) == GGML_VEC_INDEX_E_IO); reset_fault_hooks(); - CHECK(std::filesystem::file_size(compact_parent_delta_path) == 16); + CHECK(std::filesystem::file_size(compact_parent_delta_path) > 16); auto * compact_parent_replayed = ggml_vec_index_load_with_delta( compact_parent_snapshot_path.c_str(), compact_parent_delta_path.c_str()); CHECK(compact_parent_replayed != nullptr); @@ -447,6 +535,9 @@ int main() { std::filesystem::remove(cached_tail_delta_path); std::filesystem::remove(cached_tail_delta_path + ".lock"); + test_quantized_logged_faults(/*bit_width=*/8); + test_quantized_logged_faults(/*bit_width=*/4); + const std::string shared_snapshot_path = (std::filesystem::temp_directory_path() / "ggml-vector-index-shared-delta-base.tvim").string(); diff --git a/tests/test-vector-index.cpp b/tests/test-vector-index.cpp index 6a7e5fb32250..70db6f6f1408 100644 --- a/tests/test-vector-index.cpp +++ b/tests/test-vector-index.cpp @@ -198,12 +198,43 @@ void expect_corrupt_load_fails( } // namespace int main() { + CHECK(ggml_vec_index_create(0, /*bit_width=*/32) == nullptr); + CHECK(ggml_vec_index_create(-1, /*bit_width=*/32) == nullptr); + CHECK(ggml_vec_index_create(kDim, /*bit_width=*/16) == nullptr); + CHECK(ggml_vec_index_contains(nullptr, 123ULL) == 0); + CHECK(ggml_vec_index_len(nullptr) == 0); + CHECK(ggml_vec_index_dim(nullptr) == 0); + CHECK(ggml_vec_index_bit_width(nullptr) == 0); + auto * idx = ggml_vec_index_create(kDim, /*bit_width=*/32); CHECK(idx != nullptr); CHECK(ggml_vec_index_dim(idx) == kDim); CHECK(ggml_vec_index_len(idx) == 0); CHECK(ggml_vec_index_bit_width(idx) == 32); + // Zero-row adds are valid no-ops and must not create delta artifacts. + { + const std::array vector = { + 1.0f, 0.0f, 0.0f, 0.0f, + }; + const uint64_t id = 1234ULL; + CHECK(ggml_vec_index_add(idx, vector.data(), /*n=*/0, &id) + == GGML_VEC_INDEX_OK); + CHECK(ggml_vec_index_len(idx) == 0); + + const std::string zero_delta_path = + (std::filesystem::temp_directory_path() / + "ggml-vector-index-zero-add.tvid").string(); + std::filesystem::remove(zero_delta_path); + std::filesystem::remove(zero_delta_path + ".lock"); + CHECK(ggml_vec_index_add_logged( + idx, vector.data(), /*n=*/0, &id, zero_delta_path.c_str()) == + GGML_VEC_INDEX_OK); + CHECK(ggml_vec_index_len(idx) == 0); + CHECK(!std::filesystem::exists(zero_delta_path)); + std::filesystem::remove(zero_delta_path + ".lock"); + } + // Non-finite vectors are rejected without mutation. { const std::array bad_vector = { @@ -252,6 +283,50 @@ int main() { CHECK(ggml_vec_index_contains(idx, ids[0]) == 1); CHECK(ggml_vec_index_contains(idx, 999ULL) == 0); + // Zero-query searches are no-ops, while k=0 is invalid for every search mode. + { + std::array scores = { 123.0f, 456.0f, 789.0f, 101.0f }; + std::array out_ids = { 1, 2, 3, 4 }; + CHECK(ggml_vec_index_search( + idx, seeds[0].data(), /*n_q=*/0, /*k=*/1, + scores.data(), out_ids.data()) == GGML_VEC_INDEX_OK); + CHECK(scores[0] == 123.0f); + CHECK(out_ids[0] == 1); + CHECK(ggml_vec_index_search( + idx, seeds[0].data(), /*n_q=*/1, /*k=*/0, + scores.data(), out_ids.data()) == GGML_VEC_INDEX_E_INVALID_ARG); + + const std::array allowed = { ids[0], ids[2] }; + CHECK(ggml_vec_index_search_filtered( + idx, seeds[0].data(), /*n_q=*/0, /*k=*/1, + allowed.data(), static_cast(allowed.size()), + scores.data(), out_ids.data()) == GGML_VEC_INDEX_OK); + CHECK(ggml_vec_index_search_filtered( + idx, seeds[0].data(), /*n_q=*/1, /*k=*/0, + allowed.data(), static_cast(allowed.size()), + scores.data(), out_ids.data()) == GGML_VEC_INDEX_E_INVALID_ARG); + + ggml_vec_index_filter_t * filter = ggml_vec_index_filter_create( + idx, allowed.data(), static_cast(allowed.size())); + CHECK(filter != nullptr); + CHECK(ggml_vec_index_search_prepared_filtered( + idx, filter, seeds[0].data(), /*n_q=*/0, /*k=*/1, + scores.data(), out_ids.data()) == GGML_VEC_INDEX_OK); + CHECK(ggml_vec_index_search_prepared_filtered( + idx, filter, seeds[0].data(), /*n_q=*/1, /*k=*/0, + scores.data(), out_ids.data()) == GGML_VEC_INDEX_E_INVALID_ARG); + ggml_vec_index_filter_free(filter); + + CHECK(ggml_vec_index_build_ivf(idx, /*n_lists=*/2, /*n_iter=*/1) + == GGML_VEC_INDEX_OK); + CHECK(ggml_vec_index_search_ivf( + idx, seeds[0].data(), /*n_q=*/0, /*k=*/1, /*nprobe=*/1, + scores.data(), out_ids.data()) == GGML_VEC_INDEX_OK); + CHECK(ggml_vec_index_search_ivf( + idx, seeds[0].data(), /*n_q=*/1, /*k=*/0, /*nprobe=*/1, + scores.data(), out_ids.data()) == GGML_VEC_INDEX_E_INVALID_ARG); + } + // Non-finite queries are rejected before search. { const std::array bad_query = { @@ -336,6 +411,34 @@ int main() { ann, query.data(), 1, /*k=*/1, /*nprobe=*/16, ann_scores.data(), ann_ids.data()) == GGML_VEC_INDEX_OK); ggml_vec_index_free(ann); + + auto * empty_list_ann = ggml_vec_index_create(kDim, /*bit_width=*/32); + CHECK(empty_list_ann != nullptr); + const std::array empty_list_vecs = { + 3.0f, 3.0f, 0.0f, 0.0f, + -2.0f, 1.0f, 0.0f, 0.0f, + 0.0f, 1.0f, 0.0f, 0.0f, + }; + const std::array empty_list_ids = { 7101, 7102, 7103 }; + const std::array empty_list_query = { 1.0f, -3.0f, 0.0f, 0.0f }; + CHECK(ggml_vec_index_add( + empty_list_ann, + empty_list_vecs.data(), + static_cast(empty_list_ids.size()), + empty_list_ids.data()) == GGML_VEC_INDEX_OK); + CHECK(ggml_vec_index_build_ivf(empty_list_ann, /*n_lists=*/3, /*n_iter=*/1) + == GGML_VEC_INDEX_OK); + CHECK(ggml_vec_index_search_ivf( + empty_list_ann, + empty_list_query.data(), + 1, + /*k=*/1, + /*nprobe=*/1, + ann_scores.data(), + ann_ids.data()) == GGML_VEC_INDEX_OK); + CHECK(ann_ids[0] == empty_list_ids[2]); + CHECK(std::fabs(ann_scores[0] + 3.0f) < 1e-5f); + ggml_vec_index_free(empty_list_ann); } // Read-only APIs on one handle can run concurrently. @@ -420,6 +523,84 @@ int main() { ggml_vec_index_free(concurrent); } + // Mutations are serialized with readers on the same handle. + { + constexpr int n_rows = 16; + std::vector rows; + std::vector row_ids; + rows.reserve(static_cast(n_rows) * kDim); + row_ids.reserve(n_rows); + for (int row = 0; row < n_rows; ++row) { + const std::vector v = normalize({ + 1.0f, + static_cast((row % 3) - 1), + static_cast(((row + 2) % 5) - 2), + 0.5f, + }); + rows.insert(rows.end(), v.begin(), v.end()); + row_ids.push_back(static_cast(9000 + row)); + } + + auto * concurrent_mutation = ggml_vec_index_create(kDim, /*bit_width=*/32); + CHECK(concurrent_mutation != nullptr); + CHECK(ggml_vec_index_add( + concurrent_mutation, rows.data(), n_rows, row_ids.data()) == + GGML_VEC_INDEX_OK); + + std::atomic ready{ 0 }; + std::atomic start{ false }; + std::atomic done{ false }; + std::atomic failures{ 0 }; + std::vector readers; + for (int t = 0; t < 4; ++t) { + readers.emplace_back([&, t]() { + const float * query = + rows.data() + static_cast(t % n_rows) * kDim; + std::array scores{}; + std::array out_ids{}; + ready.fetch_add(1); + while (!start.load()) { + std::this_thread::yield(); + } + while (!done.load()) { + if (ggml_vec_index_search( + concurrent_mutation, query, 1, /*k=*/3, + scores.data(), out_ids.data()) != GGML_VEC_INDEX_OK) { + failures.fetch_add(1); + } + if (ggml_vec_index_len(concurrent_mutation) < n_rows) { + failures.fetch_add(1); + } + } + }); + } + + while (ready.load() != 4) { + std::this_thread::yield(); + } + start.store(true); + for (int iter = 0; iter < 100; ++iter) { + const std::vector v = normalize({ + 0.25f, + static_cast((iter % 7) - 3), + 1.0f, + -0.5f, + }); + const uint64_t id = static_cast(10000 + iter); + CHECK(ggml_vec_index_add(concurrent_mutation, v.data(), 1, &id) + == GGML_VEC_INDEX_OK); + CHECK(ggml_vec_index_remove(concurrent_mutation, id) == 1); + } + done.store(true); + for (std::thread & reader : readers) { + reader.join(); + } + CHECK(failures.load() == 0); + CHECK(ggml_vec_index_len(concurrent_mutation) == n_rows); + + ggml_vec_index_free(concurrent_mutation); + } + // Top-1 of querying with each unit vector should retrieve itself with // score very close to 1.0 (full f32, no quantization noise). { @@ -883,6 +1064,26 @@ int main() { ggml_vec_index_free(replayed_after_old_log_append); ggml_vec_index_free(compacted_with_old_log); + write_file_bytes(delta_path, pre_compact_delta); + auto * recompact_from_old_log = ggml_vec_index_load_with_delta( + snapshot_path.c_str(), delta_path.c_str()); + CHECK(recompact_from_old_log != nullptr); + CHECK(ggml_vec_index_compact_delta( + recompact_from_old_log, snapshot_path.c_str(), delta_path.c_str()) == + GGML_VEC_INDEX_OK); + CHECK(std::filesystem::file_size(delta_path) == 16); + const uint64_t post_recompact_id = (1ULL << 41) + 11ULL; + CHECK(ggml_vec_index_add_logged( + recompact_from_old_log, seeds[3].data(), 1, &post_recompact_id, delta_path.c_str()) == + GGML_VEC_INDEX_OK); + auto * replayed_after_recompact = ggml_vec_index_load_with_delta( + snapshot_path.c_str(), delta_path.c_str()); + CHECK(replayed_after_recompact != nullptr); + CHECK(ggml_vec_index_len(replayed_after_recompact) == 3); + CHECK(ggml_vec_index_contains(replayed_after_recompact, post_recompact_id) == 1); + ggml_vec_index_free(replayed_after_recompact); + ggml_vec_index_free(recompact_from_old_log); + CHECK(ggml_vec_index_compact_delta( base, snapshot_path.c_str(), delta_path.c_str()) == GGML_VEC_INDEX_OK); const uint64_t post_compact_id = (1ULL << 41) + 8ULL; @@ -1064,6 +1265,8 @@ int main() { CHECK(ggml_vec_index_remove(mapped, mmap_ids[0]) == GGML_VEC_INDEX_E_INVALID_ARG); CHECK(ggml_vec_index_compact(mapped) == GGML_VEC_INDEX_E_INVALID_ARG); + CHECK(ggml_vec_index_write(mapped, mmap_path.c_str()) + == GGML_VEC_INDEX_E_INVALID_ARG); CHECK(ggml_vec_index_write(mapped, mmap_copy_path.c_str()) == GGML_VEC_INDEX_OK); auto * copied = ggml_vec_index_load(mmap_copy_path.c_str()); CHECK(copied != nullptr); From d18858e1e55246f585495664b2108a155169c3ff Mon Sep 17 00:00:00 2001 From: Nidhin Date: Fri, 17 Jul 2026 00:21:34 +0530 Subject: [PATCH 13/13] Fix vector-index delta recovery race Keep logged mutation completion checks under the delta log lock, and clarify vector-index API docs around SIMD dispatch, scoring, IVF probes, and mmap endian requirements. --- ggml/CMakeLists.txt | 1 + ggml/include/ggml-vector-index.h | 19 ++++++---- ggml/src/ggml-vector-index.cpp | 63 ++++++++++++++++++-------------- 3 files changed, 49 insertions(+), 34 deletions(-) diff --git a/ggml/CMakeLists.txt b/ggml/CMakeLists.txt index c1e09b8282d4..e9e38adfce7f 100644 --- a/ggml/CMakeLists.txt +++ b/ggml/CMakeLists.txt @@ -449,6 +449,7 @@ if (MSVC) configure_msvc_target(ggml-base) configure_msvc_target(ggml) + configure_msvc_target(ggml-vector-index) configure_msvc_target(ggml-cpu) configure_msvc_target(ggml-cpu-x64) configure_msvc_target(ggml-cpu-sse42) diff --git a/ggml/include/ggml-vector-index.h b/ggml/include/ggml-vector-index.h index 66109617997e..05258223f63a 100644 --- a/ggml/include/ggml-vector-index.h +++ b/ggml/include/ggml-vector-index.h @@ -5,7 +5,9 @@ // This public C API supports full f32 storage (`bit_width=32`), q8 storage // (`bit_width=8`), and packed q4 // storage (`bit_width=4`) with CPU search directly against quantized codes. -// q8 uses NEON/AVX2 when available; q4 uses NEON when available. +// q8 uses NEON when available; on x86 it uses AVX2 when compiled with AVX2, +// and GCC/Clang builds can runtime-dispatch AVX2 from a non-AVX2 baseline. +// q4 uses NEON when available. // // Threading: read-only APIs on the same handle can run concurrently. Mutations, // persistence writes, compaction, and IVF builds are serialized with reads and @@ -137,10 +139,11 @@ GGML_API int ggml_vec_index_build_ivf( // against the index (does not mutate state). // // Score semantics: dot product. For f32 storage this is a full-precision dot -// product. For q4/q8 storage, the query remains f32 and each indexed component -// is scored as `q_code * per_vector_scale` without expanding the stored matrix -// back to f32. Callers that want cosine similarity must L2-normalize their -// vectors before insert AND before query; the index does NOT normalize +// product. For q4/q8 storage, the query remains f32 and the dot product is +// computed against dequantized indexed components: +// `query[i] * (q_code * per_vector_scale)`, without expanding the stored +// matrix back to f32. Callers that want cosine similarity must L2-normalize +// their vectors before insert AND before query; the index does NOT normalize // internally. All query components must be finite. GGML_API int ggml_vec_index_search( const ggml_vec_index_t * idx, @@ -186,8 +189,9 @@ GGML_API int ggml_vec_index_search_prepared_filtered( // IVF-flat ANN top-k search. `ggml_vec_index_build_ivf` must have been called // after the most recent mutation. `nprobe` controls how many centroid lists are -// searched; higher values improve recall and lower the latency win. If nprobe -// is greater than the number of built lists, all lists are searched. +// searched; higher values improve recall and lower the latency win. `nprobe` +// must be >= 1. If nprobe is greater than the number of built lists, all lists +// are searched. GGML_API int ggml_vec_index_search_ivf( const ggml_vec_index_t * idx, const float * queries, @@ -212,6 +216,7 @@ GGML_API ggml_vec_index_t * ggml_vec_index_load(const char * path); // Mutating APIs return GGML_VEC_INDEX_E_INVALID_ARG on mmap-backed handles. // `ggml_vec_index_write` can snapshot mmap-backed handles, but callers must // write to a different path than the mapped source file. +// Requires a little-endian host; use `ggml_vec_index_load` on other hosts. // Returns NULL on failure or unsupported file format. GGML_API ggml_vec_index_t * ggml_vec_index_load_mmap(const char * path); diff --git a/ggml/src/ggml-vector-index.cpp b/ggml/src/ggml-vector-index.cpp index 1c216eeeb798..3c5db81b5470 100644 --- a/ggml/src/ggml-vector-index.cpp +++ b/ggml/src/ggml-vector-index.cpp @@ -1440,7 +1440,12 @@ bool delta_log_ends_at_state( } } -int append_delta_record( +struct DeltaAppendResult { + int status = GGML_VEC_INDEX_OK; + bool record_complete = false; +}; + +DeltaAppendResult append_delta_record( ggml_vec_index & idx, const char * delta_path, uint8_t op, @@ -1449,33 +1454,33 @@ int append_delta_record( uint32_t state_crc, const std::vector & payload) { if (delta_path == nullptr) { - return GGML_VEC_INDEX_E_INVALID_ARG; + return { GGML_VEC_INDEX_E_INVALID_ARG, false }; } DeltaLogLock delta_lock(delta_path); if (!delta_lock.ok()) { - return GGML_VEC_INDEX_E_IO; + return { GGML_VEC_INDEX_E_IO, false }; } uint64_t old_size = 0; uint32_t existing_base_crc = 0; if (!validate_delta_header(delta_path, idx, old_size, existing_base_crc)) { - return GGML_VEC_INDEX_E_IO; + return { GGML_VEC_INDEX_E_IO, false }; } if (old_size != 0) { uint32_t tail_crc = 0; uint64_t complete_size = 0; if (!get_delta_tail_cache(idx, delta_path, old_size, tail_crc, complete_size) && !inspect_delta_log_tail(delta_path, idx, tail_crc, complete_size)) { - return GGML_VEC_INDEX_E_IO; + return { GGML_VEC_INDEX_E_IO, false }; } if (tail_crc != base_crc_for_new_log) { invalidate_delta_tail_cache(idx); - return GGML_VEC_INDEX_E_IO; + return { GGML_VEC_INDEX_E_IO, false }; } if (complete_size != old_size) { if (!truncate_file_to(delta_path, complete_size)) { invalidate_delta_tail_cache(idx); - return GGML_VEC_INDEX_E_INTERNAL; + return { GGML_VEC_INDEX_E_INTERNAL, false }; } old_size = complete_size; } @@ -1486,7 +1491,7 @@ int append_delta_record( // compacting the snapshot but before replacing the old delta log). if (!truncate_file_to(delta_path, 0)) { invalidate_delta_tail_cache(idx); - return GGML_VEC_INDEX_E_INTERNAL; + return { GGML_VEC_INDEX_E_INTERNAL, false }; } old_size = 0; idx.delta_log_rebase_pending = false; @@ -1498,7 +1503,7 @@ int append_delta_record( std::FILE * f = nullptr; if (!open_append_file(delta_path, &f)) { - return GGML_VEC_INDEX_E_IO; + return { GGML_VEC_INDEX_E_IO, false }; } auto close_file = [&]() { if (f != nullptr) { @@ -1506,11 +1511,16 @@ int append_delta_record( f = nullptr; } }; - auto fail_io = [&]() { + auto fail_io = [&]() -> DeltaAppendResult { close_file(); const bool truncated = truncate_file_to(delta_path, old_size); + const bool record_complete = + !truncated && delta_log_ends_at_state(delta_path, idx, state_crc); invalidate_delta_tail_cache(idx); - return truncated ? GGML_VEC_INDEX_E_IO : GGML_VEC_INDEX_E_INTERNAL; + return { + truncated ? GGML_VEC_INDEX_E_IO : GGML_VEC_INDEX_E_INTERNAL, + record_complete, + }; }; if (old_size == 0) { @@ -1542,8 +1552,13 @@ int append_delta_record( f = nullptr; if (close_result != 0 || !fsync_parent_dir(delta_path)) { const bool truncated = truncate_file_to(delta_path, old_size); + const bool record_complete = + !truncated && delta_log_ends_at_state(delta_path, idx, state_crc); invalidate_delta_tail_cache(idx); - return truncated ? GGML_VEC_INDEX_E_IO : GGML_VEC_INDEX_E_INTERNAL; + return { + truncated ? GGML_VEC_INDEX_E_IO : GGML_VEC_INDEX_E_INTERNAL, + record_complete, + }; } const uint64_t written_size = old_size + @@ -1551,7 +1566,7 @@ int append_delta_record( kTvidRecordHeaderSize + static_cast(payload.size()); set_delta_tail_cache(idx, delta_path, written_size, state_crc); - return GGML_VEC_INDEX_OK; + return { GGML_VEC_INDEX_OK, true }; } int write_empty_delta_log_unlocked(const ggml_vec_index & idx, const char * delta_path) { @@ -2084,7 +2099,7 @@ int ggml_vec_index_add_logged( added = true; const uint32_t added_state_crc = index_state_crc32c(*idx); - const int append_status = append_delta_record( + const DeltaAppendResult append_result = append_delta_record( *idx, delta_path, kTvidOpAdd, @@ -2092,11 +2107,8 @@ int ggml_vec_index_add_logged( base_crc, added_state_crc, payload); - if (append_status != GGML_VEC_INDEX_OK) { - const bool record_is_complete = - append_status == GGML_VEC_INDEX_E_INTERNAL && - delta_log_ends_at_state(delta_path, *idx, added_state_crc); - if (record_is_complete) { + if (append_result.status != GGML_VEC_INDEX_OK) { + if (append_result.record_complete) { ++idx->generation; invalidate_ivf(*idx); added = false; @@ -2105,7 +2117,7 @@ int ggml_vec_index_add_logged( rollback_appended_slots_unlocked(idx, base_slot, ids, n); } added = false; - return append_status; + return append_result.status; } ++idx->generation; invalidate_ivf(*idx); @@ -2145,7 +2157,7 @@ int ggml_vec_index_remove_logged( const std::vector payload = build_remove_delta_payload(id); const uint32_t base_crc = index_state_crc32c(*idx); const uint32_t post_remove_crc = index_state_crc32c_after_remove(*idx, id); - const int append_status = append_delta_record( + const DeltaAppendResult append_result = append_delta_record( *idx, delta_path, kTvidOpRemove, @@ -2153,14 +2165,11 @@ int ggml_vec_index_remove_logged( base_crc, post_remove_crc, payload); - if (append_status != GGML_VEC_INDEX_OK) { - const bool record_is_complete = - append_status == GGML_VEC_INDEX_E_INTERNAL && - delta_log_ends_at_state(delta_path, *idx, post_remove_crc); - if (record_is_complete) { + if (append_result.status != GGML_VEC_INDEX_OK) { + if (append_result.record_complete) { return ggml_vec_index_remove_unlocked(idx, id); } - return append_status; + return append_result.status; } return ggml_vec_index_remove_unlocked(idx, id); } catch (const std::bad_alloc &) {