GCC Code Coverage Report

Directory: src/
File: src/quant/rabitq.c
Date: 2026-10-23 09:17:32
Exec Total Coverage
Lines: 499 607 82.2%
Functions: 46 52 88.5%
Branches: 232 393 59.0%
Line Branch Exec Source
1 /*
2 * Copyright (c) 2026 Tiger Data, Inc.
3 * Licensed under the PostgreSQL License. See LICENSE for details.
4 *
5 * rabitq.c - RaBitQ (Randomized Binary Quantization) implementation
6 *
7 * Core implementation with SIMD dispatch for the binary inner product
8 * operation. The hot path during search is computing:
9 *
10 * sum = 0
11 * for each bit i:
12 * if bits[i] == 1:
13 * sum += transformed_query[i]
14 *
15 * This is optimized with AVX2/AVX-512/NEON intrinsics.
16 */
17
18 #include "vs_config.h"
19
20 #include <math.h>
21 #include <stdatomic.h>
22 #include <string.h>
23
24 #include "algo/simd_utils.h"
25 #include "algo/vecops.h"
26 #include "core/memory.h"
27 #include "core/platform.h"
28 #include "quant/matrix.h"
29 #include "quant/rabitq.h"
30 #include "types/vec16.h"
31 #include "types/vec32.h"
32
33 /*
34 * Compiler-Vectorized Implementation
35 *
36 * Uses target_clones to generate multiple versions for different ISAs.
37 * The dynamic linker selects the best version at load time.
38 *
39 * Structured for auto-vectorization: processes 8 floats per byte with
40 * explicit mask expansion that compilers can optimize.
41 */
42
43 VS_TARGET_CLONES static float
44 344 rabitq_inner_product_compiler(
45 const float *transformed, const uint8_t *bits, Dimension dim)
46 {
47 344 float sum = 0.0f;
48
49 /* Main loop: process 8 floats (1 byte) at a time.
50 * This structure helps auto-vectorization by:
51 * 1. Fixed iteration count inner loop (8 iterations)
52 * 2. Simple bit test with mask array lookup
53 * 3. Multiply instead of conditional for branchless code
54 */
55 344 Dimension i = 0;
56
2/2
✓ Branch 0 taken 10216 times.
✓ Branch 1 taken 344 times.
 10560 for (; i + 8 <= dim; i += 8)
57 {
58 10216 uint8_t byte = bits[i / 8];
59
60 /* Unrolled: multiply by bit value (0 or 1) */
61 10216 sum += transformed[i + 0] * ((byte >> 0) & 1);
62 10216 sum += transformed[i + 1] * ((byte >> 1) & 1);
63 10216 sum += transformed[i + 2] * ((byte >> 2) & 1);
64 10216 sum += transformed[i + 3] * ((byte >> 3) & 1);
65 10216 sum += transformed[i + 4] * ((byte >> 4) & 1);
66 10216 sum += transformed[i + 5] * ((byte >> 5) & 1);
67 10216 sum += transformed[i + 6] * ((byte >> 6) & 1);
68 10216 sum += transformed[i + 7] * ((byte >> 7) & 1);
69 }
70
71 /* Handle tail elements */
72
2/2
✓ Branch 0 taken 24 times.
✓ Branch 1 taken 344 times.
 368 for (; i < dim; i++)
73 {
74 24 int byte_idx = i / 8;
75 24 int bit_idx = i % 8;
76 24 int bit = (bits[byte_idx] >> bit_idx) & 1;
77 24 sum += transformed[i] * bit;
78 }
79
80 344 return sum;
81 }
82
83 /*
84 * Compiler-vectorized multi-candidate inner product baseline.
85 *
86 * Simple loop calling the single-candidate function. This is the
87 * baseline that hand-optimized vertical SIMD kernels must beat.
88 */
89 VS_TARGET_CLONES static void
90 16 rabitq_inner_product_multi_compiler(
91 const float *transformed,
92 const uint8_t *bits,
93 uint32_t stride,
94 Dimension dim,
95 uint32_t count,
96 float *results)
97 {
98
2/2
✓ Branch 0 taken 132 times.
✓ Branch 1 taken 16 times.
 148 for (uint32_t i = 0; i < count; i++)
99 132 results[i] = rabitq_inner_product_compiler(
100 132 transformed, bits + (size_t)i * stride, dim);
101 16 }
102
103 /*
104 * Function pointer dispatch for inner product and sign extraction
105 */
106 typedef float (*InnerProductFn)(const float *, const uint8_t *, Dimension);
107 typedef void (*ExtractSignsFn)(const float *, uint8_t *, Dimension);
108 static InnerProductFn g_inner_product_fn = NULL;
109 static InnerProductMultiFn g_inner_product_multi_fn = NULL;
110 static ExtractSignsFn g_extract_signs_fn = NULL;
111 static const char *g_impl_name = NULL;
112 static _Atomic(bool) g_rabitq_initialized = false;
113
114 /* Hamming distance function pointer dispatch */
115 typedef uint32_t (*HammingFn)(const uint8_t *, const uint8_t *, uint32_t);
116 typedef void (*HammingMultiFn)(
117 const uint8_t *,
118 const uint8_t *,
119 uint32_t,
120 uint32_t,
121 uint32_t,
122 uint32_t *);
123 static HammingFn g_hamming_fn = NULL;
124 static HammingMultiFn g_hamming_multi_fn = NULL;
125 static const char *g_hamming_impl_name = NULL;
126
127 /* Forward declaration for compiler-vectorized fallback */
128 VS_TARGET_CLONES static void rabitq_extract_signs_compiler(
129 const float *transformed, uint8_t *bits, Dimension dim);
130
131 /* Forward declarations for compiler-vectorized hamming */
132 VS_TARGET_CLONES static uint32_t rabitq_hamming_compiler(
133 const uint8_t *a, const uint8_t *b, uint32_t packed_bytes);
134 VS_TARGET_CLONES static void rabitq_hamming_multi_compiler(
135 const uint8_t *query_bits,
136 const uint8_t *data_bits,
137 uint32_t stride,
138 uint32_t packed_bytes,
139 uint32_t count,
140 uint32_t *results);
141
142 void
143 328 vs_rabitq_force_reinit(void)
144 {
145 328 g_rabitq_initialized = false;
146 328 g_inner_product_fn = NULL;
147 328 g_inner_product_multi_fn = NULL;
148 328 g_extract_signs_fn = NULL;
149 328 g_impl_name = NULL;
150 328 g_hamming_fn = NULL;
151 328 g_hamming_multi_fn = NULL;
152 328 g_hamming_impl_name = NULL;
153 328 }
154
155 int
156 593 vs_rabitq_init_simd(void)
157 {
158
2/2
✓ Branch 0 taken 257 times.
✓ Branch 1 taken 336 times.
 593 if (g_rabitq_initialized)
159 2 return 0;
160
161 #ifdef VS_SIMD_FULL
162 591 SimdCapability caps = vs_detect_simd();
163
164 #if defined(__x86_64__) || defined(_M_X64)
165
2/2
✓ Branch 0 taken 347 times.
✓ Branch 1 taken 244 times.
 591 if ((caps & VS_SIMD_AVX512_DQ) == VS_SIMD_AVX512_DQ)
166 {
167 347 g_inner_product_fn = vs_rabitq_inner_product_avx512;
168 347 g_inner_product_multi_fn = vs_rabitq_inner_product_multi_avx512;
169 347 g_extract_signs_fn = vs_rabitq_extract_signs_avx512;
170 347 g_impl_name = "avx512";
171 }
172
2/2
✓ Branch 0 taken 56 times.
✓ Branch 1 taken 188 times.
 244 else if (caps & SIMD_AVX2)
173 {
174 56 g_inner_product_fn = vs_rabitq_inner_product_avx2;
175 56 g_inner_product_multi_fn = vs_rabitq_inner_product_multi_avx2;
176 56 g_extract_signs_fn = vs_rabitq_extract_signs_avx2;
177 56 g_impl_name = "avx2";
178 }
179 else
180 {
181 188 g_inner_product_fn = rabitq_inner_product_compiler;
182 188 g_inner_product_multi_fn = rabitq_inner_product_multi_compiler;
183 188 g_extract_signs_fn = rabitq_extract_signs_compiler;
184 188 g_impl_name = "compiler";
185 }
186 #elif defined(__aarch64__) || defined(_M_ARM64)
187 if (caps & SIMD_NEON)
188 {
189 g_inner_product_fn = vs_rabitq_inner_product_neon;
190 g_inner_product_multi_fn = vs_rabitq_inner_product_multi_neon;
191 g_extract_signs_fn = vs_rabitq_extract_signs_neon;
192 g_impl_name = "neon";
193 }
194 else
195 {
196 g_inner_product_fn = rabitq_inner_product_compiler;
197 g_inner_product_multi_fn = rabitq_inner_product_multi_compiler;
198 g_extract_signs_fn = rabitq_extract_signs_compiler;
199 g_impl_name = "compiler";
200 }
201 #else
202 /* No hand-optimized kernels for this architecture. */
203 (void)caps;
204 g_inner_product_fn = rabitq_inner_product_compiler;
205 g_inner_product_multi_fn = rabitq_inner_product_multi_compiler;
206 g_extract_signs_fn = rabitq_extract_signs_compiler;
207 g_impl_name = "compiler";
208 #endif
209
210 /* Hamming dispatch (VPOPCNTDQ > AVX2 > compiler) */
211 #if defined(__x86_64__) || defined(_M_X64)
212
2/2
✓ Branch 0 taken 291 times.
✓ Branch 1 taken 300 times.
 591 if (caps & SIMD_AVX512_VPOPCNTDQ)
213 {
214 291 g_hamming_fn = vs_rabitq_hamming_avx512;
215 291 g_hamming_multi_fn = vs_rabitq_hamming_multi_avx512;
216 291 g_hamming_impl_name = "avx512-vpopcntdq";
217 }
218
2/2
✓ Branch 0 taken 56 times.
✓ Branch 1 taken 244 times.
 300 else if (caps & SIMD_AVX2)
219 {
220 56 g_hamming_fn = vs_rabitq_hamming_avx2;
221 56 g_hamming_multi_fn = vs_rabitq_hamming_multi_avx2;
222 56 g_hamming_impl_name = "avx2";
223 }
224 else
225 {
226 244 g_hamming_fn = rabitq_hamming_compiler;
227 244 g_hamming_multi_fn = rabitq_hamming_multi_compiler;
228 244 g_hamming_impl_name = "compiler";
229 }
230 #elif defined(__aarch64__) || defined(_M_ARM64)
231 g_hamming_fn = rabitq_hamming_compiler;
232 g_hamming_multi_fn = rabitq_hamming_multi_compiler;
233 g_hamming_impl_name = "compiler";
234 #else
235 g_hamming_fn = rabitq_hamming_compiler;
236 g_hamming_multi_fn = rabitq_hamming_multi_compiler;
237 g_hamming_impl_name = "compiler";
238 #endif
239
240 #else
241 /* simd=compiler or simd=none mode */
242 g_inner_product_fn = rabitq_inner_product_compiler;
243 g_inner_product_multi_fn = rabitq_inner_product_multi_compiler;
244 g_extract_signs_fn = rabitq_extract_signs_compiler;
245 g_impl_name = "compiler";
246 g_hamming_fn = rabitq_hamming_compiler;
247 g_hamming_multi_fn = rabitq_hamming_multi_compiler;
248 g_hamming_impl_name = "compiler";
249 #endif
250
251 591 g_rabitq_initialized = true;
252 591 return 0;
253 }
254
255 const char *
256 2 vs_rabitq_impl_name(void)
257 {
258
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 2 times.
 2 if (vs_unlikely(!g_rabitq_initialized))
259 ✗ vs_rabitq_init_simd();
260 2 return g_impl_name;
261 }
262
263 const char *
264 2 vs_rabitq_hamming_impl_name(void)
265 {
266
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 2 times.
 2 if (vs_unlikely(!g_rabitq_initialized))
267 ✗ vs_rabitq_init_simd();
268 2 return g_hamming_impl_name;
269 }
270
271 /*
272 * Internal helper: compute inner product using dispatched function
273 */
274 static inline float
275 7310 rabitq_inner_product(
276 const float *transformed, const uint8_t *bits, Dimension dim)
277 {
278
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 7310 times.
 7310 if (vs_unlikely(!g_rabitq_initialized))
279 ✗ vs_rabitq_init_simd();
280 7310 return g_inner_product_fn(transformed, bits, dim);
281 }
282
283 /*
284 * Compiler-Vectorized Sign Extraction
285 *
286 * Extracts sign bits from transformed floats. Structured for
287 * auto-vectorization by processing 8 floats at a time with explicit comparison
288 * and bit packing.
289 */
290 VS_TARGET_CLONES static void
291 218 rabitq_extract_signs_compiler(
292 const float *transformed, uint8_t *bits, Dimension dim)
293 {
294 218 Dimension packed_bytes = (dim + 7) / 8;
295 218 memset(bits, 0, packed_bytes);
296
297 /* Main loop: process 8 floats -> 1 byte at a time */
298 218 Dimension i = 0;
299
2/2
✓ Branch 0 taken 6262 times.
✓ Branch 1 taken 218 times.
 6480 for (; i + 8 <= dim; i += 8)
300 {
301 6262 uint8_t byte = 0;
302 /* Explicit bit tests help vectorization */
303 6262 byte |= (transformed[i + 0] > 0) ? (1 << 0) : 0;
304
2/2
✓ Branch 0 taken 3100 times.
✓ Branch 1 taken 3162 times.
 6262 byte |= (transformed[i + 1] > 0) ? (1 << 1) : 0;
305
2/2
✓ Branch 0 taken 3030 times.
✓ Branch 1 taken 3232 times.
 6262 byte |= (transformed[i + 2] > 0) ? (1 << 2) : 0;
306
2/2
✓ Branch 0 taken 2940 times.
✓ Branch 1 taken 3322 times.
 6262 byte |= (transformed[i + 3] > 0) ? (1 << 3) : 0;
307
2/2
✓ Branch 0 taken 2932 times.
✓ Branch 1 taken 3330 times.
 6262 byte |= (transformed[i + 4] > 0) ? (1 << 4) : 0;
308
2/2
✓ Branch 0 taken 2952 times.
✓ Branch 1 taken 3310 times.
 6262 byte |= (transformed[i + 5] > 0) ? (1 << 5) : 0;
309
2/2
✓ Branch 0 taken 2988 times.
✓ Branch 1 taken 3274 times.
 6262 byte |= (transformed[i + 6] > 0) ? (1 << 6) : 0;
310
2/2
✓ Branch 0 taken 2718 times.
✓ Branch 1 taken 3544 times.
 6262 byte |= (transformed[i + 7] > 0) ? (1 << 7) : 0;
311 6262 bits[i / 8] = byte;
312 }
313
314 /* Handle tail elements */
315
2/2
✓ Branch 0 taken 30 times.
✓ Branch 1 taken 218 times.
 248 for (; i < dim; i++)
316 {
317 30 int byte_idx = i / 8;
318 30 int bit_idx = i % 8;
319
2/2
✓ Branch 0 taken 8 times.
✓ Branch 1 taken 22 times.
 30 if (transformed[i] > 0)
320 8 bits[byte_idx] |= (1 << bit_idx);
321 }
322 218 }
323
324 /*
325 * Internal helper: extract sign bits using dispatched function
326 */
327 static inline void
328 694792 rabitq_extract_signs(const float *transformed, uint8_t *bits, Dimension dim)
329 {
330
2/2
✓ Branch 0 taken 8 times.
✓ Branch 1 taken 694784 times.
 694792 if (vs_unlikely(!g_rabitq_initialized))
331 8 vs_rabitq_init_simd();
332 694792 g_extract_signs_fn(transformed, bits, dim);
333 694792 }
334
335 /*
336 * Error bound helpers
337 *
338 * Used by all with_bound distance functions (scalar and batch, asymmetric
339 * and symmetric). Centralizes the f_error derivation and lower_bound
340 * computation that were previously duplicated across 4 functions.
341 */
342
343 static inline float
344 690534 rabitq_derive_f_error(
345 float f_add, float f_rescale, float c_error, Dimension dim)
346 {
347 690534 float f_rsq = f_rescale * f_rescale;
348 690534 float ratio = f_rsq / f_add;
349
3/4
✓ Branch 0 taken 290060 times.
✓ Branch 1 taken 400474 times.
✗ Branch 2 not taken.
✓ Branch 3 taken 255392 times.
 690534 if (ratio <= 1.0f || dim <= 1)
350 34668 return 2e-4f * sqrtf(f_add);
351 655866 return c_error * sqrtf(f_rsq - f_add);
352 }
353
354 float
355 683774 vs_rabitq_derive_f_error(float f_add, float f_rescale, Dimension dim)
356 {
357 1118916 float c_error = (dim > 1)
358 683774 ? 2.0f * VS_RABITQ_EPSILON / sqrtf((float)(dim - 1))
359
1/2
✓ Branch 0 taken 683774 times.
✗ Branch 1 not taken.
 683774 : 0.0f;
360 683774 return rabitq_derive_f_error(f_add, f_rescale, c_error, dim);
361 }
362
363 static inline Distance
364 6760 rabitq_lower_bound(
365 Distance est_dist, float f_error, float g_error, float multiplier)
366 {
367 6760 float err_margin = multiplier * f_error * g_error;
368 6760 float fp_margin = 1e-5f * fabsf(est_dist);
369 6760 return est_dist - err_margin - fp_margin;
370 }
371
372 /*
373 * Compiler-Vectorized Hamming Distance
374 *
375 * Computes XOR + popcount between two packed bit vectors.
376 * Uses target_clones to generate multiple versions for different ISAs.
377 */
378 VS_TARGET_CLONES static uint32_t
379 100 rabitq_hamming_compiler(
380 const uint8_t *a, const uint8_t *b, uint32_t packed_bytes)
381 {
382 100 uint32_t count = 0;
383
384 /* Main loop: process 8 bytes (64 bits) at a time */
385 100 uint32_t i = 0;
386
2/2
✓ Branch 0 taken 680 times.
✓ Branch 1 taken 100 times.
 780 for (; i + 8 <= packed_bytes; i += 8)
387 {
388 ✗ uint64_t va, vb;
389 680 memcpy(&va, a + i, 8);
390 680 memcpy(&vb, b + i, 8);
391 680 count += (uint32_t)__builtin_popcountll(va ^ vb);
392 }
393
394 /* Byte tail */
395
2/2
✓ Branch 0 taken 90 times.
✓ Branch 1 taken 100 times.
 190 for (; i < packed_bytes; i++)
396 90 count += (uint32_t)__builtin_popcount(a[i] ^ b[i]);
397
398 100 return count;
399 }
400
401 /*
402 * Compiler-Vectorized Multi-Candidate Hamming Distance
403 */
404 VS_TARGET_CLONES static void
405 ✗ rabitq_hamming_multi_compiler(
406 const uint8_t *query_bits,
407 const uint8_t *data_bits,
408 uint32_t stride,
409 uint32_t packed_bytes,
410 uint32_t count,
411 uint32_t *results)
412 {
413 ✗ for (uint32_t c = 0; c < count; c++)
414 {
415 ✗ results[c] = rabitq_hamming_compiler(
416 ✗ query_bits, data_bits + c * stride, packed_bytes);
417 }
418 ✗ }
419
420 /*
421 * Lifecycle functions
422 */
423
424 /*
425 * Apply the index rotation P^T*in -> out. Uses the O(d log d) Randomized
426 * Hadamard Transform when armed (supported dims), else the dense matrix.
427 */
428 static inline void
429 773925 rabitq_rotate(const RaBitQParams *p, const float *in, float *out)
430 {
431
2/2
✓ Branch 0 taken 711310 times.
✓ Branch 1 taken 62615 times.
 773925 if (p->use_fast_rotate)
432 711310 vs_fast_rotate_apply(&p->fr, in, out);
433 else
434 62615 vs_matrix_transpose_vector_mul(p->P, in, out, p->dim);
435 773925 }
436
437 RaBitQParams *
438 839 vs_rabitq_create(Dimension dim, uint64_t seed)
439 {
440 839 size_t size = VS_RABITQ_PARAMS_SIZE(dim);
441 839 RaBitQParams *params = vs_alloc(size);
442
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 839 times.
 839 if (params == NULL)
443 ✗ return NULL;
444
445
1/2
✗ Branch 1 not taken.
✓ Branch 2 taken 839 times.
 839 if (vs_rabitq_init(params, dim, seed) != 0)
446 {
447 ✗ vs_free(params);
448 ✗ return NULL;
449 }
450
451 560 return params;
452 }
453
454 RaBitQParams *
455 12 vs_rabitq_create_from_matrix(Dimension dim, uint64_t seed, const float *P)
456 {
457 /* Stores an explicit dense matrix, so always allocate the full layout. */
458 12 size_t size = VS_RABITQ_PARAMS_DENSE_SIZE(dim);
459 12 RaBitQParams *params = vs_alloc(size);
460
2/2
✓ Branch 0 taken 4 times.
✓ Branch 1 taken 8 times.
 12 if (params == NULL)
461 ✗ return NULL;
462
463 12 params->dim = dim;
464 12 params->seed = seed;
465 12 params->packed_bytes = VS_RABITQ_BYTES(dim);
466 12 params->use_fast_rotate = false;
467 12 memcpy(params->P, P, (size_t)dim * dim * sizeof(float));
468
469 12 return params;
470 }
471
472 int
473 938 vs_rabitq_init(RaBitQParams *params, Dimension dim, uint64_t seed)
474 {
475
4/4
✓ Branch 0 taken 562 times.
✓ Branch 1 taken 376 times.
✓ Branch 2 taken 2 times.
✓ Branch 3 taken 560 times.
 938 if (params == NULL || dim == 0)
476 4 return -1;
477
478 934 params->dim = dim;
479 934 params->seed = seed;
480 934 params->packed_bytes = VS_RABITQ_BYTES(dim);
481
482 /*
483 * Prefer the O(d log d) Randomized Hadamard rotation where the dim
484 * supports it: the dense P is neither built (no O(d^3) QR) nor stored
485 * (no O(d^2) per-backend matrix). Same seed, so build-encode and
486 * query-rotate stay consistent. Unsupported dims fall back to the
487 * dense random orthogonal matrix in P[].
488 */
489 934 params->use_fast_rotate = vs_fast_rotate_supported(dim);
490
2/2
✓ Branch 0 taken 708 times.
✓ Branch 1 taken 226 times.
 934 if (params->use_fast_rotate)
491 708 vs_fast_rotate_init(&params->fr, dim, seed);
492
1/2
✗ Branch 1 not taken.
✓ Branch 2 taken 226 times.
 226 else if (vs_random_orthogonal_matrix(params->P, dim, seed) != 0)
493 ✗ return -1;
494
495 560 return 0;
496 }
497
498 void
499 4 vs_rabitq_cleanup(RaBitQParams *params)
500 {
501 ✗ (void)params;
502 /* P is now inline — nothing to free */
503 4 }
504
505 void
506 272 vs_rabitq_destroy(RaBitQParams *params)
507 {
508
2/2
✓ Branch 0 taken 2 times.
✓ Branch 1 taken 270 times.
 272 if (params == NULL)
509 2 return;
510
511 270 vs_free(params);
512 }
513
514 /*
515 * Encoding functions
516 */
517
518 RaBitQData *
519 658 vs_rabitq_encode(const RaBitQParams *params, Vec32Ref input, Vec32Ref centroid)
520 {
521
6/6
✓ Branch 0 taken 656 times.
✓ Branch 1 taken 2 times.
✓ Branch 2 taken 654 times.
✓ Branch 3 taken 2 times.
✓ Branch 4 taken 2 times.
✓ Branch 5 taken 652 times.
 658 if (params == NULL || input.data == NULL || centroid.data == NULL)
522 6 return NULL;
523
524
3/4
✓ Branch 0 taken 652 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 2 times.
✓ Branch 3 taken 650 times.
 652 if (input.dim != params->dim || centroid.dim != params->dim)
525 2 return NULL;
526
527 650 size_t size = VS_RABITQ_DATA_SIZE(params->dim);
528 650 RaBitQData *output = vs_alloc(size);
529
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 650 times.
 650 if (output == NULL)
530 ✗ return NULL;
531
532
1/2
✗ Branch 1 not taken.
✓ Branch 2 taken 650 times.
 650 if (vs_rabitq_encode_into(params, input, centroid, output) != 0)
533 {
534 ✗ vs_free(output);
535 ✗ return NULL;
536 }
537
538 650 return output;
539 }
540
541 void
542 29606 vs_rabitq_scratch_init(RaBitQScratch *scratch, Dimension dim)
543 {
544 29606 scratch->residual = vs_alloc_aligned(dim * sizeof(float), 64);
545 29606 scratch->transformed = vs_alloc_aligned(dim * sizeof(float), 64);
546 29606 scratch->xu_cb = vs_alloc_aligned(dim * sizeof(float), 64);
547 29606 }
548
549 void
550 12186 vs_rabitq_scratch_cleanup(RaBitQScratch *scratch)
551 {
552 12186 vs_free_aligned(scratch->residual);
553 12186 vs_free_aligned(scratch->transformed);
554 12186 vs_free_aligned(scratch->xu_cb);
555 12186 scratch->residual = NULL;
556 12186 scratch->transformed = NULL;
557 12186 scratch->xu_cb = NULL;
558 12186 }
559
560 int
561 74116 vs_rabitq_encode_into_ex(
562 const RaBitQParams *params,
563 Vec32Ref input,
564 Vec32Ref centroid,
565 RaBitQData *output,
566 RaBitQScratch *scratch)
567 {
568
4/8
✓ Branch 0 taken 74116 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 74116 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 74116 times.
✗ Branch 5 not taken.
✓ Branch 6 taken 59416 times.
✗ Branch 7 not taken.
 74116 if (params == NULL || input.data == NULL || centroid.data == NULL ||
569
2/2
✓ Branch 0 taken 14700 times.
✓ Branch 1 taken 59416 times.
 74116 output == NULL || scratch == NULL)
570 ✗ return -1;
571
572
3/4
✓ Branch 0 taken 74116 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 14700 times.
✓ Branch 3 taken 59416 times.
 74116 if (input.dim != params->dim || centroid.dim != params->dim)
573 ✗ return -1;
574
575 74116 Dimension dim = params->dim;
576
577 74116 float *residual = scratch->residual;
578 74116 float *transformed = scratch->transformed;
579
580 /* residual = input - centroid; transformed = P^T * residual. The signs +
581 * factor math is shared with encode_from_pt (which the insert path calls
582 * directly with a pre-rotated residual). */
583 74116 vec32_sub(input.data, centroid.data, residual, dim);
584 74116 rabitq_rotate(params, residual, transformed);
585
586 74116 return vs_rabitq_encode_from_pt(params, transformed, output, scratch);
587 }
588
589 int
590 662491 vs_rabitq_encode_from_pt(
591 const RaBitQParams *params,
592 const float *pt_residual,
593 RaBitQData *output,
594 RaBitQScratch *scratch)
595 {
596
4/8
✓ Branch 0 taken 662491 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 236484 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 236484 times.
✗ Branch 5 not taken.
✗ Branch 6 not taken.
✓ Branch 7 taken 236484 times.
 662491 if (params == NULL || pt_residual == NULL || output == NULL ||
597
1/2
✓ Branch 0 taken 426007 times.
✗ Branch 1 not taken.
 426007 scratch == NULL)
598 ✗ return -1;
599
600 662491 Dimension dim = params->dim;
601 662491 float *xu_cb = scratch->xu_cb;
602
603 /* pt_residual is the rotated residual P^T*(input-centroid); everything
604 * below operates on it exactly as encode_into_ex did on `transformed`. */
605 662491 rabitq_extract_signs(pt_residual, output->bits, dim);
606
607 662491 float cb = -0.5f;
608
3/3
✓ Branch 0 taken 7203648 times.
✓ Branch 1 taken 31104811 times.
✓ Branch 2 taken 426007 times.
 38734466 for (Dimension i = 0; i < dim; i++)
609 {
610 38071975 int byte_idx = i / 8;
611 38071975 int bit_idx = i % 8;
612 38071975 int bit = (output->bits[byte_idx] >> bit_idx) & 1;
613 38071975 xu_cb[i] = (float)bit + cb;
614 }
615
616 662491 float l2_sqr = vs_l2_norm_squared(pt_residual, dim);
617 662491 float ip_resi_xucb = vs_dot_product(pt_residual, xu_cb, dim);
618
619
2/2
✓ Branch 0 taken 1639 times.
✓ Branch 1 taken 660852 times.
 662491 if (fabsf(ip_resi_xucb) < 1e-10f)
620 1639 ip_resi_xucb = 1e-10f;
621
622 662491 float sqrt_d = sqrtf((float)dim);
623 662491 float l1_norm = 2.0f * fabsf(ip_resi_xucb);
624 662491 output->f_add = l2_sqr;
625 662491 output->f_rescale = l2_sqr * sqrt_d / l1_norm;
626
627 662491 return 0;
628 }
629
630 int
631 25085 vs_rabitq_encode_into(
632 const RaBitQParams *params,
633 Vec32Ref input,
634 Vec32Ref centroid,
635 RaBitQData *output)
636 {
637
5/8
✓ Branch 0 taken 25085 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 25085 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 25085 times.
✗ Branch 5 not taken.
✓ Branch 6 taken 2 times.
✓ Branch 7 taken 25083 times.
 25085 if (params == NULL || input.data == NULL || centroid.data == NULL ||
638 output == NULL)
639 2 return -1;
640
641
4/4
✓ Branch 0 taken 14834 times.
✓ Branch 1 taken 10249 times.
✓ Branch 2 taken 2 times.
✓ Branch 3 taken 25079 times.
 25083 if (input.dim != params->dim || centroid.dim != params->dim)
642 4 return -1;
643
644 25079 Dimension dim = params->dim;
645
646 /* Allocate temporary buffers */
647 25079 float *residual = vs_alloc_aligned(dim * sizeof(float), 64);
648 25079 float *transformed = vs_alloc_aligned(dim * sizeof(float), 64);
649 25079 float *xu_cb = vs_alloc_aligned(dim * sizeof(float), 64);
650
651
4/6
✓ Branch 0 taken 25079 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 14832 times.
✓ Branch 3 taken 10247 times.
✗ Branch 4 not taken.
✓ Branch 5 taken 14832 times.
 25079 if (residual == NULL || transformed == NULL || xu_cb == NULL)
652 {
653 ✗ if (residual)
654 ✗ vs_free_aligned(residual);
655 ✗ if (transformed)
656 ✗ vs_free_aligned(transformed);
657 ✗ if (xu_cb)
658 ✗ vs_free_aligned(xu_cb);
659 ✗ return -1;
660 }
661
662 /* Step 1: Compute residual = input - centroid */
663 25079 vec32_sub(input.data, centroid.data, residual, dim);
664
665 /* Step 2: Transform residual through P^T
666 * RaBitQ applies the same rotation to all vectors (data, centroid, query).
667 * The library stores Q^T and multiplies by it, which is effectively P^T.
668 * Since P^T * (data - centroid) = P^T * data - P^T * centroid, applying
669 * P^T to the residual is equivalent to rotating both vectors then
670 * subtracting.
671 */
672 25079 rabitq_rotate(params, residual, transformed);
673
674 /* Step 3: Extract sign bits (LSB-first packing, FAISS-compatible) */
675 25079 rabitq_extract_signs(transformed, output->bits, dim);
676
677 /* Step 4: Compute xu_cb = binary_code - 0.5 (for 1-bit quantization) */
678 /* In 1-bit RaBitQ, cb = -0.5, so xu_cb[i] = bit[i] - 0.5 */
679 25079 float cb = -0.5f;
680
3/3
✓ Branch 0 taken 476772 times.
✓ Branch 1 taken 780025 times.
✓ Branch 2 taken 10247 times.
 1267044 for (Dimension i = 0; i < dim; i++)
681 {
682 1241965 int byte_idx = i / 8;
683 1241965 int bit_idx = i % 8; /* LSB-first */
684 1241965 int bit = (output->bits[byte_idx] >> bit_idx) & 1;
685 1241965 xu_cb[i] = (float)bit + cb;
686 }
687
688 /* Step 5: Compute factors for distance estimation */
689 25079 float l2_sqr = vs_l2_norm_squared(transformed, dim);
690
691 25079 float ip_resi_xucb = vs_dot_product(transformed, xu_cb, dim);
692
693 /* Handle corner case: avoid division by near-zero in f_rescale */
694
2/2
✓ Branch 0 taken 123 times.
✓ Branch 1 taken 24956 times.
 25079 if (fabsf(ip_resi_xucb) < 1e-10f)
695 123 ip_resi_xucb = 1e-10f;
696
697 /* L2 distance factors (FAISS-style formula for better accuracy)
698 *
699 * FAISS uses a correction factor based on the actual distribution of
700 * values: est_dist = ||v-c||² + ||q-c||² - 2 * dp_multiplier * final_dot
701 * where:
702 * dp_multiplier = ||v-c||² * sqrt(d) / ||v-c||_1
703 *
704 * The L1 norm is encoded in ip_resi_xucb (which equals 0.5 * ||v-c||_1
705 * because xu_cb = bit - 0.5, so the dot product accumulates ±0.5 * |val|).
706 *
707 * f_add = ||v-c||² (the L2 squared distance to centroid)
708 * f_rescale = dp_multiplier = ||v-c||² * sqrt(d) / ||v-c||_1
709 *
710 * f_error is derived at query time from f_add and f_rescale.
711 */
712 25079 float sqrt_d = sqrtf((float)dim);
713 25079 float l1_norm = 2.0f * fabsf(ip_resi_xucb); /* ||v-c||_1 */
714 25079 output->f_add = l2_sqr;
715 25079 output->f_rescale = l2_sqr * sqrt_d / l1_norm; /* dp_multiplier */
716
717 /* Cleanup */
718 25079 vs_free_aligned(residual);
719 25079 vs_free_aligned(transformed);
720 25079 vs_free_aligned(xu_cb);
721
722 25079 return 0;
723 }
724
725 /*
726 * Batch encode implementation — always_inline so that the specialized
727 * wrappers below pass a static const Vec32TypeOps from the header,
728 * enabling the compiler to inline through every vtable function pointer.
729 * VS_TARGET_CLONES on the wrappers generates AVX2/AVX-512 variants.
730 *
731 * For f32 input, to_float_block returns the input pointer (zero-copy).
732 * For f16, it bulk-converts all vectors in one SIMD-dispatched call.
733 * Bulk conversion via to_float_block amortizes call overhead for f16
734 * (one SIMD-dispatched call vs N per-vector calls).
735 */
736 __attribute__((always_inline)) static inline int
737 ✗ rabitq_encode_batch_impl(
738 const RaBitQParams *params,
739 const void *vectors,
740 Vec32Ref centroid,
741 float *f_add,
742 float *f_rescale,
743 uint8_t *bits,
744 uint16_t count,
745 const Vec32TypeOps *ops)
746 {
747 56 Dimension dim = params->dim;
748 56 uint32_t packed_bytes = params->packed_bytes;
749
750 /* Bulk-convert to f32 if needed. For f32, returns input pointer
751 * (zero-copy). For f16, converts into conv_buf via SIMD. */
752 ✗ float *conv_buf =
753 56 vs_alloc_aligned((size_t)count * dim * sizeof(float), 64);
754 56 const float *fvecs = ops->to_float_block(vectors, conv_buf, count, dim);
755
756 /* Allocate batch buffers */
757 ✗ float *residuals =
758 56 vs_alloc_aligned((size_t)count * dim * sizeof(float), 64);
759 ✗ float *transformed =
760 56 vs_alloc_aligned((size_t)count * dim * sizeof(float), 64);
761 56 float *cent_rotated = vs_alloc_aligned(dim * sizeof(float), 64);
762
763
3/18
✗ Branch 0 not taken.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
✗ Branch 6 not taken.
✗ Branch 7 not taken.
✗ Branch 8 not taken.
✗ Branch 9 not taken.
✗ Branch 10 not taken.
✗ Branch 11 not taken.
✓ Branch 12 taken 56 times.
✗ Branch 13 not taken.
✓ Branch 14 taken 56 times.
✗ Branch 15 not taken.
✗ Branch 16 not taken.
✓ Branch 17 taken 56 times.
 56 if (residuals == NULL || transformed == NULL || cent_rotated == NULL)
764 {
765 ✗ if (residuals)
766 ✗ vs_free_aligned(residuals);
767 ✗ if (transformed)
768 ✗ vs_free_aligned(transformed);
769 ✗ if (cent_rotated)
770 ✗ vs_free_aligned(cent_rotated);
771 ✗ vs_free_aligned(conv_buf);
772 ✗ return -1;
773 }
774
775 /* Step 1: Compute all residuals = vectors[i] - centroid */
776
2/6
✗ Branch 0 not taken.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✓ Branch 4 taken 460 times.
✓ Branch 5 taken 56 times.
 516 for (uint16_t i = 0; i < count; i++)
777 {
778 460 const float *vec = fvecs + i * dim;
779 460 float *res = residuals + i * dim;
780 460 vec32_sub(vec, centroid.data, res, dim);
781 }
782
783 /* Step 2: Batch transform all residuals through P^T
784 * This is the key optimization: matrix P stays in cache while
785 * processing all vectors.
786 */
787
1/6
✗ Branch 0 not taken.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✓ Branch 4 taken 56 times.
✗ Branch 5 not taken.
 56 if (params->use_fast_rotate)
788
2/6
✗ Branch 0 not taken.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✓ Branch 4 taken 460 times.
✓ Branch 5 taken 56 times.
 516 for (uint16_t i = 0; i < count; i++)
789 460 vs_fast_rotate_apply(
790 &params->fr,
791 460 residuals + (size_t)i * dim,
792 460 transformed + (size_t)i * dim);
793 else
794 ✗ vs_matrix_transpose_vector_mul_batch(
795 ✗ params->P, residuals, transformed, count, dim);
796
797 /* Step 3: Rotate centroid once (shared across all vectors) */
798 56 rabitq_rotate(params, centroid.data, cent_rotated);
799
800 /* Step 4: Process each transformed vector to extract bits and factors */
801
2/8
✗ Branch 0 not taken.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✓ Branch 4 taken 460 times.
✓ Branch 5 taken 56 times.
✗ Branch 7 not taken.
✗ Branch 8 not taken.
 516 for (uint16_t i = 0; i < count; i++)
802 {
803 460 const float *trans = transformed + i * dim;
804 460 uint8_t *vec_bits = bits + (size_t)i * packed_bytes;
805
806 /* Extract sign bits (LSB-first packing, FAISS-compatible) */
807 460 rabitq_extract_signs(trans, vec_bits, dim);
808
809 /* Compute xu_cb = binary_code - 0.5 */
810 460 float cb = -0.5f;
811 460 float ip_resi_xucb = 0.0f;
812 460 float ip_cent_xucb = 0.0f;
813 460 float l2_sqr = 0.0f;
814
815
2/8
✗ Branch 0 not taken.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✓ Branch 4 taken 110080 times.
✓ Branch 5 taken 460 times.
✗ Branch 7 not taken.
✗ Branch 8 not taken.
 110540 for (Dimension j = 0; j < dim; j++)
816 {
817 110080 int byte_idx = j / 8;
818 110080 int bit_idx = j % 8;
819 110080 int bit = (vec_bits[byte_idx] >> bit_idx) & 1;
820 110080 float xu_cb_j = (float)bit + cb;
821
822 110080 ip_resi_xucb += trans[j] * xu_cb_j;
823 110080 ip_cent_xucb += cent_rotated[j] * xu_cb_j;
824 110080 l2_sqr += trans[j] * trans[j];
825 }
826
827 /* ip_cent_xucb computed for future centered distance support */
828 ✗ (void)ip_cent_xucb;
829
830 /* Handle corner case: avoid division by near-zero in f_rescale */
831
1/6
✗ Branch 0 not taken.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✓ Branch 5 taken 460 times.
 460 if (fabsf(ip_resi_xucb) < 1e-10f)
832 ✗ ip_resi_xucb = 1e-10f;
833
834 /* L2 distance factors (FAISS-style)
835 * dp_multiplier = ||v-c||² * sqrt(d) / ||v-c||_1
836 * ip_resi_xucb = 0.5 * ||v-c||_1, so ||v-c||_1 = 2 * |ip_resi_xucb|
837 */
838 460 float sqrt_d = sqrtf((float)dim);
839 460 float l1_norm = 2.0f * fabsf(ip_resi_xucb);
840 460 f_add[i] = l2_sqr;
841 460 f_rescale[i] = l2_sqr * sqrt_d / l1_norm; /* dp_multiplier */
842 }
843
844 /* Cleanup */
845 56 vs_free_aligned(residuals);
846 56 vs_free_aligned(transformed);
847 56 vs_free_aligned(cent_rotated);
848 56 vs_free_aligned(conv_buf);
849
850 56 return 0;
851 }
852
853 /* Specialized wrappers — VS_TARGET_CLONES generates SIMD variants */
854
855 VS_TARGET_CLONES static int
856 56 rabitq_encode_batch_f32(
857 const RaBitQParams *params,
858 const void *vectors,
859 Vec32Ref centroid,
860 float *f_add,
861 float *f_rescale,
862 uint8_t *bits,
863 uint16_t count)
864 {
865 112 return rabitq_encode_batch_impl(
866 params,
867 vectors,
868 centroid,
869 f_add,
870 f_rescale,
871 bits,
872 count,
873 &vs_f32_type_ops);
874 }
875
876 VS_TARGET_CLONES static int
877 ✗ rabitq_encode_batch_f16(
878 const RaBitQParams *params,
879 const void *vectors,
880 Vec32Ref centroid,
881 float *f_add,
882 float *f_rescale,
883 uint8_t *bits,
884 uint16_t count)
885 {
886 ✗ return rabitq_encode_batch_impl(
887 params,
888 vectors,
889 centroid,
890 f_add,
891 f_rescale,
892 bits,
893 count,
894 &vs_f16_type_ops);
895 }
896
897 #if defined(VS_F16C_SUPPORT) && !defined(VS_SIMD_NONE)
898 VS_TARGET_F16C_AVX2 static int
899 ✗ rabitq_encode_batch_f16c(
900 const RaBitQParams *params,
901 const void *vectors,
902 Vec32Ref centroid,
903 float *f_add,
904 float *f_rescale,
905 uint8_t *bits,
906 uint16_t count)
907 {
908 ✗ return rabitq_encode_batch_impl(
909 params,
910 vectors,
911 centroid,
912 f_add,
913 f_rescale,
914 bits,
915 count,
916 &vs_f16c_type_ops);
917 }
918 #endif
919
920 int
921 72 vs_rabitq_encode_batch(
922 const RaBitQParams *params,
923 const void *vectors,
924 VecType vec_type,
925 Vec32Ref centroid,
926 float *f_add,
927 float *f_rescale,
928 uint8_t *bits,
929 uint16_t count)
930 {
931
8/8
✓ Branch 0 taken 70 times.
✓ Branch 1 taken 2 times.
✓ Branch 2 taken 68 times.
✓ Branch 3 taken 2 times.
✓ Branch 4 taken 66 times.
✓ Branch 5 taken 2 times.
✓ Branch 6 taken 64 times.
✓ Branch 7 taken 2 times.
 72 if (params == NULL || vectors == NULL || centroid.data == NULL ||
932
6/6
✓ Branch 0 taken 62 times.
✓ Branch 1 taken 2 times.
✓ Branch 2 taken 60 times.
✓ Branch 3 taken 2 times.
✓ Branch 4 taken 2 times.
✓ Branch 5 taken 58 times.
 64 f_add == NULL || f_rescale == NULL || bits == NULL || count == 0)
933 14 return -1;
934
935
2/2
✓ Branch 0 taken 2 times.
✓ Branch 1 taken 56 times.
 58 if (centroid.dim != params->dim)
936 2 return -1;
937
938 /* Single dispatch point — selects the inline vtable once */
939
1/3
✓ Branch 0 taken 56 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
 56 switch (vec_type)
940 {
941 56 case VS_VEC_F32:
942 56 return rabitq_encode_batch_f32(
943 params, vectors, centroid, f_add, f_rescale, bits, count);
944 #if defined(VS_F16C_SUPPORT) && !defined(VS_SIMD_NONE)
945 ✗ case VS_VEC_F16C:
946 ✗ return rabitq_encode_batch_f16c(
947 params, vectors, centroid, f_add, f_rescale, bits, count);
948 #endif
949 ✗ default:
950 ✗ return rabitq_encode_batch_f16(
951 params, vectors, centroid, f_add, f_rescale, bits, count);
952 }
953 }
954
955 /*
956 * Batch helpers
957 */
958
959 RaBitQBatch *
960 ✗ vs_rabitq_batch_create(uint16_t count, Dimension dim)
961 {
962 ✗ if (count == 0 || dim == 0)
963 ✗ return NULL;
964
965 ✗ RaBitQBatch *batch = vs_alloc(sizeof(RaBitQBatch));
966 ✗ if (batch == NULL)
967 ✗ return NULL;
968
969 ✗ uint32_t packed_bytes = VS_RABITQ_BYTES(dim);
970
971 ✗ batch->count = count;
972 ✗ batch->packed_bytes = (uint16_t)packed_bytes;
973 ✗ batch->f_add = vs_alloc(count * sizeof(float));
974 ✗ batch->f_rescale = vs_alloc(count * sizeof(float));
975 ✗ batch->bits = vs_alloc((size_t)count * packed_bytes);
976
977 ✗ if (batch->f_add == NULL || batch->f_rescale == NULL ||
978 ✗ batch->bits == NULL)
979 {
980 ✗ vs_rabitq_batch_destroy(batch);
981 ✗ return NULL;
982 }
983
984 ✗ return batch;
985 }
986
987 void
988 ✗ vs_rabitq_batch_destroy(RaBitQBatch *batch)
989 {
990 ✗ if (batch == NULL)
991 ✗ return;
992
993 ✗ if (batch->f_add)
994 ✗ vs_free(batch->f_add);
995 ✗ if (batch->f_rescale)
996 ✗ vs_free(batch->f_rescale);
997 ✗ if (batch->bits)
998 ✗ vs_free(batch->bits);
999 ✗ vs_free(batch);
1000 }
1001
1002 RaBitQBatch *
1003 ✗ vs_rabitq_encode_batch_alloc(
1004 const RaBitQParams *params,
1005 const void *vectors,
1006 VecType vec_type,
1007 Vec32Ref centroid,
1008 uint16_t count)
1009 {
1010 ✗ if (params == NULL || vectors == NULL || centroid.data == NULL ||
1011 count == 0)
1012 ✗ return NULL;
1013
1014 ✗ RaBitQBatch *batch = vs_rabitq_batch_create(count, params->dim);
1015 ✗ if (batch == NULL)
1016 ✗ return NULL;
1017
1018 ✗ if (vs_rabitq_encode_batch(
1019 params,
1020 vectors,
1021 vec_type,
1022 centroid,
1023 batch->f_add,
1024 batch->f_rescale,
1025 batch->bits,
1026 count) != 0)
1027 {
1028 ✗ vs_rabitq_batch_destroy(batch);
1029 ✗ return NULL;
1030 }
1031
1032 ✗ return batch;
1033 }
1034
1035 /*
1036 * Query preparation
1037 */
1038
1039 RaBitQQueryState *
1040 6760 vs_rabitq_prepare_query_ex(
1041 const RaBitQParams *params,
1042 Vec32Ref query,
1043 Vec32Ref centroid,
1044 VsDistanceMode mode)
1045 {
1046
6/6
✓ Branch 0 taken 6758 times.
✓ Branch 1 taken 2 times.
✓ Branch 2 taken 6756 times.
✓ Branch 3 taken 2 times.
✓ Branch 4 taken 2 times.
✓ Branch 5 taken 6754 times.
 6760 if (params == NULL || query.data == NULL || centroid.data == NULL)
1047 6 return NULL;
1048
1049
4/4
✓ Branch 0 taken 6752 times.
✓ Branch 1 taken 2 times.
✓ Branch 2 taken 2 times.
✓ Branch 3 taken 6750 times.
 6754 if (query.dim != params->dim || centroid.dim != params->dim)
1050 4 return NULL;
1051
1052 6750 Dimension dim = params->dim;
1053
1054 6750 RaBitQQueryState *state = vs_alloc(sizeof(RaBitQQueryState));
1055
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 6750 times.
 6750 if (state == NULL)
1056 ✗ return NULL;
1057
1058 6750 state->transformed = vs_alloc_aligned(dim * sizeof(float), 64);
1059
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 6750 times.
 6750 if (state->transformed == NULL)
1060 {
1061 ✗ vs_free(state);
1062 ✗ return NULL;
1063 }
1064
1065 6750 state->dim = dim;
1066
1067 /* Compute residual = query - centroid */
1068 6750 float *residual = vs_alloc_aligned(dim * sizeof(float), 64);
1069
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 6750 times.
 6750 if (residual == NULL)
1070 {
1071 ✗ vs_free_aligned(state->transformed);
1072 ✗ vs_free(state);
1073 ✗ return NULL;
1074 }
1075
1076 6750 vec32_sub(query.data, centroid.data, residual, dim);
1077
1078 /* Transform through P^T */
1079 6750 rabitq_rotate(params, residual, state->transformed);
1080
1081 /* Compute g_add = ||query - centroid||^2 */
1082 6750 state->g_add = vs_l2_norm_squared(state->transformed, dim);
1083
1084 /* Compute g_error = sqrt(g_add) for error bound */
1085 6750 state->g_error = sqrtf(state->g_add);
1086
1087 /* Compute sum of transformed values for FAISS-style distance formula */
1088 6750 state->sum_transformed = vec32_sum(state->transformed, dim);
1089
1090 /* Precompute 1/sqrt(dim) for distance formula */
1091 6750 state->inv_sqrt_d = 1.0f / sqrtf((float)dim);
1092
1093 /* Precompute C_error = 2*ε/√(d-1) for deriving f_error from compact
1094 * data */
1095
1/2
✓ Branch 0 taken 6750 times.
✗ Branch 1 not taken.
 6750 if (dim > 1)
1096 6750 state->c_error = 2.0f * VS_RABITQ_EPSILON / sqrtf((float)(dim - 1));
1097 else
1098 ✗ state->c_error = 0.0f;
1099
1100 /* Compute symmetric search fields: query sign bits and g_scale */
1101 6750 uint32_t packed_bytes = VS_RABITQ_BYTES(dim);
1102 6750 state->query_bits = vs_alloc_aligned(packed_bytes, 64);
1103
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 6750 times.
 6750 if (state->query_bits == NULL)
1104 {
1105 ✗ vs_free_aligned(state->transformed);
1106 ✗ vs_free_aligned(residual);
1107 ✗ vs_free(state);
1108 ✗ return NULL;
1109 }
1110 6750 rabitq_extract_signs(state->transformed, state->query_bits, dim);
1111
1112 /* g_scale = mean(|transformed|) = L1(transformed) / dim */
1113 6750 float l1_sum = 0.0f;
1114
2/3
✓ Branch 0 taken 783158 times.
✓ Branch 1 taken 6750 times.
✗ Branch 2 not taken.
 789908 for (Dimension i = 0; i < dim; i++)
1115 783158 l1_sum += fabsf(state->transformed[i]);
1116 6750 state->g_scale = l1_sum / (float)dim;
1117
1118 /* Set dispatch function pointers and error multiplier based on mode */
1119 6750 state->mode = mode;
1120
2/2
✓ Branch 0 taken 6 times.
✓ Branch 1 taken 6744 times.
 6750 if (mode == VS_DISTANCE_MODE_SYMMETRIC)
1121 {
1122 6 state->distance_fn = vs_rabitq_distance_symmetric;
1123 6 state->distance_with_bound_fn =
1124 vs_rabitq_distance_symmetric_with_bound;
1125 6 state->error_multiplier = 3.0f;
1126 }
1127 else
1128 {
1129 6744 state->distance_fn = vs_rabitq_distance;
1130 6744 state->distance_with_bound_fn = vs_rabitq_distance_with_bound;
1131 6744 state->error_multiplier = 1.0f;
1132 }
1133
1134 6750 vs_free_aligned(residual);
1135
1136 6750 return state;
1137 }
1138
1139 RaBitQQueryState *
1140 6750 vs_rabitq_prepare_query(
1141 const RaBitQParams *params, Vec32Ref query, Vec32Ref centroid)
1142 {
1143 6750 return vs_rabitq_prepare_query_ex(
1144 params, query, centroid, VS_DISTANCE_MODE_ASYMMETRIC);
1145 }
1146
1147 void
1148 667924 vs_rabitq_rotate(const RaBitQParams *params, const float *input, float *output)
1149 {
1150 667924 rabitq_rotate(params, input, output);
1151 667924 }
1152
1153 void
1154 36186 vs_rabitq_init_query_constants(RaBitQQueryState *state, Dimension dim)
1155 {
1156 36186 state->dim = dim;
1157 36186 state->inv_sqrt_d = 1.0f / sqrtf((float)dim);
1158
1/2
✓ Branch 0 taken 36186 times.
✗ Branch 1 not taken.
 36186 if (dim > 1)
1159 36186 state->c_error = 2.0f * VS_RABITQ_EPSILON / sqrtf((float)(dim - 1));
1160 else
1161 ✗ state->c_error = 0.0f;
1162 36186 }
1163
1164 void
1165 616469 vs_rabitq_init_query_state(
1166 RaBitQQueryState *state,
1167 const float *pt_query,
1168 const float *pt_centroid,
1169 Dimension dim,
1170 VsDistanceMode mode)
1171 {
1172 /* transformed = pt_query - pt_centroid (O(dim) vector subtraction) */
1173 616469 vec32_sub(pt_query, pt_centroid, state->transformed, dim);
1174
1175 /* Compute per-centroid scalar fields from transformed.
1176 * inv_sqrt_d and c_error are dim-dependent constants set once
1177 * via vs_rabitq_init_query_constants(). */
1178 616469 state->g_add = vs_l2_norm_squared(state->transformed, dim);
1179 616469 state->g_error = sqrtf(state->g_add);
1180 616469 state->sum_transformed = vec32_sum(state->transformed, dim);
1181
1182 /* Dispatch pointers */
1183 616469 state->mode = mode;
1184
2/2
✓ Branch 0 taken 12 times.
✓ Branch 1 taken 616457 times.
 616469 if (mode == VS_DISTANCE_MODE_SYMMETRIC)
1185 {
1186 /* Symmetric mode needs sign bits and L1 norm */
1187 12 rabitq_extract_signs(state->transformed, state->query_bits, dim);
1188
1189 12 float l1_sum = 0.0f;
1190
2/3
✓ Branch 0 taken 384 times.
✓ Branch 1 taken 12 times.
✗ Branch 2 not taken.
 396 for (Dimension i = 0; i < dim; i++)
1191 384 l1_sum += fabsf(state->transformed[i]);
1192 12 state->g_scale = l1_sum / (float)dim;
1193
1194 12 state->distance_fn = vs_rabitq_distance_symmetric;
1195 12 state->distance_with_bound_fn =
1196 vs_rabitq_distance_symmetric_with_bound;
1197 12 state->error_multiplier = 3.0f;
1198 }
1199 else
1200 {
1201 616457 state->distance_fn = vs_rabitq_distance;
1202 616457 state->distance_with_bound_fn = vs_rabitq_distance_with_bound;
1203 616457 state->error_multiplier = 1.0f;
1204 }
1205 616469 }
1206
1207 void
1208 6752 vs_rabitq_free_query(RaBitQQueryState *state)
1209 {
1210
2/2
✓ Branch 0 taken 2 times.
✓ Branch 1 taken 6750 times.
 6752 if (state == NULL)
1211 2 return;
1212
1213
1/2
✓ Branch 0 taken 6750 times.
✗ Branch 1 not taken.
 6750 if (state->transformed != NULL)
1214 6750 vs_free_aligned(state->transformed);
1215
1/2
✓ Branch 0 taken 6750 times.
✗ Branch 1 not taken.
 6750 if (state->query_bits != NULL)
1216 6750 vs_free_aligned(state->query_bits);
1217
1218 6750 vs_free(state);
1219 }
1220
1221 /*
1222 * Distance computation
1223 */
1224
1225 Distance
1226 6912 vs_rabitq_distance(
1227 const RaBitQQueryState *query_state,
1228 const RaBitQData *data,
1229 Dimension dim)
1230 {
1231
4/4
✓ Branch 0 taken 6908 times.
✓ Branch 1 taken 4 times.
✓ Branch 2 taken 4 times.
✓ Branch 3 taken 6904 times.
 6912 if (query_state == NULL || data == NULL)
1232 8 return -1.0f;
1233
1234
2/2
✓ Branch 0 taken 6 times.
✓ Branch 1 taken 6898 times.
 6904 if (query_state->dim != dim)
1235 6 return -1.0f;
1236
1237 /* Compute binary inner product: sum of transformed[i] where bit[i] = 1 */
1238 ✗ float binary_ip =
1239 6898 rabitq_inner_product(query_state->transformed, data->bits, dim);
1240
1241 /* FAISS-style distance formula (better accuracy with non-quantized query):
1242 *
1243 * est_dist = ||v-c||² + ||q-c||² - 2 * dp_multiplier * final_dot
1244 *
1245 * where:
1246 * - ||v-c||² = f_add (stored per vector)
1247 * - ||q-c||² = g_add (stored in query state)
1248 * - dp_multiplier = ||v-c||² * sqrt(d) / ||v-c||_1 = f_rescale
1249 * - final_dot = (2 * binary_ip - sum_transformed) / sqrt(d)
1250 *
1251 * FAISS normalizes dp_multiplier by the L1 norm to correct for the
1252 * distribution of residual values. This gives better distance estimates.
1253 */
1254 6898 float final_dot = (2.0f * binary_ip - query_state->sum_transformed) *
1255 6898 query_state->inv_sqrt_d;
1256
1257 6898 float est_dist = data->f_add + query_state->g_add -
1258 6898 2.0f * data->f_rescale * final_dot;
1259
1260 6898 return est_dist;
1261 }
1262
1263 /*
1264 * Shared distance formula for batch functions
1265 *
1266 * Applies: distances[i] = f_add[i] + g_add - 2 * f_rescale[i] * final_dots[i]
1267 * with optional error bound derivation using qstate->error_multiplier.
1268 *
1269 * Callers convert mode-specific intermediates (IPs or Hamming distances)
1270 * into uniform final_dots[] before calling this. For symmetric mode,
1271 * g_scale is folded into final_dots so the formula is identical.
1272 *
1273 * Marked always_inline so the compiler sees the full loop body in each
1274 * caller, preserving auto-vectorization of the f_add[]/f_rescale[] math.
1275 */
1276 __attribute__((always_inline)) static inline void
1277 293847 rabitq_apply_distances(
1278 const RaBitQQueryState *qstate,
1279 const float *f_add,
1280 const float *f_rescale,
1281 const float *final_dots,
1282 uint32_t count,
1283 Dimension dim,
1284 Distance *distances,
1285 Distance *lower_bounds)
1286 {
1287 1260755 float g_add = qstate->g_add;
1288
1289
2/2
✓ Branch 0 taken 14812285 times.
✓ Branch 1 taken 1260755 times.
 16073040 for (uint32_t i = 0; i < count; i++)
1290 {
1291 14812285 distances[i] = f_add[i] + g_add - 2.0f * f_rescale[i] * final_dots[i];
1292
1293
2/2
✓ Branch 0 taken 16 times.
✓ Branch 1 taken 14812269 times.
 14812285 if (lower_bounds != NULL)
1294 {
1295 16 float f_error = rabitq_derive_f_error(
1296 16 f_add[i], f_rescale[i], qstate->c_error, dim);
1297 16 lower_bounds[i] = rabitq_lower_bound(
1298 16 distances[i],
1299 f_error,
1300 16 qstate->g_error,
1301 16 qstate->error_multiplier);
1302 }
1303 }
1304 966908 }
1305
1306 void
1307 62 vs_rabitq_distance_batch(
1308 const RaBitQQueryState *qstate,
1309 const float *f_add,
1310 const float *f_rescale,
1311 const uint8_t *bits,
1312 uint32_t count,
1313 Dimension dim,
1314 Distance *distances)
1315 {
1316
10/10
✓ Branch 0 taken 60 times.
✓ Branch 1 taken 2 times.
✓ Branch 2 taken 58 times.
✓ Branch 3 taken 2 times.
✓ Branch 4 taken 56 times.
✓ Branch 5 taken 2 times.
✓ Branch 6 taken 54 times.
✓ Branch 7 taken 2 times.
✓ Branch 8 taken 52 times.
✓ Branch 9 taken 2 times.
 62 if (qstate == NULL || f_add == NULL || f_rescale == NULL || bits == NULL ||
1317
2/2
✓ Branch 0 taken 2 times.
✓ Branch 1 taken 50 times.
 52 distances == NULL || count == 0)
1318 12 return;
1319
1320 50 uint32_t packed_bytes = VS_RABITQ_BYTES(dim);
1321 50 float g_add = qstate->g_add;
1322 50 float sum_t = qstate->sum_transformed;
1323 50 float inv_sqrt_d = qstate->inv_sqrt_d;
1324
1325
2/2
✓ Branch 0 taken 412 times.
✓ Branch 1 taken 50 times.
 462 for (uint32_t i = 0; i < count; i++)
1326 {
1327 412 float ip = rabitq_inner_product(
1328 412 qstate->transformed, bits + (size_t)i * packed_bytes, dim);
1329 412 float final_dot = (2.0f * ip - sum_t) * inv_sqrt_d;
1330 412 distances[i] = f_add[i] + g_add - 2.0f * f_rescale[i] * final_dot;
1331 }
1332 }
1333
1334 void
1335 1289226 vs_rabitq_inner_product_multi(
1336 const float *transformed,
1337 const uint8_t *bits,
1338 uint32_t stride,
1339 Dimension dim,
1340 uint32_t count,
1341 float *results)
1342 {
1343
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 1289226 times.
 1289226 if (vs_unlikely(!g_rabitq_initialized))
1344 ✗ vs_rabitq_init_simd();
1345 1289226 g_inner_product_multi_fn(transformed, bits, stride, dim, count, results);
1346 1289226 }
1347
1348 void
1349 50 vs_rabitq_distance_batch_multi(
1350 const RaBitQQueryState *qstate,
1351 const float *f_add,
1352 const float *f_rescale,
1353 const uint8_t *bits,
1354 uint32_t count,
1355 Dimension dim,
1356 Distance *distances)
1357 {
1358 50 float *scratch = vs_alloc(count * sizeof(float));
1359 50 vs_rabitq_distance_batch_multi_with_bound(
1360 qstate,
1361 f_add,
1362 f_rescale,
1363 bits,
1364 50 VS_RABITQ_BYTES(dim),
1365 count,
1366 dim,
1367 distances,
1368 NULL,
1369 scratch);
1370 50 vs_free(scratch);
1371 50 }
1372
1373 void
1374 1260755 vs_rabitq_distance_batch_multi_with_bound(
1375 const RaBitQQueryState *qstate,
1376 const float *f_add,
1377 const float *f_rescale,
1378 const uint8_t *bits,
1379 uint32_t stride,
1380 uint32_t count,
1381 Dimension dim,
1382 Distance *distances,
1383 Distance *lower_bounds,
1384 float *scratch)
1385 {
1386
5/10
✓ Branch 0 taken 1260755 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 1260755 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 966908 times.
✗ Branch 5 not taken.
✓ Branch 6 taken 966908 times.
✗ Branch 7 not taken.
✓ Branch 8 taken 966908 times.
✗ Branch 9 not taken.
 1260755 if (qstate == NULL || f_add == NULL || f_rescale == NULL || bits == NULL ||
1387
2/2
✓ Branch 0 taken 293847 times.
✓ Branch 1 taken 966908 times.
 1260755 distances == NULL || count == 0)
1388 ✗ return;
1389
1390 /* Compute all inner products in a single multi-candidate pass */
1391 1260755 vs_rabitq_inner_product_multi(
1392 1260755 qstate->transformed, bits, stride, dim, count, scratch);
1393
1394 /* Convert IPs to final_dots in-place */
1395 1260755 float sum_t = qstate->sum_transformed;
1396 1260755 float inv_sqrt_d = qstate->inv_sqrt_d;
1397
2/2
✓ Branch 0 taken 14812285 times.
✓ Branch 1 taken 1260755 times.
 16073040 for (uint32_t i = 0; i < count; i++)
1398 14812285 scratch[i] = (2.0f * scratch[i] - sum_t) * inv_sqrt_d;
1399
1400 1260755 rabitq_apply_distances(
1401 qstate,
1402 f_add,
1403 f_rescale,
1404 scratch,
1405 count,
1406 dim,
1407 distances,
1408 lower_bounds);
1409 }
1410
1411 void
1412 12 vs_rabitq_distance_batch_symmetric_with_bound(
1413 const RaBitQQueryState *qstate,
1414 const float *f_add,
1415 const float *f_rescale,
1416 const uint8_t *bits,
1417 uint32_t stride,
1418 uint32_t count,
1419 Dimension dim,
1420 Distance *distances,
1421 Distance *lower_bounds,
1422 uint32_t *scratch)
1423 {
1424
5/10
✓ Branch 0 taken 12 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 12 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 12 times.
✗ Branch 5 not taken.
✓ Branch 6 taken 12 times.
✗ Branch 7 not taken.
✓ Branch 8 taken 12 times.
✗ Branch 9 not taken.
 12 if (qstate == NULL || f_add == NULL || f_rescale == NULL || bits == NULL ||
1425
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 12 times.
 12 distances == NULL || count == 0)
1426 ✗ return;
1427
1428 12 uint32_t packed_bytes = VS_RABITQ_BYTES(dim);
1429
1430 /* Compute all Hamming distances in a single multi-candidate pass */
1431 12 vs_rabitq_hamming_distance_multi(
1432 12 qstate->query_bits, bits, stride, packed_bytes, count, scratch);
1433
1434 /* Apply symmetric distance formula (+ optional error bounds) */
1435 12 float g_add = qstate->g_add;
1436 12 float g_scale = qstate->g_scale;
1437 12 float inv_sqrt_d = qstate->inv_sqrt_d;
1438
1439
2/2
✓ Branch 0 taken 66 times.
✓ Branch 1 taken 12 times.
 78 for (uint32_t i = 0; i < count; i++)
1440 {
1441 66 float sym_dot = (float)((int32_t)dim - 2 * (int32_t)scratch[i]);
1442 66 float final_dot = sym_dot * inv_sqrt_d * g_scale;
1443
1444 66 distances[i] = f_add[i] + g_add - 2.0f * f_rescale[i] * final_dot;
1445
1446
2/2
✓ Branch 0 taken 16 times.
✓ Branch 1 taken 50 times.
 66 if (lower_bounds != NULL)
1447 {
1448 16 float f_error = rabitq_derive_f_error(
1449 16 f_add[i], f_rescale[i], qstate->c_error, dim);
1450 16 lower_bounds[i] = rabitq_lower_bound(
1451 16 distances[i],
1452 f_error,
1453 16 qstate->g_error,
1454 16 qstate->error_multiplier);
1455 }
1456 }
1457 }
1458
1459 /*
1460 * Common scalar _with_bound implementation
1461 *
1462 * Uses qstate->distance_fn() (mode-dispatched) and
1463 * qstate->error_multiplier to compute both estimated distance
1464 * and lower bound. Both public _with_bound functions delegate here.
1465 */
1466 static void
1467 6736 rabitq_distance_with_bound_common(
1468 const RaBitQQueryState *qstate,
1469 const RaBitQData *data,
1470 Dimension dim,
1471 Distance *est_dist,
1472 Distance *lower_bound)
1473 {
1474
6/8
✓ Branch 0 taken 6732 times.
✓ Branch 1 taken 4 times.
✓ Branch 2 taken 6728 times.
✓ Branch 3 taken 4 times.
✓ Branch 4 taken 6728 times.
✗ Branch 5 not taken.
✗ Branch 6 not taken.
✓ Branch 7 taken 6728 times.
 6736 if (qstate == NULL || data == NULL || est_dist == NULL ||
1475 ✗ lower_bound == NULL)
1476 {
1477
2/2
✓ Branch 0 taken 6 times.
✓ Branch 1 taken 2 times.
 8 if (est_dist)
1478 6 *est_dist = -1.0f;
1479
2/2
✓ Branch 0 taken 6 times.
✓ Branch 1 taken 2 times.
 8 if (lower_bound)
1480 6 *lower_bound = -1.0f;
1481 8 return;
1482 }
1483
1484 6728 *est_dist = qstate->distance_fn(qstate, data, dim);
1485
1486 6728 float f_error = rabitq_derive_f_error(
1487 6728 data->f_add, data->f_rescale, qstate->c_error, dim);
1488 6728 *lower_bound = rabitq_lower_bound(
1489 6728 *est_dist, f_error, qstate->g_error, qstate->error_multiplier);
1490 }
1491
1492 void
1493 6634 vs_rabitq_distance_with_bound(
1494 const RaBitQQueryState *query_state,
1495 const RaBitQData *data,
1496 Dimension dim,
1497 Distance *est_dist,
1498 Distance *lower_bound)
1499 {
1500 6634 rabitq_distance_with_bound_common(
1501 query_state, data, dim, est_dist, lower_bound);
1502 6634 }
1503
1504 /*
1505 * Hamming distance - public API
1506 */
1507
1508 uint32_t
1509 172 vs_rabitq_hamming_distance(
1510 const uint8_t *a, const uint8_t *b, uint32_t packed_bytes)
1511 {
1512
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 172 times.
 172 if (vs_unlikely(!g_rabitq_initialized))
1513 ✗ vs_rabitq_init_simd();
1514 172 return g_hamming_fn(a, b, packed_bytes);
1515 }
1516
1517 void
1518 14 vs_rabitq_hamming_distance_multi(
1519 const uint8_t *query_bits,
1520 const uint8_t *data_bits,
1521 uint32_t stride,
1522 uint32_t packed_bytes,
1523 uint32_t count,
1524 uint32_t *results)
1525 {
1526
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 14 times.
 14 if (vs_unlikely(!g_rabitq_initialized))
1527 ✗ vs_rabitq_init_simd();
1528 14 g_hamming_multi_fn(
1529 query_bits, data_bits, stride, packed_bytes, count, results);
1530 14 }
1531
1532 /*
1533 * Symmetric distance computation
1534 *
1535 * Both query and data are 1-bit quantized. Uses Hamming distance
1536 * (XOR + popcount) instead of asymmetric inner product (mask + add).
1537 * ~32x fewer inner loop iterations at the cost of additional query
1538 * quantization error.
1539 */
1540
1541 Distance
1542 50 vs_rabitq_distance_symmetric(
1543 const RaBitQQueryState *qstate, const RaBitQData *data, Dimension dim)
1544 {
1545
4/4
✓ Branch 0 taken 48 times.
✓ Branch 1 taken 2 times.
✓ Branch 2 taken 2 times.
✓ Branch 3 taken 46 times.
 50 if (qstate == NULL || data == NULL)
1546 4 return -1.0f;
1547
1548
2/2
✓ Branch 0 taken 2 times.
✓ Branch 1 taken 44 times.
 46 if (qstate->dim != dim)
1549 2 return -1.0f;
1550
1551 44 uint32_t packed_bytes = VS_RABITQ_BYTES(dim);
1552 44 uint32_t hamming = vs_rabitq_hamming_distance(
1553 44 qstate->query_bits, data->bits, packed_bytes);
1554
1555 /* sym_dot = dim - 2 * hamming (range: [-dim, dim]) */
1556 44 float sym_dot = (float)((int32_t)dim - 2 * (int32_t)hamming);
1557 44 float final_dot = sym_dot * qstate->inv_sqrt_d;
1558
1559 /* est_dist = f_add + g_add - 2 * f_rescale * g_scale * final_dot */
1560 44 return data->f_add + qstate->g_add -
1561 44 2.0f * data->f_rescale * qstate->g_scale * final_dot;
1562 }
1563
1564 void
1565 102 vs_rabitq_distance_symmetric_with_bound(
1566 const RaBitQQueryState *qstate,
1567 const RaBitQData *data,
1568 Dimension dim,
1569 Distance *est_dist,
1570 Distance *lower_bound)
1571 {
1572 102 rabitq_distance_with_bound_common(
1573 qstate, data, dim, est_dist, lower_bound);
1574 102 }
1575
1576 void
1577 4 vs_rabitq_distance_batch_symmetric(
1578 const RaBitQQueryState *qstate,
1579 const float *f_add,
1580 const float *f_rescale,
1581 const uint8_t *bits,
1582 uint32_t count,
1583 Dimension dim,
1584 Distance *distances)
1585 {
1586 4 uint32_t *scratch = vs_alloc(count * sizeof(uint32_t));
1587 4 vs_rabitq_distance_batch_symmetric_with_bound(
1588 qstate,
1589 f_add,
1590 f_rescale,
1591 bits,
1592 4 VS_RABITQ_BYTES(dim),
1593 count,
1594 dim,
1595 distances,
1596 NULL,
1597 scratch);
1598 4 vs_free(scratch);
1599 4 }
1600