usearch 2.15.2 → 2.16.0

package/include/usearch/index_plugins.hpp

@@ -4,6 +4,7 @@
 #include <stdlib.h> // `aligned_alloc`
 
 #include <atomic>  // `std::atomic`
+#include <chrono>  // `std::chrono`
 #include <cstring> // `std::strncmp`
 #include <thread>  // `std::thread`
 
@@ -42,20 +43,38 @@
 
 #if USEARCH_USE_SIMSIMD
 // Propagate the `f16` settings
-#if !defined(SIMSIMD_NATIVE_F16)
-#define SIMSIMD_NATIVE_F16 !USEARCH_USE_FP16LIB
+#if defined(USEARCH_CAN_COMPILE_FP16) || defined(USEARCH_CAN_COMPILE_FLOAT16)
+#if USEARCH_CAN_COMPILE_FP16 || USEARCH_CAN_COMPILE_FLOAT16
+#define SIMSIMD_NATIVE_F16 1
+#else
+#define SIMSIMD_NATIVE_F16 0
+#endif
+#endif
+// Propagate the `bf16` settings
+#if defined(USEARCH_CAN_COMPILE_BF16) || defined(USEARCH_CAN_COMPILE_BFLOAT16)
+#if USEARCH_CAN_COMPILE_BF16 || USEARCH_CAN_COMPILE_BFLOAT16
+#define SIMSIMD_NATIVE_BF16 1
+#else
+#define SIMSIMD_NATIVE_BF16 0
+#endif
 #endif
-
-// Overwrite the dynamic dispatch settings
-#undef SIMSIMD_DYNAMIC_DISPATCH
-#define SIMSIMD_DYNAMIC_DISPATCH 0
 // No problem, if some of the functions are unused or undefined
 #pragma GCC diagnostic push
+#pragma GCC diagnostic ignored "-Wall"
+#pragma GCC diagnostic ignored "-Wunused"
 #pragma GCC diagnostic ignored "-Wunused-function"
+#pragma GCC diagnostic ignored "-Wunused-parameter"
+#pragma GCC diagnostic ignored "-Wunused-variable"
+#pragma GCC diagnostic ignored "-Wunused-but-set-variable"
+#if defined(_MSC_VER)
 #pragma warning(push)
-#pragma warning(disable : 4101)
+#pragma warning(disable : 4101) // "Unused variables"
+#pragma warning(disable : 4068) // "Unknown pragmas", when MSVC tries to read GCC pragmas
+#endif                          // _MSC_VER
 #include <simsimd/simsimd.h>
+#ifdef _MSC_VER
 #pragma warning(pop)
+#endif // _MSC_VER
 #pragma GCC diagnostic pop
 #endif
 
@@ -380,9 +399,10 @@ inline float f16_to_f32(std::uint16_t u16) noexcept {
 #if USEARCH_USE_FP16LIB
     return fp16_ieee_to_fp32_value(u16);
 #elif USEARCH_USE_SIMSIMD
-    return simsimd_uncompress_f16(u16);
+    return simsimd_f16_to_f32((simsimd_f16_t const*)&u16);
 #else
-    f16_native_t f16;
+#warning "It's recommended to use SimSIMD and fp16lib for half-precision numerics"
+    _Float16 f16;
     std::memcpy(&f16, &u16, sizeof(std::uint16_t));
     return float(f16);
 #endif
@@ -395,9 +415,12 @@ inline std::uint16_t f32_to_f16(float f32) noexcept {
 #if USEARCH_USE_FP16LIB
     return fp16_ieee_from_fp32_value(f32);
 #elif USEARCH_USE_SIMSIMD
-    return simsimd_compress_f16(f32);
+    std::uint16_t result;
+    simsimd_f32_to_f16(f32, (simsimd_f16_t*)&result);
+    return result;
 #else
-    f16_native_t f16 = f16_native_t(f32);
+#warning "It's recommended to use SimSIMD and fp16lib for half-precision numerics"
+    _Float16 f16 = _Float16(f32);
     std::uint16_t u16;
     std::memcpy(&u16, &f16, sizeof(std::uint16_t));
     return u16;
@@ -410,15 +433,14 @@ inline std::uint16_t f32_to_f16(float f32) noexcept {
  */
 inline float bf16_to_f32(std::uint16_t u16) noexcept {
 #if USEARCH_USE_SIMSIMD
-    return simsimd_uncompress_bf16(u16);
+    return simsimd_bf16_to_f32((simsimd_bf16_t const*)&u16);
 #else
     union float_or_unsigned_int_t {
         float f;
         unsigned int i;
-    };
-    union float_or_unsigned_int_t result_union;
-    result_union.i = u16 << 16; // Zero extends the mantissa
-    return result_union.f;
+    } conv;
+    conv.i = u16 << 16; // Zero extends the mantissa
+    return conv.f;
 #endif
 }
 
@@ -428,17 +450,18 @@ inline float bf16_to_f32(std::uint16_t u16) noexcept {
  */
 inline std::uint16_t f32_to_bf16(float f32) noexcept {
 #if USEARCH_USE_SIMSIMD
-    return simsimd_compress_bf16(f32);
+    std::uint16_t result;
+    simsimd_f32_to_bf16(f32, (simsimd_bf16_t*)&result);
+    return result;
 #else
     union float_or_unsigned_int_t {
         float f;
         unsigned int i;
-    };
-    union float_or_unsigned_int_t value;
-    value.f = f32;
-    value.i >>= 16;
-    value.i &= 0xFFFF;
-    return (unsigned short)value.i;
+    } conv;
+    conv.f = f32;
+    conv.i >>= 16;
+    conv.i &= 0xFFFF;
+    return (unsigned short)conv.i;
 #endif
 }
 
@@ -1183,6 +1206,71 @@ template <> struct cast_gt<f64_t, b1x8_t> : public cast_to_b1x8_gt<f64_t> {};
 template <> struct cast_gt<b1x8_t, i8_t> : public cast_from_b1x8_gt<i8_t> {};
 template <> struct cast_gt<i8_t, b1x8_t> : public cast_to_b1x8_gt<i8_t> {};
 
+/**
+ *  @brief  Type-punned array casting function.
+ *          Arguments: input buffer, bytes in input buffer, output buffer.
+ *          Returns `true` if the casting was performed successfully, `false` otherwise.
+ */
+using cast_punned_t = bool (*)(byte_t const*, std::size_t, byte_t*);
+
+/**
+ *  @brief  A collection of casting functions for typical vector types.
+ *          Covers to/from conversions for boolean, integer, half-precision,
+ *          single-precision, and double-precision scalars.
+ */
+struct casts_punned_t {
+    struct group_t {
+        cast_punned_t b1x8{};
+        cast_punned_t i8{};
+        cast_punned_t f16{};
+        cast_punned_t f32{};
+        cast_punned_t f64{};
+
+        cast_punned_t operator[](scalar_kind_t scalar_kind) const noexcept {
+            switch (scalar_kind) {
+            case scalar_kind_t::f64_k: return f64;
+            case scalar_kind_t::f32_k: return f32;
+            case scalar_kind_t::f16_k: return f16;
+            case scalar_kind_t::bf16_k: return f16;
+            case scalar_kind_t::i8_k: return i8;
+            case scalar_kind_t::b1x8_k: return b1x8;
+            default: return nullptr;
+            }
+        }
+
+    } from, to;
+
+    template <typename scalar_at> static casts_punned_t make() noexcept {
+        casts_punned_t result;
+
+        result.from.b1x8 = &cast_gt<b1x8_t, scalar_at>::try_;
+        result.from.i8 = &cast_gt<i8_t, scalar_at>::try_;
+        result.from.f16 = &cast_gt<f16_t, scalar_at>::try_;
+        result.from.f32 = &cast_gt<f32_t, scalar_at>::try_;
+        result.from.f64 = &cast_gt<f64_t, scalar_at>::try_;
+
+        result.to.b1x8 = &cast_gt<scalar_at, b1x8_t>::try_;
+        result.to.i8 = &cast_gt<scalar_at, i8_t>::try_;
+        result.to.f16 = &cast_gt<scalar_at, f16_t>::try_;
+        result.to.f32 = &cast_gt<scalar_at, f32_t>::try_;
+        result.to.f64 = &cast_gt<scalar_at, f64_t>::try_;
+
+        return result;
+    }
+
+    static casts_punned_t make(scalar_kind_t scalar_kind) noexcept {
+        switch (scalar_kind) {
+        case scalar_kind_t::f64_k: return casts_punned_t::make<f64_t>();
+        case scalar_kind_t::f32_k: return casts_punned_t::make<f32_t>();
+        case scalar_kind_t::f16_k: return casts_punned_t::make<f16_t>();
+        case scalar_kind_t::bf16_k: return casts_punned_t::make<bf16_t>();
+        case scalar_kind_t::i8_k: return casts_punned_t::make<i8_t>();
+        case scalar_kind_t::b1x8_k: return casts_punned_t::make<b1x8_t>();
+        default: return {};
+        }
+    }
+};
+
 /*  Don't complain if the vectorization of the inner loops fails:
  *
  *  > warning: loop not vectorized: the optimizer was unable to perform the requested transformation;
@@ -1910,32 +1998,33 @@ class metric_punned_t {
  *  @brief  View over a potentially-strided memory buffer, containing a row-major matrix.
  */
 template <typename scalar_at> //
-class vectors_view_gt {
+class matrix_slice_gt {
     using scalar_t = scalar_at;
+    using byte_addressable_t = typename std::conditional<std::is_const<scalar_t>::value, byte_t const, byte_t>::type;
 
-    scalar_t const* begin_{};
+    scalar_t* begin_{};
     std::size_t dimensions_{};
     std::size_t count_{};
     std::size_t stride_bytes_{};
 
   public:
-    vectors_view_gt() noexcept = default;
-    vectors_view_gt(vectors_view_gt const&) noexcept = default;
-    vectors_view_gt& operator=(vectors_view_gt const&) noexcept = default;
+    matrix_slice_gt() noexcept = default;
+    matrix_slice_gt(matrix_slice_gt const&) noexcept = default;
+    matrix_slice_gt& operator=(matrix_slice_gt const&) noexcept = default;
 
-    vectors_view_gt(scalar_t const* begin, std::size_t dimensions, std::size_t count = 1) noexcept
-        : vectors_view_gt(begin, dimensions, count, dimensions * sizeof(scalar_at)) {}
+    matrix_slice_gt(scalar_t* begin, std::size_t dimensions, std::size_t count = 1) noexcept
+        : matrix_slice_gt(begin, dimensions, count, dimensions * sizeof(scalar_at)) {}
 
-    vectors_view_gt(scalar_t const* begin, std::size_t dimensions, std::size_t count, std::size_t stride_bytes) noexcept
+    matrix_slice_gt(scalar_t* begin, std::size_t dimensions, std::size_t count, std::size_t stride_bytes) noexcept
         : begin_(begin), dimensions_(dimensions), count_(count), stride_bytes_(stride_bytes) {}
 
     explicit operator bool() const noexcept { return begin_; }
     std::size_t size() const noexcept { return count_; }
     std::size_t dimensions() const noexcept { return dimensions_; }
-    std::size_t stride() const noexcept { return stride_bytes_; }
-    scalar_t const* data() const noexcept { return begin_; }
-    scalar_t const* at(std::size_t i) const noexcept {
-        return reinterpret_cast<scalar_t const*>(reinterpret_cast<byte_t const*>(begin_) + i * stride_bytes_);
+    std::size_t stride_bytes() const noexcept { return stride_bytes_; }
+    scalar_t* data() const noexcept { return begin_; }
+    scalar_t* at(std::size_t i) const noexcept {
+        return reinterpret_cast<scalar_t*>(reinterpret_cast<byte_addressable_t*>(begin_) + i * stride_bytes_);
     }
 };
 
@@ -1944,7 +2033,7 @@ struct exact_offset_and_distance_t {
     f32_t distance;
 };
 
-using exact_search_results_t = vectors_view_gt<exact_offset_and_distance_t>;
+using exact_search_results_t = matrix_slice_gt<exact_offset_and_distance_t const>;
 
 /**
  *  @brief  Helper-structure for exact search operations.
@@ -1965,14 +2054,14 @@ class exact_search_t {
 
   public:
     template <typename scalar_at, typename executor_at = dummy_executor_t, typename progress_at = dummy_progress_t>
-    exact_search_results_t operator()(                                          //
-        vectors_view_gt<scalar_at> dataset, vectors_view_gt<scalar_at> queries, //
-        std::size_t wanted, metric_punned_t const& metric,                      //
+    exact_search_results_t operator()(                                                      //
+        matrix_slice_gt<scalar_at const> dataset, matrix_slice_gt<scalar_at const> queries, //
+        std::size_t wanted, metric_punned_t const& metric,                                  //
         executor_at&& executor = executor_at{}, progress_at&& progress = progress_at{}) {
-        return operator()(                                                                     //
-            metric,                                                                            //
-            reinterpret_cast<byte_t const*>(dataset.data()), dataset.size(), dataset.stride(), //
-            reinterpret_cast<byte_t const*>(queries.data()), queries.size(), queries.stride(), //
+        return operator()(                                                                           //
+            metric,                                                                                  //
+            reinterpret_cast<byte_t const*>(dataset.data()), dataset.size(), dataset.stride_bytes(), //
+            reinterpret_cast<byte_t const*>(queries.data()), queries.size(), queries.stride_bytes(), //
             wanted, executor, progress);
     }
 
@@ -2049,6 +2138,342 @@ class exact_search_t {
     }
 };
 
+struct kmeans_clustering_result_t {
+    error_t error{};
+    std::size_t computed_distances{};
+    /// @brief The number of iterations the algorithm took to converge.
+    std::size_t iterations{};
+    /// @brief The number of points that changed clusters in the last iteration.
+    std::size_t last_iteration_points_shifted{};
+    /// @brief The inertia of the last iteration (sum of squared distances to centroids).
+    f64_t last_iteration_inertia{};
+    /// @brief The total elapsed runtime of the algorithm in seconds.
+    f64_t runtime_seconds{};
+    /// @brief The total distance between the points and their assigned centroids.
+    f64_t aggregate_distance{};
+
+    explicit operator bool() const noexcept { return !error; }
+    kmeans_clustering_result_t failed(error_t message) noexcept {
+        error = std::move(message);
+        return std::move(*this);
+    }
+};
+
+/**
+ *  @brief  Helper-class for K-Means clustering of dense vectors.
+ *          Doesn't require constructing the index, but benefits from mixed-precision logic.
+ *          ! Doesn't guarantee that the clusters are balanced in size.
+ *
+ *  The algorithm is as follows:
+ *  - Initialization: Select K initial centroids (randomly or with a heuristic).
+ *  - Assignment: Assign each data point to the nearest centroid based on the Euclidean distance.
+ *  - Update: Recalculate the centroids as the mean of all points assigned to each centroid.
+ *  - Repeat: Repeat the assignment and update steps until the centroids no longer change significantly
+ *            or an early-exit condition is met.
+ */
+template <typename allocator_at = std::allocator<char>> class kmeans_clustering_gt {
+  public:
+    using distance_t = distance_punned_t;
+
+    metric_kind_t metric_kind{metric_kind_t::l2sq_k};
+    scalar_kind_t quantization_kind{scalar_kind_t::bf16_k};
+
+    static constexpr std::size_t max_iterations_default_k = 300;
+    static constexpr f64_t inertia_threshold_default_k = 1e-4;
+    static constexpr f64_t max_seconds_default_k = 60.0;
+    static constexpr f64_t min_shifts_default_k = 0.01;
+
+    /// @brief Early-exit parameter - the maximum number of iterations to perform.
+    std::size_t max_iterations{max_iterations_default_k};
+    /// @brief Early-exit parameter - the threshold for the final inertia to terminate early.
+    f64_t inertia_threshold{inertia_threshold_default_k};
+    /// @brief Early-exit parameter - the maximum runtime allowed in seconds.
+    f64_t max_seconds{max_seconds_default_k};
+    /// @brief Early-exit parameter - the minimum share of points that must change clusters per iteration.
+    f64_t min_shifts{min_shifts_default_k};
+    /// @brief The random seed to use for centroid initialization.
+    std::uint64_t seed{0};
+
+    kmeans_clustering_gt(std::uint64_t seed) noexcept : seed(seed) {}
+    kmeans_clustering_gt() noexcept(false) {
+        std::random_device random_device;
+        seed = random_device();
+    }
+
+    kmeans_clustering_gt(kmeans_clustering_gt const&) = default;
+    kmeans_clustering_gt& operator=(kmeans_clustering_gt const&) = default;
+
+    template <typename scalar_at, typename executor_at = dummy_executor_t, typename progress_at = dummy_progress_t>
+    kmeans_clustering_result_t operator()( //
+        matrix_slice_gt<scalar_at const> points, matrix_slice_gt<scalar_at> centroids,
+        span_gt<std::size_t> point_to_centroid_index, span_gt<distance_t> point_to_centroid_distance, //
+        executor_at&& executor = executor_at{}, progress_at&& progress = progress_at{}) {
+        return operator()(                                                                        //
+            reinterpret_cast<byte_t const*>(points.data()), points.size(), points.stride_bytes(), //
+            reinterpret_cast<byte_t*>(centroids.data()), centroids.size(), centroids.stride_bytes(),
+            point_to_centroid_index.data(), point_to_centroid_distance.data(), //
+            scalar_kind<scalar_at>(), points.dimensions(), executor, progress);
+    }
+
+    template <typename executor_at = dummy_executor_t, typename progress_at = dummy_progress_t>
+    kmeans_clustering_result_t operator()(                                                       //
+        byte_t const* points_data, std::size_t points_count, std::size_t points_stride_bytes,    //
+        byte_t* centroids_data, std::size_t wanted_clusters, std::size_t centroids_stride_bytes, //
+        std::size_t* point_to_centroid_index, distance_t* point_to_centroid_distance,            //
+        scalar_kind_t original_scalar_kind, std::size_t dimensions, executor_at&& executor = executor_at{},
+        progress_at&& progress = progress_at{}) {
+
+        // Perform sanity checks for algorithm settings.
+        kmeans_clustering_result_t result;
+        if (max_iterations < 1)
+            return result.failed("The number of iterations must be at least 1");
+
+        // Perform sanity checks for input arguments.
+        if (wanted_clusters < 2)
+            return result.failed("The number of clusters must be at least 2");
+        if (wanted_clusters >= points_count)
+            return result.failed("The number of clusters must be less than the number of vectors");
+
+        metric_punned_t metric = metric_punned_t::builtin(dimensions, metric_kind, quantization_kind);
+        if (!metric)
+            return result.failed("Unsupported metric or scalar kind");
+
+        // Let's allocate memory for the centroids coordinates and make sure it's
+        // rows are aligned to cache lines to avoid false sharing.
+        buffer_gt<distance_t, aligned_allocator_gt<distance_t, 64>> point_to_centroid_distance_buffer(points_count);
+        buffer_gt<std::size_t, aligned_allocator_gt<std::size_t, 64>> point_to_centroid_index_buffer(points_count);
+        buffer_gt<std::atomic<std::size_t>, aligned_allocator_gt<std::atomic<std::size_t>, 64>> cluster_sizes_buffer(
+            wanted_clusters);
+
+        // For a mixed precision computation, we keep the centroids represented in two forms -
+        // double precision and quantized the same way as in the index, to avoid paying conversion penalties.
+        // Double precision is needed to avoid accumulating errors when aggregating too many entries.
+        std::size_t const bytes_per_vector_original =
+            divide_round_up<CHAR_BIT>(dimensions * bits_per_scalar(original_scalar_kind));
+        std::size_t const bytes_per_vector_quantized = metric.bytes_per_vector();
+        std::size_t const stride_per_vector_quantized = divide_round_up<64>(bytes_per_vector_quantized) * 64;
+        buffer_gt<byte_t, aligned_allocator_gt<byte_t, 64>> points_quantized_buffer( //
+            points_count * stride_per_vector_quantized);
+        buffer_gt<byte_t, aligned_allocator_gt<byte_t, 64>> centroids_quantized_buffer( //
+            wanted_clusters * stride_per_vector_quantized);
+
+        // When aggregating centroids, we want to parallelize the operation and need more memory.
+        // For every thread we keep two double-precision vectors. One is the up-casting output buffer for quantized
+        // coordinates, and the other is the temporary buffer for the partial sums of the double-precision coordinates.
+        // The ordering:
+        //
+        //      - thread 0: [centroid 0, centroid 1, centroid 2, centroid 3, ...]
+        //      - thread 1: [centroid 0, centroid 1, centroid 2, centroid 3, ...]
+        //      - thread 2: [centroid 0, centroid 1, centroid 2, centroid 3, ...]
+        //
+        std::size_t const thread_count = executor.size();
+        buffer_gt<f64_t, aligned_allocator_gt<f64_t, 64>> centroids_precise_buffer( //
+            wanted_clusters * dimensions * thread_count);
+        buffer_gt<f64_t, aligned_allocator_gt<f64_t, 64>> points_precise_buffer( //
+            wanted_clusters * dimensions * thread_count);
+
+        // Check if all memory allocations were successful.
+        if (!centroids_precise_buffer || !points_precise_buffer || !centroids_quantized_buffer ||
+            !point_to_centroid_index_buffer || !cluster_sizes_buffer || !point_to_centroid_distance_buffer ||
+            !points_quantized_buffer)
+            return result.failed("No memory for result outputs!");
+
+        std::fill_n(point_to_centroid_index_buffer.data(), points_count, wanted_clusters);
+        std::fill_n(point_to_centroid_distance_buffer.data(), points_count, std::numeric_limits<distance_t>::max());
+
+        // Initialize the casting kernel for quantization and export.
+        casts_punned_t casts = casts_punned_t::make(quantization_kind);
+        cast_punned_t const& compress_points = casts.from[original_scalar_kind];
+        cast_punned_t const& decompress_points = casts.to[original_scalar_kind];
+        cast_punned_t const& compress_precise = casts.from.f64;
+        cast_punned_t const& decompress_precise = casts.to.f64;
+        for (std::size_t i = 0; i < points_count; i++) {
+            byte_t const* vector = points_data + i * points_stride_bytes;
+            byte_t* quantized = points_quantized_buffer.data() + i * stride_per_vector_quantized;
+            if (!compress_points(vector, dimensions, quantized))
+                std::memcpy(quantized, vector, bytes_per_vector_original);
+        }
+
+        // Initialize centroids with random points vectors.
+        std::mt19937_64 random_engine;
+        random_engine.seed(seed);
+        for (std::size_t i = 0; i < wanted_clusters; i++) {
+            // Generate the random index of the points vector,
+            // that is unique and not already used as a centroid.
+            std::size_t random_index;
+            do {
+                random_index = random_engine() % points_count;
+                bool is_unique = true;
+                for (std::size_t j = 0; j < i; j++) {
+                    if (point_to_centroid_index_buffer[j] == random_index) {
+                        is_unique = false;
+                        break;
+                    }
+                }
+                if (is_unique)
+                    break;
+            } while (true);
+
+            // Copy the vector to the centroid and quantize it.
+            byte_t const* quantized_point = points_quantized_buffer.data() + random_index * stride_per_vector_quantized;
+            byte_t* quantized_centroid = centroids_quantized_buffer.data() + i * stride_per_vector_quantized;
+            std::memcpy(quantized_centroid, quantized_point, bytes_per_vector_quantized);
+            point_to_centroid_index_buffer[random_index] = i;
+            point_to_centroid_distance_buffer[random_index] = 0;
+        }
+
+        auto start_time = std::chrono::high_resolution_clock::now();
+        std::size_t iterations = 0;
+        std::size_t const min_points_shifted_per_iteration = static_cast<std::size_t>(min_shifts * points_count);
+        f64_t last_aggregate_distance = std::numeric_limits<f64_t>::max();
+
+        while (iterations < max_iterations) {
+            iterations++;
+
+            // For every point, find the closest centroid.
+            std::atomic<std::size_t> points_shifted{0};
+            executor.dynamic(points_count, [&](std::size_t thread_idx, std::size_t points_idx) {
+                byte_t const* quantized_point =
+                    points_quantized_buffer.data() + points_idx * stride_per_vector_quantized;
+                byte_t const* quantized_centroids = centroids_quantized_buffer.data();
+                distance_t closest_distance_local = std::numeric_limits<distance_t>::max();
+                std::size_t closest_idx_local = 0;
+                for (std::size_t centroid_idx = 0; centroid_idx < wanted_clusters; centroid_idx++) {
+                    byte_t const* quantized_centroid = quantized_centroids + centroid_idx * stride_per_vector_quantized;
+                    distance_t distance = metric(quantized_point, quantized_centroid);
+                    if (distance < closest_distance_local) {
+                        closest_distance_local = distance;
+                        closest_idx_local = centroid_idx;
+                    }
+                }
+
+                distance_t& closest_distance_ref = point_to_centroid_distance_buffer[points_idx];
+                std::size_t& closest_idx_ref = point_to_centroid_index_buffer[points_idx];
+                if (closest_idx_local != closest_idx_ref) {
+                    closest_idx_ref = closest_idx_local;
+                    points_shifted.fetch_add(1, std::memory_order_relaxed);
+                }
+
+                closest_distance_ref = closest_distance_local;
+                return true;
+            });
+
+            f64_t aggregate_distance = 0.0;
+            for (std::size_t i = 0; i < points_count; i++)
+                aggregate_distance += point_to_centroid_distance_buffer[i];
+            f64_t aggregate_distance_change =
+                std::abs(aggregate_distance - last_aggregate_distance) / last_aggregate_distance;
+
+            auto current_time = std::chrono::high_resolution_clock::now();
+            std::chrono::duration<f64_t> elapsed_time = current_time - start_time;
+            result.runtime_seconds = elapsed_time.count();
+            result.last_iteration_inertia = aggregate_distance_change;
+            result.last_iteration_points_shifted = points_shifted.load(std::memory_order_relaxed);
+
+            // Check for early-exit conditions
+            if (last_aggregate_distance != 0.0 && inertia_threshold != 0.0)
+                if (aggregate_distance_change <= inertia_threshold)
+                    break;
+            if (min_points_shifted_per_iteration != 0 || result.last_iteration_points_shifted == 0)
+                if (result.last_iteration_points_shifted <= min_points_shifted_per_iteration)
+                    break;
+            if (max_seconds != 0)
+                if (result.runtime_seconds >= max_seconds)
+                    break;
+
+            // For every centroid, recalculate the mean of all points assigned to it.
+            // That part is problematic to parallelize on many-core-systems, because of the contention.
+            // Alternatively, a tree-like approach can be used, where every core accumulates it's own partial sums.
+            // And those are later aggregated by a single thread.
+            std::memset(centroids_precise_buffer.data(), 0,
+                        wanted_clusters * dimensions * thread_count * sizeof(f64_t));
+            std::memset(reinterpret_cast<byte_t*>(cluster_sizes_buffer.data()), 0,
+                        wanted_clusters * sizeof(std::atomic<std::size_t>));
+            executor.dynamic(points_count, [&](std::size_t thread_idx, std::size_t points_idx) {
+                std::size_t centroid_idx = point_to_centroid_index_buffer[points_idx];
+                byte_t const* quantized_point =
+                    points_quantized_buffer.data() + points_idx * stride_per_vector_quantized;
+                f64_t* centroid_precise = centroids_precise_buffer.data() + wanted_clusters * dimensions * thread_idx +
+                                          centroid_idx * dimensions;
+
+                // Upcast the points point into a buffer of double-precision floats.
+                f64_t* point_precise = points_precise_buffer.data() + wanted_clusters * dimensions * thread_idx +
+                                       centroid_idx * dimensions;
+                if (!decompress_precise(quantized_point, dimensions, reinterpret_cast<byte_t*>(point_precise)))
+                    std::memcpy(reinterpret_cast<byte_t*>(point_precise), quantized_point, bytes_per_vector_quantized);
+
+                // Now add the vector from the points into the centroid partial sum.
+                for (std::size_t i = 0; i < dimensions; i++)
+                    centroid_precise[i] += point_precise[i];
+
+                cluster_sizes_buffer[centroid_idx].fetch_add(1, std::memory_order_relaxed);
+                return true;
+            });
+
+            // Aggregate the partial sums into the final centroids - storing them in the high-precision
+            // buffer of the first thread. Normalization procedure is different for different metrics.
+            for (std::size_t centroid_idx = 0; centroid_idx < wanted_clusters; centroid_idx++) {
+                f64_t* centroid_precise_aggregated = centroids_precise_buffer.data() + centroid_idx * dimensions;
+                for (std::size_t thread_idx = 1; thread_idx < thread_count; thread_idx++) {
+                    f64_t* centroid_precise = centroids_precise_buffer.data() +
+                                              wanted_clusters * dimensions * thread_idx + centroid_idx * dimensions;
+                    for (std::size_t i = 0; i < dimensions; i++)
+                        centroid_precise_aggregated[i] += centroid_precise[i];
+                }
+
+                // Normalize based on the metric kind
+                if (metric_kind == metric_kind_t::l2sq_k) {
+                    // Normalize for Euclidean distance (L2)
+                    std::size_t cluster_size = cluster_sizes_buffer[centroid_idx].load(std::memory_order_relaxed);
+                    if (cluster_size > 0)
+                        for (std::size_t i = 0; i < dimensions; i++)
+                            centroid_precise_aggregated[i] /= static_cast<f64_t>(cluster_size);
+
+                } else if (metric_kind == metric_kind_t::cos_k) {
+                    // Normalize for Cosine distance
+                    f64_t norm = 0.0;
+                    for (std::size_t i = 0; i < dimensions; i++)
+                        norm += centroid_precise_aggregated[i] * centroid_precise_aggregated[i];
+                    norm = std::sqrt(norm);
+                    if (norm > 0.0)
+                        for (std::size_t i = 0; i < dimensions; i++)
+                            centroid_precise_aggregated[i] /= norm;
+                }
+
+                // Quantize the centroid after normalization for further iterations
+                byte_t* centroid_quantized =
+                    centroids_quantized_buffer.data() + centroid_idx * stride_per_vector_quantized;
+                if (!compress_precise(reinterpret_cast<byte_t*>(centroid_precise_aggregated), dimensions,
+                                      centroid_quantized))
+                    std::memcpy(centroid_quantized, reinterpret_cast<byte_t*>(centroid_precise_aggregated),
+                                bytes_per_vector_quantized);
+            }
+        }
+
+        // Export stats.
+        result.iterations = iterations;
+        result.computed_distances = points_count * wanted_clusters * iterations;
+        result.aggregate_distance = 0;
+        for (distance_t distance : point_to_centroid_distance_buffer)
+            result.aggregate_distance += distance;
+
+        // We've finished all the iterations, now we can export the centroids back to the original precision.
+        std::memcpy(point_to_centroid_index, point_to_centroid_index_buffer.data(), points_count * sizeof(std::size_t));
+        std::memcpy(point_to_centroid_distance, point_to_centroid_distance_buffer.data(),
+                    points_count * sizeof(distance_t));
+        for (std::size_t i = 0; i < wanted_clusters; i++) {
+            byte_t const* quantized_centroid = centroids_quantized_buffer.data() + i * stride_per_vector_quantized;
+            byte_t* centroid = centroids_data + i * centroids_stride_bytes;
+            if (!decompress_points(quantized_centroid, dimensions, centroid))
+                std::memcpy(centroid, quantized_centroid, bytes_per_vector_quantized);
+        }
+
+        return result;
+    }
+};
+
+using kmeans_clustering_t = kmeans_clustering_gt<>;
+
 /**
  *  @brief  C++11 Multi-Hash-Set with Linear Probing.
  *