diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index 7a70941..aacda68 100644 --- a/.github/workflows/main.yml +++ b/.github/workflows/main.yml @@ -192,6 +192,10 @@ jobs: if: contains(matrix.name, 'linux') || matrix.name == 'windows' || ( matrix.name == 'macos' && matrix.arch != 'x86_64' ) run: ${{ matrix.name == 'linux-musl' && matrix.arch == 'arm64' && 'docker exec alpine' || '' }} make test ${{ matrix.make && matrix.make || ''}} + - name: unit test sqlite-vector + if: contains(matrix.name, 'linux') || ( matrix.name == 'macos' && matrix.arch != 'x86_64' ) + run: ${{ matrix.name == 'linux-musl' && matrix.arch == 'arm64' && 'docker exec alpine' || '' }} make unittest ${{ matrix.make && matrix.make || ''}} + - uses: actions/upload-artifact@v4.6.2 if: always() with: diff --git a/API.md b/API.md index 972659e..37f0ca2 100644 --- a/API.md +++ b/API.md @@ -39,7 +39,9 @@ Returns the active backend used for vector computation. This indicates the SIMD * `CPU` – Generic fallback * `SSE2` – SIMD on Intel/AMD * `AVX2` – Advanced SIMD on modern x86 CPUs +* `AVX512` – Wide SIMD on supported x86 CPUs * `NEON` – SIMD on ARM (e.g., mobile) +* `RVV` – SIMD on supported RISC-V CPUs **Example:** @@ -50,6 +52,22 @@ SELECT vector_backend(); --- +## `vector_turboquant_backend()` + +**Returns:** `TEXT` + +**Description:** +Returns the active backend used by TurboQuant lookup-table scans. This is useful when validating that TurboQuant is using the expected SIMD path on a target runtime. + +**Example:** + +```sql +SELECT vector_turboquant_backend(); +-- e.g., 'NEON' +``` + +--- + ## `vector_init(table, column, options)` **Returns:** `NULL` @@ -101,7 +119,7 @@ SELECT vector_init('documents', 'embedding', 'dimension=384,type=FLOAT32,distanc **Returns:** `INTEGER` **Description:** -Returns the total number of succesfully quantized rows. +Returns the total number of successfully quantized rows. Performs quantization on the specified table and column. This precomputes internal data structures to support fast approximate nearest neighbor (ANN) search. Read more about quantization [here](https://github.com/sqliteai/sqlite-vector/blob/main/QUANTIZATION.md). @@ -117,13 +135,16 @@ If a quantization already exists for the specified table and column, it is repla **Available options:** * `max_memory`: Max memory to use for quantization (default: 30MB) -* `qtype`: Quantization type: `UINT8`, `INT8` or `1BIT` +* `qtype`: Quantization type: `UINT8`, `INT8`, `1BIT`, `TURBO`, `TURBO2`, `TURBO3`, or `TURBO4` +* `qbits`: TurboQuant bit width (`2`, `3`, or `4`). Defaults to `4` when `qtype=TURBO`. **Example:** ```sql SELECT vector_quantize('documents', 'embedding', 'max_memory=50MB'); SELECT vector_quantize('documents', 'embedding', 'qtype=BIT'); +SELECT vector_quantize('documents', 'embedding', 'qtype=TURBO,qbits=4'); +SELECT vector_quantize('documents', 'embedding', 'qtype=TURBO2'); ``` --- @@ -281,9 +302,9 @@ You **must run `vector_quantize()`** before using `vector_quantize_scan()` and w **Performance Highlights:** -* Handles **1M vectors** of dimension 768 in a few milliseconds. -* Uses **<50MB** of RAM. -* Achieves **>0.95 recall**. +* Supports compact SIMD 2-, 3-, and 4-bit TurboQuant scans for high-dimensional vectors. +* `qbits=2` minimizes memory; `qbits=4` usually gives the better recall/speed balance. +* Recall depends on the dataset, distance function, bit width, and `k`; validate it against `vector_full_scan()` for your workload. **Examples:** diff --git a/Package.swift b/Package.swift index 06ff107..cf86b8c 100644 --- a/Package.swift +++ b/Package.swift @@ -14,8 +14,8 @@ let package = Package( targets: [ .binaryTarget( name: "vectorBinary", - url: "https://github.com/sqliteai/sqlite-vector/releases/download/0.9.95/vector-apple-xcframework-0.9.95.zip", - checksum: "db4a3a733ff6d719c18a4692b5cbab80327daff004d8199cb53a198cb5072e85" + url: "https://github.com/sqliteai/sqlite-vector/releases/download/1.0.0/vector-apple-xcframework-1.0.0.zip", + checksum: "26962a269a0e7f5da3fca421c7c43a9d7542367ccadb898f8e4fc057ad18a18c" ), .target( name: "vector", diff --git a/QUANTIZATION.md b/QUANTIZATION.md index 7e7000d..0f938cd 100644 --- a/QUANTIZATION.md +++ b/QUANTIZATION.md @@ -16,12 +16,23 @@ This can result in a **4×–5× speedup** on nearest neighbor queries while kee #### What is Quantization? -Quantization compresses high-dimensional float vectors (e.g., `FLOAT32`) into compact representations using lower-precision formats (e.g., `UINT8`). This drastically reduces the size of the data—often by a factor of 4 to 8—making it practical to load large datasets entirely in memory, even on edge devices. +Quantization compresses high-dimensional vectors into compact representations using lower-precision formats such as `UINT8`, `INT8`, `1BIT`, and TurboQuant (`TURBO`). TurboQuant supports 2-, 3-, and 4-bit scalar codes plus one scale per vector, which is useful for large edge-oriented datasets where raw `FLOAT32` storage is too large. + +```sql +-- Default quantization. +SELECT vector_quantize('my_table', 'my_column'); + +-- TurboQuant, 4 bits per dimension. +SELECT vector_quantize('my_table', 'my_column', 'qtype=TURBO,qbits=4'); + +-- TurboQuant shorthand for 2 bits per dimension. +SELECT vector_quantize('my_table', 'my_column', 'qtype=TURBO2'); +``` #### Why is it Important? -* **Faster Searches**: With preloaded quantized vectors, distance computations are up to 5× faster. -* **Lower Memory Footprint**: Quantized vectors use significantly less RAM, allowing millions of vectors to fit in memory. +* **Faster Searches**: With quantized vectors, distance computations can be several times faster than brute force. +* **Lower Memory Footprint**: Quantized vectors use significantly less RAM and disk than raw vectors, allowing millions of vectors to fit in constrained environments. * **Edge-ready**: The reduced size and in-memory access make this ideal for mobile, embedded, and on-device AI applications. #### Estimate Memory Usage @@ -33,10 +44,11 @@ SELECT vector_quantize_memory('my_table', 'my_column'); ``` This gives you an approximate number of bytes needed to load the quantized vectors into memory. +For TurboQuant, the scan representation is approximately `rows * (8 + 4 + ceil(dimension * qbits / 8))` bytes before allocator overhead and SQLite page/cache effects. You can skip `vector_quantize_preload()` to keep the quantized data on disk and reduce resident memory, at the cost of more SQLite page reads during scans. #### Accuracy You Can Trust -Despite the compression, our quantization algorithms are finely tuned to maintain high accuracy. You can expect **recall rates greater than 0.95**, ensuring that approximate searches closely match exact results in quality. +Despite the compression, quantization is approximate. Recall depends on the dataset, vector dimensionality, distance function, bit width, and requested `k`. TurboQuant scans use SIMD lookup-table kernels where available. Prefer `qbits=4` when recall is the priority and `qbits=2` when memory is the primary constraint. #### Measuring Recall in SQLite-Vector @@ -74,3 +86,5 @@ SELECT Where `?1` is the input vector (as a BLOB) and `?2` is the number of nearest neighbors `k`. This query compares exact and quantized results and computes the recall ratio, helping you validate the quality of quantized search. + +For a reproducible real-dataset run, use `test/recall_turboquant_real.py`. It downloads the Fashion-MNIST ANN-Benchmarks HDF5 dataset, loads a configurable subset into SQLite, and reports recall@k for TurboQuant 2/3/4-bit against `vector_full_scan()`. diff --git a/README.md b/README.md index b89608b..63e4f7e 100644 --- a/README.md +++ b/README.md @@ -1,11 +1,52 @@ +
+ + SQLite AI + + +

SQLite-Vector

+

Production-grade vector search inside SQLite.
+ Exact search, SIMD distance kernels, and SIMD 2/3/4-bit TurboQuant scans — runs anywhere SQLite runs: mobile, browser, edge, server.

+ +

+ Free managed instance → · + Docs · + Website · + Blog +

+ +

+ Data: + Vector · + Sync · + Columnar · + JS +
+ AI: + AI · + Agent · + Memory · + MCP +
+

+
+ +
+ +> **Building RAG or semantic search?** SQLite-Vector ships as an extension you can drop into any SQLite app. Need it managed with sync and auth? **[SQLite Cloud free tier](https://dashboard.sqlitecloud.io/auth/sign-in)** gives you 512 MB and 20 connections, no credit card. + +--- + # SQLite Vector -**SQLite Vector** is a cross-platform, ultra-efficient SQLite extension that brings vector search capabilities to your embedded database. It works seamlessly on **iOS, Android, Windows, Linux, and macOS**, using just **30MB of memory** by default. With support for **Float32, Float16, BFloat16, Int8, UInt8 and 1Bit**, and **highly optimized distance functions**, it's the ideal solution for **Edge AI** applications. +**SQLite Vector** is a cross-platform, ultra-efficient SQLite extension that brings vector search capabilities to your embedded database. It works seamlessly on **iOS, Android, Windows, Linux, and macOS**, using just **30MB of memory** by default. With support for **Float32, Float16, BFloat16, Int8, UInt8, 1Bit, and TurboQuant 2/3/4-bit quantization**, plus **highly optimized distance functions**, it's the ideal solution for **Edge AI** applications. + +SQLite-Vector includes **TurboQuant**, a compact data-oblivious vector quantizer inspired by the Google Research paper [TurboQuant: Online Vector Quantization with Near-Optimal Distortion Rate](https://arxiv.org/abs/2504.19874). It stores each vector as low-bit scalar codes plus one scale value, then scores directly from SIMD lookup-table kernels without reconstructing full vectors. ## Highlights * **No virtual tables required** – store vectors directly as `BLOB`s in ordinary tables * **Blazing fast** – optimized C implementation with SIMD acceleration +* **TurboQuant support** – SIMD 2-, 3-, and 4-bit quantization scans with `qtype=TURBO` * **Low memory footprint** – defaults to just 30MB of RAM usage * **Zero preindexing needed** – no long preprocessing or index-building phases * **Works offline** – perfect for on-device, privacy-preserving AI workloads @@ -21,6 +62,7 @@ | Doesn't need preindexing | ✅ | ❌ (can take hours for large datasets) | | Doesn't need external server | ✅ | ❌ (often needs Redis/FAISS/Weaviate/etc.) | | Memory-efficient | ✅ | ❌ | +| TurboQuant low-bit scanning | ✅ | ❌ | | Easy to use SQL | ✅ | ❌ (often complex JOINs, subqueries) | | Offline/Edge ready | ✅ | ❌ | | Cross-platform | ✅ | ❌ | @@ -80,6 +122,9 @@ SELECT vector_init('images', 'embedding', 'type=FLOAT32,dimension=384'); -- Quantize vector SELECT vector_quantize('images', 'embedding'); +-- Or use TurboQuant for compact 2/3/4-bit quantization +SELECT vector_quantize('images', 'embedding', 'qtype=TURBO,qbits=4'); + -- Optional preload quantized version in memory (for a 4x/5x speedup) SELECT vector_quantize_preload('images', 'embedding'); @@ -96,6 +141,46 @@ SELECT e.id, v.distance FROM images AS e LIMIT 10; ``` +## TurboQuant Benchmark and Recall + +TurboQuant can be selected with `qtype=TURBO,qbits=N`, where `N` is `2`, `3`, or `4`. Shorthand aliases are also available: `TURBO2`, `TURBO3`, and `TURBO4`. + +```sql +-- Highest recall TurboQuant mode currently recommended as the default +SELECT vector_quantize('images', 'embedding', 'qtype=TURBO,qbits=4'); + +-- Smaller edge-oriented representation +SELECT vector_quantize('images', 'embedding', 'qtype=TURBO2'); +``` + +The following benchmark compares `vector_full_scan()` brute force against `vector_quantize_scan()` using TurboQuant on a synthetic dataset of **1,000,000 vectors**, **768 dimensions**, **DOT** distance, **k=10**, and **5 queries**. The database was file-backed, with raw vectors stored in SQLite and quantized data preloaded for the scan. Recall is measured as overlap with exact brute-force top-10 results. These numbers were measured on macOS ARM64 using the NEON backend; timings vary by CPU, storage, cache settings, and allocator behavior. + +| Mode | Quantized storage | Max RSS | Peak memory footprint | Full scan / query | TurboQuant / query | Speedup | Recall@10 | +| --- | ---: | ---: | ---: | ---: | ---: | ---: | ---: | +| TurboQuant 4-bit | 396 MB | ~488 MB | ~487 MB | 3248 ms | 218 ms | 14.92x | 0.84 | +| TurboQuant 3-bit | 300 MB | ~394 MB | ~393 MB | 1727 ms | 188 ms | 9.19x | 0.74 | +| TurboQuant 2-bit | 204 MB | ~310 MB | ~310 MB | 3265 ms | 85 ms | 38.27x | 0.48 | + +For comparison, the raw `FLOAT32` vectors alone are about **3.07 GB** for 1M x 768 before SQLite row/page overhead. TurboQuant 4-bit reduces the scan representation to about **13%** of that raw vector payload, TurboQuant 3-bit to about **10%**, and TurboQuant 2-bit to about **7%**. Actual resident memory depends on whether the database is in-memory or file-backed, SQLite cache settings, preloading, page cache behavior, and the host allocator. + +The TurboQuant scan backend can be checked separately from the regular distance backend: + +```sql +SELECT vector_backend(), vector_turboquant_backend(); +``` + +For edge deployments, `vector_quantize_memory(table, column)` estimates the quantized scan representation. TurboQuant stores each row as `rowid + scale + packed_codes`, roughly `rows * (8 + 4 + ceil(dim * qbits / 8))` bytes before allocator and SQLite cache overhead. The synthetic benchmark in `test/benchmark_turboquant.c` also supports `PRELOAD=0` to compare the lower-RAM, non-preloaded path. + +Real-dataset recall can be reproduced with `test/recall_turboquant_real.py`, which downloads Fashion-MNIST in the ANN-Benchmarks HDF5 format and compares TurboQuant against `vector_full_scan()` using L2 distance. Example run on macOS ARM64/NEON with 10,000 base vectors, 50 queries, and k=10: + +| Mode | Quantized storage | Full scan / query | TurboQuant / query | Speedup | Recall@10 | +| --- | ---: | ---: | ---: | ---: | ---: | +| TurboQuant 4-bit | 4.04 MB | 16.32 ms | 4.80 ms | 3.40x | 0.948 | +| TurboQuant 3-bit | 3.06 MB | 16.32 ms | 8.28 ms | 1.97x | 0.868 | +| TurboQuant 2-bit | 2.08 MB | 16.32 ms | 1.86 ms | 8.78x | 0.596 | + +`qbits=4` is the recommended starting point when recall matters. `qbits=2` is useful for tighter edge memory budgets, but should be validated on the target embeddings because recall can drop significantly depending on the dataset. + ### Swift Package You can [add this repository as a package dependency to your Swift project](https://developer.apple.com/documentation/xcode/adding-package-dependencies-to-your-app#Add-a-package-dependency). After adding the package, you'll need to set up SQLite with extension loading by following steps 4 and 5 of [this guide](https://github.com/sqliteai/sqlite-extensions-guide/blob/main/platforms/ios.md#4-set-up-sqlite-with-extension-loading). @@ -257,19 +342,32 @@ Free Use in Open-Source Projects: You may use, copy, distribute, and prepare der --- -## Part of the SQLite AI Ecosystem -This project is part of the **SQLite AI** ecosystem, a collection of extensions that bring modern AI capabilities to the world’s most widely deployed database. The goal is to make SQLite the default data and inference engine for Edge AI applications. +## ☁️ Hosted version + +Don't want to run it yourself? **[SQLite Cloud](https://sqlite.ai)** is the managed version of SQLite-Vector and the rest of the stack — with sync, backups, auth, edge functions, and multi-region support included. + +[**Start free →**](https://dashboard.sqlitecloud.io/auth/sign-in) + +--- + +## Part of the SQLite AI stack + +SQLite-Vector is one piece of a larger ecosystem that turns SQLite into a runtime for intelligent, distributed data: + +**Data layer** +- [**sqlite-vector**](https://github.com/sqliteai/sqlite-vector) — ANN vector search inside SQLite *(you are here)* +- [sqlite-sync](https://github.com/sqliteai/sqlite-sync) — Offline-first CRDT sync across devices +- [sqlite-columnar](https://github.com/sqliteai/sqlite-columnar) — Column-oriented analytics for OLAP queries +- [sqlite-js](https://github.com/sqliteai/sqlite-js) — Custom SQLite functions written in JavaScript -Other projects in the ecosystem include: +**AI layer** +- [sqlite-ai](https://github.com/sqliteai/sqlite-ai) — On-device LLM inference and embeddings +- [sqlite-agent](https://github.com/sqliteai/sqlite-agent) — Autonomous AI agents running inside SQLite +- [sqlite-memory](https://github.com/sqliteai/sqlite-memory) — Persistent, searchable memory for agents +- [sqlite-mcp](https://github.com/sqliteai/sqlite-mcp) — Call MCP tools directly from SQL queries -- **[SQLite-AI](https://github.com/sqliteai/sqlite-ai)** — On-device inference and embedding generation directly inside SQLite. -- **[SQLite-Memory](https://github.com/sqliteai/sqlite-memory)** — Markdown-based AI agent memory with semantic search. -- **[SQLite-Vector](https://github.com/sqliteai/sqlite-vector)** — Ultra-efficient vector search for embeddings stored as BLOBs in standard SQLite tables. -- **[SQLite-Sync](https://github.com/sqliteai/sqlite-sync)** — Local-first CRDT-based synchronization for seamless, conflict-free data sync and real-time collaboration across devices. -- **[SQLite-Agent](https://github.com/sqliteai/sqlite-agent)** — Run autonomous AI agents directly from within SQLite databases. -- **[SQLite-MCP](https://github.com/sqliteai/sqlite-mcp)** — Connect SQLite databases to MCP servers and invoke their tools. -- **[SQLite-JS](https://github.com/sqliteai/sqlite-js)** — Create custom SQLite functions using JavaScript. -- **[Liteparser](https://github.com/sqliteai/liteparser)** — A highly efficient and fully compliant SQLite SQL parser. +**Managed platform** +- [SQLite Cloud](https://sqlite.ai) — Hosted SQLite with sync, auth, edge functions, and analytics. [Free tier →](https://dashboard.sqlitecloud.io/auth/sign-in) -Learn more at **[SQLite AI](https://sqlite.ai)**. +Built by [SQLite AI](https://sqlite.ai). Questions? [Contact us](https://sqlite.ai/support). diff --git a/src/distance-avx2.c b/src/distance-avx2.c index c581480..8e0da29 100644 --- a/src/distance-avx2.c +++ b/src/distance-avx2.c @@ -15,6 +15,8 @@ extern distance_function_t dispatch_distance_table[VECTOR_DISTANCE_MAX][VECTOR_TYPE_MAX]; extern const char *distance_backend_name; +extern turbo_lut_dot_function_t turbo_lut_dot_function; +extern const char *turbo_lut_backend_name; #define _mm256_abs_ps(x) _mm256_andnot_ps(_mm256_set1_ps(-0.0f), (x)) @@ -996,6 +998,55 @@ float bit1_distance_hamming_avx2 (const void *v1, const void *v2, int n) { return (float)distance; } +static inline uint16_t turbo_lut3_index_avx2 (const uint8_t *packed, int row, int packed_bytes) { + size_t bit_pos = (size_t)row * 12u; + size_t byte_pos = bit_pos / 8u; + int shift = (int)(bit_pos % 8u); + uint32_t word = 0; + if ((int)byte_pos < packed_bytes) word |= packed[byte_pos]; + if ((int)byte_pos + 1 < packed_bytes) word |= (uint32_t)packed[byte_pos + 1] << 8; + return (uint16_t)((word >> shift) & 0x0fffu); +} + +float turbo_lut_dot_avx2 (const uint8_t *packed, float scale, const float *query_lut, int lut_rows, int bits, int packed_bytes) { + __m256 acc = _mm256_setzero_ps(); + const __m256i lane = _mm256_setr_epi32(0, 1, 2, 3, 4, 5, 6, 7); + int r = 0; + if (bits == 3) { + const __m256i stride = _mm256_set1_epi32(4096); + for (; r + 7 < lut_rows; r += 8) { + int idx[8]; + for (int i = 0; i < 8; ++i) idx[i] = turbo_lut3_index_avx2(packed, r + i, packed_bytes); + __m256i codes = _mm256_loadu_si256((const __m256i *)idx); + __m256i rows = _mm256_add_epi32(_mm256_set1_epi32(r), lane); + __m256i indices = _mm256_add_epi32(_mm256_mullo_epi32(rows, stride), codes); + __m256 vals = _mm256_i32gather_ps(query_lut, indices, 4); + acc = _mm256_add_ps(acc, vals); + } + } else { + const __m256i stride = _mm256_set1_epi32(256); + for (; r + 7 < lut_rows; r += 8) { + __m128i codes8 = _mm_loadl_epi64((const __m128i *)(packed + r)); + __m256i codes = _mm256_cvtepu8_epi32(codes8); + __m256i rows = _mm256_add_epi32(_mm256_set1_epi32(r), lane); + __m256i indices = _mm256_add_epi32(_mm256_mullo_epi32(rows, stride), codes); + __m256 vals = _mm256_i32gather_ps(query_lut, indices, 4); + acc = _mm256_add_ps(acc, vals); + } + } + + float partial[8]; + _mm256_storeu_ps(partial, acc); + float dot = partial[0] + partial[1] + partial[2] + partial[3] + + partial[4] + partial[5] + partial[6] + partial[7]; + if (bits == 3) { + for (; r < lut_rows; ++r) dot += query_lut[(size_t)r * 4096u + turbo_lut3_index_avx2(packed, r, packed_bytes)]; + } else { + for (; r < lut_rows; ++r) dot += query_lut[(size_t)r * 256u + packed[r]]; + } + return dot * scale; +} + #endif // MARK: - @@ -1035,5 +1086,7 @@ void init_distance_functions_avx2 (void) { dispatch_distance_table[VECTOR_DISTANCE_HAMMING][VECTOR_TYPE_BIT] = bit1_distance_hamming_avx2; distance_backend_name = "AVX2"; + turbo_lut_dot_function = turbo_lut_dot_avx2; + turbo_lut_backend_name = "AVX2"; #endif } diff --git a/src/distance-avx2.h b/src/distance-avx2.h index 20cc5a4..aa7c39c 100644 --- a/src/distance-avx2.h +++ b/src/distance-avx2.h @@ -8,8 +8,10 @@ #ifndef __VECTOR_DISTANCE_AVX2__ #define __VECTOR_DISTANCE_AVX2__ +#include #include void init_distance_functions_avx2 (void); +float turbo_lut_dot_avx2 (const uint8_t *packed, float scale, const float *query_lut, int lut_rows, int bits, int packed_bytes); #endif diff --git a/src/distance-avx512.c b/src/distance-avx512.c index cd2fdb4..11b7d2f 100644 --- a/src/distance-avx512.c +++ b/src/distance-avx512.c @@ -16,6 +16,8 @@ extern distance_function_t dispatch_distance_table[VECTOR_DISTANCE_MAX][VECTOR_TYPE_MAX]; extern const char *distance_backend_name; +extern turbo_lut_dot_function_t turbo_lut_dot_function; +extern const char *turbo_lut_backend_name; // Abs for f32 (AVX512F has native abs) #define _mm512_abs_ps(x) _mm512_abs_ps(x) @@ -917,6 +919,52 @@ static float bit1_distance_hamming_avx512(const void *v1, const void *v2, int n) return (float)distance; } +static inline uint16_t turbo_lut3_index_avx512 (const uint8_t *packed, int row, int packed_bytes) { + size_t bit_pos = (size_t)row * 12u; + size_t byte_pos = bit_pos / 8u; + int shift = (int)(bit_pos % 8u); + uint32_t word = 0; + if ((int)byte_pos < packed_bytes) word |= packed[byte_pos]; + if ((int)byte_pos + 1 < packed_bytes) word |= (uint32_t)packed[byte_pos + 1] << 8; + return (uint16_t)((word >> shift) & 0x0fffu); +} + +float turbo_lut_dot_avx512 (const uint8_t *packed, float scale, const float *query_lut, int lut_rows, int bits, int packed_bytes) { + __m512 acc = _mm512_setzero_ps(); + const __m512i lane = _mm512_setr_epi32(0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15); + int r = 0; + if (bits == 3) { + const __m512i stride = _mm512_set1_epi32(4096); + for (; r + 15 < lut_rows; r += 16) { + int idx[16]; + for (int i = 0; i < 16; ++i) idx[i] = turbo_lut3_index_avx512(packed, r + i, packed_bytes); + __m512i codes = _mm512_loadu_si512((const void *)idx); + __m512i rows = _mm512_add_epi32(_mm512_set1_epi32(r), lane); + __m512i indices = _mm512_add_epi32(_mm512_mullo_epi32(rows, stride), codes); + __m512 vals = _mm512_i32gather_ps(indices, query_lut, 4); + acc = _mm512_add_ps(acc, vals); + } + } else { + const __m512i stride = _mm512_set1_epi32(256); + for (; r + 15 < lut_rows; r += 16) { + __m128i codes8 = _mm_loadu_si128((const __m128i *)(packed + r)); + __m512i codes = _mm512_cvtepu8_epi32(codes8); + __m512i rows = _mm512_add_epi32(_mm512_set1_epi32(r), lane); + __m512i indices = _mm512_add_epi32(_mm512_mullo_epi32(rows, stride), codes); + __m512 vals = _mm512_i32gather_ps(indices, query_lut, 4); + acc = _mm512_add_ps(acc, vals); + } + } + + float dot = _mm512_reduce_add_ps(acc); + if (bits == 3) { + for (; r < lut_rows; ++r) dot += query_lut[(size_t)r * 4096u + turbo_lut3_index_avx512(packed, r, packed_bytes)]; + } else { + for (; r < lut_rows; ++r) dot += query_lut[(size_t)r * 256u + packed[r]]; + } + return dot * scale; +} + #endif // MARK: - @@ -956,5 +1004,7 @@ void init_distance_functions_avx512(void) { dispatch_distance_table[VECTOR_DISTANCE_HAMMING][VECTOR_TYPE_BIT] = bit1_distance_hamming_avx512; distance_backend_name = "AVX512"; + turbo_lut_dot_function = turbo_lut_dot_avx512; + turbo_lut_backend_name = "AVX512"; #endif } diff --git a/src/distance-avx512.h b/src/distance-avx512.h index 265e164..d8bf661 100644 --- a/src/distance-avx512.h +++ b/src/distance-avx512.h @@ -8,8 +8,10 @@ #ifndef __VECTOR_DISTANCE_AVX512__ #define __VECTOR_DISTANCE_AVX512__ +#include #include void init_distance_functions_avx512 (void); +float turbo_lut_dot_avx512 (const uint8_t *packed, float scale, const float *query_lut, int lut_rows, int bits, int packed_bytes); #endif diff --git a/src/distance-cpu.c b/src/distance-cpu.c index 7d05ee2..ba0a358 100644 --- a/src/distance-cpu.c +++ b/src/distance-cpu.c @@ -21,6 +21,8 @@ const char *distance_backend_name = "CPU"; distance_function_t dispatch_distance_table[VECTOR_DISTANCE_MAX][VECTOR_TYPE_MAX] = {0}; +const char *turbo_lut_backend_name = "CPU"; +turbo_lut_dot_function_t turbo_lut_dot_function = NULL; #define LASSQ_UPDATE(ad_) do { \ double _ad = (ad_); \ @@ -864,6 +866,31 @@ float bit1_distance_hamming_cpu (const void *v1, const void *v2, int n) { // MARK: - +static inline uint16_t turbo_lut3_index_cpu (const uint8_t *packed, int row, int packed_bytes) { + size_t bit_pos = (size_t)row * 12u; + size_t byte_pos = bit_pos / 8u; + int shift = (int)(bit_pos % 8u); + uint32_t word = 0; + if ((int)byte_pos < packed_bytes) word |= packed[byte_pos]; + if ((int)byte_pos + 1 < packed_bytes) word |= (uint32_t)packed[byte_pos + 1] << 8; + return (uint16_t)((word >> shift) & 0x0fffu); +} + +float turbo_lut_dot_cpu (const uint8_t *packed, float scale, const float *query_lut, int lut_rows, int bits, int packed_bytes) { + double dot = 0.0; + if (bits == 3) { + for (int r = 0; r < lut_rows; ++r) { + dot += (double)query_lut[(size_t)r * 4096u + turbo_lut3_index_cpu(packed, r, packed_bytes)]; + } + } else { + (void)packed_bytes; + for (int r = 0; r < lut_rows; ++r) { + dot += (double)query_lut[(size_t)r * 256u + packed[r]]; + } + } + return (float)(dot * (double)scale); +} + void init_cpu_functions (void) { distance_function_t cpu_table[VECTOR_DISTANCE_MAX][VECTOR_TYPE_MAX] = { [VECTOR_DISTANCE_L2] = { @@ -907,6 +934,8 @@ void init_cpu_functions (void) { }; memcpy(dispatch_distance_table, cpu_table, sizeof(cpu_table)); + turbo_lut_dot_function = turbo_lut_dot_cpu; + turbo_lut_backend_name = "CPU"; } void init_distance_functions (bool force_cpu) { @@ -933,4 +962,3 @@ void init_distance_functions (bool force_cpu) { } #endif } - diff --git a/src/distance-cpu.h b/src/distance-cpu.h index 653109b..82214d2 100644 --- a/src/distance-cpu.h +++ b/src/distance-cpu.h @@ -47,7 +47,8 @@ typedef enum { VECTOR_QUANT_AUTO = 0, VECTOR_QUANT_U8BIT = 1, VECTOR_QUANT_S8BIT = 2, - VECTOR_QUANT_1BIT = 3 + VECTOR_QUANT_1BIT = 3, + VECTOR_QUANT_TURBO = 4 } vector_qtype; typedef enum { @@ -61,10 +62,14 @@ typedef enum { #define VECTOR_DISTANCE_MAX 7 typedef float (*distance_function_t)(const void *v1, const void *v2, int n); +typedef float (*turbo_lut_dot_function_t)(const uint8_t *packed, float scale, const float *query_lut, int lut_rows, int bits, int packed_bytes); // ENTRYPOINT void init_distance_functions (bool force_cpu); +extern turbo_lut_dot_function_t turbo_lut_dot_function; +extern const char *turbo_lut_backend_name; + // MARK: - FLOAT16/BFLOAT16 - // typedef uint16_t bfloat16_t; // don't typedef to bfloat16_t to avoid mix with ’s native bfloat16_t diff --git a/src/distance-neon.c b/src/distance-neon.c index b5cd6b9..29653a9 100644 --- a/src/distance-neon.c +++ b/src/distance-neon.c @@ -22,6 +22,8 @@ extern distance_function_t dispatch_distance_table[VECTOR_DISTANCE_MAX][VECTOR_TYPE_MAX]; extern const char *distance_backend_name; +extern turbo_lut_dot_function_t turbo_lut_dot_function; +extern const char *turbo_lut_backend_name; // Helper function for 32-bit ARM: vmaxv_u16 is not available in ARMv7 NEON #ifdef _ARM32BIT_ @@ -1271,6 +1273,57 @@ float bit1_distance_hamming_neon (const void *v1, const void *v2, int n) { return (float)distance; } +static inline uint16_t turbo_lut3_index_neon (const uint8_t *packed, int row, int packed_bytes) { + size_t bit_pos = (size_t)row * 12u; + size_t byte_pos = bit_pos / 8u; + int shift = (int)(bit_pos % 8u); + uint32_t word = 0; + if ((int)byte_pos < packed_bytes) word |= packed[byte_pos]; + if ((int)byte_pos + 1 < packed_bytes) word |= (uint32_t)packed[byte_pos + 1] << 8; + return (uint16_t)((word >> shift) & 0x0fffu); +} + +float turbo_lut_dot_neon (const uint8_t *packed, float scale, const float *query_lut, int lut_rows, int bits, int packed_bytes) { + float32x4_t acc = vdupq_n_f32(0.0f); + int r = 0; + if (bits == 3) { + for (; r + 3 < lut_rows; r += 4) { + float tmp[4] = { + query_lut[(size_t)(r + 0) * 4096u + turbo_lut3_index_neon(packed, r + 0, packed_bytes)], + query_lut[(size_t)(r + 1) * 4096u + turbo_lut3_index_neon(packed, r + 1, packed_bytes)], + query_lut[(size_t)(r + 2) * 4096u + turbo_lut3_index_neon(packed, r + 2, packed_bytes)], + query_lut[(size_t)(r + 3) * 4096u + turbo_lut3_index_neon(packed, r + 3, packed_bytes)] + }; + acc = vaddq_f32(acc, vld1q_f32(tmp)); + } + } else { + for (; r + 3 < lut_rows; r += 4) { + float tmp[4] = { + query_lut[(size_t)(r + 0) * 256u + packed[r + 0]], + query_lut[(size_t)(r + 1) * 256u + packed[r + 1]], + query_lut[(size_t)(r + 2) * 256u + packed[r + 2]], + query_lut[(size_t)(r + 3) * 256u + packed[r + 3]] + }; + acc = vaddq_f32(acc, vld1q_f32(tmp)); + } + } + + float dot; + #if defined(__aarch64__) + dot = vaddvq_f32(acc); + #else + float partial[4]; + vst1q_f32(partial, acc); + dot = partial[0] + partial[1] + partial[2] + partial[3]; + #endif + if (bits == 3) { + for (; r < lut_rows; ++r) dot += query_lut[(size_t)r * 4096u + turbo_lut3_index_neon(packed, r, packed_bytes)]; + } else { + for (; r < lut_rows; ++r) dot += query_lut[(size_t)r * 256u + packed[r]]; + } + return dot * scale; +} + #endif // MARK: - @@ -1310,5 +1363,7 @@ void init_distance_functions_neon (void) { dispatch_distance_table[VECTOR_DISTANCE_HAMMING][VECTOR_TYPE_BIT] = bit1_distance_hamming_neon; distance_backend_name = "NEON"; + turbo_lut_dot_function = turbo_lut_dot_neon; + turbo_lut_backend_name = "NEON"; #endif } diff --git a/src/distance-neon.h b/src/distance-neon.h index df2aa78..2e190ff 100644 --- a/src/distance-neon.h +++ b/src/distance-neon.h @@ -8,8 +8,10 @@ #ifndef __VECTOR_DISTANCE_NEON__ #define __VECTOR_DISTANCE_NEON__ +#include #include void init_distance_functions_neon (void); +float turbo_lut_dot_neon (const uint8_t *packed, float scale, const float *query_lut, int lut_rows, int bits, int packed_bytes); #endif diff --git a/src/distance-rvv.c b/src/distance-rvv.c index 14788cd..e0d099a 100644 --- a/src/distance-rvv.c +++ b/src/distance-rvv.c @@ -16,6 +16,8 @@ extern distance_function_t dispatch_distance_table[VECTOR_DISTANCE_MAX][VECTOR_TYPE_MAX]; extern const char *distance_backend_name; +extern turbo_lut_dot_function_t turbo_lut_dot_function; +extern const char *turbo_lut_backend_name; // MARK: - UTILS - @@ -985,6 +987,36 @@ float bit1_distance_hamming_rvv (const void *v1, const void *v2, int n) { // Copy the accumulator back into a scalar register return (float) uint64_sum_vector_u64m8(vdistance, vl); } + +static inline uint16_t turbo_lut3_index_rvv (const uint8_t *packed, int row, int packed_bytes) { + size_t bit_pos = (size_t)row * 12u; + size_t byte_pos = bit_pos / 8u; + int shift = (int)(bit_pos % 8u); + uint32_t word = 0; + if ((int)byte_pos < packed_bytes) word |= packed[byte_pos]; + if ((int)byte_pos + 1 < packed_bytes) word |= (uint32_t)packed[byte_pos + 1] << 8; + return (uint16_t)((word >> shift) & 0x0fffu); +} + +float turbo_lut_dot_rvv (const uint8_t *packed, float scale, const float *query_lut, int lut_rows, int bits, int packed_bytes) { + size_t vlmax = __riscv_vsetvlmax_e32m8(); + vfloat32m8_t acc = __riscv_vfmv_v_f_f32m8(0.0f, vlmax); + int r = 0; + while (r < lut_rows) { + size_t n = (size_t)(lut_rows - r); + size_t vl = __riscv_vsetvl_e32m8(n); + float tmp[vl]; + for (size_t i = 0; i < vl; ++i) { + int row = r + (int)i; + if (bits == 3) tmp[i] = query_lut[(size_t)row * 4096u + turbo_lut3_index_rvv(packed, row, packed_bytes)]; + else tmp[i] = query_lut[(size_t)row * 256u + packed[row]]; + } + vfloat32m8_t vals = __riscv_vle32_v_f32m8(tmp, vl); + acc = __riscv_vfadd_vv_f32m8(acc, vals, vl); + r += (int)vl; + } + return float32_sum_vector_f32m8(acc, vlmax) * scale; +} #endif // MARK: - @@ -1024,5 +1056,7 @@ void init_distance_functions_rvv (void) { dispatch_distance_table[VECTOR_DISTANCE_HAMMING][VECTOR_TYPE_BIT] = bit1_distance_hamming_rvv; distance_backend_name = "RVV"; + turbo_lut_dot_function = turbo_lut_dot_rvv; + turbo_lut_backend_name = "RVV"; #endif } diff --git a/src/distance-rvv.h b/src/distance-rvv.h index 0d4c509..8e2fbcf 100644 --- a/src/distance-rvv.h +++ b/src/distance-rvv.h @@ -8,8 +8,10 @@ #ifndef __VECTOR_DISTANCE_RVV__ #define __VECTOR_DISTANCE_RVV__ +#include #include void init_distance_functions_rvv (void); +float turbo_lut_dot_rvv (const uint8_t *packed, float scale, const float *query_lut, int lut_rows, int bits, int packed_bytes); #endif diff --git a/src/distance-sse2.c b/src/distance-sse2.c index 426d9b6..c739469 100644 --- a/src/distance-sse2.c +++ b/src/distance-sse2.c @@ -15,6 +15,8 @@ extern distance_function_t dispatch_distance_table[VECTOR_DISTANCE_MAX][VECTOR_TYPE_MAX]; extern const char *distance_backend_name; +extern turbo_lut_dot_function_t turbo_lut_dot_function; +extern const char *turbo_lut_backend_name; // accumulate 32-bit #define ACCUMULATE(MUL, ACC) \ @@ -1073,6 +1075,52 @@ float bit1_distance_hamming_sse2 (const void *v1, const void *v2, int n) { return (float)distance; } +static inline uint16_t turbo_lut3_index_sse2 (const uint8_t *packed, int row, int packed_bytes) { + size_t bit_pos = (size_t)row * 12u; + size_t byte_pos = bit_pos / 8u; + int shift = (int)(bit_pos % 8u); + uint32_t word = 0; + if ((int)byte_pos < packed_bytes) word |= packed[byte_pos]; + if ((int)byte_pos + 1 < packed_bytes) word |= (uint32_t)packed[byte_pos + 1] << 8; + return (uint16_t)((word >> shift) & 0x0fffu); +} + +float turbo_lut_dot_sse2 (const uint8_t *packed, float scale, const float *query_lut, int lut_rows, int bits, int packed_bytes) { + __m128 acc = _mm_setzero_ps(); + int r = 0; + if (bits == 3) { + for (; r + 3 < lut_rows; r += 4) { + float tmp[4] = { + query_lut[(size_t)(r + 0) * 4096u + turbo_lut3_index_sse2(packed, r + 0, packed_bytes)], + query_lut[(size_t)(r + 1) * 4096u + turbo_lut3_index_sse2(packed, r + 1, packed_bytes)], + query_lut[(size_t)(r + 2) * 4096u + turbo_lut3_index_sse2(packed, r + 2, packed_bytes)], + query_lut[(size_t)(r + 3) * 4096u + turbo_lut3_index_sse2(packed, r + 3, packed_bytes)] + }; + acc = _mm_add_ps(acc, _mm_loadu_ps(tmp)); + } + } else { + for (; r + 3 < lut_rows; r += 4) { + float tmp[4] = { + query_lut[(size_t)(r + 0) * 256u + packed[r + 0]], + query_lut[(size_t)(r + 1) * 256u + packed[r + 1]], + query_lut[(size_t)(r + 2) * 256u + packed[r + 2]], + query_lut[(size_t)(r + 3) * 256u + packed[r + 3]] + }; + acc = _mm_add_ps(acc, _mm_loadu_ps(tmp)); + } + } + + float partial[4]; + _mm_storeu_ps(partial, acc); + float dot = partial[0] + partial[1] + partial[2] + partial[3]; + if (bits == 3) { + for (; r < lut_rows; ++r) dot += query_lut[(size_t)r * 4096u + turbo_lut3_index_sse2(packed, r, packed_bytes)]; + } else { + for (; r < lut_rows; ++r) dot += query_lut[(size_t)r * 256u + packed[r]]; + } + return dot * scale; +} + #endif // MARK: - @@ -1112,5 +1160,7 @@ void init_distance_functions_sse2 (void) { dispatch_distance_table[VECTOR_DISTANCE_HAMMING][VECTOR_TYPE_BIT] = bit1_distance_hamming_sse2; distance_backend_name = "SSE2"; + turbo_lut_dot_function = turbo_lut_dot_sse2; + turbo_lut_backend_name = "SSE2"; #endif } diff --git a/src/distance-sse2.h b/src/distance-sse2.h index 429738e..decb874 100644 --- a/src/distance-sse2.h +++ b/src/distance-sse2.h @@ -8,8 +8,10 @@ #ifndef __VECTOR_DISTANCE_SSE2__ #define __VECTOR_DISTANCE_SSE2__ +#include #include void init_distance_functions_sse2 (void); +float turbo_lut_dot_sse2 (const uint8_t *packed, float scale, const float *query_lut, int lut_rows, int bits, int packed_bytes); #endif diff --git a/src/sqlite-vector.c b/src/sqlite-vector.c index 44e4ff2..bf150d3 100644 --- a/src/sqlite-vector.c +++ b/src/sqlite-vector.c @@ -44,7 +44,6 @@ char *strcasestr(const char *haystack, const char *needle) { } #endif - #ifdef SQLITE_WASM_EXTRA_INIT #define sqlite3_mutex_alloc(_type) NULL #define sqlite3_mutex_enter(_mutex) @@ -115,11 +114,14 @@ SQLITE_EXTENSION_INIT1 #define OPTION_KEY_MAXMEMORY "max_memory" #define OPTION_KEY_DISTANCE "distance" #define OPTION_KEY_QUANTTYPE "qtype" +#define OPTION_KEY_QUANTBITS "qbits" #define OPTION_KEY_QUANTSCALE "qscale" // used only in serialize/unserialize #define OPTION_KEY_QUANTOFFSET "qoffset" // used only in serialize/unserialize #define VECTOR_INTERNAL_TABLE "CREATE TABLE IF NOT EXISTS _sqliteai_vector (tblname TEXT, colname TEXT, key TEXT, value ANY, PRIMARY KEY(tblname, colname, key));" +typedef struct turbo_rotation_plan turbo_rotation_plan; + typedef struct { vector_type v_type; // vector type int v_dim; // vector dimension @@ -127,6 +129,7 @@ typedef struct { vector_distance v_distance; // vector distance function vector_qtype q_type; // quantization type + int q_bits; // bit width for TurboQuant uint64_t max_memory; // max memory } vector_options; @@ -142,6 +145,15 @@ typedef struct { void *preloaded; int precounter; + sqlite3_int64 preloaded_bytes; + + turbo_rotation_plan *turbo_plan; + int turbo_plan_dim; + bool turbo_codebook_ready; + int turbo_codebook_dim; + int turbo_codebook_bits; + float turbo_boundaries[15]; + float turbo_centroids[16]; } table_context; typedef struct { @@ -176,6 +188,12 @@ typedef struct { int dcounter; int dindex; int is_eof; + float turbo_qnorm_sq; + int turbo_bits; + sqlite3_int64 data_bytes; + float *turbo_query_lut; + float *turbo_norm_lut; + int turbo_lut_rows; } stream; // NON-STREAMING VT INTERFACE @@ -193,6 +211,7 @@ typedef int (*vcursor_sort_callback)(vFullScanCursor *c); extern distance_function_t dispatch_distance_table[VECTOR_DISTANCE_MAX][VECTOR_TYPE_MAX]; extern const char *distance_backend_name; +extern const char *turbo_lut_backend_name; static sqlite3_mutex *qmutex; @@ -482,6 +501,11 @@ static int sqlite_unserialize (sqlite3_context *context, table_context *ctx) { ctx->options.q_type = (vector_qtype)sqlite3_column_int(vm, 1); continue; } + + if (strcmp(key, OPTION_KEY_QUANTBITS) == 0) { + ctx->options.q_bits = sqlite3_column_int(vm, 1); + continue; + } if (strcmp(key, OPTION_KEY_QUANTSCALE) == 0) { ctx->scale = (float)sqlite3_column_double(vm, 1); @@ -845,6 +869,494 @@ static void quantize_binary_i8 (const int8_t *input, uint8_t *output, int dim) { } } +// MARK: - TurboQuant - + +static inline size_t turbo_bytes_for_dim (int dim, int bits) { + return ((size_t)dim * (size_t)bits + 7u) / 8u; +} + +static inline uint64_t turbo_splitmix64 (uint64_t *state) { + uint64_t z = (*state += 0x9E3779B97F4A7C15ull); + z = (z ^ (z >> 30)) * 0xBF58476D1CE4E5B9ull; + z = (z ^ (z >> 27)) * 0x94D049BB133111EBull; + return z ^ (z >> 31); +} + +struct turbo_rotation_plan { + int dim; + int rounds; + int pair_count; + uint8_t *signs; + int *a; + int *b; + float *s; + float *c; +}; + +static void turbo_rotation_plan_free (turbo_rotation_plan *plan) { + if (!plan) return; + if (plan->signs) sqlite3_free(plan->signs); + if (plan->a) sqlite3_free(plan->a); + if (plan->b) sqlite3_free(plan->b); + if (plan->s) sqlite3_free(plan->s); + if (plan->c) sqlite3_free(plan->c); + memset(plan, 0, sizeof(*plan)); +} + +static int turbo_rotation_plan_init (turbo_rotation_plan *plan, int dim) { + memset(plan, 0, sizeof(*plan)); + plan->dim = dim; + plan->rounds = 12; + int pairs_per_round = dim / 2; + plan->pair_count = pairs_per_round * plan->rounds; + + plan->signs = (uint8_t *)sqlite3_malloc64((sqlite3_uint64)dim); + if (plan->pair_count > 0) { + plan->a = (int *)sqlite3_malloc64((sqlite3_uint64)plan->pair_count * sizeof(int)); + plan->b = (int *)sqlite3_malloc64((sqlite3_uint64)plan->pair_count * sizeof(int)); + plan->s = (float *)sqlite3_malloc64((sqlite3_uint64)plan->pair_count * sizeof(float)); + plan->c = (float *)sqlite3_malloc64((sqlite3_uint64)plan->pair_count * sizeof(float)); + } + if (!plan->signs || (plan->pair_count > 0 && (!plan->a || !plan->b || !plan->s || !plan->c))) { + turbo_rotation_plan_free(plan); + return SQLITE_NOMEM; + } + + uint64_t sign_state = 0xA5A5A5A55A5A5A5Aull ^ (uint64_t)dim; + for (int i = 0; i < dim; ++i) plan->signs[i] = (uint8_t)(turbo_splitmix64(&sign_state) & 1ull); + + int *perm = (int *)sqlite3_malloc64((sqlite3_uint64)dim * sizeof(int)); + if (!perm) { + turbo_rotation_plan_free(plan); + return SQLITE_NOMEM; + } + + int idx = 0; + for (int r = 0; r < plan->rounds; ++r) { + for (int i = 0; i < dim; ++i) perm[i] = i; + uint64_t state = 0xD1B54A32D192ED03ull ^ ((uint64_t)dim << 32) ^ (uint64_t)r; + for (int i = dim - 1; i > 0; --i) { + int j = (int)(turbo_splitmix64(&state) % (uint64_t)(i + 1)); + int tmp = perm[i]; + perm[i] = perm[j]; + perm[j] = tmp; + } + + for (int i = 0; i + 1 < dim; i += 2) { + uint64_t rnd = turbo_splitmix64(&state); + float angle = (float)((double)(rnd >> 11) * (6.28318530717958647692 / 9007199254740992.0)); + plan->a[idx] = perm[i]; + plan->b[idx] = perm[i + 1]; + plan->s[idx] = sinf(angle); + plan->c[idx] = cosf(angle); + idx++; + } + } + sqlite3_free(perm); + return SQLITE_OK; +} + +static float turbo_value_at (const void *v, vector_type type, int i) { + switch (type) { + case VECTOR_TYPE_F32: return ((const float *)v)[i]; + case VECTOR_TYPE_F16: return float16_to_float32(((const uint16_t *)v)[i]); + case VECTOR_TYPE_BF16: return bfloat16_to_float32(((const uint16_t *)v)[i]); + case VECTOR_TYPE_U8: return (float)((const uint8_t *)v)[i]; + case VECTOR_TYPE_I8: return (float)((const int8_t *)v)[i]; + case VECTOR_TYPE_BIT: return (float)((((const uint8_t *)v)[i / 8] >> (i % 8)) & 1); + } + return 0.0f; +} + +static float turbo_copy_float (const void *v, vector_type type, int dim, float *out) { + double norm_sq = 0.0; + for (int i = 0; i < dim; ++i) { + float x = turbo_value_at(v, type, i); + out[i] = x; + norm_sq += (double)x * (double)x; + } + return (float)norm_sq; +} + +static void turbo_normalize_inplace (float *v, int dim, float norm_sq) { + if (norm_sq <= 1e-20f) { + memset(v, 0, (size_t)dim * sizeof(float)); + return; + } + float inv = 1.0f / sqrtf(norm_sq); + for (int i = 0; i < dim; ++i) v[i] *= inv; +} + +static void turbo_rotate_with_plan (const float *input, float *output, const turbo_rotation_plan *plan) { + int dim = plan->dim; + memcpy(output, input, (size_t)dim * sizeof(float)); + if (dim <= 1) return; + + for (int i = 0; i < dim; ++i) if (plan->signs[i]) output[i] = -output[i]; + + for (int i = 0; i < plan->pair_count; ++i) { + int a = plan->a[i]; + int b = plan->b[i]; + float s = plan->s[i]; + float c = plan->c[i]; + float x = output[a]; + float y = output[b]; + output[a] = c * x - s * y; + output[b] = s * x + c * y; + } +} + +static inline double turbo_normal_pdf (double z) { + return 0.39894228040143267794 * exp(-0.5 * z * z); +} + +static inline double turbo_normal_cdf (double z) { + return 0.5 * erfc(-z * 0.70710678118654752440); +} + +static double turbo_beta_pdf_shifted (double x, int dim) { + if (x <= -1.0 || x >= 1.0 || dim <= 1) return 0.0; + double a = ((double)dim - 1.0) * 0.5; + double y = (x + 1.0) * 0.5; + double log_pdf = (a - 1.0) * (log(y) + log1p(-y)) + lgamma(2.0 * a) - 2.0 * lgamma(a) - log(2.0); + return exp(log_pdf); +} + +static void turbo_beta_interval_moments (double lo, double hi, int dim, double *prob_out, double *moment_out) { + static const double nodes[16] = { + 0.048307665687738316, + 0.14447196158279649, + 0.23928736225213707, + 0.33186860228212767, + 0.42135127613063533, + 0.50689990893222939, + 0.58771575724076233, + 0.66304426693021520, + 0.73218211874028968, + 0.79448379596794241, + 0.84936761373256997, + 0.89632115576605212, + 0.93490607593773969, + 0.96476225558750643, + 0.98561151154526834, + 0.99726386184948156 + }; + static const double weights[16] = { + 0.09654008851472780, + 0.09563872007927486, + 0.09384439908080457, + 0.09117387869576388, + 0.08765209300440381, + 0.08331192422694676, + 0.07819389578707031, + 0.07234579410884851, + 0.06582222277636185, + 0.05868409347853555, + 0.05099805926237618, + 0.04283589802222668, + 0.03427386291302143, + 0.02539206530926206, + 0.01627439473090567, + 0.00701861000947010 + }; + + if (lo < -1.0) lo = -1.0; + if (hi > 1.0) hi = 1.0; + if (hi <= lo) { + *prob_out = 0.0; + *moment_out = 0.0; + return; + } + + double mid = 0.5 * (lo + hi); + double half = 0.5 * (hi - lo); + double prob = 0.0; + double moment = 0.0; + for (int i = 0; i < 16; ++i) { + double dx = half * nodes[i]; + double x1 = mid - dx; + double x2 = mid + dx; + double p1 = turbo_beta_pdf_shifted(x1, dim); + double p2 = turbo_beta_pdf_shifted(x2, dim); + prob += weights[i] * (p1 + p2); + moment += weights[i] * (x1 * p1 + x2 * p2); + } + + *prob_out = half * prob; + *moment_out = half * moment; +} + +static void turbo_make_codebook (int bits, int dim, float *boundaries, float *centroids) { + int levels = 1 << bits; + double c[16]; + double next[16]; + double sigma = (dim > 0) ? 1.0 / sqrt((double)dim) : 1.0; + + for (int i = 0; i < levels; ++i) { + c[i] = (-3.0 + 6.0 * (double)i / (double)(levels - 1)) * sigma; + } + + if (dim <= 1) { + for (int it = 0; it < 80; ++it) { + double max_change = 0.0; + for (int i = 0; i < levels; ++i) { + double lo = (i == 0) ? -INFINITY : 0.5 * (c[i - 1] + c[i]); + double hi = (i == levels - 1) ? INFINITY : 0.5 * (c[i] + c[i + 1]); + double zl = lo / sigma; + double zh = hi / sigma; + double p = turbo_normal_cdf(zh) - turbo_normal_cdf(zl); + if (p <= 1e-15) { + next[i] = c[i]; + } else { + double pdf_lo = isinf(zl) ? 0.0 : turbo_normal_pdf(zl); + double pdf_hi = isinf(zh) ? 0.0 : turbo_normal_pdf(zh); + next[i] = sigma * (pdf_lo - pdf_hi) / p; + } + double change = fabs(next[i] - c[i]); + if (change > max_change) max_change = change; + } + memcpy(c, next, (size_t)levels * sizeof(double)); + if (max_change < 1e-12) break; + } + + for (int i = 0; i < levels - 1; ++i) boundaries[i] = (float)(0.5 * (c[i] + c[i + 1])); + for (int i = 0; i < levels; ++i) centroids[i] = (float)c[i]; + return; + } + + for (int it = 0; it < 200; ++it) { + double max_change = 0.0; + for (int i = 0; i < levels; ++i) { + double lo = (i == 0) ? -1.0 : 0.5 * (c[i - 1] + c[i]); + double hi = (i == levels - 1) ? 1.0 : 0.5 * (c[i] + c[i + 1]); + double p = 0.0; + double moment = 0.0; + turbo_beta_interval_moments(lo, hi, dim, &p, &moment); + if (p <= 1e-15) { + next[i] = c[i]; + } else { + next[i] = moment / p; + } + double change = fabs(next[i] - c[i]); + if (change > max_change) max_change = change; + } + memcpy(c, next, (size_t)levels * sizeof(double)); + if (max_change < 1e-12) break; + } + + for (int i = 0; i < levels - 1; ++i) boundaries[i] = (float)(0.5 * (c[i] + c[i + 1])); + for (int i = 0; i < levels; ++i) centroids[i] = (float)c[i]; +} + +static inline uint8_t turbo_code_for_value (float x, const float *boundaries, int bits) { + uint8_t code = 0; + int nboundaries = (1 << bits) - 1; + for (int i = 0; i < nboundaries; ++i) { + if (x > boundaries[i]) ++code; + } + return code; +} + +static void turbo_quantize_rotated (const float *rotated, uint8_t *packed, const float *boundaries, const float *centroids, int bits, int dim, float *inner_out) { + memset(packed, 0, turbo_bytes_for_dim(dim, bits)); + + double inner = 0.0; + for (int j = 0; j < dim; ++j) { + uint8_t code = turbo_code_for_value(rotated[j], boundaries, bits); + inner += (double)rotated[j] * (double)centroids[code]; + size_t bit_pos = (size_t)j * (size_t)bits; + size_t byte_pos = bit_pos / 8u; + int shift = (int)(bit_pos % 8u); + packed[byte_pos] |= (uint8_t)(code << shift); + if (shift + bits > 8) packed[byte_pos + 1] |= (uint8_t)(code >> (8 - shift)); + } + *inner_out = (float)inner; +} + +static inline uint8_t turbo_unpack_code (const uint8_t *packed, int bits, int dim, int j) { + (void)dim; + size_t bit_pos = (size_t)j * (size_t)bits; + size_t byte_pos = bit_pos / 8u; + int shift = (int)(bit_pos % 8u); + uint16_t value = packed[byte_pos]; + if (shift + bits > 8) value |= (uint16_t)packed[byte_pos + 1] << 8; + return (uint8_t)((value >> shift) & ((1u << bits) - 1u)); +} + +static float turbo_distance_from_rotated_query (const float *query_rot, float query_norm_sq, const uint8_t *packed, float scale, const float *centroids, int bits, int dim, vector_distance distance) { + double dot = 0.0; + double xnorm_sq = 0.0; + for (int j = 0; j < dim; ++j) { + float c = centroids[turbo_unpack_code(packed, bits, dim, j)] * scale; + dot += (double)query_rot[j] * (double)c; + xnorm_sq += (double)c * (double)c; + } + + switch (distance) { + case VECTOR_DISTANCE_DOT: + return (float)-dot; + case VECTOR_DISTANCE_COSINE: { + float d = (float)(1.0 - dot); + return d < 0.0f ? 0.0f : d; + } + case VECTOR_DISTANCE_L2: { + double d2 = (double)query_norm_sq + xnorm_sq - 2.0 * dot; + if (d2 < 0.0) d2 = 0.0; + return (float)sqrt(d2); + } + case VECTOR_DISTANCE_SQUARED_L2: { + double d2 = (double)query_norm_sq + xnorm_sq - 2.0 * dot; + return (float)(d2 < 0.0 ? 0.0 : d2); + } + default: + return INFINITY; + } +} + +static float *turbo_build_query_lut (const float *query_rot, const float *centroids, int bits, int dim, int *lut_rows_out) { + if (bits < 2 || bits > 4) { + *lut_rows_out = 0; + return NULL; + } + + int codes_per_row = (bits == 3) ? 4 : (8 / bits); + int entries_per_row = (bits == 3) ? 4096 : 256; + int rows = (dim + codes_per_row - 1) / codes_per_row; + float *lut = (float *)sqlite3_malloc64((sqlite3_uint64)rows * (sqlite3_uint64)entries_per_row * sizeof(float)); + if (!lut) { + *lut_rows_out = 0; + return NULL; + } + + uint8_t mask = (uint8_t)((1u << bits) - 1u); + for (int r = 0; r < rows; ++r) { + for (int entry = 0; entry < entries_per_row; ++entry) { + double sum = 0.0; + for (int c = 0; c < codes_per_row; ++c) { + int j = r * codes_per_row + c; + if (j >= dim) break; + int code = (entry >> (c * bits)) & mask; + sum += (double)query_rot[j] * (double)centroids[code]; + } + lut[(size_t)r * (size_t)entries_per_row + (size_t)entry] = (float)sum; + } + } + + *lut_rows_out = rows; + return lut; +} + +static float *turbo_build_norm_lut (const float *centroids, int bits, int dim, int *lut_rows_out) { + if (bits < 2 || bits > 4) { + *lut_rows_out = 0; + return NULL; + } + + int codes_per_row = (bits == 3) ? 4 : (8 / bits); + int entries_per_row = (bits == 3) ? 4096 : 256; + int rows = (dim + codes_per_row - 1) / codes_per_row; + float *lut = (float *)sqlite3_malloc64((sqlite3_uint64)rows * (sqlite3_uint64)entries_per_row * sizeof(float)); + if (!lut) { + *lut_rows_out = 0; + return NULL; + } + + uint8_t mask = (uint8_t)((1u << bits) - 1u); + for (int r = 0; r < rows; ++r) { + for (int entry = 0; entry < entries_per_row; ++entry) { + double sum = 0.0; + for (int c = 0; c < codes_per_row; ++c) { + int j = r * codes_per_row + c; + if (j >= dim) break; + int code = (entry >> (c * bits)) & mask; + double value = (double)centroids[code]; + sum += value * value; + } + lut[(size_t)r * (size_t)entries_per_row + (size_t)entry] = (float)sum; + } + } + + *lut_rows_out = rows; + return lut; +} + +static inline uint16_t turbo_lut3_index (const uint8_t *packed, int row, int packed_bytes) { + size_t bit_pos = (size_t)row * 12u; + size_t byte_pos = bit_pos / 8u; + int shift = (int)(bit_pos % 8u); + uint32_t word = 0; + if ((int)byte_pos < packed_bytes) word |= packed[byte_pos]; + if ((int)byte_pos + 1 < packed_bytes) word |= (uint32_t)packed[byte_pos + 1] << 8; + return (uint16_t)((word >> shift) & 0x0fffu); +} + +static inline float turbo_dot_from_lut (const uint8_t *packed, float scale, const float *query_lut, int lut_rows, int bits, int packed_bytes) { + if (turbo_lut_dot_function) return turbo_lut_dot_function(packed, scale, query_lut, lut_rows, bits, packed_bytes); + double dot = 0.0; + if (bits == 3) { + for (int r = 0; r < lut_rows; ++r) { + dot += (double)query_lut[(size_t)r * 4096u + turbo_lut3_index(packed, r, packed_bytes)]; + } + } else { + for (int r = 0; r < lut_rows; ++r) { + dot += (double)query_lut[(size_t)r * 256u + packed[r]]; + } + } + return (float)(dot * (double)scale); +} + +static int table_context_ensure_turbo_plan (table_context *t_ctx, int dim) { + if (t_ctx->turbo_plan && t_ctx->turbo_plan_dim == dim) return SQLITE_OK; + + if (t_ctx->turbo_plan) { + turbo_rotation_plan_free(t_ctx->turbo_plan); + sqlite3_free(t_ctx->turbo_plan); + t_ctx->turbo_plan = NULL; + t_ctx->turbo_plan_dim = 0; + } + + turbo_rotation_plan *plan = (turbo_rotation_plan *)sqlite3_malloc64(sizeof(turbo_rotation_plan)); + if (!plan) return SQLITE_NOMEM; + int rc = turbo_rotation_plan_init(plan, dim); + if (rc != SQLITE_OK) { + sqlite3_free(plan); + return rc; + } + + t_ctx->turbo_plan = plan; + t_ctx->turbo_plan_dim = dim; + return SQLITE_OK; +} + +static int table_context_ensure_turbo_codebook (table_context *t_ctx, int bits, int dim) { + if (bits < 2 || bits > 4) return SQLITE_MISUSE; + if (t_ctx->turbo_codebook_ready && t_ctx->turbo_codebook_bits == bits && t_ctx->turbo_codebook_dim == dim) return SQLITE_OK; + + turbo_make_codebook(bits, dim, t_ctx->turbo_boundaries, t_ctx->turbo_centroids); + t_ctx->turbo_codebook_ready = true; + t_ctx->turbo_codebook_bits = bits; + t_ctx->turbo_codebook_dim = dim; + return SQLITE_OK; +} + +static int table_context_require_turbo_cache (table_context *t_ctx, int bits, int dim) { + if (!t_ctx->turbo_plan || t_ctx->turbo_plan_dim != dim) return SQLITE_MISUSE; + if (!t_ctx->turbo_codebook_ready || t_ctx->turbo_codebook_bits != bits || t_ctx->turbo_codebook_dim != dim) return SQLITE_MISUSE; + return SQLITE_OK; +} + +static void table_context_free_turbo_cache (table_context *t_ctx) { + if (t_ctx->turbo_plan) { + turbo_rotation_plan_free(t_ctx->turbo_plan); + sqlite3_free(t_ctx->turbo_plan); + } + t_ctx->turbo_plan = NULL; + t_ctx->turbo_plan_dim = 0; + t_ctx->turbo_codebook_ready = false; + t_ctx->turbo_codebook_dim = 0; + t_ctx->turbo_codebook_bits = 0; +} + // MARK: - General Utils - static int vector_type_to_size (vector_type type) { @@ -893,6 +1405,7 @@ static vector_qtype quant_name_to_type (const char *qname) { if (strcasecmp(qname, "UINT8") == 0) return VECTOR_QUANT_U8BIT; if (strcasecmp(qname, "INT8") == 0) return VECTOR_QUANT_S8BIT; if (strcasecmp(qname, "1BIT") == 0 || strcasecmp(qname, "BIT") == 0 || strcasecmp(qname, "BINARY") == 0) return VECTOR_QUANT_1BIT; + if (strcasecmp(qname, "TURBO") == 0 || strcasecmp(qname, "TURBOQUANT") == 0 || strcasecmp(qname, "TURBO2") == 0 || strcasecmp(qname, "TURBO3") == 0 || strcasecmp(qname, "TURBO4") == 0) return VECTOR_QUANT_TURBO; return -1; } @@ -1086,6 +1599,19 @@ bool vector_keyvalue_callback (sqlite3_context *context, void *xdata, const char vector_qtype type = quant_name_to_type(buffer); if ((int)type == -1) return context_result_error(context, SQLITE_ERROR, "Invalid quantization type: '%s' is not a recognized or supported quantization type", buffer); options->q_type = type; + if (type == VECTOR_QUANT_TURBO) { + if (strcasecmp(buffer, "TURBO2") == 0) options->q_bits = 2; + else if (strcasecmp(buffer, "TURBO3") == 0) options->q_bits = 3; + else if (strcasecmp(buffer, "TURBO4") == 0) options->q_bits = 4; + else if (options->q_bits == 0) options->q_bits = 4; + } + return true; + } + + if (KEY_MATCH(OPTION_KEY_QUANTBITS)) { + int bits = (int)strtol(buffer, NULL, 0); + if (bits < 2 || bits > 4) return context_result_error(context, SQLITE_ERROR, "Invalid TurboQuant bit width: expected 2, 3, or 4, got '%s'", buffer); + options->q_bits = bits; return true; } @@ -1104,6 +1630,16 @@ static inline int nearly_zero_float32 (float x) { return fabsf(x) <= 8.0f * FLT_EPSILON; // tweak factor for your use } +static inline size_t quantized_vector_bytes (vector_qtype qtype, int dim, int bits) { + if (qtype == VECTOR_QUANT_1BIT) return (size_t)((dim + 7) / 8); + if (qtype == VECTOR_QUANT_TURBO) return sizeof(float) + turbo_bytes_for_dim(dim, bits); + return (size_t)dim * sizeof(uint8_t); +} + +static inline size_t quantized_row_bytes (vector_qtype qtype, int dim, int bits) { + return sizeof(int64_t) + quantized_vector_bytes(qtype, dim, bits); +} + // MARK: - SQL - static char *generate_create_quant_table (const char *table_name, const char *column_name, char sql[STATIC_SQL_SIZE]) { @@ -1152,6 +1688,7 @@ void vector_context_free (void *p) { if (ctx->tables[i].c_name) sqlite3_free(ctx->tables[i].c_name); if (ctx->tables[i].pk_name) sqlite3_free(ctx->tables[i].pk_name); if (ctx->tables[i].preloaded) sqlite3_free(ctx->tables[i].preloaded); + table_context_free_turbo_cache(&ctx->tables[i]); } sqlite3_free(p); } @@ -1205,7 +1742,18 @@ void vector_context_add (sqlite3_context *context, vector_context *ctx, const ch ctx->tables[index].options = *options; ctx->table_count++; - sqlite_unserialize(context, &ctx->tables[index]); + int rc = sqlite_unserialize(context, &ctx->tables[index]); + if (rc != SQLITE_OK) { + context_result_error(context, rc, "Unable to load vector metadata for '%s.%s'", table_name, column_name); + return; + } + if (ctx->tables[index].options.q_type == VECTOR_QUANT_TURBO) { + int bits = ctx->tables[index].options.q_bits; + int dim = ctx->tables[index].options.v_dim; + rc = table_context_ensure_turbo_codebook(&ctx->tables[index], bits, dim); + if (rc == SQLITE_OK) rc = table_context_ensure_turbo_plan(&ctx->tables[index], dim); + if (rc != SQLITE_OK) context_result_error(context, rc, "Unable to initialize TurboQuant cache for '%s.%s'", table_name, column_name); + } } void vector_options_init (vector_options *options) { @@ -1214,6 +1762,7 @@ void vector_options_init (vector_options *options) { options->v_distance = VECTOR_DISTANCE_L2; options->max_memory = DEFAULT_MAX_MEMORY; options->q_type = VECTOR_QUANT_AUTO; + options->q_bits = 4; } vector_options vector_options_create (void) { @@ -1250,12 +1799,11 @@ static int vector_serialize_quantization (sqlite3 *db, const char *table_name, c if (rc == SQLITE_DONE) rc = SQLITE_OK; vector_serialize_quantization_cleanup: - if (rc != SQLITE_OK) printf("Error in vector_serialize_quantization: %s\n", sqlite3_errmsg(db)); if (vm) sqlite3_finalize(vm); return rc; } -static int vector_rebuild_quantization (sqlite3_context *context, const char *table_name, const char *column_name, table_context *t_ctx, vector_qtype qtype, uint64_t max_memory, uint32_t *count) { +static int vector_rebuild_quantization (sqlite3_context *context, const char *table_name, const char *column_name, table_context *t_ctx, vector_qtype qtype, int q_bits, uint64_t max_memory, uint32_t *count) { int rc = SQLITE_NOMEM; sqlite3_stmt *vm = NULL; @@ -1266,10 +1814,22 @@ static int vector_rebuild_quantization (sqlite3_context *context, const char *ta const char *pk_name = t_ctx->pk_name; int dim = t_ctx->options.v_dim; vector_type type = t_ctx->options.v_type; + float *turbo_values = NULL; + float *turbo_rotated = NULL; - // compute size of a single quant, format is: rowid + quantize dimensions - size_t quant_bytes = (qtype == VECTOR_QUANT_1BIT) ? ((dim + 7) / 8) : (dim * sizeof(uint8_t)); - size_t q_size = sizeof(int64_t) + quant_bytes; + if (qtype == VECTOR_QUANT_TURBO && (q_bits < 2 || q_bits > 4)) { + context_result_error(context, SQLITE_ERROR, "TurboQuant requires qbits=2, 3, or 4"); + return SQLITE_MISUSE; + } + + if (qtype == VECTOR_QUANT_TURBO && (type == VECTOR_TYPE_BIT || t_ctx->options.v_distance == VECTOR_DISTANCE_HAMMING || t_ctx->options.v_distance == VECTOR_DISTANCE_L1)) { + context_result_error(context, SQLITE_ERROR, "TurboQuant supports FLOAT/INT vectors with DOT, COSINE, L2, or SQUARED_L2 distance"); + return SQLITE_MISUSE; + } + if (qtype != VECTOR_QUANT_TURBO) table_context_free_turbo_cache(t_ctx); + + // compute size of a single quant, format is: rowid + quantized payload + size_t q_size = quantized_row_bytes(qtype, dim, q_bits); if (q_size == 0) { sqlite3_result_error(context, "Vector dimension is zero, which is not possible", -1); return SQLITE_MISUSE; @@ -1283,8 +1843,17 @@ static int vector_rebuild_quantization (sqlite3_context *context, const char *ta if (count <= 0) { // no vectors t_ctx->options.q_type = (qtype == VECTOR_QUANT_AUTO) ? VECTOR_QUANT_U8BIT : qtype; + t_ctx->options.q_bits = q_bits; t_ctx->scale = 1.0f; t_ctx->offset = 0.0f; + if (t_ctx->options.q_type == VECTOR_QUANT_TURBO) { + rc = table_context_ensure_turbo_codebook(t_ctx, q_bits, dim); + if (rc == SQLITE_OK) rc = table_context_ensure_turbo_plan(t_ctx, dim); + if (rc != SQLITE_OK) { + context_result_error(context, rc, "Unable to initialize TurboQuant cache"); + return rc; + } + } return SQLITE_OK; } } @@ -1309,7 +1878,7 @@ static int vector_rebuild_quantization (sqlite3_context *context, const char *ta float max_val = -FLT_MAX; bool contains_negative = false; - if (qtype != VECTOR_QUANT_1BIT) { + if (qtype != VECTOR_QUANT_1BIT && qtype != VECTOR_QUANT_TURBO) { while (1) { rc = sqlite3_step(vm); if (rc == SQLITE_DONE) {rc = SQLITE_OK; break;} @@ -1378,6 +1947,7 @@ static int vector_rebuild_quantization (sqlite3_context *context, const char *ta float offset = (qtype == VECTOR_QUANT_U8BIT) ? min_val : 0.0f; t_ctx->options.q_type = qtype; + t_ctx->options.q_bits = q_bits; t_ctx->scale = scale; t_ctx->offset = offset; @@ -1389,6 +1959,15 @@ static int vector_rebuild_quantization (sqlite3_context *context, const char *ta // actual quantization (ONLY 8bit is supported in this version) uint32_t n_processed = 0; int64_t min_rowid = 0, max_rowid = 0; + if (qtype == VECTOR_QUANT_TURBO) { + rc = table_context_ensure_turbo_codebook(t_ctx, q_bits, dim); + if (rc != SQLITE_OK) goto vector_rebuild_quantization_cleanup; + rc = table_context_ensure_turbo_plan(t_ctx, dim); + if (rc != SQLITE_OK) goto vector_rebuild_quantization_cleanup; + turbo_values = (float *)sqlite3_malloc64((sqlite3_uint64)dim * sizeof(float)); + turbo_rotated = (float *)sqlite3_malloc64((sqlite3_uint64)dim * sizeof(float)); + if (!turbo_values || !turbo_rotated) { rc = SQLITE_NOMEM; goto vector_rebuild_quantization_cleanup; } + } while (1) { rc = sqlite3_step(vm); if (rc == SQLITE_DONE) {rc = SQLITE_OK; break;} @@ -1398,6 +1977,13 @@ static int vector_rebuild_quantization (sqlite3_context *context, const char *ta int64_t rowid = (int64_t)sqlite3_column_int64(vm, 0); const void *blob = sqlite3_column_blob(vm, 1); if (!blob) continue; + size_t blob_size = (size_t)sqlite3_column_bytes(vm, 1); + size_t need_bytes = vector_bytes_for_dim(type, dim); + if (blob_size < need_bytes) { + context_result_error(context, SQLITE_ERROR, "Invalid vector blob found at rowid %lld", (long long)rowid); + rc = SQLITE_ERROR; + goto vector_rebuild_quantization_cleanup; + } if (n_processed == 0) min_rowid = rowid; VECTOR_PRINT((void *)blob, type, dim); @@ -1407,7 +1993,23 @@ static int vector_rebuild_quantization (sqlite3_context *context, const char *ta data += sizeof(int64_t); // quantize vector - if (qtype == VECTOR_QUANT_1BIT) { + if (qtype == VECTOR_QUANT_TURBO) { + float norm_sq = turbo_copy_float(blob, type, dim, turbo_values); + turbo_normalize_inplace(turbo_values, dim, norm_sq); + turbo_rotate_with_plan(turbo_values, turbo_rotated, t_ctx->turbo_plan); + + float inner = 0.0f; + uint8_t *scale_ptr = data; + data += sizeof(float); + turbo_quantize_rotated(turbo_rotated, data, t_ctx->turbo_boundaries, t_ctx->turbo_centroids, q_bits, dim, &inner); + + float norm = sqrtf(norm_sq); + float vector_scale = 0.0f; + if (inner > 1e-10f) { + vector_scale = (t_ctx->options.v_distance == VECTOR_DISTANCE_COSINE || t_ctx->options.v_normalized) ? (1.0f / inner) : (norm / inner); + } + memcpy(scale_ptr, &vector_scale, sizeof(float)); + } else if (qtype == VECTOR_QUANT_1BIT) { // 1-bit quantization: convert source to binary based on type switch (type) { case VECTOR_TYPE_F32: quantize_binary((const float *)blob, data, dim, t_ctx->binary_mean); break; @@ -1434,7 +2036,7 @@ static int vector_rebuild_quantization (sqlite3_context *context, const char *ta VECTOR_PRINT((void *)data, qprint, dim); #endif - data += (qtype == VECTOR_QUANT_1BIT) ? ((dim + 7) / 8) : (dim * sizeof(uint8_t)); + data += (qtype == VECTOR_QUANT_TURBO) ? turbo_bytes_for_dim(dim, q_bits) : ((qtype == VECTOR_QUANT_1BIT) ? ((dim + 7) / 8) : (dim * sizeof(uint8_t))); max_rowid = rowid; ++n_processed; ++tot_processed; @@ -1455,7 +2057,8 @@ static int vector_rebuild_quantization (sqlite3_context *context, const char *ta } vector_rebuild_quantization_cleanup: - if (rc != SQLITE_OK) printf("Error in vector_rebuild_quantization: %s\n", sqlite3_errmsg(db)); + if (turbo_values) sqlite3_free(turbo_values); + if (turbo_rotated) sqlite3_free(turbo_rotated); if (original) sqlite3_free(original); if (vm) sqlite3_finalize(vm); if (count) *count = tot_processed; @@ -1482,8 +2085,18 @@ static void vector_quantize_preload (sqlite3_context *context, int argc, sqlite3 sqlite3_free(t_ctx->preloaded); t_ctx->preloaded = NULL; t_ctx->precounter = 0; + t_ctx->preloaded_bytes = 0; } sqlite3_mutex_leave(qmutex); + + if (t_ctx->options.q_type == VECTOR_QUANT_TURBO) { + int rc = table_context_ensure_turbo_codebook(t_ctx, t_ctx->options.q_bits, t_ctx->options.v_dim); + if (rc == SQLITE_OK) rc = table_context_ensure_turbo_plan(t_ctx, t_ctx->options.v_dim); + if (rc != SQLITE_OK) { + context_result_error(context, rc, "Unable to initialize TurboQuant cache"); + return; + } + } char sql[STATIC_SQL_SIZE]; generate_memory_quant_table(table_name, column_name, sql); @@ -1537,6 +2150,7 @@ static void vector_quantize_preload (sqlite3_context *context, int argc, sqlite3 sqlite3_mutex_enter(qmutex); t_ctx->preloaded = buffer; t_ctx->precounter = counter; + t_ctx->preloaded_bytes = required; sqlite3_mutex_leave(qmutex); } @@ -1570,13 +2184,15 @@ static int vector_quantize (sqlite3_context *context, const char *table_name, co if (res == false) {rc = SQLITE_ERROR; goto quantize_cleanup;} sqlite3_mutex_enter(qmutex); - rc = vector_rebuild_quantization(context, table_name, column_name, t_ctx, options.q_type, options.max_memory, &counter); + rc = vector_rebuild_quantization(context, table_name, column_name, t_ctx, options.q_type, options.q_bits, options.max_memory, &counter); sqlite3_mutex_leave(qmutex); if (rc != SQLITE_OK) goto quantize_cleanup; // serialize quantization options rc = sqlite_serialize(context, table_name, column_name, SQLITE_INTEGER, OPTION_KEY_QUANTTYPE, t_ctx->options.q_type, 0); if (rc != SQLITE_OK) goto quantize_cleanup; + rc = sqlite_serialize(context, table_name, column_name, SQLITE_INTEGER, OPTION_KEY_QUANTBITS, t_ctx->options.q_bits, 0); + if (rc != SQLITE_OK) goto quantize_cleanup; rc = sqlite_serialize(context, table_name, column_name, SQLITE_FLOAT, OPTION_KEY_QUANTSCALE, 0, t_ctx->scale); if (rc != SQLITE_OK) goto quantize_cleanup; rc = sqlite_serialize(context, table_name, column_name, SQLITE_FLOAT, OPTION_KEY_QUANTOFFSET, 0, t_ctx->offset); @@ -1614,7 +2230,7 @@ static void vector_quantize3 (sqlite3_context *context, int argc, sqlite3_value bool was_preloaded = false; int rc = vector_quantize(context, table_name, column_name, options, &was_preloaded); - if ((rc == SQLITE_OK) && (was_preloaded)) vector_quantize_preload(context, argc, argv); + if ((rc == SQLITE_OK) && (was_preloaded)) vector_quantize_preload(context, 2, argv); } static void vector_quantize2 (sqlite3_context *context, int argc, sqlite3_value **argv) { @@ -1661,6 +2277,7 @@ static void vector_quantize_cleanup (sqlite3_context *context, int argc, sqlite3 sqlite3_free(t_ctx->preloaded); t_ctx->preloaded = NULL; t_ctx->precounter = 0; + t_ctx->preloaded_bytes = 0; } sqlite3_mutex_leave(qmutex); @@ -1908,6 +2525,24 @@ static int vCursorFilterCommon (sqlite3_vtab_cursor *cur, int idxNum, const char vFullScanCursor *c = (vFullScanCursor *)cur; vFullScan *vtab = (vFullScan *)cur->pVtab; + if (c->stream.vm) { + sqlite3_finalize(c->stream.vm); + c->stream.vm = NULL; + } + if (c->stream.vector) { + sqlite3_free(c->stream.vector); + c->stream.vector = NULL; + } + if (c->stream.turbo_query_lut) { + sqlite3_free(c->stream.turbo_query_lut); + c->stream.turbo_query_lut = NULL; + } + if (c->stream.turbo_norm_lut) { + sqlite3_free(c->stream.turbo_norm_lut); + c->stream.turbo_norm_lut = NULL; + } + memset(&c->stream, 0, sizeof(c->stream)); + if (argc != 3 && argc != 4) { return sqlite_vtab_set_error(&vtab->base, "%s expects 3 or 4 arguments, but %d were provided", fname, argc); } @@ -2010,6 +2645,7 @@ static int vCursorFilterCommon (sqlite3_vtab_cursor *cur, int idxNum, const char int rc = run_callback(vtab->db, c, vector, vsize); if (vector_allocated) sqlite3_free((void *)vector); + if (rc != SQLITE_OK) return rc; int count = sort_callback(c); c->row_count -= count; @@ -2108,6 +2744,8 @@ static int vFullScanCursorClose (sqlite3_vtab_cursor *cur){ if (c->rowids) sqlite3_free(c->rowids); if (c->distance) sqlite3_free(c->distance); if (c->stream.vector) sqlite3_free(c->stream.vector); + if (c->stream.turbo_query_lut) sqlite3_free(c->stream.turbo_query_lut); + if (c->stream.turbo_norm_lut) sqlite3_free(c->stream.turbo_norm_lut); if (c->stream.vm) sqlite3_finalize(c->stream.vm); sqlite3_free(c); return SQLITE_OK; @@ -2154,6 +2792,93 @@ static int vFullScanCursorNext (sqlite3_vtab_cursor *cur){ } } + if (c->table->options.q_type == VECTOR_QUANT_TURBO) { + const size_t rowid_size = sizeof(int64_t); + const size_t packed_size = turbo_bytes_for_dim(dimension, c->stream.turbo_bits); + const size_t total_stride = rowid_size + sizeof(float) + packed_size; + const float *centroids = c->table->turbo_centroids; + vector_distance distance_type = c->table->options.v_distance; + + if (vm == NULL) { + if (c->stream.data == NULL) return SQLITE_MISUSE; + if (c->stream.dindex >= c->stream.dcounter) { + c->stream.is_eof = 1; + return SQLITE_OK; + } + if (c->stream.data_bytes < 0 || (sqlite3_uint64)c->stream.data_bytes < ((sqlite3_uint64)c->stream.dindex + 1u) * (sqlite3_uint64)total_stride) { + return SQLITE_CORRUPT; + } + + const uint8_t *current_data = (const uint8_t *)c->stream.data + ((size_t)c->stream.dindex * total_stride); + float scale = 0.0f; + memcpy(&scale, current_data + rowid_size, sizeof(float)); + const uint8_t *packed = current_data + rowid_size + sizeof(float); + float distance; + if (c->stream.turbo_query_lut && (distance_type == VECTOR_DISTANCE_DOT || distance_type == VECTOR_DISTANCE_COSINE)) { + float dot = turbo_dot_from_lut(packed, scale, c->stream.turbo_query_lut, c->stream.turbo_lut_rows, c->stream.turbo_bits, (int)packed_size); + distance = (distance_type == VECTOR_DISTANCE_DOT) ? -dot : (1.0f - dot); + if (distance < 0.0f && distance_type == VECTOR_DISTANCE_COSINE) distance = 0.0f; + } else if (c->stream.turbo_query_lut && c->stream.turbo_norm_lut && (distance_type == VECTOR_DISTANCE_L2 || distance_type == VECTOR_DISTANCE_SQUARED_L2)) { + float dot = turbo_dot_from_lut(packed, scale, c->stream.turbo_query_lut, c->stream.turbo_lut_rows, c->stream.turbo_bits, (int)packed_size); + float norm = turbo_dot_from_lut(packed, 1.0f, c->stream.turbo_norm_lut, c->stream.turbo_lut_rows, c->stream.turbo_bits, (int)packed_size); + double d2 = (double)c->stream.turbo_qnorm_sq + ((double)scale * (double)scale * (double)norm) - 2.0 * (double)dot; + if (d2 < 0.0) d2 = 0.0; + distance = (distance_type == VECTOR_DISTANCE_L2) ? (float)sqrt(d2) : (float)d2; + } else { + distance = turbo_distance_from_rotated_query((const float *)v1, c->stream.turbo_qnorm_sq, packed, scale, centroids, c->stream.turbo_bits, dimension, distance_type); + } + if (nearly_zero_float32(distance)) distance = 0.0f; + c->stream.distance = distance; + c->stream.rowid = INT64_FROM_INT8PTR(current_data); + c->stream.dindex++; + return SQLITE_OK; + } + + if (c->stream.dcounter == 0) { + int rc = sqlite3_step(vm); + if (rc == SQLITE_DONE) { c->stream.is_eof = 1; return SQLITE_OK; } + else if (rc != SQLITE_ROW) return rc; + + c->stream.dcounter = sqlite3_column_int(vm, 0); + c->stream.data = (uint8_t *)sqlite3_column_blob(vm, 1); + c->stream.data_bytes = sqlite3_column_bytes(vm, 1); + c->stream.dindex = 0; + if (c->stream.data == NULL || c->stream.dcounter < 0 || c->stream.data_bytes < 0 || (sqlite3_uint64)c->stream.data_bytes < (sqlite3_uint64)c->stream.dcounter * (sqlite3_uint64)total_stride) { + return SQLITE_CORRUPT; + } + } + + const uint8_t *current_data = (const uint8_t *)c->stream.data + ((size_t)c->stream.dindex * total_stride); + float scale = 0.0f; + memcpy(&scale, current_data + rowid_size, sizeof(float)); + const uint8_t *packed = current_data + rowid_size + sizeof(float); + float distance; + if (c->stream.turbo_query_lut && (distance_type == VECTOR_DISTANCE_DOT || distance_type == VECTOR_DISTANCE_COSINE)) { + float dot = turbo_dot_from_lut(packed, scale, c->stream.turbo_query_lut, c->stream.turbo_lut_rows, c->stream.turbo_bits, (int)packed_size); + distance = (distance_type == VECTOR_DISTANCE_DOT) ? -dot : (1.0f - dot); + if (distance < 0.0f && distance_type == VECTOR_DISTANCE_COSINE) distance = 0.0f; + } else if (c->stream.turbo_query_lut && c->stream.turbo_norm_lut && (distance_type == VECTOR_DISTANCE_L2 || distance_type == VECTOR_DISTANCE_SQUARED_L2)) { + float dot = turbo_dot_from_lut(packed, scale, c->stream.turbo_query_lut, c->stream.turbo_lut_rows, c->stream.turbo_bits, (int)packed_size); + float norm = turbo_dot_from_lut(packed, 1.0f, c->stream.turbo_norm_lut, c->stream.turbo_lut_rows, c->stream.turbo_bits, (int)packed_size); + double d2 = (double)c->stream.turbo_qnorm_sq + ((double)scale * (double)scale * (double)norm) - 2.0 * (double)dot; + if (d2 < 0.0) d2 = 0.0; + distance = (distance_type == VECTOR_DISTANCE_L2) ? (float)sqrt(d2) : (float)d2; + } else { + distance = turbo_distance_from_rotated_query((const float *)v1, c->stream.turbo_qnorm_sq, packed, scale, centroids, c->stream.turbo_bits, dimension, distance_type); + } + if (nearly_zero_float32(distance)) distance = 0.0f; + c->stream.distance = distance; + c->stream.rowid = INT64_FROM_INT8PTR(current_data); + c->stream.dindex++; + + if (c->stream.dindex == c->stream.dcounter) { + c->stream.dcounter = 0; + c->stream.data = NULL; + c->stream.data_bytes = 0; + } + return SQLITE_OK; + } + // QUANTIZATION sizes const size_t rowid_size = sizeof(int64_t); const size_t vector_size = (size_t)c->stream.vsize; // correctly set by caller for 1-bit or 8-bit @@ -2387,7 +3112,150 @@ static int vQuantRunMemory(vFullScanCursor *c, uint8_t *v, vector_qtype qtype, i return SQLITE_OK; } +static int vTurboPrepareQuery (vFullScanCursor *c, const void *v1, float **qrot_out, float *qnorm_sq_out) { + int dim = c->table->options.v_dim; + vector_type type = c->table->options.v_type; + float *values = (float *)sqlite3_malloc64((sqlite3_uint64)dim * sizeof(float)); + float *rotated = (float *)sqlite3_malloc64((sqlite3_uint64)dim * sizeof(float)); + if (!values || !rotated) { + if (values) sqlite3_free(values); + if (rotated) sqlite3_free(rotated); + return SQLITE_NOMEM; + } + + float norm_sq = turbo_copy_float(v1, type, dim, values); + if (c->table->options.v_distance == VECTOR_DISTANCE_COSINE) { + turbo_normalize_inplace(values, dim, norm_sq); + norm_sq = 1.0f; + } + int rc = table_context_require_turbo_cache(c->table, c->table->options.q_bits, dim); + if (rc != SQLITE_OK) { + sqlite3_free(values); + sqlite3_free(rotated); + return rc; + } + turbo_rotate_with_plan(values, rotated, c->table->turbo_plan); + sqlite3_free(values); + + *qrot_out = rotated; + *qnorm_sq_out = norm_sq; + return SQLITE_OK; +} + +static int vTurboRunPackedRows (vFullScanCursor *c, const uint8_t *data, sqlite3_int64 data_bytes, int counter, const float *qrot, float qnorm_sq, const float *centroids, const float *query_lut, const float *norm_lut, int lut_rows, int bits) { + int dim = c->table->options.v_dim; + vector_distance distance_type = c->table->options.v_distance; + size_t packed_bytes = turbo_bytes_for_dim(dim, bits); + size_t total_stride = sizeof(int64_t) + sizeof(float) + packed_bytes; + if (!data || counter < 0 || data_bytes < 0 || (sqlite3_uint64)data_bytes < (sqlite3_uint64)counter * (sqlite3_uint64)total_stride) { + return SQLITE_CORRUPT; + } + + double *distance = c->distance; + int64_t *rowids = c->rowids; + int max_index = c->max_index; + double current_max = distance[max_index]; + + for (int i = 0; i < counter; ++i) { + const uint8_t *current = data + ((size_t)i * total_stride); + float scale = 0.0f; + memcpy(&scale, current + sizeof(int64_t), sizeof(float)); + const uint8_t *packed = current + sizeof(int64_t) + sizeof(float); + + float dist; + if (query_lut && (distance_type == VECTOR_DISTANCE_DOT || distance_type == VECTOR_DISTANCE_COSINE)) { + float dot = turbo_dot_from_lut(packed, scale, query_lut, lut_rows, bits, (int)packed_bytes); + dist = (distance_type == VECTOR_DISTANCE_DOT) ? -dot : (1.0f - dot); + if (dist < 0.0f && distance_type == VECTOR_DISTANCE_COSINE) dist = 0.0f; + } else if (query_lut && norm_lut && (distance_type == VECTOR_DISTANCE_L2 || distance_type == VECTOR_DISTANCE_SQUARED_L2)) { + float dot = turbo_dot_from_lut(packed, scale, query_lut, lut_rows, bits, (int)packed_bytes); + float norm = turbo_dot_from_lut(packed, 1.0f, norm_lut, lut_rows, bits, (int)packed_bytes); + double d2 = (double)qnorm_sq + ((double)scale * (double)scale * (double)norm) - 2.0 * (double)dot; + if (d2 < 0.0) d2 = 0.0; + dist = (distance_type == VECTOR_DISTANCE_L2) ? (float)sqrt(d2) : (float)d2; + } else { + dist = turbo_distance_from_rotated_query(qrot, qnorm_sq, packed, scale, centroids, bits, dim, distance_type); + } + if (nearly_zero_float32(dist)) dist = 0.0f; + + if (dist < current_max) { + distance[max_index] = dist; + rowids[max_index] = INT64_FROM_INT8PTR(current); + max_index = vFullScanFindMaxIndex(distance, c->row_count); + current_max = distance[max_index]; + } + } + + c->max_index = max_index; + return SQLITE_OK; +} + +static int vTurboRun (sqlite3 *db, vFullScanCursor *c, const void *v1, int v1size) { + (void)v1size; + + int dim = c->table->options.v_dim; + int bits = c->table->options.q_bits; + float *qrot = NULL; + float *query_lut = NULL; + float *norm_lut = NULL; + int lut_rows = 0; + float qnorm_sq = 0.0f; + sqlite3_stmt *vm = NULL; + int rc = SQLITE_OK; + + if (bits < 2 || bits > 4) return SQLITE_MISUSE; + rc = table_context_require_turbo_cache(c->table, bits, dim); + if (rc != SQLITE_OK) goto cleanup; + rc = vTurboPrepareQuery(c, v1, &qrot, &qnorm_sq); + if (rc != SQLITE_OK) goto cleanup; + if (c->table->options.v_distance == VECTOR_DISTANCE_DOT || c->table->options.v_distance == VECTOR_DISTANCE_COSINE || + c->table->options.v_distance == VECTOR_DISTANCE_L2 || c->table->options.v_distance == VECTOR_DISTANCE_SQUARED_L2) { + query_lut = turbo_build_query_lut(qrot, c->table->turbo_centroids, bits, dim, &lut_rows); + if (!query_lut) { rc = SQLITE_NOMEM; goto cleanup; } + } + if (c->table->options.v_distance == VECTOR_DISTANCE_L2 || c->table->options.v_distance == VECTOR_DISTANCE_SQUARED_L2) { + int norm_rows = 0; + norm_lut = turbo_build_norm_lut(c->table->turbo_centroids, bits, dim, &norm_rows); + if (!norm_lut) { rc = SQLITE_NOMEM; goto cleanup; } + if (norm_rows != lut_rows) { rc = SQLITE_CORRUPT; goto cleanup; } + } + + if (c->table->preloaded) { + rc = vTurboRunPackedRows(c, (const uint8_t *)c->table->preloaded, c->table->preloaded_bytes, c->table->precounter, qrot, qnorm_sq, c->table->turbo_centroids, query_lut, norm_lut, lut_rows, bits); + goto cleanup; + } + + char sql[STATIC_SQL_SIZE]; + generate_select_quant_table(c->table->t_name, c->table->c_name, sql); + rc = sqlite3_prepare_v2(db, sql, -1, &vm, NULL); + if (rc != SQLITE_OK) goto cleanup; + + while (1) { + rc = sqlite3_step(vm); + if (rc == SQLITE_DONE) { rc = SQLITE_OK; break; } + if (rc != SQLITE_ROW) break; + + int counter = sqlite3_column_int(vm, 0); + const uint8_t *data = (const uint8_t *)sqlite3_column_blob(vm, 1); + int bytes = sqlite3_column_bytes(vm, 1); + if (data) { + rc = vTurboRunPackedRows(c, data, bytes, counter, qrot, qnorm_sq, c->table->turbo_centroids, query_lut, norm_lut, lut_rows, bits); + if (rc != SQLITE_OK) break; + } + } + +cleanup: + if (rc != SQLITE_OK && c && c->base.pVtab) sqlite_vtab_set_error(c->base.pVtab, "TurboQuant scan failed: %s", sqlite3_errmsg(db)); + if (vm) sqlite3_finalize(vm); + if (query_lut) sqlite3_free(query_lut); + if (norm_lut) sqlite3_free(norm_lut); + if (qrot) sqlite3_free(qrot); + return rc; +} + static int vQuantRun (sqlite3 *db, vFullScanCursor *c, const void *v1, int v1size) { + if (c->table->options.q_type == VECTOR_QUANT_TURBO) return vTurboRun(db, c, v1, v1size); + // quantize target vector int dimension = c->table->options.v_dim; vector_qtype qtype = c->table->options.q_type; @@ -2483,7 +3351,7 @@ static int vQuantRun (sqlite3 *db, vFullScanCursor *c, const void *v1, int v1siz rc = SQLITE_OK; vquant_run_cleanup: - if (rc != SQLITE_OK) printf("Error in vector_rebuild_quantization: %s\n", sqlite3_errmsg(db)); + if (rc != SQLITE_OK && c && c->base.pVtab) sqlite_vtab_set_error(c->base.pVtab, "Quantized scan failed: %s", sqlite3_errmsg(db)); if (vm) sqlite3_finalize(vm); if (v) sqlite3_free(v); return rc; @@ -2509,7 +3377,11 @@ static int vStreamScanCursorRun (sqlite3 *db, vFullScanCursor *c, const void *v1 c->stream.vdim = dimension; char *sql = sqlite3_mprintf("SELECT %q, %q FROM %q;", pk_name, col_name, table_name); - if (!sql) return SQLITE_NOMEM; + if (!sql) { + sqlite3_free(v); + c->stream.vector = NULL; + return SQLITE_NOMEM; + } sqlite3_stmt *vm = NULL; int rc = sqlite3_prepare_v2(db, sql, -1, &vm, NULL); @@ -2530,10 +3402,95 @@ static int vStreamScanCursorRun (sqlite3 *db, vFullScanCursor *c, const void *v1 cleanup: if (sql) sqlite3_free(sql); if (vm) sqlite3_finalize(vm); + if (v) sqlite3_free(v); + c->stream.vector = NULL; return rc; } +static int vStreamTurboCursorRun (sqlite3 *db, vFullScanCursor *c, const void *v1, int v1size) { + (void)v1size; + + float *qrot = NULL; + float qnorm_sq = 0.0f; + int dim = c->table->options.v_dim; + int bits = c->table->options.q_bits; + int rc = table_context_require_turbo_cache(c->table, bits, dim); + if (rc != SQLITE_OK) return rc; + + rc = vTurboPrepareQuery(c, v1, &qrot, &qnorm_sq); + if (rc != SQLITE_OK) return rc; + + c->stream.vector = qrot; + c->stream.vsize = (int)turbo_bytes_for_dim(dim, bits); + c->stream.vdim = dim; + c->stream.turbo_qnorm_sq = qnorm_sq; + c->stream.turbo_bits = bits; + if (c->table->options.v_distance == VECTOR_DISTANCE_DOT || c->table->options.v_distance == VECTOR_DISTANCE_COSINE || + c->table->options.v_distance == VECTOR_DISTANCE_L2 || c->table->options.v_distance == VECTOR_DISTANCE_SQUARED_L2) { + c->stream.turbo_query_lut = turbo_build_query_lut(qrot, c->table->turbo_centroids, bits, dim, &c->stream.turbo_lut_rows); + if (!c->stream.turbo_query_lut) { + sqlite3_free(qrot); + c->stream.vector = NULL; + return SQLITE_NOMEM; + } + } + if (c->table->options.v_distance == VECTOR_DISTANCE_L2 || c->table->options.v_distance == VECTOR_DISTANCE_SQUARED_L2) { + int norm_rows = 0; + int norm_rc = SQLITE_OK; + c->stream.turbo_norm_lut = turbo_build_norm_lut(c->table->turbo_centroids, bits, dim, &norm_rows); + if (!c->stream.turbo_norm_lut) norm_rc = SQLITE_NOMEM; + else if (norm_rows != c->stream.turbo_lut_rows) norm_rc = SQLITE_CORRUPT; + if (norm_rc != SQLITE_OK) { + sqlite3_free(qrot); + if (c->stream.turbo_query_lut) { + sqlite3_free(c->stream.turbo_query_lut); + c->stream.turbo_query_lut = NULL; + } + if (c->stream.turbo_norm_lut) { + sqlite3_free(c->stream.turbo_norm_lut); + c->stream.turbo_norm_lut = NULL; + } + c->stream.turbo_lut_rows = 0; + c->stream.vector = NULL; + return norm_rc; + } + } + + if (c->table->preloaded) { + c->stream.dindex = 0; + c->stream.data = c->table->preloaded; + c->stream.dcounter = c->table->precounter; + c->stream.data_bytes = c->table->preloaded_bytes; + return SQLITE_OK; + } + + char sql[STATIC_SQL_SIZE]; + generate_select_quant_table(c->table->t_name, c->table->c_name, sql); + sqlite3_stmt *vm = NULL; + rc = sqlite3_prepare_v2(db, sql, -1, &vm, NULL); + if (rc != SQLITE_OK) { + sqlite3_free(qrot); + if (c->stream.turbo_query_lut) { + sqlite3_free(c->stream.turbo_query_lut); + c->stream.turbo_query_lut = NULL; + c->stream.turbo_lut_rows = 0; + } + if (c->stream.turbo_norm_lut) { + sqlite3_free(c->stream.turbo_norm_lut); + c->stream.turbo_norm_lut = NULL; + } + c->stream.vector = NULL; + if (vm) sqlite3_finalize(vm); + return rc; + } + + c->stream.vm = vm; + return SQLITE_OK; +} + static int vStreamQuantCursorRun (sqlite3 *db, vFullScanCursor *c, const void *v1, int v1size) { + if (c->table->options.q_type == VECTOR_QUANT_TURBO) return vStreamTurboCursorRun(db, c, v1, v1size); + // quantize input vector int dimension = c->table->options.v_dim; vector_qtype qtype = c->table->options.q_type; @@ -2602,6 +3559,8 @@ static int vStreamQuantCursorRun (sqlite3 *db, vFullScanCursor *c, const void *v cleanup: if (vm) sqlite3_finalize(vm); + if (v) sqlite3_free(v); + c->stream.vector = NULL; return rc; } @@ -2726,6 +3685,10 @@ static void vector_version (sqlite3_context *context, int argc, sqlite3_value ** static void vector_backend (sqlite3_context *context, int argc, sqlite3_value **argv) { sqlite3_result_text(context, distance_backend_name, -1, NULL); } + +static void vector_turboquant_backend (sqlite3_context *context, int argc, sqlite3_value **argv) { + sqlite3_result_text(context, turbo_lut_backend_name, -1, NULL); +} // MARK: - @@ -2762,6 +3725,9 @@ SQLITE_VECTOR_API int sqlite3_vector_init (sqlite3 *db, char **pzErrMsg, const s rc = sqlite3_create_function(db, "vector_backend", 0, SQLITE_UTF8, ctx, vector_backend, NULL, NULL); if (rc != SQLITE_OK) goto cleanup; + + rc = sqlite3_create_function(db, "vector_turboquant_backend", 0, SQLITE_UTF8, ctx, vector_turboquant_backend, NULL, NULL); + if (rc != SQLITE_OK) goto cleanup; // table_name, column_name, options rc = sqlite3_create_function(db, "vector_init", 3, SQLITE_UTF8, ctx, vector_init, NULL, NULL); diff --git a/src/sqlite-vector.h b/src/sqlite-vector.h index 8f34bba..48a1029 100644 --- a/src/sqlite-vector.h +++ b/src/sqlite-vector.h @@ -24,7 +24,7 @@ extern "C" { #endif -#define SQLITE_VECTOR_VERSION "0.9.95" +#define SQLITE_VECTOR_VERSION "1.0.0" SQLITE_VECTOR_API int sqlite3_vector_init (sqlite3 *db, char **pzErrMsg, const sqlite3_api_routines *pApi); diff --git a/test/benchmark_turboquant.c b/test/benchmark_turboquant.c new file mode 100644 index 0000000..4d05581 --- /dev/null +++ b/test/benchmark_turboquant.c @@ -0,0 +1,260 @@ +/* + * benchmark_turboquant.c + * + * Synthetic brute-force vs TurboQuant benchmark for sqlite-vector. + * Build example: + * gcc -DSQLITE_CORE -O3 -Isrc -Ilibs test/benchmark_turboquant.c libs/sqlite3.c src/sqlite-vector.c ... -o build/benchmark_turboquant -lm -lpthread + */ + +#include +#include +#include +#include +#include +#include + +#include "sqlite3.h" +#include "sqlite-vector.h" + +#ifndef NVECS +#define NVECS 4000 +#endif +#ifndef NQUERIES +#define NQUERIES 40 +#endif +#ifndef DIM +#define DIM 128 +#endif +#ifndef K +#define K 10 +#endif +#ifndef Q_BITS +#define Q_BITS 4 +#endif +#ifndef KEEP_VECTORS +#define KEEP_VECTORS 1 +#endif +#ifndef DB_PATH +#define DB_PATH ":memory:" +#endif +#ifndef SQLITE_CACHE_KB +#define SQLITE_CACHE_KB 65536 +#endif +#ifndef PRELOAD +#define PRELOAD 1 +#endif + +static uint64_t rng_state = 0x123456789abcdef0ull; + +static void normalize(float *v, int dim); + +static uint64_t splitmix64(void) { + uint64_t z = (rng_state += 0x9E3779B97F4A7C15ull); + z = (z ^ (z >> 30)) * 0xBF58476D1CE4E5B9ull; + z = (z ^ (z >> 27)) * 0x94D049BB133111EBull; + return z ^ (z >> 31); +} + +static float rand_f32(void) { + return (float)((double)(splitmix64() >> 11) * (1.0 / 9007199254740992.0)); +} + +static float rand_normal(void) { + float u1 = rand_f32(); + float u2 = rand_f32(); + if (u1 < 1e-7f) u1 = 1e-7f; + return sqrtf(-2.0f * logf(u1)) * cosf(6.28318530717958647692f * u2); +} + +static uint64_t vector_seed(int idx) { + uint64_t x = 0x9E3779B97F4A7C15ull ^ ((uint64_t)(idx + 1) * 0xBF58476D1CE4E5B9ull); + x ^= (uint64_t)DIM * 0x94D049BB133111EBull; + return x; +} + +static void fill_vector_for_id(int idx, float *row) { + uint64_t saved = rng_state; + rng_state = vector_seed(idx); + for (int j = 0; j < DIM; ++j) row[j] = rand_normal(); + normalize(row, DIM); + rng_state = saved; +} + +static void normalize(float *v, int dim) { + double norm_sq = 0.0; + for (int i = 0; i < dim; ++i) norm_sq += (double)v[i] * (double)v[i]; + float inv = norm_sq > 1e-20 ? 1.0f / sqrtf((float)norm_sq) : 0.0f; + for (int i = 0; i < dim; ++i) v[i] *= inv; +} + +static double now_ms(void) { + struct timespec ts; + clock_gettime(CLOCK_MONOTONIC, &ts); + return (double)ts.tv_sec * 1000.0 + (double)ts.tv_nsec / 1000000.0; +} + +static int exec_or_die(sqlite3 *db, const char *sql) { + char *err = NULL; + int rc = sqlite3_exec(db, sql, NULL, NULL, &err); + if (rc != SQLITE_OK) { + fprintf(stderr, "SQL error: %s\nSQL: %s\n", err ? err : sqlite3_errmsg(db), sql); + sqlite3_free(err); + exit(1); + } + return rc; +} + +static void run_query(sqlite3_stmt *stmt, const float *query, int ids[K]) { + sqlite3_reset(stmt); + sqlite3_clear_bindings(stmt); + sqlite3_bind_blob(stmt, 1, query, DIM * (int)sizeof(float), SQLITE_STATIC); + + int i = 0; + while (sqlite3_step(stmt) == SQLITE_ROW && i < K) { + ids[i++] = sqlite3_column_int(stmt, 0); + } + while (i < K) ids[i++] = -1; +} + +static int overlap_at_k(const int a[K], const int b[K]) { + int count = 0; + for (int i = 0; i < K; ++i) { + for (int j = 0; j < K; ++j) { + if (a[i] == b[j]) { + count++; + break; + } + } + } + return count; +} + +int main(void) { + sqlite3 *db = NULL; + if (strcmp(DB_PATH, ":memory:") != 0) remove(DB_PATH); + if (sqlite3_open(DB_PATH, &db) != SQLITE_OK) { + fprintf(stderr, "cannot open sqlite database\n"); + return 1; + } + if (sqlite3_vector_init(db, NULL, NULL) != SQLITE_OK) { + fprintf(stderr, "cannot init sqlite-vector\n"); + return 1; + } + + float *vectors = KEEP_VECTORS ? (float *)malloc((size_t)NVECS * DIM * sizeof(float)) : NULL; + float *scratch = KEEP_VECTORS ? NULL : (float *)malloc((size_t)DIM * sizeof(float)); + float *queries = (float *)malloc((size_t)NQUERIES * DIM * sizeof(float)); + int (*exact_ids)[K] = malloc((size_t)NQUERIES * sizeof(*exact_ids)); + int (*turbo_ids)[K] = malloc((size_t)NQUERIES * sizeof(*turbo_ids)); + if ((KEEP_VECTORS && !vectors) || (!KEEP_VECTORS && !scratch) || !queries || !exact_ids || !turbo_ids) { + fprintf(stderr, "out of memory\n"); + return 1; + } + + if (KEEP_VECTORS) { + for (int i = 0; i < NVECS; ++i) { + float *row = vectors + (size_t)i * DIM; + fill_vector_for_id(i, row); + } + } + for (int i = 0; i < NQUERIES; ++i) { + int base = (int)(splitmix64() % NVECS); + float *q = queries + (size_t)i * DIM; + float *v = KEEP_VECTORS ? vectors + (size_t)base * DIM : scratch; + if (!KEEP_VECTORS) fill_vector_for_id(base, scratch); + for (int j = 0; j < DIM; ++j) q[j] = v[j] + 0.03f * rand_normal(); + normalize(q, DIM); + } + + exec_or_die(db, "PRAGMA journal_mode=OFF;"); + exec_or_die(db, "PRAGMA synchronous=OFF;"); + char pragma_sql[128]; + snprintf(pragma_sql, sizeof(pragma_sql), "PRAGMA cache_size=-%d;", SQLITE_CACHE_KB); + exec_or_die(db, pragma_sql); + exec_or_die(db, "CREATE TABLE bench(id INTEGER PRIMARY KEY, v BLOB);"); + sqlite3_stmt *insert = NULL; + sqlite3_prepare_v2(db, "INSERT INTO bench(id, v) VALUES(?1, ?2);", -1, &insert, NULL); + for (int i = 0; i < NVECS; ++i) { + sqlite3_reset(insert); + sqlite3_clear_bindings(insert); + sqlite3_bind_int(insert, 1, i + 1); + const float *row = KEEP_VECTORS ? vectors + (size_t)i * DIM : scratch; + if (!KEEP_VECTORS) fill_vector_for_id(i, scratch); + sqlite3_bind_blob(insert, 2, row, DIM * (int)sizeof(float), SQLITE_STATIC); + if (sqlite3_step(insert) != SQLITE_DONE) { + fprintf(stderr, "insert failed: %s\n", sqlite3_errmsg(db)); + return 1; + } + } + sqlite3_finalize(insert); + + char init_sql[256]; + snprintf(init_sql, sizeof(init_sql), "SELECT vector_init('bench', 'v', 'type=f32,dimension=%d,distance=DOT');", DIM); + exec_or_die(db, init_sql); + + double t0 = now_ms(); + char quant_sql[128]; + snprintf(quant_sql, sizeof(quant_sql), "SELECT vector_quantize('bench', 'v', 'qtype=TURBO,qbits=%d,max_memory=0');", Q_BITS); + exec_or_die(db, quant_sql); + double quant_ms = now_ms() - t0; +#if PRELOAD + exec_or_die(db, "SELECT vector_quantize_preload('bench', 'v');"); +#endif + + sqlite3_stmt *full = NULL; + sqlite3_stmt *turbo = NULL; + sqlite3_prepare_v2(db, "SELECT id FROM vector_full_scan('bench', 'v', ?1, 10);", -1, &full, NULL); + sqlite3_prepare_v2(db, "SELECT id FROM vector_quantize_scan('bench', 'v', ?1, 10);", -1, &turbo, NULL); + + t0 = now_ms(); + for (int i = 0; i < NQUERIES; ++i) run_query(full, queries + (size_t)i * DIM, exact_ids[i]); + double full_ms = now_ms() - t0; + + t0 = now_ms(); + for (int i = 0; i < NQUERIES; ++i) run_query(turbo, queries + (size_t)i * DIM, turbo_ids[i]); + double turbo_ms = now_ms() - t0; + + int overlap = 0; + for (int i = 0; i < NQUERIES; ++i) overlap += overlap_at_k(exact_ids[i], turbo_ids[i]); + double recall = (double)overlap / (double)(NQUERIES * K); + + sqlite3_finalize(full); + sqlite3_finalize(turbo); + + int open_stmts = 0; + for (sqlite3_stmt *stmt = sqlite3_next_stmt(db, NULL); stmt; stmt = sqlite3_next_stmt(db, stmt)) open_stmts++; + + sqlite3_int64 memory = 0; + sqlite3_stmt *mem = NULL; + sqlite3_prepare_v2(db, "SELECT vector_quantize_memory('bench', 'v');", -1, &mem, NULL); + if (sqlite3_step(mem) == SQLITE_ROW) memory = sqlite3_column_int64(mem, 0); + sqlite3_finalize(mem); + + const unsigned char *distance_backend = (const unsigned char *)"unknown"; + const unsigned char *turbo_backend = (const unsigned char *)"unknown"; + sqlite3_stmt *backend = NULL; + if (sqlite3_prepare_v2(db, "SELECT vector_backend(), vector_turboquant_backend();", -1, &backend, NULL) == SQLITE_OK && + sqlite3_step(backend) == SQLITE_ROW) { + const unsigned char *v = sqlite3_column_text(backend, 0); + if (v) distance_backend = v; + v = sqlite3_column_text(backend, 1); + if (v) turbo_backend = v; + } + + printf("dataset vectors=%d dim=%d queries=%d k=%d qbits=%d preload=%d\n", NVECS, DIM, NQUERIES, K, Q_BITS, PRELOAD); + printf("backend distance=%s turboquant=%s\n", distance_backend, turbo_backend); + printf("turboquant build_ms=%.3f storage_bytes=%lld\n", quant_ms, (long long)memory); + printf("full_scan_ms=%.3f per_query_ms=%.3f\n", full_ms, full_ms / NQUERIES); + printf("turboquant_ms=%.3f per_query_ms=%.3f speedup=%.2fx\n", turbo_ms, turbo_ms / NQUERIES, full_ms / turbo_ms); + printf("recall@%d=%.4f open_statements=%d\n", K, recall, open_stmts); + + if (backend) sqlite3_finalize(backend); + free(vectors); + free(scratch); + free(queries); + free(exact_ids); + free(turbo_ids); + sqlite3_close(db); + if (strcmp(DB_PATH, ":memory:") != 0) remove(DB_PATH); + return open_stmts == 0 ? 0 : 1; +} diff --git a/test/recall_turboquant_real.py b/test/recall_turboquant_real.py new file mode 100644 index 0000000..461d204 --- /dev/null +++ b/test/recall_turboquant_real.py @@ -0,0 +1,178 @@ +#!/usr/bin/env python3 +"""Recall benchmark for TurboQuant on a real ANN-Benchmarks dataset.""" + +import argparse +import os +import sqlite3 +import sys +import tempfile +import time +import urllib.request + + +DATASET_URLS = ( + "http://ann-benchmarks.com/fashion-mnist-784-euclidean.hdf5", + "https://huggingface.co/datasets/hhy3/ann-datasets/resolve/main/fashion-mnist-784-euclidean.hdf5", +) +DATASET_NAME = "fashion-mnist-784-euclidean.hdf5" + + +def require_deps(): + try: + import h5py # noqa: F401 + import numpy as np # noqa: F401 + except ImportError as exc: + print(f"missing Python dependency: {exc}", file=sys.stderr) + print("install with: python3 -m pip install h5py numpy", file=sys.stderr) + sys.exit(2) + + +def default_extension_path(): + root = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) + for name in ("vector.dylib", "vector.so", "vector.dll", "vector"): + path = os.path.join(root, "dist", name) + if os.path.exists(path): + return path + return os.path.join(root, "dist", "vector") + + +def download_dataset(cache_dir): + os.makedirs(cache_dir, exist_ok=True) + path = os.path.join(cache_dir, DATASET_NAME) + if not os.path.exists(path): + last_error = None + for url in DATASET_URLS: + try: + print(f"downloading {url}") + req = urllib.request.Request(url, headers={"User-Agent": "sqlite-vector-recall/1.0"}) + with urllib.request.urlopen(req) as response, open(path, "wb") as out: + while True: + chunk = response.read(1024 * 1024) + if not chunk: + break + out.write(chunk) + last_error = None + break + except Exception as exc: + last_error = exc + try: + os.remove(path) + except OSError: + pass + if last_error is not None: + raise last_error + return path + + +def blob(row): + import numpy as np + + return sqlite3.Binary(np.ascontiguousarray(row, dtype=" #include +#include #include #include #include +#include #include "sqlite3.h" #include "sqlite-vector.h" @@ -70,6 +72,37 @@ static int setup_table(sqlite3 *db, const char *tbl, const char *type, return 0; } +static int setup_f32_blob_table(sqlite3 *db, const char *tbl, const char *distance, + int dim, const float *rows, int n) { + char sql[512]; + sqlite3_stmt *stmt = NULL; + + snprintf(sql, sizeof(sql), "CREATE TABLE \"%s\" (id INTEGER PRIMARY KEY, v BLOB);", tbl); + if (exec_sql(db, sql) != SQLITE_OK) return -1; + + snprintf(sql, sizeof(sql), "INSERT INTO \"%s\" (id, v) VALUES (?1, ?2);", tbl); + int rc = sqlite3_prepare_v2(db, sql, -1, &stmt, NULL); + if (rc != SQLITE_OK) return -1; + + for (int i = 0; i < n; ++i) { + sqlite3_reset(stmt); + sqlite3_clear_bindings(stmt); + sqlite3_bind_int(stmt, 1, i + 1); + sqlite3_bind_blob(stmt, 2, rows + (size_t)i * dim, dim * (int)sizeof(float), SQLITE_STATIC); + rc = sqlite3_step(stmt); + if (rc != SQLITE_DONE) { + sqlite3_finalize(stmt); + return -1; + } + } + sqlite3_finalize(stmt); + + snprintf(sql, sizeof(sql), + "SELECT vector_init('%s', 'v', 'type=f32,dimension=%d,distance=%s');", + tbl, dim, distance); + return exec_sql(db, sql) == SQLITE_OK ? 0 : -1; +} + /* ---------- Callback helpers for querying results ---------- */ typedef struct { @@ -100,6 +133,45 @@ static int scan_cb_col0(void *ctx, int ncols, char **vals, char **names) { return 0; } +static int exec_scan_sql(sqlite3 *db, const char *sql, sqlite3_callback cb, void *ctx) { + char *err = NULL; + int rc = sqlite3_exec(db, sql, cb, ctx, &err); + if (rc != SQLITE_OK) { + printf(" SQL error (%d): %s\n Statement: %s\n", rc, err ? err : "unknown", sql); + sqlite3_free(err); + } + return rc; +} + +static int exec_bound_scan_sql(sqlite3 *db, const char *sql, const float *query, int dim, scan_result *r) { + sqlite3_stmt *stmt = NULL; + int rc = sqlite3_prepare_v2(db, sql, -1, &stmt, NULL); + if (rc != SQLITE_OK) return rc; + + rc = sqlite3_bind_blob(stmt, 1, query, dim * (int)sizeof(float), SQLITE_STATIC); + if (rc == SQLITE_OK) { + while ((rc = sqlite3_step(stmt)) == SQLITE_ROW) { + if (r->count < 64) { + r->ids[r->count] = sqlite3_column_int(stmt, 0); + r->distances[r->count] = sqlite3_column_double(stmt, 1); + } + r->count++; + } + if (rc == SQLITE_DONE) rc = SQLITE_OK; + } + + int frc = sqlite3_finalize(stmt); + return rc == SQLITE_OK ? frc : rc; +} + +static int open_stmt_count(sqlite3 *db) { + int count = 0; + for (sqlite3_stmt *stmt = sqlite3_next_stmt(db, NULL); stmt; stmt = sqlite3_next_stmt(db, stmt)) { + count++; + } + return count; +} + /* ---------- Test: basics ---------- */ static void test_basics(sqlite3 *db) { @@ -132,6 +204,20 @@ static void test_basics(sqlite3 *db) { } sqlite3_finalize(stmt); } + + /* vector_turboquant_backend() */ + { + sqlite3_stmt *stmt; + int rc = sqlite3_prepare_v2(db, "SELECT vector_turboquant_backend();", -1, &stmt, NULL); + ASSERT(rc == SQLITE_OK, "vector_turboquant_backend() prepares"); + if (rc == SQLITE_OK) { + rc = sqlite3_step(stmt); + ASSERT(rc == SQLITE_ROW, "vector_turboquant_backend() returns a row"); + const char *v = (const char *)sqlite3_column_text(stmt, 0); + ASSERT(v != NULL && strlen(v) > 0, "vector_turboquant_backend() returns non-empty text"); + } + sqlite3_finalize(stmt); + } } /* ---------- Test: vector_full_scan for a given (type, distance) pair ---------- */ @@ -275,6 +361,287 @@ static void test_quantize_scan(sqlite3 *db, const char *type, const char *qtype, } } +static void test_turboquant(sqlite3 *db) { + printf("\n=== TurboQuant ===\n"); + + char primary_tbl[64] = "tq_f32_DOT_4"; + const char *tbl = primary_tbl; + const int dim = 16; + const char *vecs[] = { + "[1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0]", + "[-1, -1, -1, -1, -1, -1, -1, -1, 0, 0, 0, 0, 0, 0, 0, 0]", + "[0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1]", + "[0, 0, 0, 0, 0, 0, 0, 0, -1, -1, -1, -1, -1, -1, -1, -1]", + "[1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0]", + "[0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1]", + "[2, 2, 2, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]", + "[0, 0, 0, 0, 0, 0, 0, 0, 2, 2, 2, 2, 0, 0, 0, 0]" + }; + const char *query = "[1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0]"; + char sql[1024]; + + ASSERT(open_stmt_count(db) == 0, "no open statements before TurboQuant test"); + + const char *distances[] = {"DOT", "COSINE", "L2", "SQUARED_L2"}; + const int bits_values[] = {2, 3, 4}; + for (int d = 0; d < 4; ++d) { + for (int b = 0; b < 3; ++b) { + int bits = bits_values[b]; + char tname[64]; + char msg[160]; + snprintf(tname, sizeof(tname), "tq_f32_%s_%d", distances[d], bits); + if (setup_table(db, tname, "f32", distances[d], dim, vecs, 8) != 0) { + snprintf(msg, sizeof(msg), "TurboQuant setup %s qbits=%d", distances[d], bits); + ASSERT(0, msg); + return; + } + + snprintf(sql, sizeof(sql), "SELECT vector_quantize('%s', 'v', 'qtype=TURBO,qbits=%d,max_memory=96');", tname, bits); + snprintf(msg, sizeof(msg), "vector_quantize TurboQuant %s qbits=%d executes", distances[d], bits); + ASSERT(exec_sql(db, sql) == SQLITE_OK, msg); + + { + sqlite3_stmt *stmt = NULL; + snprintf(sql, sizeof(sql), "SELECT vector_quantize_memory('%s', 'v');", tname); + int rc = sqlite3_prepare_v2(db, sql, -1, &stmt, NULL); + ASSERT(rc == SQLITE_OK, "TurboQuant memory prepares"); + if (rc == SQLITE_OK && sqlite3_step(stmt) == SQLITE_ROW) { + sqlite3_int64 bytes = sqlite3_column_int64(stmt, 0); + sqlite3_int64 expected = (sqlite3_int64)8 * (sqlite3_int64)(sizeof(int64_t) + sizeof(float) + ((dim * bits + 7) / 8)); + snprintf(msg, sizeof(msg), "TurboQuant qbits=%d uses expected compact storage", bits); + ASSERT(bytes == expected, msg); + } else { + ASSERT(0, "TurboQuant memory returns a row"); + } + sqlite3_finalize(stmt); + } + + scan_result exact = {0}; + snprintf(sql, sizeof(sql), + "SELECT id, distance FROM vector_full_scan('%s', 'v', vector_as_f32('%s'), 1);", + tname, query); + snprintf(msg, sizeof(msg), "TurboQuant exact top-1 executes %s qbits=%d", distances[d], bits); + ASSERT(sqlite3_exec(db, sql, scan_cb, &exact, NULL) == SQLITE_OK, msg); + + scan_result approx = {0}; + snprintf(sql, sizeof(sql), + "SELECT id, distance FROM vector_quantize_scan('%s', 'v', vector_as_f32('%s'), 1);", + tname, query); + snprintf(msg, sizeof(msg), "TurboQuant top-1 executes %s qbits=%d", distances[d], bits); + ASSERT(sqlite3_exec(db, sql, scan_cb, &approx, NULL) == SQLITE_OK, msg); + snprintf(msg, sizeof(msg), "TurboQuant top-1 matches full scan %s qbits=%d", distances[d], bits); + ASSERT(exact.count == 1 && approx.count == 1 && exact.ids[0] == approx.ids[0], msg); + if (strcmp(distances[d], "DOT") == 0) { + ASSERT(approx.distances[0] < 0.0, "TurboQuant DOT distance uses negative score convention"); + } + + if (strcmp(distances[d], "DOT") == 0 && bits == 4) { + snprintf(primary_tbl, sizeof(primary_tbl), "%s", tname); + tbl = primary_tbl; + snprintf(sql, sizeof(sql), "SELECT vector_quantize_preload('%s', 'v');", tbl); + ASSERT(exec_sql(db, sql) == SQLITE_OK, "TurboQuant preload executes"); + + memset(&approx, 0, sizeof(approx)); + snprintf(sql, sizeof(sql), + "SELECT id, distance FROM vector_quantize_scan('%s', 'v', vector_as_f32('%s'), 3);", + tbl, query); + ASSERT(exec_scan_sql(db, sql, scan_cb, &approx) == SQLITE_OK, "preloaded TurboQuant top-k executes"); + ASSERT(approx.count == 3, "preloaded TurboQuant returns k rows"); + + scan_result streamed = {0}; + snprintf(sql, sizeof(sql), + "SELECT id, distance FROM vector_quantize_scan('%s', 'v', vector_as_f32('%s')) ORDER BY distance LIMIT 3;", + tbl, query); + ASSERT(exec_scan_sql(db, sql, scan_cb, &streamed) == SQLITE_OK, "TurboQuant streaming scan executes"); + ASSERT(streamed.count == 3, "TurboQuant streaming scan returns rows"); + } + } + } + + snprintf(sql, sizeof(sql), "SELECT vector_quantize('%s', 'v', 'qtype=TURBO2');", tbl); + ASSERT(exec_sql(db, sql) == SQLITE_OK, "vector_quantize qtype=TURBO2 executes"); + + snprintf(sql, sizeof(sql), "SELECT vector_quantize('%s', 'v', 'qtype=TURBO,qbits=5');", tbl); + { + char *err = NULL; + int rc = sqlite3_exec(db, sql, NULL, NULL, &err); + ASSERT(rc != SQLITE_OK, "TurboQuant rejects unsupported qbits"); + sqlite3_free(err); + } + + { + const char *wr = "tq_without_rowid"; + snprintf(sql, sizeof(sql), "CREATE TABLE \"%s\" (id INTEGER PRIMARY KEY, v BLOB) WITHOUT ROWID;", wr); + ASSERT(exec_sql(db, sql) == SQLITE_OK, "TurboQuant WITHOUT ROWID table creates"); + for (int i = 0; i < 8; ++i) { + snprintf(sql, sizeof(sql), "INSERT INTO \"%s\" (id, v) VALUES (%d, vector_as_f32('%s'));", wr, i + 1, vecs[i]); + ASSERT(exec_sql(db, sql) == SQLITE_OK, "TurboQuant WITHOUT ROWID insert"); + } + snprintf(sql, sizeof(sql), "SELECT vector_init('%s', 'v', 'type=f32,dimension=%d,distance=DOT');", wr, dim); + ASSERT(exec_sql(db, sql) == SQLITE_OK, "TurboQuant WITHOUT ROWID vector_init"); + snprintf(sql, sizeof(sql), "SELECT vector_quantize('%s', 'v', 'qtype=TURBO,qbits=2');", wr); + ASSERT(exec_sql(db, sql) == SQLITE_OK, "TurboQuant WITHOUT ROWID quantize"); + scan_result wr_result = {0}; + snprintf(sql, sizeof(sql), "SELECT id, distance FROM vector_quantize_scan('%s', 'v', vector_as_f32('%s'), 1);", wr, query); + ASSERT(sqlite3_exec(db, sql, scan_cb, &wr_result, NULL) == SQLITE_OK, "TurboQuant WITHOUT ROWID scan"); + ASSERT(wr_result.count == 1 && wr_result.ids[0] == 1, "TurboQuant WITHOUT ROWID top-1"); + } + + { + const char *corrupt = "tq_corrupt"; + if (setup_table(db, corrupt, "f32", "DOT", dim, vecs, 8) == 0) { + snprintf(sql, sizeof(sql), "SELECT vector_quantize('%s', 'v', 'qtype=TURBO,qbits=4,max_memory=0');", corrupt); + ASSERT(exec_sql(db, sql) == SQLITE_OK, "TurboQuant corrupt-fixture quantize"); + snprintf(sql, sizeof(sql), "UPDATE \"vector0_%s_v\" SET data = substr(data, 1, 3) WHERE rowid = (SELECT rowid FROM \"vector0_%s_v\" LIMIT 1);", corrupt, corrupt); + ASSERT(exec_sql(db, sql) == SQLITE_OK, "TurboQuant corrupt-fixture truncates blob"); + snprintf(sql, sizeof(sql), "SELECT id, distance FROM vector_quantize_scan('%s', 'v', vector_as_f32('%s'), 1);", corrupt, query); + char *err = NULL; + int rc = sqlite3_exec(db, sql, NULL, NULL, &err); + ASSERT(rc != SQLITE_OK, "TurboQuant rejects undersized quant blob"); + sqlite3_free(err); + } else { + ASSERT(0, "TurboQuant corrupt-fixture setup"); + } + } + + ASSERT(open_stmt_count(db) == 0, "no open statements after TurboQuant test"); +} + +static void test_turboquant_edge_dimensions(sqlite3 *db) { + printf("\n=== TurboQuant Edge Dimensions ===\n"); + + sqlite3 *edge_db = NULL; + char *errmsg = NULL; + int rc = sqlite3_open(":memory:", &edge_db); + ASSERT(rc == SQLITE_OK, "TurboQuant edge-dimension database opens"); + if (rc != SQLITE_OK) return; + ASSERT(sqlite3_vector_init(edge_db, &errmsg, NULL) == SQLITE_OK, "TurboQuant edge-dimension extension init"); + sqlite3_free(errmsg); + db = edge_db; + + const int dims[] = {1, 3, 7, 15, 17, 769}; + const int bits_values[] = {2, 3, 4}; + const char *distances[] = {"DOT", "COSINE", "L2", "SQUARED_L2"}; + const int nrows = 4; + char sql[512]; + char msg[192]; + + ASSERT(open_stmt_count(db) == 0, "no open statements before TurboQuant edge-dimension test"); + + for (int di = 0; di < (int)(sizeof(dims) / sizeof(dims[0])); ++di) { + int dim = dims[di]; + float *rows = (float *)sqlite3_malloc64((sqlite3_uint64)nrows * (sqlite3_uint64)dim * sizeof(float)); + if (!rows) { + ASSERT(0, "TurboQuant edge-dimension fixture allocates rows"); + sqlite3_close(edge_db); + return; + } + + for (int j = 0; j < dim; ++j) { + rows[j] = 1.0f; + rows[dim + j] = -1.0f; + rows[2 * dim + j] = 0.0f; + rows[3 * dim + j] = 0.0f; + } + rows[3 * dim + (dim > 1 ? dim - 1 : 0)] = dim > 1 ? 1.0f : -0.5f; + + for (int d = 0; d < 4; ++d) { + for (int b = 0; b < 3; ++b) { + int bits = bits_values[b]; + char tname[80]; + snprintf(tname, sizeof(tname), "tq_edge_%s_%d_%d", distances[d], dim, bits); + + if (setup_f32_blob_table(db, tname, distances[d], dim, rows, nrows) != 0) { + snprintf(msg, sizeof(msg), "TurboQuant edge setup %s dim=%d qbits=%d", distances[d], dim, bits); + ASSERT(0, msg); + sqlite3_free(rows); + sqlite3_close(edge_db); + return; + } + + snprintf(sql, sizeof(sql), "SELECT vector_quantize('%s', 'v', 'qtype=TURBO,qbits=%d,max_memory=0');", tname, bits); + snprintf(msg, sizeof(msg), "TurboQuant edge quantize %s dim=%d qbits=%d", distances[d], dim, bits); + ASSERT(exec_sql(db, sql) == SQLITE_OK, msg); + + scan_result exact = {0}; + snprintf(sql, sizeof(sql), "SELECT id, distance FROM vector_full_scan('%s', 'v', ?1, 1);", tname); + snprintf(msg, sizeof(msg), "TurboQuant edge exact scan %s dim=%d qbits=%d", distances[d], dim, bits); + ASSERT(exec_bound_scan_sql(db, sql, rows, dim, &exact) == SQLITE_OK, msg); + snprintf(msg, sizeof(msg), "TurboQuant edge exact returns one row %s dim=%d qbits=%d", distances[d], dim, bits); + ASSERT(exact.count == 1, msg); + + scan_result approx = {0}; + snprintf(sql, sizeof(sql), "SELECT id, distance FROM vector_quantize_scan('%s', 'v', ?1, 1);", tname); + snprintf(msg, sizeof(msg), "TurboQuant edge scan %s dim=%d qbits=%d", distances[d], dim, bits); + ASSERT(exec_bound_scan_sql(db, sql, rows, dim, &approx) == SQLITE_OK, msg); + + snprintf(msg, sizeof(msg), "TurboQuant edge scan returns one row %s dim=%d qbits=%d", distances[d], dim, bits); + ASSERT(approx.count == 1 && approx.ids[0] >= 1 && approx.ids[0] <= nrows, msg); + + scan_result streamed = {0}; + snprintf(sql, sizeof(sql), "SELECT id, distance FROM vector_quantize_scan('%s', 'v', ?1) LIMIT 2;", tname); + snprintf(msg, sizeof(msg), "TurboQuant edge streaming scan %s dim=%d qbits=%d", distances[d], dim, bits); + ASSERT(exec_bound_scan_sql(db, sql, rows, dim, &streamed) == SQLITE_OK, msg); + snprintf(msg, sizeof(msg), "TurboQuant edge streaming returns rows %s dim=%d qbits=%d", distances[d], dim, bits); + ASSERT(streamed.count == 2, msg); + } + } + + sqlite3_free(rows); + } + + ASSERT(open_stmt_count(db) == 0, "no open statements after TurboQuant edge-dimension test"); + sqlite3_close(edge_db); +} + +static void test_turboquant_reopen(void) { + printf("\n=== TurboQuant Reopen ===\n"); + + char path[256]; + snprintf(path, sizeof(path), "/tmp/sqlite_vector_turboquant_%ld.db", (long)getpid()); + unlink(path); + + const char *vecs[] = { + "[1, 1, 1, 1, 0, 0, 0, 0]", + "[0, 0, 0, 0, 1, 1, 1, 1]", + "[-1, -1, -1, -1, 0, 0, 0, 0]" + }; + const char *query = "[1, 1, 1, 1, 0, 0, 0, 0]"; + const char *tbl = "tq_reopen"; + char sql[1024]; + sqlite3 *db = NULL; + char *errmsg = NULL; + + int rc = sqlite3_open(path, &db); + ASSERT(rc == SQLITE_OK, "TurboQuant reopen database opens"); + ASSERT(sqlite3_vector_init(db, &errmsg, NULL) == SQLITE_OK, "TurboQuant reopen extension init"); + if (setup_table(db, tbl, "f32", "DOT", 8, vecs, 3) == 0) { + snprintf(sql, sizeof(sql), "SELECT vector_quantize('%s', 'v', 'qtype=TURBO,qbits=2,max_memory=0');", tbl); + ASSERT(exec_sql(db, sql) == SQLITE_OK, "TurboQuant reopen quantize"); + } else { + ASSERT(0, "TurboQuant reopen setup"); + } + ASSERT(open_stmt_count(db) == 0, "TurboQuant reopen no open statements before close"); + sqlite3_close(db); + db = NULL; + + rc = sqlite3_open(path, &db); + ASSERT(rc == SQLITE_OK, "TurboQuant reopen database reopens"); + errmsg = NULL; + ASSERT(sqlite3_vector_init(db, &errmsg, NULL) == SQLITE_OK, "TurboQuant reopen extension re-init"); + snprintf(sql, sizeof(sql), "SELECT vector_init('%s', 'v', 'type=f32,dimension=8,distance=DOT');", tbl); + ASSERT(exec_sql(db, sql) == SQLITE_OK, "TurboQuant reopen vector_init reloads metadata"); + + scan_result r = {0}; + snprintf(sql, sizeof(sql), "SELECT id, distance FROM vector_quantize_scan('%s', 'v', vector_as_f32('%s'), 1);", tbl, query); + ASSERT(sqlite3_exec(db, sql, scan_cb, &r, NULL) == SQLITE_OK, "TurboQuant reopen scan executes"); + ASSERT(r.count == 1 && r.ids[0] == 1, "TurboQuant reopen top-1"); + ASSERT(open_stmt_count(db) == 0, "TurboQuant reopen no open statements after scan"); + + sqlite3_close(db); + unlink(path); +} + /* ---------- Test vectors ---------- */ /* 4-dimensional float vectors for numeric types */ @@ -715,6 +1082,10 @@ int main(void) { test_quantize_scan(db, "bit", "1BIT", 8, bit_vecs, bit_nvecs, bit_query); } + test_turboquant(db); + test_turboquant_edge_dimensions(db); + test_turboquant_reopen(); + /* 4. Streaming ORDER BY (regression test for issue #43) */ printf("\n=== Streaming ORDER BY ===\n"); {