diff --git a/ggml/CMakeLists.txt b/ggml/CMakeLists.txt index 1f534ad7230d..e9e38adfce7f 100644 --- a/ggml/CMakeLists.txt +++ b/ggml/CMakeLists.txt @@ -343,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( @@ -448,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/cmake/ggml-config.cmake.in b/ggml/cmake/ggml-config.cmake.in index 75f824225bd8..cae4feab2973 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) @@ -201,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 new file mode 100644 index 000000000000..05258223f63a --- /dev/null +++ b/ggml/include/ggml-vector-index.h @@ -0,0 +1,308 @@ +#pragma once +// +// ggml-vector-index: vector-index C API. +// +// 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 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 +// 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. + +#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 + +// Opaque handle to a vector index instance. +struct ggml_vec_index; +typedef struct ggml_vec_index ggml_vec_index_t; + +// 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; + +// 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_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` 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); + +// 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). +// All vector components must be finite. UINT64_MAX is reserved for search +// 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, + int n, + const uint64_t * ids); + +// 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`. +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); + +// 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. 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. +// 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 +// 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: dot product. For f32 storage this is a full-precision dot +// 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, + const float * queries, + int n_q, + int k, + 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. 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, + 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. `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, + 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, + const char * path); + +// 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); + +// 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 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); + +// 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. 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); + +// File format (.tvim version 2, all little-endian): +// +// offset size field +// ------ ----- ------------------------------------------------------- +// 0 4 magic = "TVPI" (bytes 0x54, 0x56, 0x50, 0x49) +// 4 1 version = 2 +// 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 (0 for packed q4, 1 for q8, 4 for f32) +// 28 4 reserved (zero) +// 32 ... qparams: +// - f32: empty +// - 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 +// ... 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. q8 uses symmetric per-vector quantization: +// scale = max(abs(v)) / 127, code = round(v / scale) clamped to [-127, 127]. +// 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 +} +#endif diff --git a/ggml/src/CMakeLists.txt b/ggml/src/CMakeLists.txt index 4abc1e301dfe..7876be7ec7e4 100644 --- a/ggml/src/CMakeLists.txt +++ b/ggml/src/CMakeLists.txt @@ -261,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 PRIVATE Threads::Threads) + function(ggml_add_backend_library backend) if (GGML_BACKEND_DL) add_library(${backend} MODULE ${ARGN}) @@ -492,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() @@ -514,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 new file mode 100644 index 000000000000..3c5db81b5470 --- /dev/null +++ b/ggml/src/ggml-vector-index.cpp @@ -0,0 +1,3848 @@ +// 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 +// std::unordered_map for lookup and a parallel vector for +// 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 +// AVX2 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 +#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 + +#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 +#include +#include +#else +#ifndef NOMINMAX +#define NOMINMAX +#endif +#include +#include +#include +#include +#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); +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 + +namespace { + +constexpr uint8_t kTvimMagic[4] = { 'T', 'V', 'P', 'I' }; +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; +constexpr size_t kMaxIndexLen = static_cast(std::numeric_limits::max()); + +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 }; +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(); + 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; +} + +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) { +#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); + 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; +} + +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)); + return bits; +} + +float u32_to_float(uint32_t bits) { + float v; + std::memcpy(&v, &bits, sizeof(v)); + return v; +} + +bool write_u32_le(std::FILE * f, uint32_t v) { + uint8_t bytes[4]; + put_u32_le(bytes, v); + return write_bytes(f, bytes, sizeof(bytes)); +} + +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; +} + +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_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 (!write_bytes(f, bytes, 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 (!write_bytes(f, bytes, sizeof(bytes))) { + return false; + } + crc = crc32c_update(crc, bytes, sizeof(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; + } +}; + +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])) { + 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 (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, 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; +} + +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) { + 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; +} + +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) { + 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 GGML_VEC_INDEX_TEST_HOOKS + if (g_test_parent_fsync_fail.load()) { + return false; + } +#endif +#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 + +#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); +} + +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); +} + +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 + +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 { + mutable std::shared_mutex mutex; + + 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; + 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; + std::string mapped_source_path; + 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; + + // Flat row-major q8 storage for bit_width=8 plus one scale per vector. + 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; + std::vector slot_active; + size_t n_active = 0; + + // 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 { + const ggml_vec_index_t * owner = nullptr; + uint64_t owner_cookie = 0; + 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_q4(idx) ? kStorageQ4 : (is_q8(idx) ? kStorageQ8 : kStorageF32); +} + +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); +} + +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) { + 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 = round_nearest_even(scaled); + q = std::max(-127, std::min(127, q)); + dst[i] = static_cast(q); + } +} + +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 = 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]; + 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(active_count(idx))); + if (is_q4(idx)) { + 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])); + } + } + 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)) { + 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])); + } + } + 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 { + 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 (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; +} + +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 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(active_count(idx) - 1)); + + if (is_q4(idx)) { + const size_t row_bytes = q4_row_bytes(dim_sz); + 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 < 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)) { + 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 < 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 { + 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(data[slot * dim_sz + i])); + } + } + } + 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; +} + +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 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(); +} + +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; +} + +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; +} + +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) { + 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(), + 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::shared_ptr process_mutex; + 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; + 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 { +#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; + } + 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) { +#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; + } + 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; +} + +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; + } +} + +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, + 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, false }; + } + DeltaLogLock delta_lock(delta_path); + if (!delta_lock.ok()) { + 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, 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, false }; + } + if (tail_crc != base_crc_for_new_log) { + invalidate_delta_tail_cache(idx); + 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, false }; + } + 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)) { + invalidate_delta_tail_cache(idx); + return { GGML_VEC_INDEX_E_INTERNAL, false }; + } + 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(); + + std::FILE * f = nullptr; + if (!open_append_file(delta_path, &f)) { + return { GGML_VEC_INDEX_E_IO, false }; + } + auto close_file = [&]() { + if (f != nullptr) { + std::fclose(f); + f = nullptr; + } + }; + 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, + record_complete, + }; + }; + + 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)) { + 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, + record_complete, + }; + } + 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, true }; +} + +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; + } + 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; + } +} + +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, + 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 (!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; + } + 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 +// --------------------------------------------------------------------------- + +ggml_vec_index_t * ggml_vec_index_create(int dim, int bit_width) { + try { + if (dim <= 0) { + return nullptr; + } + if (!is_supported_bit_width(bit_width)) { + return nullptr; + } + auto * idx = new (std::nothrow) ggml_vec_index(); + if (idx == nullptr) { + return nullptr; + } + 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; + } +} + +void ggml_vec_index_free(ggml_vec_index_t * idx) { + delete 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, + bool finalize) { + + size_t base_slot = 0; + size_t dim_sz = 0; + bool resized = false; + + auto rollback = [&]() noexcept { + if (idx == nullptr || !resized) { + return; + } + rollback_appended_slots_unlocked(idx, base_slot, ids, n); + }; + + 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; + } + if (n == 0) { + return GGML_VEC_INDEX_OK; + } + + // 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) { + 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; + } + if (!batch_ids.insert(ids[i]).second) { + 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 > 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; + } + 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; + } + if (!all_finite(vectors, n_sz * dim_sz)) { + return GGML_VEC_INDEX_E_INVALID_ARG; + } + + resized = true; + 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 { + 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); + + 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_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, + 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->slot_active[slot] = 1; + test_maybe_throw_bad_alloc(); + idx->id_to_slot.emplace(ids[i], slot); + } + idx->n_active += n_sz; + if (finalize) { + ++idx->generation; + invalidate_ivf(*idx); + } + } catch (const std::bad_alloc &) { + rollback(); + return GGML_VEC_INDEX_E_OOM; + } catch (...) { + rollback(); + return GGML_VEC_INDEX_E_INTERNAL; + } + return GGML_VEC_INDEX_OK; +} + +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, true); + } 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; + } + if (!is_valid_id(id)) { + 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; + 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; + } +} + +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, + row_bytes * sizeof(uint8_t)); + 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, + dim_sz * sizeof(int8_t)); + 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, + 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); + } + idx->data.swap(new_data); + } + + 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 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_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_compact_unlocked(idx); + } 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; + size_t base_slot = 0; + 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 (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; + } + 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); + 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 DeltaAppendResult append_result = append_delta_record( + *idx, + delta_path, + kTvidOpAdd, + static_cast(n), + base_crc, + added_state_crc, + payload); + if (append_result.status != GGML_VEC_INDEX_OK) { + if (append_result.record_complete) { + ++idx->generation; + invalidate_ivf(*idx); + added = false; + return GGML_VEC_INDEX_OK; + } else { + rollback_appended_slots_unlocked(idx, base_slot, ids, n); + } + added = false; + return append_result.status; + } + ++idx->generation; + invalidate_ivf(*idx); + added = false; + return GGML_VEC_INDEX_OK; + } catch (const std::bad_alloc &) { + if (added) { + rollback_appended_slots_unlocked(idx, base_slot, ids, n); + } + return GGML_VEC_INDEX_E_OOM; + } catch (...) { + if (added) { + rollback_appended_slots_unlocked(idx, base_slot, ids, n); + } + 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 (!is_valid_id(id)) { + return GGML_VEC_INDEX_E_INVALID_ARG; + } + if (idx->read_only_mmap) { + return GGML_VEC_INDEX_E_INVALID_ARG; + } + 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 DeltaAppendResult append_result = append_delta_record( + *idx, + delta_path, + kTvidOpRemove, + 1, + base_crc, + post_remove_crc, + payload); + if (append_result.status != GGML_VEC_INDEX_OK) { + if (append_result.record_complete) { + return ggml_vec_index_remove_unlocked(idx, id); + } + return append_result.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; + } + 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*/) { + // 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; +} + +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; +} + +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) { +#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; +} + +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 + +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 +} + +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, + q4_data_ptr(idx) + slot * q4_row_bytes(static_cast(dim)), + idx.q4_scale[slot], + dim) : + is_q8(idx) ? + dot_q8( + query, + q8_data_ptr(idx) + slot * static_cast(dim), + idx.q8_scale[slot], + dim) : + dot( + query, + f32_data_ptr(idx) + 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 = + 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]; + 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 = 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; + } + } else { + std::memcpy( + dst, + f32_data_ptr(idx) + 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) { + 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] }); + } 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() && slot_is_active(idx, it->second)) { + 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 size_t n_live = active_count(*idx); + const int dim = idx->dim; + 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_live)); + 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)); + 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 = 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); + } + + 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 : 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; + 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 : 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); + } + + 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->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) { + 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; + } + + 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 GGML_VEC_INDEX_OK; +} + +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); +} + +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); +} + +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->owner = idx; + owned->owner_cookie = idx->filter_cookie; + 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; + } +} + +void ggml_vec_index_filter_free(ggml_vec_index_filter_t * filter) { + delete filter; +} + +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); +} + +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) { + + if (idx == nullptr || queries == nullptr || + out_scores == nullptr || out_ids == nullptr) { + return GGML_VEC_INDEX_E_INVALID_ARG; + } + if (n_q < 0 || k <= 0 || nprobe <= 0) { + 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; + } + 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) { + 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; + }); + + std::vector selected_lists; + selected_lists.reserve(static_cast(probe_count)); + size_t candidate_count = 0; + 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 (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); + } + return GGML_VEC_INDEX_OK; + } catch (const std::bad_alloc &) { + return GGML_VEC_INDEX_E_OOM; + } catch (...) { + return GGML_VEC_INDEX_E_INTERNAL; + } +} + +// --------------------------------------------------------------------------- +// Persistence +// --------------------------------------------------------------------------- + +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; + } + 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; + } + 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_slots > std::numeric_limits::max() / dim_sz) { + return GGML_VEC_INDEX_E_INTERNAL; + } + 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; + } + + 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); + temp.stream = nullptr; + } + remove_temp_file(temp); + return GGML_VEC_INDEX_E_IO; + }; + std::FILE * f = temp.stream; + + // 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] = storage_kind(*idx); + header[7] = kFlagCRC32C; + const uint32_t dim_le = static_cast(idx->dim); + 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); + 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))) { + 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_quantized(*idx)) { + const std::vector & scales = is_q4(*idx) ? idx->q4_scale : idx->q8_scale; + 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)) { + 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 { + 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); + } + } + } else { + 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 (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(); + } + } + + 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; + } catch (...) { + return GGML_VEC_INDEX_E_INTERNAL; + } +} + +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) { + return nullptr; + } + 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), kTvimV1HeaderSize); + if (!f || f.gcount() != static_cast(kTvimV1HeaderSize)) { + return nullptr; + } + if (std::memcmp(header, kTvimMagic, 4) != 0) { + return nullptr; + } + + 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 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 == 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( + 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( + 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 > kMaxIndexLen) { + return nullptr; + } + if (n != 0 && dim_sz > std::numeric_limits::max() / n) { + return nullptr; + } + + test_maybe_throw_bad_alloc(); + 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 { + 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; + 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_quantized(*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; + } + } + 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_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) : + read_u32_le(f, bits); + if (!read_ok) { + return nullptr; + } + scale = u32_to_float(bits); + if (!std::isfinite(scale) || scale <= 0.0f) { + 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()); + } + } + 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; + } + } + } 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 { + 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; + } + } + } + + for (uint64_t & id : idx->slot_to_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 (!is_valid_id(id)) { + 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; + } + } + + 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; + } +} + +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); + 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); + 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)); + if (!is_valid_id(id)) { + return nullptr; + } + idx->slot_to_id[slot] = id; + if (!idx->id_to_slot.emplace(id, slot).second) { + return nullptr; + } + } + + 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(); + } catch (...) { + return nullptr; + } +} + +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()); + if (!is_valid_id(id)) { + return false; + } + 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; + } + + 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 + +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; + } + 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; + } + 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 (...) { + return GGML_VEC_INDEX_E_INTERNAL; + } +} + +// --------------------------------------------------------------------------- +// Stats +// --------------------------------------------------------------------------- + +int ggml_vec_index_len(const ggml_vec_index_t * idx) { + if (idx == nullptr) { + return 0; + } + try { + std::shared_lock lock(idx->mutex); + const size_t n = active_count(*idx); + return n > kMaxIndexLen ? std::numeric_limits::max() : static_cast(n); + } catch (...) { + return 0; + } +} + +int ggml_vec_index_dim(const ggml_vec_index_t * idx) { + 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) { + if (idx == nullptr) { + return 0; + } + try { + std::shared_lock lock(idx->mutex); + return idx->bit_width; + } catch (...) { + return 0; + } +} diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index b61f1581cfb1..35566699463e 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -144,6 +144,39 @@ 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 Threads::Threads) +if (LLAMA_TESTS_INSTALL) + install(TARGETS test-vector-index RUNTIME) +endif() +llama_test(test-vector-index) + +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-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) +target_compile_features(bench-vector-index PRIVATE cxx_std_17) +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) endif () diff --git a/tests/bench-vector-index.cpp b/tests/bench-vector-index.cpp new file mode 100644 index 000000000000..113cbbf8c096 --- /dev/null +++ b/tests/bench-vector-index.cpp @@ -0,0 +1,1114 @@ +// 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 +#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; + int delta_ops = 256; + int ivf_lists = 64; + int ivf_iters = 4; + int ivf_nprobe = 4; + int delete_stride = 2; +}; + +struct TimedSearch { + double ms = 0.0; + std::vector scores; + 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; +}; + +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) { + 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"; +#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 +} + +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); + + 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; +} + +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, + 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( + idx, + queries.data(), + n_query, + k, + result.scores.data(), + result.ids.data()) == GGML_VEC_INDEX_OK); + }); + return result; +} + +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); + }); + 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; +} + +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; +} + +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 = round_nearest_even(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); + 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); +} + +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; + 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; +}; + +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, + 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() { + 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); + 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 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)); + 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)); + 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, + 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)); + 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 = + 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); + 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); + + 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 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) + + 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); + 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 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(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( + " 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 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, + q4_prepared_filtered[i].ms); + std::printf( + " 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, + 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", + 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( + " 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); + 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); + return 0; +} diff --git a/tests/test-vector-index-faults.cpp b/tests/test-vector-index-faults.cpp new file mode 100644 index 000000000000..ec4230fed3a3 --- /dev/null +++ b/tests/test-vector-index-faults.cpp @@ -0,0 +1,618 @@ +#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); +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 { + +#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); + 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) { + 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); + } +} + +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() { + 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(); + } + + { + 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(); + 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); + + { + 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_E_IO); + 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(); + 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); + } + + 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); + + 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) == 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()) == 1); + 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 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_E_IO); + 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"); + + 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(); + 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); + 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 new file mode 100644 index 000000000000..70db6f6f1408 --- /dev/null +++ b/tests/test-vector-index.cpp @@ -0,0 +1,1820 @@ +// 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 +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#ifndef _WIN32 +#include +#endif + +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; +} + +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()); + + 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 : round_nearest_even(vector[i] / scale); + code = std::max(-127, std::min(127, code)); + acc += query[i] * (static_cast(code) * scale); + } + return acc; +} + +float q4_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 / 7.0f; + + float acc = 0.0f; + for (size_t i = 0; i < vector.size(); ++i) { + int code = max_abs == 0.0f ? + 0 : round_nearest_even(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()); + 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_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))); + } +} + +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() { + 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 = { + 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); + } + + // 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 bugs when this codepath is called from bindings. + 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); + + // 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 = { + 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] }; + 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); + } + + // 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); + + 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. + { + 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); + } + + // 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). + { + 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); + } + } + + // 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); + + 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( + 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); + 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); +#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 + + 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()); + 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); + + // 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 = + (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 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); + 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); + 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); + + 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); + + 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); + + 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; + 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); + } + + // 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. + { + 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_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)); + + 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); + } + } + + // 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_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); + 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; + 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); + } + + // q4 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); + } + + // 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 = + (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); + 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 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 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(); + 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]; + } + }); + 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) { + 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; +}