From e0353ddfebc0ae5c51ea4a5ae8dc0208c01558cf Mon Sep 17 00:00:00 2001 From: Marco Bambini Date: Mon, 24 Aug 2026 18:22:56 +0200 Subject: [PATCH 1/5] fix: memory safety, SQL injection and ORDER BY in the extension core Eight defects found by a full read of the sources, each reproduced by compiling and running the tree before being fixed. Crashes: * An unsupported (distance, type) pair called a NULL function pointer. The dispatch table only implements HAMMING for BIT vectors, but distance=hamming was accepted for any type. vector_init now rejects the pair up front, and all five dispatch sites go through a bounds-checked lookup so a future gap in the table becomes an error rather than a segfault. * A query vector supplied as a BLOB was never length-checked, so the kernels read v_dim elements from whatever the caller passed. The JSON branch already validated the dimension; the BLOB branch now does too. Memory safety (all four reproduced under ASan): * The U8/I8/1BIT quantized scans decoded rows on trust. vector0_* is an ordinary writable table, so a truncated blob or an inflated counter walked off the end. All four paths now carry the same length guard the TurboQuant paths had, and the non-TurboQuant streaming cursor - which never recorded the buffer length at all - now tracks it. * vector_quantize_preload sized its buffer from SUM(LENGTH(data)) and filled it with no per-row bound and an int offset. No race is needed to trigger it: LENGTH() counts characters on a TEXT value while sqlite3_column_bytes() returns UTF-8 bytes, so multi-byte text in the data column makes the SUM under-report - ASan reports a 2400-byte write past a 1200-byte buffer. Copies are now bounded, offsets are 64-bit, and the loaded rows must hold the number of vectors they claim before the index is published. * The preloaded index was written and freed under qmutex but read without it, and locking the read would not have been enough: a streaming cursor holds the pointer across many xNext calls, so vector_quantize_cleanup() freed it mid-scan on a single thread. The index is now reference counted; scans hold a reference while they walk it and preload/cleanup swap the table's copy without touching what a live cursor is reading. * sqlite3_vector_init freed the context by hand on its error path even though sqlite3_create_function_v2() had already taken ownership - and SQLite invokes the destructor when that call itself fails, so the second free landed immediately. Injecting an allocation failure at each of 400 points inside init crashed at 23 of them; the sweep is now clean. Correctness: * Identifiers were interpolated with %q, which escapes string literals and does nothing for ';' or '"'. A table named x;DROP TABLE victim;-- injected statements through vector_quantize_cleanup, which runs its DROP outside any savepoint. All nine identifier positions now use %w inside explicit double quotes; generate_quant_table_name switched to %s because it builds a bare name that is bound as a parameter. Generated SQL for ordinary names is unchanged, and names like order-items or "my vec" now work instead of failing with a syntax error. * vFullScanBestIndex asserted orderByConsumed unconditionally, so SQLite dropped the sorter and every ORDER BY on the table-valued function was silently ignored - including ORDER BY distance DESC. It is now claimed only for a single ascending ORDER BY on the distance column. Also adds the cosine shortcut behind normalized=1: for unit-length FLOAT32 vectors cosine distance is 1 - dot, which drops two thirds of the arithmetic from the inner loop. The query is normalized once per scan so the result is exact rather than merely rank-equivalent, and quantized scans ignore the flag because the quantized index holds scaled integers whose norm is not 1. Co-Authored-By: Claude Opus 5 --- src/sqlite-vector.c | 342 +++++++++++++++++++++++++++++++++++--------- 1 file changed, 271 insertions(+), 71 deletions(-) diff --git a/src/sqlite-vector.c b/src/sqlite-vector.c index bf150d3..7acb3c6 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) { @@ -1434,6 +1487,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 +1736,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 +1781,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); @@ -1837,7 +1931,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) { @@ -2079,15 +2173,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 +2189,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 +2207,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 +2223,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,7 +2307,11 @@ 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: { @@ -2271,15 +2378,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 +2642,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 +2697,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 +2717,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); @@ -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; } @@ -3290,9 +3453,11 @@ static int vQuantRun (sqlite3 *db, vFullScanCursor *c, const void *v1, int v1siz } } - 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; } @@ -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; } From 7ea60662c13dac9c7941d108e897aa67d41d0e1b Mon Sep 17 00:00:00 2001 From: Marco Bambini Date: Mon, 24 Aug 2026 18:23:21 +0200 Subject: [PATCH 2/5] perf: build the SIMD kernels that were being compiled out, and unstarve them Every x86 build was shipping scalar code. The AVX2 and AVX-512 kernels sit behind #if defined(__AVX2__) / __AVX512F__, but CFLAGS never enabled those ISAs and CI just runs `make extension`, so both files compiled to nothing and init_distance_functions_avx2() was an empty stub. The dispatch chain made it worse: an if/else if ladder meant a CPU reporting AVX2 called the empty stub and never fell through to the SSE2 kernels that were compiled in. Forcing cpu_supports_avx2() true on a stock build reported vector_backend() = CPU. * The Makefile now probes the compiler and gives distance-avx2.c and distance-avx512.c their own -mavx2 -mfma / -mavx512{f,bw,vl,dq}, leaving every other translation unit at the baseline target. The probe makes this a no-op on non-x86 targets and on multi-arch builds. * Each init_distance_functions_* now returns whether its kernels were compiled in, and the ladder walks down instead of stopping at an empty stub. Turning the flags on showed why nobody had noticed: distance-avx512.c had never been compiled by anyone. It called _mm512_mask_set1_ps and _mm512_mask_set1_pd in ten places and neither intrinsic exists; the intent - keep where the mask is set, zero elsewhere - is _mm512_maskz_mov_ps/pd. It also needs AVX512DQ for _mm512_extractf32x8_ps, which cpu_supports_avx512() did not check for. On top of that, the f32 kernels kept a single vector accumulator, which makes the loop one dependency chain: an FMA has roughly four cycles of latency, so it retires one vector every four cycles however many FMA ports the core has. L2, dot and L1 now carry four independent accumulators over 16 lanes per iteration and use true FMA; cosine carries four per quantity on NEON and AVX-512, two on AVX2 where sixteen YMM registers would spill, and SSE2 stays at two so 32-bit x86 does not spill either. AVX2 and AVX-512 cosine were also implemented as three separate calls to the dot kernel - three passes over the data - and are now one fused pass. Measured on Apple M-series, dim 768, f32, L2-resident, best of 300 runs (Mvec/s): DOT 11.1 -> 25.6, SQUARED_L2 11.1 -> 25.6, L1 14.2 -> 28.4, COSINE 9.8 -> 17.1. End to end through the virtual table a full cosine scan goes from 4.35 to 7.02 Mvec/s - 1.6x, not 2.3x - because about a third of the time is SQLite's per-row step and past a few megabytes the scan is bound by memory bandwidth. Accuracy improved rather than degraded: shorter summation chains cut the worst error against a double-precision reference (dot at 1536 dims, 4.1e-6 -> 1.8e-6). Verification: the full 1447-test suite passes with the AVX2 backend active (Rosetta executes AVX2, so these kernels - never compiled before - are exercised for the first time), and every kernel is checked against a double-precision reference across 25 dimensions including tails on CPU, NEON, SSE2 and AVX2. AVX-512 is COMPILE-VERIFIED ONLY. No AVX-512 hardware or emulator was available: Docker's x86_64 emulation on Apple Silicon stops at AVX2 and SIGILLs on an EVEX-encoded vfmadd231ps, and QEMU TCG does not implement AVX-512 either. Disassembly confirms the intended shape - four distinct zmm accumulators in the dot loop, six in cosine, no stack spills, no calls back into the dot kernel - but `make unittest` should be run on real AVX-512 hardware before a release. Co-Authored-By: Claude Opus 5 --- Makefile | 14 ++++- src/distance-avx2.c | 139 ++++++++++++++++++++++++++++++------------ src/distance-avx2.h | 4 +- src/distance-avx512.c | 127 +++++++++++++++++++++++++++----------- src/distance-avx512.h | 4 +- src/distance-cpu.c | 22 ++++--- src/distance-neon.c | 132 ++++++++++++++++++++++++--------------- src/distance-neon.h | 4 +- src/distance-rvv.c | 5 +- src/distance-rvv.h | 4 +- src/distance-sse2.c | 111 ++++++++++++++++++--------------- src/distance-sse2.h | 4 +- 12 files changed, 380 insertions(+), 190 deletions(-) diff --git a/Makefile b/Makefile index a7132af..37eab8a 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();" 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 From 673732b460c0835febb9eb6a80ec34bf3c2a1b4f Mon Sep 17 00:00:00 2001 From: Marco Bambini Date: Mon, 24 Aug 2026 18:23:21 +0200 Subject: [PATCH 3/5] docs: document the normalized option normalized= was parsed and enforced for consistency across vector_init calls but had no effect and appeared in neither README.md nor API.md. It now drives the cosine shortcut, so the contract needs stating: it is an assertion, not a request - if the stored vectors are not unit length the distances will be wrong - and quantized scans ignore it. Also notes that HAMMING is only valid with type=1BIT, which is now enforced. Co-Authored-By: Claude Opus 5 --- API.md | 12 +++++++++++- README.md | 4 ++++ 2 files changed, 15 insertions(+), 1 deletion(-) 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/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'); From 394e64fa5892da3e68275dc4d61bcb41671aef21 Mon Sep 17 00:00:00 2001 From: Marco Bambini Date: Mon, 24 Aug 2026 21:41:34 +0200 Subject: [PATCH 4/5] fix: the four remaining low-severity defects * xFilter returned SQLITE_DONE for k=0. Any non-OK return from xFilter is an error code to SQLite; this one only looked harmless because 101 happens to be the value sqlite3_step() reports at end of results. It now sets an empty result and returns SQLITE_OK, and negative k takes the same path instead of relying on sqlite3_malloc() rejecting a negative size. * The 4-wide bodies of quantize_float32_to_{un,}signed8bit cast to int and clamped afterwards, which is undefined for NaN and for magnitudes past INT_MAX - the scalar helpers right above them already clamped in float first, as does every other quantizer in the file. They now call q_round_u8 / q_round_s8 too. Behaviour is unchanged on arm64 and x86-64 (both happen to land on 0 for NaN), but -fsanitize=float-cast-overflow reported "nan is outside the range of representable values of type 'int'" at sqlite-vector.c:612 before and is clean after. * sqlite_get_int_prikey_column bound a parameter to a statement that has none (silently returning SQLITE_RANGE) and read bare type/name columns alongside COUNT(*), which SQLite takes from an arbitrary row of the group. It now selects one row per primary-key column and requires exactly one, of INTEGER affinity, copying the name before the second step invalidates it. * A BIT column under 8-bit quantization copied (dim+7)/8 bytes into a dim-byte slot and left the rest uninitialised. On a populated table this failed anyway, but with the unrelated message "not an error"; on an empty one it silently recorded qtype=UINT8, which then applied to rows inserted later. qtype=AUTO on a BIT column now means the identity 1BIT and an explicit 8-bit request is refused with a message that says so. The three copy sites zero the slot first, since an index written by an older build can still hold that shape. The last one also fixes error reporting in vector_quantize: the cleanup label overwrote whatever the callee had put on the context with sqlite3_errmsg(db), which reported "not an error" whenever the failure came from a context error rather than from SQLite. It now keeps the callee's message unless SQLite actually has one, captured before the rollback resets it. Co-Authored-By: Claude Opus 5 --- src/sqlite-vector.c | 126 ++++++++++++++++++++++---------------------- 1 file changed, 63 insertions(+), 63 deletions(-) diff --git a/src/sqlite-vector.c b/src/sqlite-vector.c index 7acb3c6..81b9347 100644 --- a/src/sqlite-vector.c +++ b/src/sqlite-vector.c @@ -406,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; } } } @@ -601,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); } } @@ -713,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); } } @@ -1922,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) { @@ -2121,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; } } @@ -2315,13 +2304,19 @@ static int vector_quantize (sqlite3_context *context, const char *table_name, co 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; } @@ -2747,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) { @@ -3449,7 +3449,7 @@ 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 } } @@ -3706,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 } } From 61ca8f0c18eded260bc06bcfcf0eef733544cc6a Mon Sep 17 00:00:00 2001 From: Marco Bambini Date: Mon, 24 Aug 2026 21:55:48 +0200 Subject: [PATCH 5/5] ci: actually build and assert the SIMD kernels, including AVX-512 The unittest target builds every source in a single invocation, which leaves __AVX2__ and __AVX512F__ undefined: those kernels compile to nothing and the suite exercises the scalar fallback instead. Every green CI run so far has been testing the SIMD backends by not testing them - the same mechanism that made every shipped x86 build scalar. * unittest-simd compiles per translation unit the way the extension target does, so the SIMD kernels are actually linked in. RUNNER wraps the binaries for an emulator, EXPECT_BACKEND asserts which tier got installed. * test/backend.c prints the installed backends and, given an argument, exits non-zero unless that is the one that was installed. Without this a silent fallback is indistinguishable from a pass: the suite runs, 1447 tests go green, and the kernels under test were never reached. * A new avx512 job runs the suite on those kernels. 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) - so the job uses the hardware when it is there and Intel SDE when it is not, emulating Skylake-X because that is exactly the F/BW/VL/DQ set cpu_supports_avx512() requires. SDE is pulled from Intel's own download mirror with a pinned SHA-256; a mismatch fails the job rather than running an unverified binary. Verified on linux/amd64: unittest-simd installs the AVX2 backend and passes all 1447 tests (the plain unittest target reports CPU there), and EXPECT_BACKEND fails loudly on a mismatch. The SDE path could not be exercised locally - Pin aborts under nested emulation on Apple Silicon - but by construction that path cannot pass without AVX-512 having run. Co-Authored-By: Claude Opus 5 --- .github/workflows/main.yml | 63 ++++++++++++++++++++++++++++++++++++++ Makefile | 30 ++++++++++++++++++ test/backend.c | 59 +++++++++++++++++++++++++++++++++++ 3 files changed, 152 insertions(+) create mode 100644 test/backend.c 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/Makefile b/Makefile index 37eab8a..8c08ae7 100644 --- a/Makefile +++ b/Makefile @@ -151,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/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; +}