GCC Code Coverage Report

Directory: src/
File: src/algo/kmeans.c
Date: 2026-10-23 09:17:32
Exec Total Coverage
Lines: 408 452 90.3%
Functions: 31 33 93.9%
Branches: 244 441 55.3%
Line Branch Exec Source
1 /*
2 * Copyright (c) 2026 Tiger Data, Inc.
3 * Licensed under the PostgreSQL License. See LICENSE for details.
4 *
5 * kmeans.c - K-means clustering orchestration
6 *
7 * Common infrastructure shared by all k-means variants:
8 * - k-means++ initialization (D²-weighted sampling)
9 * - Centroid update step (mean of assigned vectors)
10 * - Empty cluster handling (largest-cluster splitting)
11 * - Convergence checking (max centroid shift)
12 * - State management and public API
13 *
14 * Algorithm-specific assignment steps live in separate files:
15 * - kmeans_lloyd.c: brute-force (CBLAS sgemm or builtin batch)
16 * - kmeans_hamerly.c: Hamerly's single-bound acceleration
17 * - kmeans_elkan.c: Elkan's per-centroid bounds
18 */
19
20 #include "vs_config.h"
21
22 #include <float.h>
23 #include <math.h>
24 #include <stdint.h>
25 #include <stdlib.h>
26 #include <string.h>
27
28 #include "algo/kmeans_elkan.h"
29 #include "algo/kmeans_hamerly.h"
30 #include "algo/kmeans_internal.h"
31 #include "algo/kmeans_lloyd.h"
32 #include "algo/vecops.h"
33 #include "core/log.h"
34 #include "core/memory.h"
35
36 /*
37 * Xoshiro256** PRNG (same as matrix.c - duplicated to keep kmeans.c
38 * self-contained without exposing PRNG in a shared header).
39 */
40 typedef struct
41 {
42 uint64_t s[4];
43 } Xoshiro256State;
44
45 static inline uint64_t
46 22860 xo_rotl(uint64_t x, int k)
47 {
48 22860 return (x << k) | (x >> (64 - k));
49 }
50
51 static uint64_t
52 16884 xo_next(Xoshiro256State *state)
53 {
54 16884 uint64_t *s = state->s;
55 16884 uint64_t result = xo_rotl(s[1] * 5, 7) * 9;
56 16884 uint64_t t = s[1] << 17;
57
58 16884 s[2] ^= s[0];
59 16884 s[3] ^= s[1];
60 16884 s[1] ^= s[2];
61 16884 s[0] ^= s[3];
62 16884 s[2] ^= t;
63 16884 s[3] = xo_rotl(s[3], 45);
64
65 16884 return result;
66 }
67
68 static void
69 2683 xo_seed(Xoshiro256State *state, uint64_t seed)
70 {
71
2/2
✓ Branch 0 taken 10732 times.
✓ Branch 1 taken 2683 times.
 13415 for (int i = 0; i < 4; i++)
72 {
73 10732 seed += 0x9e3779b97f4a7c15ULL;
74 10732 uint64_t z = seed;
75 10732 z = (z ^ (z >> 30)) * 0xbf58476d1ce4e5b9ULL;
76 10732 z = (z ^ (z >> 27)) * 0x94d049bb133111ebULL;
77 10732 state->s[i] = z ^ (z >> 31);
78 }
79 2683 }
80
81 /* Random double in [0, 1) */
82 static double
83 16884 xo_uniform(Xoshiro256State *state)
84 {
85 27792 uint64_t x = xo_next(state) >> 11;
86 16884 return (double)x / (double)(1ULL << 53);
87 }
88
89 /* CBLAS runtime toggle (same pattern as matrix.c) */
90 static bool g_use_cblas = true;
91
92 void
93 6 vs_kmeans_set_use_cblas(bool use_cblas)
94 {
95 6 g_use_cblas = use_cblas;
96 6 }
97
98 bool
99 2 vs_kmeans_get_use_cblas(void)
100 {
101 #ifdef VS_HAVE_CBLAS
102 return g_use_cblas;
103 #else
104 2 return false;
105 #endif
106 }
107
108 const char *
109 2 vs_kmeans_impl_name(void)
110 {
111 #ifdef VS_HAVE_CBLAS
112 return g_use_cblas ? "cblas" : "builtin";
113 #else
114 2 return "builtin";
115 #endif
116 }
117
118 static const char *algo_names[] = {
119 [KMEANS_ALGO_AUTO] = "auto",
120 [KMEANS_ALGO_LLOYD] = "lloyd",
121 [KMEANS_ALGO_HAMERLY] = "hamerly",
122 [KMEANS_ALGO_ELKAN] = "elkan",
123 [KMEANS_ALGO_CBLAS] = "lloyd(cblas)",
124 };
125
126 const char *
127 12 vs_kmeans_algo_name(KMeansAlgorithm algo)
128 {
129
2/2
✓ Branch 0 taken 10 times.
✓ Branch 1 taken 2 times.
 12 if ((unsigned)algo <= KMEANS_ALGO_CBLAS)
130 10 return algo_names[algo];
131 2 return "unknown";
132 }
133
134 bool
135 2 vs_cblas_is_single_threaded(void)
136 {
137 2 const char *v = getenv("OMP_NUM_THREADS");
138
1/6
✗ Branch 0 not taken.
✓ Branch 1 taken 2 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
 2 return v != NULL && v[0] == '1' && v[1] == '\0';
139 }
140
141 /*
142 * Pin BLAS to single-threaded operation via the vendor's runtime API.
143 *
144 * Only symbols for the BLAS vendor detected at configure time are
145 * referenced. Declaring all vendors' externs and relying on weak
146 * symbols doesn't survive macOS's linker (Mach-O has no portable
147 * "undefined weak reference resolves to NULL" that holds up under
148 * LTO).
149 *
150 * Supported vendors: OpenBLAS (openblas_set_num_threads),
151 * BLIS/AOCL (bli_thread_set_num_threads). Others (Accelerate,
152 * MKL, ARM PL) must be pinned via env vars.
153 */
154 #ifdef VS_BLAS_OPENBLAS
155 extern void openblas_set_num_threads(int);
156 #endif
157 #ifdef VS_BLAS_BLIS
158 extern void bli_thread_set_num_threads(long);
159 #endif
160
161 void
162 255 vs_cblas_pin_single_thread(void)
163 {
164 #ifdef VS_BLAS_OPENBLAS
165 openblas_set_num_threads(1);
166 #endif
167 #ifdef VS_BLAS_BLIS
168 bli_thread_set_num_threads(1);
169 #endif
170 255 }
171
172 /*
173 * Precompute ||x||^2 for all input vectors.
174 */
175 __attribute__((always_inline)) static inline void
176 precompute_norms_x_impl(KMeansState *st, const Vec32TypeOps *ops)
177 {
178 1120 size_t esz = ops->element_size;
179
4/6
✗ Branch 0 not taken.
✗ Branch 1 not taken.
✓ Branch 2 taken 2000 times.
✓ Branch 3 taken 12 times.
✓ Branch 4 taken 790118 times.
✓ Branch 5 taken 2573 times.
 794703 for (uint32_t i = 0; i < st->nvecs; i++)
180
4/7
✗ Branch 0 not taken.
✗ Branch 1 not taken.
✗ Branch 3 not taken.
✓ Branch 4 taken 2000 times.
✓ Branch 5 taken 300851 times.
✓ Branch 6 taken 401637 times.
✓ Branch 7 taken 87630 times.
 1331393 st->norms_x[i] = ops->norm_sq(km_get_vector(st, i, esz), st->dim);
181 1120 }
182
183 /*
184 * Compute distance from a typed vector to a float32 centroid.
185 */
186 __attribute__((always_inline)) static inline float
187 5934327 vector_centroid_distance_impl(
188 DistanceMetric metric,
189 const void *vec,
190 const float *centroid,
191 Dimension dim,
192 const Vec32TypeOps *ops)
193 {
194 8337713 switch (metric)
195 {
196 8050419 case DISTANCE_L2:
197 8050419 return ops->l2_squared(vec, centroid, dim);
198 101600 case DISTANCE_INNER_PRODUCT:
199 101600 return -ops->dot_product(vec, centroid, dim);
200 185694 case DISTANCE_COSINE:
201 185694 return 1.0f - ops->dot_product(vec, centroid, dim);
202 }
203 ✗ return FLT_MAX;
204 }
205
206 /*
207 * k-means++ initialization.
208 */
209 __attribute__((always_inline)) static inline void
210 1517 kmeans_init_plusplus_impl(
211 KMeansState *st, uint64_t seed, const Vec32TypeOps *ops)
212 {
213 1517 Xoshiro256State rng;
214 2683 xo_seed(&rng, seed);
215
216 2683 Dimension dim = st->dim;
217 2683 uint32_t nvecs = st->nvecs;
218 2683 uint32_t nlist = st->nlist;
219 2683 size_t esz = ops->element_size;
220 2683 float *dists = vs_alloc(nvecs * sizeof(float));
221
222 /* 1. First centroid: random vector */
223 2683 uint32_t idx = (uint32_t)(xo_uniform(&rng) * nvecs);
224
2/6
✗ Branch 0 not taken.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 12 times.
✗ Branch 4 not taken.
✓ Branch 5 taken 2671 times.
 2683 if (idx >= nvecs)
225 ✗ idx = nvecs - 1;
226
4/8
✗ Branch 0 not taken.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✓ Branch 4 taken 1276 times.
✓ Branch 5 taken 253 times.
✓ Branch 6 taken 928 times.
✓ Branch 7 taken 226 times.
 3849 ops->to_float_one(km_get_vector(st, idx, esz), st->centroids, dim);
227
228 /* Initialize distances to first centroid */
229
6/8
✗ Branch 0 not taken.
✗ Branch 1 not taken.
✓ Branch 2 taken 2000 times.
✓ Branch 3 taken 12 times.
✓ Branch 4 taken 253784 times.
✓ Branch 5 taken 1154 times.
✓ Branch 7 taken 569714 times.
✓ Branch 8 taken 1517 times.
 828181 for (uint32_t i = 0; i < nvecs; i++)
230
4/12
✗ Branch 0 not taken.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✓ Branch 4 taken 2000 times.
✗ Branch 5 not taken.
✗ Branch 6 not taken.
✗ Branch 7 not taken.
✓ Branch 8 taken 790118 times.
✓ Branch 9 taken 15200 times.
✓ Branch 10 taken 18180 times.
✗ Branch 11 not taken.
 1650996 dists[i] = vector_centroid_distance_impl(
231 st->metric,
232 km_get_vector(st, i, esz),
233
3/6
✗ Branch 0 not taken.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 2000 times.
✓ Branch 4 taken 470865 times.
✓ Branch 5 taken 352633 times.
 825498 st->centroids,
234 dim,
235 ops);
236
237 /* 2. Pick remaining centroids */
238
4/6
✗ Branch 0 not taken.
✗ Branch 1 not taken.
✓ Branch 2 taken 32 times.
✓ Branch 3 taken 12 times.
✓ Branch 4 taken 14169 times.
✓ Branch 5 taken 2671 times.
 16884 for (uint32_t k = 1; k < nlist; k++)
239 {
240 /* Compute cumulative distribution */
241 4810 double total = 0.0;
242
4/6
✗ Branch 0 not taken.
✗ Branch 1 not taken.
✓ Branch 2 taken 5600 times.
✓ Branch 3 taken 32 times.
✓ Branch 4 taken 7506615 times.
✓ Branch 5 taken 14169 times.
 7526416 for (uint32_t i = 0; i < nvecs; i++)
243 7512215 total += (double)dists[i];
244
245 /* Sample from distribution */
246 14201 double r = xo_uniform(&rng) * total;
247 14201 double cum = 0.0;
248 14201 uint32_t chosen = 0;
249
2/6
✗ Branch 0 not taken.
✗ Branch 1 not taken.
✓ Branch 2 taken 4164 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 4832602 times.
✗ Branch 5 not taken.
 4836766 for (uint32_t i = 0; i < nvecs; i++)
250 {
251 4836766 cum += (double)dists[i];
252
4/6
✗ Branch 0 not taken.
✗ Branch 1 not taken.
✓ Branch 2 taken 32 times.
✓ Branch 3 taken 4132 times.
✓ Branch 4 taken 3564758 times.
✓ Branch 5 taken 1267844 times.
 4836766 if (cum >= r)
253 {
254 4810 chosen = i;
255 4810 break;
256 }
257 }
258
259 /* Copy chosen vector as new centroid (convert to f32) */
260 23592 ops->to_float_one(
261 km_get_vector(st, chosen, esz),
262
3/6
✗ Branch 0 not taken.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 32 times.
✓ Branch 4 taken 10161 times.
✓ Branch 5 taken 4008 times.
 14201 st->centroids + (size_t)k * dim,
263 dim);
264
265 /* Update min distances */
266
6/8
✗ Branch 0 not taken.
✗ Branch 1 not taken.
✓ Branch 2 taken 5600 times.
✓ Branch 3 taken 32 times.
✓ Branch 4 taken 2142002 times.
✓ Branch 5 taken 4778 times.
✓ Branch 7 taken 5364613 times.
✓ Branch 8 taken 9391 times.
 7526416 for (uint32_t i = 0; i < nvecs; i++)
267 {
268
4/12
✗ Branch 0 not taken.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✓ Branch 4 taken 5600 times.
✗ Branch 5 not taken.
✗ Branch 6 not taken.
✗ Branch 7 not taken.
✓ Branch 8 taken 7252701 times.
✓ Branch 9 taken 86400 times.
✓ Branch 10 taken 167514 times.
✗ Branch 11 not taken.
 7512215 float d = vector_centroid_distance_impl(
269 st->metric,
270 km_get_vector(st, i, esz),
271
3/6
✗ Branch 0 not taken.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 5600 times.
✓ Branch 4 taken 3378220 times.
✓ Branch 5 taken 4128395 times.
 7512215 st->centroids + (size_t)k * dim,
272 dim,
273 ops);
274
4/6
✗ Branch 0 not taken.
✗ Branch 1 not taken.
✓ Branch 2 taken 1920 times.
✓ Branch 3 taken 3680 times.
✓ Branch 4 taken 1249377 times.
✓ Branch 5 taken 6257238 times.
 7512215 if (d < dists[i])
275 1251297 dists[i] = d;
276 }
277 }
278
279 2683 vs_free(dists);
280 2683 }
281
282 /*
283 * Update step: recompute centroids as mean of assigned vectors.
284 */
285 __attribute__((always_inline)) static inline void
286 ✗ kmeans_update_centroids_impl(KMeansState *st, const Vec32TypeOps *ops)
287 {
288 208 uint32_t nlist = st->nlist;
289 208 Dimension dim = st->dim;
290 208 size_t esz = ops->element_size;
291
292 /* Zero accumulators */
293 208 memset(st->new_centroids, 0, (size_t)nlist * dim * sizeof(float));
294 208 memset(st->cluster_sizes, 0, nlist * sizeof(uint32_t));
295
296 /* Accumulate */
297
4/6
✗ Branch 0 not taken.
✗ Branch 1 not taken.
✓ Branch 2 taken 13600 times.
✓ Branch 3 taken 72 times.
✓ Branch 4 taken 22800 times.
✓ Branch 5 taken 136 times.
 36608 for (uint32_t i = 0; i < st->nvecs; i++)
298 {
299 36400 ClusterId c = st->assignments[i];
300
2/6
✗ Branch 0 not taken.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 13600 times.
✗ Branch 4 not taken.
✓ Branch 5 taken 22800 times.
 36400 st->cluster_sizes[c]++;
301
0/6
✗ Branch 0 not taken.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
 36400 const void *vec = km_get_vector(st, i, esz);
302 36400 float *cent = st->new_centroids + (size_t)c * dim;
303 36400 ops->sum_to_float(vec, cent, dim);
304 }
305
306 /* Divide by cluster size */
307
4/6
✗ Branch 0 not taken.
✗ Branch 1 not taken.
✓ Branch 2 taken 328 times.
✓ Branch 3 taken 72 times.
✓ Branch 4 taken 560 times.
✓ Branch 5 taken 136 times.
 1096 for (uint32_t j = 0; j < nlist; j++)
308 {
309
2/6
✗ Branch 0 not taken.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 328 times.
✗ Branch 4 not taken.
✓ Branch 5 taken 560 times.
 888 if (st->cluster_sizes[j] == 0)
310 ✗ continue;
311
312 888 float inv_size = 1.0f / (float)st->cluster_sizes[j];
313 888 vec32_scale(
314 888 st->new_centroids + (size_t)j * dim,
315 inv_size,
316 888 st->new_centroids + (size_t)j * dim,
317 dim);
318 }
319
320 /* For cosine: normalize centroids to unit length */
321
2/6
✗ Branch 0 not taken.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 72 times.
✗ Branch 4 not taken.
✓ Branch 5 taken 136 times.
 208 if (st->metric == DISTANCE_COSINE)
322 {
323 ✗ for (uint32_t j = 0; j < nlist; j++)
324 {
325 ✗ if (st->cluster_sizes[j] == 0)
326 ✗ continue;
327 ✗ float *cent = st->new_centroids + (size_t)j * dim;
328 ✗ float norm = vs_l2_norm(cent, dim);
329 ✗ if (norm > 1e-10f)
330 ✗ vec32_scale(cent, 1.0f / norm, cent, dim);
331 }
332 }
333
334 /* Swap new centroids into place */
335 208 float *tmp = st->centroids;
336 208 st->centroids = st->new_centroids;
337 208 st->new_centroids = tmp;
338 208 }
339
340 /*
341 * Handle empty clusters by splitting the largest cluster.
342 *
343 * FAISS approach: replace empty centroid with a perturbation of the
344 * largest cluster's centroid.
345 */
346 static void
347 2835 kmeans_handle_empty_clusters(KMeansState *st)
348 {
349 2835 uint32_t dim = st->dim;
350 2835 uint32_t nlist = st->nlist;
351
352
2/2
✓ Branch 0 taken 17568 times.
✓ Branch 1 taken 2835 times.
 20403 for (uint32_t j = 0; j < nlist; j++)
353 {
354
2/2
✓ Branch 0 taken 17441 times.
✓ Branch 1 taken 127 times.
 17568 if (st->cluster_sizes[j] > 0)
355 17441 continue;
356
357 /* Find largest cluster */
358 127 uint32_t largest = 0;
359 127 uint32_t largest_size = st->cluster_sizes[0];
360
2/2
✓ Branch 0 taken 2055 times.
✓ Branch 1 taken 127 times.
 2182 for (uint32_t k = 1; k < nlist; k++)
361 {
362
2/2
✓ Branch 0 taken 93 times.
✓ Branch 1 taken 1962 times.
 2055 if (st->cluster_sizes[k] > largest_size)
363 {
364 93 largest = k;
365 93 largest_size = st->cluster_sizes[k];
366 }
367 }
368
369
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 127 times.
 127 if (largest_size <= 1)
370 ✗ continue; /* Can't split a cluster of size 1 */
371
372 /* Perturb: empty = largest * (1 + eps), largest *= (1 - eps) */
373 127 float *c_empty = st->centroids + (size_t)j * dim;
374 127 float *c_largest = st->centroids + (size_t)largest * dim;
375
376
2/2
✓ Branch 0 taken 2198 times.
✓ Branch 1 taken 127 times.
 2325 for (uint32_t d = 0; d < dim; d++)
377 {
378 2198 float val = c_largest[d];
379 2198 c_empty[d] = val * (1.0f + 1e-4f);
380 2198 c_largest[d] = val * (1.0f - 1e-4f);
381 }
382
383 /* Split size estimate (roughly half each) */
384 127 uint32_t half = largest_size / 2;
385 127 st->cluster_sizes[j] = half;
386 127 st->cluster_sizes[largest] -= half;
387 }
388 2835 }
389
390 /*
391 * Max centroid movement between two centroid arrays.
392 * Returns the maximum squared L2 shift.
393 */
394 float
395 19826 kmeans_max_centroid_shift_between(
396 const float *a, const float *b, uint32_t nlist, Dimension dim)
397 {
398 19826 float max_shift = 0.0f;
399
400
2/2
✓ Branch 0 taken 142884 times.
✓ Branch 1 taken 19826 times.
 162710 for (uint32_t j = 0; j < nlist; j++)
401 {
402 242866 float shift = vs_l2_distance_squared(
403 142884 a + (size_t)j * dim, b + (size_t)j * dim, dim);
404
2/2
✓ Branch 0 taken 37146 times.
✓ Branch 1 taken 105738 times.
 142884 if (shift > max_shift)
405 37146 max_shift = shift;
406 }
407
408 19826 return max_shift;
409 }
410
411 VS_TARGET_CLONES void
412 3389 kmeans_assign_accumulate(
413 const float *vectors,
414 const uint32_t *indices,
415 uint32_t start,
416 uint32_t end,
417 const float *centroids,
418 const float *norms_c,
419 uint32_t k,
420 Dimension dim,
421 DistanceMetric metric,
422 const uint32_t *filter,
423 uint32_t filter_val,
424 float *out_sums,
425 uint32_t *out_cnts,
426 float *out_cost)
427 {
428 3389 float cost = 0.0f;
429
430
2/2
✓ Branch 0 taken 3317643 times.
✓ Branch 1 taken 3389 times.
 3321032 for (uint32_t i = start; i < end; i++)
431 {
432
1/4
✗ Branch 0 not taken.
✓ Branch 1 taken 3317643 times.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
 3317643 if (filter != NULL && filter[i] != filter_val)
433 ✗ continue;
434
435
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 3317643 times.
 3317643 uint32_t idx = indices ? indices[i] : i;
436 3317643 const float *vec = vectors + (size_t)idx * dim;
437 3317643 float best_d = __FLT_MAX__;
438 3317643 uint32_t best_c = 0;
439
440
2/2
✓ Branch 0 taken 32520054 times.
✓ Branch 1 taken 3317643 times.
 35837697 for (uint32_t c = 0; c < k; c++)
441 {
442 32520054 const float *cent = centroids + (size_t)c * dim;
443 16063490 float d;
444
445
3/3
✓ Branch 0 taken 31449326 times.
✓ Branch 1 taken 360000 times.
✓ Branch 2 taken 710728 times.
 32520054 switch (metric)
446 {
447 31449326 default: /* L2 */
448 {
449 31449326 float nx = vs_l2_norm_squared(vec, dim);
450 31449326 float dot = vs_dot_product(vec, cent, dim);
451 31449326 d = nx + norms_c[c] - 2.0f * dot;
452
2/2
✓ Branch 0 taken 61 times.
✓ Branch 1 taken 31449265 times.
 31449326 if (d < 0.0f)
453 61 d = 0.0f;
454 16052964 break;
455 }
456 360000 case DISTANCE_INNER_PRODUCT:
457 360000 d = -vs_dot_product(vec, cent, dim);
458 360000 break;
459 710728 case DISTANCE_COSINE:
460 710728 d = 1.0f - vs_dot_product(vec, cent, dim);
461 710728 break;
462 }
463
464
2/2
✓ Branch 0 taken 9491477 times.
✓ Branch 1 taken 23028577 times.
 32520054 if (d < best_d)
465 {
466 9491477 best_d = d;
467 9491477 best_c = c;
468 }
469 }
470
471 3317643 cost += best_d;
472 3317643 out_cnts[best_c]++;
473 3317643 float *sum = out_sums + (size_t)best_c * dim;
474
2/2
✓ Branch 0 taken 114863651 times.
✓ Branch 1 taken 3317643 times.
 118181294 for (uint32_t d = 0; d < dim; d++)
475 114863651 sum[d] += vec[d];
476 }
477
478 3389 *out_cost += cost;
479 3389 }
480
481 VS_TARGET_CLONES void
482 326 kmeans_assign(
483 const float *vectors,
484 uint32_t start,
485 uint32_t end,
486 const float *centroids,
487 const float *norms_c,
488 uint32_t k,
489 Dimension dim,
490 DistanceMetric metric,
491 uint32_t *out_assignments)
492 {
493
2/2
✓ Branch 0 taken 226330 times.
✓ Branch 1 taken 326 times.
 226656 for (uint32_t i = start; i < end; i++)
494 {
495 226330 const float *vec = vectors + (size_t)i * dim;
496 226330 float best_d = __FLT_MAX__;
497 226330 uint32_t best_c = 0;
498
499
2/2
✓ Branch 0 taken 1769443 times.
✓ Branch 1 taken 226330 times.
 1995773 for (uint32_t c = 0; c < k; c++)
500 {
501 1769443 const float *cent = centroids + (size_t)c * dim;
502 885561 float d;
503
504
3/3
✓ Branch 0 taken 1709823 times.
✓ Branch 1 taken 18000 times.
✓ Branch 2 taken 41620 times.
 1769443 switch (metric)
505 {
506 1709823 default: /* L2 */
507 {
508 1709823 float nx = vs_l2_norm_squared(vec, dim);
509 1709823 float dot = vs_dot_product(vec, cent, dim);
510 1709823 d = nx + norms_c[c] - 2.0f * dot;
511
2/2
✓ Branch 0 taken 3 times.
✓ Branch 1 taken 1709820 times.
 1709823 if (d < 0.0f)
512 3 d = 0.0f;
513 862082 break;
514 }
515 18000 case DISTANCE_INNER_PRODUCT:
516 18000 d = -vs_dot_product(vec, cent, dim);
517 18000 break;
518 41620 case DISTANCE_COSINE:
519 41620 d = 1.0f - vs_dot_product(vec, cent, dim);
520 41620 break;
521 }
522
523
2/2
✓ Branch 0 taken 553634 times.
✓ Branch 1 taken 1215809 times.
 1769443 if (d < best_d)
524 {
525 553634 best_d = d;
526 553634 best_c = c;
527 }
528 }
529
530 226330 out_assignments[i] = best_c;
531 }
532 326 }
533
534 float
535 1202 kmeans_merge_centroids(
536 float *centroids,
537 float *norms_c,
538 uint32_t nlist,
539 Dimension dim,
540 DistanceMetric metric,
541 const KMeansReduce *reduce,
542 float *out_total_cost)
543 {
544 1202 const uint32_t nworkers = reduce->nworkers;
545 1202 uint32_t *cluster_vector_counts = reduce->cluster_vector_counts;
546
547 /* Sum costs */
548 1202 float total_cost = 0.0f;
549
2/2
✓ Branch 0 taken 3393 times.
✓ Branch 1 taken 1202 times.
 4595 for (uint32_t t = 0; t < nworkers; t++)
550 3393 total_cost += reduce->worker_costs[t];
551 1202 *out_total_cost = total_cost;
552
553 /* Merge per-worker accumulators */
554 1202 float *new_cents = centroids;
555 1202 memset(new_cents, 0, (size_t)nlist * dim * sizeof(float));
556 1202 memset(cluster_vector_counts, 0, (size_t)nlist * sizeof(uint32_t));
557
558
2/2
✓ Branch 0 taken 3393 times.
✓ Branch 1 taken 1202 times.
 4595 for (uint32_t t = 0; t < nworkers; t++)
559 {
560 3393 const float *vector_sums = reduce->worker_vector_sums[t];
561 3393 const uint32_t *vector_counts = reduce->worker_vector_counts[t];
562
563
2/2
✓ Branch 0 taken 29972 times.
✓ Branch 1 taken 3393 times.
 33365 for (uint32_t c = 0; c < nlist; c++)
564 {
565 29972 cluster_vector_counts[c] += vector_counts[c];
566 29972 float *dst = new_cents + (size_t)c * dim;
567 29972 const float *src = vector_sums + (size_t)c * dim;
568
2/2
✓ Branch 0 taken 2652152 times.
✓ Branch 1 taken 29972 times.
 2682124 for (uint32_t d = 0; d < dim; d++)
569 2652152 dst[d] += src[d];
570 }
571 }
572
573 /* Divide by cluster size to get mean */
574
2/2
✓ Branch 0 taken 10056 times.
✓ Branch 1 taken 1202 times.
 11258 for (uint32_t c = 0; c < nlist; c++)
575 {
576
2/2
✓ Branch 0 taken 9 times.
✓ Branch 1 taken 10047 times.
 10056 if (cluster_vector_counts[c] == 0)
577 9 continue;
578 10047 float inv = 1.0f / (float)cluster_vector_counts[c];
579 10047 float *cent = new_cents + (size_t)c * dim;
580
2/2
✓ Branch 0 taken 882871 times.
✓ Branch 1 taken 10047 times.
 892918 for (uint32_t d = 0; d < dim; d++)
581 882871 cent[d] *= inv;
582 }
583
584 /* Normalize for cosine metric */
585
2/2
✓ Branch 0 taken 87 times.
✓ Branch 1 taken 1115 times.
 1202 if (metric == DISTANCE_COSINE)
586 {
587
2/2
✓ Branch 0 taken 820 times.
✓ Branch 1 taken 87 times.
 907 for (uint32_t c = 0; c < nlist; c++)
588 {
589
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 820 times.
 820 if (cluster_vector_counts[c] == 0)
590 ✗ continue;
591 820 float *cent = new_cents + (size_t)c * dim;
592 820 float norm = vs_l2_norm(cent, dim);
593
1/2
✓ Branch 0 taken 820 times.
✗ Branch 1 not taken.
 820 if (norm > 1e-10f)
594 820 vec32_scale(cent, 1.0f / norm, cent, dim);
595 }
596 }
597
598 /* Convergence: max centroid shift (squared) */
599 1700 float shift_sq = kmeans_max_centroid_shift_between(
600 1202 centroids, reduce->prev_centroids, nlist, dim);
601
602 /* Precompute centroid norms for next iteration */
603
4/4
✓ Branch 0 taken 1145 times.
✓ Branch 1 taken 57 times.
✓ Branch 2 taken 630 times.
✓ Branch 3 taken 72 times.
 1202 if (norms_c != NULL && metric == DISTANCE_L2)
604 {
605
2/2
✓ Branch 0 taken 9110 times.
✓ Branch 1 taken 1073 times.
 10183 for (uint32_t j = 0; j < nlist; j++)
606 9110 norms_c[j] = vs_l2_norm_squared(centroids + (size_t)j * dim, dim);
607 }
608
609 1202 return shift_sq;
610 }
611
612 /*
613 * Allocate working state for one k-means run.
614 *
615 * Creates an arena context and allocates everything (including the
616 * struct itself) within it. kmeans_state_destroy() bulk-frees all
617 * memory by deleting the arena.
618 */
619 static KMeansState *
620 2683 kmeans_state_create(
621 const void *vectors,
622 const uint32_t *indices,
623 VecType vec_type,
624 uint32_t nvecs,
625 Dimension dim,
626 uint32_t nlist,
627 DistanceMetric metric)
628 {
629 2683 VsMemCtx ctx = vs_memctx_create(NULL, "kmeans_state");
630 2683 VsMemCtx old_ctx = vs_memctx_switch(ctx);
631
632 2683 KMeansState *st = vs_alloc0(sizeof(KMeansState));
633
634 2683 st->memctx = ctx;
635 2683 st->vectors = vectors;
636 2683 st->indices = indices;
637 2683 st->vec_type = vec_type;
638 2683 st->nvecs = nvecs;
639 2683 st->nlist = nlist;
640 2683 st->dim = dim;
641 2683 st->metric = metric;
642
643 2683 st->centroids = vs_alloc((size_t)nlist * dim * sizeof(float));
644 2683 st->assignments = vs_alloc(nvecs * sizeof(ClusterId));
645 2683 st->cluster_sizes = vs_alloc(nlist * sizeof(uint32_t));
646 2683 st->new_centroids = vs_alloc((size_t)nlist * dim * sizeof(float));
647
648 2683 uint32_t block = KMEANS_BLOCK_SIZE;
649
1/2
✓ Branch 0 taken 1166 times.
✗ Branch 1 not taken.
 2683 if (block > nvecs)
650 1166 block = nvecs;
651 2683 st->dist_block = vs_alloc((size_t)block * nlist * sizeof(float));
652
653 /*
654 * Allocate conversion buffer for non-f32 types, or for f32 with
655 * indices (Lloyd gather needs a contiguous block).
656 */
657
5/5
✓ Branch 0 taken 1517 times.
✓ Branch 1 taken 1154 times.
✓ Branch 2 taken 1276 times.
✓ Branch 3 taken 1181 times.
✓ Branch 4 taken 226 times.
 2683 if (vs_vec_element_size(vec_type) != sizeof(float) || indices != NULL)
658 2204 st->vec_block = vs_alloc((size_t)block * dim * sizeof(float));
659
660
2/2
✓ Branch 0 taken 2585 times.
✓ Branch 1 taken 98 times.
 2683 if (metric == DISTANCE_L2)
661 {
662 2585 st->norms_x = vs_alloc(nvecs * sizeof(float));
663 2585 st->norms_c = vs_alloc(nlist * sizeof(float));
664 }
665
666 2683 vs_memctx_switch(old_ctx);
667 2683 return st;
668 }
669
670 static void
671 2683 kmeans_state_destroy(KMeansState *st)
672 {
673
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 1166 times.
 2683 if (st == NULL)
674 ✗ return;
675 2683 VsMemCtx ctx = (VsMemCtx)st->memctx;
676 2683 vs_memctx_delete(ctx); /* st is now invalid */
677 }
678
679 /*
680 * Algorithm vtable for the k-means iteration loop.
681 *
682 * Each variant provides: create/destroy for per-algorithm state,
683 * assign for the assignment step, and optionally update_bounds
684 * for bound-accelerated algorithms (Hamerly, Elkan).
685 */
686 typedef struct
687 {
688 void *(*create)(const KMeansState *st);
689 void (*destroy)(void *algo_state);
690 void (*assign)(KMeansState *st, void *algo_state);
691 void (*update_bounds)(
692 KMeansState *st, void *algo_state, const float *old_cents);
693 } KMeansAlgoOps;
694
695 /* Lloyd wrappers (no per-algorithm state) */
696
697 static void
698 12 lloyd_plain_assign(KMeansState *st, void *state)
699 {
700 ✗ (void)state;
701 12 lloyd_assign(st, false);
702 12 }
703
704 static void
705 ✗ lloyd_cblas_assign(KMeansState *st, void *state)
706 {
707 ✗ (void)state;
708 ✗ lloyd_assign(st, true);
709 ✗ }
710
711 static const KMeansAlgoOps lloyd_ops = {
712 .assign = lloyd_plain_assign,
713 };
714
715 static const KMeansAlgoOps lloyd_cblas_ops = {
716 .assign = lloyd_cblas_assign,
717 };
718
719 /* Hamerly wrappers */
720
721 static void *
722 26 hamerly_wrap_create(const KMeansState *st)
723 {
724 26 return hamerly_create(st->nvecs, st->nlist, st->dim);
725 }
726
727 static void
728 126 hamerly_wrap_assign(KMeansState *st, void *state)
729 {
730 126 hamerly_assign(st, state);
731 126 }
732
733 static void
734 100 hamerly_wrap_bounds(KMeansState *st, void *state, const float *old_cents)
735 {
736 100 hamerly_update_bounds(st, state, old_cents);
737 100 }
738
739 static void
740 26 hamerly_wrap_destroy(void *state)
741 {
742 26 hamerly_destroy(state);
743 26 }
744
745 static const KMeansAlgoOps hamerly_ops = {
746 .create = hamerly_wrap_create,
747 .destroy = hamerly_wrap_destroy,
748 .assign = hamerly_wrap_assign,
749 .update_bounds = hamerly_wrap_bounds,
750 };
751
752 /* Elkan wrappers */
753
754 static void *
755 26 elkan_wrap_create(const KMeansState *st)
756 {
757 26 return elkan_create(st->nvecs, st->nlist, st->dim);
758 }
759
760 static void
761 126 elkan_wrap_assign(KMeansState *st, void *state)
762 {
763 126 elkan_assign(st, state);
764 126 }
765
766 static void
767 100 elkan_wrap_bounds(KMeansState *st, void *state, const float *old_cents)
768 {
769 100 elkan_update_bounds(st, state, old_cents);
770 100 }
771
772 static void
773 26 elkan_wrap_destroy(void *state)
774 {
775 26 elkan_destroy(state);
776 26 }
777
778 static const KMeansAlgoOps elkan_ops = {
779 .create = elkan_wrap_create,
780 .destroy = elkan_wrap_destroy,
781 .assign = elkan_wrap_assign,
782 .update_bounds = elkan_wrap_bounds,
783 };
784
785 /*
786 * Run one complete k-means attempt (init + iterate to convergence).
787 *
788 * always_inline — the specialized wrappers below pass a static const
789 * Vec32TypeOps from the header, so the compiler inlines through
790 * every vtable function pointer. VS_TARGET_CLONES on the wrappers
791 * generates AVX2/AVX-512 variants of the entire inlined body.
792 */
793 __attribute__((always_inline)) static inline void
794 1517 kmeans_run_one_impl(
795 KMeansState *st,
796 const KMeansOptions *opts,
797 uint64_t seed,
798 const KMeansAlgoOps *algo,
799 const Vec32TypeOps *ops)
800 {
801 /* Precompute input vector norms for L2 */
802 2683 if (st->metric == DISTANCE_L2)
803 1517 precompute_norms_x_impl(st, ops);
804
805 2683 size_t cent_bytes = (size_t)st->nlist * st->dim * sizeof(float);
806
4/6
✗ Branch 0 not taken.
✗ Branch 1 not taken.
✓ Branch 3 taken 8 times.
✓ Branch 4 taken 4 times.
✓ Branch 6 taken 44 times.
✓ Branch 7 taken 2627 times.
 2683 void *algo_state = algo->create ? algo->create(st) : NULL;
807
4/6
✗ Branch 0 not taken.
✗ Branch 1 not taken.
✓ Branch 3 taken 8 times.
✓ Branch 4 taken 4 times.
✓ Branch 6 taken 44 times.
✓ Branch 7 taken 2627 times.
 2683 float *old_cents = algo->update_bounds ? vs_alloc(cent_bytes) : NULL;
808
809
2/6
✗ Branch 0 not taken.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 12 times.
✗ Branch 4 not taken.
✓ Branch 5 taken 2671 times.
 2683 if (opts->initial_centroids != NULL)
810 {
811 ✗ memcpy(st->centroids,
812 ✗ opts->initial_centroids,
813 ✗ (size_t)st->nlist * st->dim * sizeof(float));
814 }
815 else
816 {
817 3034 kmeans_init_plusplus_impl(st, seed, ops);
818 }
819
820 /* Use fused iterate path for Lloyd with f32 vectors.
821 * Works with or without a thread pool — serial fallback
822 * calls work+reduce in a plain loop.
823 * Hamerly/Elkan use the generic loop below (needs bounds). */
824
3/6
✗ Branch 0 not taken.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 12 times.
✓ Branch 4 taken 1154 times.
✓ Branch 5 taken 1517 times.
 3837 bool use_iterate = st->vec_type == VS_VEC_F32 &&
825
2/6
✗ Branch 0 not taken.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✓ Branch 4 taken 1110 times.
✓ Branch 5 taken 1561 times.
 2671 algo->update_bounds == NULL;
826
827
3/6
✗ Branch 0 not taken.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 12 times.
✓ Branch 4 taken 1110 times.
✓ Branch 5 taken 44 times.
 2683 if (use_iterate)
828 {
829 2627 bool use_cblas = (algo == &lloyd_cblas_ops);
830 2627 lloyd_iterate(st, use_cblas, opts);
831 2627 kmeans_handle_empty_clusters(st);
832
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 2627 times.
 2627 if (algo->destroy)
833 ✗ algo->destroy(algo_state);
834
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 2627 times.
 2627 if (old_cents)
835 ✗ vs_free(old_cents);
836 1110 return;
837 }
838
839
2/6
✗ Branch 0 not taken.
✗ Branch 1 not taken.
✓ Branch 2 taken 72 times.
✗ Branch 3 not taken.
✓ Branch 4 taken 136 times.
✗ Branch 5 not taken.
 208 for (uint32_t iter = 0; iter < opts->max_iterations; iter++)
840 {
841 /* Save old centroids for convergence check */
842
0/6
✗ Branch 0 not taken.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
 208 memcpy(st->new_centroids, st->centroids, cent_bytes);
843
844 /* Save separately for bounds update (Hamerly/Elkan) */
845
3/6
✗ Branch 0 not taken.
✗ Branch 1 not taken.
✓ Branch 2 taken 64 times.
✓ Branch 3 taken 8 times.
✓ Branch 4 taken 136 times.
✗ Branch 5 not taken.
 208 if (old_cents)
846 200 memcpy(old_cents, st->centroids, cent_bytes);
847
848 208 algo->assign(st, algo_state);
849 ✗ kmeans_update_centroids_impl(st, ops);
850 208 kmeans_handle_empty_clusters(st);
851
852
3/6
✗ Branch 0 not taken.
✗ Branch 1 not taken.
✓ Branch 2 taken 64 times.
✓ Branch 3 taken 8 times.
✓ Branch 4 taken 136 times.
✗ Branch 5 not taken.
 208 if (algo->update_bounds)
853 200 algo->update_bounds(st, algo_state, old_cents);
854
855 208 float shift_sq = kmeans_max_centroid_shift_between(
856 208 st->centroids, st->new_centroids, st->nlist, st->dim);
857
858
3/6
✗ Branch 0 not taken.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✓ Branch 3 taken 72 times.
✓ Branch 4 taken 16 times.
✓ Branch 5 taken 120 times.
 208 if (opts->verbose)
859
1/9
✗ Branch 0 not taken.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✓ Branch 6 taken 16 times.
✗ Branch 7 not taken.
✗ Branch 11 not taken.
✗ Branch 12 not taken.
 16 vs_log(" iter %u: cost=%.4f, max_shift=%.6f\n",
860 iter,
861 st->total_cost,
862 sqrtf(shift_sq));
863
864 208 st->total_cost = 0.0f;
865
866 208 float tol_sq = opts->tolerance * opts->tolerance;
867
4/6
✗ Branch 0 not taken.
✗ Branch 1 not taken.
✓ Branch 2 taken 12 times.
✓ Branch 3 taken 60 times.
✓ Branch 4 taken 44 times.
✓ Branch 5 taken 92 times.
 208 if (shift_sq < tol_sq)
868 {
869 56 algo->assign(st, algo_state);
870
3/6
✗ Branch 0 not taken.
✗ Branch 1 not taken.
✓ Branch 2 taken 8 times.
✓ Branch 3 taken 4 times.
✓ Branch 4 taken 44 times.
✗ Branch 5 not taken.
 56 if (algo->destroy)
871 52 algo->destroy(algo_state);
872
3/6
✗ Branch 0 not taken.
✗ Branch 1 not taken.
✓ Branch 2 taken 8 times.
✓ Branch 3 taken 4 times.
✓ Branch 4 taken 44 times.
✗ Branch 5 not taken.
 56 if (old_cents)
873 52 vs_free(old_cents);
874 56 return;
875 }
876 }
877
878 /* Final assignment after max iterations */
879 ✗ algo->assign(st, algo_state);
880 ✗ if (algo->destroy)
881 ✗ algo->destroy(algo_state);
882 ✗ if (old_cents)
883 ✗ vs_free(old_cents);
884 }
885
886 /* Specialized wrappers — VS_TARGET_CLONES generates SIMD variants */
887
888 VS_TARGET_CLONES static void
889
2/2
✓ Branch 0 taken 1108 times.
✓ Branch 1 taken 46 times.
 2671 kmeans_run_one_f32(
890 KMeansState *st,
891 const KMeansOptions *opts,
892 uint64_t seed,
893 const KMeansAlgoOps *algo)
894 {
895
2/2
✓ Branch 0 taken 1465 times.
✓ Branch 1 taken 52 times.
 1517 kmeans_run_one_impl(st, opts, seed, algo, &vs_f32_type_ops);
896 2671 }
897
898 VS_TARGET_CLONES static void
899
1/2
✓ Branch 0 taken 12 times.
✗ Branch 1 not taken.
 12 kmeans_run_one_f16(
900 KMeansState *st,
901 const KMeansOptions *opts,
902 uint64_t seed,
903 const KMeansAlgoOps *algo)
904 {
905 ✗ kmeans_run_one_impl(st, opts, seed, algo, &vs_f16_type_ops);
906 12 }
907
908 #if defined(VS_F16C_SUPPORT) && !defined(VS_SIMD_NONE)
909 VS_TARGET_F16C_AVX2 static void
910 ✗ kmeans_run_one_f16c(
911 KMeansState *st,
912 const KMeansOptions *opts,
913 uint64_t seed,
914 const KMeansAlgoOps *algo)
915 {
916 ✗ kmeans_run_one_impl(st, opts, seed, algo, &vs_f16c_type_ops);
917 ✗ }
918 #endif
919
920 /* Single dispatch point — selects the inline vtable once */
921 static void
922 2683 kmeans_run_one(
923 KMeansState *st,
924 const KMeansOptions *opts,
925 uint64_t seed,
926 const KMeansAlgoOps *algo)
927 {
928
2/3
✓ Branch 0 taken 2671 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 12 times.
 2683 switch (st->vec_type)
929 {
930 2671 case VS_VEC_F32:
931 2671 kmeans_run_one_f32(st, opts, seed, algo);
932 2671 break;
933 #if defined(VS_F16C_SUPPORT) && !defined(VS_SIMD_NONE)
934 ✗ case VS_VEC_F16C:
935 ✗ kmeans_run_one_f16c(st, opts, seed, algo);
936 ✗ break;
937 #endif
938 12 default:
939 12 kmeans_run_one_f16(st, opts, seed, algo);
940 12 break;
941 }
942 2683 }
943
944 /*
945 * Public API
946 */
947
948 KMeansResult *
949 2664 vs_kmeans(
950 const void *vectors,
951 const uint32_t *indices,
952 VecType vec_type,
953 uint32_t nvecs,
954 Dimension dim,
955 uint32_t nlist,
956 DistanceMetric metric,
957 const KMeansOptions *options)
958 {
959
8/8
✓ Branch 0 taken 2662 times.
✓ Branch 1 taken 2 times.
✓ Branch 2 taken 2660 times.
✓ Branch 3 taken 2 times.
✓ Branch 4 taken 1150 times.
✓ Branch 5 taken 2 times.
✓ Branch 6 taken 2 times.
✓ Branch 7 taken 1148 times.
 2664 if (vectors == NULL || nvecs == 0 || dim == 0 || nlist == 0)
960 8 return NULL;
961
962 /* Clamp nlist to nvecs */
963
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 1148 times.
 2656 if (nlist > nvecs)
964 ✗ nlist = nvecs;
965
966 2656 KMeansOptions opts = VS_KMEANS_OPTIONS_DEFAULT;
967
2/2
✓ Branch 0 taken 2654 times.
✓ Branch 1 taken 2 times.
 2656 if (options != NULL)
968 2654 opts = *options;
969
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 2656 times.
 2656 if (opts.nredo == 0)
970 ✗ opts.nredo = 1;
971
972 2656 KMeansResult *best = NULL;
973 2656 float best_cost = FLT_MAX;
974
975 /* Resolve AUTO to a concrete algorithm */
976 2656 KMeansAlgorithm algo = opts.algorithm;
977
2/2
✓ Branch 0 taken 566 times.
✓ Branch 1 taken 582 times.
 2656 if (algo == KMEANS_ALGO_AUTO)
978 {
979 #ifdef VS_HAVE_CBLAS
980 algo = g_use_cblas ? KMEANS_ALGO_CBLAS : KMEANS_ALGO_LLOYD;
981 #else
982 566 algo = KMEANS_ALGO_LLOYD;
983 #endif
984 }
985
986 /* Hamerly/Elkan only support L2 — fall back to Lloyd for others */
987
6/6
✓ Branch 0 taken 1126 times.
✓ Branch 1 taken 1530 times.
✓ Branch 2 taken 22 times.
✓ Branch 3 taken 1104 times.
✓ Branch 4 taken 4 times.
✓ Branch 5 taken 40 times.
 2656 if ((algo == KMEANS_ALGO_HAMERLY || algo == KMEANS_ALGO_ELKAN) &&
988 metric != DISTANCE_L2)
989 4 algo = KMEANS_ALGO_LLOYD;
990
991 /* Select algorithm vtable */
992 1508 const KMeansAlgoOps *algo_ops;
993
4/4
✓ Branch 0 taken 20 times.
✓ Branch 1 taken 20 times.
✓ Branch 2 taken 1508 times.
✓ Branch 3 taken 1108 times.
 2656 switch (algo)
994 {
995 20 case KMEANS_ALGO_HAMERLY:
996 20 algo_ops = &hamerly_ops;
997 20 break;
998 20 case KMEANS_ALGO_ELKAN:
999 20 algo_ops = &elkan_ops;
1000 20 break;
1001 ✗ case KMEANS_ALGO_CBLAS:
1002 ✗ algo_ops = &lloyd_cblas_ops;
1003 ✗ break;
1004 2616 case KMEANS_ALGO_LLOYD:
1005 default:
1006 2616 algo_ops = &lloyd_ops;
1007 2616 break;
1008 }
1009
1010 2656 size_t cent_sz = (size_t)nlist * dim * sizeof(float);
1011 2656 size_t assign_sz = nvecs * sizeof(ClusterId);
1012 2656 size_t clsize_sz = nlist * sizeof(uint32_t);
1013
1014
2/2
✓ Branch 0 taken 2683 times.
✓ Branch 1 taken 2656 times.
 5339 for (uint32_t redo = 0; redo < opts.nredo; redo++)
1015 {
1016 2683 KMeansState *st = kmeans_state_create(
1017 vectors, indices, vec_type, nvecs, dim, nlist, metric);
1018
1019 2683 uint64_t seed = opts.seed + redo;
1020
1021 /* Run in arena context so per-iteration temps land there */
1022 2683 VsMemCtx run_ctx = vs_memctx_switch((VsMemCtx)st->memctx);
1023 2683 kmeans_run_one(st, &opts, seed, algo_ops);
1024
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 1517 times.
 2683 vs_memctx_switch(run_ctx);
1025
1026
3/4
✓ Branch 0 taken 12 times.
✓ Branch 1 taken 2671 times.
✓ Branch 2 taken 12 times.
✗ Branch 3 not taken.
 2683 if (opts.verbose && opts.nredo > 1)
1027
3/5
✓ Branch 0 taken 12 times.
✗ Branch 1 not taken.
✓ Branch 2 taken 10 times.
✓ Branch 3 taken 2 times.
✗ Branch 4 not taken.
 12 vs_log("redo %u/%u: cost=%.4f%s\n",
1028 redo + 1,
1029 opts.nredo,
1030 st->total_cost,
1031 st->total_cost < best_cost ? " (best)" : "");
1032
1033
2/2
✓ Branch 0 taken 2672 times.
✓ Branch 1 taken 11 times.
 2683 if (st->total_cost < best_cost)
1034 {
1035 /* Save as best — copy out of arena into caller ctx */
1036
2/2
✓ Branch 0 taken 16 times.
✓ Branch 1 taken 2656 times.
 2672 if (best != NULL)
1037 16 vs_kmeans_result_destroy(best);
1038
1039 2672 best = vs_alloc0(sizeof(KMeansResult));
1040 2672 best->nlist = nlist;
1041 2672 best->dim = dim;
1042 2672 best->total_cost = st->total_cost;
1043
1044 2672 best->centroids = vs_alloc(cent_sz);
1045 2672 memcpy(best->centroids, st->centroids, cent_sz);
1046
1047 2672 best->assignments = vs_alloc(assign_sz);
1048 2672 memcpy(best->assignments, st->assignments, assign_sz);
1049
1050 2672 best->cluster_sizes = vs_alloc(clsize_sz);
1051 2672 memcpy(best->cluster_sizes, st->cluster_sizes, clsize_sz);
1052
1053 2672 best_cost = st->total_cost;
1054 }
1055
1056 2683 kmeans_state_destroy(st);
1057 }
1058
1059 1148 return best;
1060 }
1061
1062 void
1063 2672 vs_kmeans_result_destroy(KMeansResult *result)
1064 {
1065
2/2
✓ Branch 0 taken 1512 times.
✓ Branch 1 taken 1160 times.
 2672 if (result == NULL)
1066 ✗ return;
1067 2672 vs_free(result->centroids);
1068 2672 vs_free(result->assignments);
1069 2672 vs_free(result->cluster_sizes);
1070 2672 vs_free(result);
1071 }
1072