| 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_lloyd.c - Lloyd's k-means assignment step | ||
| 6 | * | ||
| 7 | * Brute-force assignment: compute the full N×K distance matrix each | ||
| 8 | * iteration. Two backends share the same block-processing structure: | ||
| 9 | * | ||
| 10 | * CBLAS path: | ||
| 11 | * Uses sgemm for the [block × dim] × [dim × nlist] dot product | ||
| 12 | * matrix. BLAS cache-tiling makes this 10-50x faster than naive | ||
| 13 | * loops for large K. | ||
| 14 | * | ||
| 15 | * Builtin path: | ||
| 16 | * Batch dot-product kernel with VS_TARGET_CLONES for AVX-512/AVX2 | ||
| 17 | * auto-vectorization. FMA generation requires -ffp-contract=fast | ||
| 18 | * (set in meson.build). | ||
| 19 | * | ||
| 20 | * Both use the decomposition: | ||
| 21 | * ||x - c||² = ||x||² + ||c||² - 2⟨x,c⟩ | ||
| 22 | * The ⟨x,c⟩ term is computed in batch, norms are precomputed. | ||
| 23 | */ | ||
| 24 | |||
| 25 | #include "vs_config.h" | ||
| 26 | |||
| 27 | #include <math.h> | ||
| 28 | #include <stddef.h> | ||
| 29 | #include <stdint.h> | ||
| 30 | #include <string.h> | ||
| 31 | |||
| 32 | #ifdef VS_HAVE_CBLAS | ||
| 33 | /* See matrix.c for the rationale on the Apple branch. */ | ||
| 34 | #ifdef __APPLE__ | ||
| 35 | #include <vecLib/cblas_new.h> | ||
| 36 | #else | ||
| 37 | #include <cblas.h> | ||
| 38 | #endif | ||
| 39 | #endif | ||
| 40 | |||
| 41 | #include "algo/kmeans_lloyd.h" | ||
| 42 | #include "algo/vecops.h" | ||
| 43 | #include "core/log.h" | ||
| 44 | #include "core/memory.h" | ||
| 45 | #include "types/vec16.h" | ||
| 46 | |||
| 47 | /* | ||
| 48 | * Precompute ||c||² for all centroids (L2 only). | ||
| 49 | */ | ||
| 50 | static void | ||
| 51 | 17536 | precompute_norms_c(KMeansState *st) | |
| 52 | { | ||
| 53 |
2/2✓ Branch 0 taken 128643 times.
✓ Branch 1 taken 17536 times.
|
146179 | for (uint32_t j = 0; j < st->nlist; j++) |
| 54 | 128643 | st->norms_c[j] = vs_l2_norm_squared( | |
| 55 | 128643 | st->centroids + (size_t)j * st->dim, st->dim); | |
| 56 | 17536 | } | |
| 57 | |||
| 58 | /* | ||
| 59 | * Assignment step: CBLAS path for one block. | ||
| 60 | * | ||
| 61 | * Computes the N_block x K distance matrix using sgemm, then finds | ||
| 62 | * the nearest centroid per vector (argmin per row). | ||
| 63 | * | ||
| 64 | * For L2: dist[i][j] = ||x_i||² + ||c_j||² - 2⟨x_i, c_j⟩ | ||
| 65 | * For IP: dist[i][j] = -⟨x_i, c_j⟩ | ||
| 66 | * For cos: dist[i][j] = 1 - ⟨x_i, c_j⟩ | ||
| 67 | */ | ||
| 68 | #ifdef VS_HAVE_CBLAS | ||
| 69 | __attribute__((always_inline)) static inline void | ||
| 70 | lloyd_assign_block_cblas_impl( | ||
| 71 | KMeansState *st, | ||
| 72 | uint32_t block_start, | ||
| 73 | uint32_t block_count, | ||
| 74 | const Vec32TypeOps *ops) | ||
| 75 | { | ||
| 76 | uint32_t nlist = st->nlist; | ||
| 77 | uint32_t dim = st->dim; | ||
| 78 | float *dist = st->dist_block; | ||
| 79 | |||
| 80 | /* Get float32 view of this block */ | ||
| 81 | const float *block_vecs; | ||
| 82 | if (st->indices != NULL) | ||
| 83 | { | ||
| 84 | /* Gather indexed vectors into vec_block */ | ||
| 85 | size_t esz = ops->element_size; | ||
| 86 | for (uint32_t i = 0; i < block_count; i++) | ||
| 87 | { | ||
| 88 | uint32_t idx = st->indices[block_start + i]; | ||
| 89 | const void *src = (const char *)st->vectors + | ||
| 90 | (size_t)idx * dim * esz; | ||
| 91 | /* Unreachable | ||
| 92 | * unless the standalone arena is exhausted, which no | ||
| 93 | * caller handles. */ | ||
| 94 | /* NOLINTNEXTLINE(clang-analyzer-core.NullPointerArithm) */ | ||
| 95 | ops->to_float_one(src, st->vec_block + (size_t)i * dim, dim); | ||
| 96 | } | ||
| 97 | block_vecs = st->vec_block; | ||
| 98 | } | ||
| 99 | else | ||
| 100 | { | ||
| 101 | /* Zero-copy for f32, converts for f16 */ | ||
| 102 | const void *raw = (const char *)st->vectors + | ||
| 103 | (size_t)block_start * dim * ops->element_size; | ||
| 104 | block_vecs = ops->to_float_block(raw, st->vec_block, block_count, dim); | ||
| 105 | } | ||
| 106 | |||
| 107 | /* Compute dot products via sgemm */ | ||
| 108 | float alpha = -2.0f; | ||
| 109 | if (st->metric == DISTANCE_INNER_PRODUCT || st->metric == DISTANCE_COSINE) | ||
| 110 | alpha = -1.0f; | ||
| 111 | |||
| 112 | cblas_sgemm( | ||
| 113 | CblasRowMajor, | ||
| 114 | CblasNoTrans, | ||
| 115 | CblasTrans, | ||
| 116 | (int)block_count, /* M: rows */ | ||
| 117 | (int)nlist, /* N: cols */ | ||
| 118 | (int)dim, /* K: inner dim */ | ||
| 119 | alpha, | ||
| 120 | block_vecs, | ||
| 121 | (int)dim, | ||
| 122 | st->centroids, | ||
| 123 | (int)dim, | ||
| 124 | 0.0f, | ||
| 125 | dist, | ||
| 126 | (int)nlist); | ||
| 127 | |||
| 128 | /* Add norms / offset */ | ||
| 129 | switch (st->metric) | ||
| 130 | { | ||
| 131 | case DISTANCE_L2: | ||
| 132 | for (uint32_t i = 0; i < block_count; i++) | ||
| 133 | { | ||
| 134 | float nx = st->norms_x[block_start + i]; | ||
| 135 | for (uint32_t j = 0; j < nlist; j++) | ||
| 136 | dist[(size_t)i * nlist + j] += nx + st->norms_c[j]; | ||
| 137 | } | ||
| 138 | break; | ||
| 139 | case DISTANCE_COSINE: | ||
| 140 | for (uint32_t i = 0; i < block_count; i++) | ||
| 141 | for (uint32_t j = 0; j < nlist; j++) | ||
| 142 | dist[(size_t)i * nlist + j] += 1.0f; | ||
| 143 | break; | ||
| 144 | case DISTANCE_INNER_PRODUCT: | ||
| 145 | /* dist = -⟨x,c⟩, already correct */ | ||
| 146 | break; | ||
| 147 | } | ||
| 148 | |||
| 149 | /* Find argmin per row */ | ||
| 150 | for (uint32_t i = 0; i < block_count; i++) | ||
| 151 | { | ||
| 152 | uint32_t vec_idx = block_start + i; | ||
| 153 | float min_dist = dist[(size_t)i * nlist]; | ||
| 154 | uint32_t min_j = 0; | ||
| 155 | |||
| 156 | for (uint32_t j = 1; j < nlist; j++) | ||
| 157 | { | ||
| 158 | float d = dist[(size_t)i * nlist + j]; | ||
| 159 | if (d < min_dist) | ||
| 160 | { | ||
| 161 | min_dist = d; | ||
| 162 | min_j = j; | ||
| 163 | } | ||
| 164 | } | ||
| 165 | |||
| 166 | st->assignments[vec_idx] = min_j; | ||
| 167 | st->total_cost += min_dist; | ||
| 168 | } | ||
| 169 | } | ||
| 170 | |||
| 171 | static void | ||
| 172 | lloyd_assign_block_cblas_f32( | ||
| 173 | KMeansState *st, uint32_t block_start, uint32_t block_count) | ||
| 174 | { | ||
| 175 | lloyd_assign_block_cblas_impl( | ||
| 176 | st, block_start, block_count, &vs_f32_type_ops); | ||
| 177 | } | ||
| 178 | |||
| 179 | static void | ||
| 180 | lloyd_assign_block_cblas_f16( | ||
| 181 | KMeansState *st, uint32_t block_start, uint32_t block_count) | ||
| 182 | { | ||
| 183 | lloyd_assign_block_cblas_impl( | ||
| 184 | st, block_start, block_count, &vs_f16_type_ops); | ||
| 185 | } | ||
| 186 | |||
| 187 | #if defined(VS_F16C_SUPPORT) && !defined(VS_SIMD_NONE) | ||
| 188 | static void | ||
| 189 | lloyd_assign_block_cblas_f16c( | ||
| 190 | KMeansState *st, uint32_t block_start, uint32_t block_count) | ||
| 191 | { | ||
| 192 | lloyd_assign_block_cblas_impl( | ||
| 193 | st, block_start, block_count, &vs_f16c_type_ops); | ||
| 194 | } | ||
| 195 | #endif | ||
| 196 | #endif /* VS_HAVE_CBLAS */ | ||
| 197 | |||
| 198 | /* | ||
| 199 | * Batch dot-product matrix: dots[i*nlist + j] = dot(vecs[i], cents[j]) | ||
| 200 | * | ||
| 201 | * VS_TARGET_CLONES generates AVX-512, AVX2, and default versions. | ||
| 202 | * The innermost loop over dim is auto-vectorized by the compiler, | ||
| 203 | * eliminating per-vector function pointer dispatch overhead. | ||
| 204 | * | ||
| 205 | * FMA generation requires -ffp-contract=fast (set in meson.build). | ||
| 206 | * Without it, GCC in an ISO C mode generates separate vmulps + | ||
| 207 | * horizontal scalar adds instead of vfmadd231ps accumulate. | ||
| 208 | */ | ||
| 209 | VS_TARGET_CLONES static void | ||
| 210 | 21748 | lloyd_compute_dot_products( | |
| 211 | const float *vecs, | ||
| 212 | const float *centroids, | ||
| 213 | float *dots, | ||
| 214 | uint32_t block_count, | ||
| 215 | uint32_t nlist, | ||
| 216 | uint32_t dim) | ||
| 217 | { | ||
| 218 |
2/2✓ Branch 0 taken 13744612 times.
✓ Branch 1 taken 21748 times.
|
13766360 | for (uint32_t i = 0; i < block_count; i++) |
| 219 | { | ||
| 220 | 13744612 | const float *v = vecs + (size_t)i * dim; | |
| 221 |
2/2✓ Branch 0 taken 134367756 times.
✓ Branch 1 taken 13744612 times.
|
148112368 | for (uint32_t j = 0; j < nlist; j++) |
| 222 | { | ||
| 223 | 134367756 | const float *c = centroids + (size_t)j * dim; | |
| 224 | 134367756 | float dot = 0.0f; | |
| 225 |
2/2✓ Branch 0 taken 9032265479 times.
✓ Branch 1 taken 134367756 times.
|
9166633235 | for (uint32_t d = 0; d < dim; d++) |
| 226 | /* Unreachable | ||
| 227 | * unless the standalone arena is exhausted, which no | ||
| 228 | * caller handles. */ | ||
| 229 | /* NOLINTNEXTLINE(clang-analyzer-core.NullDereference) */ | ||
| 230 | 9032265479 | dot += v[d] * c[d]; | |
| 231 | 134367756 | dots[(size_t)i * nlist + j] = dot; | |
| 232 | } | ||
| 233 | } | ||
| 234 | 21748 | } | |
| 235 | |||
| 236 | /* | ||
| 237 | * Convert dot products to distances and find argmin per row. | ||
| 238 | */ | ||
| 239 | static void | ||
| 240 | 21748 | lloyd_dots_to_assignments( | |
| 241 | KMeansState *st, | ||
| 242 | float *dist, | ||
| 243 | uint32_t block_start, | ||
| 244 | uint32_t block_count) | ||
| 245 | { | ||
| 246 | 21748 | uint32_t nlist = st->nlist; | |
| 247 | |||
| 248 | /* Convert dot products to distances based on metric */ | ||
| 249 |
3/4✓ Branch 0 taken 8226 times.
✓ Branch 1 taken 374 times.
✓ Branch 2 taken 13148 times.
✗ Branch 3 not taken.
|
21748 | switch (st->metric) |
| 250 | { | ||
| 251 | 7528 | case DISTANCE_L2: | |
| 252 | /* dist = ||x||² + ||c||² - 2*dot(x,c) */ | ||
| 253 |
2/2✓ Branch 0 taken 13226788 times.
✓ Branch 1 taken 20456 times.
|
13247244 | for (uint32_t i = 0; i < block_count; i++) |
| 254 | { | ||
| 255 | 13226788 | float nx = st->norms_x[block_start + i]; | |
| 256 |
2/2✓ Branch 0 taken 131028649 times.
✓ Branch 1 taken 13226788 times.
|
144255437 | for (uint32_t j = 0; j < nlist; j++) |
| 257 | { | ||
| 258 | 131028649 | size_t idx = (size_t)i * nlist + j; | |
| 259 | 131028649 | dist[idx] = nx + st->norms_c[j] - 2.0f * dist[idx]; | |
| 260 | } | ||
| 261 | } | ||
| 262 | 7528 | break; | |
| 263 | 368 | case DISTANCE_INNER_PRODUCT: | |
| 264 | /* dist = -dot(x,c) */ | ||
| 265 |
2/2✓ Branch 0 taken 233784 times.
✓ Branch 1 taken 6368 times.
|
240152 | for (uint32_t i = 0; i < block_count; i++) |
| 266 |
2/2✓ Branch 0 taken 887334 times.
✓ Branch 1 taken 239778 times.
|
1127112 | for (uint32_t j = 0; j < nlist; j++) |
| 267 | { | ||
| 268 | 887334 | size_t idx = (size_t)i * nlist + j; | |
| 269 | 887334 | dist[idx] = -dist[idx]; | |
| 270 | } | ||
| 271 | 368 | break; | |
| 272 | 220 | case DISTANCE_COSINE: | |
| 273 | /* dist = 1 - dot(x,c) (vectors pre-normalized) */ | ||
| 274 |
2/2✓ Branch 0 taken 58408 times.
✓ Branch 1 taken 220556 times.
|
278964 | for (uint32_t i = 0; i < block_count; i++) |
| 275 |
2/2✓ Branch 0 taken 2451773 times.
✓ Branch 1 taken 278046 times.
|
2729819 | for (uint32_t j = 0; j < nlist; j++) |
| 276 | { | ||
| 277 | 2451773 | size_t idx = (size_t)i * nlist + j; | |
| 278 | 2451773 | dist[idx] = 1.0f - dist[idx]; | |
| 279 | } | ||
| 280 | 220 | break; | |
| 281 | } | ||
| 282 | |||
| 283 | /* Find argmin per row */ | ||
| 284 |
2/2✓ Branch 0 taken 13744612 times.
✓ Branch 1 taken 21748 times.
|
13766360 | for (uint32_t i = 0; i < block_count; i++) |
| 285 | { | ||
| 286 | 13744612 | uint32_t vec_idx = block_start + i; | |
| 287 | 13744612 | float min_dist = dist[(size_t)i * nlist]; | |
| 288 | 13744612 | uint32_t min_j = 0; | |
| 289 | |||
| 290 |
2/2✓ Branch 0 taken 120623144 times.
✓ Branch 1 taken 13744612 times.
|
134367756 | for (uint32_t j = 1; j < nlist; j++) |
| 291 | { | ||
| 292 | 120623144 | float d = dist[(size_t)i * nlist + j]; | |
| 293 |
2/2✓ Branch 0 taken 22280160 times.
✓ Branch 1 taken 98342984 times.
|
120623144 | if (d < min_dist) |
| 294 | { | ||
| 295 | 22280160 | min_dist = d; | |
| 296 | 22280160 | min_j = j; | |
| 297 | } | ||
| 298 | } | ||
| 299 | |||
| 300 | 13744612 | st->assignments[vec_idx] = min_j; | |
| 301 | 13744612 | st->total_cost += min_dist; | |
| 302 | } | ||
| 303 | 21748 | } | |
| 304 | |||
| 305 | /* | ||
| 306 | * Assignment step: builtin fallback for one block. | ||
| 307 | * | ||
| 308 | * For f32 input: uses batch dot-product + norm decomposition. | ||
| 309 | * For f16 input: preconverts block to f32, then uses the same f32 kernel. | ||
| 310 | */ | ||
| 311 | static void | ||
| 312 | 21736 | lloyd_assign_block_builtin_f32( | |
| 313 | KMeansState *st, uint32_t block_start, uint32_t block_count) | ||
| 314 | { | ||
| 315 | 21736 | uint32_t nlist = st->nlist; | |
| 316 | 21736 | uint32_t dim = st->dim; | |
| 317 | 21736 | float *dist = st->dist_block; | |
| 318 | |||
| 319 | 13632 | const float *block_vecs; | |
| 320 |
2/2✓ Branch 0 taken 17321 times.
✓ Branch 1 taken 4415 times.
|
21736 | if (st->indices != NULL) |
| 321 | { | ||
| 322 | /* Gather indexed vectors into contiguous block */ | ||
| 323 |
2/2✓ Branch 0 taken 8058587 times.
✓ Branch 1 taken 17321 times.
|
8075908 | for (uint32_t i = 0; i < block_count; i++) |
| 324 | { | ||
| 325 | 8058587 | uint32_t idx = st->indices[block_start + i]; | |
| 326 | /* Unreachable unless the standalone arena is exhausted, which | ||
| 327 | * no caller handles. */ | ||
| 328 | /* NOLINTNEXTLINE(clang-analyzer-core.NonNullParamChecker) */ | ||
| 329 | 8058587 | memcpy(st->vec_block + (size_t)i * dim, | |
| 330 | 8058587 | (const float *)st->vectors + (size_t)idx * dim, | |
| 331 | dim * sizeof(float)); | ||
| 332 | } | ||
| 333 | 17321 | block_vecs = st->vec_block; | |
| 334 | } | ||
| 335 | else | ||
| 336 | { | ||
| 337 | 4415 | block_vecs = (const float *)st->vectors + (size_t)block_start * dim; | |
| 338 | } | ||
| 339 | |||
| 340 | 21736 | lloyd_compute_dot_products( | |
| 341 | 21736 | block_vecs, st->centroids, dist, block_count, nlist, dim); | |
| 342 | 21736 | lloyd_dots_to_assignments(st, dist, block_start, block_count); | |
| 343 | 21736 | } | |
| 344 | |||
| 345 | static void | ||
| 346 | 12 | lloyd_assign_block_builtin_f16( | |
| 347 | KMeansState *st, uint32_t block_start, uint32_t block_count) | ||
| 348 | { | ||
| 349 | 12 | uint32_t nlist = st->nlist; | |
| 350 | 12 | uint32_t dim = st->dim; | |
| 351 | |||
| 352 | /* Preconvert f16 block to f32 — O(block*dim), saves O(block*K*dim) */ | ||
| 353 |
1/2✗ Branch 0 not taken.
✓ Branch 1 taken 12 times.
|
12 | if (st->indices != NULL) |
| 354 | { | ||
| 355 | ✗ | for (uint32_t i = 0; i < block_count; i++) | |
| 356 | { | ||
| 357 | ✗ | uint32_t idx = st->indices[block_start + i]; | |
| 358 | ✗ | const half *src = (const half *)st->vectors + (size_t)idx * dim; | |
| 359 | /* Unreachable | ||
| 360 | * unless the standalone arena is exhausted, which no | ||
| 361 | * caller handles. */ | ||
| 362 | /* NOLINTNEXTLINE(clang-analyzer-core.NullPointerArithm) */ | ||
| 363 | ✗ | vs_half_to_float_array(src, st->vec_block + (size_t)i * dim, dim); | |
| 364 | } | ||
| 365 | } | ||
| 366 | else | ||
| 367 | { | ||
| 368 | 12 | const half *src = (const half *)st->vectors + | |
| 369 | 12 | (size_t)block_start * dim; | |
| 370 | 12 | vs_half_to_float_array(src, st->vec_block, block_count * dim); | |
| 371 | } | ||
| 372 | |||
| 373 | 12 | lloyd_compute_dot_products( | |
| 374 | 12 | st->vec_block, | |
| 375 | 12 | st->centroids, | |
| 376 | st->dist_block, | ||
| 377 | block_count, | ||
| 378 | nlist, | ||
| 379 | dim); | ||
| 380 | 12 | lloyd_dots_to_assignments(st, st->dist_block, block_start, block_count); | |
| 381 | 12 | } | |
| 382 | |||
| 383 | /* Block assignment function pointer — selected once per lloyd_assign call */ | ||
| 384 | typedef void (*lloyd_block_fn)(KMeansState *, uint32_t, uint32_t); | ||
| 385 | |||
| 386 | static lloyd_block_fn | ||
| 387 | 2639 | lloyd_select_block_fn(KMeansState *st, bool use_cblas) | |
| 388 | { | ||
| 389 | #ifdef VS_HAVE_CBLAS | ||
| 390 | if (use_cblas) | ||
| 391 | { | ||
| 392 | switch (st->vec_type) | ||
| 393 | { | ||
| 394 | case VS_VEC_F32: | ||
| 395 | return lloyd_assign_block_cblas_f32; | ||
| 396 | #if defined(VS_F16C_SUPPORT) && !defined(VS_SIMD_NONE) | ||
| 397 | case VS_VEC_F16C: | ||
| 398 | return lloyd_assign_block_cblas_f16c; | ||
| 399 | #endif | ||
| 400 | default: | ||
| 401 | return lloyd_assign_block_cblas_f16; | ||
| 402 | } | ||
| 403 | } | ||
| 404 | #endif | ||
| 405 | 1517 | (void)use_cblas; | |
| 406 |
2/2✓ Branch 0 taken 1110 times.
✓ Branch 1 taken 12 times.
|
2639 | switch (st->vec_type) |
| 407 | { | ||
| 408 | 1110 | case VS_VEC_F32: | |
| 409 | 1110 | return lloyd_assign_block_builtin_f32; | |
| 410 | 12 | default: | |
| 411 | 12 | return lloyd_assign_block_builtin_f16; | |
| 412 | } | ||
| 413 | } | ||
| 414 | |||
| 415 | /* Cap on the distance buffer each thread touches per block. */ | ||
| 416 | #define LLOYD_DIST_BUF_MAX (8 * 1024 * 1024) | ||
| 417 | |||
| 418 | /* | ||
| 419 | * Adaptive block size for the parallel builtin path. | ||
| 420 | * | ||
| 421 | * The distance buffer is [block × nlist] floats per thread. With | ||
| 422 | * fixed block=4096, large nlist blows past L3 (e.g., 4096×8000×4 | ||
| 423 | * = 128MB per thread). Capping it at LLOYD_DIST_BUF_MAX keeps the | ||
| 424 | * distance matrix, vector block, and centroid tile near L3 instead | ||
| 425 | * of growing with nlist. | ||
| 426 | */ | ||
| 427 | static uint32_t | ||
| 428 | 2627 | lloyd_parallel_block_size(uint32_t nlist, uint32_t nvecs) | |
| 429 | { | ||
| 430 | 2627 | uint32_t block = KMEANS_BLOCK_SIZE; | |
| 431 | 2627 | uint32_t buf_sz = block * nlist * (uint32_t)sizeof(float); | |
| 432 |
1/2✗ Branch 0 not taken.
✓ Branch 1 taken 1110 times.
|
2627 | if (buf_sz > LLOYD_DIST_BUF_MAX) |
| 433 | { | ||
| 434 | ✗ | block = LLOYD_DIST_BUF_MAX / (nlist * (uint32_t)sizeof(float)); | |
| 435 | ✗ | if (block < 64) | |
| 436 | ✗ | block = 64; | |
| 437 | } | ||
| 438 |
1/2✓ Branch 0 taken 1110 times.
✗ Branch 1 not taken.
|
2627 | if (block > nvecs) |
| 439 | 1110 | block = nvecs; | |
| 440 | 2627 | return block; | |
| 441 | } | ||
| 442 | |||
| 443 | static void | ||
| 444 | 21055 | lloyd_assign_range( | |
| 445 | KMeansState *st, | ||
| 446 | lloyd_block_fn block_fn, | ||
| 447 | float *dist_buf, | ||
| 448 | uint32_t block_size, | ||
| 449 | uint32_t range_start, | ||
| 450 | uint32_t range_end, | ||
| 451 | float *cost_out) | ||
| 452 | { | ||
| 453 | 21055 | float *saved_dist = st->dist_block; | |
| 454 | 21055 | st->dist_block = dist_buf; | |
| 455 | |||
| 456 | 21055 | float cost = 0.0f; | |
| 457 |
2/2✓ Branch 0 taken 21748 times.
✓ Branch 1 taken 21055 times.
|
42803 | for (uint32_t start = range_start; start < range_end; start += block_size) |
| 458 | { | ||
| 459 | 21748 | uint32_t count = range_end - start; | |
| 460 |
1/2✗ Branch 0 not taken.
✓ Branch 1 taken 8116 times.
|
21748 | if (count > block_size) |
| 461 | ✗ | count = block_size; | |
| 462 | 21748 | st->total_cost = 0.0f; | |
| 463 | 21748 | block_fn(st, start, count); | |
| 464 | 21748 | cost += st->total_cost; | |
| 465 | } | ||
| 466 | 21055 | *cost_out = cost; | |
| 467 | 21055 | st->dist_block = saved_dist; | |
| 468 | 21055 | } | |
| 469 | |||
| 470 | typedef struct | ||
| 471 | { | ||
| 472 | KMeansState *st; | ||
| 473 | lloyd_block_fn block_fn; | ||
| 474 | float *dist_bufs; /* [nthreads * block * nlist] */ | ||
| 475 | float *costs; /* [nthreads] */ | ||
| 476 | float *vec_bufs; /* [nthreads * block * dim] or NULL */ | ||
| 477 | uint32_t block; | ||
| 478 | } LloydParCtx; | ||
| 479 | |||
| 480 | static void | ||
| 481 | 2627 | lloyd_par_worker(uint32_t thread_id, uint32_t start, uint32_t end, void *arg) | |
| 482 | { | ||
| 483 | 2627 | LloydParCtx *ctx = (LloydParCtx *)arg; | |
| 484 | 2627 | uint32_t nlist = ctx->st->nlist; | |
| 485 | 2627 | uint32_t dim = ctx->st->dim; | |
| 486 | 2627 | uint32_t block = ctx->block; | |
| 487 | |||
| 488 | 2627 | KMeansState local = *ctx->st; | |
| 489 | 2627 | local.dist_block = ctx->dist_bufs + (size_t)thread_id * block * nlist; | |
| 490 |
2/2✓ Branch 0 taken 2192 times.
✓ Branch 1 taken 435 times.
|
2627 | if (ctx->vec_bufs) |
| 491 | 2192 | local.vec_block = ctx->vec_bufs + (size_t)thread_id * block * dim; | |
| 492 | |||
| 493 | 2627 | lloyd_assign_range( | |
| 494 | &local, | ||
| 495 | ctx->block_fn, | ||
| 496 | local.dist_block, | ||
| 497 | ctx->block, | ||
| 498 | start, | ||
| 499 | end, | ||
| 500 | 2627 | &ctx->costs[thread_id]); | |
| 501 | 2627 | } | |
| 502 | |||
| 503 | void | ||
| 504 | 12 | lloyd_assign(KMeansState *st, bool use_cblas) | |
| 505 | { | ||
| 506 | 12 | st->total_cost = 0.0f; | |
| 507 | |||
| 508 |
1/2✓ Branch 0 taken 12 times.
✗ Branch 1 not taken.
|
12 | if (st->metric == DISTANCE_L2) |
| 509 | 12 | precompute_norms_c(st); | |
| 510 | |||
| 511 |
0/2✗ Branch 0 not taken.
✗ Branch 1 not taken.
|
12 | lloyd_block_fn block_fn = lloyd_select_block_fn(st, use_cblas); |
| 512 | |||
| 513 | 12 | lloyd_assign_range( | |
| 514 | st, | ||
| 515 | block_fn, | ||
| 516 | st->dist_block, | ||
| 517 | KMEANS_BLOCK_SIZE, | ||
| 518 | 0, | ||
| 519 | st->nvecs, | ||
| 520 | &st->total_cost); | ||
| 521 | 12 | } | |
| 522 | |||
| 523 | /* ---------------------------------------------------------------- | ||
| 524 | * Fused iterative assign + update via iterate callback | ||
| 525 | * ---------------------------------------------------------------- */ | ||
| 526 | |||
| 527 | typedef struct | ||
| 528 | { | ||
| 529 | KMeansState *st; | ||
| 530 | lloyd_block_fn block_fn; | ||
| 531 | uint32_t nthreads; | ||
| 532 | uint32_t block; | ||
| 533 | |||
| 534 | /* Per-thread buffers */ | ||
| 535 | float *dist_bufs; /* [nthreads * block * nlist] */ | ||
| 536 | float *vec_bufs; /* [nthreads * block * dim] or NULL */ | ||
| 537 | float *centroid_sums; /* [nthreads * nlist * dim] */ | ||
| 538 | uint32_t *centroid_cnts; /* [nthreads * nlist] */ | ||
| 539 | float *costs; /* [nthreads] */ | ||
| 540 | |||
| 541 | /* Convergence */ | ||
| 542 | float *old_centroids; /* [nlist * dim] saved before iteration */ | ||
| 543 | float tolerance; | ||
| 544 | bool verbose; | ||
| 545 | uint32_t completed_iters; | ||
| 546 | } LloydIterCtx; | ||
| 547 | |||
| 548 | /* | ||
| 549 | * Work function: assign vectors [start, end) to nearest centroids, | ||
| 550 | * then accumulate per-thread centroid sums for the update step. | ||
| 551 | * | ||
| 552 | * Each thread gets its own dist/vec buffers and centroid accumulators | ||
| 553 | * so there is no shared mutable state during the parallel phase. | ||
| 554 | */ | ||
| 555 | static void | ||
| 556 | 18416 | lloyd_iter_work(uint32_t thread_id, uint32_t start, uint32_t end, void *arg) | |
| 557 | { | ||
| 558 | 18416 | LloydIterCtx *ctx = (LloydIterCtx *)arg; | |
| 559 | 18416 | KMeansState *st = ctx->st; | |
| 560 | 18416 | uint32_t nlist = st->nlist; | |
| 561 | 18416 | uint32_t dim = st->dim; | |
| 562 | 18416 | uint32_t block = ctx->block; | |
| 563 | |||
| 564 | /* 1. Set up thread-local KMeansState with private buffers */ | ||
| 565 | 18416 | KMeansState local = *st; | |
| 566 | 18416 | local.dist_block = ctx->dist_bufs + (size_t)thread_id * block * nlist; | |
| 567 |
2/2✓ Branch 0 taken 15129 times.
✓ Branch 1 taken 3287 times.
|
18416 | if (ctx->vec_bufs) |
| 568 | 15129 | local.vec_block = ctx->vec_bufs + (size_t)thread_id * block * dim; | |
| 569 | |||
| 570 | /* 2. Assign: compute distances and find nearest centroid */ | ||
| 571 | 18416 | lloyd_assign_range( | |
| 572 | &local, | ||
| 573 | ctx->block_fn, | ||
| 574 | local.dist_block, | ||
| 575 | block, | ||
| 576 | start, | ||
| 577 | end, | ||
| 578 | 18416 | &ctx->costs[thread_id]); | |
| 579 | |||
| 580 | /* 3. Accumulate: add each vector to its assigned centroid's sum */ | ||
| 581 | 18416 | float *my_sums = ctx->centroid_sums + (size_t)thread_id * nlist * dim; | |
| 582 | 18416 | uint32_t *my_cnts = ctx->centroid_cnts + (size_t)thread_id * nlist; | |
| 583 | |||
| 584 |
2/2✓ Branch 0 taken 12925914 times.
✓ Branch 1 taken 18416 times.
|
12944330 | for (uint32_t i = start; i < end; i++) |
| 585 | { | ||
| 586 | 12925914 | ClusterId c = st->assignments[i]; | |
| 587 |
2/2✓ Branch 0 taken 7587722 times.
✓ Branch 1 taken 5338192 times.
|
12925914 | uint32_t idx = st->indices ? st->indices[i] : i; |
| 588 | 12925914 | const float *vec = (const float *)st->vectors + (size_t)idx * dim; | |
| 589 | 12925914 | float *sum = my_sums + (size_t)c * dim; | |
| 590 |
2/2✓ Branch 0 taken 783501828 times.
✓ Branch 1 taken 12925914 times.
|
796427742 | for (uint32_t d = 0; d < dim; d++) |
| 591 | 783501828 | sum[d] += vec[d]; | |
| 592 | 12925914 | my_cnts[c]++; | |
| 593 | } | ||
| 594 | 18416 | } | |
| 595 | |||
| 596 | /* | ||
| 597 | * Reduce function: merge per-thread results, update centroids, | ||
| 598 | * and check convergence. Runs on the leader thread between | ||
| 599 | * barrier-synchronized iterations. | ||
| 600 | * | ||
| 601 | * Returns true to continue iterating, false to stop. | ||
| 602 | */ | ||
| 603 | static bool | ||
| 604 | 18416 | lloyd_iter_reduce(void *arg, uint32_t iteration) | |
| 605 | { | ||
| 606 | 18416 | LloydIterCtx *ctx = (LloydIterCtx *)arg; | |
| 607 | 18416 | KMeansState *st = ctx->st; | |
| 608 | 18416 | uint32_t nlist = st->nlist; | |
| 609 | 18416 | uint32_t dim = st->dim; | |
| 610 | 18416 | uint32_t nt = ctx->nthreads; | |
| 611 | |||
| 612 | /* 1. Sum per-thread costs into total */ | ||
| 613 | 18416 | st->total_cost = 0.0f; | |
| 614 |
2/2✓ Branch 0 taken 18416 times.
✓ Branch 1 taken 18416 times.
|
36832 | for (uint32_t t = 0; t < nt; t++) |
| 615 | 18416 | st->total_cost += ctx->costs[t]; | |
| 616 | |||
| 617 | /* 2. Merge per-thread centroid accumulators into new_centroids */ | ||
| 618 | 18416 | memset(st->new_centroids, 0, (size_t)nlist * dim * sizeof(float)); | |
| 619 | 18416 | memset(st->cluster_sizes, 0, nlist * sizeof(uint32_t)); | |
| 620 | |||
| 621 |
2/2✓ Branch 0 taken 18416 times.
✓ Branch 1 taken 18416 times.
|
36832 | for (uint32_t t = 0; t < nt; t++) |
| 622 | { | ||
| 623 | 18416 | float *sums = ctx->centroid_sums + (size_t)t * nlist * dim; | |
| 624 | 18416 | uint32_t *cnts = ctx->centroid_cnts + (size_t)t * nlist; | |
| 625 | |||
| 626 |
2/2✓ Branch 0 taken 131940 times.
✓ Branch 1 taken 18416 times.
|
150356 | for (uint32_t c = 0; c < nlist; c++) |
| 627 | { | ||
| 628 | 131940 | st->cluster_sizes[c] += cnts[c]; | |
| 629 | 131940 | float *dst = st->new_centroids + (size_t)c * dim; | |
| 630 | 131940 | float *src = sums + (size_t)c * dim; | |
| 631 |
2/2✓ Branch 0 taken 10187524 times.
✓ Branch 1 taken 131940 times.
|
10319464 | for (uint32_t d = 0; d < dim; d++) |
| 632 | 10187524 | dst[d] += src[d]; | |
| 633 | } | ||
| 634 | } | ||
| 635 | |||
| 636 | /* 3. Compute new centroids as mean of assigned vectors */ | ||
| 637 |
2/2✓ Branch 0 taken 131940 times.
✓ Branch 1 taken 18416 times.
|
150356 | for (uint32_t c = 0; c < nlist; c++) |
| 638 | { | ||
| 639 |
2/2✓ Branch 0 taken 554 times.
✓ Branch 1 taken 131386 times.
|
131940 | if (st->cluster_sizes[c] == 0) |
| 640 | 554 | continue; | |
| 641 | 131386 | float inv = 1.0f / (float)st->cluster_sizes[c]; | |
| 642 | 131386 | float *cent = st->new_centroids + (size_t)c * dim; | |
| 643 |
2/2✓ Branch 0 taken 10175528 times.
✓ Branch 1 taken 131386 times.
|
10306914 | for (uint32_t d = 0; d < dim; d++) |
| 644 | 10175528 | cent[d] *= inv; | |
| 645 | } | ||
| 646 | |||
| 647 | /* 4. Re-normalize centroids for cosine metric */ | ||
| 648 |
2/2✓ Branch 0 taken 848 times.
✓ Branch 1 taken 17568 times.
|
18416 | if (st->metric == DISTANCE_COSINE) |
| 649 | { | ||
| 650 |
2/2✓ Branch 0 taken 4712 times.
✓ Branch 1 taken 848 times.
|
5560 | for (uint32_t c = 0; c < nlist; c++) |
| 651 | { | ||
| 652 |
1/2✗ Branch 0 not taken.
✓ Branch 1 taken 4712 times.
|
4712 | if (st->cluster_sizes[c] == 0) |
| 653 | ✗ | continue; | |
| 654 | 4712 | float *cent = st->new_centroids + (size_t)c * dim; | |
| 655 | 4712 | float norm = vs_l2_norm(cent, dim); | |
| 656 |
1/2✓ Branch 0 taken 4712 times.
✗ Branch 1 not taken.
|
4712 | if (norm > 1e-10f) |
| 657 | 4712 | vec32_scale(cent, 1.0f / norm, cent, dim); | |
| 658 | } | ||
| 659 | } | ||
| 660 | |||
| 661 | /* 5. Swap old and new centroids */ | ||
| 662 | 18416 | float *tmp = st->centroids; | |
| 663 | 18416 | st->centroids = st->new_centroids; | |
| 664 | 18416 | st->new_centroids = tmp; | |
| 665 | |||
| 666 | /* 6. Check convergence: max centroid movement (squared) */ | ||
| 667 | 29838 | float shift_sq = kmeans_max_centroid_shift_between( | |
| 668 | 18416 | st->centroids, ctx->old_centroids, nlist, dim); | |
| 669 | |||
| 670 |
2/2✓ Branch 0 taken 8 times.
✓ Branch 1 taken 18408 times.
|
18416 | if (ctx->verbose) |
| 671 |
1/3✓ Branch 0 taken 8 times.
✗ Branch 1 not taken.
✗ Branch 2 not taken.
|
8 | vs_log(" iter %u: cost=%.4f, max_shift=%.6f\n", |
| 672 | iteration, | ||
| 673 | st->total_cost, | ||
| 674 | sqrtf(shift_sq)); | ||
| 675 | |||
| 676 | 18416 | ctx->completed_iters = iteration + 1; | |
| 677 | |||
| 678 | 18416 | float tol_sq = ctx->tolerance * ctx->tolerance; | |
| 679 |
2/2✓ Branch 0 taken 11141 times.
✓ Branch 1 taken 7275 times.
|
18416 | if (shift_sq < tol_sq) |
| 680 | 1012 | return false; | |
| 681 | |||
| 682 | /* 7. Prepare for next iteration: save centroids, reset accumulators */ | ||
| 683 | 16111 | memcpy(ctx->old_centroids, | |
| 684 |
2/2✓ Branch 0 taken 9521 times.
✓ Branch 1 taken 608 times.
|
16111 | st->centroids, |
| 685 | 5982 | (size_t)nlist * dim * sizeof(float)); | |
| 686 |
2/2✓ Branch 0 taken 9521 times.
✓ Branch 1 taken 608 times.
|
16111 | memset(ctx->centroid_sums, 0, (size_t)nt * nlist * dim * sizeof(float)); |
| 687 | 16111 | memset(ctx->centroid_cnts, 0, (size_t)nt * nlist * sizeof(uint32_t)); | |
| 688 | 16111 | memset(ctx->costs, 0, nt * sizeof(float)); | |
| 689 | |||
| 690 | /* 8. Precompute centroid norms for next iteration's distance calc */ | ||
| 691 |
2/2✓ Branch 0 taken 14995 times.
✓ Branch 1 taken 1116 times.
|
16111 | if (st->metric == DISTANCE_L2) |
| 692 | 14995 | precompute_norms_c(st); | |
| 693 | |||
| 694 | 5982 | return true; | |
| 695 | } | ||
| 696 | |||
| 697 | uint32_t | ||
| 698 | 2627 | lloyd_iterate(KMeansState *st, bool use_cblas, const KMeansOptions *opts) | |
| 699 | { | ||
| 700 | 1517 | (void)use_cblas; | |
| 701 | |||
| 702 | 2627 | uint32_t nt = 1; /* serial: the driver owns parallelism, not k-means */ | |
| 703 | 2627 | uint32_t nlist = st->nlist; | |
| 704 | 2627 | uint32_t dim = st->dim; | |
| 705 |
1/2✗ Branch 0 not taken.
✓ Branch 1 taken 1517 times.
|
2627 | uint32_t block = lloyd_parallel_block_size(nlist, st->nvecs); |
| 706 | |||
| 707 |
1/2✗ Branch 0 not taken.
✓ Branch 1 taken 1517 times.
|
2627 | lloyd_block_fn block_fn = lloyd_select_block_fn(st, use_cblas); |
| 708 | |||
| 709 | /* Precompute norms before first iteration */ | ||
| 710 |
2/2✓ Branch 0 taken 2529 times.
✓ Branch 1 taken 98 times.
|
2627 | if (st->metric == DISTANCE_L2) |
| 711 | 2529 | precompute_norms_c(st); | |
| 712 | |||
| 713 | 2627 | VS_MEMCTX_SCOPE(iter_ctx); | |
| 714 | 2627 | VsMemCtx old_ctx = vs_memctx_switch(iter_ctx); | |
| 715 | |||
| 716 | 2627 | float *dist_bufs = vs_alloc((size_t)nt * block * nlist * sizeof(float)); | |
| 717 | 2627 | float *vec_bufs = NULL; | |
| 718 |
2/2✓ Branch 0 taken 2192 times.
✓ Branch 1 taken 435 times.
|
2627 | if (st->vec_block != NULL) |
| 719 | 2192 | vec_bufs = vs_alloc((size_t)nt * block * dim * sizeof(float)); | |
| 720 | |||
| 721 | 2627 | float *centroid_sums = vs_alloc0((size_t)nt * nlist * dim * sizeof(float)); | |
| 722 | 2627 | uint32_t *centroid_cnts = vs_alloc0((size_t)nt * nlist * sizeof(uint32_t)); | |
| 723 | 2627 | float *costs = vs_alloc0(nt * sizeof(float)); | |
| 724 | 2627 | float *old_centroids = vs_alloc((size_t)nlist * dim * sizeof(float)); | |
| 725 | |||
| 726 | 2627 | vs_memctx_switch(old_ctx); | |
| 727 | |||
| 728 | 2627 | memcpy(old_centroids, st->centroids, (size_t)nlist * dim * sizeof(float)); | |
| 729 | |||
| 730 | 2627 | LloydIterCtx ctx = { | |
| 731 | .st = st, | ||
| 732 | .block_fn = block_fn, | ||
| 733 | .nthreads = nt, | ||
| 734 | .block = block, | ||
| 735 | .dist_bufs = dist_bufs, | ||
| 736 | .vec_bufs = vec_bufs, | ||
| 737 | .centroid_sums = centroid_sums, | ||
| 738 | .centroid_cnts = centroid_cnts, | ||
| 739 | .costs = costs, | ||
| 740 | .old_centroids = old_centroids, | ||
| 741 | 2627 | .tolerance = opts->tolerance, | |
| 742 | 2627 | .verbose = opts->verbose, | |
| 743 | .completed_iters = 0, | ||
| 744 | }; | ||
| 745 | |||
| 746 |
2/2✓ Branch 0 taken 18416 times.
✓ Branch 1 taken 322 times.
|
18738 | for (uint32_t iter = 0; iter < opts->max_iterations; iter++) |
| 747 | { | ||
| 748 | 18416 | lloyd_iter_work(0, 0, st->nvecs, &ctx); | |
| 749 |
2/2✓ Branch 1 taken 11141 times.
✓ Branch 2 taken 7275 times.
|
18416 | if (!lloyd_iter_reduce(&ctx, iter)) |
| 750 | 1012 | break; | |
| 751 | } | ||
| 752 | |||
| 753 | /* Final assignment (reduce already did the last centroid update) */ | ||
| 754 | 2627 | memset(costs, 0, nt * sizeof(float)); | |
| 755 | |||
| 756 | 2627 | LloydParCtx par_ctx = { | |
| 757 | .st = st, | ||
| 758 | .block_fn = block_fn, | ||
| 759 | .dist_bufs = dist_bufs, | ||
| 760 | .costs = costs, | ||
| 761 | .vec_bufs = vec_bufs, | ||
| 762 | .block = block, | ||
| 763 | }; | ||
| 764 | |||
| 765 | 2627 | lloyd_par_worker(0, 0, st->nvecs, &par_ctx); | |
| 766 | 2627 | st->total_cost = 0.0f; | |
| 767 |
2/2✓ Branch 0 taken 2627 times.
✓ Branch 1 taken 2627 times.
|
5254 | for (uint32_t t = 0; t < nt; t++) |
| 768 | 2627 | st->total_cost += costs[t]; | |
| 769 | |||
| 770 |
1/2✓ Branch 0 taken 1517 times.
✗ Branch 1 not taken.
|
2627 | return ctx.completed_iters; |
| 771 | } | ||
| 772 |