| 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.h - K-means clustering with BLAS-accelerated assignment | ||
| 6 | * | ||
| 7 | * Standard Lloyd's algorithm with k-means++ initialization and optional | ||
| 8 | * CBLAS sgemm for the assignment step. The assignment converts O(N*K) | ||
| 9 | * individual distance computations into a single matrix multiply. | ||
| 10 | * | ||
| 11 | * Supports L2, inner product, and cosine distance metrics. | ||
| 12 | * For cosine: input vectors must be pre-normalized (unit length). | ||
| 13 | */ | ||
| 14 | |||
| 15 | #ifndef VS_KMEANS_H | ||
| 16 | #define VS_KMEANS_H | ||
| 17 | |||
| 18 | #include <stdbool.h> | ||
| 19 | #include <stdint.h> | ||
| 20 | |||
| 21 | #include "core/types.h" | ||
| 22 | |||
| 23 | /* | ||
| 24 | * KMeansResult - Output of k-means clustering | ||
| 25 | * | ||
| 26 | * All arrays are heap-allocated and owned by the result. Call | ||
| 27 | * vs_kmeans_result_destroy() to free. | ||
| 28 | */ | ||
| 29 | typedef struct KMeansResult | ||
| 30 | { | ||
| 31 | float *centroids; /* [nlist * dim] row-major */ | ||
| 32 | ClusterId *assignments; /* [nvecs] cluster id per vector */ | ||
| 33 | uint32_t *cluster_sizes; /* [nlist] vectors per cluster */ | ||
| 34 | uint32_t nlist; | ||
| 35 | Dimension dim; | ||
| 36 | uint32_t iterations; /* actual iterations in best run */ | ||
| 37 | float total_cost; /* sum of distances to assigned centroids */ | ||
| 38 | } KMeansResult; | ||
| 39 | |||
| 40 | /* | ||
| 41 | * KMeansAlgorithm - Assignment step algorithm | ||
| 42 | */ | ||
| 43 | typedef enum KMeansAlgorithm | ||
| 44 | { | ||
| 45 | KMEANS_ALGO_AUTO = 0, /* CBLAS if available, else Hamerly/Lloyd */ | ||
| 46 | KMEANS_ALGO_LLOYD, /* Standard Lloyd's (block dot product) */ | ||
| 47 | KMEANS_ALGO_HAMERLY, /* Hamerly's accelerated (L2 only) */ | ||
| 48 | KMEANS_ALGO_ELKAN, /* Elkan's accelerated (L2 only, O(nk) mem) */ | ||
| 49 | KMEANS_ALGO_CBLAS, /* CBLAS sgemm (requires BLAS library) */ | ||
| 50 | } KMeansAlgorithm; | ||
| 51 | |||
| 52 | /* | ||
| 53 | * KMeansOptions - Configuration for k-means | ||
| 54 | */ | ||
| 55 | typedef struct KMeansOptions | ||
| 56 | { | ||
| 57 | uint32_t max_iterations; /* default: 20 */ | ||
| 58 | float tolerance; /* convergence threshold, default: 1e-4 */ | ||
| 59 | uint64_t seed; /* random seed, default: 42 */ | ||
| 60 | uint32_t nredo; /* number of restarts, default: 1 */ | ||
| 61 | bool verbose; /* print per-iteration stats */ | ||
| 62 | KMeansAlgorithm algorithm; /* assignment algorithm, default: auto */ | ||
| 63 | const float *initial_centroids; /* skip init, use these */ | ||
| 64 | } KMeansOptions; | ||
| 65 | |||
| 66 | #define VS_KMEANS_OPTIONS_DEFAULT \ | ||
| 67 | {.max_iterations = 20, \ | ||
| 68 | .tolerance = 1e-4f, \ | ||
| 69 | .seed = 42, \ | ||
| 70 | .nredo = 1, \ | ||
| 71 | .verbose = false, \ | ||
| 72 | .algorithm = KMEANS_ALGO_AUTO} | ||
| 73 | |||
| 74 | /* | ||
| 75 | * Run k-means clustering on typed vectors. | ||
| 76 | * | ||
| 77 | * vec_type selects how input vectors are accessed (float32, float16, | ||
| 78 | * or float16 with hand-written F16C SIMD). Centroids are always | ||
| 79 | * float32. For f16 input, distance computation uses mixed-type | ||
| 80 | * operations (half × float32) to avoid bulk conversion. | ||
| 81 | * | ||
| 82 | * For DISTANCE_COSINE: input vectors MUST be pre-normalized (unit | ||
| 83 | * length). The function normalizes centroids after each update step | ||
| 84 | * internally. The caller should normalize vectors during sampling. | ||
| 85 | * | ||
| 86 | * Runs nredo independent attempts and returns the best (lowest cost). | ||
| 87 | * | ||
| 88 | * Parameters: | ||
| 89 | * vectors: [nvecs * dim] row-major input vectors (or full array | ||
| 90 | * when indices is non-NULL) | ||
| 91 | * indices: optional index array [nvecs] mapping logical positions | ||
| 92 | * to rows in vectors. NULL = contiguous [0..nvecs). | ||
| 93 | * vec_type: element type (VS_VEC_F32, VS_VEC_F16, VS_VEC_F16C) | ||
| 94 | * nvecs: number of vectors (or indices) to cluster | ||
| 95 | * dim: vector dimension | ||
| 96 | * nlist: number of clusters (K) | ||
| 97 | * metric: distance metric (L2, IP, or cosine) | ||
| 98 | * options: configuration (NULL for defaults) | ||
| 99 | * | ||
| 100 | * Returns allocated result on success, NULL on failure. | ||
| 101 | */ | ||
| 102 | KMeansResult *vs_kmeans( | ||
| 103 | const void *vectors, | ||
| 104 | const uint32_t *indices, | ||
| 105 | VecType vec_type, | ||
| 106 | uint32_t nvecs, | ||
| 107 | Dimension dim, | ||
| 108 | uint32_t nlist, | ||
| 109 | DistanceMetric metric, | ||
| 110 | const KMeansOptions *options); | ||
| 111 | |||
| 112 | /* | ||
| 113 | * Convenience wrapper for float32 vectors. | ||
| 114 | */ | ||
| 115 | static inline KMeansResult * | ||
| 116 | 242 | vs_kmeans_f32( | |
| 117 | const float *vectors, | ||
| 118 | uint32_t nvecs, | ||
| 119 | Dimension dim, | ||
| 120 | uint32_t nlist, | ||
| 121 | DistanceMetric metric, | ||
| 122 | const KMeansOptions *options) | ||
| 123 | { | ||
| 124 | 242 | return vs_kmeans( | |
| 125 | vectors, NULL, VS_VEC_F32, nvecs, dim, nlist, metric, options); | ||
| 126 | } | ||
| 127 | |||
| 128 | /* | ||
| 129 | * Free a k-means result and all owned arrays. | ||
| 130 | */ | ||
| 131 | void vs_kmeans_result_destroy(KMeansResult *result); | ||
| 132 | |||
| 133 | /* | ||
| 134 | * Get human-readable algorithm name. | ||
| 135 | */ | ||
| 136 | const char *vs_kmeans_algo_name(KMeansAlgorithm algo); | ||
| 137 | |||
| 138 | /* | ||
| 139 | * CBLAS runtime control (legacy, used by PostgreSQL extension). | ||
| 140 | * For benchmarks, prefer KMeansOptions.algorithm instead. | ||
| 141 | */ | ||
| 142 | void vs_kmeans_set_use_cblas(bool use_cblas); | ||
| 143 | bool vs_kmeans_get_use_cblas(void); | ||
| 144 | const char *vs_kmeans_impl_name(void); | ||
| 145 | |||
| 146 | /* | ||
| 147 | * Check if BLAS is configured for single-threaded operation. | ||
| 148 | * | ||
| 149 | * Returns true if OMP_NUM_THREADS=1 is set in the environment. | ||
| 150 | * BLAS libraries (BLIS, OpenBLAS, MKL) read this at load time to | ||
| 151 | * size their thread pools. Must be set before process start. | ||
| 152 | * | ||
| 153 | * PostgreSQL: set in systemd unit or pg_ctl environment. | ||
| 154 | * CLI: prefix command with OMP_NUM_THREADS=1. | ||
| 155 | */ | ||
| 156 | bool vs_cblas_is_single_threaded(void); | ||
| 157 | |||
| 158 | /* | ||
| 159 | * Pin BLAS to single-threaded via runtime APIs (BLIS, OpenBLAS). | ||
| 160 | * Call once at startup. Uses weak symbols — safe when the linked | ||
| 161 | * BLAS doesn't provide the function. | ||
| 162 | */ | ||
| 163 | void vs_cblas_pin_single_thread(void); | ||
| 164 | |||
| 165 | #endif /* VS_KMEANS_H */ | ||
| 166 |