Source audit: 13 defects fixed, SIMD kernels actually built and tested

.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

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');
 ```
 
 ---

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

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');
 

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
 }

src/distance-avx2.h

@@ -8,10 +8,12 @@
 #ifndef __VECTOR_DISTANCE_AVX2__
 #define __VECTOR_DISTANCE_AVX2__
 
+#include <stdbool.h>
 #include <stdint.h>
 #include <stdio.h>
 
-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