diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index aacda68..46c3aaa 100644 --- a/.github/workflows/main.yml +++ b/.github/workflows/main.yml @@ -203,6 +203,69 @@ jobs: path: dist/vector.* if-no-files-found: error + avx512: + # The AVX-512 kernels are compiled out of every other job: the plain `unittest` target + # builds all sources in one invocation, so __AVX512F__ is undefined and the suite + # silently exercises the scalar fallback. This job builds them for real and asserts + # which backend was installed, so a fallback fails instead of passing quietly. + # + # GitHub's hosted fleet is mixed - some runners have AVX-512, some do not, and there + # is no way to request one (actions/runner#1069). When this runner has it we run the + # kernels natively; otherwise we run them under Intel SDE, which emulates the ISA + # deterministically. Either way EXPECT_BACKEND makes the job fail if AVX-512 is not + # what actually ran. + name: avx512 kernels + if: ${{ !contains(github.event.head_commit.message, '[auto-update]') }} + runs-on: ubuntu-24.04 + timeout-minutes: 30 + env: + # Intel SDE, from Intel's own download mirror. Bump both together; the job fails on + # a checksum mismatch rather than running an unverified binary. + SDE_URL: https://downloadmirror.intel.com/924984/sde-external-10.13.1-2026-07-28-lin.tar.xz + SDE_SHA256: 94e97d623fec54385686e1e7ba65ebc9941748c05ee451423948334892bf2b50 + steps: + - uses: actions/checkout@v4.2.2 + + - name: does this runner have AVX-512? + id: cpu + run: | + grep -m1 '^model name' /proc/cpuinfo + missing="" + for f in avx512f avx512bw avx512vl avx512dq; do + grep -qw "$f" /proc/cpuinfo || missing="$missing $f" + done + if [ -z "$missing" ]; then + echo "native=true" >> "$GITHUB_OUTPUT" + echo "::notice title=AVX-512::this runner has AVX-512, running the kernels on hardware" + else + echo "native=false" >> "$GITHUB_OUTPUT" + echo "::notice title=AVX-512::this runner is missing$missing, running the kernels under Intel SDE" + fi + + - name: install Intel SDE + if: steps.cpu.outputs.native == 'false' + run: | + curl -fsSL -o /tmp/sde.tar.xz "$SDE_URL" + echo "$SDE_SHA256 /tmp/sde.tar.xz" | sha256sum -c - + mkdir -p /tmp/sde + tar -xJf /tmp/sde.tar.xz -C /tmp/sde --strip-components=1 + /tmp/sde/sde64 --version | head -2 + + - name: run the suite on the AVX-512 kernels + run: | + if [ "${{ steps.cpu.outputs.native }}" = "true" ]; then + make unittest-simd EXPECT_BACKEND=AVX512 + else + make unittest-simd EXPECT_BACKEND=AVX512 RUNNER="/tmp/sde/sde64 -skx --" + fi + + - name: run the suite on the AVX2 kernels + # Only meaningful where AVX-512 is absent: the runtime check prefers AVX-512 + # whenever the CPU has it, so asserting AVX2 there would correctly fail. Cheap + # extra coverage either way - until now no job built these kernels at all. + if: steps.cpu.outputs.native == 'false' + run: make unittest-simd EXPECT_BACKEND=AVX2 + release: runs-on: ubuntu-22.04 name: release diff --git a/API.md b/API.md index 37f0ca2..ec6c8d2 100644 --- a/API.md +++ b/API.md @@ -104,12 +104,22 @@ This ensures that each vector can be uniquely identified and efficiently referen * `COSINE` * `DOT` * `L1` - * `HAMMING` + * `HAMMING` (only valid with `type=1BIT`) +* `normalized`: Set to `1` to declare that every stored vector is unit length. With + `type=FLOAT32` and `distance=COSINE` this lets a full-precision scan compute + `1 - dot` instead of the full cosine, dropping two thirds of the arithmetic from the + inner loop; the query vector is normalized once per scan, so the reported distances are + unchanged. It is an assertion, not a request: if the stored vectors are *not* unit + length the distances will be wrong. Quantized scans ignore it, because the quantized + index holds scaled integers whose norm is not 1. Default `0`. **Example:** ```sql SELECT vector_init('documents', 'embedding', 'dimension=384,type=FLOAT32,distance=cosine'); + +-- embeddings already normalized by the model: faster cosine, same results +SELECT vector_init('documents', 'embedding', 'dimension=384,type=FLOAT32,distance=cosine,normalized=1'); ``` --- diff --git a/Makefile b/Makefile index a7132af..8c08ae7 100644 --- a/Makefile +++ b/Makefile @@ -102,6 +102,18 @@ ifneq (,$(findstring rv64,$(ARCH))) CFLAGS += -march=$(ARCH) endif +# The AVX2 and AVX-512 kernels are guarded by __AVX2__ / __AVX512F__, so without these +# flags they compile to nothing and every x86 build falls back to the scalar kernels. +# Enable the ISA for those two translation units only, so the baseline target of the +# rest of the extension is unchanged: the runtime check in init_distance_functions() +# still decides which set gets installed. Probing the compiler keeps this a no-op on +# non-x86 targets and on multi-arch (universal) builds. +AVX2_CFLAGS := $(shell $(CC) $(CFLAGS) -mavx2 -mfma -E -x c /dev/null >/dev/null 2>&1 && echo -mavx2 -mfma) +AVX512_CFLAGS := $(shell $(CC) $(CFLAGS) -mavx512f -mavx512bw -mavx512vl -mavx512dq -E -x c /dev/null >/dev/null 2>&1 && echo -mavx512f -mavx512bw -mavx512vl -mavx512dq) + +$(BUILD_DIR)/distance-avx2.o: ISA_CFLAGS := $(AVX2_CFLAGS) +$(BUILD_DIR)/distance-avx512.o: ISA_CFLAGS := $(AVX512_CFLAGS) + # Windows .def file generation $(DEF_FILE): ifeq ($(PLATFORM),windows) @@ -129,7 +141,7 @@ endif # Object files $(BUILD_DIR)/%.o: %.c - $(CC) $(CFLAGS) -O3 -fPIC -c $< -o $@ + $(CC) $(CFLAGS) $(ISA_CFLAGS) -O3 -fPIC -c $< -o $@ test: $(TARGET) $(SQLITE3) ":memory:" -cmd ".bail on" ".load ./dist/vector" "SELECT vector_version();" @@ -139,6 +151,36 @@ unittest: $(CC) $(CFLAGS) -DSQLITE_CORE -O2 $(TEST_SRC) -o $(BUILD_DIR)/test_vector -lm -lpthread ./$(BUILD_DIR)/test_vector +# The unittest target above builds every source in a single invocation, which leaves +# __AVX2__ and __AVX512F__ undefined: those kernels compile to nothing and the suite +# silently exercises the scalar fallback instead. This target compiles per translation +# unit the way the extension does, so the SIMD backends are actually under test. +# +# make unittest-simd run on whatever this CPU supports +# make unittest-simd EXPECT_BACKEND=AVX512 fail unless AVX-512 was installed +# make unittest-simd RUNNER="sde64 -spr --" run under an emulator +UNITTEST_OBJ = $(patsubst %.c, $(BUILD_DIR)/ut-%.o, $(notdir $(SRC_FILES))) $(BUILD_DIR)/ut-sqlite3.o + +$(BUILD_DIR)/ut-distance-avx2.o: ISA_CFLAGS := $(AVX2_CFLAGS) +$(BUILD_DIR)/ut-distance-avx512.o: ISA_CFLAGS := $(AVX512_CFLAGS) + +$(BUILD_DIR)/ut-%.o: %.c + $(CC) $(CFLAGS) $(ISA_CFLAGS) -DSQLITE_CORE -O2 -c $< -o $@ + +$(BUILD_DIR)/backend: test/backend.c $(UNITTEST_OBJ) + $(CC) $(CFLAGS) -DSQLITE_CORE -O2 $< $(UNITTEST_OBJ) -o $@ -lm -lpthread + +$(BUILD_DIR)/test_vector_simd: test/test_vector.c $(UNITTEST_OBJ) + $(CC) $(CFLAGS) -DSQLITE_CORE -O2 $< $(UNITTEST_OBJ) -o $@ -lm -lpthread + +# RUNNER wraps both binaries, so an emulator sees the same build the assertion checked +RUNNER ?= +EXPECT_BACKEND ?= + +unittest-simd: $(BUILD_DIR)/backend $(BUILD_DIR)/test_vector_simd + $(RUNNER) ./$(BUILD_DIR)/backend $(EXPECT_BACKEND) + $(RUNNER) ./$(BUILD_DIR)/test_vector_simd + # Clean up generated files clean: rm -rf $(BUILD_DIR)/* $(DIST_DIR)/* *.gcda *.gcno *.gcov *.sqlite diff --git a/README.md b/README.md index 63e4f7e..99298a2 100644 --- a/README.md +++ b/README.md @@ -119,6 +119,10 @@ INSERT INTO images (embedding, label) VALUES (vector_as_f32('[0.3, 1.0, 0.9, 3.2 -- distance=L1, distance=COSINE, distance=DOT, distance=SQUARED_L2, or distance=HAMMING. SELECT vector_init('images', 'embedding', 'type=FLOAT32,dimension=384'); +-- If your embeddings are already unit length, say so: FLOAT32 cosine scans then compute +-- 1 - dot instead of the full cosine, with the same results. +-- SELECT vector_init('images', 'embedding', 'type=FLOAT32,dimension=384,distance=COSINE,normalized=1'); + -- Quantize vector SELECT vector_quantize('images', 'embedding'); diff --git a/src/distance-avx2.c b/src/distance-avx2.c index 4108467..243d772 100644 --- a/src/distance-avx2.c +++ b/src/distance-avx2.c @@ -20,11 +20,6 @@ extern const char *turbo_lut_backend_name; #define _mm256_abs_ps(x) _mm256_andnot_ps(_mm256_set1_ps(-0.0f), (x)) -static inline __m256 mm256_abs_ps(__m256 x) { - const __m256 mask = _mm256_castsi256_ps(_mm256_set1_epi32(0x7FFFFFFF)); - return _mm256_and_ps(x, mask); -} - static inline double hsum256d(__m256d v) { __m128d lo = _mm256_castpd256_pd128(v); __m128d hi = _mm256_extractf128_pd(v, 1); @@ -66,31 +61,54 @@ static inline bool block_has_l2_inf_mismatch_bf16_8(const uint16_t* a, const uin // MARK: - FLOAT32 - +// A single accumulator makes the loop one dependency chain: an FMA has around four cycles +// of latency, so it retires one vector every four cycles however many FMA ports the core +// has. Four independent accumulators keep them fed and still fit sixteen YMM registers. +#if defined(__FMA__) +#define MM256_FMA_PS(_acc, _x, _y) _mm256_fmadd_ps((_x), (_y), (_acc)) +#else +#define MM256_FMA_PS(_acc, _x, _y) _mm256_add_ps((_acc), _mm256_mul_ps((_x), (_y))) +#endif + +static inline float hsum256_ps (__m256 v) { + __m128 lo = _mm256_castps256_ps128(v); + __m128 hi = _mm256_extractf128_ps(v, 1); + __m128 s = _mm_add_ps(lo, hi); + s = _mm_add_ps(s, _mm_movehl_ps(s, s)); + s = _mm_add_ss(s, _mm_shuffle_ps(s, s, 0x55)); + return _mm_cvtss_f32(s); +} + static inline float float32_distance_l2_impl_avx2 (const void *v1, const void *v2, int n, bool use_sqrt) { const float *a = (const float *)v1; const float *b = (const float *)v2; - - __m256 acc = _mm256_setzero_ps(); + + __m256 acc0 = _mm256_setzero_ps(), acc1 = acc0, acc2 = acc0, acc3 = acc0; int i = 0; + for (; i <= n - 32; i += 32) { + __m256 d0 = _mm256_sub_ps(_mm256_loadu_ps(a + i ), _mm256_loadu_ps(b + i )); + __m256 d1 = _mm256_sub_ps(_mm256_loadu_ps(a + i + 8), _mm256_loadu_ps(b + i + 8)); + __m256 d2 = _mm256_sub_ps(_mm256_loadu_ps(a + i + 16), _mm256_loadu_ps(b + i + 16)); + __m256 d3 = _mm256_sub_ps(_mm256_loadu_ps(a + i + 24), _mm256_loadu_ps(b + i + 24)); + acc0 = MM256_FMA_PS(acc0, d0, d0); + acc1 = MM256_FMA_PS(acc1, d1, d1); + acc2 = MM256_FMA_PS(acc2, d2, d2); + acc3 = MM256_FMA_PS(acc3, d3, d3); + } for (; i <= n - 8; i += 8) { - __m256 va = _mm256_loadu_ps(a + i); - __m256 vb = _mm256_loadu_ps(b + i); - __m256 diff = _mm256_sub_ps(va, vb); - acc = _mm256_add_ps(acc, _mm256_mul_ps(diff, diff)); + __m256 d = _mm256_sub_ps(_mm256_loadu_ps(a + i), _mm256_loadu_ps(b + i)); + acc0 = MM256_FMA_PS(acc0, d, d); } - float temp[8]; - _mm256_storeu_ps(temp, acc); - float total = temp[0] + temp[1] + temp[2] + temp[3] + - temp[4] + temp[5] + temp[6] + temp[7]; + float total = hsum256_ps(_mm256_add_ps(_mm256_add_ps(acc0, acc1), _mm256_add_ps(acc2, acc3))); for (; i < n; ++i) { float d = a[i] - b[i]; total += d * d; } - return use_sqrt ? sqrtf((float)total) : (float)total; + return use_sqrt ? sqrtf(total) : total; } float float32_distance_l2_avx2 (const void *v1, const void *v2, int n) { @@ -104,21 +122,26 @@ float float32_distance_l2_squared_avx2 (const void *v1, const void *v2, int n) { float float32_distance_l1_avx2 (const void *v1, const void *v2, int n) { const float *a = (const float *)v1; const float *b = (const float *)v2; - - __m256 acc = _mm256_setzero_ps(); + + __m256 acc0 = _mm256_setzero_ps(), acc1 = acc0, acc2 = acc0, acc3 = acc0; int i = 0; + for (; i <= n - 32; i += 32) { + __m256 d0 = _mm256_sub_ps(_mm256_loadu_ps(a + i ), _mm256_loadu_ps(b + i )); + __m256 d1 = _mm256_sub_ps(_mm256_loadu_ps(a + i + 8), _mm256_loadu_ps(b + i + 8)); + __m256 d2 = _mm256_sub_ps(_mm256_loadu_ps(a + i + 16), _mm256_loadu_ps(b + i + 16)); + __m256 d3 = _mm256_sub_ps(_mm256_loadu_ps(a + i + 24), _mm256_loadu_ps(b + i + 24)); + acc0 = _mm256_add_ps(acc0, _mm256_abs_ps(d0)); + acc1 = _mm256_add_ps(acc1, _mm256_abs_ps(d1)); + acc2 = _mm256_add_ps(acc2, _mm256_abs_ps(d2)); + acc3 = _mm256_add_ps(acc3, _mm256_abs_ps(d3)); + } for (; i <= n - 8; i += 8) { - __m256 va = _mm256_loadu_ps(a + i); - __m256 vb = _mm256_loadu_ps(b + i); - __m256 diff = _mm256_sub_ps(va, vb); - acc = _mm256_add_ps(acc, _mm256_abs_ps(diff)); + __m256 d = _mm256_sub_ps(_mm256_loadu_ps(a + i), _mm256_loadu_ps(b + i)); + acc0 = _mm256_add_ps(acc0, _mm256_abs_ps(d)); } - float temp[8]; - _mm256_storeu_ps(temp, acc); - float total = temp[0] + temp[1] + temp[2] + temp[3] + - temp[4] + temp[5] + temp[6] + temp[7]; + float total = hsum256_ps(_mm256_add_ps(_mm256_add_ps(acc0, acc1), _mm256_add_ps(acc2, acc3))); for (; i < n; ++i) { total += fabsf(a[i] - b[i]); @@ -130,20 +153,21 @@ float float32_distance_l1_avx2 (const void *v1, const void *v2, int n) { float float32_distance_dot_avx2 (const void *v1, const void *v2, int n) { const float *a = (const float *)v1; const float *b = (const float *)v2; - - __m256 acc = _mm256_setzero_ps(); + + __m256 acc0 = _mm256_setzero_ps(), acc1 = acc0, acc2 = acc0, acc3 = acc0; int i = 0; + for (; i <= n - 32; i += 32) { + acc0 = MM256_FMA_PS(acc0, _mm256_loadu_ps(a + i ), _mm256_loadu_ps(b + i )); + acc1 = MM256_FMA_PS(acc1, _mm256_loadu_ps(a + i + 8), _mm256_loadu_ps(b + i + 8)); + acc2 = MM256_FMA_PS(acc2, _mm256_loadu_ps(a + i + 16), _mm256_loadu_ps(b + i + 16)); + acc3 = MM256_FMA_PS(acc3, _mm256_loadu_ps(a + i + 24), _mm256_loadu_ps(b + i + 24)); + } for (; i <= n - 8; i += 8) { - __m256 va = _mm256_loadu_ps(a + i); - __m256 vb = _mm256_loadu_ps(b + i); - acc = _mm256_add_ps(acc, _mm256_mul_ps(va, vb)); + acc0 = MM256_FMA_PS(acc0, _mm256_loadu_ps(a + i), _mm256_loadu_ps(b + i)); } - float temp[8]; - _mm256_storeu_ps(temp, acc); - float total = temp[0] + temp[1] + temp[2] + temp[3] + - temp[4] + temp[5] + temp[6] + temp[7]; + float total = hsum256_ps(_mm256_add_ps(_mm256_add_ps(acc0, acc1), _mm256_add_ps(acc2, acc3))); for (; i < n; ++i) { total += a[i] * b[i]; @@ -153,13 +177,45 @@ float float32_distance_dot_avx2 (const void *v1, const void *v2, int n) { } float float32_distance_cosine_avx2 (const void *a, const void *b, int n) { - float dot = -float32_distance_dot_avx2(a, b, n); - float norm_a = sqrtf(-float32_distance_dot_avx2(a, a, n)); - float norm_b = sqrtf(-float32_distance_dot_avx2(b, b, n)); + const float *x = (const float *)a; + const float *y = (const float *)b; + + // one fused pass, not three calls to the dot kernel: the data is read once instead of + // three times, which is what actually costs on anything larger than L1 + __m256 dot0 = _mm256_setzero_ps(), dot1 = dot0; + __m256 na0 = dot0, na1 = dot0; + __m256 nb0 = dot0, nb1 = dot0; + int i = 0; + + for (; i <= n - 16; i += 16) { + __m256 a0 = _mm256_loadu_ps(x + i), a1 = _mm256_loadu_ps(x + i + 8); + __m256 b0 = _mm256_loadu_ps(y + i), b1 = _mm256_loadu_ps(y + i + 8); + dot0 = MM256_FMA_PS(dot0, a0, b0); dot1 = MM256_FMA_PS(dot1, a1, b1); + na0 = MM256_FMA_PS(na0, a0, a0); na1 = MM256_FMA_PS(na1, a1, a1); + nb0 = MM256_FMA_PS(nb0, b0, b0); nb1 = MM256_FMA_PS(nb1, b1, b1); + } + for (; i <= n - 8; i += 8) { + __m256 va = _mm256_loadu_ps(x + i), vb = _mm256_loadu_ps(y + i); + dot0 = MM256_FMA_PS(dot0, va, vb); + na0 = MM256_FMA_PS(na0, va, va); + nb0 = MM256_FMA_PS(nb0, vb, vb); + } + + float dot = hsum256_ps(_mm256_add_ps(dot0, dot1)); + float norm_a = hsum256_ps(_mm256_add_ps(na0, na1)); + float norm_b = hsum256_ps(_mm256_add_ps(nb0, nb1)); + + for (; i < n; ++i) { + float ai = x[i]; + float bi = y[i]; + dot += ai * bi; + norm_a += ai * ai; + norm_b += bi * bi; + } if (norm_a == 0.0f || norm_b == 0.0f) return 1.0f; - float cosine_similarity = dot / (norm_a * norm_b); + float cosine_similarity = dot / (sqrtf(norm_a) * sqrtf(norm_b)); if (cosine_similarity > 1.0f) cosine_similarity = 1.0f; if (cosine_similarity < -1.0f) cosine_similarity = -1.0f; return 1.0f - cosine_similarity; @@ -1052,7 +1108,7 @@ float turbo_lut_dot_avx2 (const uint8_t *packed, float scale, const float *query // MARK: - -void init_distance_functions_avx2 (void) { +bool init_distance_functions_avx2 (void) { #if defined(__AVX2__) || (defined(_MSC_VER) && defined(__AVX2__)) dispatch_distance_table[VECTOR_DISTANCE_L2][VECTOR_TYPE_F32] = float32_distance_l2_avx2; dispatch_distance_table[VECTOR_DISTANCE_L2][VECTOR_TYPE_F16] = float16_distance_l2_avx2; @@ -1089,5 +1145,8 @@ void init_distance_functions_avx2 (void) { distance_backend_name = "AVX2"; turbo_lut_dot_function = turbo_lut_dot_avx2; turbo_lut_backend_name = "AVX2"; + return true; +#else + return false; #endif } diff --git a/src/distance-avx2.h b/src/distance-avx2.h index aa7c39c..f3f1c55 100644 --- a/src/distance-avx2.h +++ b/src/distance-avx2.h @@ -8,10 +8,12 @@ #ifndef __VECTOR_DISTANCE_AVX2__ #define __VECTOR_DISTANCE_AVX2__ +#include #include #include -void init_distance_functions_avx2 (void); +// returns true when the AVX2 kernels were compiled into this build +bool 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 11b7d2f..aff8be3 100644 --- a/src/distance-avx512.c +++ b/src/distance-avx512.c @@ -9,7 +9,7 @@ #include "distance-cpu.h" // Check for AVX512 Foundation (F) and Byte/Word (BW) which are standard on Skylake-X/IceLake+ -#if defined(__AVX512F__) && defined(__AVX512BW__) +#if defined(__AVX512F__) && defined(__AVX512BW__) && defined(__AVX512VL__) && defined(__AVX512DQ__) #include #include #include @@ -75,22 +75,32 @@ static inline bool block_has_l2_inf_mismatch_bf16_16(const uint16_t* a, const ui // MARK: - FLOAT32 - +// A single accumulator makes the loop one dependency chain: an FMA has around four cycles +// of latency, so it retires one vector every four cycles however many FMA ports the core +// has. Four independent accumulators keep them fed; 32 ZMM registers hold them easily. static inline float float32_distance_l2_impl_avx512(const void* v1, const void* v2, int n, bool use_sqrt) { const float* a = (const float*)v1; const float* b = (const float*)v2; - __m512 acc = _mm512_setzero_ps(); + __m512 acc0 = _mm512_setzero_ps(), acc1 = acc0, acc2 = acc0, acc3 = acc0; int i = 0; - // Stride 16 for AVX-512 + for (; i <= n - 64; i += 64) { + __m512 d0 = _mm512_sub_ps(_mm512_loadu_ps(a + i ), _mm512_loadu_ps(b + i )); + __m512 d1 = _mm512_sub_ps(_mm512_loadu_ps(a + i + 16), _mm512_loadu_ps(b + i + 16)); + __m512 d2 = _mm512_sub_ps(_mm512_loadu_ps(a + i + 32), _mm512_loadu_ps(b + i + 32)); + __m512 d3 = _mm512_sub_ps(_mm512_loadu_ps(a + i + 48), _mm512_loadu_ps(b + i + 48)); + acc0 = _mm512_fmadd_ps(d0, d0, acc0); + acc1 = _mm512_fmadd_ps(d1, d1, acc1); + acc2 = _mm512_fmadd_ps(d2, d2, acc2); + acc3 = _mm512_fmadd_ps(d3, d3, acc3); + } for (; i <= n - 16; i += 16) { - __m512 va = _mm512_loadu_ps(a + i); - __m512 vb = _mm512_loadu_ps(b + i); - __m512 diff = _mm512_sub_ps(va, vb); - acc = _mm512_fmadd_ps(diff, diff, acc); + __m512 d = _mm512_sub_ps(_mm512_loadu_ps(a + i), _mm512_loadu_ps(b + i)); + acc0 = _mm512_fmadd_ps(d, d, acc0); } - float total = hsum512_ps(acc); + float total = hsum512_ps(_mm512_add_ps(_mm512_add_ps(acc0, acc1), _mm512_add_ps(acc2, acc3))); for (; i < n; ++i) { float d = a[i] - b[i]; @@ -112,17 +122,25 @@ float float32_distance_l1_avx512(const void* v1, const void* v2, int n) { const float* a = (const float*)v1; const float* b = (const float*)v2; - __m512 acc = _mm512_setzero_ps(); + __m512 acc0 = _mm512_setzero_ps(), acc1 = acc0, acc2 = acc0, acc3 = acc0; int i = 0; + for (; i <= n - 64; i += 64) { + __m512 d0 = _mm512_sub_ps(_mm512_loadu_ps(a + i ), _mm512_loadu_ps(b + i )); + __m512 d1 = _mm512_sub_ps(_mm512_loadu_ps(a + i + 16), _mm512_loadu_ps(b + i + 16)); + __m512 d2 = _mm512_sub_ps(_mm512_loadu_ps(a + i + 32), _mm512_loadu_ps(b + i + 32)); + __m512 d3 = _mm512_sub_ps(_mm512_loadu_ps(a + i + 48), _mm512_loadu_ps(b + i + 48)); + acc0 = _mm512_add_ps(acc0, _mm512_abs_ps(d0)); + acc1 = _mm512_add_ps(acc1, _mm512_abs_ps(d1)); + acc2 = _mm512_add_ps(acc2, _mm512_abs_ps(d2)); + acc3 = _mm512_add_ps(acc3, _mm512_abs_ps(d3)); + } for (; i <= n - 16; i += 16) { - __m512 va = _mm512_loadu_ps(a + i); - __m512 vb = _mm512_loadu_ps(b + i); - __m512 diff = _mm512_sub_ps(va, vb); - acc = _mm512_add_ps(acc, _mm512_abs_ps(diff)); + __m512 d = _mm512_sub_ps(_mm512_loadu_ps(a + i), _mm512_loadu_ps(b + i)); + acc0 = _mm512_add_ps(acc0, _mm512_abs_ps(d)); } - float total = hsum512_ps(acc); + float total = hsum512_ps(_mm512_add_ps(_mm512_add_ps(acc0, acc1), _mm512_add_ps(acc2, acc3))); for (; i < n; ++i) { total += fabsf(a[i] - b[i]); @@ -135,16 +153,20 @@ float float32_distance_dot_avx512(const void* v1, const void* v2, int n) { const float* a = (const float*)v1; const float* b = (const float*)v2; - __m512 acc = _mm512_setzero_ps(); + __m512 acc0 = _mm512_setzero_ps(), acc1 = acc0, acc2 = acc0, acc3 = acc0; int i = 0; + for (; i <= n - 64; i += 64) { + acc0 = _mm512_fmadd_ps(_mm512_loadu_ps(a + i ), _mm512_loadu_ps(b + i ), acc0); + acc1 = _mm512_fmadd_ps(_mm512_loadu_ps(a + i + 16), _mm512_loadu_ps(b + i + 16), acc1); + acc2 = _mm512_fmadd_ps(_mm512_loadu_ps(a + i + 32), _mm512_loadu_ps(b + i + 32), acc2); + acc3 = _mm512_fmadd_ps(_mm512_loadu_ps(a + i + 48), _mm512_loadu_ps(b + i + 48), acc3); + } for (; i <= n - 16; i += 16) { - __m512 va = _mm512_loadu_ps(a + i); - __m512 vb = _mm512_loadu_ps(b + i); - acc = _mm512_fmadd_ps(va, vb, acc); + acc0 = _mm512_fmadd_ps(_mm512_loadu_ps(a + i), _mm512_loadu_ps(b + i), acc0); } - float total = hsum512_ps(acc); + float total = hsum512_ps(_mm512_add_ps(_mm512_add_ps(acc0, acc1), _mm512_add_ps(acc2, acc3))); for (; i < n; ++i) { total += a[i] * b[i]; @@ -154,13 +176,45 @@ float float32_distance_dot_avx512(const void* v1, const void* v2, int n) { } float float32_distance_cosine_avx512(const void* a, const void* b, int n) { - float dot = -float32_distance_dot_avx512(a, b, n); - float norm_a = sqrtf(-float32_distance_dot_avx512(a, a, n)); - float norm_b = sqrtf(-float32_distance_dot_avx512(b, b, n)); + const float* x = (const float*)a; + const float* y = (const float*)b; + + // one fused pass, not three calls to the dot kernel: the data is read once instead of + // three times, which is what actually costs on anything larger than L1 + __m512 dot0 = _mm512_setzero_ps(), dot1 = dot0; + __m512 na0 = dot0, na1 = dot0; + __m512 nb0 = dot0, nb1 = dot0; + int i = 0; + + for (; i <= n - 32; i += 32) { + __m512 a0 = _mm512_loadu_ps(x + i), a1 = _mm512_loadu_ps(x + i + 16); + __m512 b0 = _mm512_loadu_ps(y + i), b1 = _mm512_loadu_ps(y + i + 16); + dot0 = _mm512_fmadd_ps(a0, b0, dot0); dot1 = _mm512_fmadd_ps(a1, b1, dot1); + na0 = _mm512_fmadd_ps(a0, a0, na0); na1 = _mm512_fmadd_ps(a1, a1, na1); + nb0 = _mm512_fmadd_ps(b0, b0, nb0); nb1 = _mm512_fmadd_ps(b1, b1, nb1); + } + for (; i <= n - 16; i += 16) { + __m512 va = _mm512_loadu_ps(x + i), vb = _mm512_loadu_ps(y + i); + dot0 = _mm512_fmadd_ps(va, vb, dot0); + na0 = _mm512_fmadd_ps(va, va, na0); + nb0 = _mm512_fmadd_ps(vb, vb, nb0); + } + + float dot = hsum512_ps(_mm512_add_ps(dot0, dot1)); + float norm_a = hsum512_ps(_mm512_add_ps(na0, na1)); + float norm_b = hsum512_ps(_mm512_add_ps(nb0, nb1)); + + for (; i < n; ++i) { + float ai = x[i]; + float bi = y[i]; + dot += ai * bi; + norm_a += ai * ai; + norm_b += bi * bi; + } if (norm_a == 0.0f || norm_b == 0.0f) return 1.0f; - float cosine_similarity = dot / (norm_a * norm_b); + float cosine_similarity = dot / (sqrtf(norm_a) * sqrtf(norm_b)); if (cosine_similarity > 1.0f) cosine_similarity = 1.0f; if (cosine_similarity < -1.0f) cosine_similarity = -1.0f; return 1.0f - cosine_similarity; @@ -203,7 +257,7 @@ static inline float float16_distance_l2_impl_avx512(const void* v1, const void* __mmask16 mask_valid = mask_a & mask_b; // If not valid, set d to 0.0 - d = _mm512_mask_set1_ps(d, ~mask_valid, 0.0f); + d = _mm512_maskz_mov_ps(mask_valid, d); // Widen to f64 and accumulate __m256 d_lo = _mm512_castps512_ps256(d); @@ -260,7 +314,7 @@ float float16_distance_l1_avx512(const void* v1, const void* v2, int n) { // Zero out NaNs __mmask16 mask_a = _mm512_cmp_ps_mask(va, va, _CMP_ORD_Q); __mmask16 mask_b = _mm512_cmp_ps_mask(vb, vb, _CMP_ORD_Q); - d = _mm512_mask_set1_ps(d, ~(mask_a & mask_b), 0.0f); + d = _mm512_maskz_mov_ps(mask_a & mask_b, d); // Convert to double to accumulate __m256 d_lo = _mm512_castps512_ps256(d); @@ -317,8 +371,8 @@ float float16_distance_dot_avx512(const void* v1, const void* v2, int n) { __mmask16 mask_a = _mm512_cmp_ps_mask(va, va, _CMP_ORD_Q); __mmask16 mask_b = _mm512_cmp_ps_mask(vb, vb, _CMP_ORD_Q); - va = _mm512_mask_set1_ps(va, ~mask_a, 0.0f); - vb = _mm512_mask_set1_ps(vb, ~mask_b, 0.0f); + va = _mm512_maskz_mov_ps(mask_a, va); + vb = _mm512_maskz_mov_ps(mask_b, vb); // This multiply might generate Infs, but we checked scalar first. // We still need to handle the case where standard float math generates Inf from finite * finite? @@ -414,8 +468,8 @@ static inline float bfloat16_distance_l2_impl_avx512(const void* v1, const void* /* zero-out NaNs */ __mmask8 m0 = _mm512_cmp_pd_mask(d0, d0, _CMP_ORD_Q); __mmask8 m1 = _mm512_cmp_pd_mask(d1, d1, _CMP_ORD_Q); - d0 = _mm512_mask_set1_pd(d0, ~m0, 0.0); - d1 = _mm512_mask_set1_pd(d1, ~m1, 0.0); + d0 = _mm512_maskz_mov_pd(m0, d0); + d1 = _mm512_maskz_mov_pd(m1, d1); acc0 = _mm512_fmadd_pd(d0, d0, acc0); acc1 = _mm512_fmadd_pd(d1, d1, acc1); @@ -470,8 +524,8 @@ float bfloat16_distance_l1_avx512(const void* v1, const void* v2, int n) { // NaN -> 0 __mmask8 m0 = _mm512_cmp_pd_mask(d0, d0, _CMP_ORD_Q); __mmask8 m1 = _mm512_cmp_pd_mask(d1, d1, _CMP_ORD_Q); - d0 = _mm512_mask_set1_pd(d0, ~m0, 0.0); - d1 = _mm512_mask_set1_pd(d1, ~m1, 0.0); + d0 = _mm512_maskz_mov_pd(m0, d0); + d1 = _mm512_maskz_mov_pd(m1, d1); acc0 = _mm512_add_pd(acc0, d0); acc1 = _mm512_add_pd(acc1, d1); @@ -518,8 +572,8 @@ float bfloat16_distance_dot_avx512(const void* v1, const void* v2, int n) { // NaN -> 0 __mmask16 ma = _mm512_cmp_ps_mask(af, af, _CMP_ORD_Q); __mmask16 mb = _mm512_cmp_ps_mask(bf, bf, _CMP_ORD_Q); - af = _mm512_mask_set1_ps(af, ~ma, 0.0f); - bf = _mm512_mask_set1_ps(bf, ~mb, 0.0f); + af = _mm512_maskz_mov_ps(ma, af); + bf = _mm512_maskz_mov_ps(mb, bf); __m512 prod = _mm512_mul_ps(af, bf); @@ -969,8 +1023,8 @@ float turbo_lut_dot_avx512 (const uint8_t *packed, float scale, const float *que // MARK: - -void init_distance_functions_avx512(void) { -#if defined(__AVX512F__) && defined(__AVX512BW__) +bool init_distance_functions_avx512(void) { +#if defined(__AVX512F__) && defined(__AVX512BW__) && defined(__AVX512VL__) && defined(__AVX512DQ__) dispatch_distance_table[VECTOR_DISTANCE_L2][VECTOR_TYPE_F32] = float32_distance_l2_avx512; dispatch_distance_table[VECTOR_DISTANCE_L2][VECTOR_TYPE_F16] = float16_distance_l2_avx512; dispatch_distance_table[VECTOR_DISTANCE_L2][VECTOR_TYPE_BF16] = bfloat16_distance_l2_avx512; @@ -1006,5 +1060,8 @@ void init_distance_functions_avx512(void) { distance_backend_name = "AVX512"; turbo_lut_dot_function = turbo_lut_dot_avx512; turbo_lut_backend_name = "AVX512"; + return true; +#else + return false; #endif } diff --git a/src/distance-avx512.h b/src/distance-avx512.h index d8bf661..1eddb78 100644 --- a/src/distance-avx512.h +++ b/src/distance-avx512.h @@ -8,10 +8,12 @@ #ifndef __VECTOR_DISTANCE_AVX512__ #define __VECTOR_DISTANCE_AVX512__ +#include #include #include -void init_distance_functions_avx512 (void); +// returns true when the AVX512 kernels were compiled into this build +bool 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 ba0a358..978cde6 100644 --- a/src/distance-cpu.c +++ b/src/distance-cpu.c @@ -822,7 +822,10 @@ float bit1_distance_hamming_cpu (const void *v1, const void *v2, int n) { bool has_avx512bw = (cpu_info[1] & (1 << 30)); bool has_avx512vl = (cpu_info[1] & (1 << 31)); - return has_avx512f && has_avx512bw && has_avx512vl; + // EBX Bit 17: AVX512DQ, needed by _mm512_extractf32x8_ps in the f16/bf16 kernels + bool has_avx512dq = (cpu_info[1] & (1 << 17)); + + return has_avx512f && has_avx512bw && has_avx512vl && has_avx512dq; #endif } @@ -942,16 +945,15 @@ void init_distance_functions (bool force_cpu) { init_cpu_functions(); if (force_cpu) return; + // each backend reports whether its kernels were actually compiled into this build: + // a tier whose ISA was not enabled at compile time installs nothing and we must keep + // walking down, otherwise an AVX2-capable CPU would end up on the scalar fallback + // even though the SSE2 kernels are available #if defined(__x86_64__) || defined(_M_X64) || defined(__i386__) || defined(_M_IX86) - if (cpu_supports_avx512()) { - init_distance_functions_avx512(); - } - else if (cpu_supports_avx2()) { - init_distance_functions_avx2(); - } - else if (cpu_supports_sse2()) { - init_distance_functions_sse2(); - } + bool installed = false; + if (!installed && cpu_supports_avx512()) installed = init_distance_functions_avx512(); + if (!installed && cpu_supports_avx2()) installed = init_distance_functions_avx2(); + if (!installed && cpu_supports_sse2()) installed = init_distance_functions_sse2(); #elif defined(__ARM_NEON) || defined(__aarch64__) if (cpu_supports_neon()) { init_distance_functions_neon(); diff --git a/src/distance-neon.c b/src/distance-neon.c index 29653a9..3298146 100644 --- a/src/distance-neon.c +++ b/src/distance-neon.c @@ -38,27 +38,48 @@ static inline uint16_t vmaxv_u16_compat(uint16x4_t v) { // MARK: FLOAT32 - +// One accumulator turns the loop into a single dependency chain: FMLA has around four +// cycles of latency, so the loop retires one vector every four cycles however many FMA +// pipes the core has. Four independent accumulators keep them all fed. +#if defined(__aarch64__) || defined(__ARM_FEATURE_FMA) +#define VFMA_F32(_acc, _x, _y) vfmaq_f32((_acc), (_x), (_y)) +#else +#define VFMA_F32(_acc, _x, _y) vmlaq_f32((_acc), (_x), (_y)) +#endif + +static inline float hsum_f32x4 (float32x4_t v) { + #if defined(__aarch64__) + return vaddvq_f32(v); + #else + float t[4]; + vst1q_f32(t, v); + return t[0] + t[1] + t[2] + t[3]; + #endif +} + float float32_distance_l2_impl_neon (const void *v1, const void *v2, int n, bool use_sqrt) { const float *a = (const float *)v1; const float *b = (const float *)v2; - - float32x4_t acc = vdupq_n_f32(0.0f); + + float32x4_t acc0 = vdupq_n_f32(0.0f), acc1 = acc0, acc2 = acc0, acc3 = acc0; int i = 0; + for (; i <= n - 16; i += 16) { + float32x4_t d0 = vsubq_f32(vld1q_f32(a + i ), vld1q_f32(b + i )); + float32x4_t d1 = vsubq_f32(vld1q_f32(a + i + 4), vld1q_f32(b + i + 4)); + float32x4_t d2 = vsubq_f32(vld1q_f32(a + i + 8), vld1q_f32(b + i + 8)); + float32x4_t d3 = vsubq_f32(vld1q_f32(a + i + 12), vld1q_f32(b + i + 12)); + acc0 = VFMA_F32(acc0, d0, d0); + acc1 = VFMA_F32(acc1, d1, d1); + acc2 = VFMA_F32(acc2, d2, d2); + acc3 = VFMA_F32(acc3, d3, d3); + } for (; i <= n - 4; i += 4) { - float32x4_t va = vld1q_f32(a + i); - float32x4_t vb = vld1q_f32(b + i); - float32x4_t d = vsubq_f32(va, vb); - acc = vmlaq_f32(acc, d, d); // acc += d * d + float32x4_t d = vsubq_f32(vld1q_f32(a + i), vld1q_f32(b + i)); + acc0 = VFMA_F32(acc0, d, d); } - float sum; - #if defined(__aarch64__) - sum = vaddvq_f32(acc); // fast horizontal add on arm64 - #else - float tmp[4]; vst1q_f32(tmp, acc); - sum = tmp[0] + tmp[1] + tmp[2] + tmp[3]; - #endif + float sum = hsum_f32x4(vaddq_f32(vaddq_f32(acc0, acc1), vaddq_f32(acc2, acc3))); for (; i < n; ++i) { float d = a[i] - b[i]; @@ -79,29 +100,36 @@ float float32_distance_l2_squared_neon (const void *v1, const void *v2, int n) { float float32_distance_cosine_neon (const void *v1, const void *v2, int n) { const float *a = (const float *)v1; const float *b = (const float *)v2; - - float32x4_t acc_dot = vdupq_n_f32(0.0f); - float32x4_t acc_a2 = vdupq_n_f32(0.0f); - float32x4_t acc_b2 = vdupq_n_f32(0.0f); + + // three quantities x four accumulators: twelve independent chains, which aarch64's + // 32 vector registers hold without spilling + float32x4_t dot0 = vdupq_n_f32(0.0f), dot1 = dot0, dot2 = dot0, dot3 = dot0; + float32x4_t na0 = dot0, na1 = dot0, na2 = dot0, na3 = dot0; + float32x4_t nb0 = dot0, nb1 = dot0, nb2 = dot0, nb3 = dot0; int i = 0; + for (; i <= n - 16; i += 16) { + float32x4_t a0 = vld1q_f32(a + i), a1 = vld1q_f32(a + i + 4); + float32x4_t a2 = vld1q_f32(a + i + 8), a3 = vld1q_f32(a + i + 12); + float32x4_t b0 = vld1q_f32(b + i), b1 = vld1q_f32(b + i + 4); + float32x4_t b2 = vld1q_f32(b + i + 8), b3 = vld1q_f32(b + i + 12); + dot0 = VFMA_F32(dot0, a0, b0); dot1 = VFMA_F32(dot1, a1, b1); + dot2 = VFMA_F32(dot2, a2, b2); dot3 = VFMA_F32(dot3, a3, b3); + na0 = VFMA_F32(na0, a0, a0); na1 = VFMA_F32(na1, a1, a1); + na2 = VFMA_F32(na2, a2, a2); na3 = VFMA_F32(na3, a3, a3); + nb0 = VFMA_F32(nb0, b0, b0); nb1 = VFMA_F32(nb1, b1, b1); + nb2 = VFMA_F32(nb2, b2, b2); nb3 = VFMA_F32(nb3, b3, b3); + } for (; i <= n - 4; i += 4) { - float32x4_t va = vld1q_f32(a + i); - float32x4_t vb = vld1q_f32(b + i); - - acc_dot = vmlaq_f32(acc_dot, va, vb); // dot += a * b - acc_a2 = vmlaq_f32(acc_a2, va, va); // norm_a += a * a - acc_b2 = vmlaq_f32(acc_b2, vb, vb); // norm_b += b * b + float32x4_t va = vld1q_f32(a + i), vb = vld1q_f32(b + i); + dot0 = VFMA_F32(dot0, va, vb); + na0 = VFMA_F32(na0, va, va); + nb0 = VFMA_F32(nb0, vb, vb); } - float d[4], a2[4], b2[4]; - vst1q_f32(d, acc_dot); - vst1q_f32(a2, acc_a2); - vst1q_f32(b2, acc_b2); - - float dot = d[0] + d[1] + d[2] + d[3]; - float norm_a = a2[0] + a2[1] + a2[2] + a2[3]; - float norm_b = b2[0] + b2[1] + b2[2] + b2[3]; + float dot = hsum_f32x4(vaddq_f32(vaddq_f32(dot0, dot1), vaddq_f32(dot2, dot3))); + float norm_a = hsum_f32x4(vaddq_f32(vaddq_f32(na0, na1), vaddq_f32(na2, na3))); + float norm_b = hsum_f32x4(vaddq_f32(vaddq_f32(nb0, nb1), vaddq_f32(nb2, nb3))); for (; i < n; ++i) { float ai = a[i]; @@ -121,19 +149,21 @@ float float32_distance_cosine_neon (const void *v1, const void *v2, int n) { float float32_distance_dot_neon (const void *v1, const void *v2, int n) { const float *a = (const float *)v1; const float *b = (const float *)v2; - - float32x4_t acc = vdupq_n_f32(0.0f); + + float32x4_t acc0 = vdupq_n_f32(0.0f), acc1 = acc0, acc2 = acc0, acc3 = acc0; int i = 0; + for (; i <= n - 16; i += 16) { + acc0 = VFMA_F32(acc0, vld1q_f32(a + i ), vld1q_f32(b + i )); + acc1 = VFMA_F32(acc1, vld1q_f32(a + i + 4), vld1q_f32(b + i + 4)); + acc2 = VFMA_F32(acc2, vld1q_f32(a + i + 8), vld1q_f32(b + i + 8)); + acc3 = VFMA_F32(acc3, vld1q_f32(a + i + 12), vld1q_f32(b + i + 12)); + } for (; i <= n - 4; i += 4) { - float32x4_t va = vld1q_f32(a + i); - float32x4_t vb = vld1q_f32(b + i); - acc = vmlaq_f32(acc, va, vb); // acc += a * b + acc0 = VFMA_F32(acc0, vld1q_f32(a + i), vld1q_f32(b + i)); } - float tmp[4]; - vst1q_f32(tmp, acc); - float dot = tmp[0] + tmp[1] + tmp[2] + tmp[3]; + float dot = hsum_f32x4(vaddq_f32(vaddq_f32(acc0, acc1), vaddq_f32(acc2, acc3))); for (; i < n; ++i) { dot += a[i] * b[i]; @@ -145,20 +175,21 @@ float float32_distance_dot_neon (const void *v1, const void *v2, int n) { float float32_distance_l1_neon (const void *v1, const void *v2, int n) { const float *a = (const float *)v1; const float *b = (const float *)v2; - - float32x4_t acc = vdupq_n_f32(0.0f); + + float32x4_t acc0 = vdupq_n_f32(0.0f), acc1 = acc0, acc2 = acc0, acc3 = acc0; int i = 0; + for (; i <= n - 16; i += 16) { + acc0 = vaddq_f32(acc0, vabdq_f32(vld1q_f32(a + i ), vld1q_f32(b + i ))); + acc1 = vaddq_f32(acc1, vabdq_f32(vld1q_f32(a + i + 4), vld1q_f32(b + i + 4))); + acc2 = vaddq_f32(acc2, vabdq_f32(vld1q_f32(a + i + 8), vld1q_f32(b + i + 8))); + acc3 = vaddq_f32(acc3, vabdq_f32(vld1q_f32(a + i + 12), vld1q_f32(b + i + 12))); + } for (; i <= n - 4; i += 4) { - float32x4_t va = vld1q_f32(a + i); - float32x4_t vb = vld1q_f32(b + i); - float32x4_t d = vabdq_f32(va, vb); // |a - b| - acc = vaddq_f32(acc, d); + acc0 = vaddq_f32(acc0, vabdq_f32(vld1q_f32(a + i), vld1q_f32(b + i))); } - float tmp[4]; - vst1q_f32(tmp, acc); - float sum = tmp[0] + tmp[1] + tmp[2] + tmp[3]; + float sum = hsum_f32x4(vaddq_f32(vaddq_f32(acc0, acc1), vaddq_f32(acc2, acc3))); for (; i < n; ++i) { sum += fabsf(a[i] - b[i]); @@ -1328,7 +1359,7 @@ float turbo_lut_dot_neon (const uint8_t *packed, float scale, const float *query // MARK: - -void init_distance_functions_neon (void) { +bool init_distance_functions_neon (void) { #if defined(__ARM_NEON) || defined(__ARM_NEON__) dispatch_distance_table[VECTOR_DISTANCE_L2][VECTOR_TYPE_F32] = float32_distance_l2_neon; dispatch_distance_table[VECTOR_DISTANCE_L2][VECTOR_TYPE_F16] = float16_distance_l2_neon; @@ -1365,5 +1396,8 @@ void init_distance_functions_neon (void) { distance_backend_name = "NEON"; turbo_lut_dot_function = turbo_lut_dot_neon; turbo_lut_backend_name = "NEON"; + return true; +#else + return false; #endif } diff --git a/src/distance-neon.h b/src/distance-neon.h index 2e190ff..6baa3c7 100644 --- a/src/distance-neon.h +++ b/src/distance-neon.h @@ -8,10 +8,12 @@ #ifndef __VECTOR_DISTANCE_NEON__ #define __VECTOR_DISTANCE_NEON__ +#include #include #include -void init_distance_functions_neon (void); +// returns true when the NEON kernels were compiled into this build +bool 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 e0d099a..95cb2fe 100644 --- a/src/distance-rvv.c +++ b/src/distance-rvv.c @@ -1021,7 +1021,7 @@ float turbo_lut_dot_rvv (const uint8_t *packed, float scale, const float *query_ // MARK: - -void init_distance_functions_rvv (void) { +bool init_distance_functions_rvv (void) { #if defined(__riscv_v_intrinsic) dispatch_distance_table[VECTOR_DISTANCE_L2][VECTOR_TYPE_F32] = float32_distance_l2_rvv; dispatch_distance_table[VECTOR_DISTANCE_L2][VECTOR_TYPE_F16] = float16_distance_l2_rvv; @@ -1058,5 +1058,8 @@ void init_distance_functions_rvv (void) { distance_backend_name = "RVV"; turbo_lut_dot_function = turbo_lut_dot_rvv; turbo_lut_backend_name = "RVV"; + return true; +#else + return false; #endif } diff --git a/src/distance-rvv.h b/src/distance-rvv.h index 8e2fbcf..5b09e1e 100644 --- a/src/distance-rvv.h +++ b/src/distance-rvv.h @@ -8,10 +8,12 @@ #ifndef __VECTOR_DISTANCE_RVV__ #define __VECTOR_DISTANCE_RVV__ +#include #include #include -void init_distance_functions_rvv (void); +// returns true when the RVV kernels were compiled into this build +bool 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 c739469..e5439c3 100644 --- a/src/distance-sse2.c +++ b/src/distance-sse2.c @@ -65,31 +65,41 @@ static inline __m128 bf16x4_to_f32x4_loadu(const uint16_t* p) { // MARK: - FLOAT32 - +// A single accumulator makes the loop one dependency chain, so it retires one vector per +// add latency however many ports the core has. Two independent accumulators double that +// while still fitting the eight XMM registers available on 32-bit x86. +static inline float hsum128_ps (__m128 v) { + __m128 s = _mm_add_ps(v, _mm_movehl_ps(v, v)); + s = _mm_add_ss(s, _mm_shuffle_ps(s, s, 0x55)); + return _mm_cvtss_f32(s); +} + static inline float float32_distance_l2_impl_sse2 (const void *v1, const void *v2, int n, bool use_sqrt) { const float *a = (const float *)v1; const float *b = (const float *)v2; - - __m128 acc = _mm_setzero_ps(); + + __m128 acc0 = _mm_setzero_ps(), acc1 = _mm_setzero_ps(); int i = 0; + for (; i <= n - 8; i += 8) { + __m128 d0 = _mm_sub_ps(_mm_loadu_ps(a + i ), _mm_loadu_ps(b + i )); + __m128 d1 = _mm_sub_ps(_mm_loadu_ps(a + i + 4), _mm_loadu_ps(b + i + 4)); + acc0 = _mm_add_ps(acc0, _mm_mul_ps(d0, d0)); + acc1 = _mm_add_ps(acc1, _mm_mul_ps(d1, d1)); + } for (; i <= n - 4; i += 4) { - __m128 va = _mm_loadu_ps(a + i); - __m128 vb = _mm_loadu_ps(b + i); - __m128 diff = _mm_sub_ps(va, vb); - __m128 sq = _mm_mul_ps(diff, diff); - acc = _mm_add_ps(acc, sq); + __m128 d = _mm_sub_ps(_mm_loadu_ps(a + i), _mm_loadu_ps(b + i)); + acc0 = _mm_add_ps(acc0, _mm_mul_ps(d, d)); } - float partial[4]; - _mm_storeu_ps(partial, acc); - float total = partial[0] + partial[1] + partial[2] + partial[3]; + float total = hsum128_ps(_mm_add_ps(acc0, acc1)); for (; i < n; ++i) { float d = a[i] - b[i]; total += d * d; } - return use_sqrt ? sqrtf((float)total) : (float)total; + return use_sqrt ? sqrtf(total) : total; } float float32_distance_l2_sse2 (const void *v1, const void *v2, int n) { @@ -103,21 +113,23 @@ float float32_distance_l2_squared_sse2 (const void *v1, const void *v2, int n) { float float32_distance_l1_sse2 (const void *v1, const void *v2, int n) { const float *a = (const float *)v1; const float *b = (const float *)v2; - - __m128 acc = _mm_setzero_ps(); + + const __m128 sign = _mm_set1_ps(-0.0f); + __m128 acc0 = _mm_setzero_ps(), acc1 = _mm_setzero_ps(); int i = 0; + for (; i <= n - 8; i += 8) { + __m128 d0 = _mm_sub_ps(_mm_loadu_ps(a + i ), _mm_loadu_ps(b + i )); + __m128 d1 = _mm_sub_ps(_mm_loadu_ps(a + i + 4), _mm_loadu_ps(b + i + 4)); + acc0 = _mm_add_ps(acc0, _mm_andnot_ps(sign, d0)); // abs using bitmask + acc1 = _mm_add_ps(acc1, _mm_andnot_ps(sign, d1)); + } for (; i <= n - 4; i += 4) { - __m128 va = _mm_loadu_ps(a + i); - __m128 vb = _mm_loadu_ps(b + i); - __m128 diff = _mm_sub_ps(va, vb); - __m128 abs_diff = _mm_andnot_ps(_mm_set1_ps(-0.0f), diff); // abs using bitmask - acc = _mm_add_ps(acc, abs_diff); + __m128 d = _mm_sub_ps(_mm_loadu_ps(a + i), _mm_loadu_ps(b + i)); + acc0 = _mm_add_ps(acc0, _mm_andnot_ps(sign, d)); } - float partial[4]; - _mm_storeu_ps(partial, acc); - float total = partial[0] + partial[1] + partial[2] + partial[3]; + float total = hsum128_ps(_mm_add_ps(acc0, acc1)); for (; i < n; ++i) { total += fabsf(a[i] - b[i]); @@ -129,20 +141,19 @@ float float32_distance_l1_sse2 (const void *v1, const void *v2, int n) { float float32_distance_dot_sse2 (const void *v1, const void *v2, int n) { const float *a = (const float *)v1; const float *b = (const float *)v2; - - __m128 acc = _mm_setzero_ps(); + + __m128 acc0 = _mm_setzero_ps(), acc1 = _mm_setzero_ps(); int i = 0; + for (; i <= n - 8; i += 8) { + acc0 = _mm_add_ps(acc0, _mm_mul_ps(_mm_loadu_ps(a + i ), _mm_loadu_ps(b + i ))); + acc1 = _mm_add_ps(acc1, _mm_mul_ps(_mm_loadu_ps(a + i + 4), _mm_loadu_ps(b + i + 4))); + } for (; i <= n - 4; i += 4) { - __m128 va = _mm_loadu_ps(a + i); - __m128 vb = _mm_loadu_ps(b + i); - __m128 prod = _mm_mul_ps(va, vb); - acc = _mm_add_ps(acc, prod); + acc0 = _mm_add_ps(acc0, _mm_mul_ps(_mm_loadu_ps(a + i), _mm_loadu_ps(b + i))); } - float partial[4]; - _mm_storeu_ps(partial, acc); - float total = partial[0] + partial[1] + partial[2] + partial[3]; + float total = hsum128_ps(_mm_add_ps(acc0, acc1)); for (; i < n; ++i) { total += a[i] * b[i]; @@ -154,7 +165,9 @@ float float32_distance_dot_sse2 (const void *v1, const void *v2, int n) { float float32_distance_cosine_sse2 (const void *v1, const void *v2, int n) { const float *a = (const float *)v1; const float *b = (const float *)v2; - + + // the three quantities are already independent chains, so one accumulator each keeps + // the register file within reach of 32-bit x86 __m128 acc_dot = _mm_setzero_ps(); __m128 acc_a2 = _mm_setzero_ps(); __m128 acc_b2 = _mm_setzero_ps(); @@ -169,27 +182,24 @@ float float32_distance_cosine_sse2 (const void *v1, const void *v2, int n) { acc_b2 = _mm_add_ps(acc_b2, _mm_mul_ps(vb, vb)); } - float dot[4], a2[4], b2[4]; - _mm_storeu_ps(dot, acc_dot); - _mm_storeu_ps(a2, acc_a2); - _mm_storeu_ps(b2, acc_b2); - - float total_dot = dot[0] + dot[1] + dot[2] + dot[3]; - float total_a2 = a2[0] + a2[1] + a2[2] + a2[3]; - float total_b2 = b2[0] + b2[1] + b2[2] + b2[3]; + float dot = hsum128_ps(acc_dot); + float norm_a = hsum128_ps(acc_a2); + float norm_b = hsum128_ps(acc_b2); for (; i < n; ++i) { - total_dot += a[i] * b[i]; - total_a2 += a[i] * a[i]; - total_b2 += b[i] * b[i]; + float ai = a[i]; + float bi = b[i]; + dot += ai * bi; + norm_a += ai * ai; + norm_b += bi * bi; } - float denom = sqrtf(total_a2 * total_b2); - if (denom == 0.0f) return 1.0f; - float cosine_sim = total_dot / denom; - if (cosine_sim > 1.0f) cosine_sim = 1.0f; - if (cosine_sim < -1.0f) cosine_sim = -1.0f; - return 1.0f - cosine_sim; + if (norm_a == 0.0f || norm_b == 0.0f) return 1.0f; + + float cosine_similarity = dot / (sqrtf(norm_a) * sqrtf(norm_b)); + if (cosine_similarity > 1.0f) cosine_similarity = 1.0f; + if (cosine_similarity < -1.0f) cosine_similarity = -1.0f; + return 1.0f - cosine_similarity; } // MARK: - FLOAT16 - @@ -1125,7 +1135,7 @@ float turbo_lut_dot_sse2 (const uint8_t *packed, float scale, const float *query // MARK: - -void init_distance_functions_sse2 (void) { +bool init_distance_functions_sse2 (void) { #if defined(__SSE2__) || (defined(_MSC_VER) && (defined(_M_X64) || (_M_IX86_FP >= 2))) dispatch_distance_table[VECTOR_DISTANCE_L2][VECTOR_TYPE_F32] = float32_distance_l2_sse2; dispatch_distance_table[VECTOR_DISTANCE_L2][VECTOR_TYPE_F16] = float16_distance_l2_sse2; @@ -1162,5 +1172,8 @@ void init_distance_functions_sse2 (void) { distance_backend_name = "SSE2"; turbo_lut_dot_function = turbo_lut_dot_sse2; turbo_lut_backend_name = "SSE2"; + return true; +#else + return false; #endif } diff --git a/src/distance-sse2.h b/src/distance-sse2.h index decb874..6bccb0f 100644 --- a/src/distance-sse2.h +++ b/src/distance-sse2.h @@ -8,10 +8,12 @@ #ifndef __VECTOR_DISTANCE_SSE2__ #define __VECTOR_DISTANCE_SSE2__ +#include #include #include -void init_distance_functions_sse2 (void); +// returns true when the SSE2 kernels were compiled into this build +bool 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 bf150d3..81b9347 100644 --- a/src/sqlite-vector.c +++ b/src/sqlite-vector.c @@ -122,6 +122,17 @@ SQLITE_EXTENSION_INIT1 typedef struct turbo_rotation_plan turbo_rotation_plan; +// Reference-counted snapshot of the in-memory quantized index. A scan holds a reference +// for as long as it walks the buffer - streaming cursors keep it across many xNext calls - +// so vector_quantize_preload() and vector_quantize_cleanup() can replace or drop the +// table's copy without freeing memory a running scan is still reading. +typedef struct { + int refcount; // guarded by qmutex + int counter; // number of quantized vectors in data + sqlite3_int64 bytes; // usable bytes in data + uint8_t data[]; +} preload_index; + typedef struct { vector_type v_type; // vector type int v_dim; // vector dimension @@ -143,9 +154,7 @@ typedef struct { float offset; // computed value by quantization bool binary_mean; // binary mean option for 1BIT quantization - void *preloaded; - int precounter; - sqlite3_int64 preloaded_bytes; + preload_index *preloaded; // owned reference, guarded by qmutex turbo_rotation_plan *turbo_plan; int turbo_plan_dim; @@ -194,6 +203,7 @@ typedef struct { float *turbo_query_lut; float *turbo_norm_lut; int turbo_lut_rows; + preload_index *preload_ref; // reference held while data points into it } stream; // NON-STREAMING VT INTERFACE @@ -215,6 +225,49 @@ extern const char *turbo_lut_backend_name; static sqlite3_mutex *qmutex; +// MARK: - Preloaded Index - + +static preload_index *preload_index_new (sqlite3_int64 bytes) { + if (bytes <= 0) return NULL; + preload_index *idx = (preload_index *)sqlite3_malloc64(sizeof(preload_index) + (sqlite3_uint64)bytes); + if (!idx) return NULL; + + idx->refcount = 1; + idx->counter = 0; + idx->bytes = bytes; + return idx; +} + +static void preload_index_release (preload_index *idx) { + if (!idx) return; + + sqlite3_mutex_enter(qmutex); + int refs = --idx->refcount; + sqlite3_mutex_leave(qmutex); + + if (refs == 0) sqlite3_free(idx); +} + +// borrow the table's index for the duration of a scan (NULL if nothing is preloaded) +static preload_index *preload_index_acquire (table_context *t_ctx) { + sqlite3_mutex_enter(qmutex); + preload_index *idx = t_ctx->preloaded; + if (idx) ++idx->refcount; + sqlite3_mutex_leave(qmutex); + + return idx; +} + +// hand a new index (or NULL) to the table and drop the reference to the previous one +static void preload_index_install (table_context *t_ctx, preload_index *idx) { + sqlite3_mutex_enter(qmutex); + preload_index *old = t_ctx->preloaded; + t_ctx->preloaded = idx; + sqlite3_mutex_leave(qmutex); + + preload_index_release(old); +} + // MARK: - SQLite Utils - bool sqlite_system_exists (sqlite3 *db, const char *name, const char *type) { @@ -353,21 +406,25 @@ static bool sqlite_table_is_without_rowid (sqlite3 *db, const char *table_name) static char *sqlite_get_int_prikey_column (sqlite3 *db, const char *table_name) { char sql[STATIC_SQL_SIZE]; - sqlite3_snprintf(sizeof(sql), sql, "SELECT COUNT(*), type, name FROM pragma_table_info('%q') WHERE pk > 0;", table_name); + // one row per PRIMARY KEY column. The statement takes no parameters: the old version + // both bound one anyway and read bare type/name columns alongside COUNT(*), which + // SQLite takes from an arbitrary row of the group. + sqlite3_snprintf(sizeof(sql), sql, "SELECT name, type FROM pragma_table_info('%q') WHERE pk > 0;", table_name); char *prikey = NULL; sqlite3_stmt *stmt = NULL; if (sqlite3_prepare_v2(db, sql, -1, &stmt, NULL) == SQLITE_OK) { - sqlite3_bind_text(stmt, 1, table_name, -1, SQLITE_STATIC); - if (sqlite3_step(stmt) == SQLITE_ROW) { - int count = sqlite3_column_int(stmt, 0); - if (count == 1) { - const char *decl_type = (const char *)sqlite3_column_text(stmt, 1); - // see https://www.sqlite.org/datatype3.html (Determination Of Column Affinity) - if (decl_type && strcasestr(decl_type, "INT")) { - prikey = sqlite_strdup((const char *)sqlite3_column_text(stmt, 2)); - } + const char *name = (const char *)sqlite3_column_text(stmt, 0); + const char *decl_type = (const char *)sqlite3_column_text(stmt, 1); + // see https://www.sqlite.org/datatype3.html (Determination Of Column Affinity) + if (name && decl_type && strcasestr(decl_type, "INT")) prikey = sqlite_strdup(name); + + // a composite primary key gives more than one row and cannot stand in for the + // rowid (step() invalidates name, so the copy has to be taken first) + if (prikey && sqlite3_step(stmt) != SQLITE_DONE) { + sqlite3_free(prikey); + prikey = NULL; } } } @@ -548,36 +605,23 @@ static inline int8_t q_round_s8 (float s) { return (int8_t)(int)r; } +// NOTE: these go through q_round_u8/q_round_s8 rather than casting first and clamping +// after. Converting a float to int is undefined when the value is NaN or outside the +// int range, and it does differ in practice: arm64 saturates, x86 yields INT_MIN. The +// helpers clamp in float and only then cast, which is also what every other quantizer +// below already did. static inline void quantize_float32_to_unsigned8bit (const float *v, uint8_t *q, float offset, float scale, int n) { int i = 0; for (; i + 3 < n; i += 4) { - float s0 = (v[i] - offset) * scale; - float s1 = (v[i + 1] - offset) * scale; - float s2 = (v[i + 2] - offset) * scale; - float s3 = (v[i + 3] - offset) * scale; - - int r0 = (int)(s0 + 0.5f * (1.0f - 2.0f * (s0 < 0.0f))); - int r1 = (int)(s1 + 0.5f * (1.0f - 2.0f * (s1 < 0.0f))); - int r2 = (int)(s2 + 0.5f * (1.0f - 2.0f * (s2 < 0.0f))); - int r3 = (int)(s3 + 0.5f * (1.0f - 2.0f * (s3 < 0.0f))); - - r0 = r0 > 255 ? 255 : (r0 < 0 ? 0 : r0); - r1 = r1 > 255 ? 255 : (r1 < 0 ? 0 : r1); - r2 = r2 > 255 ? 255 : (r2 < 0 ? 0 : r2); - r3 = r3 > 255 ? 255 : (r3 < 0 ? 0 : r3); - - q[i] = (uint8_t)r0; - q[i + 1] = (uint8_t)r1; - q[i + 2] = (uint8_t)r2; - q[i + 3] = (uint8_t)r3; + q[i] = q_round_u8((v[i] - offset) * scale); + q[i + 1] = q_round_u8((v[i + 1] - offset) * scale); + q[i + 2] = q_round_u8((v[i + 2] - offset) * scale); + q[i + 3] = q_round_u8((v[i + 3] - offset) * scale); } // Handle remaining elements for (; i < n; ++i) { - float scaled = (v[i] - offset) * scale; - int rounded = (int)(scaled + 0.5f * (1.0f - 2.0f * (scaled < 0.0f))); - rounded = rounded > 255 ? 255 : (rounded < 0 ? 0 : rounded); - q[i] = (uint8_t)rounded; + q[i] = q_round_u8((v[i] - offset) * scale); } } @@ -660,32 +704,14 @@ static inline void quantize_i8_to_unsigned8bit (const int8_t *v, uint8_t *q, flo static inline void quantize_float32_to_signed8bit (const float *v, int8_t *q, float offset, float scale, int n) { int i = 0; for (; i + 3 < n; i += 4) { - float s0 = (v[i] - offset) * scale; - float s1 = (v[i + 1] - offset) * scale; - float s2 = (v[i + 2] - offset) * scale; - float s3 = (v[i + 3] - offset) * scale; - - int r0 = (int)(s0 + 0.5f * (1.0f - 2.0f * (s0 < 0.0f))); - int r1 = (int)(s1 + 0.5f * (1.0f - 2.0f * (s1 < 0.0f))); - int r2 = (int)(s2 + 0.5f * (1.0f - 2.0f * (s2 < 0.0f))); - int r3 = (int)(s3 + 0.5f * (1.0f - 2.0f * (s3 < 0.0f))); - - r0 = r0 > 127 ? 127 : (r0 < -128 ? -128 : r0); - r1 = r1 > 127 ? 127 : (r1 < -128 ? -128 : r1); - r2 = r2 > 127 ? 127 : (r2 < -128 ? -128 : r2); - r3 = r3 > 127 ? 127 : (r3 < -128 ? -128 : r3); - - q[i] = (int8_t)r0; - q[i + 1] = (int8_t)r1; - q[i + 2] = (int8_t)r2; - q[i + 3] = (int8_t)r3; + q[i] = q_round_s8((v[i] - offset) * scale); + q[i + 1] = q_round_s8((v[i + 1] - offset) * scale); + q[i + 2] = q_round_s8((v[i + 2] - offset) * scale); + q[i + 3] = q_round_s8((v[i + 3] - offset) * scale); } for (; i < n; ++i) { - float scaled = (v[i] - offset) * scale; - int rounded = (int)(scaled + 0.5f * (1.0f - 2.0f * (scaled < 0.0f))); - rounded = rounded > 127 ? 127 : (rounded < -128 ? -128 : rounded); - q[i] = (int8_t)rounded; + q[i] = q_round_s8((v[i] - offset) * scale); } } @@ -1434,6 +1460,46 @@ const char *vector_distance_to_name (vector_distance type) { return "N/A"; } +// HAMMING is implemented for BIT vectors only, every other type implements +// everything but HAMMING. BIT columns are always scanned with HAMMING (the scan +// forces it), so any distance is accepted for them. +static bool vector_distance_is_supported (vector_distance distance, vector_type type) { + if (type == VECTOR_TYPE_BIT) return true; + return (distance != VECTOR_DISTANCE_HAMMING); +} + +// bounds-checked lookup: returns NULL instead of an out-of-range or unpopulated +// entry, so a gap in the dispatch table becomes an error and not a crash +static distance_function_t vector_lookup_distance_function (vector_distance distance, vector_type type) { + if ((int)distance <= 0 || (int)distance >= VECTOR_DISTANCE_MAX) return NULL; + if ((int)type <= 0 || (int)type >= VECTOR_TYPE_MAX) return NULL; + return dispatch_distance_table[distance][type]; +} + +// For unit-length vectors cosine distance is exactly 1 - dot, so the two norm +// accumulators - two thirds of the arithmetic in the cosine kernels - drop out of the +// inner loop. The DOT kernels return the negated dot product, hence the addition. +static float cosine_normalized_f32 (const void *v1, const void *v2, int n) { + distance_function_t dot_fn = dispatch_distance_table[VECTOR_DISTANCE_DOT][VECTOR_TYPE_F32]; + float d = 1.0f + dot_fn(v1, v2, n); + if (d < 0.0f) return 0.0f; + if (d > 2.0f) return 2.0f; + return d; +} + +// The stored vectors are unit length only because the caller said so with normalized=1, +// and only in full precision: the quantized index holds scaled integers whose norm is +// whatever the scale made it, so the shortcut is not valid there. The query is normalized +// once per scan in vCursorFilterCommon, which makes the result exact rather than merely +// rank-equivalent. +static bool vector_use_normalized_cosine (const table_context *t_ctx, bool quantized) { + if (quantized) return false; + if (!t_ctx->options.v_normalized) return false; + if (t_ctx->options.v_distance != VECTOR_DISTANCE_COSINE) return false; + if (t_ctx->options.v_type != VECTOR_TYPE_F32) return false; + return (dispatch_distance_table[VECTOR_DISTANCE_DOT][VECTOR_TYPE_F32] != NULL); +} + #if DEBUG_VECTOR_SERIALIZATION static void vector_print (void *buf, vector_type type, int n) { printf("type: %s - dim: %d [", vector_type_to_name(type), n); @@ -1643,31 +1709,32 @@ static inline size_t quantized_row_bytes (vector_qtype qtype, int dim, int bits) // MARK: - SQL - static char *generate_create_quant_table (const char *table_name, const char *column_name, char sql[STATIC_SQL_SIZE]) { - return sqlite3_snprintf(STATIC_SQL_SIZE, sql, "CREATE TABLE IF NOT EXISTS vector0_%q_%q (rowid1 INTEGER, rowid2 INTEGER, counter INTEGER, data BLOB);", table_name, column_name); + return sqlite3_snprintf(STATIC_SQL_SIZE, sql, "CREATE TABLE IF NOT EXISTS \"vector0_%w_%w\" (rowid1 INTEGER, rowid2 INTEGER, counter INTEGER, data BLOB);", table_name, column_name); } static char *generate_drop_quant_table (const char *table_name, const char *column_name, char sql[STATIC_SQL_SIZE]) { - return sqlite3_snprintf(STATIC_SQL_SIZE, sql, "DROP TABLE IF EXISTS vector0_%q_%q;", table_name, column_name); + return sqlite3_snprintf(STATIC_SQL_SIZE, sql, "DROP TABLE IF EXISTS \"vector0_%w_%w\";", table_name, column_name); } static char *generate_select_from_table (const char *table_name, const char *column_name, const char *pk_name, char sql[STATIC_SQL_SIZE]) { - return sqlite3_snprintf(STATIC_SQL_SIZE, sql, "SELECT %q, %q FROM %q ORDER BY %q;", pk_name, column_name, table_name, pk_name); + return sqlite3_snprintf(STATIC_SQL_SIZE, sql, "SELECT \"%w\", \"%w\" FROM \"%w\" ORDER BY \"%w\";", pk_name, column_name, table_name, pk_name); } static char *generate_select_quant_table (const char *table_name, const char *column_name, char sql[STATIC_SQL_SIZE]) { - return sqlite3_snprintf(STATIC_SQL_SIZE, sql, "SELECT counter, data FROM vector0_%q_%q;", table_name, column_name); + return sqlite3_snprintf(STATIC_SQL_SIZE, sql, "SELECT counter, data FROM \"vector0_%w_%w\";", table_name, column_name); } static char *generate_memory_quant_table (const char *table_name, const char *column_name, char sql[STATIC_SQL_SIZE]) { - return sqlite3_snprintf(STATIC_SQL_SIZE, sql, "SELECT SUM(LENGTH(data)) FROM vector0_%q_%q;", table_name, column_name); + return sqlite3_snprintf(STATIC_SQL_SIZE, sql, "SELECT SUM(LENGTH(data)) FROM \"vector0_%w_%w\";", table_name, column_name); } static char *generate_insert_quant_table (const char *table_name, const char *column_name, char sql[STATIC_SQL_SIZE]) { - return sqlite3_snprintf(STATIC_SQL_SIZE, sql, "INSERT INTO vector0_%q_%q (rowid1, rowid2, counter, data) VALUES (?, ?, ?, ?);", table_name, column_name); + return sqlite3_snprintf(STATIC_SQL_SIZE, sql, "INSERT INTO \"vector0_%w_%w\" (rowid1, rowid2, counter, data) VALUES (?, ?, ?, ?);", table_name, column_name); } static char *generate_quant_table_name (const char *table_name, const char *column_name, char sql[STATIC_SQL_SIZE]) { - return sqlite3_snprintf(STATIC_SQL_SIZE, sql, "vector0_%q_%q", table_name, column_name); + // NOTE: a plain name, not SQL - it is bound as a parameter, so no escaping here + return sqlite3_snprintf(STATIC_SQL_SIZE, sql, "vector0_%s_%s", table_name, column_name); } // MARK: - Vector Context and Options - @@ -1687,7 +1754,7 @@ void vector_context_free (void *p) { if (ctx->tables[i].t_name) sqlite3_free(ctx->tables[i].t_name); 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); + preload_index_release(ctx->tables[i].preloaded); table_context_free_turbo_cache(&ctx->tables[i]); } sqlite3_free(p); @@ -1828,6 +1895,19 @@ static int vector_rebuild_quantization (sqlite3_context *context, const char *ta } if (qtype != VECTOR_QUANT_TURBO) table_context_free_turbo_cache(t_ctx); + // A BIT column is already binary, so 8-bit quantization has nothing to scale: it wrote + // (dim+7)/8 bytes into a dim-byte slot and left the rest uninitialised. AUTO now means + // the identity 1BIT, and an explicit 8-bit request is refused. This used to fail with + // an unrelated message on a populated table and to silently record qtype=UINT8 on an + // empty one, which then applied to rows inserted later. + if (type == VECTOR_TYPE_BIT) { + if (qtype == VECTOR_QUANT_AUTO) qtype = VECTOR_QUANT_1BIT; + if (qtype != VECTOR_QUANT_1BIT) { + context_result_error(context, SQLITE_ERROR, "BIT vectors can only be quantized with qtype=1BIT"); + return SQLITE_MISUSE; + } + } + // 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) { @@ -1837,7 +1917,7 @@ static int vector_rebuild_quantization (sqlite3_context *context, const char *ta // max_memory == 0 means use all required memory if (max_memory == 0) { - sqlite3_snprintf(sizeof(sql), sql, "SELECT COUNT(*) FROM %q;", table_name); + sqlite3_snprintf(sizeof(sql), sql, "SELECT COUNT(*) FROM \"%w\";", table_name); int64_t count = sqlite_read_int64(db, sql); max_memory = (count == 0) ? DEFAULT_MAX_MEMORY : (uint64_t)count * (uint64_t)q_size; if (count <= 0) { @@ -2027,7 +2107,10 @@ static int vector_rebuild_quantization (sqlite3_context *context, const char *ta case VECTOR_TYPE_BF16: quantize_bfloat16((const uint16_t *)blob, data, offset, scale, dim, qtype); break; case VECTOR_TYPE_U8: quantize_u8((const uint8_t *)blob, data, offset, scale, dim, qtype); break; case VECTOR_TYPE_I8: quantize_i8((const int8_t *)blob, data, offset, scale, dim, qtype); break; - case VECTOR_TYPE_BIT: memcpy(data, blob, (dim + 7) / 8); break; // BIT to 8-bit: just copy + // unreachable for new indexes (see the guard above), but one written by an + // older build can still be loaded: zero the slot so the bytes past the + // packed bits are never fed to a distance kernel uninitialised + case VECTOR_TYPE_BIT: memset(data, 0, (size_t)dim); memcpy(data, blob, (dim + 7) / 8); break; } } @@ -2079,15 +2162,8 @@ static void vector_quantize_preload (sqlite3_context *context, int argc, sqlite3 return; } - // free previous preload (if any) - sqlite3_mutex_enter(qmutex); - if (t_ctx->preloaded) { - sqlite3_free(t_ctx->preloaded); - t_ctx->preloaded = NULL; - t_ctx->precounter = 0; - t_ctx->preloaded_bytes = 0; - } - sqlite3_mutex_leave(qmutex); + // drop the previous preload: scans already walking it keep their own reference + preload_index_install(t_ctx, NULL); 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); @@ -2102,17 +2178,17 @@ static void vector_quantize_preload (sqlite3_context *context, int argc, sqlite3 generate_memory_quant_table(table_name, column_name, sql); sqlite3 *db = sqlite3_context_db_handle(context); sqlite3_int64 required = sqlite_read_int64(db, sql); - if (required == 0) { + if (required <= 0) { context_result_error(context, SQLITE_ERROR, "Unable to read data from database. Ensure that vector_quantize() has been called before using vector_quantize_preload()"); return; } - int counter = 0; - void *buffer = (void *)sqlite3_malloc64(required); - if (!buffer) { + preload_index *idx = preload_index_new(required); + if (!idx) { context_result_error(context, SQLITE_NOMEM, "Out of memory: unable to allocate %lld bytes for quant buffer", (long long)required); return; } + uint8_t *buffer = idx->data; sqlite3_stmt *vm = NULL; generate_select_quant_table(table_name, column_name, sql); @@ -2120,11 +2196,15 @@ static void vector_quantize_preload (sqlite3_context *context, int argc, sqlite3 if (rc != SQLITE_OK) { context_result_error(context, rc, "Internal statement error: %s", sqlite3_errmsg(db)); sqlite3_finalize(vm); - sqlite3_free(buffer); + preload_index_release(idx); return; } - int seek = 0; + // the shadow table is an ordinary writable table and can change between the SUM that + // sized the buffer and the scan below, so bound every copy against what was allocated + const size_t row_stride = quantized_row_bytes(t_ctx->options.q_type, t_ctx->options.v_dim, t_ctx->options.q_bits); + sqlite3_int64 seek = 0; + sqlite3_int64 counter = 0; while (1) { rc = sqlite3_step(vm); if (rc == SQLITE_DONE) {rc = SQLITE_OK; break;} // return error: rebuild must be call (only if first time run) @@ -2132,26 +2212,38 @@ static void vector_quantize_preload (sqlite3_context *context, int argc, sqlite3 int n = sqlite3_column_int(vm, 0); int bytes = sqlite3_column_bytes(vm, 1); - uint8_t *data = (uint8_t *)sqlite3_column_blob(vm, 1); + const uint8_t *data = (const uint8_t *)sqlite3_column_blob(vm, 1); + + if (!data || n < 0 || bytes < 0 || (sqlite3_int64)bytes > required - seek) { + rc = SQLITE_CORRUPT; + break; + } - // no check here because I am sure quantization was performed only on non NULL data - memcpy((uint8_t *)buffer + seek, data, bytes); + memcpy(buffer + seek, data, (size_t)bytes); seek += bytes; counter += n; } sqlite3_finalize(vm); + // the loaded rows must really hold the number of vectors they claim: this is the same + // invariant the scans verify per chunk, checked once here so a malformed index is + // rejected at preload time instead of at query time + if ((rc == SQLITE_OK) && ((counter > INT_MAX) || ((sqlite3_uint64)seek < (sqlite3_uint64)counter * (sqlite3_uint64)row_stride))) { + rc = SQLITE_CORRUPT; + } + if (rc != SQLITE_OK) { - sqlite3_free(buffer); - context_result_error(context, rc, "vector_quantize_preload failed: %s", sqlite3_errmsg(db)); + preload_index_release(idx); + if (rc == SQLITE_CORRUPT) context_result_error(context, rc, "vector_quantize_preload failed: inconsistent quantization data for '%s.%s'", table_name, column_name); + else context_result_error(context, rc, "vector_quantize_preload failed: %s", sqlite3_errmsg(db)); return; } - sqlite3_mutex_enter(qmutex); - t_ctx->preloaded = buffer; - t_ctx->precounter = counter; - t_ctx->preloaded_bytes = required; - sqlite3_mutex_leave(qmutex); + idx->counter = (int)counter; + // seek, not required: rows removed between the two queries leave the tail of the + // buffer uninitialised, and the scans bound themselves with this length + idx->bytes = seek; + preload_index_install(t_ctx, idx); } static int vector_quantize (sqlite3_context *context, const char *table_name, const char *column_name, const char *arg_options, bool *was_preloaded) { @@ -2204,17 +2296,27 @@ static int vector_quantize (sqlite3_context *context, const char *table_name, co // success: returns the total number of quantized rows sqlite3_result_int64(context, (sqlite3_int64)counter); - if (was_preloaded) *was_preloaded = (t_ctx->preloaded != NULL); + if (was_preloaded) { + sqlite3_mutex_enter(qmutex); + *was_preloaded = (t_ctx->preloaded != NULL); + sqlite3_mutex_leave(qmutex); + } return SQLITE_OK; quantize_cleanup: { - const char *errmsg = sqlite3_errmsg(db); + // capture before the rollback, which resets the connection's error state + char *errmsg = (sqlite3_errcode(db) != SQLITE_OK) ? sqlite3_mprintf("%s", sqlite3_errmsg(db)) : NULL; if (savepoint_open) { sqlite3_exec(db, "ROLLBACK TO quantize;", NULL, NULL, NULL); sqlite3_exec(db, "RELEASE quantize;", NULL, NULL, NULL); } - sqlite3_result_error(context, errmsg, -1); + // only replace the message when SQLite actually has one: the callees set their own + // through context_result_error, and overwriting it reported "not an error" + if (errmsg) { + sqlite3_result_error(context, errmsg, -1); + sqlite3_free(errmsg); + } sqlite3_result_error_code(context, rc); return rc; } @@ -2271,15 +2373,9 @@ static void vector_quantize_cleanup (sqlite3_context *context, int argc, sqlite3 table_context *t_ctx = vector_context_lookup(v_ctx, table_name, column_name); if (!t_ctx) return; // if no table context exists then do nothing - // release any memory used in quantization - sqlite3_mutex_enter(qmutex); - if (t_ctx->preloaded) { - sqlite3_free(t_ctx->preloaded); - t_ctx->preloaded = NULL; - t_ctx->precounter = 0; - t_ctx->preloaded_bytes = 0; - } - sqlite3_mutex_leave(qmutex); + // release any memory used in quantization: scans still walking the index keep it + // alive through their own reference and free it when they are done + preload_index_install(t_ctx, NULL); // drop quant table (if any) char sql[STATIC_SQL_SIZE]; @@ -2541,6 +2637,10 @@ static int vCursorFilterCommon (sqlite3_vtab_cursor *cur, int idxNum, const char sqlite3_free(c->stream.turbo_norm_lut); c->stream.turbo_norm_lut = NULL; } + if (c->stream.preload_ref) { + preload_index_release(c->stream.preload_ref); + c->stream.preload_ref = NULL; + } memset(&c->stream, 0, sizeof(c->stream)); if (argc != 3 && argc != 4) { @@ -2592,6 +2692,13 @@ static int vCursorFilterCommon (sqlite3_vtab_cursor *cur, int idxNum, const char vector = (const void *)sqlite3_value_blob(argv[2]); vsize = sqlite3_value_bytes(argv[2]); if (!vector) return sqlite_vtab_set_error(&vtab->base, "%s: input vector cannot be NULL", fname); + + // the JSON branch above validates the dimension inside vector_from_json, the BLOB + // branch must do it here: the distance functions read v_dim elements unconditionally + size_t expected_bytes = vector_bytes_for_dim(t_ctx->options.v_type, t_ctx->options.v_dim); + if ((size_t)vsize != expected_bytes) { + return sqlite_vtab_set_error(&vtab->base, "%s: input vector must be %lld bytes (%d dimensions of type %s), but %d were provided", fname, (long long)expected_bytes, t_ctx->options.v_dim, vector_type_to_name(t_ctx->options.v_type), vsize); + } } VECTOR_PRINT((void*)vector, t_ctx->options.v_type, t_ctx->options.v_dim); @@ -2605,6 +2712,26 @@ static int vCursorFilterCommon (sqlite3_vtab_cursor *cur, int idxNum, const char } } + // with the 1 - dot shortcut the query has to be unit length too, and normalizing it + // once per scan costs one pass over a single vector + if (vector_use_normalized_cosine(t_ctx, quantized)) { + int dim = t_ctx->options.v_dim; + float *qn = (float *)sqlite_memdup(vector, vsize); + if (!qn) { + if (vector_allocated) sqlite3_free((void *)vector); + return SQLITE_NOMEM; + } + double norm_sq = 0.0; + for (int i = 0; i < dim; ++i) norm_sq += (double)qn[i] * (double)qn[i]; + if (norm_sq > 0.0) { + float inv = (float)(1.0 / sqrt(norm_sq)); + for (int i = 0; i < dim; ++i) qn[i] *= inv; + } + if (vector_allocated) sqlite3_free((void *)vector); + vector = qn; + vector_allocated = true; + } + c->table = t_ctx; if (is_streaming) { int rc = stream_callback(vtab->db, c, vector, vsize); @@ -2615,9 +2742,14 @@ static int vCursorFilterCommon (sqlite3_vtab_cursor *cur, int idxNum, const char // non-streaming flow int k = sqlite3_value_int(argv[3]); - if (k == 0) { + if (k <= 0) { + // an empty result, not an error: any non-OK return from xFilter is an error code + // to SQLite, and SQLITE_DONE only looked harmless because it happens to be the + // value sqlite3_step() reports at end of results if (vector_allocated) sqlite3_free((void *)vector); - return SQLITE_DONE; + c->row_index = 0; + c->row_count = 0; + return SQLITE_OK; } if (c->row_count != k) { @@ -2718,8 +2850,13 @@ static int vFullScanBestIndex (sqlite3_vtab *tab, sqlite3_index_info *pIdxInfo) // top-k mode: 4 positional args, argv[3] has the k integer pIdxInfo->estimatedCost = (double)1; pIdxInfo->estimatedRows = 100; - pIdxInfo->orderByConsumed = 1; pIdxInfo->idxNum = 1; + + // rows are emitted in ascending distance order, so that is the only ORDER BY we + // may claim: telling SQLite otherwise makes it drop a sorter we do not replace + pIdxInfo->orderByConsumed = (pIdxInfo->nOrderBy == 1 && + pIdxInfo->aOrderBy[0].iColumn == VECTOR_COLUMN_DISTANCE && + pIdxInfo->aOrderBy[0].desc == 0); } else { // streaming mode: 3 positional args, no sorting guaranteed pIdxInfo->estimatedCost = 1e8; @@ -2747,6 +2884,7 @@ static int vFullScanCursorClose (sqlite3_vtab_cursor *cur){ 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); + preload_index_release(c->stream.preload_ref); sqlite3_free(c); return SQLITE_OK; } @@ -2893,6 +3031,9 @@ static int vFullScanCursorNext (sqlite3_vtab_cursor *cur){ 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 *data = (const uint8_t *)c->stream.data; size_t i = (size_t)c->stream.dindex; @@ -2919,7 +3060,12 @@ static int vFullScanCursorNext (sqlite3_vtab_cursor *cur){ 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; // reset index for the new chunk + 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 *data = (const uint8_t *)c->stream.data; @@ -2939,6 +3085,7 @@ static int vFullScanCursorNext (sqlite3_vtab_cursor *cur){ // finished current chunk; force reload on next call c->stream.dcounter = 0; c->stream.data = NULL; // clear stale pointer to blob memory + c->stream.data_bytes = 0; } return SQLITE_OK; @@ -3021,7 +3168,7 @@ static int vFullScanRun (sqlite3 *db, vFullScanCursor *c, const void *v1, int v1 const char *table_name = c->table->t_name; int dimension = c->table->options.v_dim; - char *sql = sqlite3_mprintf("SELECT %q, %q FROM %q;", pk_name, col_name, table_name); + char *sql = sqlite3_mprintf("SELECT \"%w\", \"%w\" FROM \"%w\";", pk_name, col_name, table_name); if (!sql) return SQLITE_NOMEM; sqlite3_stmt *vm = NULL; @@ -3032,7 +3179,12 @@ static int vFullScanRun (sqlite3 *db, vFullScanCursor *c, const void *v1, int v1 vector_distance vd = c->table->options.v_distance; vector_type vt = c->table->options.v_type; if (vt == VECTOR_TYPE_BIT) vd = VECTOR_DISTANCE_HAMMING; // Force Hamming for BIT type - distance_function_t distance_fn = dispatch_distance_table[vd][vt]; + distance_function_t distance_fn = vector_lookup_distance_function(vd, vt); + if (!distance_fn) { + rc = sqlite_vtab_set_error(c->base.pVtab, "Distance '%s' is not supported for vector type '%s'", vector_distance_to_name(vd), vector_type_to_name(vt)); + goto cleanup; + } + if (vector_use_normalized_cosine(c->table, false)) distance_fn = cosine_normalized_f32; int dist_size = (vt == VECTOR_TYPE_BIT) ? ((dimension + 7) / 8) : dimension; size_t expected_bytes = vector_bytes_for_dim(vt, dimension); @@ -3070,13 +3222,21 @@ static int vFullScanCursorFilter (sqlite3_vtab_cursor *cur, int idxNum, const ch // MARK: - -static int vQuantRunMemory(vFullScanCursor *c, uint8_t *v, vector_qtype qtype, int dim) { - const int counter = c->table->precounter; - const uint8_t *data = c->table->preloaded; +static int vQuantRunMemory(vFullScanCursor *c, const preload_index *idx, uint8_t *v, vector_qtype qtype, int dim) { + const int counter = idx->counter; + const uint8_t *data = idx->data; const size_t rowid_size = sizeof(int64_t); const size_t vector_size = (qtype == VECTOR_QUANT_1BIT) ? ((dim + 7) / 8) : (dim * sizeof(uint8_t)); const size_t total_stride = rowid_size + vector_size; + // the preloaded index is built from an ordinary writable table, so never trust its + // row count against its length (the TurboQuant paths already do this) + const sqlite3_int64 data_bytes = idx->bytes; + if (!data || counter < 0 || data_bytes < 0 || (sqlite3_uint64)data_bytes < (sqlite3_uint64)counter * (sqlite3_uint64)total_stride) { + sqlite_vtab_set_error(c->base.pVtab, "Corrupted quantization data preloaded for '%s.%s'", c->table->t_name, c->table->c_name); + return SQLITE_CORRUPT; + } + double *distance = c->distance; int64_t *rowids = (int64_t *)c->rowids; int max_index = c->max_index; @@ -3089,7 +3249,8 @@ static int vQuantRunMemory(vFullScanCursor *c, uint8_t *v, vector_qtype qtype, i vt = VECTOR_TYPE_BIT; vd = VECTOR_DISTANCE_HAMMING; } - distance_function_t distance_fn = dispatch_distance_table[vd][vt]; + distance_function_t distance_fn = vector_lookup_distance_function(vd, vt); + if (!distance_fn) return sqlite_vtab_set_error(c->base.pVtab, "Distance '%s' is not supported for vector type '%s'", vector_distance_to_name(vd), vector_type_to_name(vt)); for (int i = 0; i < counter; ++i) { const uint8_t *current_data = data + (i * total_stride); @@ -3220,8 +3381,10 @@ static int vTurboRun (sqlite3 *db, vFullScanCursor *c, const void *v1, int v1siz 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); + preload_index *idx = preload_index_acquire(c->table); + if (idx) { + rc = vTurboRunPackedRows(c, idx->data, idx->bytes, idx->counter, qrot, qnorm_sq, c->table->turbo_centroids, query_lut, norm_lut, lut_rows, bits); + preload_index_release(idx); goto cleanup; } @@ -3286,13 +3449,15 @@ static int vQuantRun (sqlite3 *db, vFullScanCursor *c, const void *v1, int v1siz case VECTOR_TYPE_BF16: quantize_bfloat16((const uint16_t *)v1, v, offset, scale, dimension, qtype); break; case VECTOR_TYPE_U8: quantize_u8((const uint8_t *)v1, v, offset, scale, dimension, qtype); break; case VECTOR_TYPE_I8: quantize_i8((const int8_t *)v1, v, offset, scale, dimension, qtype); break; - case VECTOR_TYPE_BIT: memcpy(v, v1, (dimension + 7) / 8); break; // BIT to 8-bit: just copy + case VECTOR_TYPE_BIT: memset(v, 0, (size_t)dimension); memcpy(v, v1, (dimension + 7) / 8); break; // see vector_rebuild_quantization } } - if (c->table->preloaded) { - int rc = vQuantRunMemory(c, v, qtype, dimension); - if (v) sqlite3_free(v); + preload_index *idx = preload_index_acquire(c->table); + if (idx) { + int rc = vQuantRunMemory(c, idx, v, qtype, dimension); + preload_index_release(idx); + sqlite3_free(v); return rc; } #if DEBUG_VECTOR_SERIALIZATION @@ -3319,7 +3484,13 @@ static int vQuantRun (sqlite3 *db, vFullScanCursor *c, const void *v1, int v1siz vt = VECTOR_TYPE_BIT; vd = VECTOR_DISTANCE_HAMMING; } - distance_function_t distance_fn = dispatch_distance_table[vd][vt]; + distance_function_t distance_fn = vector_lookup_distance_function(vd, vt); + if (!distance_fn) { + sqlite_vtab_set_error(c->base.pVtab, "Distance '%s' is not supported for vector type '%s'", vector_distance_to_name(vd), vector_type_to_name(vt)); + sqlite3_finalize(vm); + sqlite3_free(v); + return SQLITE_ERROR; + } while (1) { rc = sqlite3_step(vm); @@ -3328,6 +3499,13 @@ static int vQuantRun (sqlite3 *db, vFullScanCursor *c, const void *v1, int v1siz int counter = sqlite3_column_int(vm, 0); uint8_t *data = (uint8_t *)sqlite3_column_blob(vm, 1); + int bytes = sqlite3_column_bytes(vm, 1); + if (!data || counter < 0 || bytes < 0 || (sqlite3_uint64)bytes < (sqlite3_uint64)counter * (sqlite3_uint64)total_stride) { + sqlite_vtab_set_error(c->base.pVtab, "Corrupted quantization data for '%s.%s'", c->table->t_name, c->table->c_name); + sqlite3_finalize(vm); + sqlite3_free(v); + return SQLITE_CORRUPT; + } // cache the maximum value to avoid repeated memory accesses double current_max_distance = c->distance[c->max_index]; @@ -3376,7 +3554,7 @@ static int vStreamScanCursorRun (sqlite3 *db, vFullScanCursor *c, const void *v1 c->stream.vsize = v1size; c->stream.vdim = dimension; - char *sql = sqlite3_mprintf("SELECT %q, %q FROM %q;", pk_name, col_name, table_name); + char *sql = sqlite3_mprintf("SELECT \"%w\", \"%w\" FROM \"%w\";", pk_name, col_name, table_name); if (!sql) { sqlite3_free(v); c->stream.vector = NULL; @@ -3391,7 +3569,12 @@ static int vStreamScanCursorRun (sqlite3 *db, vFullScanCursor *c, const void *v1 vector_distance vd = c->table->options.v_distance; vector_type vt = c->table->options.v_type; if (vt == VECTOR_TYPE_BIT) vd = VECTOR_DISTANCE_HAMMING; // Force Hamming for BIT type - distance_function_t distance_fn = dispatch_distance_table[vd][vt]; + distance_function_t distance_fn = vector_lookup_distance_function(vd, vt); + if (!distance_fn) { + rc = sqlite_vtab_set_error(c->base.pVtab, "Distance '%s' is not supported for vector type '%s'", vector_distance_to_name(vd), vector_type_to_name(vt)); + goto cleanup; + } + if (vector_use_normalized_cosine(c->table, false)) distance_fn = cosine_normalized_f32; c->stream.distance_fn = distance_fn; c->stream.vm = vm; @@ -3456,11 +3639,13 @@ static int vStreamTurboCursorRun (sqlite3 *db, vFullScanCursor *c, const void *v } } - if (c->table->preloaded) { + preload_index *idx = preload_index_acquire(c->table); + if (idx) { + c->stream.preload_ref = idx; c->stream.dindex = 0; - c->stream.data = c->table->preloaded; - c->stream.dcounter = c->table->precounter; - c->stream.data_bytes = c->table->preloaded_bytes; + c->stream.data = idx->data; + c->stream.dcounter = idx->counter; + c->stream.data_bytes = idx->bytes; return SQLITE_OK; } @@ -3521,7 +3706,7 @@ static int vStreamQuantCursorRun (sqlite3 *db, vFullScanCursor *c, const void *v case VECTOR_TYPE_BF16: quantize_bfloat16((const uint16_t *)v1, v, offset, scale, dimension, qtype); break; case VECTOR_TYPE_U8: quantize_u8((const uint8_t *)v1, v, offset, scale, dimension, qtype); break; case VECTOR_TYPE_I8: quantize_i8((const int8_t *)v1, v, offset, scale, dimension, qtype); break; - case VECTOR_TYPE_BIT: memcpy(v, v1, (dimension + 7) / 8); break; // BIT to 8-bit: just copy + case VECTOR_TYPE_BIT: memset(v, 0, (size_t)dimension); memcpy(v, v1, (dimension + 7) / 8); break; // see vector_rebuild_quantization } } @@ -3537,14 +3722,22 @@ static int vStreamQuantCursorRun (sqlite3 *db, vFullScanCursor *c, const void *v vt = VECTOR_TYPE_BIT; vd = VECTOR_DISTANCE_HAMMING; } - distance_function_t distance_fn = dispatch_distance_table[vd][vt]; + distance_function_t distance_fn = vector_lookup_distance_function(vd, vt); + if (!distance_fn) { + sqlite3_free(v); + c->stream.vector = NULL; + return sqlite_vtab_set_error(c->base.pVtab, "Distance '%s' is not supported for vector type '%s'", vector_distance_to_name(vd), vector_type_to_name(vt)); + } c->stream.distance_fn = distance_fn; // check if quant representation was preloaded - if (c->table->preloaded) { + preload_index *idx = preload_index_acquire(c->table); + if (idx) { + c->stream.preload_ref = idx; c->stream.dindex = 0; - c->stream.data = c->table->preloaded; - c->stream.dcounter = c->table->precounter; + c->stream.data = idx->data; + c->stream.dcounter = idx->counter; + c->stream.data_bytes = idx->bytes; return SQLITE_OK; } @@ -3651,6 +3844,11 @@ static void vector_init (sqlite3_context *context, int argc, sqlite3_value **arg return; } + if (vector_distance_is_supported(options.v_distance, options.v_type) == false) { + context_result_error(context, SQLITE_ERROR, "Distance '%s' is not supported for vector type '%s'", vector_distance_to_name(options.v_distance), vector_type_to_name(options.v_type)); + return; + } + // check if table is already loaded vector_context *v_ctx = (vector_context *)sqlite3_user_data(context); table_context *t_ctx = vector_context_lookup(v_ctx, table_name, column_name); @@ -3798,6 +3996,8 @@ SQLITE_VECTOR_API int sqlite3_vector_init (sqlite3 *db, char **pzErrMsg, const s return SQLITE_OK; cleanup: - vector_context_free(ctx); + // do NOT free ctx here: it is owned by the destructor registered with + // sqlite3_create_function_v2() above, which SQLite invokes both when that call itself + // fails and when the connection is closed. Freeing it again would be a double free. return rc; } diff --git a/test/backend.c b/test/backend.c new file mode 100644 index 0000000..ec99c92 --- /dev/null +++ b/test/backend.c @@ -0,0 +1,59 @@ +// +// backend.c +// sqlitevector +// +// Reports which distance kernels the extension actually installed, and optionally +// asserts that it is the expected one. A scan that silently falls back to a lower tier +// passes every test, so CI needs a way to tell "ran on AVX-512" from "meant to". +// +// ./backend print the installed backends +// ./backend AVX512 print them, and exit non-zero unless the distance backend is AVX512 +// + +#include +#include + +#include "sqlite3.h" + +extern int sqlite3_vector_init (sqlite3 *db, char **pzErrMsg, const sqlite3_api_routines *pApi); + +int main (int argc, char **argv) { + sqlite3 *db = NULL; + if (sqlite3_open(":memory:", &db) != SQLITE_OK) { + fprintf(stderr, "unable to open an in-memory database\n"); + return 2; + } + + int rc = sqlite3_vector_init(db, NULL, NULL); + if (rc != SQLITE_OK) { + fprintf(stderr, "sqlite3_vector_init failed (%d)\n", rc); + sqlite3_close(db); + return 2; + } + + sqlite3_stmt *stmt = NULL; + if (sqlite3_prepare_v2(db, "SELECT vector_backend(), vector_turboquant_backend();", -1, &stmt, NULL) != SQLITE_OK || + sqlite3_step(stmt) != SQLITE_ROW) { + fprintf(stderr, "unable to read the installed backends: %s\n", sqlite3_errmsg(db)); + sqlite3_finalize(stmt); + sqlite3_close(db); + return 2; + } + + const char *distance = (const char *)sqlite3_column_text(stmt, 0); + const char *turbo = (const char *)sqlite3_column_text(stmt, 1); + if (!distance) distance = "?"; + if (!turbo) turbo = "?"; + printf("distance backend: %s\n", distance); + printf("turboquant backend: %s\n", turbo); + + int result = 0; + if (argc > 1) { + result = (strcmp(distance, argv[1]) == 0) ? 0 : 1; + if (result) fprintf(stderr, "expected the %s backend, but %s was installed\n", argv[1], distance); + } + + sqlite3_finalize(stmt); + sqlite3_close(db); + return result; +}