| 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_internal.h - Shared internal types for k-means implementations | ||
| 6 | * | ||
| 7 | * This header is shared between kmeans.c (Lloyd's / CBLAS), and variant | ||
| 8 | * algorithm files (kmeans_hamerly.c, kmeans_elkan.c, etc.). | ||
| 9 | * | ||
| 10 | * Not part of the public API — do not include from outside src/algo/. | ||
| 11 | */ | ||
| 12 | |||
| 13 | #ifndef VS_KMEANS_INTERNAL_H | ||
| 14 | #define VS_KMEANS_INTERNAL_H | ||
| 15 | |||
| 16 | #include <float.h> | ||
| 17 | #include <stdint.h> | ||
| 18 | |||
| 19 | #include "algo/kmeans.h" | ||
| 20 | #include "types/vec16.h" | ||
| 21 | #include "types/vec32.h" | ||
| 22 | |||
| 23 | /* Block size for assignment step (matches FAISS) */ | ||
| 24 | #define KMEANS_BLOCK_SIZE 4096 | ||
| 25 | |||
| 26 | /* | ||
| 27 | * Redefine VS_TARGET_CLONES locally to avoid pulling in <immintrin.h> | ||
| 28 | * from simd_utils.h, which can affect codegen for the CBLAS path. | ||
| 29 | */ | ||
| 30 | #ifndef VS_TARGET_CLONES | ||
| 31 | #if !defined(VS_SIMD_NONE) && !defined(VS_COVERAGE) && \ | ||
| 32 | __has_attribute(target_clones) && \ | ||
| 33 | (defined(__x86_64__) || defined(__i386__)) | ||
| 34 | /* Same list as simd_utils.h, kept local to avoid immintrin.h. */ | ||
| 35 | #if defined(__clang__) || __GNUC__ >= 12 | ||
| 36 | #define VS_TARGET_CLONES \ | ||
| 37 | __attribute__(( \ | ||
| 38 | target_clones("default", "arch=x86-64-v3", "arch=x86-64-v4"))) | ||
| 39 | #else | ||
| 40 | #define VS_TARGET_CLONES \ | ||
| 41 | __attribute__(( \ | ||
| 42 | target_clones("default", "arch=haswell", "arch=skylake-avx512"))) | ||
| 43 | #endif | ||
| 44 | #else | ||
| 45 | #define VS_TARGET_CLONES | ||
| 46 | #endif | ||
| 47 | #endif | ||
| 48 | |||
| 49 | /* | ||
| 50 | * Internal state for one k-means run. | ||
| 51 | * | ||
| 52 | * Shared across all algorithm variants. The core fields (vectors through | ||
| 53 | * total_cost) are used by every variant. Algorithm-specific state is | ||
| 54 | * stored externally (e.g., HamerlyState) and passed alongside. | ||
| 55 | */ | ||
| 56 | typedef struct KMeansState | ||
| 57 | { | ||
| 58 | void *memctx; /* Arena context (VsMemCtx) — used by kmeans.c */ | ||
| 59 | |||
| 60 | /* Input (not owned) */ | ||
| 61 | const void *vectors; | ||
| 62 | const uint32_t *indices; /* NULL = identity mapping [0..nvecs) */ | ||
| 63 | VecType vec_type; /* element type (f32, f16, f16c) */ | ||
| 64 | uint32_t nvecs; | ||
| 65 | uint32_t nlist; | ||
| 66 | Dimension dim; | ||
| 67 | DistanceMetric metric; | ||
| 68 | |||
| 69 | /* Working state (owned) — always float32 */ | ||
| 70 | float *centroids; /* [nlist * dim] */ | ||
| 71 | float *vec_block; /* [BLOCK_SIZE * dim] BLAS convert buf */ | ||
| 72 | ClusterId *assignments; /* [nvecs] */ | ||
| 73 | uint32_t *cluster_sizes; /* [nlist] */ | ||
| 74 | float *norms_x; /* [nvecs] precomputed ||x||^2 (L2) */ | ||
| 75 | float *norms_c; /* [nlist] precomputed ||c||^2 (L2) */ | ||
| 76 | float *dist_block; /* [KMEANS_BLOCK_SIZE * nlist] work buf */ | ||
| 77 | float *new_centroids; /* [nlist * dim] accumulator */ | ||
| 78 | float total_cost; | ||
| 79 | } KMeansState; | ||
| 80 | |||
| 81 | /* Get pointer to vector i in the input array */ | ||
| 82 | __attribute__((always_inline)) static inline const void * | ||
| 83 | 6498929 | km_get_vector(const KMeansState *st, uint32_t i, size_t elem_size) | |
| 84 | { | ||
| 85 |
21/36✗ Branch 0 not taken.
✓ Branch 1 taken 1400 times.
✗ Branch 2 not taken.
✓ Branch 3 taken 8718 times.
✗ Branch 4 not taken.
✓ Branch 5 taken 6600 times.
✗ Branch 6 not taken.
✓ Branch 7 taken 8720 times.
✗ Branch 8 not taken.
✗ Branch 9 not taken.
✗ Branch 10 not taken.
✗ Branch 11 not taken.
✗ Branch 12 not taken.
✓ Branch 13 taken 2000 times.
✗ Branch 14 not taken.
✓ Branch 15 taken 12 times.
✗ Branch 16 not taken.
✓ Branch 17 taken 2000 times.
✗ Branch 18 not taken.
✓ Branch 19 taken 32 times.
✗ Branch 20 not taken.
✓ Branch 21 taken 5600 times.
✗ Branch 22 not taken.
✓ Branch 23 taken 13600 times.
✓ Branch 24 taken 449645 times.
✓ Branch 25 taken 340473 times.
✓ Branch 26 taken 2192 times.
✓ Branch 27 taken 479 times.
✓ Branch 28 taken 470865 times.
✓ Branch 29 taken 352633 times.
✓ Branch 30 taken 10161 times.
✓ Branch 31 taken 4008 times.
✓ Branch 32 taken 3378220 times.
✓ Branch 33 taken 4128395 times.
✗ Branch 34 not taken.
✓ Branch 35 taken 22800 times.
|
9208553 | uint32_t idx = st->indices ? st->indices[i] : i; |
| 86 |
6/24✗ Branch 2 not taken.
✗ Branch 3 not taken.
✗ Branch 4 not taken.
✗ Branch 5 not taken.
✗ Branch 7 not taken.
✗ Branch 8 not taken.
✗ Branch 9 not taken.
✗ Branch 10 not taken.
✗ Branch 13 not taken.
✗ Branch 14 not taken.
✗ Branch 15 not taken.
✗ Branch 16 not taken.
✗ Branch 18 not taken.
✗ Branch 19 not taken.
✗ Branch 20 not taken.
✗ Branch 21 not taken.
✓ Branch 24 taken 553694 times.
✓ Branch 25 taken 2000 times.
✓ Branch 26 taken 14020 times.
✗ Branch 27 not taken.
✓ Branch 29 taken 5146219 times.
✓ Branch 30 taken 60000 times.
✓ Branch 31 taken 158394 times.
✗ Branch 32 not taken.
|
9208553 | return (const char *)st->vectors + (size_t)idx * st->dim * elem_size; |
| 87 | } | ||
| 88 | |||
| 89 | /* | ||
| 90 | * Max centroid movement between two centroid arrays. | ||
| 91 | * Returns the maximum squared L2 shift. | ||
| 92 | */ | ||
| 93 | float kmeans_max_centroid_shift_between( | ||
| 94 | const float *a, const float *b, uint32_t nlist, Dimension dim); | ||
| 95 | |||
| 96 | /* | ||
| 97 | * The per-worker accumulators the reduce combines, plus the scratch it | ||
| 98 | * needs to do so. | ||
| 99 | * | ||
| 100 | * Each worker assigned its own slice of the vectors to the nearest | ||
| 101 | * centroid and accumulated, per cluster, the component-wise sum of the | ||
| 102 | * vectors that landed there and how many there were. Summing those is | ||
| 103 | * what produces the new centroids: sums and counts add across workers, | ||
| 104 | * whereas per-worker means could not be combined without re-weighting | ||
| 105 | * them by their counts. | ||
| 106 | */ | ||
| 107 | typedef struct KMeansReduce | ||
| 108 | { | ||
| 109 | /* [nworkers][nlist * dim] component-wise sum of the vectors each | ||
| 110 | * worker assigned to each cluster. */ | ||
| 111 | const float *const *worker_vector_sums; | ||
| 112 | |||
| 113 | /* [nworkers][nlist] how many vectors each worker assigned to each | ||
| 114 | * cluster. Dividing a cluster's summed vector by the total is what | ||
| 115 | * yields its mean. */ | ||
| 116 | const uint32_t *const *worker_vector_counts; | ||
| 117 | |||
| 118 | /* [nworkers] each worker's summed distance from its vectors to the | ||
| 119 | * centroids they were assigned to -- the k-means objective. The | ||
| 120 | * reduce sums these into out_total_cost; the parallel build | ||
| 121 | * discards the result, so nothing reads it on this path yet. */ | ||
| 122 | const float *worker_costs; | ||
| 123 | |||
| 124 | uint32_t nworkers; | ||
| 125 | |||
| 126 | /* [nlist * dim] centroids as they stood before this iteration. | ||
| 127 | * Caller-owned: the reduce reads it to report how far the furthest | ||
| 128 | * centroid moved, which is the convergence test. */ | ||
| 129 | const float *prev_centroids; | ||
| 130 | |||
| 131 | /* [nlist] vectors per cluster once the workers are summed. Zeroed | ||
| 132 | * on entry. Caller-owned so one allocation serves every iteration, | ||
| 133 | * and so the bound on nlist sits with the caller that knows it -- | ||
| 134 | * this layer accepts any nlist. */ | ||
| 135 | uint32_t *cluster_vector_counts; | ||
| 136 | } KMeansReduce; | ||
| 137 | |||
| 138 | /* | ||
| 139 | * Merge the per-worker accumulators and update the centroids in place. | ||
| 140 | * | ||
| 141 | * This is the reduce step of parallel k-means (BSP pattern). | ||
| 142 | * Called by the leader between barrier-synchronized iterations. | ||
| 143 | * | ||
| 144 | * Only the root level of the tree reaches here. The root is one | ||
| 145 | * clustering over the whole sample, so it can only be split by data, | ||
| 146 | * which is what makes a reduce necessary. Below the root the samples | ||
| 147 | * are already partitioned, so each participant builds a whole child | ||
| 148 | * subtree alone through the serial path and nothing needs merging. | ||
| 149 | * In practice that means nlist is the tree's fan-out -- a few hundred, | ||
| 150 | * not the leaf count -- though nothing here depends on that. | ||
| 151 | * | ||
| 152 | * Outputs: | ||
| 153 | * centroids: [nlist * dim] updated positions (in-place) | ||
| 154 | * norms_c: [nlist] updated centroid norms (L2, may be NULL) | ||
| 155 | * | ||
| 156 | * Returns the maximum squared centroid shift (for convergence check). | ||
| 157 | */ | ||
| 158 | float kmeans_merge_centroids( | ||
| 159 | float *centroids, | ||
| 160 | float *norms_c, | ||
| 161 | uint32_t nlist, | ||
| 162 | Dimension dim, | ||
| 163 | DistanceMetric metric, | ||
| 164 | const KMeansReduce *reduce, | ||
| 165 | float *out_total_cost); | ||
| 166 | |||
| 167 | /* | ||
| 168 | * Scalar assign+accumulate kernel for a range of vectors. | ||
| 169 | * | ||
| 170 | * For each vector in [start, end): find nearest centroid, add | ||
| 171 | * to per-worker sums/counts. Supports optional root-assignment | ||
| 172 | * filtering for hierarchical child k-means — when filter is | ||
| 173 | * non-NULL, only vectors where filter[i] == filter_val are | ||
| 174 | * processed. | ||
| 175 | * | ||
| 176 | * Used by both standalone (via Lloyd iterate) and PG parallel | ||
| 177 | * workers (via Barrier iterate). The CBLAS path in Lloyd uses | ||
| 178 | * its own batched sgemm kernel instead. | ||
| 179 | */ | ||
| 180 | void kmeans_assign_accumulate( | ||
| 181 | const float *vectors, | ||
| 182 | const uint32_t *indices, | ||
| 183 | uint32_t start, | ||
| 184 | uint32_t end, | ||
| 185 | const float *centroids, | ||
| 186 | const float *norms_c, | ||
| 187 | uint32_t k, | ||
| 188 | Dimension dim, | ||
| 189 | DistanceMetric metric, | ||
| 190 | const uint32_t *filter, | ||
| 191 | uint32_t filter_val, | ||
| 192 | float *out_sums, | ||
| 193 | uint32_t *out_cnts, | ||
| 194 | float *out_cost); | ||
| 195 | |||
| 196 | /* | ||
| 197 | * Assignment-only kernel for a range of vectors. | ||
| 198 | * | ||
| 199 | * For each vector in [start, end): find nearest centroid and | ||
| 200 | * write the centroid index to out_assignments[i]. Unlike | ||
| 201 | * kmeans_assign_accumulate, this does not accumulate sums | ||
| 202 | * or counts — it only outputs assignments. | ||
| 203 | * | ||
| 204 | * Used for root assignment in hierarchical k-means (both | ||
| 205 | * standalone and PG paths). | ||
| 206 | */ | ||
| 207 | void kmeans_assign( | ||
| 208 | const float *vectors, | ||
| 209 | uint32_t start, | ||
| 210 | uint32_t end, | ||
| 211 | const float *centroids, | ||
| 212 | const float *norms_c, | ||
| 213 | uint32_t k, | ||
| 214 | Dimension dim, | ||
| 215 | DistanceMetric metric, | ||
| 216 | uint32_t *out_assignments); | ||
| 217 | |||
| 218 | #endif /* VS_KMEANS_INTERNAL_H */ | ||
| 219 |