From b56d283aca333ed4025cd2dad5898acac2334517 Mon Sep 17 00:00:00 2001 From: lkuffo Date: Tue, 13 Oct 2026 13:21:26 +0200 Subject: [PATCH] 2-Means Specialization --- include/superkmeans/two_means.h | 173 ++++++++++++++++++++++++++++++++ tests/CMakeLists.txt | 5 + tests/test_two_means.cpp | 117 +++++++++++++++++++++ 3 files changed, 295 insertions(+) create mode 100644 include/superkmeans/two_means.h create mode 100644 tests/test_two_means.cpp diff --git a/include/superkmeans/two_means.h b/include/superkmeans/two_means.h new file mode 100644 index 0000000..0ed6612 --- /dev/null +++ b/include/superkmeans/two_means.h @@ -0,0 +1,173 @@ +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "superkmeans/allocator.h" +#include "superkmeans/common.h" +#include "superkmeans/executor.h" +#include "superkmeans/profiler.h" + +namespace skmeans { + +/** + * @brief The closer of two centroids to a vector. |vec|^2 is common to both distances, so the + * centroids compare |c|^2 - 2 vec.c, both dot products in one pass over the vector. + */ +inline uint32_t ClosestOfTwoCentroids( + const float* SKM_RESTRICT vec, + const float* SKM_RESTRICT centroids, + const float* centroid_norms, + const size_t d +) { + const float* centroid_0 = centroids; + const float* centroid_1 = centroids + d; + float dot_0 = 0.0f; + float dot_1 = 0.0f; + SKM_VECTORIZE_LOOP + for (size_t j = 0; j < d; ++j) { + dot_0 += vec[j] * centroid_0[j]; + dot_1 += vec[j] * centroid_1[j]; + } + return centroid_norms[0] - 2.0f * dot_0 <= centroid_norms[1] - 2.0f * dot_1 ? 0 : 1; +} + +/** + * @brief k-means with 2 clusters, e.g. to split one cluster. Each iteration is one pass over the + * vectors: a vector picks the closer centroid and is added to its accumulator while still in + * cache, without the GEMM tiles and PDX layouts that make SuperKMeans with 2 clusters slower. + * + * Picks the initial centroids like SuperKMeans (the first two vectors of a seeded shuffle). An + * empty cluster keeps its previous centroid. + * + * @param data Vectors to cluster (n × d, row-major) + * @param n Number of vectors (at least 2) + * @param d Dimensionality + * @param iters Maximum number of iterations + * @param tol Stops once the total squared centroid shift of an iteration falls below it + * @param angular Normalizes the centroids after every update (spherical k-means) + * @param seed Seed of the shuffle that picks the initial centroids + * @param out_centroids Output centroids (2 × d, row-major) + * @param out_assignments Each vector's cluster (0 or 1) for the final centroids, or nullptr + * @param executor Runs each pass over the vectors, one range of vectors per worker + * @param allocator Allocates the buffers sized by n and by the number of workers + */ +inline void TwoMeans( + const float* SKM_RESTRICT data, + const size_t n, + const size_t d, + const uint32_t iters, + const float tol, + const bool angular, + const uint32_t seed, + float* SKM_RESTRICT out_centroids, + uint32_t* SKM_RESTRICT out_assignments, + ParallelExecutor& executor, + MemoryAllocator& allocator +) { + SKM_PROFILE_SCOPE("two_means"); + if (n < 2) { + throw std::invalid_argument("TwoMeans needs at least two vectors"); + } + { + std::mt19937 rng(seed); + auto indices = AllocateBuffer(allocator, n); + for (size_t i = 0; i < n; ++i) { + indices[i] = i; + } + std::shuffle(indices.get(), indices.get() + n, rng); + std::memcpy(out_centroids, data + indices[0] * d, d * sizeof(float)); + std::memcpy(out_centroids + d, data + indices[1] * d, d * sizeof(float)); + } + + const size_t n_workers = executor.NumWorkers(); + auto worker_centroid_accumulators = AllocateBuffer(allocator, n_workers * 2 * d); + std::vector worker_cluster_sizes(n_workers * 2); + float centroid_norms[2]; + const auto compute_centroid_norms = [&]() { + for (size_t ci = 0; ci < 2; ++ci) { + const float* centroid = out_centroids + ci * d; + float sum = 0.0f; + SKM_VECTORIZE_LOOP + for (size_t j = 0; j < d; ++j) { + sum += centroid[j] * centroid[j]; + } + centroid_norms[ci] = sum; + } + }; + + for (uint32_t iter = 0; iter < iters; ++iter) { + compute_centroid_norms(); + std::memset(worker_centroid_accumulators.get(), 0, n_workers * 2 * d * sizeof(float)); + std::fill(worker_cluster_sizes.begin(), worker_cluster_sizes.end(), 0); + executor.ParallelFor(n, [&](size_t begin, size_t end, size_t worker) { + float* centroid_accumulators = worker_centroid_accumulators.get() + worker * 2 * d; + uint32_t* cluster_sizes = worker_cluster_sizes.data() + worker * 2; + for (size_t i = begin; i < end; ++i) { + const float* vec = data + i * d; + const uint32_t ci = ClosestOfTwoCentroids(vec, out_centroids, centroid_norms, d); + float* acc = centroid_accumulators + ci * d; + cluster_sizes[ci] += 1; + SKM_VECTORIZE_LOOP + for (size_t j = 0; j < d; ++j) { + acc[j] += vec[j]; + } + } + }); + + float shift = 0.0f; + for (size_t ci = 0; ci < 2; ++ci) { + // Worker 0's accumulator gathers the others' and becomes the new centroid + float* new_centroid = worker_centroid_accumulators.get() + ci * d; + uint32_t cluster_size = worker_cluster_sizes[ci]; + for (size_t worker = 1; worker < n_workers; ++worker) { + const float* acc = worker_centroid_accumulators.get() + (worker * 2 + ci) * d; + SKM_VECTORIZE_LOOP + for (size_t j = 0; j < d; ++j) { + new_centroid[j] += acc[j]; + } + cluster_size += worker_cluster_sizes[worker * 2 + ci]; + } + if (cluster_size == 0) { + continue; + } + float mult = 1.0f / static_cast(cluster_size); + if (angular) { + float sum = 0.0f; + SKM_VECTORIZE_LOOP + for (size_t j = 0; j < d; ++j) { + sum += new_centroid[j] * new_centroid[j]; + } + mult = 1.0f / std::sqrt(sum); + } + float* centroid = out_centroids + ci * d; + SKM_VECTORIZE_LOOP + for (size_t j = 0; j < d; ++j) { + const float delta = new_centroid[j] * mult - centroid[j]; + shift += delta * delta; + centroid[j] = new_centroid[j] * mult; + } + } + if (shift < tol) { + break; + } + } + + if (out_assignments != nullptr) { + compute_centroid_norms(); + executor.ParallelFor(n, [&](size_t begin, size_t end, size_t) { + for (size_t i = begin; i < end; ++i) { + out_assignments[i] = + ClosestOfTwoCentroids(data + i * d, out_centroids, centroid_norms, d); + } + }); + } +} + +} // namespace skmeans diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 5fa3781..06fa544 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -85,6 +85,9 @@ target_compile_definitions(test_quantized.out PRIVATE CMAKE_SOURCE_DIR="${CMAKE_ add_executable(test_executor.out test_executor.cpp) target_link_libraries(test_executor.out PRIVATE ${TEST_COMMON_LIBS}) +add_executable(test_two_means.out test_two_means.cpp) +target_link_libraries(test_two_means.out PRIVATE ${TEST_COMMON_LIBS}) + set(CMAKE_GTEST_DISCOVER_TESTS_DISCOVERY_MODE PRE_TEST) include(GoogleTest) gtest_discover_tests(test_rotation.out) @@ -99,6 +102,7 @@ gtest_discover_tests(test_quantized_lvq4.out) gtest_discover_tests(test_quantized_rabitq.out) gtest_discover_tests(test_quantized.out) gtest_discover_tests(test_executor.out) +gtest_discover_tests(test_two_means.out) add_custom_target(tests DEPENDS @@ -114,4 +118,5 @@ add_custom_target(tests test_quantized_rabitq.out test_quantized.out test_executor.out + test_two_means.out ) diff --git a/tests/test_two_means.cpp b/tests/test_two_means.cpp new file mode 100644 index 0000000..36c2022 --- /dev/null +++ b/tests/test_two_means.cpp @@ -0,0 +1,117 @@ +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "superkmeans/allocator.h" +#include "superkmeans/executor.h" +#include "superkmeans/two_means.h" + +namespace { + +constexpr size_t N_VECTORS = 1000; +constexpr size_t DIMENSIONS = 32; +constexpr uint32_t ITERS = 25; +constexpr float TOL = 1e-6f; +constexpr uint32_t SEED = 42; + +// Two Gaussian blobs around orthogonal centers: vector i belongs to blob i % 2 +std::vector MakeTwoBlobs() { + std::mt19937 rng(SEED); + std::normal_distribution noise(0.0f, 0.5f); + std::vector data(N_VECTORS * DIMENSIONS); + for (size_t i = 0; i < N_VECTORS; ++i) { + for (size_t j = 0; j < DIMENSIONS; ++j) { + data[i * DIMENSIONS + j] = noise(rng); + } + data[i * DIMENSIONS + i % 2] += 5.0f; + } + return data; +} + +// Each blob becomes one cluster, whose centroid is the mean (angular: the unit mean) of its +// vectors, on any executor. +TEST(TwoMeansTest, SeparatesTwoBlobs) { + const std::vector data = MakeTwoBlobs(); + skmeans::VanillaAllocator allocator; + std::vector> executors; + executors.push_back(std::make_unique()); + executors.push_back(skmeans::MakeDefaultExecutor(4)); + for (auto& executor : executors) { + for (const bool angular : {false, true}) { + SCOPED_TRACE( + "workers=" + std::to_string(executor->NumWorkers()) + + " angular=" + std::to_string(angular) + ); + std::vector centroids(2 * DIMENSIONS); + std::vector assignments(N_VECTORS); + skmeans::TwoMeans( + data.data(), + N_VECTORS, + DIMENSIONS, + ITERS, + TOL, + angular, + SEED, + centroids.data(), + assignments.data(), + *executor, + allocator + ); + ASSERT_NE(assignments[0], assignments[1]); + for (size_t i = 0; i < N_VECTORS; ++i) { + ASSERT_EQ(assignments[i], assignments[i % 2]) << "vector " << i; + } + for (uint32_t ci = 0; ci < 2; ++ci) { + std::vector mean(DIMENSIONS, 0.0); + for (size_t i = 0; i < N_VECTORS; ++i) { + if (assignments[i] != ci) { + continue; + } + for (size_t j = 0; j < DIMENSIONS; ++j) { + mean[j] += data[i * DIMENSIONS + j]; + } + } + double norm = 0.0; + for (size_t j = 0; j < DIMENSIONS; ++j) { + norm += mean[j] * mean[j]; + } + const double mult = angular ? 1.0 / std::sqrt(norm) : 2.0 / N_VECTORS; + for (size_t j = 0; j < DIMENSIONS; ++j) { + EXPECT_NEAR(centroids[ci * DIMENSIONS + j], mean[j] * mult, 1e-4) + << "centroid " << ci << " dimension " << j; + } + } + } + } +} + +TEST(TwoMeansTest, FewerThanTwoVectorsThrows) { + const std::vector data = MakeTwoBlobs(); + skmeans::VanillaAllocator allocator; + skmeans::SerialExecutor executor; + std::vector centroids(2 * DIMENSIONS); + EXPECT_THROW( + skmeans::TwoMeans( + data.data(), + 1, + DIMENSIONS, + ITERS, + TOL, + false, + SEED, + centroids.data(), + nullptr, + executor, + allocator + ), + std::invalid_argument + ); +} + +} // namespace