| Line | Branch | Exec | Source |
|---|---|---|---|
| 1 | /* | ||
| 2 | * Copyright (c) 2026 Tiger Data, Inc. | ||
| 3 | * Licensed under the PostgreSQL License. See LICENSE for details. | ||
| 4 | * | ||
| 5 | * distance.c - Distance computation with SIMD dispatch | ||
| 6 | * | ||
| 7 | * Implements runtime CPU detection and function pointer dispatch for | ||
| 8 | * optimal SIMD implementations, with compiler-vectorized fallback. | ||
| 9 | * | ||
| 10 | * SIMD modes (set via meson -Dsimd=): | ||
| 11 | * - full: Hand-optimized SIMD with compiler-vectorized fallback (default) | ||
| 12 | * - compiler: Compiler-vectorized only (target_clones for ISA selection) | ||
| 13 | * - none: Truly scalar (no vectorization, for debugging/baseline) | ||
| 14 | */ | ||
| 15 | |||
| 16 | /* Include generated config first for SIMD mode defines */ | ||
| 17 | #include "vs_config.h" | ||
| 18 | |||
| 19 | #include <math.h> | ||
| 20 | #include <stddef.h> | ||
| 21 | |||
| 22 | #include "algo/distance.h" | ||
| 23 | #include "algo/simd_utils.h" | ||
| 24 | #include "algo/vecops.h" | ||
| 25 | #include "core/platform.h" | ||
| 26 | |||
| 27 | /* | ||
| 28 | * Forward declarations for hand-optimized SIMD implementations. | ||
| 29 | * These are only available when VS_SIMD_FULL is defined (simd=full mode). | ||
| 30 | */ | ||
| 31 | |||
| 32 | #ifdef VS_SIMD_FULL | ||
| 33 | |||
| 34 | #if defined(__x86_64__) || defined(_M_X64) | ||
| 35 | /* AVX-512 implementations (x86-64 only) */ | ||
| 36 | Distance vs_distance_l2_avx512(Vec32Ref a, Vec32Ref b); | ||
| 37 | Distance vs_distance_ip_avx512(Vec32Ref a, Vec32Ref b); | ||
| 38 | Distance vs_distance_cosine_avx512(Vec32Ref a, Vec32Ref b); | ||
| 39 | |||
| 40 | int vs_distance_batch_l2_avx512( | ||
| 41 | Vec32Ref query, | ||
| 42 | const float *vectors, | ||
| 43 | uint32_t count, | ||
| 44 | Dimension dim, | ||
| 45 | Distance *distances); | ||
| 46 | int vs_distance_batch_ip_avx512( | ||
| 47 | Vec32Ref query, | ||
| 48 | const float *vectors, | ||
| 49 | uint32_t count, | ||
| 50 | Dimension dim, | ||
| 51 | Distance *distances); | ||
| 52 | int vs_distance_batch_cosine_avx512( | ||
| 53 | Vec32Ref query, | ||
| 54 | const float *vectors, | ||
| 55 | uint32_t count, | ||
| 56 | Dimension dim, | ||
| 57 | Distance *distances); | ||
| 58 | |||
| 59 | /* AVX2 implementations (x86-64 only) */ | ||
| 60 | Distance vs_distance_l2_avx2(Vec32Ref a, Vec32Ref b); | ||
| 61 | Distance vs_distance_ip_avx2(Vec32Ref a, Vec32Ref b); | ||
| 62 | Distance vs_distance_cosine_avx2(Vec32Ref a, Vec32Ref b); | ||
| 63 | |||
| 64 | int vs_distance_batch_l2_avx2( | ||
| 65 | Vec32Ref query, | ||
| 66 | const float *vectors, | ||
| 67 | uint32_t count, | ||
| 68 | Dimension dim, | ||
| 69 | Distance *distances); | ||
| 70 | int vs_distance_batch_ip_avx2( | ||
| 71 | Vec32Ref query, | ||
| 72 | const float *vectors, | ||
| 73 | uint32_t count, | ||
| 74 | Dimension dim, | ||
| 75 | Distance *distances); | ||
| 76 | int vs_distance_batch_cosine_avx2( | ||
| 77 | Vec32Ref query, | ||
| 78 | const float *vectors, | ||
| 79 | uint32_t count, | ||
| 80 | Dimension dim, | ||
| 81 | Distance *distances); | ||
| 82 | #endif | ||
| 83 | |||
| 84 | #if defined(__aarch64__) || defined(_M_ARM64) | ||
| 85 | /* NEON implementations (ARM only) */ | ||
| 86 | Distance vs_distance_l2_neon(Vec32Ref a, Vec32Ref b); | ||
| 87 | Distance vs_distance_ip_neon(Vec32Ref a, Vec32Ref b); | ||
| 88 | Distance vs_distance_cosine_neon(Vec32Ref a, Vec32Ref b); | ||
| 89 | |||
| 90 | int vs_distance_batch_l2_neon( | ||
| 91 | Vec32Ref query, | ||
| 92 | const float *vectors, | ||
| 93 | uint32_t count, | ||
| 94 | Dimension dim, | ||
| 95 | Distance *distances); | ||
| 96 | int vs_distance_batch_ip_neon( | ||
| 97 | Vec32Ref query, | ||
| 98 | const float *vectors, | ||
| 99 | uint32_t count, | ||
| 100 | Dimension dim, | ||
| 101 | Distance *distances); | ||
| 102 | int vs_distance_batch_cosine_neon( | ||
| 103 | Vec32Ref query, | ||
| 104 | const float *vectors, | ||
| 105 | uint32_t count, | ||
| 106 | Dimension dim, | ||
| 107 | Distance *distances); | ||
| 108 | #endif | ||
| 109 | |||
| 110 | #endif /* VS_SIMD_FULL */ | ||
| 111 | |||
| 112 | /* | ||
| 113 | * Compiler-Vectorized Implementations | ||
| 114 | * | ||
| 115 | * These use target_clones to generate multiple versions for different ISAs. | ||
| 116 | * The dynamic linker selects the best version at load time. | ||
| 117 | * | ||
| 118 | * These serve as: | ||
| 119 | * 1. Fallback when hand-optimized SIMD isn't available for the CPU | ||
| 120 | * 2. The only implementation when simd=compiler | ||
| 121 | * 3. Reference for comparing hand-optimized vs compiler-generated code | ||
| 122 | */ | ||
| 123 | |||
| 124 | /* | ||
| 125 | * L2 and IP distance use vecops for the core computation. | ||
| 126 | * This reduces code duplication and ensures consistent SIMD optimization. | ||
| 127 | */ | ||
| 128 | |||
| 129 | Distance | ||
| 130 | 122 | vs_distance_l2_compiler(Vec32Ref a, Vec32Ref b) | |
| 131 | { | ||
| 132 |
7/12✓ Branch 0 taken 122 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 122 times.
✓ Branch 4 taken 122 times.
✗ Branch 5 not taken.
✗ Branch 6 not taken.
✓ Branch 7 taken 122 times.
✓ Branch 8 taken 122 times.
✗ Branch 9 not taken.
✓ Branch 10 taken 2 times.
✓ Branch 11 taken 120 times.
|
122 | if (vs_unlikely( |
| 133 | a.dim != b.dim || a.dim == 0 || a.data == NULL || | ||
| 134 | b.data == NULL)) | ||
| 135 | 2 | return -1.0f; | |
| 136 | |||
| 137 | 120 | return vs_l2_distance_squared(a.data, b.data, a.dim); | |
| 138 | } | ||
| 139 | |||
| 140 | Distance | ||
| 141 | 56 | vs_distance_ip_compiler(Vec32Ref a, Vec32Ref b) | |
| 142 | { | ||
| 143 |
7/12✓ Branch 0 taken 56 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 56 times.
✓ Branch 4 taken 56 times.
✗ Branch 5 not taken.
✗ Branch 6 not taken.
✓ Branch 7 taken 56 times.
✓ Branch 8 taken 56 times.
✗ Branch 9 not taken.
✓ Branch 10 taken 2 times.
✓ Branch 11 taken 54 times.
|
56 | if (vs_unlikely( |
| 144 | a.dim != b.dim || a.dim == 0 || a.data == NULL || | ||
| 145 | b.data == NULL)) | ||
| 146 | 2 | return -1.0f; | |
| 147 | |||
| 148 | 54 | return -vs_dot_product(a.data, b.data, a.dim); | |
| 149 | } | ||
| 150 | |||
| 151 | /* Static helper with target_clones - returns similarity (not distance) */ | ||
| 152 | VS_TARGET_CLONES static double | ||
| 153 | 62 | compiler_cosine_similarity(int dim, const float *pa, const float *pb) | |
| 154 | { | ||
| 155 | 62 | float dot = 0.0f; | |
| 156 | 62 | float norm_a = 0.0f; | |
| 157 | 62 | float norm_b = 0.0f; | |
| 158 | |||
| 159 | /* Auto-vectorized - match pgvector style exactly */ | ||
| 160 |
2/2✓ Branch 0 taken 7194 times.
✓ Branch 1 taken 62 times.
|
7256 | for (int i = 0; i < dim; i++) |
| 161 | { | ||
| 162 | 7194 | dot += pa[i] * pb[i]; | |
| 163 | 7194 | norm_a += pa[i] * pa[i]; | |
| 164 | 7194 | norm_b += pb[i] * pb[i]; | |
| 165 | } | ||
| 166 | |||
| 167 | /* Use sqrt(a * b) over sqrt(a) * sqrt(b), with double precision */ | ||
| 168 | 62 | return (double)dot / sqrt((double)norm_a * (double)norm_b); | |
| 169 | } | ||
| 170 | |||
| 171 | Distance | ||
| 172 | 54 | vs_distance_cosine_compiler(Vec32Ref a, Vec32Ref b) | |
| 173 | { | ||
| 174 |
7/12✓ Branch 0 taken 54 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 54 times.
✓ Branch 4 taken 54 times.
✗ Branch 5 not taken.
✗ Branch 6 not taken.
✓ Branch 7 taken 54 times.
✓ Branch 8 taken 54 times.
✗ Branch 9 not taken.
✓ Branch 10 taken 2 times.
✓ Branch 11 taken 52 times.
|
54 | if (vs_unlikely( |
| 175 | a.dim != b.dim || a.dim == 0 || a.data == NULL || | ||
| 176 | b.data == NULL)) | ||
| 177 | 2 | return -1.0f; | |
| 178 | |||
| 179 | 52 | float similarity = compiler_cosine_similarity(a.dim, a.data, b.data); | |
| 180 | |||
| 181 | /* Handle edge cases */ | ||
| 182 |
3/4✓ Branch 0 taken 50 times.
✓ Branch 1 taken 2 times.
✗ Branch 2 not taken.
✓ Branch 3 taken 50 times.
|
52 | if (isnan(similarity) || similarity < -1.0f) |
| 183 | 2 | return 1.0f; /* Maximum distance for zero/invalid vectors */ | |
| 184 |
1/2✗ Branch 0 not taken.
✓ Branch 1 taken 50 times.
|
50 | if (similarity > 1.0f) |
| 185 | ✗ | similarity = 1.0f; | |
| 186 | |||
| 187 | 50 | return 1.0f - similarity; | |
| 188 | } | ||
| 189 | |||
| 190 | /* | ||
| 191 | * Compiler-Vectorized Batch Implementations | ||
| 192 | * | ||
| 193 | * Dedicated batch functions that avoid function call overhead in the inner | ||
| 194 | * loop. | ||
| 195 | */ | ||
| 196 | |||
| 197 | int | ||
| 198 | 2 | vs_distance_batch_l2_compiler( | |
| 199 | Vec32Ref query, | ||
| 200 | const float *vectors, | ||
| 201 | uint32_t count, | ||
| 202 | Dimension dim, | ||
| 203 | Distance *distances) | ||
| 204 | { | ||
| 205 |
6/12✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 2 times.
✓ Branch 4 taken 2 times.
✗ Branch 5 not taken.
✗ Branch 6 not taken.
✓ Branch 7 taken 2 times.
✓ Branch 8 taken 2 times.
✗ Branch 9 not taken.
✗ Branch 10 not taken.
✓ Branch 11 taken 2 times.
|
2 | if (vs_unlikely( |
| 206 | query.dim != dim || query.data == NULL || vectors == NULL || | ||
| 207 | distances == NULL)) | ||
| 208 | ✗ | return -1; | |
| 209 | |||
| 210 | 2 | const float *q = query.data; | |
| 211 | |||
| 212 |
2/2✓ Branch 0 taken 10 times.
✓ Branch 1 taken 2 times.
|
12 | for (uint32_t i = 0; i < count; i++) |
| 213 | { | ||
| 214 | 10 | distances[i] = vs_l2_distance_squared(q, vectors + i * dim, dim); | |
| 215 | } | ||
| 216 | |||
| 217 | 2 | return 0; | |
| 218 | } | ||
| 219 | |||
| 220 | int | ||
| 221 | 2 | vs_distance_batch_ip_compiler( | |
| 222 | Vec32Ref query, | ||
| 223 | const float *vectors, | ||
| 224 | uint32_t count, | ||
| 225 | Dimension dim, | ||
| 226 | Distance *distances) | ||
| 227 | { | ||
| 228 |
6/12✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 2 times.
✓ Branch 4 taken 2 times.
✗ Branch 5 not taken.
✗ Branch 6 not taken.
✓ Branch 7 taken 2 times.
✓ Branch 8 taken 2 times.
✗ Branch 9 not taken.
✗ Branch 10 not taken.
✓ Branch 11 taken 2 times.
|
2 | if (vs_unlikely( |
| 229 | query.dim != dim || query.data == NULL || vectors == NULL || | ||
| 230 | distances == NULL)) | ||
| 231 | ✗ | return -1; | |
| 232 | |||
| 233 | 2 | const float *q = query.data; | |
| 234 | |||
| 235 |
2/2✓ Branch 0 taken 10 times.
✓ Branch 1 taken 2 times.
|
12 | for (uint32_t i = 0; i < count; i++) |
| 236 | { | ||
| 237 | 10 | distances[i] = -vs_dot_product(q, vectors + i * dim, dim); | |
| 238 | } | ||
| 239 | |||
| 240 | 2 | return 0; | |
| 241 | } | ||
| 242 | |||
| 243 | int | ||
| 244 | 2 | vs_distance_batch_cosine_compiler( | |
| 245 | Vec32Ref query, | ||
| 246 | const float *vectors, | ||
| 247 | uint32_t count, | ||
| 248 | Dimension dim, | ||
| 249 | Distance *distances) | ||
| 250 | { | ||
| 251 |
6/12✓ Branch 0 taken 2 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 2 times.
✓ Branch 4 taken 2 times.
✗ Branch 5 not taken.
✗ Branch 6 not taken.
✓ Branch 7 taken 2 times.
✓ Branch 8 taken 2 times.
✗ Branch 9 not taken.
✗ Branch 10 not taken.
✓ Branch 11 taken 2 times.
|
2 | if (vs_unlikely( |
| 252 | query.dim != dim || query.data == NULL || vectors == NULL || | ||
| 253 | distances == NULL)) | ||
| 254 | ✗ | return -1; | |
| 255 | |||
| 256 | 2 | const float *q = query.data; | |
| 257 | |||
| 258 |
2/2✓ Branch 0 taken 10 times.
✓ Branch 1 taken 2 times.
|
12 | for (uint32_t i = 0; i < count; i++) |
| 259 | { | ||
| 260 | 10 | float similarity = | |
| 261 | 10 | compiler_cosine_similarity(dim, q, vectors + i * dim); | |
| 262 | |||
| 263 | /* Handle edge cases */ | ||
| 264 |
2/4✓ Branch 0 taken 10 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 10 times.
|
10 | if (isnan(similarity) || similarity < -1.0f) |
| 265 | ✗ | distances[i] = 1.0f; | |
| 266 |
1/2✗ Branch 0 not taken.
✓ Branch 1 taken 10 times.
|
10 | else if (similarity > 1.0f) |
| 267 | ✗ | distances[i] = 0.0f; | |
| 268 | else | ||
| 269 | 10 | distances[i] = 1.0f - similarity; | |
| 270 | } | ||
| 271 | |||
| 272 | 2 | return 0; | |
| 273 | } | ||
| 274 | |||
| 275 | /* | ||
| 276 | * Function Pointer Dispatch | ||
| 277 | * | ||
| 278 | * Global function pointers are initialized once at startup based on detected | ||
| 279 | * CPU capabilities. After initialization, dispatch has zero overhead (direct | ||
| 280 | * function pointer call). | ||
| 281 | */ | ||
| 282 | |||
| 283 | /* Single-pair function pointers */ | ||
| 284 | typedef Distance (*DistanceFn)(Vec32Ref, Vec32Ref); | ||
| 285 | static DistanceFn g_distance_l2_fn = NULL; | ||
| 286 | static DistanceFn g_distance_ip_fn = NULL; | ||
| 287 | static DistanceFn g_distance_cosine_fn = NULL; | ||
| 288 | |||
| 289 | /* Batch function pointers */ | ||
| 290 | typedef int (*BatchDistanceFn)( | ||
| 291 | Vec32Ref, const float *, uint32_t, Dimension, Distance *); | ||
| 292 | static BatchDistanceFn g_batch_l2_fn = NULL; | ||
| 293 | static BatchDistanceFn g_batch_ip_fn = NULL; | ||
| 294 | static BatchDistanceFn g_batch_cosine_fn = NULL; | ||
| 295 | |||
| 296 | /* Implementation name and initialization flag */ | ||
| 297 | static const char *g_impl_name = NULL; | ||
| 298 | static _Atomic(bool) g_initialized = false; | ||
| 299 | |||
| 300 | void | ||
| 301 | 290 | vs_distance_force_reinit(void) | |
| 302 | { | ||
| 303 | 290 | g_initialized = false; | |
| 304 | 290 | g_distance_l2_fn = NULL; | |
| 305 | 290 | g_distance_ip_fn = NULL; | |
| 306 | 290 | g_distance_cosine_fn = NULL; | |
| 307 | 290 | g_batch_l2_fn = NULL; | |
| 308 | 290 | g_batch_ip_fn = NULL; | |
| 309 | 290 | g_batch_cosine_fn = NULL; | |
| 310 | 290 | g_impl_name = NULL; | |
| 311 | 290 | } | |
| 312 | |||
| 313 | int | ||
| 314 | 646 | vs_distance_init(void) | |
| 315 | { | ||
| 316 | /* Double-checked locking pattern for thread-safe lazy init */ | ||
| 317 |
2/2✓ Branch 0 taken 352 times.
✓ Branch 1 taken 294 times.
|
646 | if (g_initialized) |
| 318 | 97 | return 0; | |
| 319 | |||
| 320 | #ifdef VS_SIMD_FULL | ||
| 321 | /* Full mode: prefer hand-optimized, fall back to compiler-vectorized */ | ||
| 322 | 549 | SimdCapability caps = vs_detect_simd(); | |
| 323 | |||
| 324 | #if defined(__x86_64__) || defined(_M_X64) | ||
| 325 |
1/2✗ Branch 0 not taken.
✓ Branch 1 taken 549 times.
|
549 | if ((caps & VS_SIMD_AVX512_DQ) == VS_SIMD_AVX512_DQ) |
| 326 | { | ||
| 327 | ✗ | g_distance_l2_fn = vs_distance_l2_avx512; | |
| 328 | ✗ | g_distance_ip_fn = vs_distance_ip_avx512; | |
| 329 | ✗ | g_distance_cosine_fn = vs_distance_cosine_avx512; | |
| 330 | ✗ | g_batch_l2_fn = vs_distance_batch_l2_avx512; | |
| 331 | ✗ | g_batch_ip_fn = vs_distance_batch_ip_avx512; | |
| 332 | ✗ | g_batch_cosine_fn = vs_distance_batch_cosine_avx512; | |
| 333 | ✗ | g_impl_name = "avx512"; | |
| 334 | } | ||
| 335 |
2/2✓ Branch 0 taken 409 times.
✓ Branch 1 taken 140 times.
|
549 | else if (caps & SIMD_AVX2) |
| 336 | { | ||
| 337 | 409 | g_distance_l2_fn = vs_distance_l2_avx2; | |
| 338 | 409 | g_distance_ip_fn = vs_distance_ip_avx2; | |
| 339 | 409 | g_distance_cosine_fn = vs_distance_cosine_avx2; | |
| 340 | 409 | g_batch_l2_fn = vs_distance_batch_l2_avx2; | |
| 341 | 409 | g_batch_ip_fn = vs_distance_batch_ip_avx2; | |
| 342 | 409 | g_batch_cosine_fn = vs_distance_batch_cosine_avx2; | |
| 343 | 409 | g_impl_name = "avx2"; | |
| 344 | } | ||
| 345 | else | ||
| 346 | { | ||
| 347 | 140 | g_distance_l2_fn = vs_distance_l2_compiler; | |
| 348 | 140 | g_distance_ip_fn = vs_distance_ip_compiler; | |
| 349 | 140 | g_distance_cosine_fn = vs_distance_cosine_compiler; | |
| 350 | 140 | g_batch_l2_fn = vs_distance_batch_l2_compiler; | |
| 351 | 140 | g_batch_ip_fn = vs_distance_batch_ip_compiler; | |
| 352 | 140 | g_batch_cosine_fn = vs_distance_batch_cosine_compiler; | |
| 353 | 140 | g_impl_name = "compiler"; | |
| 354 | } | ||
| 355 | #elif defined(__aarch64__) || defined(_M_ARM64) | ||
| 356 | if (caps & SIMD_NEON) | ||
| 357 | { | ||
| 358 | g_distance_l2_fn = vs_distance_l2_neon; | ||
| 359 | g_distance_ip_fn = vs_distance_ip_neon; | ||
| 360 | g_distance_cosine_fn = vs_distance_cosine_neon; | ||
| 361 | g_batch_l2_fn = vs_distance_batch_l2_neon; | ||
| 362 | g_batch_ip_fn = vs_distance_batch_ip_neon; | ||
| 363 | g_batch_cosine_fn = vs_distance_batch_cosine_neon; | ||
| 364 | g_impl_name = "neon"; | ||
| 365 | } | ||
| 366 | else | ||
| 367 | { | ||
| 368 | g_distance_l2_fn = vs_distance_l2_compiler; | ||
| 369 | g_distance_ip_fn = vs_distance_ip_compiler; | ||
| 370 | g_distance_cosine_fn = vs_distance_cosine_compiler; | ||
| 371 | g_batch_l2_fn = vs_distance_batch_l2_compiler; | ||
| 372 | g_batch_ip_fn = vs_distance_batch_ip_compiler; | ||
| 373 | g_batch_cosine_fn = vs_distance_batch_cosine_compiler; | ||
| 374 | g_impl_name = "compiler"; | ||
| 375 | } | ||
| 376 | #else | ||
| 377 | /* No hand-optimized kernels for this architecture: use the | ||
| 378 | * compiler-vectorized ones regardless of what the CPU offers. */ | ||
| 379 | (void)caps; | ||
| 380 | g_distance_l2_fn = vs_distance_l2_compiler; | ||
| 381 | g_distance_ip_fn = vs_distance_ip_compiler; | ||
| 382 | g_distance_cosine_fn = vs_distance_cosine_compiler; | ||
| 383 | g_batch_l2_fn = vs_distance_batch_l2_compiler; | ||
| 384 | g_batch_ip_fn = vs_distance_batch_ip_compiler; | ||
| 385 | g_batch_cosine_fn = vs_distance_batch_cosine_compiler; | ||
| 386 | g_impl_name = "compiler"; | ||
| 387 | #endif | ||
| 388 | |||
| 389 | #elif defined(VS_SIMD_NONE) | ||
| 390 | /* None mode: truly scalar (no vectorization) */ | ||
| 391 | g_distance_l2_fn = vs_distance_l2_compiler; | ||
| 392 | g_distance_ip_fn = vs_distance_ip_compiler; | ||
| 393 | g_distance_cosine_fn = vs_distance_cosine_compiler; | ||
| 394 | g_batch_l2_fn = vs_distance_batch_l2_compiler; | ||
| 395 | g_batch_ip_fn = vs_distance_batch_ip_compiler; | ||
| 396 | g_batch_cosine_fn = vs_distance_batch_cosine_compiler; | ||
| 397 | g_impl_name = "none"; | ||
| 398 | |||
| 399 | #else | ||
| 400 | /* Compiler mode (simd=compiler): use target_clones for ISA selection */ | ||
| 401 | g_distance_l2_fn = vs_distance_l2_compiler; | ||
| 402 | g_distance_ip_fn = vs_distance_ip_compiler; | ||
| 403 | g_distance_cosine_fn = vs_distance_cosine_compiler; | ||
| 404 | g_batch_l2_fn = vs_distance_batch_l2_compiler; | ||
| 405 | g_batch_ip_fn = vs_distance_batch_ip_compiler; | ||
| 406 | g_batch_cosine_fn = vs_distance_batch_cosine_compiler; | ||
| 407 | g_impl_name = "compiler"; | ||
| 408 | #endif | ||
| 409 | |||
| 410 | 549 | g_initialized = true; | |
| 411 | 549 | return 0; | |
| 412 | } | ||
| 413 | |||
| 414 | /* | ||
| 415 | * Public API - Single-Pair Distance Functions | ||
| 416 | * | ||
| 417 | * Use IFUNC (indirect function) for zero-overhead dispatch when available. | ||
| 418 | * Falls back to manual dispatch for testing or non-GNU toolchains. | ||
| 419 | * | ||
| 420 | * IFUNC is only used in simd=full mode where we need runtime dispatch. | ||
| 421 | * In simd=compiler mode, target_clones handles ISA selection directly. | ||
| 422 | */ | ||
| 423 | |||
| 424 | /* | ||
| 425 | * Use IFUNC only when: | ||
| 426 | * - simd=full mode (VS_SIMD_FULL) | ||
| 427 | * - not explicitly disabled (VS_DISABLE_IFUNC) | ||
| 428 | * - not a sanitizer build (VS_SANITIZER) - resolver runs before sanitizer | ||
| 429 | * init | ||
| 430 | * - not a coverage build (VS_COVERAGE) - resolver code skews coverage stats | ||
| 431 | * - supported platform (x86-64 or AArch64 with GCC) | ||
| 432 | */ | ||
| 433 | #if defined(VS_SIMD_FULL) && defined(__GNUC__) && \ | ||
| 434 | !defined(VS_DISABLE_IFUNC) && !defined(VS_SANITIZER) && \ | ||
| 435 | !defined(VS_COVERAGE) && \ | ||
| 436 | (defined(__x86_64__) || defined(__aarch64__)) | ||
| 437 | #define USE_IFUNC 1 | ||
| 438 | #else | ||
| 439 | #define USE_IFUNC 0 | ||
| 440 | #endif | ||
| 441 | |||
| 442 | #if USE_IFUNC | ||
| 443 | |||
| 444 | /* IFUNC resolvers - called once at program load */ | ||
| 445 | |||
| 446 | static Distance (*resolve_distance_l2(void))(Vec32Ref, Vec32Ref) | ||
| 447 | { | ||
| 448 | SimdCapability caps = vs_detect_simd(); | ||
| 449 | |||
| 450 | #if defined(__x86_64__) || defined(_M_X64) | ||
| 451 | if ((caps & VS_SIMD_AVX512_DQ) == VS_SIMD_AVX512_DQ) | ||
| 452 | return vs_distance_l2_avx512; | ||
| 453 | else if (caps & SIMD_AVX2) | ||
| 454 | return vs_distance_l2_avx2; | ||
| 455 | #elif defined(__aarch64__) || defined(_M_ARM64) | ||
| 456 | if (caps & SIMD_NEON) | ||
| 457 | return vs_distance_l2_neon; | ||
| 458 | #endif | ||
| 459 | |||
| 460 | return vs_distance_l2_compiler; | ||
| 461 | } | ||
| 462 | |||
| 463 | static Distance (*resolve_distance_ip(void))(Vec32Ref, Vec32Ref) | ||
| 464 | { | ||
| 465 | SimdCapability caps = vs_detect_simd(); | ||
| 466 | |||
| 467 | #if defined(__x86_64__) || defined(_M_X64) | ||
| 468 | if ((caps & VS_SIMD_AVX512_DQ) == VS_SIMD_AVX512_DQ) | ||
| 469 | return vs_distance_ip_avx512; | ||
| 470 | else if (caps & SIMD_AVX2) | ||
| 471 | return vs_distance_ip_avx2; | ||
| 472 | #elif defined(__aarch64__) || defined(_M_ARM64) | ||
| 473 | if (caps & SIMD_NEON) | ||
| 474 | return vs_distance_ip_neon; | ||
| 475 | #endif | ||
| 476 | |||
| 477 | return vs_distance_ip_compiler; | ||
| 478 | } | ||
| 479 | |||
| 480 | static Distance (*resolve_distance_cosine(void))(Vec32Ref, Vec32Ref) | ||
| 481 | { | ||
| 482 | SimdCapability caps = vs_detect_simd(); | ||
| 483 | |||
| 484 | #if defined(__x86_64__) || defined(_M_X64) | ||
| 485 | if ((caps & VS_SIMD_AVX512_DQ) == VS_SIMD_AVX512_DQ) | ||
| 486 | return vs_distance_cosine_avx512; | ||
| 487 | else if (caps & SIMD_AVX2) | ||
| 488 | return vs_distance_cosine_avx2; | ||
| 489 | #elif defined(__aarch64__) || defined(_M_ARM64) | ||
| 490 | if (caps & SIMD_NEON) | ||
| 491 | return vs_distance_cosine_neon; | ||
| 492 | #endif | ||
| 493 | |||
| 494 | return vs_distance_cosine_compiler; | ||
| 495 | } | ||
| 496 | |||
| 497 | /* Public API with IFUNC attribute */ | ||
| 498 | |||
| 499 | Distance vs_distance_l2(Vec32Ref a, Vec32Ref b) | ||
| 500 | __attribute__((ifunc("resolve_distance_l2"))); | ||
| 501 | |||
| 502 | Distance vs_distance_ip(Vec32Ref a, Vec32Ref b) | ||
| 503 | __attribute__((ifunc("resolve_distance_ip"))); | ||
| 504 | |||
| 505 | Distance vs_distance_cosine(Vec32Ref a, Vec32Ref b) | ||
| 506 | __attribute__((ifunc("resolve_distance_cosine"))); | ||
| 507 | |||
| 508 | #else | ||
| 509 | |||
| 510 | /* Fallback: manual dispatch (for testing or non-GNU toolchains) */ | ||
| 511 | |||
| 512 | Distance | ||
| 513 | 739999 | vs_distance_l2(Vec32Ref a, Vec32Ref b) | |
| 514 | { | ||
| 515 |
2/2✓ Branch 0 taken 1 times.
✓ Branch 1 taken 739998 times.
|
739999 | if (vs_unlikely(!g_initialized)) |
| 516 | 1 | vs_distance_init(); | |
| 517 | 739999 | return g_distance_l2_fn(a, b); | |
| 518 | } | ||
| 519 | |||
| 520 | Distance | ||
| 521 | 152 | vs_distance_ip(Vec32Ref a, Vec32Ref b) | |
| 522 | { | ||
| 523 |
1/2✗ Branch 0 not taken.
✓ Branch 1 taken 152 times.
|
152 | if (vs_unlikely(!g_initialized)) |
| 524 | ✗ | vs_distance_init(); | |
| 525 | 152 | return g_distance_ip_fn(a, b); | |
| 526 | } | ||
| 527 | |||
| 528 | Distance | ||
| 529 | 16101 | vs_distance_cosine(Vec32Ref a, Vec32Ref b) | |
| 530 | { | ||
| 531 |
1/2✗ Branch 0 not taken.
✓ Branch 1 taken 16101 times.
|
16101 | if (vs_unlikely(!g_initialized)) |
| 532 | ✗ | vs_distance_init(); | |
| 533 | 16101 | return g_distance_cosine_fn(a, b); | |
| 534 | } | ||
| 535 | |||
| 536 | #endif /* USE_IFUNC */ | ||
| 537 | |||
| 538 | Distance | ||
| 539 | 677793 | vs_distance(Vec32Ref a, Vec32Ref b, DistanceMetric metric) | |
| 540 | { | ||
| 541 |
4/4✓ Branch 0 taken 668405 times.
✓ Branch 1 taken 2 times.
✓ Branch 2 taken 9384 times.
✓ Branch 3 taken 2 times.
|
677793 | switch (metric) |
| 542 | { | ||
| 543 | 668405 | case DISTANCE_L2: | |
| 544 | 668405 | return vs_distance_l2(a, b); | |
| 545 | 2 | case DISTANCE_INNER_PRODUCT: | |
| 546 | 2 | return vs_distance_ip(a, b); | |
| 547 | 9384 | case DISTANCE_COSINE: | |
| 548 | 9384 | return vs_distance_cosine(a, b); | |
| 549 | 2 | default: | |
| 550 | 2 | return -1.0f; | |
| 551 | } | ||
| 552 | } | ||
| 553 | |||
| 554 | /* | ||
| 555 | * Public API - Batch Distance Functions | ||
| 556 | */ | ||
| 557 | |||
| 558 | int | ||
| 559 | 16 | vs_distance_batch_l2( | |
| 560 | Vec32Ref query, | ||
| 561 | const float *vectors, | ||
| 562 | uint32_t count, | ||
| 563 | Dimension dim, | ||
| 564 | Distance *distances) | ||
| 565 | { | ||
| 566 |
1/2✗ Branch 0 not taken.
✓ Branch 1 taken 16 times.
|
16 | if (vs_unlikely(!g_initialized)) |
| 567 | ✗ | vs_distance_init(); | |
| 568 | |||
| 569 |
12/12✓ Branch 0 taken 14 times.
✓ Branch 1 taken 2 times.
✓ Branch 2 taken 2 times.
✓ Branch 3 taken 12 times.
✓ Branch 4 taken 12 times.
✓ Branch 5 taken 4 times.
✓ Branch 6 taken 2 times.
✓ Branch 7 taken 10 times.
✓ Branch 8 taken 10 times.
✓ Branch 9 taken 6 times.
✓ Branch 10 taken 2 times.
✓ Branch 11 taken 8 times.
|
16 | if (vs_unlikely( |
| 570 | query.dim != dim || query.data == NULL || vectors == NULL || | ||
| 571 | distances == NULL)) | ||
| 572 | 8 | return -1; | |
| 573 | |||
| 574 | /* Use SIMD batch implementation if available */ | ||
| 575 |
1/2✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
|
8 | if (g_batch_l2_fn != NULL) |
| 576 | 8 | return g_batch_l2_fn(query, vectors, count, dim, distances); | |
| 577 | |||
| 578 | /* Fallback: loop over single-pair function */ | ||
| 579 | ✗ | for (uint32_t i = 0; i < count; i++) | |
| 580 | { | ||
| 581 | ✗ | Vec32Ref vec = {.data = vectors + i * dim, .dim = dim}; | |
| 582 | ✗ | distances[i] = vs_distance_l2(query, vec); | |
| 583 | } | ||
| 584 | |||
| 585 | ✗ | return 0; | |
| 586 | } | ||
| 587 | |||
| 588 | int | ||
| 589 | 10 | vs_distance_batch_ip( | |
| 590 | Vec32Ref query, | ||
| 591 | const float *vectors, | ||
| 592 | uint32_t count, | ||
| 593 | Dimension dim, | ||
| 594 | Distance *distances) | ||
| 595 | { | ||
| 596 |
1/2✗ Branch 0 not taken.
✓ Branch 1 taken 10 times.
|
10 | if (vs_unlikely(!g_initialized)) |
| 597 | ✗ | vs_distance_init(); | |
| 598 | |||
| 599 |
8/12✓ Branch 0 taken 10 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 10 times.
✓ Branch 4 taken 10 times.
✗ Branch 5 not taken.
✓ Branch 6 taken 2 times.
✓ Branch 7 taken 8 times.
✓ Branch 8 taken 8 times.
✓ Branch 9 taken 2 times.
✗ Branch 10 not taken.
✓ Branch 11 taken 8 times.
|
10 | if (vs_unlikely( |
| 600 | query.dim != dim || query.data == NULL || vectors == NULL || | ||
| 601 | distances == NULL)) | ||
| 602 | 2 | return -1; | |
| 603 | |||
| 604 |
1/2✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
|
8 | if (g_batch_ip_fn != NULL) |
| 605 | 8 | return g_batch_ip_fn(query, vectors, count, dim, distances); | |
| 606 | |||
| 607 | ✗ | for (uint32_t i = 0; i < count; i++) | |
| 608 | { | ||
| 609 | ✗ | Vec32Ref vec = {.data = vectors + i * dim, .dim = dim}; | |
| 610 | ✗ | distances[i] = vs_distance_ip(query, vec); | |
| 611 | } | ||
| 612 | |||
| 613 | ✗ | return 0; | |
| 614 | } | ||
| 615 | |||
| 616 | int | ||
| 617 | 10 | vs_distance_batch_cosine( | |
| 618 | Vec32Ref query, | ||
| 619 | const float *vectors, | ||
| 620 | uint32_t count, | ||
| 621 | Dimension dim, | ||
| 622 | Distance *distances) | ||
| 623 | { | ||
| 624 |
1/2✗ Branch 0 not taken.
✓ Branch 1 taken 10 times.
|
10 | if (vs_unlikely(!g_initialized)) |
| 625 | ✗ | vs_distance_init(); | |
| 626 | |||
| 627 |
7/12✓ Branch 0 taken 10 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 10 times.
✓ Branch 4 taken 10 times.
✗ Branch 5 not taken.
✗ Branch 6 not taken.
✓ Branch 7 taken 10 times.
✓ Branch 8 taken 10 times.
✗ Branch 9 not taken.
✓ Branch 10 taken 2 times.
✓ Branch 11 taken 8 times.
|
10 | if (vs_unlikely( |
| 628 | query.dim != dim || query.data == NULL || vectors == NULL || | ||
| 629 | distances == NULL)) | ||
| 630 | 2 | return -1; | |
| 631 | |||
| 632 |
1/2✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
|
8 | if (g_batch_cosine_fn != NULL) |
| 633 | 8 | return g_batch_cosine_fn(query, vectors, count, dim, distances); | |
| 634 | |||
| 635 | ✗ | for (uint32_t i = 0; i < count; i++) | |
| 636 | { | ||
| 637 | ✗ | Vec32Ref vec = {.data = vectors + i * dim, .dim = dim}; | |
| 638 | ✗ | distances[i] = vs_distance_cosine(query, vec); | |
| 639 | } | ||
| 640 | |||
| 641 | ✗ | return 0; | |
| 642 | } | ||
| 643 | |||
| 644 | const char * | ||
| 645 | 18 | vs_distance_impl_name(void) | |
| 646 | { | ||
| 647 |
1/2✗ Branch 0 not taken.
✓ Branch 1 taken 18 times.
|
18 | if (vs_unlikely(!g_initialized)) |
| 648 | ✗ | vs_distance_init(); | |
| 649 | 18 | return g_impl_name; | |
| 650 | } | ||
| 651 |