GCC Code Coverage Report

Directory: src/
File: src/algo/kmeans_hamerly.c
Date: 2026-09-30 11:11:31
Exec Total Coverage
Lines: 155 175 88.6%
Functions: 8 8 100.0%
Branches: 64 90 71.1%
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_hamerly.c - Hamerly's accelerated k-means assignment
6 *
7 * Hamerly's algorithm maintains two bounds per vector:
8 * upper[i] = upper bound on d(x_i, c_assigned)
9 * lower[i] = lower bound on d(x_i, second_nearest_centroid)
10 *
11 * After centroid update, bounds are loosened:
12 * upper[i] += delta[assigned[i]] (assigned centroid may have moved away)
13 * lower[i] -= delta_max (closest other centroid may have moved
14 * closer)
15 *
16 * If upper[i] <= lower[i], the assignment cannot change and the vector
17 * is skipped entirely. This is the key optimization — in later iterations,
18 * 90%+ of vectors are skipped.
19 *
20 * For vectors that can't be skipped, we first tighten the upper bound
21 * by computing the exact distance to the assigned centroid. If still
22 * upper > lower, we do a full search over all centroids.
23 *
24 * Distance computation uses dot product decomposition:
25 * d^2(x, c) = ||x||^2 + ||c||^2 - 2 * dot(x, c)
26 * This needs only 1 FMA per dimension (vs 2 for direct L2).
27 */
28
29 #include "vs_config.h"
30
31 #include <math.h>
32 #include <string.h>
33
34 #include "algo/kmeans_hamerly.h"
35 #include "algo/vecops.h"
36 #include "core/memory.h"
37
38 struct HamerlyState
39 {
40 float *upper_bound; /* [nvecs] d(x, assigned centroid) */
41 float *lower_bound; /* [nvecs] d(x, 2nd nearest centroid) */
42 uint32_t nvecs;
43 uint32_t nlist;
44 uint32_t dim;
45 bool bounds_valid;
46 };
47
48 HamerlyState *
49 26 hamerly_create(uint32_t nvecs, uint32_t nlist, uint32_t dim)
50 {
51 26 HamerlyState *hs = vs_alloc0(sizeof(HamerlyState));
52 26 hs->nvecs = nvecs;
53 26 hs->nlist = nlist;
54 26 hs->dim = dim;
55 26 hs->upper_bound = vs_alloc(nvecs * sizeof(float));
56 26 hs->lower_bound = vs_alloc(nvecs * sizeof(float));
57 26 return hs;
58 }
59
60 void
61 26 hamerly_destroy(HamerlyState *hs)
62 {
63
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 26 times.
 26 if (hs == NULL)
64 ✗ return;
65 26 vs_free(hs->upper_bound);
66 26 vs_free(hs->lower_bound);
67 26 vs_free(hs);
68 }
69
70 /*
71 * Compute L2 squared distance using dot product decomposition.
72 * Requires precomputed norms: norm_a = ||a||^2, norm_b = ||b||^2.
73 */
74 static inline float
75 42496 l2_from_dot(float norm_a, float norm_b, float dot)
76 {
77 42496 float d = norm_a + norm_b - 2.0f * dot;
78
2/2
✓ Branch 0 taken 42400 times.
✓ Branch 1 taken 96 times.
 42496 return d > 0.0f ? d : 0.0f;
79 }
80
81 /*
82 * Full assignment for one vector: find nearest and second-nearest.
83 */
84 __attribute__((always_inline)) static inline void
85 ✗ assign_full_impl(
86 const void *vec,
87 float norm_x,
88 const KMeansState *st,
89 uint32_t *out_j1,
90 float *out_dist1,
91 float *out_dist2,
92 const Vec32TypeOps *ops)
93 {
94 8194 uint32_t nlist = st->nlist;
95 8194 Dimension dim = st->dim;
96 8194 const float *cents = st->centroids;
97
98 8194 float d1 = FLT_MAX, d2 = FLT_MAX;
99 8194 uint32_t j1 = 0;
100
101
8/8
✓ Branch 0 taken 2900 times.
✓ Branch 1 taken 700 times.
✓ Branch 2 taken 10480 times.
✓ Branch 3 taken 2096 times.
✓ Branch 4 taken 12300 times.
✓ Branch 5 taken 3300 times.
✓ Branch 6 taken 10490 times.
✓ Branch 7 taken 2098 times.
 44364 for (uint32_t j = 0; j < nlist; j++)
102 {
103 36170 float dp = ops->dot_product(vec, cents + (size_t)j * dim, dim);
104
0/8
✗ 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.
 36170 float d = l2_from_dot(norm_x, st->norms_c[j], dp);
105
106
8/8
✓ Branch 0 taken 1460 times.
✓ Branch 1 taken 1440 times.
✓ Branch 2 taken 4844 times.
✓ Branch 3 taken 5636 times.
✓ Branch 4 taken 6496 times.
✓ Branch 5 taken 5804 times.
✓ Branch 6 taken 4848 times.
✓ Branch 7 taken 5642 times.
 36170 if (d < d1)
107 {
108 17648 d2 = d1;
109 17648 d1 = d;
110 17648 j1 = j;
111 }
112
8/8
✓ Branch 0 taken 826 times.
✓ Branch 1 taken 614 times.
✓ Branch 2 taken 2696 times.
✓ Branch 3 taken 2940 times.
✓ Branch 4 taken 3486 times.
✓ Branch 5 taken 2318 times.
✓ Branch 6 taken 2692 times.
✓ Branch 7 taken 2950 times.
 18522 else if (d < d2)
113 {
114 9700 d2 = d;
115 }
116 }
117
118 8194 *out_j1 = j1;
119 8194 *out_dist1 = d1;
120 8194 *out_dist2 = d2;
121 8194 }
122
123 /*
124 * Precompute centroid norms (||c||^2) into st->norms_c.
125 */
126 static void
127 126 precompute_norms_c(KMeansState *st)
128 {
129
2/2
✓ Branch 0 taken 528 times.
✓ Branch 1 taken 126 times.
 654 for (uint32_t j = 0; j < st->nlist; j++)
130 528 st->norms_c[j] = vs_l2_norm_squared(
131 528 st->centroids + (size_t)j * st->dim, st->dim);
132 126 }
133
134 /*
135 * Main Hamerly assignment loop.
136 *
137 * always_inline — the specialized wrappers below pass a static const
138 * Vec32TypeOps from the header, so the compiler inlines through
139 * every vtable function pointer. VS_TARGET_CLONES on the wrappers
140 * generates AVX2/AVX-512 variants of the entire inlined body.
141 */
142 __attribute__((always_inline)) static inline void
143 ✗ hamerly_assign_impl(KMeansState *st, HamerlyState *hs, const Vec32TypeOps *ops)
144 {
145 90 uint32_t nvecs = st->nvecs;
146 90 Dimension dim = st->dim;
147 90 size_t esz = ops->element_size;
148 90 const float *cents = st->centroids;
149
150 90 if (!hs->bounds_valid)
151 {
152 /*
153 * First call: full assignment, initialize bounds.
154 * Find nearest and second-nearest for each vector.
155 */
156
2/2
✓ Branch 0 taken 3300 times.
✓ Branch 1 taken 22 times.
 3322 for (uint32_t i = 0; i < nvecs; i++)
157 {
158
0/2
✗ Branch 0 not taken.
✗ Branch 1 not taken.
 3300 const void *v = km_get_vector(st, i, esz);
159 3300 float nx = st->norms_x[i];
160
161 ✗ uint32_t j1;
162 ✗ float d1, d2;
163 ✗ assign_full_impl(v, nx, st, &j1, &d1, &d2, ops);
164
165 3300 st->assignments[i] = j1;
166 3300 st->total_cost += d1;
167 3300 hs->upper_bound[i] = sqrtf(d1);
168 3300 hs->lower_bound[i] = sqrtf(d2);
169 }
170
171 22 hs->bounds_valid = true;
172 22 return;
173 }
174
175 /*
176 * Subsequent calls: use bounds to skip vectors.
177 */
178
2/2
✓ Branch 0 taken 11400 times.
✓ Branch 1 taken 68 times.
 11468 for (uint32_t i = 0; i < nvecs; i++)
179 {
180 11400 float ub = hs->upper_bound[i];
181 11400 float lb = hs->lower_bound[i];
182
183 /* Step 1: bound check — can we skip this vector? */
184
2/2
✓ Branch 0 taken 8236 times.
✓ Branch 1 taken 3164 times.
 11400 if (ub <= lb)
185 {
186 /* Assignment can't change. Accumulate cost from bound. */
187 8236 st->total_cost += ub * ub;
188 9302 continue;
189 }
190
191 /* Step 2: tighten upper bound with exact distance */
192
0/2
✗ Branch 0 not taken.
✗ Branch 1 not taken.
 3164 const void *v = km_get_vector(st, i, esz);
193 3164 float nx = st->norms_x[i];
194 3164 uint32_t prev = st->assignments[i];
195 3164 float dot_p = ops->dot_product(v, cents + (size_t)prev * dim, dim);
196
0/2
✗ Branch 0 not taken.
✗ Branch 1 not taken.
 3164 float d_sq = l2_from_dot(nx, st->norms_c[prev], dot_p);
197 3164 ub = sqrtf(d_sq);
198
199 3164 hs->upper_bound[i] = ub;
200
201 /* Check again with tightened bound */
202
2/2
✓ Branch 0 taken 1066 times.
✓ Branch 1 taken 2098 times.
 3164 if (ub <= lb)
203 {
204 1066 st->total_cost += d_sq;
205 1066 continue;
206 }
207
208 /* Step 3: full search required */
209 ✗ uint32_t j1;
210 ✗ float d1, d2;
211 ✗ assign_full_impl(v, nx, st, &j1, &d1, &d2, ops);
212
213 2098 st->assignments[i] = j1;
214 2098 st->total_cost += d1;
215 2098 hs->upper_bound[i] = sqrtf(d1);
216 2098 hs->lower_bound[i] = sqrtf(d2);
217 }
218 }
219
220 /*
221 * Preconvert f16 variant: converts each f16 vector to f32 once before
222 * the centroid loop, then uses the f32 dot product kernel. Saves
223 * O(K*dim) f16→f32 conversions per vector (one conversion vs K).
224 * Vectors skipped by bounds check are never converted.
225 */
226 __attribute__((always_inline)) static inline void
227 ✗ hamerly_assign_preconvert_impl(KMeansState *st, HamerlyState *hs, size_t esz)
228 {
229 36 uint32_t nvecs = st->nvecs;
230 36 Dimension dim = st->dim;
231 36 const float *cents = st->centroids;
232 36 float *buf = st->vec_block;
233
234 36 const Vec32TypeOps *f32ops = &vs_f32_type_ops;
235
236 36 if (!hs->bounds_valid)
237 {
238
2/2
✓ Branch 0 taken 700 times.
✓ Branch 1 taken 4 times.
 704 for (uint32_t i = 0; i < nvecs; i++)
239 {
240
0/2
✗ Branch 0 not taken.
✗ Branch 1 not taken.
 700 const void *raw = km_get_vector(st, i, esz);
241 700 vs_half_to_float_array((const half *)raw, buf, dim);
242 700 float nx = st->norms_x[i];
243
244 ✗ uint32_t j1;
245 ✗ float d1, d2;
246 ✗ assign_full_impl(buf, nx, st, &j1, &d1, &d2, f32ops);
247
248 700 st->assignments[i] = j1;
249 700 st->total_cost += d1;
250 700 hs->upper_bound[i] = sqrtf(d1);
251 700 hs->lower_bound[i] = sqrtf(d2);
252 }
253
254 4 hs->bounds_valid = true;
255 4 return;
256 }
257
258
2/2
✓ Branch 0 taken 6200 times.
✓ Branch 1 taken 32 times.
 6232 for (uint32_t i = 0; i < nvecs; i++)
259 {
260 6200 float ub = hs->upper_bound[i];
261 6200 float lb = hs->lower_bound[i];
262
263
2/2
✓ Branch 0 taken 3038 times.
✓ Branch 1 taken 3162 times.
 6200 if (ub <= lb)
264 {
265 3038 st->total_cost += ub * ub;
266 4104 continue;
267 }
268
269 /* Convert once — only for vectors that need distance computation */
270
0/2
✗ Branch 0 not taken.
✗ Branch 1 not taken.
 3162 const void *raw = km_get_vector(st, i, esz);
271 3162 vs_half_to_float_array((const half *)raw, buf, dim);
272
273 3162 float nx = st->norms_x[i];
274 3162 uint32_t prev = st->assignments[i];
275 ✗ float dot_p =
276 3162 f32ops->dot_product(buf, cents + (size_t)prev * dim, dim);
277
0/2
✗ Branch 0 not taken.
✗ Branch 1 not taken.
 3162 float d_sq = l2_from_dot(nx, st->norms_c[prev], dot_p);
278 3162 ub = sqrtf(d_sq);
279
280 3162 hs->upper_bound[i] = ub;
281
282
2/2
✓ Branch 0 taken 1066 times.
✓ Branch 1 taken 2096 times.
 3162 if (ub <= lb)
283 {
284 1066 st->total_cost += d_sq;
285 1066 continue;
286 }
287
288 ✗ uint32_t j1;
289 ✗ float d1, d2;
290 ✗ assign_full_impl(buf, nx, st, &j1, &d1, &d2, f32ops);
291
292 2096 st->assignments[i] = j1;
293 2096 st->total_cost += d1;
294 2096 hs->upper_bound[i] = sqrtf(d1);
295 2096 hs->lower_bound[i] = sqrtf(d2);
296 }
297 }
298
299 /* Specialized wrappers — VS_TARGET_CLONES generates SIMD variants */
300 VS_TARGET_CLONES static void
301
2/2
✓ Branch 0 taken 22 times.
✓ Branch 1 taken 68 times.
 90 hamerly_assign_f32(KMeansState *st, HamerlyState *hs)
302 {
303 ✗ hamerly_assign_impl(st, hs, &vs_f32_type_ops);
304 90 }
305
306 VS_TARGET_CLONES static void
307
2/2
✓ Branch 0 taken 4 times.
✓ Branch 1 taken 32 times.
 36 hamerly_assign_f16(KMeansState *st, HamerlyState *hs)
308 {
309 ✗ hamerly_assign_preconvert_impl(st, hs, sizeof(half));
310 36 }
311
312 /* Public entry — dispatches once based on type */
313 void
314 126 hamerly_assign(KMeansState *st, HamerlyState *hs)
315 {
316 126 st->total_cost = 0.0f;
317 126 precompute_norms_c(st);
318
319
2/2
✓ Branch 0 taken 90 times.
✓ Branch 1 taken 36 times.
 126 switch (st->vec_type)
320 {
321 90 case VS_VEC_F32:
322 90 hamerly_assign_f32(st, hs);
323 90 break;
324 36 default:
325 36 hamerly_assign_f16(st, hs);
326 36 break;
327 }
328 126 }
329
330 void
331 100 hamerly_update_bounds(
332 KMeansState *st, HamerlyState *hs, const float *old_centroids)
333 {
334 100 uint32_t nlist = st->nlist;
335 100 uint32_t dim = st->dim;
336 100 uint32_t nvecs = st->nvecs;
337
338 /*
339 * Compute per-centroid movement (Euclidean distance) and
340 * find the two largest movements: delta_max and delta_second.
341 *
342 * delta_max is used to loosen all lower bounds.
343 * We track which centroid moved the most (p_max) so that
344 * vectors assigned to it use delta_second instead.
345 * (Hamerly's original paper optimization.)
346 */
347 100 float delta_max = 0.0f, delta_second = 0.0f;
348 100 uint32_t p_max = 0;
349
350
2/2
✓ Branch 0 taken 432 times.
✓ Branch 1 taken 100 times.
 532 for (uint32_t j = 0; j < nlist; j++)
351 {
352 432 float d_sq = vs_l2_distance_squared(
353 432 st->centroids + (size_t)j * dim,
354 432 old_centroids + (size_t)j * dim,
355 dim);
356 432 float d = sqrtf(d_sq);
357
358
2/2
✓ Branch 0 taken 210 times.
✓ Branch 1 taken 222 times.
 432 if (d >= delta_max)
359 {
360 210 delta_second = delta_max;
361 210 delta_max = d;
362 210 p_max = j;
363 }
364
2/2
✓ Branch 0 taken 122 times.
✓ Branch 1 taken 100 times.
 222 else if (d > delta_second)
365 {
366 122 delta_second = d;
367 }
368 }
369
370 /*
371 * Also need per-centroid delta for the upper bound update.
372 * We compute these on the fly to avoid an extra allocation.
373 * Store them temporarily in norms_c (will be recomputed at
374 * next assign call anyway).
375 */
376 100 float *delta = st->norms_c; /* temporary reuse */
377
2/2
✓ Branch 0 taken 432 times.
✓ Branch 1 taken 100 times.
 532 for (uint32_t j = 0; j < nlist; j++)
378 {
379 432 float d_sq = vs_l2_distance_squared(
380 432 st->centroids + (size_t)j * dim,
381 432 old_centroids + (size_t)j * dim,
382 dim);
383 432 delta[j] = sqrtf(d_sq);
384 }
385
386 /* Update bounds for each vector */
387
2/2
✓ Branch 0 taken 17600 times.
✓ Branch 1 taken 100 times.
 17700 for (uint32_t i = 0; i < nvecs; i++)
388 {
389 17600 uint32_t a = st->assignments[i];
390
391 17600 hs->upper_bound[i] += delta[a];
392
393 /*
394 * Use delta_second if this vector's assigned centroid is
395 * the one that moved the most. Otherwise use delta_max.
396 */
397
2/2
✓ Branch 0 taken 4124 times.
✓ Branch 1 taken 13476 times.
 17600 float lb_delta = (a == p_max) ? delta_second : delta_max;
398 17600 hs->lower_bound[i] -= lb_delta;
399
1/2
✗ Branch 0 not taken.
✓ Branch 1 taken 17600 times.
 17600 if (hs->lower_bound[i] < 0.0f)
400 ✗ hs->lower_bound[i] = 0.0f;
401 }
402 100 }
403