GCC Code Coverage Report

Directory: src/
File: src/algo/kmeans_elkan.c
Date: 2026-09-30 11:11:31
Exec Total Coverage
Lines: 250 259 96.5%
Functions: 11 11 100.0%
Branches: 85 114 74.6%
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_elkan.c - Elkan's accelerated k-means assignment
6 *
7 * Elkan's algorithm maintains K lower bounds per vector (one per
8 * centroid) plus a single upper bound on the distance to the assigned
9 * centroid. The triangle inequality enables three pruning rules:
10 *
11 * 1. If upper[i] <= s[assigned[i]], skip entirely.
12 * (s[j] = min half-distance from centroid j to any other centroid)
13 *
14 * 2. If upper[i] <= lower[i][k], skip centroid k.
15 * (lower bound is per-centroid, so tighter than Hamerly's global)
16 *
17 * 3. If upper[i] <= halfcdist[assigned[i]][k], skip centroid k.
18 * (centroid-centroid distance provides a tighter filter)
19 *
20 * After centroid update, bounds are adjusted per-centroid:
21 * lower[i][k] -= movement[k] (only loosened by that centroid)
22 * upper[i] += movement[assigned[i]]
23 *
24 * This is the key advantage over Hamerly: each lower bound is only
25 * affected by its own centroid's movement, not by the global maximum.
26 *
27 * Distance computation uses dot product decomposition:
28 * d(x, c) = sqrt(||x||^2 + ||c||^2 - 2 * dot(x, c))
29 */
30
31 #include "vs_config.h"
32
33 #include <math.h>
34 #include <string.h>
35
36 #include "algo/kmeans_elkan.h"
37 #include "algo/vecops.h"
38 #include "core/memory.h"
39
40 struct ElkanState
41 {
42 float *lower; /* [nvecs * nlist] per-centroid lower bounds */
43 float *upper; /* [nvecs] upper bound on assigned distance */
44 float *halfcdist; /* [nlist * nlist] half centroid-centroid dist */
45 float *s; /* [nlist] min half-dist to any other centroid */
46 float *cdist; /* [nlist] centroid movement distances */
47 uint32_t nvecs;
48 uint32_t nlist;
49 uint32_t dim;
50 bool bounds_valid;
51 };
52
53 ElkanState *
54 26 elkan_create(uint32_t nvecs, uint32_t nlist, uint32_t dim)
55 {
56 26 ElkanState *es = vs_alloc0(sizeof(ElkanState));
57 26 es->nvecs = nvecs;
58 26 es->nlist = nlist;
59 26 es->dim = dim;
60 26 es->lower = vs_alloc0((size_t)nvecs * nlist * sizeof(float));
61 26 es->upper = vs_alloc(nvecs * sizeof(float));
62 26 es->halfcdist = vs_alloc((size_t)nlist * nlist * sizeof(float));
63 26 es->s = vs_alloc(nlist * sizeof(float));
64 26 es->cdist = vs_alloc(nlist * sizeof(float));
65 26 return es;
66 }
67
68 void
69 26 elkan_destroy(ElkanState *es)
70 {
71
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 26 times.
 26 if (es == NULL)
72 ✗ return;
73 26 vs_free(es->lower);
74 26 vs_free(es->upper);
75 26 vs_free(es->halfcdist);
76 26 vs_free(es->s);
77 26 vs_free(es->cdist);
78 26 vs_free(es);
79 }
80
81 /*
82 * Compute dot product between two float32 vectors.
83 * Used for centroid-centroid distances (always float32).
84 */
85 VS_TARGET_CLONES static float
86 948 dot_product(const float *a, const float *b, uint32_t dim)
87 {
88 948 float dot = 0.0f;
89
2/2
✓ Branch 0 taken 6372 times.
✓ Branch 1 taken 948 times.
 7320 for (uint32_t d = 0; d < dim; d++)
90 6372 dot += a[d] * b[d];
91 948 return dot;
92 }
93
94 /*
95 * Compute squared L2 distance using dot product decomposition.
96 * Requires precomputed norms: norm_a = ||a||^2, norm_b = ||b||^2.
97 */
98 static inline float
99 30180 l2_sq_from_dot(float norm_a, float norm_b, float dot)
100 {
101 30180 float d_sq = norm_a + norm_b - 2.0f * dot;
102
2/2
✓ Branch 0 taken 30084 times.
✓ Branch 1 taken 96 times.
 30180 return d_sq > 0.0f ? d_sq : 0.0f;
103 }
104
105 /*
106 * Compute Euclidean distance using dot product decomposition.
107 */
108 static inline float
109 948 l2_dist_from_dot(float norm_a, float norm_b, float dot)
110 {
111 948 return sqrtf(l2_sq_from_dot(norm_a, norm_b, dot));
112 }
113
114 /*
115 * Compute centroid-centroid half-distances and min separations.
116 *
117 * halfcdist[j][k] = 0.5 * d(c_j, c_k)
118 * s[j] = min_k(halfcdist[j][k]) for k != j
119 */
120 static void
121 126 compute_centroid_dists(ElkanState *es, const KMeansState *st)
122 {
123 126 uint32_t nlist = st->nlist;
124 126 uint32_t dim = st->dim;
125 126 const float *cents = st->centroids;
126 126 float *hcd = es->halfcdist;
127
128 /* Initialize s to infinity */
129
2/2
✓ Branch 0 taken 528 times.
✓ Branch 1 taken 126 times.
 654 for (uint32_t j = 0; j < nlist; j++)
130 528 es->s[j] = FLT_MAX;
131
132
2/2
✓ Branch 0 taken 528 times.
✓ Branch 1 taken 126 times.
 654 for (uint32_t j = 0; j < nlist; j++)
133 {
134 528 hcd[(size_t)j * nlist + j] = 0.0f;
135 528 const float *cj = cents + (size_t)j * dim;
136 528 float norm_j = st->norms_c[j];
137
138
2/2
✓ Branch 0 taken 948 times.
✓ Branch 1 taken 528 times.
 1476 for (uint32_t k = j + 1; k < nlist; k++)
139 {
140 948 float dot = dot_product(cj, cents + (size_t)k * dim, dim);
141
0/2
✗ Branch 0 not taken.
✗ Branch 1 not taken.
 948 float d = l2_dist_from_dot(norm_j, st->norms_c[k], dot);
142 948 float hd = 0.5f * d;
143
144 948 hcd[(size_t)j * nlist + k] = hd;
145 948 hcd[(size_t)k * nlist + j] = hd;
146
147
2/2
✓ Branch 0 taken 430 times.
✓ Branch 1 taken 518 times.
 948 if (hd < es->s[j])
148 430 es->s[j] = hd;
149
2/2
✓ Branch 0 taken 568 times.
✓ Branch 1 taken 380 times.
 948 if (hd < es->s[k])
150 568 es->s[k] = hd;
151 }
152 }
153 126 }
154
155 /*
156 * Precompute centroid norms (||c||^2) into st->norms_c.
157 */
158 static void
159 126 precompute_norms_c(KMeansState *st)
160 {
161
2/2
✓ Branch 0 taken 528 times.
✓ Branch 1 taken 126 times.
 654 for (uint32_t j = 0; j < st->nlist; j++)
162 528 st->norms_c[j] = vs_l2_norm_squared(
163 528 st->centroids + (size_t)j * st->dim, st->dim);
164 126 }
165
166 /*
167 * Full initial assignment — always_inline, specialized by ops vtable.
168 */
169 __attribute__((always_inline)) static inline void
170 elkan_initial_assign_impl(
171 KMeansState *st, ElkanState *es, const Vec32TypeOps *ops)
172 {
173 22 uint32_t nvecs = st->nvecs;
174 22 uint32_t nlist = st->nlist;
175 22 Dimension dim = st->dim;
176 22 size_t esz = ops->element_size;
177 22 const float *cents = st->centroids;
178
179
2/2
✓ Branch 0 taken 3300 times.
✓ Branch 1 taken 22 times.
 3322 for (uint32_t i = 0; i < nvecs; i++)
180 {
181
0/2
✗ Branch 0 not taken.
✗ Branch 1 not taken.
 3300 const void *v = km_get_vector(st, i, esz);
182 3300 float nx = st->norms_x[i];
183 3300 float *lb = es->lower + (size_t)i * nlist;
184
185 3300 float best_d_sq = FLT_MAX;
186 3300 float best_d = FLT_MAX;
187 3300 uint32_t best_j = 0;
188
189
2/2
✓ Branch 0 taken 12300 times.
✓ Branch 1 taken 3300 times.
 15600 for (uint32_t j = 0; j < nlist; j++)
190 {
191 12300 float dp = ops->dot_product(v, cents + (size_t)j * dim, dim);
192
0/2
✗ Branch 0 not taken.
✗ Branch 1 not taken.
 12300 float d_sq = l2_sq_from_dot(nx, st->norms_c[j], dp);
193 12300 float d = sqrtf(d_sq);
194 12300 lb[j] = d;
195
196
2/2
✓ Branch 0 taken 6496 times.
✓ Branch 1 taken 5804 times.
 12300 if (d_sq < best_d_sq)
197 {
198 6496 best_d_sq = d_sq;
199 6496 best_d = d;
200 6496 best_j = j;
201 }
202 }
203
204 3300 st->assignments[i] = best_j;
205 3300 es->upper[i] = best_d;
206 3300 st->total_cost += best_d_sq;
207 }
208 22 }
209
210 /*
211 * Main Elkan assignment loop — always_inline, specialized by ops vtable.
212 *
213 * Centroid-centroid distances use the separate dot_product() function
214 * (always float32 × float32). Only vector-centroid uses the typed ops.
215 */
216 __attribute__((always_inline)) static inline void
217 ✗ elkan_assign_impl(KMeansState *st, ElkanState *es, const Vec32TypeOps *ops)
218 {
219 90 uint32_t nvecs = st->nvecs;
220 90 uint32_t nlist = st->nlist;
221 90 Dimension dim = st->dim;
222 90 size_t esz = ops->element_size;
223 90 const float *cents = st->centroids;
224
225 90 if (!es->bounds_valid)
226 {
227 ✗ elkan_initial_assign_impl(st, es, ops);
228 22 compute_centroid_dists(es, st);
229 22 es->bounds_valid = true;
230 22 return;
231 }
232
233 /* Compute centroid-centroid distances for this iteration */
234 68 compute_centroid_dists(es, st);
235
236
2/3
✓ Branch 0 taken 11400 times.
✓ Branch 1 taken 68 times.
✗ Branch 2 not taken.
 11468 for (uint32_t i = 0; i < nvecs; i++)
237 {
238 11400 float ub = es->upper[i];
239 11400 uint32_t asgn = st->assignments[i];
240 11400 float *lb = es->lower + (size_t)i * nlist;
241
242 /* Step 1: skip if upper bound <= min half-centroid-distance */
243
2/2
✓ Branch 0 taken 5844 times.
✓ Branch 1 taken 5556 times.
 11400 if (ub <= es->s[asgn])
244 {
245 5844 st->total_cost += ub * ub;
246 5844 continue;
247 }
248
249
0/2
✗ Branch 0 not taken.
✗ Branch 1 not taken.
 5556 const void *v = km_get_vector(st, i, esz);
250 5556 float nx = st->norms_x[i];
251 5556 float ub_sq = ub * ub;
252 5556 bool recompute = true;
253
254
2/2
✓ Branch 0 taken 27780 times.
✓ Branch 1 taken 5556 times.
 33336 for (uint32_t k = 0; k < nlist; k++)
255 {
256
2/2
✓ Branch 0 taken 5434 times.
✓ Branch 1 taken 22346 times.
 27780 if (k == asgn)
257 5434 continue;
258
259 /* Step 3: skip centroid via lower bound or half-dist */
260
2/2
✓ Branch 0 taken 17170 times.
✓ Branch 1 taken 5176 times.
 22346 if (ub <= lb[k])
261 17170 continue;
262
263 5176 float *hcd_row = es->halfcdist + (size_t)asgn * nlist;
264
2/2
✓ Branch 0 taken 24 times.
✓ Branch 1 taken 5152 times.
 5176 if (ub <= hcd_row[k])
265 24 continue;
266
267 /* Step 3a: tighten upper bound if not yet done */
268
2/2
✓ Branch 0 taken 3662 times.
✓ Branch 1 taken 1490 times.
 5152 if (recompute)
269 {
270 ✗ float dp_a =
271 3662 ops->dot_product(v, cents + (size_t)asgn * dim, dim);
272
0/2
✗ Branch 0 not taken.
✗ Branch 1 not taken.
 3662 float d_sq_a = l2_sq_from_dot(nx, st->norms_c[asgn], dp_a);
273 3662 float d_a = sqrtf(d_sq_a);
274 3662 lb[asgn] = d_a;
275 3662 ub = d_a;
276 3662 ub_sq = d_sq_a;
277 3662 es->upper[i] = d_a;
278 3662 recompute = false;
279
280 /* Re-check with tightened upper bound */
281
2/2
✓ Branch 0 taken 1786 times.
✓ Branch 1 taken 1876 times.
 3662 if (ub <= lb[k])
282 1786 continue;
283
2/2
✓ Branch 0 taken 12 times.
✓ Branch 1 taken 1864 times.
 1876 if (ub <= hcd_row[k])
284 12 continue;
285 }
286
287 /* Step 3b: compute actual distance to centroid k */
288 3354 float dp_k = ops->dot_product(v, cents + (size_t)k * dim, dim);
289
0/2
✗ Branch 0 not taken.
✗ Branch 1 not taken.
 3354 float d_sq_k = l2_sq_from_dot(nx, st->norms_c[k], dp_k);
290 3354 float d_k = sqrtf(d_sq_k);
291 3354 lb[k] = d_k;
292
293 /* Compare squared distances (matches Lloyd/Hamerly) */
294
2/2
✓ Branch 0 taken 240 times.
✓ Branch 1 taken 3114 times.
 3354 if (d_sq_k < ub_sq)
295 {
296 240 asgn = k;
297 240 ub = d_k;
298 240 ub_sq = d_sq_k;
299 240 es->upper[i] = d_k;
300 240 st->assignments[i] = k;
301 }
302 }
303
304 5556 st->total_cost += ub_sq;
305 }
306 }
307
308 /*
309 * Preconvert f16 variants: convert each f16 vector to f32 once before
310 * the centroid loop, then use the f32 dot product kernel. Saves
311 * O(K*dim) f16→f32 conversions per vector (one conversion vs K).
312 * Vectors pruned by bounds are never converted.
313 */
314 __attribute__((always_inline)) static inline void
315 elkan_initial_assign_preconvert_impl(
316 KMeansState *st, ElkanState *es, size_t esz)
317 {
318 4 uint32_t nvecs = st->nvecs;
319 4 uint32_t nlist = st->nlist;
320 4 Dimension dim = st->dim;
321 4 const float *cents = st->centroids;
322 4 float *buf = st->vec_block;
323
324 4 const Vec32TypeOps *f32ops = &vs_f32_type_ops;
325
326
2/2
✓ Branch 0 taken 700 times.
✓ Branch 1 taken 4 times.
 704 for (uint32_t i = 0; i < nvecs; i++)
327 {
328
0/2
✗ Branch 0 not taken.
✗ Branch 1 not taken.
 700 const void *raw = km_get_vector(st, i, esz);
329 700 vs_half_to_float_array((const half *)raw, buf, dim);
330 700 float nx = st->norms_x[i];
331 700 float *lb = es->lower + (size_t)i * nlist;
332
333 700 float best_d_sq = FLT_MAX;
334 700 float best_d = FLT_MAX;
335 700 uint32_t best_j = 0;
336
337
2/2
✓ Branch 0 taken 2900 times.
✓ Branch 1 taken 700 times.
 3600 for (uint32_t j = 0; j < nlist; j++)
338 {
339 2900 float dp = f32ops->dot_product(buf, cents + (size_t)j * dim, dim);
340
0/2
✗ Branch 0 not taken.
✗ Branch 1 not taken.
 2900 float d_sq = l2_sq_from_dot(nx, st->norms_c[j], dp);
341 2900 float d = sqrtf(d_sq);
342 2900 lb[j] = d;
343
344
2/2
✓ Branch 0 taken 1460 times.
✓ Branch 1 taken 1440 times.
 2900 if (d_sq < best_d_sq)
345 {
346 1460 best_d_sq = d_sq;
347 1460 best_d = d;
348 1460 best_j = j;
349 }
350 }
351
352 700 st->assignments[i] = best_j;
353 700 es->upper[i] = best_d;
354 700 st->total_cost += best_d_sq;
355 }
356 4 }
357
358 __attribute__((always_inline)) static inline void
359 ✗ elkan_assign_preconvert_impl(KMeansState *st, ElkanState *es, size_t esz)
360 {
361 36 uint32_t nvecs = st->nvecs;
362 36 uint32_t nlist = st->nlist;
363 36 Dimension dim = st->dim;
364 36 const float *cents = st->centroids;
365 36 float *buf = st->vec_block;
366
367 36 const Vec32TypeOps *f32ops = &vs_f32_type_ops;
368
369 36 if (!es->bounds_valid)
370 {
371 ✗ elkan_initial_assign_preconvert_impl(st, es, esz);
372 4 compute_centroid_dists(es, st);
373 4 es->bounds_valid = true;
374 4 return;
375 }
376
377 32 compute_centroid_dists(es, st);
378
379
2/3
✓ Branch 0 taken 6200 times.
✓ Branch 1 taken 32 times.
✗ Branch 2 not taken.
 6232 for (uint32_t i = 0; i < nvecs; i++)
380 {
381 6200 float ub = es->upper[i];
382 6200 uint32_t asgn = st->assignments[i];
383 6200 float *lb = es->lower + (size_t)i * nlist;
384
385
2/2
✓ Branch 0 taken 644 times.
✓ Branch 1 taken 5556 times.
 6200 if (ub <= es->s[asgn])
386 {
387 644 st->total_cost += ub * ub;
388 644 continue;
389 }
390
391 /* Convert once — only for vectors that pass the s-bound */
392
0/2
✗ Branch 0 not taken.
✗ Branch 1 not taken.
 5556 const void *raw = km_get_vector(st, i, esz);
393 5556 vs_half_to_float_array((const half *)raw, buf, dim);
394
395 5556 float nx = st->norms_x[i];
396 5556 float ub_sq = ub * ub;
397 5556 bool recompute = true;
398
399
2/2
✓ Branch 0 taken 27780 times.
✓ Branch 1 taken 5556 times.
 33336 for (uint32_t k = 0; k < nlist; k++)
400 {
401
2/2
✓ Branch 0 taken 5434 times.
✓ Branch 1 taken 22346 times.
 27780 if (k == asgn)
402 5434 continue;
403
404
2/2
✓ Branch 0 taken 17170 times.
✓ Branch 1 taken 5176 times.
 22346 if (ub <= lb[k])
405 17170 continue;
406
407 5176 float *hcd_row = es->halfcdist + (size_t)asgn * nlist;
408
2/2
✓ Branch 0 taken 24 times.
✓ Branch 1 taken 5152 times.
 5176 if (ub <= hcd_row[k])
409 24 continue;
410
411
2/2
✓ Branch 0 taken 3662 times.
✓ Branch 1 taken 1490 times.
 5152 if (recompute)
412 {
413 3662 float dp_a = f32ops->dot_product(
414 3662 buf, cents + (size_t)asgn * dim, dim);
415
0/2
✗ Branch 0 not taken.
✗ Branch 1 not taken.
 3662 float d_sq_a = l2_sq_from_dot(nx, st->norms_c[asgn], dp_a);
416 3662 float d_a = sqrtf(d_sq_a);
417 3662 lb[asgn] = d_a;
418 3662 ub = d_a;
419 3662 ub_sq = d_sq_a;
420 3662 es->upper[i] = d_a;
421 3662 recompute = false;
422
423
2/2
✓ Branch 0 taken 1786 times.
✓ Branch 1 taken 1876 times.
 3662 if (ub <= lb[k])
424 1786 continue;
425
2/2
✓ Branch 0 taken 12 times.
✓ Branch 1 taken 1864 times.
 1876 if (ub <= hcd_row[k])
426 12 continue;
427 }
428
429 ✗ float dp_k =
430 3354 f32ops->dot_product(buf, cents + (size_t)k * dim, dim);
431
0/2
✗ Branch 0 not taken.
✗ Branch 1 not taken.
 3354 float d_sq_k = l2_sq_from_dot(nx, st->norms_c[k], dp_k);
432 3354 float d_k = sqrtf(d_sq_k);
433 3354 lb[k] = d_k;
434
435
2/2
✓ Branch 0 taken 240 times.
✓ Branch 1 taken 3114 times.
 3354 if (d_sq_k < ub_sq)
436 {
437 240 asgn = k;
438 240 ub = d_k;
439 240 ub_sq = d_sq_k;
440 240 es->upper[i] = d_k;
441 240 st->assignments[i] = k;
442 }
443 }
444
445 5556 st->total_cost += ub_sq;
446 }
447 }
448
449 /* Specialized wrappers — VS_TARGET_CLONES generates SIMD variants */
450 VS_TARGET_CLONES static void
451
2/2
✓ Branch 0 taken 22 times.
✓ Branch 1 taken 68 times.
 90 elkan_assign_f32(KMeansState *st, ElkanState *es)
452 {
453 ✗ elkan_assign_impl(st, es, &vs_f32_type_ops);
454 90 }
455
456 VS_TARGET_CLONES static void
457
2/2
✓ Branch 0 taken 4 times.
✓ Branch 1 taken 32 times.
 36 elkan_assign_f16(KMeansState *st, ElkanState *es)
458 {
459 ✗ elkan_assign_preconvert_impl(st, es, sizeof(half));
460 36 }
461
462 /* Public entry — dispatches once based on type */
463 void
464 126 elkan_assign(KMeansState *st, ElkanState *es)
465 {
466 126 st->total_cost = 0.0f;
467 126 precompute_norms_c(st);
468
469
2/2
✓ Branch 0 taken 90 times.
✓ Branch 1 taken 36 times.
 126 switch (st->vec_type)
470 {
471 90 case VS_VEC_F32:
472 90 elkan_assign_f32(st, es);
473 90 break;
474 36 default:
475 36 elkan_assign_f16(st, es);
476 36 break;
477 }
478 126 }
479
480 void
481 100 elkan_update_bounds(
482 KMeansState *st, ElkanState *es, const float *old_centroids)
483 {
484 100 uint32_t nlist = st->nlist;
485 100 uint32_t dim = st->dim;
486 100 uint32_t nvecs = st->nvecs;
487
488 /* Compute per-centroid movement distances */
489
2/2
✓ Branch 0 taken 432 times.
✓ Branch 1 taken 100 times.
 532 for (uint32_t j = 0; j < nlist; j++)
490 {
491 432 float d_sq = vs_l2_distance_squared(
492 432 st->centroids + (size_t)j * dim,
493 432 old_centroids + (size_t)j * dim,
494 dim);
495 432 es->cdist[j] = sqrtf(d_sq);
496 }
497
498 /* Step 5: update lower bounds (loosen per-centroid) */
499
2/2
✓ Branch 0 taken 17600 times.
✓ Branch 1 taken 100 times.
 17700 for (uint32_t i = 0; i < nvecs; i++)
500 {
501 17600 float *lb = es->lower + (size_t)i * nlist;
502
2/2
✓ Branch 0 taken 78400 times.
✓ Branch 1 taken 17600 times.
 96000 for (uint32_t k = 0; k < nlist; k++)
503 {
504 78400 lb[k] -= es->cdist[k];
505
2/2
✓ Branch 0 taken 824 times.
✓ Branch 1 taken 77576 times.
 78400 if (lb[k] < 0.0f)
506 824 lb[k] = 0.0f;
507 }
508 }
509
510 /* Step 6: update upper bounds (loosen by assigned movement) */
511
2/2
✓ Branch 0 taken 17600 times.
✓ Branch 1 taken 100 times.
 17700 for (uint32_t i = 0; i < nvecs; i++)
512 17600 es->upper[i] += es->cdist[st->assignments[i]];
513 100 }
514